friday updates

notebooks/community/ml_ops/stage2/get_started_vertex_tensorboard.ipynb

@@ -130,6 +130,7 @@
         "    ! pip3 install -U tensorflow-transform==1.2 $USER_FLAG\n",
         "    ! pip3 install -U tensorflow-io==0.18 $USER_FLAG\n",
         "    ! pip3 install --upgrade google-cloud-aiplatform[tensorboard] $USER_FLAG\n",
+        "    ! pip3 install --upgrade google-cloud-pipeline-components $USER_FLAG\n",
         "    ! pip3 install --upgrade google-cloud-bigquery $USER_FLAG\n",
         "    ! pip3 install --upgrade google-cloud-logging $USER_FLAG\n",
         "    ! pip3 install --upgrade apache-beam[gcp] $USER_FLAG\n",
@@ -232,7 +233,7 @@
         "\n",
         "You may not use a multi-regional bucket for training with Vertex AI. Not all regions provide support for all Vertex AI services.\n",
         "\n",
-        "Learn more about [Vertex AI regions](https://cloud.google.com/vertex-ai/docs/general/locations)"
+        "Learn more about [Vertex AI regions](https://cloud.google.com/vertex-ai/docs/general/locations)."
       ]
     },
     {
@@ -1107,6 +1108,28 @@
         "Alternatively, you can navigate to the Experiments tab and view the list of all experiments. Your experiment will have the same name as the training job."
       ]
     },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "delete_tensorboard"
