notebooks/official/custom/custom_training_container_and_model_registry.ipynb

{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "copyright"
      },
      "outputs": [],
      "source": [
        "# Copyright 2021 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",
      "metadata": {
        "id": "title"
      },
      "source": [
        "# Custom training with custom container image and automatic model upload to Vertex AI Model Registry\n",
        "\n",
        "<table align=\"left\">\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/custom/custom_training_container_and_model_registry.ipynb\">\n",
        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Google Colaboratory logo\"><br> Open in Colab\n",
        "    </a>\n",
        "  </td>\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://console.cloud.google.com/vertex-ai/colab/import/https:%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%2Fvertex-ai-samples%2Fmain%2Fnotebooks%2Fofficial%2Fcustom%2Fcustom_training_container_and_model_registry.ipynb\">\n",
        "      <img width=\"32px\" src=\"https://cloud.google.com/ml-engine/images/colab-enterprise-logo-32px.png\" alt=\"Google Cloud Colab Enterprise logo\"><br> Open in Colab Enterprise\n",
        "    </a>\n",
        "  </td>    \n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://console.cloud.google.com/vertex-ai/workbench/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/notebooks/official/custom/custom_training_container_and_model_registry.ipynb\">\n",
        "      <img src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" alt=\"Vertex AI logo\"><br> Open in Workbench\n",
        "    </a>\n",
        "  </td>\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/custom/custom_training_container_and_model_registry.ipynb\">\n",
        "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\"><br> View on GitHub\n",
        "    </a>\n",
        "  </td>\n",
        "</table>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "overview:custom,training,custom_container"
      },
      "source": [
        "## Overview\n",
        "\n",
        "This tutorial demonstrates how to use the Vertex AI SDK for Python to train using a custom container image and automatically register the model in Vertex AI Model Registry.\n",
        "\n",
        "Learn more about [Custom training](https://cloud.google.com/vertex-ai/docs/training/custom-training)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "objective:custom,training,custom_container,online_prediction"
      },
      "source": [
        "### Objective\n",
        "\n",
        "In this tutorial, you train a machine learning model custom container image approach for custom training in Vertex AI. The trained model is further registered in the Vertex AI Model Registry automatically. You can alternatively create custom models using `gcloud` command-line tool or online using Cloud Console.\n",
        "\n",
        "This tutorial uses the following Vertex AI services and resources:\n",
        "\n",
        "- Vertex AI Model Registry\n",
        "- Vertex AI Training\n",
        "\n",
        "\n",
        "The steps performed include:\n",
        "\n",
        "- Create a Vertex AI custom job for training a model.\n",
        "- Train and register a TensorFlow model using a custom container.\n",
        "- List the registered model in the Vertex AI Model Registry."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dataset:custom,cifar10,icn"
      },
      "source": [
        "### Dataset\n",
        "\n",
        "The dataset used for this tutorial is the [CIFAR10 dataset](https://www.tensorflow.org/datasets/catalog/cifar10) from [TensorFlow Datasets](https://www.tensorflow.org/datasets/catalog/overview). The version of the dataset you use is built into TensorFlow. The trained model predicts which class an image is from the ten classes: airplane, automobile, bird, cat, deer, dog, frog, horse, ship, truck."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "costs"
      },
      "source": [
        "### Costs\n",
        "\n",
        "This tutorial uses billable components of Google Cloud:\n",
        "\n",
        "* Vertex AI\n",
        "* Cloud Storage\n",
        "* Artifact Registry\n",
        "* Cloud Build\n",
        "\n",
        "Learn about [Vertex AI\n",
        "pricing](https://cloud.google.com/vertex-ai/pricing), [Cloud Storage\n",
        "pricing](https://cloud.google.com/storage/pricing), [Artifact Registry pricing](https://cloud.google.com/artifact-registry/pricing), and [Cloud Build pricing](https://cloud.google.com/build/pricing) and use the [Pricing\n",
        "Calculator](https://cloud.google.com/products/calculator/)\n",
        "to generate a cost estimate based on your projected usage."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "f0316df526f8"
      },
      "source": [
        "## Get started"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "a2c2cb2109a0"
      },
      "source": [
        "### Install Vertex AI SDK for Python and other required packages\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "d679a60ed556"
      },
      "outputs": [],
      "source": [
        "! pip3 install --upgrade google-cloud-aiplatform --quiet"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "R5Xep4W9lq-Z"
      },
      "source": [
        "### Restart runtime (Colab only)\n",
        "\n",
        "To use the newly installed packages, you must restart the runtime on Google Colab."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "XRvKdaPDTznN"
      },
      "outputs": [],
      "source": [
        "import sys\n",
        "\n",
        "if \"google.colab\" in sys.modules:\n",
        "\n",
        "    import IPython\n",
        "\n",
        "    app = IPython.Application.instance()\n",
        "    app.kernel.do_shutdown(True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "SbmM4z7FOBpM"
      },
      "source": [
        "<div class=\"alert alert-block alert-warning\">\n",
        "<b>⚠️ The kernel is going to restart. Wait until it's finished before continuing to the next step. ⚠️</b>\n",
        "</div>\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dmWOrTJ3gx13"
      },
      "source": [
        "### Authenticate your notebook environment (Colab only)\n",
        "\n",
        "Authenticate your environment on Google Colab.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "NyKGtVQjgx13"
      },
      "outputs": [],
      "source": [
        "import sys\n",
        "\n",
        "if \"google.colab\" in sys.modules:\n",
        "\n",
        "    from google.colab import auth\n",
        "\n",
        "    auth.authenticate_user()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DF4l8DTdWgPY"
      },
      "source": [
        "### Set Google Cloud project information\n",
        "\n",
        "To run this tutorial, you must have an existing Google Cloud project. Learn more about [setting up a project and a development environment](https://cloud.google.com/vertex-ai/docs/start/cloud-environment)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Nqwi-5ufWp_B"
      },
      "outputs": [],
      "source": [
        "PROJECT_ID = \"[your-project-id]\"  # @param {type:\"string\"}\n",
        "LOCATION = \"us-central1\"  # @param {type:\"string\"}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "OKtKGmr9pfr6"
      },
      "source": [
        "### Create a Cloud Storage bucket\n",
        "\n",
        "Create a storage bucket to store artifacts such as datasets and trained model files."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "In3aQanwYjFB"
      },
      "outputs": [],
      "source": [
        "BUCKET_URI = f\"gs://your-bucket-name-{PROJECT_ID}-unique\"  # @param {type:\"string\"}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "create_bucket"
      },
      "source": [
        "**If your bucket doesn't already exist**: Run the following cell to create your Cloud Storage bucket."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "create_bucket"
      },
      "outputs": [],
      "source": [
        "! gsutil mb -l $LOCATION -p $PROJECT_ID $BUCKET_URI"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "init_aip:mbsdk"
      },
      "source": [
        "### Initialize Vertex AI SDK for Python\n",
        "\n",
        "To get started using Vertex AI, you must enable the [Vertex AI API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com).\n",
        "\n",
        "Initialize the Vertex AI SDK for Python for your project and corresponding bucket."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "init_aip:mbsdk"
      },
      "outputs": [],
      "source": [
        "from google.cloud import aiplatform\n",
        "\n",
        "aiplatform.init(project=PROJECT_ID, location=LOCATION, staging_bucket=BUCKET_URI)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "accelerators:training,prediction,cpu,mbsdk"
      },
      "source": [
        "### Set hardware accelerators\n",
        "\n",
        "You can set hardware accelerators for training and prediction.\n",
        "\n",
        "Set the variables `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 T4 GPUs allocated to each VM, you would specify:\n",
        "\n",
        "    (aiplatform.gapic.AcceleratorType.NVIDIA_TESLA_T4, 4)\n",
        "\n",
        "\n",
        "Otherwise specify `(None, None)` to use a container image to run on a CPU.\n",
        "\n",
        "Learn more [hardware accelerator support](https://cloud.google.com/vertex-ai/docs/general/locations#accelerators) for your location."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "accelerators:training,prediction,cpu,mbsdk"
      },
      "outputs": [],
      "source": [
        "TRAIN_GPU, TRAIN_NGPU = (None, None)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "container:prediction"
      },
      "source": [
        "### Set pre-built container image\n",
        "\n",
        "Set the pre-built Docker container image for prediction.\n",
        "\n",
        "Set the variable `TF` to the TensorFlow version of the container image. Replace dot with hyphen in the version number for specifying the image version. For example, `2-1` indicates version 2.1, and `1-15` indicates version 1.15. \n",
        "\n",
        "For the latest list of pre-built images, 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:prediction"
      },
      "outputs": [],
      "source": [
        "# Set TF version\n",
        "TF = \"2-7\"\n",
        "\n",
        "# Format the deployment image path\n",
        "DEPLOY_VERSION = \"tf2-cpu.{}\".format(TF)\n",
        "DEPLOY_IMAGE = \"{}-docker.pkg.dev/vertex-ai/prediction/{}:latest\".format(\n",
        "    LOCATION.split(\"-\")[0], DEPLOY_VERSION\n",
        ")\n",
        "\n",
        "print(\"Deployment:\", DEPLOY_IMAGE)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "machine:training,prediction"
      },
      "source": [
        "### Set machine type\n",
        "\n",
        "Next, set the machine type to use for training and prediction.\n",
        "\n",
        " - Set the variables `TRAIN_COMPUTE` and `DEPLOY_COMPUTE` to configure the compute of the VMs used for training and prediction respectively.\n",
        " - Set `machine type`\n",
        "     - `n1-standard`: 3.75GB of memory per vCPU.\n",
        "     - `n1-highmem`: 6.5GB of memory per vCPU\n",
        "     - `n1-highcpu`: 0.9 GB of memory per vCPU\n",
        " - `vCPUs`: number of \\[2, 4, 8, 16, 32, 64, 96 \\]\n",
        "\n",
        "**Note**: The following isn't supported for training:\n",
        "\n",
        " - `standard`: 2 vCPUs\n",
        " - `highcpu`: 2, 4 and 8 vCPUs\n",
        "\n",
        "**Note**: You may also use n2 and e2 machine types for training and deployment, but they don't support GPUs."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "machine:training,prediction"
      },
      "outputs": [],
      "source": [
        "TRAIN_COMPUTE = \"n1-standard-4\"\n",
        "print(\"Train machine type\", TRAIN_COMPUTE)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "train_custom_model"
      },
      "source": [
        "## Create the custom container image for training\n",
        "\n",
        "Now, you're ready for training your custom model on CIFAR10 data. There are two ways you can train a custom model using a container image:\n",
        "\n",
        "- **Use a Google Cloud prebuilt container**. If you use a prebuilt container, you provide a Python package that runs your training code inside the pre-built container.\n",
        "\n",
        "- **Use your own custom container image**. If you use your own container, you need to build the container image that fetches and your code for training a custom model."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "create_docker_container:training"
      },
      "source": [
        "### Create a training folder\n",
        "\n",
        "In this tutorial, you train a CIFAR10 model using your own custom container image. To run the training in your container, you create a Python training script that trains the model.\n",
        "\n",
        "First, create a directory for storing the training scripts and other components. Then, create a subdirectory(`trainer/`) for storing your training scripts separately. This subdirectory should include a `__init__.py` to make it a module. Learn more about defining [Python modules](https://cloud.google.com/vertex-ai/docs/training/create-python-pre-built-container#python-modules).\n",
        "\n",
        "Eventually, your directory structure looks like:\n",
        "\n",
        "```\n",
        "- custom/\n",
        "    - Dockerfile\n",
        "    - trainer/\n",
        "        - __init__.py\n",
        "        - task.py\n",
        "```"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "e5fc0fd1bfc2"
      },
      "outputs": [],
      "source": [
        "# Set the name of your app directory\n",
        "APPLICATION_DIR = \"custom\"\n",
        "# Remove if there's any such folder already\n",
        "! rm -rf $APPLICATION_DIR\n",
        "# Create your app directory\n",
        "! mkdir $APPLICATION_DIR\n",
        "# Create a subdirectory for store the training scripts\n",
        "! mkdir $APPLICATION_DIR/trainer\n",
        "# Create the init file\n",
        "! touch $APPLICATION_DIR/trainer/__init__.py"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "taskpy_contents:cifar10"
      },
      "source": [
        "### Create the training script\n",
        "\n",
        "In the next cell, define your training script `task.py` inside your training folder. \n",
        "\n",
        "In the training script, you perform the following steps sequentially:\n",
        "\n",
        "1. Load CIFAR10 dataset from TF Datasets (tfds).\n",
        "1. Build a model using TF.Keras model API.\n",
        "1. Compile the model (`compile()`).\n",
        "1. Set a training distribution strategy according to the argument `args.distribute`.\n",
        "1. Train the model (`fit()`) with epochs and steps according to the arguments `args.epochs` and `args.steps`\n",
        "1. Save the trained model to the specified base output directory which is accessed by the environment variable `AIP_MODEL_DIR`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "taskpy_contents:cifar10"
      },
      "outputs": [],
      "source": [
        "%%writefile $APPLICATION_DIR/trainer/task.py\n",
        "# Single, Mirror and Multi-Machine Distributed Training for CIFAR-10\n",
        "\n",
        "import tensorflow_datasets as tfds\n",
        "import tensorflow as tf\n",
        "from tensorflow.python.client import device_lib\n",
        "import argparse\n",
        "import os\n",
        "import sys\n",
        "tfds.disable_progress_bar()\n",
        "\n",
        "parser = argparse.ArgumentParser()\n",
        "\n",
        "parser.add_argument('--lr', dest='lr',\n",
        "                    default=0.01, type=float,\n",
        "                    help='Learning rate.')\n",
        "parser.add_argument('--epochs', dest='epochs',\n",
        "                    default=10, type=int,\n",
        "                    help='Number of epochs.')\n",
        "parser.add_argument('--steps', dest='steps',\n",
        "                    default=200, 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",
        "args = parser.parse_args()\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",
        "print('DEVICES', device_lib.list_local_devices())\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",
        "# Preparing dataset\n",
        "BUFFER_SIZE = 10000\n",
        "BATCH_SIZE = 64\n",
        "\n",
        "def make_datasets_unbatched():\n",
        "\n",
        "  # Scaling CIFAR10 data from (0, 255] to (0., 1.]\n",
        "  def scale(image, label):\n",
        "    image = tf.cast(image, tf.float32)\n",
        "    image /= 255.0\n",
        "    return image, label\n",
        "\n",
        "\n",
        "  datasets, info = tfds.load(name='cifar10',\n",
        "                            with_info=True,\n",
        "                            as_supervised=True)\n",
        "  return datasets['train'].map(scale).cache().shuffle(BUFFER_SIZE).repeat()\n",
        "\n",
        "\n",
        "# Build the Keras model\n",
        "def build_and_compile_cnn_model():\n",
        "  model = tf.keras.Sequential([\n",
        "      tf.keras.layers.Conv2D(32, 3, activation='relu', input_shape=(32, 32, 3)),\n",
        "      tf.keras.layers.MaxPooling2D(),\n",
        "      tf.keras.layers.Conv2D(32, 3, activation='relu'),\n",
        "      tf.keras.layers.MaxPooling2D(),\n",
        "      tf.keras.layers.Flatten(),\n",
        "      tf.keras.layers.Dense(10, activation='softmax')\n",
        "  ])\n",
        "  model.compile(\n",
        "      loss=tf.keras.losses.sparse_categorical_crossentropy,\n",
        "      optimizer=tf.keras.optimizers.SGD(learning_rate=args.lr),\n",
        "      metrics=['accuracy'])\n",
        "  return model\n",
        "\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 = make_datasets_unbatched().batch(GLOBAL_BATCH_SIZE)\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_cnn_model()\n",
        "\n",
        "MODEL_DIR = os.getenv(\"AIP_MODEL_DIR\")\n",
        "model.fit(x=train_dataset, epochs=args.epochs, steps_per_epoch=args.steps)\n",
        "model.save(MODEL_DIR)\n",
        "print (\"Save the model to\",MODEL_DIR)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "write_docker_file:training,tf-dlvm"
      },
      "source": [
        "### Write the contents of Dockerfile\n",
        "\n",
        "To containerize your code, you need to create a Dockerfile. In the Dockerfile you define all the steps needed to run your container. These steps include:\n",
        "\n",
        "1. Install a pre-defined container image from TensorFlow repository for deep learning images.\n",
        "2. Copy in the Python training code to the container.\n",
        "3. Set the entry into the Python training script as `trainer/task.py`. \n",
        "\n",
        "**Note**: The extension `.py` is dropped for `task.py` in the ENTRYPOINT command, as it is implied."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "write_docker_file:training,tf-dlvm"
      },
      "outputs": [],
      "source": [
        "%%writefile $APPLICATION_DIR/Dockerfile\n",
        "# Fetch the base image\n",
        "FROM gcr.io/deeplearning-platform-release/tf2-cpu.2-3\n",
        "\n",
        "# Set the working dir for the rest of the commands\n",
        "WORKDIR /\n",
        "\n",
        "# Copies the trainer code to the docker image.\n",
        "COPY trainer /trainer\n",
        "\n",
        "# Sets up the entry point to invoke the trainer.\n",
        "ENTRYPOINT [\"python\", \"-m\", \"trainer.task\"]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "a2ed53b6e259"
      },
      "source": [
        "### Enable Artifact Registry API\n",
        "\n",
        "To use your container image with Vertex AI, you need to upload your image to the Artifact Registry. Before you can push your image, you must enable the Artifact Registry API service for your project.\n",
        "\n",
        "Learn more about [enabling the Artifact Registry service](https://cloud.google.com/artifact-registry/docs/enable-service)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "13970c971b51"
      },
      "outputs": [],
      "source": [
        "! gcloud services enable artifactregistry.googleapis.com"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "616150778c61"
      },
      "source": [
        "### Create a repository in Artifact Registry"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "8d79730faf62"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "\n",
        "# Set the repository name\n",
        "REPO_NAME='custom-train-cifar10'\n",
        "\n",
        "if os.getenv(\"IS_TESTING\"):\n",
        "    ! sudo apt-get update --yes && sudo apt-get --only-upgrade --yes install google-cloud-sdk-cloud-run-proxy google-cloud-sdk-harbourbridge google-cloud-sdk-cbt google-cloud-sdk-gke-gcloud-auth-plugin google-cloud-sdk-kpt google-cloud-sdk-local-extract google-cloud-sdk-minikube google-cloud-sdk-app-engine-java google-cloud-sdk-app-engine-go google-cloud-sdk-app-engine-python google-cloud-sdk-spanner-emulator google-cloud-sdk-bigtable-emulator google-cloud-sdk-nomos google-cloud-sdk-package-go-module google-cloud-sdk-firestore-emulator kubectl google-cloud-sdk-datastore-emulator google-cloud-sdk-app-engine-python-extras google-cloud-sdk-cloud-build-local google-cloud-sdk-kubectl-oidc google-cloud-sdk-anthos-auth google-cloud-sdk-app-engine-grpc google-cloud-sdk-pubsub-emulator google-cloud-sdk-datalab google-cloud-sdk-skaffold google-cloud-sdk google-cloud-sdk-terraform-tools google-cloud-sdk-config-connector\n",
        "    ! gcloud components update --quiet\n",
        "\n",
        "# Create a repository in the Artifact Registry\n",
        "!gcloud artifacts repositories create $REPO_NAME --repository-format=docker \\\n",
        "--location=$LOCATION --description=\"Docker repository\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "94f77f953470"
      },
      "source": [
        "### Configure authentication to your private repo\n",
        "\n",
        "Before you push or pull container images, configure Docker to use the `gcloud` command-line tool to authenticate requests to Artifact Registry for your region."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "3a47ce54c71d"
      },
      "outputs": [],
      "source": [
        "! gcloud auth configure-docker {LOCATION}-docker.pkg.dev --quiet"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5fb8be3937a3"
      },
      "source": [
        "### Push your container image to Artifact Registry\n",
        "\n",
        "Specify the image path in the repository as `IMAGE_URI` and use the image path as a tag with the `docker build` command. Then, use `docker push` command to push your container image to Artifact Registry. Artifact registry creates the image in the repository based on the tag specified when you build the image.\n",
        "\n",
        "**Note**: As Docker is currently not fully supported in Colab, you use Cloud Build to push your container image to Artifact Registry. Learn more about [`gcloud builds submit`](https://cloud.google.com/sdk/gcloud/reference/builds/submit) command."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "4ab255f5f0ed"
      },
      "outputs": [],
      "source": [
        "import sys\n",
        "\n",
        "IS_COLAB = \"google.colab\" in sys.modules\n",
        "\n",
        "# Specify the image uri for Artifact Registry\n",
        "IMAGE_URI = f\"{LOCATION}-docker.pkg.dev/{PROJECT_ID}/{REPO_NAME}/cifar10:latest\"\n",
        "\n",
        "# If the env isn't Colab, run the docker commands\n",
        "if not IS_COLAB:\n",
        "    ! docker build -f $APPLICATION_DIR/Dockerfile -t $IMAGE_URI custom\n",
        "    ! docker push $IMAGE_URI\n",
        "# If using Colab, run the Cloud Build command\n",
        "else:\n",
        "    ! gcloud builds submit {APPLICATION_DIR}/ --region={LOCATION} --tag={IMAGE_URI}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "create_custom_training_job:mbsdk,no_model,custom"
      },
      "source": [
        "## Create a 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",
        "- `model_serving_container_image_uri`: The container image uri for deploying the model."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "create_custom_training_job:mbsdk,no_model,custom"
      },
      "outputs": [],
      "source": [
        "# Create the custom container training job\n",
        "job = aiplatform.CustomContainerTrainingJob(\n",
        "    display_name=\"cifar10-training\",\n",
        "    container_uri=IMAGE_URI,\n",
        "    model_serving_container_image_uri=DEPLOY_IMAGE,\n",
        ")\n",
        "\n",
        "print(job)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "prepare_custom_cmdargs"
      },
      "source": [
        "### Set the command-line arguments for training script\n",
        "\n",
        "In this step, you define the command-line arguments for running your training script. \n",
        "\n",
        "In this example, you pass the following two arguments:\n",
        "\n",
        "  - `epochs`: The number of epochs for training as defined in *EPOCHS*.\n",
        "  - `steps`: The number of steps per epoch as defined in *STEPS*.\n",
        "  \n",
        "Refer to the defined training script for more supported arguments."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "prepare_custom_cmdargs"
      },
      "outputs": [],
      "source": [
        "# Set the number of epochs\n",
        "EPOCHS = 1\n",
        "# Set the number of steps\n",
        "STEPS = 1\n",
        "\n",
        "CMDARGS = [\n",
        "    \"--epochs=\" + str(EPOCHS),\n",
        "    \"--steps=\" + str(STEPS),\n",
        "]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "run_custom_job:mbsdk,no_model"
      },
      "source": [
        "### Run the custom training job\n",
        "\n",
        "Next, you run the custom job using the `run()` method, with the following parameters:\n",
        "\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 type of hardware accelerator to be used.\n",
        "- `accelerator_count`: The number of accelerators to attach to a worker replica.\n",
        "- `base_output_dir`: The Cloud Storage location where the model artifacts need to be saved.\n",
        "- `model_display_name`: A human readable name for the registered model.\n",
        "- `sync`: Whether to block until completion of the job (Boolean)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "run_custom_job:mbsdk,no_model"
      },
      "outputs": [],
      "source": [
        "# Run the training with GPU\n",
        "if TRAIN_GPU:\n",
        "    model = job.run(\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=BUCKET_URI,\n",
        "        model_display_name=\"cifar10\",\n",
        "        sync=True,\n",
        "    )\n",
        "# Run the training with CPU\n",
        "else:\n",
        "    model = job.run(\n",
        "        args=CMDARGS,\n",
        "        replica_count=1,\n",
        "        machine_type=TRAIN_COMPUTE,\n",
        "        base_output_dir=BUCKET_URI,\n",
        "        model_display_name=\"cifar10\",\n",
        "        sync=True,\n",
        "    )"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "929696a237fe"
      },
      "source": [
        "### View the model in the Model Registry\n",
        "\n",
        "The `run()` method returns a Vertex AI model resource which indicates that your model is successfully registered in the Model Registry."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "33a1644f99f1"
      },
      "outputs": [],
      "source": [
        "print(model.gca_resource)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cleanup:mbsdk"
      },
      "source": [
        "# Cleaning up\n",
        "\n",
        "To clean up all Google Cloud resources used in this project, you can [delete the Google Cloud\n",
        "project](https://cloud.google.com/resource-manager/docs/creating-managing-projects#shutting_down_projects) you used for the tutorial.\n",
        "\n",
        "Otherwise, you can delete the individual resources you created in this tutorial."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "cleanup:mbsdk"
      },
      "outputs": [],
      "source": [
        "# Delete the artifact registry repo\n",
        "! gcloud artifacts repositories delete $REPO_NAME --location $LOCATION --quiet\n",
        "\n",
        "# Delete the custom training job\n",
        "job.delete()\n",
        "\n",
        "# Delete model\n",
        "model.delete()\n",
        "\n",
        "# Delete the Cloud Storage bucket\n",
        "delete_bucket = True\n",
        "if delete_bucket:\n",
        "    ! gsutil -m rm -r $BUCKET_URI\n",
        "\n",
        "# Delete application directory\n",
        "!rm -rf $APPLICATION_DIR"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "name": "custom_training_container_and_model_registry.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}