examples/training/02_tensorflow_transfer_learning_gpu_example.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "12b45b0b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Copyright 2024 Google LLC\n",
    "#\n",
    "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
    "# you may not use this file except in compliance with the License.\n",
    "# You may obtain a copy of the License at\n",
    "#\n",
    "#     https://www.apache.org/licenses/LICENSE-2.0\n",
    "#\n",
    "# Unless required by applicable law or agreed to in writing, software\n",
    "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
    "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
    "# See the License for the specific language governing permissions and\n",
    "# limitations under the License."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8790d7ed",
   "metadata": {},
   "source": [
    "# AutoMLOps - Tensorflow Transfer Learning GPU Example\n",
    "\n",
    "<table align=\"left\">\n",
    "  <td>\n",
    "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/automlops/blob/main/examples/training/02_tensorflow_transfer_learning_gpu_example.ipynb\">\n",
    "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Colab logo\"> Run in Colab\n",
    "    </a>\n",
    "  </td>\n",
    "  <td>\n",
    "    <a href=\"https://github.com/GoogleCloudPlatform/automlops/blob/main/examples/training/02_tensorflow_transfer_learning_gpu_example.ipynb\">\n",
    "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\">\n",
    "      View on GitHub\n",
    "    </a>\n",
    "  </td>\n",
    "  <td>\n",
    "    <a href=\"https://console.cloud.google.com/vertex-ai/workbench/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/automlops/main/examples/training/02_tensorflow_transfer_learning_gpu_example.ipynb\">\n",
    "        <img src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" alt=\"Vertex AI logo\">\n",
    "      Open in Vertex AI Workbench\n",
    "    </a>\n",
    "  </td>\n",
    "</table>\n",
    "<br/><br/><br/>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6f938540",
   "metadata": {},
   "source": [
    "# Overview\n",
    "\n",
    "In this tutorial you'll use transfer learning to train an image classification model on the cassava dataset from TensorFlow Datasets. The architecture you'll use is a ResNet50 model from the tf.keras.applications library pretrained on the Imagenet dataset. This tutorial will walk you through how to use AutoMLOps to define, create and run a MLOps pipeline around this model training. For uptraining the ResNet50 model, we will use a GPU."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "22881a6a",
   "metadata": {},
   "source": [
    "# Objective\n",
    "In this tutorial, you will learn how to create and run MLOps pipelines integrated with CI/CD. This tutorial goes through training a tensorflow model using accelerators; the pipeline goes through the following workflow:\n",
    "1. importer: Google cloud pipeline component for importing tensorflow models into Vertex Model Registry\n",
    "2. custom_train_model: A custom component that trains a tensorflow model.\n",
    "3. model_upload: Google cloud pipeline component that executes an upload operation.\n",
    "\n",
    "# Prerequisites\n",
    "\n",
    "In order to use AutoMLOps, the following are required:\n",
    "\n",
    "- Python 3.7 - 3.10\n",
    "- [Google Cloud SDK 407.0.0](https://cloud.google.com/sdk/gcloud/reference)\n",
    "- [beta 2022.10.21](https://cloud.google.com/sdk/gcloud/reference/beta)\n",
    "- `git` installed\n",
    "- `git` logged-in:\n",
    "```\n",
    "  git config --global user.email \"you@example.com\"\n",
    "  git config --global user.name \"Your Name\"\n",
    "```\n",
    "- [Application Default Credentials (ADC)](https://cloud.google.com/docs/authentication/provide-credentials-adc) are setup. This can be done through the following commands:\n",
    "```\n",
    "gcloud auth application-default login\n",
    "gcloud config set account <account@example.com>\n",
    "```\n",
    "\n",
    "# APIs & IAM\n",
    "Based on the user options selection, AutoMLOps will enable up to the following APIs during the provision step:\n",
    "- [aiplatform.googleapis.com](https://cloud.google.com/vertex-ai/docs/reference/rest)\n",
    "- [artifactregistry.googleapis.com](https://cloud.google.com/artifact-registry/docs/reference/rest)\n",
    "- [cloudbuild.googleapis.com](https://cloud.google.com/build/docs/api/reference/rest)\n",
    "- [cloudfunctions.googleapis.com](https://cloud.google.com/functions/docs/reference/rest)\n",
    "- [cloudresourcemanager.googleapis.com](https://cloud.google.com/resource-manager/reference/rest)\n",
    "- [cloudscheduler.googleapis.com](https://cloud.google.com/scheduler/docs/reference/rest)\n",
    "- [compute.googleapis.com](https://cloud.google.com/compute/docs/reference/rest/v1)\n",
    "- [iam.googleapis.com](https://cloud.google.com/iam/docs/reference/rest)\n",
    "- [iamcredentials.googleapis.com](https://cloud.google.com/iam/docs/reference/credentials/rest)\n",
    "- [logging.googleapis.com](https://cloud.google.com/logging/docs/reference/v2/rest)\n",
    "- [pubsub.googleapis.com](https://cloud.google.com/pubsub/docs/reference/rest)\n",
    "- [run.googleapis.com](https://cloud.google.com/run/docs/reference/rest)\n",
    "- [storage.googleapis.com](https://cloud.google.com/storage/docs/apis)\n",
    "- [sourcerepo.googleapis.com](https://cloud.google.com/source-repositories/docs/reference/rest)\n",
    "\n",
    "\n",
    "AutoMLOps will create the following service account and update [IAM permissions](https://cloud.google.com/iam/docs/understanding-roles) during the provision step:\n",
    "1. Pipeline Runner Service Account (defaults to: vertex-pipelines@PROJECT_ID.iam.gserviceaccount.com). Roles added:\n",
    "- roles/aiplatform.user\n",
    "- roles/artifactregistry.reader\n",
    "- roles/bigquery.user\n",
    "- roles/bigquery.dataEditor\n",
    "- roles/iam.serviceAccountUser\n",
    "- roles/storage.admin\n",
    "- roles/cloudfunctions.admin\n",
    "\n",
    "# User Guide\n",
    "\n",
    "For a user-guide, please view these [slides](../../AutoMLOps_User_Guide.pdf).\n",
    "\n",
    "# Costs\n",
    "\n",
    "This tutorial uses billable components of Google Cloud:\n",
    "- Vertex AI\n",
    "- Artifact Registry\n",
    "- Cloud Storage\n",
    "- Cloud Source Repository\n",
    "- Cloud Build\n",
    "- Cloud Run\n",
    "- Cloud Scheduler\n",
    "- Cloud Pub/Sub\n",
    "\n",
    "Learn about [Vertex AI pricing](https://cloud.google.com/vertex-ai/pricing), and use the [Pricing Calculator](https://cloud.google.com/products/calculator/) to generate a cost estimate based on your projected usage.\n",
    "\n",
    "# Ground-rules for using AutoMLOps\n",
    "1. Do not use variables, functions, code, etc. not defined within the scope of a custom component. These custom components will become containers and will have no reference to the out of scope code.\n",
    "2. Import statements and helper functions must be added inside the function. Provide parameter type hints.\n",
    "3. Test each of your components for accuracy and correctness before running them using AutoMLOps. We cannot fix bugs automatically; bugs are much more difficult to fix once they are made into pipelines.\n",
    "4. If you are using Kubeflow, be sure to define all the requirements needed to run the custom component - it can be easy to leave out packages which will cause the container to fail when running within a pipeline. \n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "12381413",
   "metadata": {},
   "source": [
    "# Dataset\n",
    "For training data, we are using the [cassava dataset](https://www.tensorflow.org/datasets/catalog/cassava) from [TensorFlow Datasets](https://www.tensorflow.org/datasets). This dataset consists of leaf images for the cassava plant depicting healthy and four (4) disease conditions; Cassava Mosaic Disease (CMD), Cassava Bacterial Blight (CBB), Cassava Greem Mite (CGM) and Cassava Brown Streak Disease (CBSD). Dataset consists of a total of 9430 labelled images. The 9430 labelled images are split into a training set (5656), a test set(1885) and a validation set (1889). The number of images per class are unbalanced with the two disease classes CMD and CBSD having 72% of the images."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a231b629",
   "metadata": {},
   "source": [
    "# Setup Git\n",
    "Set up your git configuration below"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "92f90b44",
   "metadata": {},
   "outputs": [],
   "source": [
    "!git config --global user.email 'you@example.com'\n",
    "!git config --global user.name 'Your Name'"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d5e4d190",
   "metadata": {},
   "source": [
    "# Install AutoMLOps\n",
    "\n",
    "Install AutoMLOps from [PyPI](https://pypi.org/project/google-cloud-automlops/), or locally by cloning the repo and running `pip install .`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "94451868",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "!pip3 install google-cloud-automlops --user"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "17db55d5",
   "metadata": {},
   "source": [
    "# Restart the kernel\n",
    "Once you've installed the AutoMLOps package, you need to restart the notebook kernel so it can find the package.\n",
    "\n",
    "**Note: Once this cell has finished running, continue on. You do not need to re-run any of the cells above.**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "57c53b86",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "if not os.getenv('IS_TESTING'):\n",
    "    # Automatically restart kernel after installs\n",
    "    import IPython\n",
    "\n",
    "    app = IPython.Application.instance()\n",
    "    app.kernel.do_shutdown(True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "250d511b",
   "metadata": {},
   "source": [
    "# Set variables\n",
    "Set variables. If you don't know your project ID, leave the field blank and the following cells may be able to find it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "931ff517",
   "metadata": {},
   "outputs": [],
   "source": [
    "PROJECT_ID = '[your-project-id]'  # @param {type:\"string\"}\n",
    "\n",
    "BUCKET_NAME = 'automlops-sandbox-bucket'  # @param {type:\"string\"}\n",
    "BUCKET_URI = f'gs://{BUCKET_NAME}'\n",
    "MODEL_DIR = BUCKET_URI + '/tensorflow_model'\n",
    "\n",
    "TRAINING_IMAGE = 'us-docker.pkg.dev/vertex-ai/training/tf-gpu.2-11.py310:latest' # includes required cuda packages\n",
    "SERVING_IMAGE = 'us-docker.pkg.dev/vertex-ai/prediction/tf-gpu.2-11.py310:latest'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "be0be295",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Project ID: automlops-sandbox\n"
     ]
    }
   ],
   "source": [
    "if PROJECT_ID == '' or PROJECT_ID is None or PROJECT_ID == '[your-project-id]':\n",
    "    # Get your GCP project id from gcloud\n",
    "    shell_output = !gcloud config list --format 'value(core.project)' 2>/dev/null\n",
    "    PROJECT_ID = shell_output[0]\n",
    "    print('Project ID:', PROJECT_ID)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "2c36482d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Updated property [core/project].\r\n"
     ]
    }
   ],
   "source": [
    "! gcloud config set project $PROJECT_ID"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "557dd504",
   "metadata": {},
   "source": [
    "Set your Model_ID below:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "92019c46",
   "metadata": {},
   "outputs": [],
   "source": [
    "MODEL_ID = 'cassava-resnet-50'"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4ecec8ba",
   "metadata": {},
   "source": [
    "# AutoMLOps Tensorflow Example\n",
    "This workflow will generate a pipeline using Kubeflow spec. AutoMLOps provides 2 functions for defining MLOps pipelines:\n",
    "\n",
    "- `AutoMLOps.component(...)`: Defines a component, which is a containerized python function.\n",
    "- `AutoMLOps.pipeline(...)`: Defines a pipeline, which is a series of components.\n",
    "\n",
    "AutoMLOps provides 6 functions for building and maintaining MLOps pipelines:\n",
    "\n",
    "- `AutoMLOps.generate(...)`: Generates the MLOps codebase. Users can specify the tooling and technologies they would like to use in their MLOps pipeline.\n",
    "- `AutoMLOps.provision(...)`: Runs provisioning scripts to create and maintain necessary infra for MLOps.\n",
    "- `AutoMLOps.deprovision(...)`: Runs deprovisioning scripts to tear down MLOps infra created using AutoMLOps.\n",
    "- `AutoMLOps.deploy(...)`: Builds and pushes component container, then triggers the pipeline job.\n",
    "- `AutoMLOps.launchAll(...)`: Runs `generate()`, `provision()`, and `deploy()` all in succession.\n",
    "- `AutoMLOps.monitor(...)`: Creates model monitoring jobs on deployed endpoints.\n",
    "\n",
    "Please see the [readme](https://github.com/GoogleCloudPlatform/automlops/blob/main/README.md) for more information.\n",
    "\n",
    "**Note: This workflow requires python package `kfp>=2.0.0`.**"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a55db8c4",
   "metadata": {},
   "source": [
    "## Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "781e2b21",
   "metadata": {},
   "outputs": [],
   "source": [
    "from kfp.dsl import Metrics, Model, Output\n",
    "from google_cloud_automlops import AutoMLOps"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c51d294e",
   "metadata": {},
   "source": [
    "## Model Training\n",
    "Define a custom component for training a model. The architecture you'll use is a ResNet50 model from the tf.keras.applications library pretrained on the [Imagenet dataset](https://www.image-net.org/)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "7f6fd905",
   "metadata": {},
   "outputs": [],
   "source": [
    "@AutoMLOps.component(\n",
    "    packages_to_install=[\n",
    "        'tensorflow',\n",
    "        'tensorflow_datasets',\n",
    "        'opencv-python-headless'\n",
    "    ]\n",
    ")\n",
    "def custom_train_model(\n",
    "    metrics: Output[Metrics],\n",
    "    model_dir: str,\n",
    "    output_model: Output[Model],\n",
    "    lr: float = 0.001,\n",
    "    epochs: int = 10,\n",
    "    steps: int = 200,\n",
    "    distribute: str = 'single'\n",
    "):\n",
    "    import faulthandler\n",
    "    import os\n",
    "    import sys\n",
    "\n",
    "    import tensorflow as tf\n",
    "    import tensorflow_datasets as tfds\n",
    "    from tensorflow.python.client import device_lib\n",
    "\n",
    "    faulthandler.enable()\n",
    "    tfds.disable_progress_bar()\n",
    "\n",
    "    print('Component start')\n",
    "\n",
    "    print(f'Python Version = {sys.version}')\n",
    "    print(f'TensorFlow Version = {tf.__version__}')\n",
    "    print(f'''TF_CONFIG = {os.environ.get('TF_CONFIG', 'Not found')}''')\n",
    "    print(f'DEVICES = {device_lib.list_local_devices()}')\n",
    "\n",
    "    # Single Machine, single compute device\n",
    "    if 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 distribute == 'mirror':\n",
    "        strategy = tf.distribute.MirroredStrategy()\n",
    "    # Multiple Machine, multiple compute device\n",
    "    elif distribute == 'multi':\n",
    "        strategy = tf.distribute.experimental.MultiWorkerMirroredStrategy()\n",
    "\n",
    "    # Multi-worker configuration\n",
    "    print(f'num_replicas_in_sync = {strategy.num_replicas_in_sync}')\n",
    "\n",
    "    # Preparing dataset\n",
    "    BUFFER_SIZE = 10000\n",
    "    BATCH_SIZE = 64\n",
    "\n",
    "    def preprocess_data(image, label):\n",
    "        '''Resizes and scales images.'''\n",
    "\n",
    "        image = tf.image.resize(image, (300,300))\n",
    "        return tf.cast(image, tf.float32) / 255., label\n",
    "\n",
    "    def create_dataset(batch_size: int):\n",
    "        '''Loads Cassava dataset and preprocesses data.'''\n",
    "\n",
    "        data, info = tfds.load(name='cassava', as_supervised=True, with_info=True)\n",
    "        number_of_classes = info.features['label'].num_classes\n",
    "        train_data = data['train'].map(preprocess_data,\n",
    "                                       num_parallel_calls=tf.data.experimental.AUTOTUNE)\n",
    "        train_data  = train_data.cache().shuffle(BUFFER_SIZE).repeat()\n",
    "        train_data  = train_data.batch(batch_size)\n",
    "        train_data  = train_data.prefetch(tf.data.experimental.AUTOTUNE)\n",
    "\n",
    "        # Set AutoShardPolicy\n",
    "        options = tf.data.Options()\n",
    "        options.experimental_distribute.auto_shard_policy = tf.data.experimental.AutoShardPolicy.DATA\n",
    "        train_data = train_data.with_options(options)\n",
    "\n",
    "        return train_data, number_of_classes\n",
    "\n",
    "    # Build the ResNet50 Keras model    \n",
    "    def create_model(number_of_classes: int, lr: int = 0.001):\n",
    "        '''Creates and compiles pretrained ResNet50 model.'''\n",
    "\n",
    "        base_model = tf.keras.applications.ResNet50(weights='imagenet', include_top=False)\n",
    "        x = base_model.output\n",
    "        x = tf.keras.layers.GlobalAveragePooling2D()(x)\n",
    "        x = tf.keras.layers.Dense(1016, activation='relu')(x)\n",
    "        predictions = tf.keras.layers.Dense(number_of_classes, activation='softmax')(x)\n",
    "        model = tf.keras.Model(inputs=base_model.input, outputs=predictions)\n",
    "\n",
    "        model.compile(\n",
    "            loss=tf.keras.losses.sparse_categorical_crossentropy,\n",
    "            optimizer=tf.keras.optimizers.Adam(lr),\n",
    "            metrics=['accuracy'])\n",
    "        return model\n",
    "\n",
    "    # Train the model\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",
    "    GLOBAL_BATCH_SIZE = BATCH_SIZE * NUM_WORKERS\n",
    "    train_dataset, number_of_classes = create_dataset(GLOBAL_BATCH_SIZE)\n",
    "\n",
    "    with strategy.scope():\n",
    "        # Creation of dataset, and model building/compiling need to be within `strategy.scope()`.\n",
    "        resnet_model = create_model(number_of_classes, lr)\n",
    "\n",
    "    h = resnet_model.fit(x=train_dataset, epochs=epochs, steps_per_epoch=steps)\n",
    "    acc = h.history['accuracy'][-1]\n",
    "    resnet_model.save(model_dir)\n",
    "    \n",
    "    output_model.path = model_dir\n",
    "    metrics.log_metric('accuracy', (acc * 100.0))\n",
    "    metrics.log_metric('framework', 'Tensorflow')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3602954f",
   "metadata": {},
   "source": [
    "## Define the Pipeline\n",
    "Define your pipeline using `@AutoMLOps.pipeline`. You can optionally give the pipeline a name and description. Define the structure by listing the components to be called in your pipeline; use `.after` to specify the order of execution."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "01996d4c",
   "metadata": {},
   "outputs": [],
   "source": [
    "@AutoMLOps.pipeline(name='tensorflow-gpu-example')\n",
    "def pipeline(\n",
    "    project_id: str,\n",
    "    model_dir: str,\n",
    "    lr: float,\n",
    "    epochs: int,\n",
    "    steps: int,\n",
    "    serving_image: str,\n",
    "    distribute: str,\n",
    "):\n",
    "    from google_cloud_pipeline_components.types import artifact_types\n",
    "    from google_cloud_pipeline_components.v1.model import ModelUploadOp\n",
    "    from kfp.v2.components import importer_node\n",
    "\n",
    "    custom_train_model_task = custom_train_model(\n",
    "        model_dir=model_dir,\n",
    "        lr=lr,\n",
    "        epochs=epochs,\n",
    "        steps=steps,\n",
    "        distribute=distribute\n",
    "    )\n",
    "\n",
    "    unmanaged_model_importer = importer_node.importer(\n",
    "        artifact_uri=model_dir,\n",
    "        artifact_class=artifact_types.UnmanagedContainerModel,\n",
    "        metadata={\n",
    "            'containerSpec': {\n",
    "                'imageUri': serving_image\n",
    "            }\n",
    "        },\n",
    "    )\n",
    "\n",
    "    model_upload_op = ModelUploadOp(\n",
    "        project=project_id,\n",
    "        display_name='tensorflow_gpu_example',\n",
    "        unmanaged_container_model=unmanaged_model_importer.outputs['artifact'],\n",
    "    )\n",
    "    model_upload_op.after(custom_train_model_task)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6fe0b0f2",
   "metadata": {},
   "source": [
    "## Define the Pipeline Arguments"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "bc244ca7",
   "metadata": {},
   "outputs": [],
   "source": [
    "pipeline_params = {\n",
    "    'project_id': PROJECT_ID,\n",
    "    'model_dir': MODEL_DIR,\n",
    "    'lr': 0.01,\n",
    "    'epochs': 10,\n",
    "    'steps': 200,\n",
    "    'serving_image': SERVING_IMAGE,\n",
    "    'distribute': 'single'\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1b0c815c",
   "metadata": {},
   "source": [
    "## Generate and Run the pipeline\n",
    "`AutoMLOps.launchAll(...)` runs `generate()`, `provision()`, and `deploy()` all in succession. In this case, we are specifying a custom job spec, where we will use an Nvidia A100 GPU to accelerate the training of the model. \n",
    "\n",
    "*Note: if you run this cell below without a larger container, the training job will run out of memory and fail:*\n",
    "```\n",
    "The replica workerpool0-0 ran out-of-memory and exited with a non-zero status of 137(SIGKILL). To find out more about why your job exited please check the logs:\n",
    "```\n",
    "\n",
    "This use case is an ideal example for where specifying `custom_training_job_specs` for AutoMLOps is useful and necessary."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "25cfeff0-c5e9-4861-83f0-844dac5e153a",
   "metadata": {},
   "outputs": [],
   "source": [
    "AutoMLOps.launchAll(project_id=PROJECT_ID, \n",
    "                    pipeline_params=pipeline_params, \n",
    "                    use_ci=True,\n",
    "                    schedule_pattern='59 11 * * 0', # retrain every Sunday at Midnight\n",
    "                    base_image=TRAINING_IMAGE,\n",
    "                    naming_prefix=MODEL_ID,\n",
    "                    custom_training_job_specs = [{\n",
    "                       'component_spec': 'custom_train_model',\n",
    "                       'display_name': 'train-model-accelerated',\n",
    "                       'machine_type': 'a2-highgpu-1g',\n",
    "                       'accelerator_type': 'NVIDIA_TESLA_V100',\n",
    "                       'accelerator_count': 1\n",
    "                    }]\n",
    ")"
   ]
  }
 ],
 "metadata": {
  "environment": {
   "kernel": "python3",
   "name": "common-cpu.m102",
   "type": "gcloud",
   "uri": "gcr.io/deeplearning-platform-release/base-cpu:m102"
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.9.6"
  },
  "vscode": {
   "interpreter": {
    "hash": "31f2aee4e71d21fbe5cf8b01ff0e069b9275f58929596ceb00d14d90e3e16cd6"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}