+      },
+      "source": [
+        "### Delete the TensorBoard instance\n",
+        "\n",
+        "Next, delete the TensorBoard instance."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "delete_tensorboard"
+      },
+      "outputs": [],
+      "source": [
+        "tensorboard.delete()"
+      ]
+    },
     {
       "cell_type": "markdown",
       "metadata": {

notebooks/community/ml_ops/stage2/mlops_experimentation.ipynb

@@ -783,12 +783,15 @@
       "source": [
         "import json\n",
         "\n",
-        "with tf.io.gfile.GFile(\n",
-        "    \"gs://\" + dataset.labels[\"user_metadata\"] + \"/metadata.jsonl\", \"r\"\n",
-        ") as f:\n",
-        "    metadata = json.load(f)\n",
+        "try:\n",
+        "    with tf.io.gfile.GFile(\n",
+        "        \"gs://\" + dataset.labels[\"user_metadata\"] + \"/metadata.jsonl\", \"r\"\n",
+        "    ) as f:\n",
+        "        metadata = json.load(f)\n",
         "\n",
-        "print(metadata)"
+        "    print(metadata)\n",
+        "except:\n",
+        "    print(\"no metadata\")"
       ]
     },
     {
@@ -1164,7 +1167,7 @@
         "    display_name=\"chicago_\" + TIMESTAMP,\n",
         "    artifact_uri=MODEL_DIR,\n",
         "    serving_container_image_uri=DEPLOY_IMAGE,\n",
-        "    labels={\"base_model\": 1},\n",
+        "    labels={\"base_model\": \"1\"},\n",
         "    sync=True,\n",
         ")"
       ]
@@ -1213,7 +1216,7 @@
         "setup_cfg = \"[egg_info]\\n\\ntag_build =\\n\\ntag_date = 0\"\n",
         "! echo \"$setup_cfg\" > custom/setup.cfg\n",
         "\n",
-        "setup_py = \"import setuptools\\n\\nsetuptools.setup(\\n\\n    install_requires=[\\n\\n        'google-cloud-aiplatform',\\n\\n        'cloudml-hypertune',\\n\\n        'tensorflow_datasets==1.3.0',\\n\\n    ],\\n\\n    packages=setuptools.find_packages())\"\n",
+        "setup_py = \"import setuptools\\n\\nsetuptools.setup(\\n\\n    install_requires=[\\n\\n        'google-cloud-aiplatform',\\n\\n        'cloudml-hypertune',\\n\\n        'tensorflow_datasets==1.3.0',\\n\\n        'tensorflow_data_validation==1.2',\\n\\n    ],\\n\\n    packages=setuptools.find_packages())\"\n",
         "! echo \"$setup_py\" > custom/setup.py\n",
         "\n",
         "pkg_info = \"Metadata-Version: 1.0\\n\\nName: Chicago Taxi tabular binary classifier\\n\\nVersion: 0.0.0\\n\\nSummary: Demostration training script\\n\\nHome-page: www.google.com\\n\\nAuthor: Google\\n\\nAuthor-email: cdpe@google.com\\n\\nLicense: Public\\n\\nDescription: Demo\\n\\nPlatform: Vertex AI\"\n",
@@ -1415,13 +1418,16 @@
         "        batch_size=hyperparams[\"batch_size\"],\n",
         "    )\n",
         "\n",
-        "    tensorboard = tf.keras.callbacks.TensorBoard(log_dir=log_dir)\n",
-        "\n",
         "    early_stop = tf.keras.callbacks.EarlyStopping(\n",
         "        monitor=hyperparams[\"early_stop\"][\"monitor\"], patience=hyperparams[\"early_stop\"][\"patience\"], restore_best_weights=True\n",
         "    )\n",
         "\n",
-        "    callbacks=[tensorboard, early_stop]\n",
+        "    callbacks=[early_stop]\n",
+        "\n",
+        "    if log_dir:\n",
+        "        tensorboard = tf.keras.callbacks.TensorBoard(log_dir=log_dir)\n",
+        "\n",
+        "        callbacks=callbacks.append(tensorboard)\n",
         "\n",
         "    if tuning:\n",
         "        # Instantiate the HyperTune reporting object\n",
@@ -1659,7 +1665,7 @@
         "from trainer import data\n",
         "from trainer import model as model_\n",
         "from trainer import train\n",
-        "\n",
+        "from trainer import serving\n",
         "\n",
         "parser = argparse.ArgumentParser()\n",
         "parser.add_argument('--model-dir', dest='model_dir',\n",
@@ -1683,7 +1689,7 @@
         "parser.add_argument('--distribute', dest='distribute', type=str, default='single',\n",
         "                    help='distributed training strategy')\n",
         "parser.add_argument('--tensorboard-log-dir', dest='tensorboard_log_dir',\n",
-        "                    default='/tmp/logs', type=str,\n",
+        "                    default=os.getenv('AIP_TENSORBOARD_LOG_DIR'), type=str,\n",
         "                    help='Output file for tensorboard logs')\n",
         "parser.add_argument('--experiment', dest='experiment',\n",
         "                    default=None, type=str,\n",
@@ -1697,6 +1703,9 @@
         "parser.add_argument('--evaluate', dest='evaluate',\n",
         "                    default=False, type=bool,\n",
         "                    help='Whether to perform evaluation')\n",
+        "parser.add_argument('--serving', dest='serving',\n",
+        "                    default=False, type=bool,\n",
+        "                    help='Whether to attach the serving function')\n",
         "parser.add_argument('--tuning', dest='tuning',\n",
         "                    default=False, type=bool,\n",
         "                    help='Whether to perform hyperparameter tuning')\n",
@@ -1730,10 +1739,12 @@
         "    aip.init(experiment=args.experiment, project=args.project)\n",
         "    aip.start_run(args.run)\n",
         "\n",
+        "metadata = {}\n",
+        "\n",
         "def get_data():\n",
         "    ''' Get the preprocessed training data '''\n",
         "    global train_data_file_pattern, val_data_file_pattern, test_data_file_pattern\n",
-        "    global label_column, transform_feature_spec\n",
+        "    global label_column, transform_feature_spec, metadata\n",
         "\n",
         "    dataset = aip.TabularDataset(args.dataset_id)\n",
         "    METADATA = 'gs://' + dataset.labels['user_metadata'] + \"/metadata.jsonl\"\n",
@@ -1764,6 +1775,10 @@
         "    if args.experiment:\n",
         "        aip.log_params(hyperparams)\n",
         "\n",
+        "    metadata.update(hyperparams)\n",
+        "    with tf.io.gfile.GFile(os.path.join(args.model_dir, \"metrics.txt\"), \"w\") as f:\n",
+        "        f.write(json.dumps(metadata))\n",
+        "\n",
         "    train.compile(model, hyperparams)\n",
         "    return model\n",
         "\n",
@@ -1775,6 +1790,11 @@
         "    trainparams[\"early_stop\"] = {\"monitor\": \"val_loss\", \"patience\": 5}\n",
         "    if args.experiment:\n",
         "        aip.log_params(trainparams)\n",
+        "\n",
+        "    metadata.update(trainparams)\n",
+        "    with tf.io.gfile.GFile(os.path.join(args.model_dir, \"metrics.txt\"), \"w\") as f:\n",
+        "        f.write(json.dumps(metadata))\n",
+        "\n",
         "    train.train(model, trainparams, train_data_file_pattern, val_data_file_pattern, label_column, transform_feature_spec, args.tensorboard_log_dir, args.tuning)\n",
         "    return model\n",
         "\n",
@@ -1783,8 +1803,10 @@
         "    evalparams = {}\n",
         "    evalparams[\"batch_size\"] = args.batch_size\n",
         "    metrics = train.evaluate(model, evalparams, test_data_file_pattern, label_column, transform_feature_spec)\n",
-        "    with tf.io.gfile.GFile(os.path.join(args.model_dir, \"metrics.txt\", \"w\")) as f:\n",
-        "        f.write(str(metrics))\n",
+        "\n",
+        "    metadata.update({'metrics': metrics})\n",
+        "    with tf.io.gfile.GFile(os.path.join(args.model_dir, \"metrics.txt\"), \"w\") as f:\n",
+        "        f.write(json.dumps(metadata))\n",
         "\n",
         "get_data()\n",
         "with strategy.scope():\n",
@@ -1794,8 +1816,16 @@
         "if args.evaluate:\n",
         "    evaluate_model(model)\n",
         "\n",
-        "logging.info('Save trained model to: ' + args.model_dir)\n",
-        "model.save(args.model_dir)"
+        "if args.serving:\n",
+        "    logging.info('Save serving model to: ' + args.model_dir)\n",
+        "    serving.construct_serving_model(\n",
+        "        model=model,\n",
+        "        serving_model_dir=args.model_dir,\n",
+        "        metadata=metadata\n",
+        "    )\n",
+        "else:\n",
+        "    logging.info('Save trained model to: ' + args.model_dir)\n",
+        "    model.save(args.model_dir)"
       ]
     },
     {
@@ -1966,6 +1996,7 @@
         "    accelerator_count=TRAIN_NGPU,\n",
         "    base_output_dir=MODEL_DIR,\n",
         "    service_account=SERVICE_ACCOUNT,\n",
+        "    tensorboard=tensorboard,\n",
         "    sync=True,\n",
         ")"
       ]
