03-training-formalization.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "33b0bbba",
   "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",
    "1. Receive hyperparameters using `hyperparam_gen` custom Python component.\n",
    "2. Extract data from BigQuery using `BigQueryExampleGen` component.\n",
    "3. Validate the raw data using `StatisticsGen` and `ExampleValidator` components.\n",
    "4. Process the data using `Transform` component.\n",
    "5. Train a custom model using `Trainer` component.\n",
    "7. Evaluate and Validate the custom model using `ModelEvaluator` component.\n",
    "7. Save the blessed to model registry location using `Pusher` component.\n",
    "8. Upload the model to Vertex AI using `vertex_model_pusher` custom Python component\n",
    "\n",
    "The custom components are implemented in the [tfx_pipeline/components.py](tfx_pipeline/components) module."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "be9f73ad",
   "metadata": {},
   "source": [
    "## Setup"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2e9788da",
   "metadata": {},
   "source": [
    "### Import libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a419a0b3",
   "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",
    "logging.getLogger().setLevel(logging.ERROR)\n",
    "tf.get_logger().setLevel('ERROR')\n",
    "\n",
    "print(\"TFX Version:\", tfx.__version__)\n",
    "print(\"Tensorflow Version:\", tf.__version__)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "239e86df",
   "metadata": {},
   "source": [
    "### Setup Google Cloud project"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "993bb57e",
   "metadata": {},
   "outputs": [],
   "source": [
    "PROJECT = '[your-project-id]' # Change to your project id.\n",
    "REGION = 'europe-west1' # Change to your region.\n",
    "BUCKET =  '[your-bucket-name]' # Change to your bucket name.\n",
    "SERVICE_ACCOUNT = \"[your-service-account]\"\n",
    "\n",
    "if PROJECT == \"\" or PROJECT is None or PROJECT == \"[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 = 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",
    "    # Get your bucket name to GCP project id\n",
    "    BUCKET = PROJECT\n",
    "    # Try to create the bucket if it doesn't exists\n",
    "    ! gsutil mb -l $REGION gs://$BUCKET\n",
    "    print(\"\")\n",
    "    \n",
    "PARENT = f\"projects/{PROJECT}/locations/{REGION}\"\n",
    "    \n",
    "print(\"Project ID:\", PROJECT)\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": "b960e578",
   "metadata": {},
   "source": [
    "### Set configurations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "05e18592",
   "metadata": {},
   "outputs": [],
   "source": [
    "VERSION = 'v01'\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_interactive')\n",
    "MODEL_REGISTRY = os.path.join(WORKSPACE, 'model_registry')\n",
    "PIPELINE_NAME = f'{MODEL_DISPLAY_NAME}-train-pipeline'\n",
    "PIPELINE_ROOT = os.path.join(ARTIFACT_STORE, PIPELINE_NAME)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a7d23048",
   "metadata": {},
   "source": [
    "## Create Interactive Context"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c70d501f",
   "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": "c6f5f232",
   "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": "b9aeaf26",
   "metadata": {},
   "source": [
    "## 1. Hyperparameter generation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9edebd29",
   "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": "06bffdc0",
   "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": "3bac1aa8",
   "metadata": {},
   "source": [
    "## 2. Data extraction"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "67173d8b",
   "metadata": {},
   "outputs": [],
   "source": [
    "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"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "88c32135",
   "metadata": {},
   "source": [
    "### Extract train and eval splits"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "972a4e86",
   "metadata": {},
   "outputs": [],
   "source": [
    "sql_query = datasource_utils.get_training_source_query(\n",
    "    PROJECT, REGION, DATASET_DISPLAY_NAME, ml_use='UNASSIGNED', limit=5000)\n",
    "\n",
    "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(query=sql_query, output_config=output_config)\n",
    "\n",
    "beam_pipeline_args=[\n",
    "    f\"--project={PROJECT}\",\n",
    "    f\"--temp_location={os.path.join(WORKSPACE, 'tmp')}\"\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": "ea0e6124",
   "metadata": {},
   "source": [
    "### Extract test split"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5436acbc",
   "metadata": {},
   "outputs": [],
   "source": [
    "sql_query = datasource_utils.get_training_source_query(\n",
    "    PROJECT, REGION, DATASET_DISPLAY_NAME, ml_use='TEST', limit=1000)\n",
    "\n",
    "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(query=sql_query, output_config=output_config)\n",
    "\n",
    "beam_pipeline_args=[\n",
    "    f\"--project={PROJECT}\",\n",
    "    f\"--temp_location={os.path.join(WORKSPACE, 'tmp')}\"\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": "261e271b",
   "metadata": {},
   "source": [
    "### Read sample extract tfrecords"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "201664c8",
   "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": "230eeb07",
   "metadata": {},
   "source": [
    "## 3. Data validation"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d2c1f758",
   "metadata": {},
   "source": [
    "### Import raw schema"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f54cec4a",
   "metadata": {},
   "outputs": [],
   "source": [
    "schema_importer = tfx.dsl.components.common.importer.Importer(\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": "136d69ce",
   "metadata": {},
   "source": [
    "### Generate statistics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "62a346b5",
   "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": "e90b160c",
   "metadata": {},
   "outputs": [],
   "source": [
    "!rm -r {RAW_SCHEMA_DIR}/.ipynb_checkpoints/"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7fcfa0b4",
   "metadata": {},
   "source": [
    "### Validate statistics against schema"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "831154fc",
   "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": "2441243b",
   "metadata": {},
   "outputs": [],
   "source": [
    "context.show(example_validator.outputs['anomalies'])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a31ff6af",
   "metadata": {},
   "source": [
    "## 4. Data transformation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "047a783c",
   "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": "265371cd",
   "metadata": {},
   "source": [
    "### Read sample transformed tfrecords"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f0cee0a6",
   "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": "c2fd431a",
   "metadata": {},
   "source": [
    "## 5. Model training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c023e9cf",
   "metadata": {},
   "outputs": [],
   "source": [
    "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": "89595398",
   "metadata": {},
   "source": [
    "### Get the latest model to warm start"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c261be21",
   "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": "ab0b69e7",
   "metadata": {},
   "source": [
    "### Train the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3606d40e",
   "metadata": {},
   "outputs": [],
   "source": [
    "_train_module_file = 'src/model_training/runner.py'\n",
    "\n",
    "trainer = tfx.components.Trainer(\n",
    "    module_file=_train_module_file,\n",
    "    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",
    "    hyperparameters=hyperparams_gen.outputs['hyperparameters'],\n",
    ")\n",
    "\n",
    "context.run(trainer, enable_cache=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1001fc8c",
   "metadata": {},
   "source": [
    "## 6. Model evaluation"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "12cba33d",
   "metadata": {},
   "source": [
    "### Get the latest blessed model for model validation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c3164f3f",
   "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": "766a7478",
   "metadata": {},
   "source": [
    "### Evaluate and validate the model against the baseline model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dc718215",
   "metadata": {},
   "outputs": [],
   "source": [
    "from tfx.components import Evaluator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1d67d478",
   "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 will be ignored if there is no\n",
    "                        # baseline model resolved from MLMD (first run).\n",
    "                        change_threshold=tfma.GenericChangeThreshold(\n",
    "                            direction=tfma.MetricDirection.HIGHER_IS_BETTER,\n",
    "                            absolute={'value': -1e-10}))),\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": "e3e939ef",
   "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": "dcbe53b9",
   "metadata": {},
   "source": [
    "## 7. Model pushing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "410efc86",
   "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": "09239795",
   "metadata": {},
   "source": [
    "## 8. Model Upload to Vertex AI"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3b98410b",
   "metadata": {},
   "outputs": [],
   "source": [
    "serving_runtime = 'tf2-cpu.2-5'\n",
    "serving_image_uri = f\"us-docker.pkg.dev/vertex-ai/prediction/{serving_runtime}:latest\"\n",
    "\n",
    "labels = {\n",
    "    'dataset_name': DATASET_DISPLAY_NAME,\n",
    "    'pipeline_name': PIPELINE_NAME\n",
    "}\n",
    "labels = json.dumps(labels)\n",
    "\n",
    "vertex_model_uploader = components.vertex_model_uploader(\n",
    "    project=PROJECT,\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",
    "    model_blessing=evaluator.outputs['blessing'],\n",
    "    explanation_config='',\n",
    "    labels=labels\n",
    ")\n",
    "\n",
    "context.run(vertex_model_uploader, enable_cache=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fe6c3107",
   "metadata": {},
   "outputs": [],
   "source": [
    "vertex_model_uploader.outputs['uploaded_model'].get()[0].get_string_custom_property('model_uri')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "378b5c18",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "environment": {
   "kernel": "python3",
   "name": "common-cpu.m87",
   "type": "gcloud",
   "uri": "gcr.io/deeplearning-platform-release/base-cpu:m87"
  },
  "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.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}