Merge branch 'master' into ml_ops

Merge branch 'ml_ops' of https://github.com/GoogleCloudPlatform/vertex-ai-samples into ml_ops

notebooks/community/ml_ops/stage1/get_started_bq_datasets.ipynb

@@ -167,7 +167,8 @@
         "    ! 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",
-        "    ! pip3 install --upgrade pyarrow $USER_FLAG"
+        "    ! pip3 install --upgrade pyarrow $USER_FLAG\n",
+        "    ! pip3 install --upgrade cloudml-hypertune $USER_FLAG"
       ]
     },
     {
@@ -1191,14 +1192,14 @@
         "    except Exception as e:\n",
         "        print(e)\n",
         "\n",
-        "    # Delete the AutoML or Pipeline trainig job\n",
+        "    # Delete the AutoML or Pipeline training job\n",
         "    try:\n",
         "        if \"dag\" in globals():\n",
         "            dag.delete()\n",
         "    except Exception as e:\n",
         "        print(e)\n",
         "\n",
-        "    # Delete the custom trainig job\n",
+        "    # Delete the custom training job\n",
         "    try:\n",
         "        if \"job\" in globals():\n",
         "            job.delete()\n",

notebooks/community/ml_ops/stage1/get_started_dataflow.ipynb

@@ -160,7 +160,8 @@
         "    ! 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",
-        "    ! pip3 install --upgrade pyarrow $USER_FLAG"
+        "    ! pip3 install --upgrade pyarrow $USER_FLAG\n",
+        "    ! pip3 install --upgrade cloudml-hypertune $USER_FLAG"
       ]
     },
     {
@@ -1204,14 +1205,14 @@
         "    except Exception as e:\n",
         "        print(e)\n",
         "\n",
-        "    # Delete the AutoML or Pipeline trainig job\n",
+        "    # Delete the AutoML or Pipeline training job\n",
         "    try:\n",
         "        if \"dag\" in globals():\n",
         "            dag.delete()\n",
         "    except Exception as e:\n",
         "        print(e)\n",
         "\n",
-        "    # Delete the custom trainig job\n",
+        "    # Delete the custom training job\n",
         "    try:\n",
         "        if \"job\" in globals():\n",
         "            job.delete()\n",

notebooks/community/ml_ops/stage1/get_started_vertex_datasets.ipynb

@@ -162,7 +162,8 @@
         "    ! 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",
-        "    ! pip3 install --upgrade pyarrow $USER_FLAG"
+        "    ! pip3 install --upgrade pyarrow $USER_FLAG\n",
+        "    ! pip3 install --upgrade cloudml-hypertune $USER_FLAG"
       ]
     },
     {
@@ -1571,14 +1572,14 @@
         "    except Exception as e:\n",
         "        print(e)\n",
         "\n",
-        "    # Delete the AutoML or Pipeline trainig job\n",
+        "    # Delete the AutoML or Pipeline training job\n",
         "    try:\n",
         "        if \"dag\" in globals():\n",
         "            dag.delete()\n",
         "    except Exception as e:\n",
         "        print(e)\n",
         "\n",
-        "    # Delete the custom trainig job\n",
+        "    # Delete the custom training job\n",
         "    try:\n",
         "        if \"job\" in globals():\n",
         "            job.delete()\n",

notebooks/community/ml_ops/stage1/mlops_data_management.ipynb

@@ -106,7 +106,7 @@
       "source": [
         "### Recommendations\n",
         "\n",
-        "When doing E2E MLOps on Google Cloud, the following best practices with structured (tabular) data are recommended:\n",
+        "When doing E2E MLOps on Google Cloud for data management, the following best practices with structured (tabular) data are recommended:\n",
         "\n",
         " - For large amounts of data, use BigQuery table. Otherwise, use a CSV file stored in Cloud Storage.\n",
         " - When storing a large amount of data in CSV file, shard the data at 10,000 rows per shard.\n",
@@ -143,7 +143,8 @@
         "    ! 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",
-        "    ! pip3 install --upgrade pyarrow $USER_FLAG"
+        "    ! pip3 install --upgrade pyarrow $USER_FLAG\n",
+        "    ! pip3 install --upgrade cloudml-hypertune $USER_FLAG"
       ]
     },
     {
@@ -655,7 +656,7 @@
       "source": [
         "BQ_DATASET = BQ_TABLE.split(\".\")[1]\n",
         "BQ_TABLE_COPY = f\"{PROJECT_ID}.{BQ_DATASET}.taxi_trips\"\n",
-        "LIMIT = 20000\n",
+        "LIMIT = 300000\n",
         "YEAR = 2020\n",
         "\n",
         "query = f\"\"\"\n",
@@ -691,7 +692,7 @@
         "      EXTRACT(DAY from trip_start_timestamp) as trip_day,\n",
         "      EXTRACT(DAYOFWEEK from trip_start_timestamp) as trip_day_of_week,\n",
         "      EXTRACT(HOUR from trip_start_timestamp) as trip_hour,\n",
-        "      trip_seconds,\n",
+        "      CAST(trip_seconds AS FLOAT64) as trip_seconds,\n",
         "      trip_miles,\n",
         "      payment_type,\n",
         "      ST_AsText(\n",
@@ -793,7 +794,7 @@
         "# Download a table.\n",
         "table = bigquery.TableReference.from_string(BQ_TABLE)\n",
         "\n",
-        "rows = bqclient.list_rows(table, max_results=20000)\n",
+        "rows = bqclient.list_rows(table, max_results=300000)\n",
         "\n",
         "dataframe = rows.to_dataframe()\n",
         "print(dataframe.head())"
@@ -1172,7 +1173,16 @@
         "- `Numeric`: Rescale the values with `tft.scale_to_z_score`.\n",
         "- `Categorical`: Encode as a categorical column with `tft.compute_and_apply_vocabulary`.\n",
         "\n",
-        "For expendiency, the number of examples from the dataset is limited to 500."
+        "In addition to the preprocessed (transformed) data, raw versions of the test data are generated in both tf.Example and JSONL format. The tranform artifacts are stored as well to be used by subsequent serving function for transforming raw data into transformed data.\n",
+        "\n",
+        "In summary, the outputs produced are:\n",
+        "\n",
+        "- transformed training data\n",
+        "- transformed validation data\n",
+        "- tranformed test dats\n",
+        "- raw test data as JSONL\n",
+        "- raw test data as tf.Example\n",
+        "- transform function artifacts"
       ]
     },
     {
@@ -1213,6 +1223,13 @@
         "    return len(ratio) - 1\n",
         "\n",
         "\n",
+        "def convert_to_jsonl(data):\n",
+        "    \"\"\"Converts a parsed record to JSON\"\"\"\n",
+        "    import json\n",
+        "\n",
+        "    return json.dumps(data)\n",
+        "\n",
+        "\n",
         "def run_pipeline(args):\n",
         "    \"\"\"Runs a Beam pipeline to split the dataset\"\"\"\n",
         "\n",
@@ -1221,6 +1238,8 @@
         "    raw_data_query = args[\"raw_data_query\"]\n",
         "    transformed_data_prefix = args[\"transformed_data_prefix\"]\n",
         "    transform_artifact_dir = args[\"transform_artifact_dir\"]\n",
+        "    exported_jsonl_prefix = args[\"exported_jsonl_prefix\"]\n",
+        "    exported_tfrec_prefix = args[\"exported_tfrec_prefix\"]\n",
         "    temp_location = args[\"temp_location\"]\n",
         "    project = args[\"project\"]\n",
         "\n",
@@ -1310,7 +1329,6 @@
         "            # Get data from the transformed_test_dataset.\n",
         "            transformed_test_data, _ = transformed_test_dataset\n",
         "\n",
-
         "            # write transformed test data.\n",
         "            _ = (\n",
         "                transformed_test_data\n",
@@ -1327,12 +1345,34 @@
         "                transform_artifact_dir\n",
         "            )\n",
         "\n",
+        "            # Convert raw test data to JSON (for batch prediction)\n",
+        "            json_test_data = (raw_test_data) | \"Convert Test Data\" >> beam.Map(\n",
+        "                convert_to_jsonl\n",
+        "            )\n",
         "\n",
-        "EXPORTED_DATA_PREFIX = os.path.join(BUCKET_NAME, \"exported_data\")\n",
+        "            # Write raw test data to GCS as JSONL files.\n",
+        "            _ = json_test_data | \"Write JSONL Test Data\" >> beam.io.WriteToText(\n",
+        "                file_path_prefix=exported_jsonl_prefix, file_name_suffix=\".jsonl\"\n",
+        "            )\n",
+        "\n",
+        "            # Write raw test data to GCS as TF Records\n",
+        "            _ = (\n",
+        "                raw_test_data\n",
+        "                | \"Write TF Test Data\"\n",
+        "                >> beam.io.tfrecordio.WriteToTFRecord(\n",
+        "                    file_path_prefix=os.path.join(exported_tfrec_prefix, \"data\"),\n",
+        "                    file_name_suffix=\".tfrecord\",\n",
+        "                    coder=tft.coders.ExampleProtoCoder(raw_metadata.schema),\n",
+        "                )\n",
+        "            )\n",
+        "\n",
+        "\n",
+        "EXPORTED_JSONL_PREFIX = os.path.join(BUCKET_NAME, \"exported_data/jsonl\")\n",
+        "EXPORTED_TFREC_PREFIX = os.path.join(BUCKET_NAME, \"exported_data/tfrec\")\n",
         "TRANSFORMED_DATA_PREFIX = os.path.join(BUCKET_NAME, \"transformed_data\")\n",
         "TRANSFORM_ARTIFACTS_DIR = os.path.join(BUCKET_NAME, \"transformed_artifacts\")\n",
         "\n",
-        "QUERY_STRING = \"SELECT * FROM {} LIMIT 20000\".format(BQ_TABLE)\n",
+        "QUERY_STRING = \"SELECT * FROM {} LIMIT 300000\".format(BQ_TABLE)\n",
         "JOB_NAME = \"chicago\" + TIMESTAMP\n",
         "\n",
         "args = {\n",
@@ -1340,6 +1380,8 @@
         "    \"raw_data_query\": QUERY_STRING,\n",
         "    \"transformed_data_prefix\": TRANSFORMED_DATA_PREFIX,\n",
         "    \"transform_artifact_dir\": TRANSFORM_ARTIFACTS_DIR,\n",
+        "    \"exported_jsonl_prefix\": EXPORTED_JSONL_PREFIX,\n",
+        "    \"exported_tfrec_prefix\": EXPORTED_TFREC_PREFIX,\n",
         "    \"temp_location\": os.path.join(BUCKET_NAME, \"temp\"),\n",
         "    \"project\": PROJECT_ID,\n",
         "    \"region\": REGION,\n",
@@ -1353,7 +1395,9 @@
         "! gsutil ls $TRANSFORMED_DATA_PREFIX/train\n",
         "! gsutil ls $TRANSFORMED_DATA_PREFIX/val\n",
         "! gsutil ls $TRANSFORMED_DATA_PREFIX/test\n",
-        "! gsutil ls $TRANSFORM_ARTIFACTS_DIR"
+        "! gsutil ls $TRANSFORM_ARTIFACTS_DIR\n",
+        "! gsutil ls $EXPORTED_JSONL_PREFIX\n",
+        "! gsutil ls $EXPORTED_TFREC_PREFIX"
       ]
     },
     {
@@ -1364,7 +1408,11 @@
       "source": [
         "#### Retain transformed dataset references\n",
         "\n",
-        "Next, you retain the Cloud Storage location of the transformed data with the dataset. In this example, you add it to the user-defined metadata for this dataset, which is stored in the dataset's Cloud Storage bucket."
+        "Next, you retain the Cloud Storage location of the transformed data with the dataset. In this example, you add it to the user-defined metadata for this dataset, which is stored in the dataset's Cloud Storage bucket.\n",
+        "\n",
+        "During the transformation of the data, the transformat function calculated the number of unique occurrences per feature. Some of the categorical values (string, int) may have a large number of unique values. In this case, its better to reduce their dimensionality by moving them from being categorical to embedding feature.\n",
+        "\n",
+        "The code uses a rule-of-thumb that the embedding size should be the sqrt() of the number of unique values."
       ]
     },
     {
@@ -1380,8 +1428,26 @@
         ") as f:\n",
         "    metadata = json.load(f)\n",
         "\n",
+        "tft_output = tft.TFTransformOutput(TRANSFORM_ARTIFACTS_DIR)\n",
+        "\n",
+        "CATEGORICAL_FEATURES = []\n",
+        "EMBEDDING_FEATURES = []\n",
+        "categorical_features = metadata[\"categorical_features\"]\n",
+        "for feature in categorical_features:\n",
+        "    unique = tft_output.vocabulary_size_by_name(feature)\n",
+        "    if unique > 10:\n",
+        "        EMBEDDING_FEATURES.append(feature)\n",
+        "        print(\"Convert to embedding\", feature, unique)\n",
+        "    else:\n",
+        "        CATEGORICAL_FEATURES.append(feature)\n",
+        "\n",
+        "metadata[\"categorical_features\"] = CATEGORICAL_FEATURES\n",
+        "metadata[\"embedding_features\"] = EMBEDDING_FEATURES\n",
+        "\n",
         "metadata[\"transformed_data_prefix\"] = TRANSFORMED_DATA_PREFIX\n",
         "metadata[\"transform_artifacts_dir\"] = TRANSFORM_ARTIFACTS_DIR\n",
+        "metadata[\"exported_jsonl_prefix\"] = EXPORTED_JSONL_PREFIX\n",
+        "metadata[\"exported_tfrec_prefix\"] = EXPORTED_TFREC_PREFIX\n",
         "with tf.io.gfile.GFile(\n",
         "    \"gs://\" + dataset.labels[\"user_metadata\"] + \"/metadata.jsonl\", \"w\"\n",
         ") as f:\n",
@@ -1418,11 +1484,11 @@
       "cell_type": "code",
       "execution_count": null,
       "metadata": {
-        "id": "cleanup:mbsdk"
+        "id": "cleanup:stage1"
       },
       "outputs": [],
       "source": [
-        "delete_all = True\n",
+        "delete_all = False\n",
         "\n",
         "if delete_all:\n",
         "    # Delete the dataset using the Vertex dataset object\n",
@@ -1432,48 +1498,6 @@
         "    except Exception as e:\n",
         "        print(e)\n",
         "\n",
-        "    # Delete the model using the Vertex model object\n",
-        "    try:\n",
-        "        if \"model\" in globals():\n",
-        "            model.delete()\n",
-        "    except Exception as e:\n",
-        "        print(e)\n",
-        "\n",
-        "    # Delete the endpoint using the Vertex endpoint object\n",
-        "    try:\n",
-        "        if \"endpoint\" in globals():\n",
-        "            endpoint.delete()\n",
-        "    except Exception as e:\n",
-        "        print(e)\n",
-        "\n",
-        "    # Delete the AutoML or Pipeline trainig job\n",
-        "    try:\n",
-        "        if \"dag\" in globals():\n",
-        "            dag.delete()\n",
-        "    except Exception as e:\n",
-        "        print(e)\n",
-        "\n",
-        "    # Delete the custom trainig job\n",
-        "    try:\n",
-        "        if \"job\" in globals():\n",
-        "            job.delete()\n",
-        "    except Exception as e:\n",
-        "        print(e)\n",
-        "\n",
-        "    # Delete the batch prediction job using the Vertex batch prediction object\n",
-        "    try:\n",
-        "        if \"batch_predict_job\" in globals():\n",
-        "            batch_predict_job.delete()\n",
-        "    except Exception as e:\n",
-        "        print(e)\n",
-        "\n",
-        "    # Delete the hyperparameter tuning job using the Vertex hyperparameter tuning object\n",
-        "    try:\n",
-        "        if \"hpt_job\" in globals():\n",
-        "            hpt_job.delete()\n",
-        "    except Exception as e:\n",
-        "        print(e)\n",
-        "\n",
         "    if \"BUCKET_NAME\" in globals():\n",
         "        ! gsutil rm -r $BUCKET_NAME"
       ]

notebooks/community/ml_ops/stage2/get_started_vertex_experiments.ipynb

@@ -137,7 +137,8 @@
         "    ! 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",
-        "    ! pip3 install --upgrade pyarrow $USER_FLAG"
+        "    ! pip3 install --upgrade pyarrow $USER_FLAG\n",
+        "    ! pip3 install --upgrade cloudml-hypertune $USER_FLAG"
       ]
     },
     {
@@ -701,11 +702,9 @@
         "id": "get_experiment"
       },
       "source": [
-        "### Get logged output for an experiment\n",
+        "### Get the experiment results\n",
         "\n",
-        "You retrieve the preserved output for an experiment using the `get_experiment_df()` method, with the following parameters:\n",
-        "\n",
-        "-`experiment`: The name of the experiment."
+        "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."
       ]
     },
     {
@@ -716,8 +715,9 @@
       },
       "outputs": [],
       "source": [
-        "dataframe = aip.get_experiment_df(experiment=EXPERIMENT_NAME)\n",
-        "print(dataframe)"
+        "experiment_df = aip.get_experiment_df()\n",
+        "experiment_df = experiment_df[experiment_df.experiment_name == \"example\"]\n",
+        "experiment_df.T"
       ]
     },
     {
@@ -776,14 +776,14 @@
         "    except Exception as e:\n",
         "        print(e)\n",
         "\n",
-        "    # Delete the AutoML or Pipeline trainig job\n",
+        "    # Delete the AutoML or Pipeline training job\n",
         "    try:\n",
         "        if \"dag\" in globals():\n",
         "            dag.delete()\n",
         "    except Exception as e:\n",
         "        print(e)\n",
         "\n",
-        "    # Delete the custom trainig job\n",
+        "    # Delete the custom training job\n",
         "    try:\n",
         "        if \"job\" in globals():\n",
         "            job.delete()\n",

notebooks/community/ml_ops/stage2/get_started_vertex_tensorboard.ipynb

@@ -133,7 +133,8 @@
         "    ! 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",
-        "    ! pip3 install --upgrade pyarrow $USER_FLAG"
+        "    ! pip3 install --upgrade pyarrow $USER_FLAG\n",
+        "    ! pip3 install --upgrade cloudml-hypertune $USER_FLAG"
       ]
     },
     {
@@ -819,7 +820,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        '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        'tensorflow==2.5.0',\\n\\n        'tensorflow_datasets==1.3.0',\\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: example image classification\\n\\nVersion: 0.0.0\\n\\nSummary: Demostration training script\\n\\nHome-page: www.google.com\\n\\nAuthor: Google\\n\\nAuthor-email: aferlitsch@google.com\\n\\nLicense: Public\\n\\nDescription: Demo\\n\\nPlatform: Vertex\"\n",
@@ -1161,14 +1162,14 @@
         "    except Exception as e:\n",
         "        print(e)\n",
         "\n",
-        "    # Delete the AutoML or Pipeline trainig job\n",
+        "    # Delete the AutoML or Pipeline training job\n",
         "    try:\n",
         "        if \"dag\" in globals():\n",
         "            dag.delete()\n",
         "    except Exception as e:\n",
         "        print(e)\n",
         "\n",
-        "    # Delete the custom trainig job\n",
+        "    # Delete the custom training job\n",
         "    try:\n",
         "        if \"job\" in globals():\n",
         "            job.delete()\n",

notebooks/community/ml_ops/stage2/get_started_vertex_training.ipynb

@@ -86,7 +86,10 @@
         "\n",
         "The steps performed include:\n",
         "\n",
-        "- ??"
+        "- Training using a single Python script.\n",
+        "- Training using a Python package.\n",
+        "- Training using a custom training image.\n",
+        "- Laying out a training package."
       ]
     },
     {
@@ -97,7 +100,7 @@
       "source": [
         "### Recommendations\n",
         "\n",
-        "When doing E2E MLOps on Google Cloud, the following best practices for selecting SDK methods for custom model training:\n",
+        "When doing E2E MLOps on Google Cloud, the following are best practices for selecting SDK methods for custom model training:\n",
         "\n",
         "**You have a python training package and using a prebuilt Google container**\n",
         "\n",
@@ -145,7 +148,8 @@
         "    ! 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",
-        "    ! pip3 install --upgrade pyarrow $USER_FLAG"
+        "    ! pip3 install --upgrade pyarrow $USER_FLAG\n",
+        "    ! pip3 install --upgrade cloudml-hypertune $USER_FLAG"
       ]
     },
     {
@@ -406,27 +410,30 @@
     {
       "cell_type": "markdown",
       "metadata": {
-        "id": "accelerators:training,mbsdk,cpu"
+        "id": "accelerators:training,cpu,prediction,cpu,mbsdk"
       },
       "source": [
         "#### Set hardware accelerators\n",
         "\n",
-        "You can set hardware accelerators for training.\n",
+        "You can set hardware accelerators for training and prediction.\n",
         "\n",
-        "Set the variable `TRAIN_GPU/TRAIN_NGPU` to use a container image supporting a GPU and the number of GPUs allocated to the virtual machine (VM) instance. For example, to use a GPU container image with 4 Nvidia Telsa K80 GPUs allocated to each VM, you would specify:\n",
+        "Set the variables `TRAIN_GPU/TRAIN_NGPU` and `DEPLOY_GPU/DEPLOY_NGPU` to use a container image supporting a GPU and the number of GPUs allocated to the virtual machine (VM) instance. For example, to use a GPU container image with 4 Nvidia Telsa K80 GPUs allocated to each VM, you would specify:\n",
         "\n",
         "    (aip.AcceleratorType.NVIDIA_TESLA_K80, 4)\n",
         "\n",
+        "\n",
         "Otherwise specify `(None, None)` to use a container image to run on a CPU.\n",
         "\n",
-        "Learn more [here](https://cloud.google.com/vertex-ai/docs/general/locations#accelerators) hardware accelerator support for your region"
+        "Learn more [here](https://cloud.google.com/vertex-ai/docs/general/locations#accelerators) hardware accelerator support for your region\n",
+        "\n",
+        "*Note*: TF releases before 2.3 for GPU support will fail to load the custom model in this tutorial. It is a known issue and fixed in TF 2.3 -- which is caused by static graph ops that are generated in the serving function. If you encounter this issue on your own custom models, use a container image for TF 2.3 with GPU support."
       ]
     },
     {
       "cell_type": "code",
       "execution_count": null,
       "metadata": {
-        "id": "accelerators:training,mbsdk,cpu"
+        "id": "accelerators:training,cpu,prediction,cpu,mbsdk"
       },
       "outputs": [],
       "source": [
@@ -436,30 +443,39 @@
         "        int(os.getenv(\"IS_TESTING_TRAIN_GPU\")),\n",
         "    )\n",
         "else:\n",
-        "    TRAIN_GPU, TRAIN_NGPU = (None, None)"
+        "    TRAIN_GPU, TRAIN_NGPU = (None, None)\n",
+        "\n",
+        "if os.getenv(\"IS_TESTING_DEPLOY_GPU\"):\n",
+        "    DEPLOY_GPU, DEPLOY_NGPU = (\n",
+        "        aip.gapic.AcceleratorType.NVIDIA_TESLA_K80,\n",
+        "        int(os.getenv(\"IS_TESTING_DEPLOY_GPU\")),\n",
+        "    )\n",
+        "else:\n",
+        "    DEPLOY_GPU, DEPLOY_NGPU = (None, None)"
       ]
     },
     {
       "cell_type": "markdown",
       "metadata": {
-        "id": "container:training"
+        "id": "container:training,prediction"
       },
       "source": [
         "#### Set pre-built containers\n",
         "\n",
-        "Set the pre-built Docker container image for training.\n",
-        "\n",
-        "- Set the variable `TF` to the TensorFlow version of the container image. For example, `2-1` would be version 2.1, and `1-15` would be version 1.15. The following list shows some of the pre-built images available:\n",
+        "Set the pre-built Docker container image for training and prediction.\n",
         "\n",
         "\n",
-        "For the latest list, see [Pre-built containers for training](https://cloud.google.com/ai-platform-unified/docs/training/pre-built-containers)."
+        "For the latest list, see [Pre-built containers for training](https://cloud.google.com/ai-platform-unified/docs/training/pre-built-containers).\n",
+        "\n",
+        "\n",
+        "For the latest list, see [Pre-built containers for prediction](https://cloud.google.com/ai-platform-unified/docs/predictions/pre-built-containers)."
       ]
     },
     {
       "cell_type": "code",
       "execution_count": null,
       "metadata": {
-        "id": "container:training"
+        "id": "container:training,prediction"
       },
       "outputs": [],
       "source": [
@@ -473,18 +489,29 @@
         "        TRAIN_VERSION = \"tf-gpu.{}\".format(TF)\n",
         "    else:\n",
         "        TRAIN_VERSION = \"tf-cpu.{}\".format(TF)\n",
+        "    if DEPLOY_GPU:\n",
+        "        DEPLOY_VERSION = \"tf2-gpu.{}\".format(TF)\n",
+        "    else:\n",
+        "        DEPLOY_VERSION = \"tf2-cpu.{}\".format(TF)\n",
         "else:\n",
         "    if TRAIN_GPU:\n",
         "        TRAIN_VERSION = \"tf-gpu.{}\".format(TF)\n",
         "    else:\n",
         "        TRAIN_VERSION = \"tf-cpu.{}\".format(TF)\n",
+        "    if DEPLOY_GPU:\n",
+        "        DEPLOY_VERSION = \"tf-gpu.{}\".format(TF)\n",
+        "    else:\n",
+        "        DEPLOY_VERSION = \"tf-cpu.{}\".format(TF)\n",
         "\n",
         "TRAIN_IMAGE = \"{}-docker.pkg.dev/vertex-ai/training/{}:latest\".format(\n",
         "    REGION.split(\"-\")[0], TRAIN_VERSION\n",
         ")\n",
+        "DEPLOY_IMAGE = \"{}-docker.pkg.dev/vertex-ai/prediction/{}:latest\".format(\n",
+        "    REGION.split(\"-\")[0], DEPLOY_VERSION\n",
+        ")\n",
         "\n",
-        "print(\"Training:\", TRAIN_IMAGE, TRAIN_GPU, TRAIN_NGPU)"
-      ]
+        "print(\"Training:\", TRAIN_IMAGE, TRAIN_GPU, TRAIN_NGPU)\n",
+        "print(\"Deployment:\", DEPLOY_IMAGE, DEPLOY_GPU, DEPLOY_NGPU)"
     },
     {
       "cell_type": "markdown",
@@ -732,7 +759,7 @@
       "cell_type": "code",
       "execution_count": null,
       "metadata": {
-        "id": "custom_pp_training_job:simple"
+        "id": "create_custom_pp_training_job:mbsdk,no_model"
       },
       "outputs": [],
       "source": [
@@ -793,7 +820,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        '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        'tensorflow==2.5.0',\\n\\n        'tensorflow_datasets==1.3.0',\\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: Boston Housing tabular regression\\n\\nVersion: 0.0.0\\n\\nSummary: Demostration training script\\n\\nHome-page: www.google.com\\n\\nAuthor: Google\\n\\nAuthor-email: aferlitsch@google.com\\n\\nLicense: Public\\n\\nDescription: Demo\\n\\nPlatform: Vertex\"\n",
@@ -990,7 +1017,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        '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        'tensorflow==2.5.0',\\n\\n        'tensorflow_datasets==1.3.0',\\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: Boston Housing tabular regression\\n\\nVersion: 0.0.0\\n\\nSummary: Demostration training script\\n\\nHome-page: www.google.com\\n\\nAuthor: Google\\n\\nAuthor-email: aferlitsch@google.com\\n\\nLicense: Public\\n\\nDescription: Demo\\n\\nPlatform: Vertex\"\n",
@@ -1143,7 +1170,7 @@
       },
       "outputs": [],
       "source": [
-        "! docker run $TRAIN_IMAGE --epochs=5"
+        "! docker run $TRAIN_IMAGE --epochs=5 --model-dir=./"
       ]
     },
     {
@@ -1168,6 +1195,46 @@
         "! docker push $TRAIN_IMAGE"
       ]
     },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "create_custom_container_training_job:mbsdk,no_model"