@@ -2404,6 +2435,7 @@
         "    accelerator_count=TRAIN_NGPU,\n",
         "    base_output_dir=MODEL_DIR,\n",
         "    service_account=SERVICE_ACCOUNT,\n",
+        "    tensorboard=tensorboard,\n",
         "    sync=True,\n",
         ")"
       ]
@@ -2449,8 +2481,10 @@
       },
       "outputs": [],
       "source": [
+        "EXPERIMENT_NAME = \"chicago\"\n",
+        "\n",
         "experiment_df = aip.get_experiment_df()\n",
-        "experiment_df = experiment_df[experiment_df.experiment_name == \"chicago\"]\n",
+        "experiment_df = experiment_df[experiment_df.experiment_name == EXPERIMENT_NAME]\n",
         "experiment_df.T"
       ]
     },
@@ -2477,6 +2511,28 @@
         "! gsutil cat $METRICS"
       ]
     },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "delete_tensorboard"
+      },
+      "source": [
+        "### Delete the TensorBoard instance\n",
+        "\n",
+        "Next, delete the TensorBoard instance."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "delete_tensorboard"
+      },
+      "outputs": [],
+      "source": [
+        "tensorboard.delete()"
+      ]
+    },
     {
       "cell_type": "code",
       "execution_count": null,
@@ -2508,6 +2564,13 @@
       },
       "outputs": [],
       "source": [
+        "%%writefile custom/trainer/serving.py\n",
+        "\n",
+        "import tensorflow as tf\n",
+        "import tensorflow_data_validation as tfdv\n",
+        "import tensorflow_transform as tft\n",
+        "import logging\n",
+        "\n",
         "def _get_serve_features_fn(model, tft_output):\n",
         "    \"\"\"Returns a function that accept a dictionary of features and applies TFT.\"\"\"\n",
         "\n",
@@ -2521,8 +2584,8 @@
         "        probabilities = model(transformed_features)\n",
         "        return {\"scores\": probabilities}\n",
         "\n",
-        "    return serve_features_fn\n",
         "\n",
+        "    return serve_features_fn\n",
         "\n",
         "def _get_serve_tf_examples_fn(model, tft_output, feature_spec):\n",
         "    \"\"\"Returns a function that parses a serialized tf.Example and applies TFT.\"\"\"\n",
@@ -2544,23 +2607,18 @@
         "\n",
         "    return serve_tf_examples_fn\n",
         "\n",
-        "\n",
-        "def construct_serving_model(model, serving_model_dir, metadata):\n",
+        "def construct_serving_model(\n",
+        "    model, serving_model_dir, metadata\n",
+        "):\n",
         "    global features\n",
         "\n",
-        "    schema_location = metadata[\"schema\"]\n",
-        "    features = (\n",
-        "        metadata[\"numeric_features\"]\n",
-        "        + metadata[\"categorical_features\"]\n",
-        "        + metadata[\"embedding_features\"]\n",
-        "    )\n",
+        "    schema_location = metadata['schema']\n",
+        "    features = metadata['numeric_features'] + metadata['categorical_features'] + metadata['embedding_features']\n",
         "    print(\"FEATURES\", features)\n",
         "    tft_output_dir = metadata[\"transform_artifacts_dir\"]\n",
         "\n",
         "    schema = tfdv.load_schema_text(schema_location)\n",
-        "    feature_spec = tft.tf_metadata.schema_utils.schema_as_feature_spec(\n",
-        "        schema\n",
-        "    ).feature_spec\n",
+        "    feature_spec = tft.tf_metadata.schema_utils.schema_as_feature_spec(schema).feature_spec\n",
         "\n",
         "    tft_output = tft.TFTransformOutput(tft_output_dir)\n",
         "\n",
@@ -2608,13 +2666,19 @@
       },
       "outputs": [],
       "source": [
+        "os.chdir(\"custom\")\n",
+        "\n",
+        "from trainer import serving\n",
+        "\n",
         "SERVING_MODEL_DIR = BUCKET_NAME + \"/serving_model\"\n",
         "\n",
-        "construct_serving_model(\n",
+        "serving.construct_serving_model(\n",
         "    model=model, serving_model_dir=SERVING_MODEL_DIR, metadata=metadata\n",
         ")\n",
         "\n",
-        "serving_model = tf.keras.models.load_model(SERVING_MODEL_DIR)"
+        "serving_model = tf.keras.models.load_model(SERVING_MODEL_DIR)\n",
+        "\n",
+        "os.chdir(\"..\")"
       ]
     },
     {

notebooks/community/ml_ops/stage3/mlops_formalization.ipynb

@@ -822,6 +822,41 @@
         "print(baseline_model)"
       ]
     },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "load_model_user_metadata:baseline"
