03-training-formalization.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "2cfc650e",
   "metadata": {},
   "source": [
    "# 03 - `TFX` interactive training pipeline execution\n",
    "\n",
    "The purpose of this notebook is to interactively run the following `TFX` pipeline steps:\n",
    "\n",
    "1. Receive hyperparameters using `hyperparam_gen` custom Python component.\n",
    "2. Extract data from `BigQuery` using the `BigQueryExampleGen` component.\n",
    "3. Validate the raw data using the `StatisticsGen` and `ExampleValidator` components.\n",
    "4. Process the data using the `Transform` component.\n",
    "5. Train a custom model using the `Trainer` component.\n",
    "7. Evaluate and validate the custom model using the `ModelEvaluator` component.\n",
    "7. Save the blessed to the model registry location using the `Pusher` component.\n",
    "8. Upload the blessed model as a `Vertex Model` resource using the `vertex_model_pusher` custom Python component\n",
    "\n",
    "The TFX custom Python function components are implemented in the [tfx_pipeline/components.py](tfx_pipeline/components) module.\n",
    "\n",
    "Learn about [TFX](https://www.tensorflow.org/tfx/) and [Custom Python function components](https://www.tensorflow.org/tfx/guide/custom_function_component)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e97f1df8",
   "metadata": {},
   "source": [
    "## Setup"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "89b55b85",
   "metadata": {},
   "source": [
    "### Import libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "85a0d1b8",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import json\n",
    "import numpy as np\n",
    "import tfx\n",
    "import tensorflow as tf\n",
    "import tensorflow_transform as tft\n",
    "import tensorflow_data_validation as tfdv\n",
    "import tensorflow_model_analysis as tfma\n",
    "from tensorflow_transform.tf_metadata import schema_utils\n",
    "import logging\n",
    "\n",
    "from src.common import features\n",
    "from src.model_training import data\n",
    "from src.tfx_pipelines import components\n",
    "\n",
    "from src.common import datasource_utils\n",
    "from tfx.extensions.google_cloud_big_query.example_gen.component import BigQueryExampleGen\n",
    "from tfx.proto import example_gen_pb2, transform_pb2\n",
    "\n",
    "logging.getLogger().setLevel(logging.ERROR)\n",
    "tf.get_logger().setLevel('ERROR')\n",
    "\n",
    "print('TensorFlow:', tf.__version__)\n",
    "print('TFX:', tfx.__version__)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a682184c",
   "metadata": {},
   "source": [
    "### Setup Google Cloud project"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "941c59e7",
   "metadata": {},
   "outputs": [],
   "source": [
    "PROJECT_ID = '[your-project-id]' # Change to your project id.\n",
    "REGION = 'us-central1' # Change to your region.\n",
    "BUCKET = '[your-bucket-name]'  # Change to your bucket name.\n",
    "SERVICE_ACCOUNT = '[your-service-account]'\n",
    "\n",
    "if PROJECT_ID == '' or PROJECT_ID is None or PROJECT_ID == '[your-project-id]':\n",
    "    # Get your GCP project id from gcloud\n",
    "    shell_output = !gcloud config list --format 'value(core.project)' 2>/dev/null\n",
    "    PROJECT_ID = shell_output[0]\n",
    "    \n",
    "if SERVICE_ACCOUNT == '' or SERVICE_ACCOUNT is None or SERVICE_ACCOUNT == '[your-service-account]':\n",
    "    # Get your GCP project id from gcloud\n",
    "    shell_output = !gcloud config list --format 'value(core.account)' 2>/dev/null\n",
    "    SERVICE_ACCOUNT = shell_output[0]\n",
    "    \n",
    "if BUCKET == '' or BUCKET is None or BUCKET == '[your-bucket-name]':\n",
    "    # Set your bucket name using your GCP project id\n",
    "    BUCKET = PROJECT_ID\n",
    "    # Try to create the bucket if it doesn'exists\n",
    "    ! gsutil mb -l $REGION gs://$BUCKET\n",
    "    print('')\n",
    "    \n",
    "PARENT = f'projects/{PROJECT_ID}/locations/{REGION}'\n",
    "    \n",
    "print('Project ID:', PROJECT_ID)\n",
    "print('Region:', REGION)\n",
    "print('Bucket name:', BUCKET)\n",
    "print('Service Account:', SERVICE_ACCOUNT)\n",
    "print('Vertex API Parent URI:', PARENT)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3447bb1e",
   "metadata": {},
   "source": [
    "### Set configurations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8b1a6f6f",
   "metadata": {},
   "outputs": [],
   "source": [
    "VERSION = 'v1'\n",
    "DATASET_DISPLAY_NAME = 'chicago-taxi-tips'\n",
    "MODEL_DISPLAY_NAME = f'{DATASET_DISPLAY_NAME}-classifier-{VERSION}'\n",
    "\n",
    "WORKSPACE = f'gs://{BUCKET}/{DATASET_DISPLAY_NAME}'\n",
    "RAW_SCHEMA_DIR = 'src/raw_schema'\n",
    "\n",
    "MLMD_SQLLITE = 'mlmd.sqllite'\n",
    "ARTIFACT_STORE = os.path.join(WORKSPACE, 'tfx_artifacts')\n",
    "MODEL_REGISTRY = os.path.join(WORKSPACE, 'model_registry')\n",
    "PIPELINE_NAME = f'{MODEL_DISPLAY_NAME}-training-pipeline'\n",
    "PIPELINE_ROOT = os.path.join(ARTIFACT_STORE, PIPELINE_NAME)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "074295d5",
   "metadata": {},
   "source": [
    "## Create an interactive context"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5d0bdd06",
   "metadata": {},
   "outputs": [],
   "source": [
    "REMOVE_ARTIFACTS = True\n",
    "\n",
    "if tf.io.gfile.exists(ARTIFACT_STORE) and REMOVE_ARTIFACTS:\n",
    "    print('Removing previous artifacts...')\n",
    "    tf.io.gfile.rmtree(ARTIFACT_STORE)\n",
    "    \n",
    "if tf.io.gfile.exists(MLMD_SQLLITE) and REMOVE_ARTIFACTS:\n",
    "    print('Deleting previous mlmd.sqllite...')\n",
    "    tf.io.gfile.rmtree(MLMD_SQLLITE)\n",
    "    \n",
    "print(f'Pipeline artifacts directory: {PIPELINE_ROOT}')\n",
    "print(f'Local metadata SQLlit path: {MLMD_SQLLITE}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "35e4bab6",
   "metadata": {},
   "outputs": [],
   "source": [
    "import ml_metadata as mlmd\n",
    "from ml_metadata.proto import metadata_store_pb2\n",
    "from tfx.orchestration.experimental.interactive.interactive_context import InteractiveContext\n",
    "\n",
    "connection_config = metadata_store_pb2.ConnectionConfig()\n",
    "connection_config.sqlite.filename_uri = MLMD_SQLLITE\n",
    "connection_config.sqlite.connection_mode = 3 # READWRITE_OPENCREATE\n",
    "mlmd_store = mlmd.metadata_store.MetadataStore(connection_config)\n",
    "\n",
    "context = InteractiveContext(\n",
    "  pipeline_name=PIPELINE_NAME,\n",
    "  pipeline_root=PIPELINE_ROOT,\n",
    "  metadata_connection_config=connection_config\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "982ab3a4",
   "metadata": {},
   "source": [
    "## 1. Generate the hyperparameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "160caff0",
   "metadata": {},
   "outputs": [],
   "source": [
    "hyperparams_gen = components.hyperparameters_gen(\n",
    "    num_epochs=5,\n",
    "    learning_rate=0.001,\n",
    "    batch_size=512,\n",
    "    hidden_units='64,64',\n",
    ")\n",
    "\n",
    "context.run(hyperparams_gen, enable_cache=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7de56caf",
   "metadata": {},
   "outputs": [],
   "source": [
    "json.load(\n",
    "    tf.io.gfile.GFile(\n",
    "        os.path.join(\n",
    "            hyperparams_gen.outputs.hyperparameters.get()[0].uri, 'hyperparameters.json')\n",
    "    )\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "57f60939",
   "metadata": {},
   "source": [
    "## 2. Extract data"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0a06ad47",
   "metadata": {},
   "source": [
    "### Extract the train and eval splits"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7d5e6d33",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_sql_query = datasource_utils.get_training_source_query(\n",
    "    project=PROJECT_ID, \n",
    "    region=REGION, \n",
    "    dataset_display_name=DATASET_DISPLAY_NAME, \n",
    "    ml_use='UNASSIGNED',\n",
    "    limit=5000\n",
    ")\n",
    "\n",
    "train_output_config = example_gen_pb2.Output(\n",
    "    split_config=example_gen_pb2.SplitConfig(\n",
    "        splits=[\n",
    "            example_gen_pb2.SplitConfig.Split(name='train', hash_buckets=4),\n",
    "            example_gen_pb2.SplitConfig.Split(name='eval', hash_buckets=1),\n",
    "        ]\n",
    "    )\n",
    ")\n",
    "\n",
    "train_example_gen = BigQueryExampleGen(\n",
    "    query=train_sql_query, \n",
    "    output_config=train_output_config\n",
    ")\n",
    "\n",
    "beam_pipeline_args=[\n",
    "    f'--project={PROJECT_ID}',\n",
    "    f'--temp_location={WORKSPACE}/temp'\n",
    "]\n",
    "\n",
    "context.run(\n",
    "    train_example_gen,\n",
    "    beam_pipeline_args=beam_pipeline_args,\n",
    "    enable_cache=False\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e7796bf5",
   "metadata": {},
   "source": [
    "### Extract the test split"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4deba492",
   "metadata": {},
   "outputs": [],
   "source": [
    "test_sql_query = datasource_utils.get_training_source_query(\n",
    "    project=PROJECT_ID, \n",
    "    region=REGION, \n",
    "    dataset_display_name=DATASET_DISPLAY_NAME, \n",
    "    ml_use='test',\n",
    "    limit=1000\n",
    ")\n",
    "\n",
    "test_output_config = example_gen_pb2.Output(\n",
    "    split_config=example_gen_pb2.SplitConfig(\n",
    "        splits=[\n",
    "            example_gen_pb2.SplitConfig.Split(name='test', hash_buckets=1),\n",
    "        ]\n",
    "    )\n",
    ")\n",
    "\n",
    "test_example_gen = BigQueryExampleGen(\n",
    "    query=test_sql_query, \n",
    "    output_config=test_output_config\n",
    ")\n",
    "\n",
    "beam_pipeline_args=[\n",
    "    f'--project={PROJECT_ID}',\n",
    "    f'--temp_location={WORKSPACE}/temp'\n",
    "]\n",
    "\n",
    "context.run(\n",
    "    test_example_gen,\n",
    "    beam_pipeline_args=beam_pipeline_args,\n",
    "    enable_cache=False\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "414788d7",
   "metadata": {},
   "source": [
    "### Read some sample extracted TFRecords"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f2d2af73",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_uri = os.path.join(train_example_gen.outputs.examples.get()[0].uri, 'Split-train/*')\n",
    "source_raw_schema = tfdv.load_schema_text(os.path.join(RAW_SCHEMA_DIR, 'schema.pbtxt'))\n",
    "raw_feature_spec = schema_utils.schema_as_feature_spec(source_raw_schema).feature_spec\n",
    "\n",
    "def _parse_tf_example(tfrecord):\n",
    "    return tf.io.parse_single_example(tfrecord, raw_feature_spec)\n",
    "\n",
    "tfrecord_filenames = tf.data.Dataset.list_files(train_uri)\n",
    "dataset = tf.data.TFRecordDataset(tfrecord_filenames, compression_type='GZIP')\n",
    "dataset = dataset.map(_parse_tf_example)\n",
    "\n",
    "for raw_features in dataset.shuffle(1000).batch(3).take(1):\n",
    "    for key in raw_features:\n",
    "        print(f'{key}: {np.squeeze(raw_features[key], -1)}')\n",
    "    print('')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d45fd3cc",
   "metadata": {},
   "source": [
    "## 3. Validate the data"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7b22b1ed",
   "metadata": {},
   "source": [
    "### Import the raw schema"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b9fc2db4",
   "metadata": {},
   "outputs": [],
   "source": [
    "schema_importer = tfx.components.common_nodes.importer_node.ImporterNode(\n",
    "    source_uri=RAW_SCHEMA_DIR,\n",
    "    artifact_type=tfx.types.standard_artifacts.Schema,\n",
    "    reimport=False\n",
    ")\n",
    "\n",
    "context.run(schema_importer)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f2328e26",
   "metadata": {},
   "source": [
    "### Generate statistics for the dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7fbf4219",
   "metadata": {},
   "outputs": [],
   "source": [
    "statistics_gen = tfx.components.StatisticsGen(\n",
    "    examples=train_example_gen.outputs.examples)\n",
    "context.run(statistics_gen)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d30e6e5e",
   "metadata": {},
   "outputs": [],
   "source": [
    "!rm -r {RAW_SCHEMA_DIR}/.ipynb_checkpoints/"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "98a4ee9b",
   "metadata": {},
   "source": [
    "### Validate statistics against the schema"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cbca5969",
   "metadata": {},
   "outputs": [],
   "source": [
    "example_validator = tfx.components.ExampleValidator(\n",
    "    statistics=statistics_gen.outputs.statistics,\n",
    "    schema=schema_importer.outputs.result,\n",
    ")\n",
    "\n",
    "context.run(example_validator)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1aff8087",
   "metadata": {},
   "outputs": [],
   "source": [
    "context.show(example_validator.outputs.anomalies)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d421e2c7",
   "metadata": {},
   "source": [
    "## 4. Transform the data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a9f2fbd3",
   "metadata": {},
   "outputs": [],
   "source": [
    "_transform_module_file = 'src/preprocessing/transformations.py'\n",
    "\n",
    "transform = tfx.components.Transform(\n",
    "    examples=train_example_gen.outputs.examples,\n",
    "    schema=schema_importer.outputs.result,\n",
    "    module_file=_transform_module_file,\n",
    "    splits_config=transform_pb2.SplitsConfig(\n",
    "        analyze=['train'], transform=['train', 'eval']),\n",
    ")\n",
    "\n",
    "context.run(transform, enable_cache=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "77048aee",
   "metadata": {},
   "source": [
    "### Read a sample of transformed TFrecords"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fbd04f3c",
   "metadata": {},
   "outputs": [],
   "source": [
    "transformed_train_uri = os.path.join(transform.outputs.transformed_examples.get()[0].uri, 'Split-train/*')\n",
    "transform_graph_uri = transform.outputs.transform_graph.get()[0].uri\n",
    "\n",
    "tft_output = tft.TFTransformOutput(transform_graph_uri)\n",
    "transform_feature_spec = tft_output.transformed_feature_spec()\n",
    "\n",
    "for input_features, target in data.get_dataset(\n",
    "    transformed_train_uri, transform_feature_spec, batch_size=3).take(1):\n",
    "    for key in input_features:\n",
    "        print(f'{key} ({input_features[key].dtype}): {input_features[key].numpy().tolist()}')\n",
    "    print(f'target: {target.numpy().tolist()}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "15a98bd9",
   "metadata": {},
   "source": [
    "## 5. Train the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0c4eada2",
   "metadata": {},
   "outputs": [],
   "source": [
    "from tfx.components.base import executor_spec\n",
    "from tfx.components.trainer import executor as trainer_executor\n",
    "from tfx.dsl.components.common.resolver import Resolver\n",
    "from tfx.dsl.experimental import latest_artifacts_resolver\n",
    "from tfx.dsl.experimental import latest_blessed_model_resolver"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7a61aebf",
   "metadata": {},
   "source": [
    "### Get the latest model to warm start the training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "12269fa3",
   "metadata": {},
   "outputs": [],
   "source": [
    "latest_model_resolver = Resolver(\n",
    "    strategy_class=latest_artifacts_resolver.LatestArtifactsResolver,\n",
    "    latest_model=tfx.types.Channel(type=tfx.types.standard_artifacts.Model)\n",
    ")\n",
    "\n",
    "context.run(latest_model_resolver, enable_cache=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "15c7bcf0",
   "metadata": {},
   "source": [
    "### Train the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "34c66a3f",
   "metadata": {},
   "outputs": [],
   "source": [
    "_train_module_file = 'src/model_training/runner.py'\n",
    "\n",
    "trainer = tfx.components.Trainer(\n",
    "    custom_executor_spec=executor_spec.ExecutorClassSpec(trainer_executor.GenericExecutor),\n",
    "    module_file=_train_module_file,\n",
    "    transformed_examples=transform.outputs.transformed_examples,\n",
    "    schema=schema_importer.outputs.result,\n",
    "    base_model=latest_model_resolver.outputs.latest_model,\n",
    "    transform_graph=transform.outputs.transform_graph,\n",
    "    train_args=tfx.proto.trainer_pb2.TrainArgs(num_steps=0),\n",
    "    eval_args=tfx.proto.trainer_pb2.EvalArgs(num_steps=None),\n",
    "    hyperparameters=hyperparams_gen.outputs.hyperparameters,\n",
    ")\n",
    "\n",
    "context.run(trainer, enable_cache=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c7d45b19",
   "metadata": {},
   "source": [
    "## 6. Evaluate the model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e6d3dc22",
   "metadata": {},
   "source": [
    "### Get the latest blessed model for model validation as the baseline model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a9fe863e",
   "metadata": {},
   "outputs": [],
   "source": [
    "blessed_model_resolver = Resolver(\n",
    "    strategy_class=latest_blessed_model_resolver.LatestBlessedModelResolver,\n",
    "    model=tfx.types.Channel(type=tfx.types.standard_artifacts.Model),\n",
    "    model_blessing=tfx.types.Channel(type=tfx.types.standard_artifacts.ModelBlessing)\n",
    ")\n",
    "\n",
    "context.run(blessed_model_resolver, enable_cache=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e0006fd3",
   "metadata": {},
   "source": [
    "### Evaluate and validate the model against the baseline model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eadd9842",
   "metadata": {},
   "outputs": [],
   "source": [
    "from tfx.components import Evaluator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3e79529a",
   "metadata": {},
   "outputs": [],
   "source": [
    "eval_config = tfma.EvalConfig(\n",
    "    model_specs=[\n",
    "        tfma.ModelSpec(\n",
    "            signature_name='serving_tf_example',\n",
    "            label_key=features.TARGET_FEATURE_NAME,\n",
    "            prediction_key='probabilities')\n",
    "    ],\n",
    "    slicing_specs=[\n",
    "        tfma.SlicingSpec(),\n",
    "    ],\n",
    "    metrics_specs=[\n",
    "        tfma.MetricsSpec(\n",
    "            metrics=[   \n",
    "                tfma.MetricConfig(class_name='ExampleCount'),\n",
    "                tfma.MetricConfig(\n",
    "                    class_name='BinaryAccuracy',\n",
    "                    threshold=tfma.MetricThreshold(\n",
    "                        value_threshold=tfma.GenericValueThreshold(\n",
    "                            lower_bound={'value': 0.8}),\n",
    "                        change_threshold=tfma.GenericChangeThreshold(\n",
    "                            direction=tfma.MetricDirection.HIGHER_IS_BETTER,\n",
    "                            absolute={'value': -1e-10})\n",
    "                    )),\n",
    "        ])\n",
    "    ])\n",
    "\n",
    "\n",
    "evaluator = Evaluator(\n",
    "    examples=test_example_gen.outputs.examples,\n",
    "    example_splits=['test'],\n",
    "    model=trainer.outputs.model,\n",
    "    baseline_model=blessed_model_resolver.outputs.model,\n",
    "    eval_config=eval_config,\n",
    "    schema=schema_importer.outputs.result\n",
    ")\n",
    "\n",
    "context.run(evaluator, enable_cache=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "32720329",
   "metadata": {},
   "outputs": [],
   "source": [
    "evaluation_results = evaluator.outputs.evaluation.get()[0].uri\n",
    "print('validation_ok:', tfma.load_validation_result(evaluation_results).validation_ok, '\\n')\n",
    "\n",
    "for entry in list(tfma.load_metrics(evaluation_results))[0].metric_keys_and_values:\n",
    "    value = entry.value.double_value.value\n",
    "    if value:\n",
    "        print(entry.key.name, ':', round(entry.value.double_value.value, 3))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cf6981b3",
   "metadata": {},
   "source": [
    "## 7. Push the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "88b49413",
   "metadata": {},
   "outputs": [],
   "source": [
    "exported_model_location = os.path.join(MODEL_REGISTRY, MODEL_DISPLAY_NAME)\n",
    "\n",
    "push_destination=tfx.proto.pusher_pb2.PushDestination(\n",
    "    filesystem=tfx.proto.pusher_pb2.PushDestination.Filesystem(\n",
    "        base_directory=exported_model_location,\n",
    "    )\n",
    ")\n",
    "\n",
    "pusher = tfx.components.Pusher(\n",
    "    model=trainer.outputs.model,\n",
    "    model_blessing=evaluator.outputs.blessing,\n",
    "    push_destination=push_destination\n",
    ")\n",
    "\n",
    "context.run(pusher, enable_cache=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7667d363",
   "metadata": {},
   "source": [
    "## 8. Upload the model as a `Vertex Model` resource"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8997d0dd",
   "metadata": {},
   "outputs": [],
   "source": [
    "serving_runtime = 'tf2-cpu.2-4'\n",
    "serving_image_uri = f'gcr.io/cloud-aiplatform/prediction/{serving_runtime}:latest'\n",
    "\n",
    "vertex_model_uploader = components.vertex_model_uploader(\n",
    "    project=PROJECT_ID,\n",
    "    region=REGION,\n",
    "    model_display_name=MODEL_DISPLAY_NAME,\n",
    "    pushed_model_location=exported_model_location,\n",
    "    serving_image_uri=serving_image_uri,\n",
    "    explanation_config=''\n",
    ")\n",
    "\n",
    "context.run(vertex_model_uploader, enable_cache=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d7500e3a",
   "metadata": {},
   "outputs": [],
   "source": [
    "vertex_model_uploader.outputs.uploaded_model.get()[0].get_string_custom_property('model_uri')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dcdcd889",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "environment": {
   "name": "common-cpu.m73",
   "type": "gcloud",
   "uri": "gcr.io/deeplearning-platform-release/base-cpu:m73"
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}