+      },
+      "source": [
+        "### Create and run custom training job\n",
+        "\n",
+        "\n",
+        "To train a custom model, you perform two steps: 1) create a custom training job, and 2) run the job.\n",
+        "\n",
+        "#### Create custom training job\n",
+        "\n",
+        "A custom training job is created with the `CustomTrainingJob` class, with the following parameters:\n",
+        "\n",
+        "- `display_name`: The human readable name for the custom training job.\n",
+        "- `container_uri`: The training container image.\n",
+        "\n",
+        "- `command`: The command (e.g., interpreter) and script to invokee within the container.\n",
+        "\n",
+        "*Note:* The interpreter and script to invoke is overridable within the container (i.e., ENTRYPOINT)."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "create_custom_container_training_job:mbsdk,no_model"
+      },
+      "outputs": [],
+      "source": [
+        "job = aip.CustomContainerTrainingJob(\n",
+        "    display_name=\"boston_\" + TIMESTAMP,\n",
+        "    container_uri=TRAIN_IMAGE,\n",
+        "    command=[\"python3\", \"trainer/task.py\"],\n",
+        ")\n",
+        "\n",
+        "print(job)"
+      ]
+    },
     {
       "cell_type": "markdown",
       "metadata": {
@@ -1216,6 +1283,467 @@
         "job.delete()"
       ]
     },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "make_model_custom_training_intro"