+      },
+      "source": [
+        "### Load baseline models's user metadata\n",
+        "\n",
+        "Load the user metadata for the baseline model."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "load_model_user_metadata:baseline"
+      },
+      "outputs": [],
+      "source": [
+        "import json\n",
+        "\n",
+        "try:\n",
+        "    with tf.io.gfile.GFile(\n",
+        "        \"gs://\" + baseline_model.labels[\"user_metadata\"] + \"/metadata.jsonl\", \"r\"\n",
+        "    ) as f:\n",
+        "        baseline_metadata = json.load(f)\n",
+        "        print(baseline_metadata)\n",
+        "\n",
+        "        with tf.io.gfile.GFile(baseline_metadata[\"train_eval_metrics\"], \"r\") as f:\n",
+        "            baseline_metrics = json.load(f)\n",
+        "            print(baseline_metrics)\n",
+        "except:\n",
+        "    print(\"no metadata\")"
+      ]
+    },
     {
       "cell_type": "markdown",
       "metadata": {
@@ -907,11 +942,14 @@
         "def make_chicago_bq_dataset(bq_table: str, year: int, limit: int, project: str) -> str:\n",
         "    from google.cloud import bigquery\n",
         "\n",
-        "    bqclient = bigquery.Client()\n",
+        "    bqclient = bigquery.Client(project=project)\n",
         "\n",
         "    BQ_DATASET = bq_table.split(\".\")[1]\n",
         "    BQ_TABLE_COPY = f\"{project}.{BQ_DATASET}.taxi_trips\"\n",
         "\n",
+        "    if bq_table.startswith(\"bq://\"):\n",
+        "        bq_table = bq_table[5:]\n",
+        "\n",
         "    query = f\"\"\"\n",
         "    CREATE OR REPLACE TABLE `{BQ_TABLE_COPY}`\n",
         "    AS (\n",
@@ -971,6 +1009,9 @@
         "    )\n",
         "    \"\"\"\n",
         "\n",
+        "    response = bqclient.query(query)\n",
+        "    _ = response.result()\n",
+        "\n",
         "    return BQ_TABLE_COPY"
       ]
     },
@@ -2015,7 +2056,7 @@
         "    bucket: str,\n",
         "    project: str = PROJECT_ID,\n",
         "    region: str = REGION,\n",
-        "    labels: dict = {\"base_model\": 1},\n",
+        "    labels: dict = {\"base_model\": \"1\"},\n",
         "):\n",
         "    from google_cloud_pipeline_components import aiplatform as gcc_aip\n",
         "\n",
@@ -2130,7 +2171,7 @@
         "### Train the model task\n",
         "\n",
         "- Retrieve the model architecture.\n",
-        "- ?? Hypertune BLAH\n",
+        "- Retrieve the hyperparameters from the baseline model\n",
         "- Train the model\n",
         "- Evaluate the model\n",
         "\n",
@@ -2146,6 +2187,7 @@
     {
       "cell_type": "markdown",
       "metadata": {
+        "id": "construct_training_package"
       },
       "source": [
         "### Construct the training package\n",
@@ -2186,7 +2228,7 @@
         "setup_cfg = \"[egg_info]\\n\\ntag_build =\\n\\ntag_date = 0\"\n",
         "! echo \"$setup_cfg\" > custom/setup.cfg\n",
         "\n",
-        "setup_py = \"import setuptools\\n\\nsetuptools.setup(\\n\\n    install_requires=[\\n\\n        'google-cloud-aiplatform',\\n\\n        'cloudml-hypertune',\\n\\n        'tensorflow_datasets==1.3.0',\\n\\n    ],\\n\\n    packages=setuptools.find_packages())\"\n",
+        "setup_py = \"import setuptools\\n\\nsetuptools.setup(\\n\\n    install_requires=[\\n\\n        'google-cloud-aiplatform',\\n\\n        'cloudml-hypertune',\\n\\n        'tensorflow_datasets==1.3.0',\\n\\n        'tensorflow_data_validation==1.2',\\n\\n    ],\\n\\n    packages=setuptools.find_packages())\"\n",
         "! echo \"$setup_py\" > custom/setup.py\n",
         "\n",
         "pkg_info = \"Metadata-Version: 1.0\\n\\nName: Chicago Taxi tabular binary classification\\n\\nVersion: 0.0.0\\n\\nSummary: Demostration training script\\n\\nHome-page: www.google.com\\n\\nAuthor: Google\\n\\nAuthor-email: cdpe@google.com\\n\\nLicense: Public\\n\\nDescription: Demo\\n\\nPlatform: Vertex AI\"\n",
@@ -2314,13 +2356,16 @@
         "        batch_size=hyperparams[\"batch_size\"],\n",
         "    )\n",
         "\n",
-        "    tensorboard = tf.keras.callbacks.TensorBoard(log_dir=log_dir)\n",
-        "\n",
         "    early_stop = tf.keras.callbacks.EarlyStopping(\n",
         "        monitor=hyperparams[\"early_stop\"][\"monitor\"], patience=hyperparams[\"early_stop\"][\"patience\"], restore_best_weights=True\n",
         "    )\n",
         "\n",
-        "    callbacks=[tensorboard, early_stop]\n",
+        "    callbacks=[early_stop]\n",
+        "\n",
+        "    if log_dir:\n",
+        "        tensorboard = tf.keras.callbacks.TensorBoard(log_dir=log_dir)\n",
+        "\n",
+        "        callbacks=callbacks.append(tensorboard)\n",
         "\n",
         "    if tuning:\n",
         "        # Instantiate the HyperTune reporting object\n",
@@ -2370,6 +2415,108 @@
         "    return evaluation_metrics"
       ]
     },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "serving_function:chicago"