+      },
+      "source": [
+        "## Model saving and conversion to Model Resource.\n",
+        "\n",
+        "Vertex Training will optional automatically convert your trained model into a Vertex Model resource, by adding the following additional steps:\n",
+        "\n",
+        "- Specify the following additional parameters when creating the custom training job (e.g., CustomTrainingJob):\n",
+        "\n",
+        "    - `base_output_dir`: The Cloud Storage location for storing the trained model artifacts.\n",
+        "    - `model_serving_container_image_uri`: The deployment container for the model.\n",
+        "\n",
+        "- Specify the following additional parameters when running the job:\n",
+        "\n",
+        "    - `model_display_name`: The display name for the Model resource.\n",
+        "\n",
+        "- In the training script, save the model to the location specified by the environment variable `AIP_MODEL_DIR`.\n",
+        "\n",
+        "The service will set the value of the environment variable `AIP_MODEL_DIR` relative to the Cloud Storage location of `base_output_dir`. Upon successful completion of the training script, the service will look at this location for the model artifacts, and upload the model artifacts into a Vertex Model resource.\n",
+        "\n",
+        "The response from the `run()` method will be the Model resource, when ran synchronously. When ran asynchronously, the response is a LRO which one can do a wait() on. Upon completion, the response object is redirected to the Model resource."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "taskpy_contents:boston"
+      },
+      "source": [
+        "#### Task.py contents\n",
+        "\n",
+        "In the next cell, you write the contents of the training script task.py. I won't go into detail, it's just there for you to browse. In summary:\n",
+        "\n",
+        "- Get the directory where to save the model artifacts from the command line (`--model_dir`), and if not specified, then from the environment variable `AIP_MODEL_DIR`.\n",
+        "- Loads Boston Housing dataset from TF.Keras builtin datasets\n",
+        "- Builds a simple deep neural network model using TF.Keras model API.\n",
+        "- Compiles the model (`compile()`).\n",
+        "- Sets a training distribution strategy according to the argument `args.distribute`.\n",
+        "- Trains the model (`fit()`) with epochs specified by `args.epochs`.\n",
+        "- Saves the trained model (`save(args.model_dir)`) to the specified model directory.\n",
+        "- Saves the maximum value for each feature `f.write(str(params))` to the specified parameters file."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "taskpy_contents:boston"
+      },
+      "outputs": [],
+      "source": [
+        "%%writefile custom/trainer/task.py\n",
+        "# Single, Mirror and Multi-Machine Distributed Training for Boston Housing\n",
+        "\n",
+        "import tensorflow_datasets as tfds\n",
+        "import tensorflow as tf\n",
+        "from tensorflow.python.client import device_lib\n",
+        "import numpy as np\n",
+        "import argparse\n",
+        "import os\n",
+        "import sys\n",
+        "tfds.disable_progress_bar()\n",
+        "\n",
+        "parser = argparse.ArgumentParser()\n",
+        "parser.add_argument('--model-dir', dest='model_dir',\n",
+        "                    default=os.getenv('AIP_MODEL_DIR'), type=str, help='Model dir.')\n",
+        "parser.add_argument('--lr', dest='lr',\n",
+        "                    default=0.001, type=float,\n",
+        "                    help='Learning rate.')\n",
+        "parser.add_argument('--epochs', dest='epochs',\n",
+        "                    default=20, type=int,\n",
+        "                    help='Number of epochs.')\n",
+        "parser.add_argument('--steps', dest='steps',\n",
+        "                    default=100, type=int,\n",
+        "                    help='Number of steps per epoch.')\n",
+        "parser.add_argument('--distribute', dest='distribute', type=str, default='single',\n",
+        "                    help='distributed training strategy')\n",
+        "parser.add_argument('--param-file', dest='param_file',\n",
+        "                    default='/tmp/param.txt', type=str,\n",
+        "                    help='Output file for parameters')\n",
+        "args = parser.parse_args()\n",
+        "\n",
+        "print('Python Version = {}'.format(sys.version))\n",
+        "print('TensorFlow Version = {}'.format(tf.__version__))\n",
+        "print('TF_CONFIG = {}'.format(os.environ.get('TF_CONFIG', 'Not found')))\n",
+        "\n",
+        "# Single Machine, single compute device\n",
+        "if args.distribute == 'single':\n",
+        "    if tf.test.is_gpu_available():\n",
+        "        strategy = tf.distribute.OneDeviceStrategy(device=\"/gpu:0\")\n",
+        "    else:\n",
+        "        strategy = tf.distribute.OneDeviceStrategy(device=\"/cpu:0\")\n",
+        "# Single Machine, multiple compute device\n",
+        "elif args.distribute == 'mirror':\n",
+        "    strategy = tf.distribute.MirroredStrategy()\n",
+        "# Multiple Machine, multiple compute device\n",
+        "elif args.distribute == 'multi':\n",
+        "    strategy = tf.distribute.experimental.MultiWorkerMirroredStrategy()\n",
+        "\n",
+        "# Multi-worker configuration\n",
+        "print('num_replicas_in_sync = {}'.format(strategy.num_replicas_in_sync))\n",
+        "\n",
+        "\n",
+        "def make_dataset():\n",
+        "\n",
+        "  # Scaling Boston Housing data features\n",
+        "  def scale(feature):\n",
+        "    max = np.max(feature)\n",
+        "    feature = (feature / max).astype(np.float)\n",
+        "    return feature, max\n",
+        "\n",
+        "  (x_train, y_train), (x_test, y_test) = tf.keras.datasets.boston_housing.load_data(\n",
+        "    path=\"boston_housing.npz\", test_split=0.2, seed=113\n",
+        "  )\n",
+        "  params = []\n",
+        "  for _ in range(13):\n",
+        "    x_train[_], max = scale(x_train[_])\n",
+        "    x_test[_], _ = scale(x_test[_])\n",
+        "    params.append(max)\n",
+        "\n",
+        "  # store the normalization (max) value for each feature\n",
+        "  with tf.io.gfile.GFile(args.param_file, 'w') as f:\n",
+        "    f.write(str(params))\n",
+        "  return (x_train, y_train), (x_test, y_test)\n",
+        "\n",
+        "\n",
+        "# Build the Keras model\n",
+        "def build_and_compile_dnn_model():\n",
+        "  model = tf.keras.Sequential([\n",
+        "      tf.keras.layers.Dense(128, activation='relu', input_shape=(13,)),\n",
+        "      tf.keras.layers.Dense(128, activation='relu'),\n",
+        "      tf.keras.layers.Dense(1, activation='linear')\n",
+        "  ])\n",
+        "  model.compile(\n",
+        "      loss='mse',\n",
+        "      optimizer=tf.keras.optimizers.RMSprop(learning_rate=args.lr))\n",
+        "  return model\n",
+        "\n",
+        "NUM_WORKERS = strategy.num_replicas_in_sync\n",
+        "# Here the batch size scales up by number of workers since\n",
+        "# `tf.data.Dataset.batch` expects the global batch size.\n",
+        "BATCH_SIZE = 16\n",
+        "GLOBAL_BATCH_SIZE = BATCH_SIZE * NUM_WORKERS\n",
+        "\n",
+        "with strategy.scope():\n",
+        "  # Creation of dataset, and model building/compiling need to be within\n",
+        "  # `strategy.scope()`.\n",
+        "  model = build_and_compile_dnn_model()\n",
+        "\n",
+        "# Train the model\n",
+        "(x_train, y_train), (x_test, y_test) = make_dataset()\n",
+        "model.fit(x_train, y_train, epochs=args.epochs, batch_size=GLOBAL_BATCH_SIZE)\n",
+        "model.save(args.model_dir)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "tarball_training_script"
+      },
+      "source": [
+        "#### Store training script on your Cloud Storage bucket\n",
+        "\n",
+        "Next, you package the training folder into a compressed tar ball, and then store it in your Cloud Storage bucket."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "tarball_training_script"
+      },
+      "outputs": [],
+      "source": [
+        "! rm -f custom.tar custom.tar.gz\n",
+        "! tar cvf custom.tar custom\n",
+        "! gzip custom.tar\n",
+        "! gsutil cp custom.tar.gz $BUCKET_NAME/trainer_boston.tar.gz"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "container:training"
+      },
+      "source": [
+        "#### Set pre-built containers\n",
+        "\n",
+        "Set the pre-built Docker container image for training.\n",
+        "\n",
+        "- Set the variable `TF` to the TensorFlow version of the container image. For example, `2-1` would be version 2.1, and `1-15` would be version 1.15. The following list shows some of the pre-built images available:\n",
+        "\n",
+        "\n",
+        "For the latest list, see [Pre-built containers for training](https://cloud.google.com/ai-platform-unified/docs/training/pre-built-containers)."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "container:training"
+      },
+      "outputs": [],
+      "source": [
+        "if os.getenv(\"IS_TESTING_TF\"):\n",
+        "    TF = os.getenv(\"IS_TESTING_TF\")\n",
+        "else:\n",
+        "    TF = \"2.1\".replace(\".\", \"-\")\n",
+        "\n",
+        "if TF[0] == \"2\":\n",
+        "    if TRAIN_GPU:\n",
+        "        TRAIN_VERSION = \"tf-gpu.{}\".format(TF)\n",
+        "    else:\n",
+        "        TRAIN_VERSION = \"tf-cpu.{}\".format(TF)\n",
+        "else:\n",
+        "    if TRAIN_GPU:\n",
+        "        TRAIN_VERSION = \"tf-gpu.{}\".format(TF)\n",
+        "    else:\n",
+        "        TRAIN_VERSION = \"tf-cpu.{}\".format(TF)\n",
+        "\n",
+        "TRAIN_IMAGE = \"{}-docker.pkg.dev/vertex-ai/training/{}:latest\".format(\n",
+        "    REGION.split(\"-\")[0], TRAIN_VERSION\n",
+        ")\n",
+        "\n",
+        "print(\"Training:\", TRAIN_IMAGE, TRAIN_GPU, TRAIN_NGPU)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "create_custom_pp_training_job:mbsdk"
+      },
+      "source": [
+        "### Create and run custom training job\n",
+        "\n",
+        "\n",
+        "To train a custom model, you perform two steps: 1) create a custom training job, and 2) run the job.\n",
+        "\n",
+        "#### Create custom training job\n",
+        "\n",
+        "A custom training job is created with the `CustomTrainingJob` class, with the following parameters:\n",
+        "\n",
+        "- `display_name`: The human readable name for the custom training job.\n",
+        "- `container_uri`: The training container image.\n",
+        "\n",
+        "- `python_package_gcs_uri`: The location of the Python training package as a tarball.\n",
+        "- `python_module_name`: The relative path to the training script in the Python package.\n",
+        "- `model_serving_container_uri`: The container image for deploying the model.\n",
+        "\n",
+        "*Note:* There is no requirements parameter. You specify any requirements in the `setup.py` script in your Python package."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "create_custom_pp_training_job:mbsdk"
+      },
+      "outputs": [],
+      "source": [
+        "DISPLAY_NAME = \"boston_\" + TIMESTAMP\n",
+        "\n",
+        "job = aip.CustomPythonPackageTrainingJob(\n",
+        "    display_name=DISPLAY_NAME,\n",
+        "    python_package_gcs_uri=f\"{BUCKET_NAME}/trainer_boston.tar.gz\",\n",
+        "    python_module_name=\"trainer.task\",\n",
+        "    container_uri=TRAIN_IMAGE,\n",
+        "    model_serving_container_image_uri=DEPLOY_IMAGE,\n",
+        "    project=PROJECT_ID,\n",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "prepare_custom_cmdargs"
+      },
+      "source": [
+        "### Prepare your command-line arguments\n",
+        "\n",
+        "Now define the command-line arguments for your custom training container:\n",
+        "\n",
+        "- `args`: The command-line arguments to pass to the executable that is set as the entry point into the container.\n",
+        "  - `--model-dir` : For our demonstrations, we use this command-line argument to specify where to store the model artifacts.\n",
+        "      - direct: You pass the Cloud Storage location as a command line argument to your training script (set variable `DIRECT = True`), or\n",
+        "      - indirect: The service passes the Cloud Storage location as the environment variable `AIP_MODEL_DIR` to your training script (set variable `DIRECT = False`). In this case, you tell the service the model artifact location in the job specification.\n",
+        "  - `\"--epochs=\" + EPOCHS`: The number of epochs for training.\n",
+        "  - `\"--steps=\" + STEPS`: The number of steps per epoch."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "prepare_custom_cmdargs"
+      },
+      "outputs": [],
+      "source": [
+        "MODEL_DIR = \"{}/{}\".format(BUCKET_NAME, TIMESTAMP)\n",
+        "\n",
+        "EPOCHS = 20\n",
+        "STEPS = 100\n",
+        "\n",
+        "DIRECT = True\n",
+        "if DIRECT:\n",
+        "    CMDARGS = [\n",
+        "        \"--model-dir=\" + MODEL_DIR,\n",
+        "        \"--epochs=\" + str(EPOCHS),\n",
+        "        \"--steps=\" + str(STEPS),\n",
+        "    ]\n",
+        "else:\n",
+        "    CMDARGS = [\n",
+        "        \"--epochs=\" + str(EPOCHS),\n",
+        "        \"--steps=\" + str(STEPS),\n",
+        "    ]"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "run_custom_job:mbsdk"
+      },
+      "source": [
+        "#### Run the custom training job\n",
+        "\n",
+        "Next, you run the custom job to start the training job by invoking the method `run`, with the following parameters:\n",
+        "\n",
+        "- `model_display_name`: The human readable name for the `Model` resource.\n",
+        "- `args`: The command-line arguments to pass to the training script.\n",
+        "- `replica_count`: The number of compute instances for training (replica_count = 1 is single node training).\n",
+        "- `machine_type`: The machine type for the compute instances.\n",
+        "- `accelerator_type`: The hardware accelerator type.\n",
+        "- `accelerator_count`: The number of accelerators to attach to a worker replica.\n",
+        "- `base_output_dir`: The Cloud Storage location to write the model artifacts to.\n",
+        "- `sync`: Whether to block until completion of the job."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "run_custom_job:mbsdk"
+      },
+      "outputs": [],
+      "source": [
+        "if TRAIN_GPU:\n",
+        "    model = job.run(\n",
+        "        model_display_name=\"boston_\" + TIMESTAMP,\n",
+        "        args=CMDARGS,\n",
+        "        replica_count=1,\n",
+        "        machine_type=TRAIN_COMPUTE,\n",
+        "        accelerator_type=TRAIN_GPU.name,\n",
+        "        accelerator_count=TRAIN_NGPU,\n",
+        "        base_output_dir=MODEL_DIR,\n",
+        "        sync=False,\n",
+        "    )\n",
+        "else:\n",
+        "    model = job.run(\n",
+        "        model_display_name=\"boston_\" + TIMESTAMP,\n",
+        "        args=CMDARGS,\n",
+        "        replica_count=1,\n",
+        "        machine_type=TRAIN_COMPUTE,\n",
+        "        base_output_dir=MODEL_DIR,\n",
+        "        sync=False,\n",
+        "    )\n",
+        "\n",
+        "model_path_to_deploy = MODEL_DIR"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "list_job"
+      },
+      "source": [
+        "### List a custom training job"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "list_job"
+      },
+      "outputs": [],
+      "source": [
+        "_job = job.list(filter=f\"display_name={DISPLAY_NAME}\")\n",
+        "print(_job)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "custom_job_wait:mbsdk"
+      },
+      "source": [
+        "### Wait for completion of custom training job\n",
+        "\n",
+        "Next, wait for the custom training job to complete. Alternatively, one can set the parameter `sync` to `True` in the `run()` method to block until the custom training job is completed."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "custom_job_wait:mbsdk"
+      },
+      "outputs": [],
+      "source": [
+        "model.wait()"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "train_script_layout"
+      },
+      "source": [
+        "## Custom training script recommendations\n",
+        "\n",
+        "As you transition from experimenting to formalization, it is common to construct training scripts in a manner that makes them easier to convert to automated pipelines. The typical practice is to separate different tasks within the custom training script into separate Python modules, such as:\n",
+        "\n",
+        "- `data.py`: The data pipeline.\n",
+        "- `model.py`: Build or retrieve model architecture.\n",
+        "- `train.py': The training pipeline.\n",
+        "- `eval.py`: Evaluate the train model.\n",
+        "- `task.py`: Entry point and script driver."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "train_script_layout"
+      },
+      "outputs": [],
+      "source": [
+        "%%writefile custom/trainer/task.py\n",
+        "from . import data\n",
+        "from . import model\n",
+        "from . import train\n",
+        "from . import eval\n",
+        "\n",
+        "parser = argparse.ArgumentParser()\n",
+        "# ...\n",
+        "args = parser.parse_args()\n",
+        "\n",
+        "train, val, test = data.get(args.batch_size, ...)\n",
+        "\n",
+        "model = model.get(args.lr, ... )\n",
+        "\n",
+        "history = train.train(model, train, val, args.epochs, ...)\n",
+        "\n",
+        "metrics = eval.evaluate(model, eval)\n",
+        "\n",
+        "model.save(args.model_dir)"
+      ]
+    },
     {
       "cell_type": "markdown",
       "metadata": {
@@ -1272,14 +1800,14 @@
         "    except Exception as e:\n",
         "        print(e)\n",
         "\n",
-        "    # Delete the AutoML or Pipeline trainig job\n",
+        "    # Delete the AutoML or Pipeline training job\n",
         "    try:\n",
         "        if \"dag\" in globals():\n",
         "            dag.delete()\n",
         "    except Exception as e:\n",
         "        print(e)\n",
         "\n",
-        "    # Delete the custom trainig job\n",
+        "    # Delete the custom training job\n",
         "    try:\n",
         "        if \"job\" in globals():\n",
         "            job.delete()\n",