+      },
+      "source": [
+        "## Add a serving function\n",
+        "\n",
+        "Next, you add a serving function to your model for online and batch prediction. This allows prediction requests to be sent in raw format (unpreprocessed), either as a serialized TF.Example or JSONL object. The serving function will then preprocess the prediction request into the transformed format expected by the model."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "serving_function:chicago"
+      },
+      "outputs": [],
+      "source": [
+        "%%writefile custom/trainer/serving.py\n",
+        "\n",
+        "import tensorflow as tf\n",
+        "import tensorflow_data_validation as tfdv\n",
+        "import tensorflow_transform as tft\n",
+        "import logging\n",
+        "\n",
+        "def _get_serve_features_fn(model, tft_output):\n",
+        "    \"\"\"Returns a function that accept a dictionary of features and applies TFT.\"\"\"\n",
+        "\n",
+        "    model.tft_layer = tft_output.transform_features_layer()\n",
+        "\n",
+        "    @tf.function\n",
+        "    def serve_features_fn(raw_features):\n",
+        "        \"\"\"Returns the output to be used in the serving signature.\"\"\"\n",
+        "\n",
+        "        transformed_features = model.tft_layer(raw_features)\n",
+        "        probabilities = model(transformed_features)\n",
+        "        return {\"scores\": probabilities}\n",
+        "\n",
+        "\n",
+        "    return serve_features_fn\n",
+        "\n",
+        "def _get_serve_tf_examples_fn(model, tft_output, feature_spec):\n",
+        "    \"\"\"Returns a function that parses a serialized tf.Example and applies TFT.\"\"\"\n",
+        "\n",
+        "    model.tft_layer = tft_output.transform_features_layer()\n",
+        "\n",
+        "    @tf.function\n",
+        "    def serve_tf_examples_fn(serialized_tf_examples):\n",
+        "        \"\"\"Returns the output to be used in the serving signature.\"\"\"\n",
+        "        for key in list(feature_spec.keys()):\n",
+        "            if key not in features:\n",
+        "                feature_spec.pop(key)\n",
+        "\n",
+        "        parsed_features = tf.io.parse_example(serialized_tf_examples, feature_spec)\n",
+        "\n",
+        "        transformed_features = model.tft_layer(parsed_features)\n",
+        "        probabilities = model(transformed_features)\n",
+        "        return {\"scores\": probabilities}\n",
+        "\n",
+        "    return serve_tf_examples_fn\n",
+        "\n",
+        "def construct_serving_model(\n",
+        "    model, serving_model_dir, metadata\n",
+        "):\n",
+        "    global features\n",
+        "\n",
+        "    schema_location = metadata['schema']\n",
+        "    features = metadata['numeric_features'] + metadata['categorical_features'] + metadata['embedding_features']\n",
+        "    print(\"FEATURES\", features)\n",
+        "    tft_output_dir = metadata[\"transform_artifacts_dir\"]\n",
+        "\n",
+        "    schema = tfdv.load_schema_text(schema_location)\n",
+        "    feature_spec = tft.tf_metadata.schema_utils.schema_as_feature_spec(schema).feature_spec\n",
+        "\n",
+        "    tft_output = tft.TFTransformOutput(tft_output_dir)\n",
+        "\n",
+        "    # Drop features that were not used in training\n",
+        "    features_input_signature = {\n",
+        "        feature_name: tf.TensorSpec(\n",
+        "            shape=(None, 1), dtype=spec.dtype, name=feature_name\n",
+        "        )\n",
+        "        for feature_name, spec in feature_spec.items()\n",
+        "        if feature_name in features\n",
+        "    }\n",
+        "\n",
+        "    signatures = {\n",
+        "        \"serving_default\": _get_serve_features_fn(\n",
+        "            model, tft_output\n",
+        "        ).get_concrete_function(features_input_signature),\n",
+        "        \"serving_tf_example\": _get_serve_tf_examples_fn(\n",
+        "            model, tft_output, feature_spec\n",
+        "        ).get_concrete_function(\n",
+        "            tf.TensorSpec(shape=[None], dtype=tf.string, name=\"examples\")\n",
+        "        ),\n",
+        "    }\n",
+        "\n",
+        "    logging.info(\"Model export started...\")\n",
+        "    model.save(serving_model_dir, signatures=signatures)\n",
+        "    logging.info(\"Model export completed.\")"
+      ]
+    },
     {
       "cell_type": "markdown",
       "metadata": {
@@ -2461,7 +2608,7 @@
         "from trainer import data\n",
         "from trainer import model as model_\n",
         "from trainer import train\n",
-        "\n",
+        "from trainer import serving\n",
         "\n",
         "parser = argparse.ArgumentParser()\n",
         "parser.add_argument('--model-dir', dest='model_dir',\n",
@@ -2485,7 +2632,7 @@
         "parser.add_argument('--distribute', dest='distribute', type=str, default='single',\n",
         "                    help='distributed training strategy')\n",
         "parser.add_argument('--tensorboard-log-dir', dest='tensorboard_log_dir',\n",
-        "                    default='/tmp/logs', type=str,\n",
+        "                    default=os.getenv('AIP_TENSORBOARD_LOG_DIR'), type=str,\n",
         "                    help='Output file for tensorboard logs')\n",
         "parser.add_argument('--experiment', dest='experiment',\n",
         "                    default=None, type=str,\n",
@@ -2499,6 +2646,9 @@
         "parser.add_argument('--evaluate', dest='evaluate',\n",
         "                    default=False, type=bool,\n",
         "                    help='Whether to perform evaluation')\n",
+        "parser.add_argument('--serving', dest='serving',\n",
+        "                    default=False, type=bool,\n",
+        "                    help='Whether to attach the serving function')\n",
         "parser.add_argument('--tuning', dest='tuning',\n",
         "                    default=False, type=bool,\n",
         "                    help='Whether to perform hyperparameter tuning')\n",
@@ -2532,10 +2682,12 @@
         "    aip.init(experiment=args.experiment, project=args.project)\n",
         "    aip.start_run(args.run)\n",
         "\n",
+        "metadata = {}\n",
+        "\n",
         "def get_data():\n",
         "    ''' Get the preprocessed training data '''\n",
         "    global train_data_file_pattern, val_data_file_pattern, test_data_file_pattern\n",
-        "    global label_column, transform_feature_spec\n",
+        "    global label_column, transform_feature_spec, metadata\n",
         "\n",
         "    dataset = aip.TabularDataset(args.dataset_id)\n",
         "    METADATA = 'gs://' + dataset.labels['user_metadata'] + \"/metadata.jsonl\"\n",
@@ -2566,6 +2718,10 @@
         "    if args.experiment:\n",
         "        aip.log_params(hyperparams)\n",
         "\n",
+        "    metadata.update(hyperparams)\n",
+        "    with tf.io.gfile.GFile(os.path.join(args.model_dir, \"metrics.txt\"), \"w\") as f:\n",
+        "        f.write(json.dumps(metadata))\n",
+        "\n",
         "    train.compile(model, hyperparams)\n",
         "    return model\n",
         "\n",
@@ -2577,6 +2733,11 @@
         "    trainparams[\"early_stop\"] = {\"monitor\": \"val_loss\", \"patience\": 5}\n",
         "    if args.experiment:\n",
         "        aip.log_params(trainparams)\n",
+        "\n",
+        "    metadata.update(trainparams)\n",
+        "    with tf.io.gfile.GFile(os.path.join(args.model_dir, \"metrics.txt\"), \"w\") as f:\n",
+        "        f.write(json.dumps(metadata))\n",
+        "\n",
         "    train.train(model, trainparams, train_data_file_pattern, val_data_file_pattern, label_column, transform_feature_spec, args.tensorboard_log_dir, args.tuning)\n",
         "    return model\n",
         "\n",
@@ -2585,8 +2746,10 @@
         "    evalparams = {}\n",
         "    evalparams[\"batch_size\"] = args.batch_size\n",
         "    metrics = train.evaluate(model, evalparams, test_data_file_pattern, label_column, transform_feature_spec)\n",
+        "\n",
+        "    metadata.update({'metrics': metrics})\n",
         "    with tf.io.gfile.GFile(os.path.join(args.model_dir, \"metrics.txt\"), \"w\") as f:\n",
-        "        f.write(str(metrics))\n",
+        "        f.write(json.dumps(metadata))\n",
         "\n",
         "get_data()\n",
         "with strategy.scope():\n",
@@ -2596,8 +2759,16 @@
         "if args.evaluate:\n",
         "    evaluate_model(model)\n",
         "\n",
-        "logging.info('Save trained model to: ' + args.model_dir)\n",
-        "model.save(args.model_dir)"
+        "if args.serving:\n",
+        "    logging.info('Save serving model to: ' + args.model_dir)\n",
+        "    serving.construct_serving_model(\n",
+        "        model=model,\n",
+        "        serving_model_dir=args.model_dir,\n",
+        "        metadata=metadata\n",
+        "    )\n",
+        "else:\n",
+        "    logging.info('Save trained model to: ' + args.model_dir)\n",
+        "    model.save(args.model_dir)"
       ]
     },
     {
@@ -2621,7 +2792,7 @@
       "source": [
         "DATASET_ID = dataset_id\n",
         "MODEL_ID = model_id\n",
-        "!cd custom; python3 -m trainer.task --model-id={MODEL_ID} --dataset-id={DATASET_ID} --experiment='chicago' --run='test' --project={PROJECT_ID} --epochs=5 --model-dir=/tmp --evaluate=True"
+        "!cd custom; python3 -m trainer.task --model-id={MODEL_ID} --dataset-id={DATASET_ID} --experiment='chicago' --run='test' --project={PROJECT_ID} --epochs=5 --model-dir=/tmp --evaluate=True --serving=True"
       ]
     },
     {
@@ -2715,6 +2886,9 @@
         "    model_serving_container_image_uri: str,\n",
         "    machine_type: str,\n",
         "    bucket: str,\n",
+        "    tensorboard: str = None,\n",
+        "    service_account: str = None,\n",
+        "    label: str = str({\"candidate_model\": \"1\"}).replace(\"'\", '\"'),\n",
         "    replica_count: int = 1,\n",
         "    accelerator_type: str = None,\n",
         "    accelerator_count: int = 0,\n",
@@ -2736,9 +2910,12 @@
         "        machine_type=machine_type,\n",
         "        accelerator_type=accelerator_type,\n",
         "        accelerator_count=accelerator_count,\n",
+        "        tensorboard=tensorboard,\n",
+        "        service_account=service_account,\n",
         "        # Serving - As part of this operation, the model is registered to Vertex AI\n",
         "        model_serving_container_image_uri=model_serving_container_image_uri,\n",
         "        model_display_name=display_name,\n",
+        "        labels=label,\n",
         "    )\n",
         "\n",
         "    endpoint_op = gcc_aip.EndpointCreateOp(\n",
@@ -2756,6 +2933,33 @@
         "    )"
       ]
     },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "create_tensorboard_instance"
+      },
+      "source": [
+        "### Create a Vertex AI TensorBoard instance\n",
+        "\n",
+        "Create a Vertex AI TensorBoard instance to use TensorBoard in conjunction with Vertex AI Training for custom model training.\n",
+        "\n",
+        "Learn more about [Get started with Vertex AI TensorBoard](https://cloud.google.com/vertex-ai/docs/experiments/tensorboard-overview)."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "create_tensorboard_instance"
+      },
+      "outputs": [],
+      "source": [
+        "TENSORBOARD_DISPLAY_NAME = \"chicago_\" + TIMESTAMP\n",
+        "tensorboard = aip.Tensorboard.create(display_name=TENSORBOARD_DISPLAY_NAME)\n",
+        "tensorboard_resource_name = tensorboard.gca_resource.name\n",
+        "print(\"TensorBoard resource name:\", tensorboard_resource_name)"
+      ]
+    },
     {
       "cell_type": "markdown",
       "metadata": {
@@ -2772,6 +2976,7 @@
         "- `python_package`: The Python package for the custom training job.\n",
         "- `python_module`: The Python module in the package to execute.\n",
         "- `args`: The command line arguments to pass to the Python module.\n",
+        "    - *Note*: The pipeline uses the hyperparameters from the baseline model. Alternatively, one could use the hyperparameters from the current blessed model, or repeat hyperparameter tuning.\n",
         "- `container_uri`: The training container image.\n",
         "- `model_serving_container_image_uri`: The associated deployment container image.\n",
         "- `machine_type`: The VM for executing the training job.\n",
@@ -2779,6 +2984,8 @@
         "- `accelerator_type`: The type of HW accelerators -- if any.\n",
         "- `accelerator_count`: The number of HW accelerators -- if any.\n",
         "- `bucket`: The Cloud Storage location to store the model artifacts.\n",
+        "- `tensorboard`: The full resource name of a Vertex AI Tensorboard.\n",
+        "- `service_account`: The service account for the Tensorboard instance.\n",
         "- `project`: The project ID.\n",
         "- `region`: The region."
       ]
@@ -2811,17 +3018,25 @@
         "            \"--model-id\",\n",
         "            model_id,\n",
         "            \"--experiment\",\n",
-        "            \"chicago\",\n",
+        "            \"chicago\" + TIMESTAMP,\n",
         "            \"--run\",\n",
         "            \"retrain\",\n",
         "            \"--epochs\",\n",
-        "            \"50\",\n",
+        "            str(int(baseline_metrics[\"num_epochs\"])),\n",
+        "            \"--batch_size\",\n",
+        "            str(int(baseline_metrics[\"batch_size\"])),\n",
+        "            \"--lr\",\n",
+        "            baseline_metrics[\"learning_rate\"],\n",
         "            \"--evaluate\",\n",
         "            \"True\",\n",
+        "            \"--serving\",\n",
+        "            \"True\",\n",
         "            \"--project\",\n",
         "            PROJECT_ID,\n",
         "        ],\n",
         "        \"container_uri\": TRAIN_IMAGE,\n",
+        "        \"tensorboard\": tensorboard.gca_resource.name,\n",
+        "        \"service_account\": SERVICE_ACCOUNT,\n",
         "        \"model_serving_container_image_uri\": DEPLOY_IMAGE,\n",
         "        \"machine_type\": TRAIN_COMPUTE,\n",
         "        \"replica_count\": 1,\n",
@@ -2856,29 +3071,70 @@
       "outputs": [],
       "source": [
         "print(\"custompythonpackagetrainingjob-run\")\n",
-        "artifacts = print_pipeline_output(pipeline, 'custompythonpackagetrainingjob-run')\n",
-        "print('\\n')print(\"custompythonpackagetrainingjob-run\")\n",
-        "artifacts = print_pipeline_output(pipeline, 'custompythonpackagetrainingjob-run')\n",
-        "print('\\n')\n",
+        "artifacts = print_pipeline_output(pipeline, \"custompythonpackagetrainingjob-run\")\n",
+        "print(\"\\n\")\n",
+        "artifacts = print_pipeline_output(pipeline, \"custompythonpackagetrainingjob-run\")\n",
+        "print(\"\\n\")\n",
         "output = !gsutil cat $artifacts\n",
         "output = json.loads(output[0])\n",
-        "model_id = output['artifacts']['model']['artifacts'][0]['metadata']['resourceName']\n",
-        "print('\\n')\n",
+        "model_id = output[\"artifacts\"][\"model\"][\"artifacts\"][0][\"metadata\"][\"resourceName\"]\n",
+        "print(\"\\n\")\n",
         "print(model_id)\n",
-        "print('\\n')\n",
+        "print(\"\\n\")\n",
         "\n",
         "print(\"endpoint-create\")\n",
-        "artifacts = print_pipeline_output(pipeline, 'endpoint-create')\n",
-        "print('\\n')\n",
+        "artifacts = print_pipeline_output(pipeline, \"endpoint-create\")\n",
+        "print(\"\\n\")\n",
         "print(\"model-deploy\")\n",
-        "artifacts = print_pipeline_output(pipeline, 'model-deploy')\n",
-        "print('\\n')\n",
+        "artifacts = print_pipeline_output(pipeline, \"model-deploy\")\n",
+        "print(\"\\n\")\n",
         "print(\"endpoint-create\")\n",
-        "artifacts = print_pipeline_output(pipeline, 'endpoint-create')\n",
-        "print('\\n')\n",
+        "artifacts = print_pipeline_output(pipeline, \"endpoint-create\")\n",
+        "print(\"\\n\")\n",
         "print(\"model-deploy\")\n",
-        "artifacts = print_pipeline_output(pipeline, 'model-deploy')\n",
-        "print('\\n')"
+        "artifacts = print_pipeline_output(pipeline, \"model-deploy\")\n",
+        "print(\"\\n\")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "get_experiment:init"
+      },
+      "source": [
+        "### Get the experiment results\n",
+        "\n",
+        "Next, you use the experiment name as a parameter to the method `get_experiment_df()` to get the results of the experiment as a pandas dataframe."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "get_experiment:init"
+      },
+      "outputs": [],
+      "source": [
+        "EXPERIMENT_NAME = \"chicago\" + TIMESTAMP\n",
+        "\n",
+        "aip.init(experiment=EXPERIMENT_NAME)\n",
+        "experiment_df = aip.get_experiment_df()\n",
+        "experiment_df = experiment_df[experiment_df.experiment_name == EXPERIMENT_NAME]\n",
+        "experiment_df.T"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "get_metrics_trained_model"
+      },
+      "outputs": [],
+      "source": [
+        "model = aip.Model(model_id)\n",
+        "model_artifacts = model.gca_resource.artifact_uri\n",
+        "\n",
+        "!gsutil cat {model_artifacts}/metrics.txt"
       ]
     },
     {