notebooks/official/ray_on_vertex_ai/get_started_with_pytorch_rov.ipynb

{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ur8xi4C7S06n"
      },
      "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",
      "metadata": {
        "id": "JAPoU8Sm5E6e"
      },
      "source": [
        "# Get started with PyTorch on Ray on Vertex AI\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/ray_on_vertex_ai/get_started_with_pytorch_rov.ipynb\">\n",
        "      <img width=\"32px\" 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%2Fray_on_vertex_ai%2Fget_started_with_pytorch_rov.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/blob/main/notebooks/official/ray_on_vertex_ai/get_started_with_pytorch_rov.ipynb\">\n",
        "      <img src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" alt=\"Vertex AI logo\"><br> Open in Vertex AI 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/ray_on_vertex_ai/get_started_with_pytorch_rov.ipynb\">\n",
        "      <img width=\"32px\" 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": "24743cf4a1e1"
      },
      "source": [
        "**_NOTE_**: This notebook has been tested in the following environment:\n",
        "\n",
        "* Python version = 3.9"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tvgnzT1CKxrO"
      },
      "source": [
        "## Overview\n",
        "\n",
        "This tutorial demonstrates how to use Ray on Vertex AI SDK and Vertex AI SDK for training and serving an PyTorch image classification model.\n",
        "\n",
        "Learn more about [Ray on Vertex AI overview](https://cloud.google.com/vertex-ai/docs/open-source/ray-on-vertex-ai/overview)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "d975e698c9a4"
      },
      "source": [
        "### Objective\n",
        "\n",
        "In this tutorial, you learn how to efficiently distribute the training process of a PyTorch image classification model by leveraging Ray on Vertex AI. Furthermore, you learn how to deploy the trained model seamlessly to Vertex AI Endpoint.\n",
        "\n",
        "This tutorial uses the following Google Cloud ML services and resources:\n",
        "\n",
        "- Ray on Vertex AI\n",
        "- Vertex AI Model Registry\n",
        "- Vertex AI Prediction\n",
        "\n",
        "\n",
        "The steps performed include:\n",
        "\n",
        "- Prepare the training script\n",
        "- Submit a Ray job using the Ray Jobs API\n",
        "- Download a trained image model from PyTorch\n",
        "- Create a custom model handler\n",
        "- Package model artifacts in a model archive file\n",
        "- Register model in Vertex AI Model Registry\n",
        "- Deploy model in Vertex AI Endpoint\n",
        "- Make online predictions"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "08d289fa873f"
      },
      "source": [
        "### Dataset\n",
        "\n",
        "This tutorial uses the [CIFAR-10 dataset](https://pytorch.org/vision/stable/generated/torchvision.datasets.CIFAR10.html) which consists of 60000 32x32 colour images in 10 classes, with 6000 images per class."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "aed92deeb4a0"
      },
      "source": [
        "### Costs\n",
        "\n",
        "This tutorial uses billable components of Google Cloud:\n",
        "\n",
        "* Vertex AI\n",
        "* Cloud Storage\n",
        "\n",
        "Learn about [Vertex AI pricing](https://cloud.google.com/vertex-ai/pricing),\n",
        "and [Cloud Storage pricing](https://cloud.google.com/storage/pricing),\n",
        "and use the [Pricing Calculator](https://cloud.google.com/products/calculator/)\n",
        "to generate a cost estimate based on your projected usage."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "i7EUnXsZhAGF"
      },
      "source": [
        "## Installation\n",
        "\n",
        "Install the following packages required to execute this notebook."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "2b4ef9b72d43"
      },
      "outputs": [],
      "source": [
        "# Install the packages\n",
        "import os\n",
        "\n",
        "if not os.getenv(\"IS_TESTING\"):\n",
        "    USER = \"--user\"\n",
        "else:\n",
        "    USER = \"\"\n",
        "\n",
        "! pip3 install {USER} google-cloud-aiplatform[ray]==1.56.0 ray[data]==2.9.3 ray[train]==2.9.3 ray[tune]==2.9.3 -q --no-warn-conflicts\n",
        "! pip3 install {USER} torch==2.1.2 torchvision==0.16.2 torchmetrics==1.2.1 torchserve==0.9.0 torch-model-archiver==0.9.0 -q --no-warn-conflicts\n",
        "! pip3 install {USER} google-auth==2.27.0 etils==1.5.2 -q --no-warn-conflicts"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "58707a750154"
      },
      "source": [
        "### Colab only: Uncomment the following cell to restart the kernel."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "f200f10a1da3"
      },
      "outputs": [],
      "source": [
        "# import IPython\n",
        "\n",
        "# app = IPython.Application.instance()\n",
        "# app.kernel.do_shutdown(True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BF1j6f9HApxa"
      },
      "source": [
        "## Before you begin\n",
        "\n",
        "### Set up your Google Cloud project\n",
        "\n",
        "**The following steps are required, regardless of your notebook environment.**\n",
        "\n",
        "1. [Select or create a Google Cloud project](https://console.cloud.google.com/cloud-resource-manager). When you first create an account, you get a $300 free credit towards your compute/storage costs.\n",
        "\n",
        "2. [Make sure that billing is enabled for your project](https://cloud.google.com/billing/docs/how-to/modify-project).\n",
        "\n",
        "3. [Enable the Vertex AI API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com).\n",
        "\n",
        "4. If you are running this notebook locally, you need to install the [Cloud SDK](https://cloud.google.com/sdk)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WReHDGG5g0XY"
      },
      "source": [
        "#### Set your project ID\n",
        "\n",
        "**If you don't know your project ID**, try the following:\n",
        "* Run `gcloud config list`.\n",
        "* Run `gcloud projects list`.\n",
        "* See the support page: [Locate the project ID](https://support.google.com/googleapi/answer/7014113)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "oM1iC_MfAts1"
      },
      "outputs": [],
      "source": [
        "PROJECT_ID = \"[your-project-id]\"  # @param {type:\"string\"}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "region"
      },
      "source": [
        "#### Region\n",
        "\n",
        "You can also change the `REGION` variable used by Vertex AI. Learn more about [Vertex AI regions](https://cloud.google.com/vertex-ai/docs/general/locations)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "fPrDj6HE9_EU"
      },
      "outputs": [],
      "source": [
        "REGION = \"us-central1\"  # @param {type: \"string\"}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "timestamp"
      },
      "source": [
        "#### Timestamp\n",
        "\n",
        "If you are in a live tutorial session, you might be using a shared test account or project. To avoid name collisions between users on resources created, you create a timestamp for each instance session, and append the timestamp onto the name of resources you create in this tutorial."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "W6Le1schAziq"
      },
      "outputs": [],
      "source": [
        "from datetime import datetime\n",
        "\n",
        "TIMESTAMP = datetime.now().strftime(\"%Y%m%d%H%M%S\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "sBCra4QMA2wR"
      },
      "source": [
        "### Authenticate your Google Cloud account\n",
        "\n",
        "Depending on your Jupyter environment, you may have to manually authenticate. Follow the relevant instructions below."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "74ccc9e52986"
      },
      "source": [
        "**1. Vertex AI Workbench**\n",
        "* Do nothing as you are already authenticated."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "de775a3773ba"
      },
      "source": [
        "**2. Local JupyterLab instance, uncomment and run:**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "254614fa0c46"
      },
      "outputs": [],
      "source": [
        "# ! gcloud auth login"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ef21552ccea8"
      },
      "source": [
        "**3. Colab, uncomment and run:**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "603adbbf0532"
      },
      "outputs": [],
      "source": [
        "# from google.colab import auth\n",
        "# auth.authenticate_user()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "f6b2ccc891ed"
      },
      "source": [
        "**4. Service account or other**\n",
        "* See how to grant Cloud Storage permissions to your service account at https://cloud.google.com/storage/docs/gsutil/commands/iam#ch-examples."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zgPO1eR3CYjk"
      },
      "source": [
        "### Create a Cloud Storage bucket\n",
        "\n",
        "Create a storage bucket to store intermediate artifacts such as datasets."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "MzGDU7TWdts_"
      },
      "outputs": [],
      "source": [
        "BUCKET_URI = f\"gs://your-bucket-name-{PROJECT_ID}-unique\"  # @param {type:\"string\"}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-EcIXiGsCePi"
      },
      "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": "NIq7R4HZCfIc"
      },
      "outputs": [],
      "source": [
        "! gsutil mb -l {REGION} -p {PROJECT_ID} {BUCKET_URI}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ObZfmnat7BaA"
      },
      "source": [
        "### Add torch-model-archiver to the PATH\n",
        "\n",
        "Update the `PATH` environment variable to add `torch-model-archiver`"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "yRRsfkBu7C8M"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "\n",
        "os.environ[\"PATH\"] = f'{os.environ.get(\"PATH\")}:~/.local/bin'"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "k9ryiScCEapt"
      },
      "source": [
        "### Set a Ray cluster on Vertex AI"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MwmK5nBBCgJa"
      },
      "source": [
        "Before running the code below, make sure to [set up](https://cloud.google.com/vertex-ai/docs/open-source/ray-on-vertex-ai/set-up) Ray on Vertex AI and [create](https://cloud.google.com/vertex-ai/docs/open-source/ray-on-vertex-ai/create-cluster) at least one Ray cluster on Vertex AI."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_mlq_1NLEonF"
      },
      "outputs": [],
      "source": [
        "import vertex_ray\n",
        "from google.cloud import aiplatform as vertex_ai\n",
        "from vertex_ray import Resources"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dorIZFvjnGKL"
      },
      "source": [
        "#### Initialize Vertex AI SDK for Python\n",
        "\n",
        "Initialize the Vertex AI SDK for Python for your project."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "VOOgvRJoQ6Xj"
      },
      "outputs": [],
      "source": [
        "vertex_ai.init(project=PROJECT_ID, location=REGION, staging_bucket=BUCKET_URI)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "15PYjIyOcx1d"
      },
      "source": [
        "#### Define cluster configuration\n",
        "\n",
        "To provision a Ray cluster on Vertex AI, you can use a default provisioning request or you can specify the replica count (number of nodes), machine type, disk_spec, and accelerator as needed."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "LjLG3rsRbiHr"
      },
      "outputs": [],
      "source": [
        "head_node_type = Resources(\n",
        "    machine_type=\"n1-standard-16\",\n",
        "    node_count=1,\n",
        ")\n",
        "\n",
        "worker_node_types = [\n",
        "    Resources(\n",
        "        machine_type=\"n1-standard-16\",\n",
        "        node_count=2,\n",
        "    )\n",
        "]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_O1xUMt7Z6r0"
      },
      "source": [
        "#### Create the Ray cluster\n",
        "\n",
        "Create the Ray cluster using the Vertex AI SDK for Python version used with Ray."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "eZkHOH3v2i1p"
      },
      "outputs": [],
      "source": [
        "cluster_name = f\"ray-cluster-{TIMESTAMP}\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "J-g6kLwqUj5n"
      },
      "outputs": [],
      "source": [
        "ray_cluster_name = vertex_ray.create_ray_cluster(\n",
        "    head_node_type=head_node_type,\n",
        "    worker_node_types=worker_node_types,\n",
        "    cluster_name=cluster_name,\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xmBlsbHAc2uO"
      },
      "source": [
        "#### Get the Ray cluster\n",
        "\n",
        "Use the Vertex AI SDK for Python to get the Ray cluster."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9UzG2WyXbZJi"
      },
      "outputs": [],
      "source": [
        "ray_cluster = vertex_ray.get_ray_cluster(ray_cluster_name)\n",
        "print(\"Ray cluster on Vertex AI:\", ray_cluster_name)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Ek1-iTbPjzdJ"
      },
      "source": [
        "### Set tutorial folder\n",
        "\n",
        "Set up the folder you use in this tutorial."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "BbfKRabXj3la"
      },
      "outputs": [],
      "source": [
        "from pathlib import Path as path\n",
        "\n",
        "root_path = path.cwd()\n",
        "tutorial_path = root_path / \"tutorial\"\n",
        "src_path = tutorial_path / \"src\"\n",
        "deliverables_path = tutorial_path / \"deliverables\"\n",
        "build_path = tutorial_path / \"build\"\n",
        "tests_path = tutorial_path / \"tests\"\n",
        "\n",
        "src_path.mkdir(parents=True, exist_ok=True)\n",
        "deliverables_path.mkdir(parents=True, exist_ok=True)\n",
        "build_path.mkdir(parents=True, exist_ok=True)\n",
        "tests_path.mkdir(parents=True, exist_ok=True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "960505627ddf"
      },
      "source": [
        "### Import libraries"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "PyQmSRbKA8r-"
      },
      "outputs": [],
      "source": [
        "import base64\n",
        "import io\n",
        "import os\n",
        "import random\n",
        "import shutil\n",
        "import string\n",
        "import time\n",
        "\n",
        "# Ray - Training\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.nn.functional as F\n",
        "import torchvision.datasets as datasets\n",
        "import torchvision.transforms as transforms\n",
        "# General\n",
        "from etils import epath\n",
        "from matplotlib import pyplot as plt\n",
        "from ray.job_submission import JobStatus, JobSubmissionClient\n",
        "# Serving\n",
        "from ray.tune import ExperimentAnalysis\n",
        "from vertex_ray.predict import torch as ray_torch"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xFgb-sZbBi8i"
      },
      "source": [
        "### Set variables"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "zykxFjqUB9jt"
      },
      "outputs": [],
      "source": [
        "# training\n",
        "LOGGING_URI = epath.Path(BUCKET_URI) / \"logs\"\n",
        "EXPERIMENT_NAME = \"torch_on_rov\"\n",
        "TRAINING_URI = LOGGING_URI / EXPERIMENT_NAME\n",
        "TRAINING_PATH = deliverables_path / EXPERIMENT_NAME\n",
        "\n",
        "# serving\n",
        "DELIVERABLES_PATH = str(deliverables_path)\n",
        "BUILD_URI = str(epath.Path(BUCKET_URI) / \"build\")\n",
        "DEPLOY_IMAGE_URI = \"us-docker.pkg.dev/vertex-ai/prediction/pytorch-cpu.1-11:latest\"\n",
        "MODEL_NAME = \"torch_on_rov_cifar10\"\n",
        "DEPLOY_COMPUTE = \"n1-standard-4\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "mYQNboaOBk6B"
      },
      "source": [
        "### Define helpers"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "2-bl6h54bqFz"
      },
      "outputs": [],
      "source": [
        "def get_id(k=3):\n",
        "    return \"\".join(random.choices(string.ascii_lowercase + string.digits, k=k))\n",
        "\n",
        "\n",
        "def plot_image_sample(test_dataset):\n",
        "    \"\"\"Plots a sample image from the CIFAR-10 dataset.\"\"\"\n",
        "\n",
        "    sample_idx = random.randrange(0, len(test_dataset))\n",
        "    image, _ = test_dataset[sample_idx]\n",
        "    pil_image = transforms.ToPILImage()(image)\n",
        "    plt.imshow(pil_image)\n",
        "    plt.show()\n",
        "\n",
        "    return pil_image\n",
        "\n",
        "\n",
        "def predict_from_image(pil_image, endpoint):\n",
        "    \"\"\"Predicts the class of an image using the given endpoint.\"\"\"\n",
        "    buffered_image = io.BytesIO()\n",
        "    pil_image.save(buffered_image, format=\"JPEG\")\n",
        "\n",
        "    data = {\"data\": base64.b64encode(buffered_image.getvalue()).decode(\"utf-8\")}\n",
        "    response = endpoint.predict(instances=[data])\n",
        "\n",
        "    return response.predictions"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8BH0TlP3PtTB"
      },
      "source": [
        "## Training a PyTorch model\n",
        "\n",
        "In this tutorial, you train a custom image classification model using Ray on Vertex AI."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Qh2BbfWPXglw"
      },
      "source": [
        "#### Prepare the training application\n",
        "\n",
        "Before you start the training, let's take a look at how a Ray job might be assembled to distribute your training.\n",
        "\n",
        "Ray 2.4.0 uses `train_loop_per_worker` fuction to a distributed multi-worker training function.\n",
        "\n",
        "After you set up your dataset and your model, define your single-worker PyTorch training function and then convert it to distributed multi-worker training function as followed:\n",
        "\n",
        "1. Use the `ray.train.torch.prepare_data_loader` to wrap your data with  `DistributedSampler` for distributed training.\n",
        "\n",
        "2. Use the `ray.train.torch.prepare_model` function to wrap your model with  `DistributedDataParallel` for distributed training.\n",
        "\n",
        "After you have your multi-worker training function, you need to define the `ScalingConfig` to specify the desired number of workers and indicate whether the distributed training process requires GPUs.\n",
        "\n",
        "Additionally, you can define a `RunConfig` to specify checkpointing and synchronization behaviors along the distributed training workload and some additional training loop parameters.\n",
        "\n",
        "Finally, pass everything to `TorchTrainer`, which Ray uses to distribute your training utilizing Distributed Data Parallelism (using PyTorch’s Distributed backend)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pznqJKCpEcL3"
      },
      "source": [
        "### Prepare the training script\n",
        "\n",
        "The file `src/task.py` is the Python script for executing the Ray distributed training job."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "rCFfhM7RP4RI"
      },
      "outputs": [],
      "source": [
        "training_script = \"\"\"\n",
        "# libraries\n",
        "\n",
        "# general libraries\n",
        "import os\n",
        "import argparse\n",
        "from etils import epath\n",
        "import tempfile\n",
        "\n",
        "# training libraries\n",
        "import torch\n",
        "from torch import Tensor\n",
        "import torch.nn as nn\n",
        "import torch.nn.functional as F\n",
        "import torch.utils.data as data\n",
        "import torchvision.transforms as transforms\n",
        "import torchvision.datasets as datasets\n",
        "from torchmetrics.aggregation import MeanMetric\n",
        "from torchmetrics.classification.accuracy import Accuracy\n",
        "\n",
        "# ray libraries\n",
        "import ray\n",
        "from ray import train\n",
        "from ray.train import ScalingConfig, RunConfig, CheckpointConfig, Checkpoint\n",
        "from ray.train.torch import TorchTrainer, TorchCheckpoint\n",
        "\n",
        "\n",
        "# helpers\n",
        "def get_args():\n",
        "    parser = argparse.ArgumentParser(description='PyTorch ImageNet Training')\n",
        "    parser.add_argument('--batch-size', dest='batch_size',\n",
        "                        type=int, default=16, help='Batch size')\n",
        "    parser.add_argument('--epochs', dest='epochs',\n",
        "                        type=int, default=10, help='Number of epochs')\n",
        "    parser.add_argument('--lr', dest='lr',\n",
        "                        type=int, default=1e-3, help='Learning rate')\n",
        "    parser.add_argument('--num-workers', dest='num_workers',\n",
        "                        type=int, default=1, help='Number of workers')\n",
        "    parser.add_argument('--use-gpu', dest='use_gpu', action='store_true',\n",
        "                        default=False, help='Use GPU')\n",
        "    parser.add_argument('--experiment-name', dest='experiment_name', type=str,\n",
        "                        default='cifar10-torch', help='Experiment name')\n",
        "    parser.add_argument('--logging-dir', dest='logging_dir',\n",
        "                        type=str, default='./logs', help='Logging directory')\n",
        "    args = parser.parse_args()\n",
        "    return args\n",
        "    \n",
        "# Create a simple model\n",
        "class Cifar10Model(nn.Module):\n",
        "    def __init__(self):\n",
        "        super(Cifar10Model, self).__init__()\n",
        "        self.conv1 = nn.Conv2d(3, 16, 3)\n",
        "        self.pool = nn.MaxPool2d(2, 2)\n",
        "        self.conv2 = nn.Conv2d(16, 16, 5)\n",
        "        self.fc1 = nn.Linear(16 * 5 * 5, 120)\n",
        "        self.fc2 = nn.Linear(120, 84)\n",
        "        self.fc3 = nn.Linear(84, 10)\n",
        "\n",
        "    def forward(self, x):\n",
        "        x = self.pool(F.relu(self.conv1(x)))\n",
        "        x = self.pool(F.relu(self.conv2(x)))\n",
        "        x = torch.flatten(x, 1)  # flatten all dimensions except batch\n",
        "        x = F.relu(self.fc1(x))\n",
        "        x = F.relu(self.fc2(x))\n",
        "        x = self.fc3(x)\n",
        "        return x\n",
        "\n",
        "def train_epoch(train_loader, model, loss_fn, optimizer, device):\n",
        "    # initiate training\n",
        "    model.train()\n",
        "    train_loss = MeanMetric()\n",
        "    train_accuracy = Accuracy(task=\"multiclass\", num_classes=10)\n",
        "\n",
        "    # load metrics to device\n",
        "    train_loss.to(device)\n",
        "    train_accuracy.to(device)\n",
        "\n",
        "    # training loop\n",
        "    for batch, (data, target) in enumerate(train_loader):\n",
        "        # move data to device\n",
        "        data = data.to(device)\n",
        "        target = target.to(device)\n",
        "\n",
        "        # compute error\n",
        "        output = model(data)\n",
        "        loss = loss_fn(output, target)\n",
        "        train_loss.update(loss)\n",
        "        train_accuracy.update(output, target)\n",
        "\n",
        "        # backpropagation\n",
        "        optimizer.zero_grad()\n",
        "        loss.backward()\n",
        "        optimizer.step()\n",
        "\n",
        "        # print results\n",
        "        if batch % 100 == 0:\n",
        "            print(f'batch {batch}, loss {loss.item():.4f}')\n",
        "\n",
        "    # compute loss and metrics\n",
        "    train_loss = Tensor.numpy(train_loss.compute(), force=True).item()\n",
        "    train_accuracy = Tensor.numpy(train_accuracy.compute(), force=True).item()\n",
        "\n",
        "    return train_loss, train_accuracy\n",
        "\n",
        "\n",
        "def evaluate_epoch(test_loader, model, loss_fn, device):\n",
        "    # initiate evaluation\n",
        "    model.eval()\n",
        "    test_loss = MeanMetric()\n",
        "    test_accuracy = Accuracy(task=\"multiclass\", num_classes=10)\n",
        "\n",
        "    # load metrics to device\n",
        "    test_loss.to(device)\n",
        "    test_accuracy.to(device)\n",
        "\n",
        "    # evaluation loop\n",
        "    for batch, (data, target) in enumerate(test_loader):\n",
        "        with torch.no_grad():\n",
        "            # move data to device\n",
        "            data = data.to(device)\n",
        "            target = target.to(device)\n",
        "\n",
        "            # get loss and accuracy\n",
        "            output = model(data)\n",
        "            test_loss.update(loss_fn(output, target))\n",
        "            test_accuracy.update(output, target)\n",
        "\n",
        "    # compute loss and metrics\n",
        "    test_loss = Tensor.numpy(test_loss.compute(), force=True).item()\n",
        "    test_accuracy = Tensor.numpy(test_accuracy.compute(), force=True).item()\n",
        "\n",
        "    return test_loss, test_accuracy\n",
        "\n",
        "\n",
        "def train_loop_per_worker(config):\n",
        "    # set configuration\n",
        "    batch_size = config[\"batch_size\"]\n",
        "    epochs = config[\"epochs\"]\n",
        "    learning_rate = config[\"learning_rate\"]\n",
        "\n",
        "    # get device\n",
        "    device = train.torch.get_device() if torch.cuda.is_available() else torch.device(\"cpu\")\n",
        "\n",
        "    # read dataset\n",
        "    normalize = transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])\n",
        "    transform = transforms.Compose([transforms.ToTensor(), normalize])\n",
        "\n",
        "    train_dataset = datasets.CIFAR10(root=\"./train\",\n",
        "                                     transform=transform,\n",
        "                                     train=True, download=True)\n",
        "\n",
        "    test_dataset = datasets.CIFAR10(root=\"./test\",\n",
        "                                    transform=transform,\n",
        "                                    train=False, download=True)\n",
        "\n",
        "    # create data loaders\n",
        "    train_loader = data.DataLoader(train_dataset, batch_size=batch_size)\n",
        "    test_loader = data.DataLoader(test_dataset, batch_size=batch_size)\n",
        "    train_loader = train.torch.prepare_data_loader(train_loader)\n",
        "    test_loader = train.torch.prepare_data_loader(test_loader)\n",
        "\n",
        "    # create model\n",
        "    model = Cifar10Model()\n",
        "    model = train.torch.prepare_model(model)\n",
        "    loss_fn = nn.CrossEntropyLoss()\n",
        "    optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)\n",
        "\n",
        "    # train model\n",
        "    for epoch in range(1, epochs + 1):\n",
        "        train_loss, train_accuracy = train_epoch(train_loader, model, loss_fn, optimizer, device)\n",
        "        test_loss, test_accuracy = evaluate_epoch(test_loader, model, loss_fn, device)\n",
        "\n",
        "        # report metrics and model checkpoint\n",
        "        train.report(\n",
        "            metrics={\"train_loss\": train_loss, \"train_accuracy\": train_accuracy,\n",
        "                     \"test_loss\": test_loss, \"test_accuracy\": test_accuracy},\n",
        "            checkpoint=TorchCheckpoint.from_state_dict(model.state_dict())\n",
        "        )\n",
        "            \n",
        "            \n",
        "def main():\n",
        "    # set configuration\n",
        "    args = get_args()\n",
        "    config = vars(args)\n",
        "\n",
        "    # initialize ray session\n",
        "    ray.init()\n",
        "\n",
        "    # train model\n",
        "    train_loop_config = {\"learning_rate\": config['lr'], \"batch_size\": config['batch_size'],\n",
        "                         \"epochs\": config['epochs']}\n",
        "    scaling_config = ScalingConfig(num_workers=config['num_workers'], use_gpu=config['use_gpu'])\n",
        "    run_config = RunConfig(checkpoint_config=CheckpointConfig(num_to_keep=1),\n",
        "                           storage_path=config['logging_dir'],\n",
        "                           name=config['experiment_name'])\n",
        "\n",
        "    trainer = TorchTrainer(\n",
        "        train_loop_per_worker=train_loop_per_worker,\n",
        "        train_loop_config=train_loop_config,\n",
        "        run_config=run_config,\n",
        "        scaling_config=scaling_config\n",
        "    )\n",
        "    result = trainer.fit()\n",
        "    print(f\"Last result: {result.metrics}\")\n",
        "\n",
        "\n",
        "if __name__ == \"__main__\":\n",
        "    main()\n",
        "\"\"\"\n",
        "\n",
        "with open(src_path / \"task.py\", \"w\") as f:\n",
        "    f.write(training_script)\n",
        "f.close()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-AOOmNti23qc"
      },
      "source": [
        "### Prepare the `requirements` file\n",
        "\n",
        "The file `requirements.txt` includes the dependencies your Ray application needs to run."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "83KUQQylbrJR"
      },
      "outputs": [],
      "source": [
        "requirements = \"\"\"\n",
        "importlib_resources==6.1.1\n",
        "etils==1.5.2\n",
        "ray[data]==2.9.3\n",
        "ray[train]==2.9.3\n",
        "ray[tune]==2.9.3\n",
        "torch==2.1.2\n",
        "torchvision==0.16.2\n",
        "torchmetrics==1.2.1\n",
        "\"\"\"\n",
        "\n",
        "with open(tutorial_path / \"requirements.txt\", \"w\") as f:\n",
        "    f.write(requirements)\n",
        "f.close()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "PYcmfEvZ3C1i"
      },
      "source": [
        "### Submit a Ray job using the Ray Jobs API\n",
        "\n",
        "Submit the script to the Ray cluster on Vertex AI using the Ray Jobs API with  the public Ray dashboard address.\n",
        "\n",
        "It's important to highlight that Ray Jobs API is the prefered option if you'd rather submit jobs programmatically. You can also use the Ray on Vertex AI SDK if you prefer an interactive Python development environment."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "PPRtGm9IJ09g"
      },
      "source": [
        "Initiate the client to submit the job."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "FljDHRQ63EP4"
      },
      "outputs": [],
      "source": [
        "client = JobSubmissionClient(address=f\"vertex_ray://{ray_cluster.dashboard_address}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "uDHIlGlQJ2oi"
      },
      "source": [
        "Submit the job."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "myMOWdlV3rp_"
      },
      "outputs": [],
      "source": [
        "id = get_id()\n",
        "\n",
        "job_id = client.submit_job(\n",
        "    submission_id=f\"ray-job-{TIMESTAMP}-{id}\",\n",
        "    entrypoint=f\"python3 task.py --experiment-name={EXPERIMENT_NAME} --num-workers=2 --logging-dir={LOGGING_URI}\",\n",
        "    runtime_env={\n",
        "        \"pip\": {\"packages\": str(tutorial_path / \"requirements.txt\")},\n",
        "        \"working_dir\": str(src_path),\n",
        "    },\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "viL5MsRTJ5Df"
      },
      "source": [
        "Check the status of the job."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "TnQfy6VK5pvr"
      },
      "outputs": [],
      "source": [
        "while True:\n",
        "    job_status = client.get_job_status(job_id)\n",
        "    if job_status == JobStatus.SUCCEEDED:\n",
        "        print(\"Job succeeded!\")\n",
        "        break\n",
        "    else:\n",
        "        if job_status == JobStatus.FAILED:\n",
        "            print(\"Job failed!\")\n",
        "            break\n",
        "        else:\n",
        "            print(\"Job is running...\")\n",
        "            time.sleep(60)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "67FI31SqKFdF"
      },
      "source": [
        "### Check model artifacts\n",
        "\n",
        "When the Ray training job has completed, check the model artifacts in the Cloud Storage location.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Q_JOYhXQKK8U"
      },
      "outputs": [],
      "source": [
        "! gsutil ls -l {TRAINING_URI}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "XKUOXlN-MBaU"
      },
      "source": [
        "## Serving a PyTorch model\n",
        "\n",
        "You can serve a PyTorch model on Vertex AI using TorchServe as follows:\n",
        "\n",
        "1. Download Ray training checkpoints.\n",
        "2. Get PyTorch model from the Ray TorchCheckpoint.\n",
        "3. Package the trained model artifacts including the model artifact, the model module and the custom handler by creating an archive file using the Torch Model Archiver tool.\n",
        "4. Register model in Vertex AI Model Registry.\n",
        "5. Deploy model to Vertex AI endpoint for predictions."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "fN5YQ60PXbYq"
      },
      "source": [
        "### Download Ray training checkpoints\n",
        "\n",
        "Download all resulting checkpoints from Ray training job."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "uBf0_agLhrLV"
      },
      "outputs": [],
      "source": [
        "! gsutil -q cp -r {TRAINING_URI} {DELIVERABLES_PATH}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "VavpxaY1hlSN"
      },
      "source": [
        "### Get the best training checkpoint\n",
        "\n",
        "Use the `ExperimentAnalysis` to retrive the best checkpoint according to relevant metrics and mode."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "nUBEKpK1gVx9"
      },
      "outputs": [],
      "source": [
        "experiment_analysis = ExperimentAnalysis(TRAINING_PATH)\n",
        "log_path = experiment_analysis.get_best_trial(metric=\"test_accuracy\", mode=\"max\")\n",
        "best_checkpoint = experiment_analysis.get_best_checkpoint(\n",
        "    log_path, metric=\"test_accuracy\", mode=\"max\"\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Bf2YPC5nlzjL"
      },
      "source": [
        "### Get PyTorch Model from Ray TorchCheckpoint\n",
        "\n",
        "Convert a TorchCheckpoint to Pytorch Model."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "78577f79de54"
      },
      "outputs": [],
      "source": [
        "class Cifar10Model(nn.Module):\n",
        "    def __init__(self):\n",
        "        super(Cifar10Model, self).__init__()\n",
        "        self.conv1 = nn.Conv2d(3, 16, 3)\n",
        "        self.pool = nn.MaxPool2d(2, 2)\n",
        "        self.conv2 = nn.Conv2d(16, 16, 5)\n",
        "        self.fc1 = nn.Linear(16 * 5 * 5, 120)\n",
        "        self.fc2 = nn.Linear(120, 84)\n",
        "        self.fc3 = nn.Linear(84, 10)\n",
        "\n",
        "    def forward(self, x):\n",
        "        x = self.pool(F.relu(self.conv1(x)))\n",
        "        x = self.pool(F.relu(self.conv2(x)))\n",
        "        x = torch.flatten(x, 1)  # flatten all dimensions except batch\n",
        "        x = F.relu(self.fc1(x))\n",
        "        x = F.relu(self.fc2(x))\n",
        "        x = self.fc3(x)\n",
        "        return x\n",
        "\n",
        "\n",
        "model_definition = Cifar10Model()\n",
        "torch_model = ray_torch.get_pytorch_model_from(best_checkpoint, model=model_definition)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "CY5AV09xm0dl"
      },
      "source": [
        "### Build PyTorch model archive (.mar) file\n",
        "\n",
        "TorchServe lets you to serve the Torch model by packaging all model artifacts into a single model archive file. In this case, the following information is required to create a standalone model archive:\n",
        "\n",
        "1. Serialized file\n",
        "2. Model file\n",
        "3. Handler"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kvfiYnrZ63iN"
      },
      "source": [
        "#### Save the model\n",
        "\n",
        "The `model.pt` contains the model state_dict.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "fLR_MZS_tXuq"
      },
      "outputs": [],
      "source": [
        "torch.save(torch_model.state_dict(), build_path / \"model.pt\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vl-xT_zDmOvg"
      },
      "source": [
        "#### Create the `model` module\n",
        "\n",
        "The `model.py` file should contain the model architecture.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "nT9Szm-LA8Uh"
      },
      "outputs": [],
      "source": [
        "model_script = \"\"\"\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.nn.functional as F\n",
        "\n",
        "\n",
        "class Cifar10Model(nn.Module):\n",
        "    def __init__(self):\n",
        "        super(Cifar10Model, self).__init__()\n",
        "        self.conv1 = nn.Conv2d(3, 16, 3)\n",
        "        self.pool = nn.MaxPool2d(2, 2)\n",
        "        self.conv2 = nn.Conv2d(16, 16, 5)\n",
        "        self.fc1 = nn.Linear(16 * 5 * 5, 120)\n",
        "        self.fc2 = nn.Linear(120, 84)\n",
        "        self.fc3 = nn.Linear(84, 10)\n",
        "\n",
        "    def forward(self, x):\n",
        "        x = self.pool(F.relu(self.conv1(x)))\n",
        "        x = self.pool(F.relu(self.conv2(x)))\n",
        "        x = torch.flatten(x, 1)  # flatten all dimensions except batch\n",
        "        x = F.relu(self.fc1(x))\n",
        "        x = F.relu(self.fc2(x))\n",
        "        x = self.fc3(x)\n",
        "        return x\n",
        "\"\"\"\n",
        "\n",
        "with open(build_path / \"model.py\", \"w\") as model_file:\n",
        "    model_file.write(model_script)\n",
        "model_file.close()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "h5Jc76AmBP-C"
      },
      "source": [
        "#### Create the `custom_handler` module\n",
        "\n",
        "The `custom_handler.py` file uses the TorchServe's inbuilt `image_classifier` handler name to handle custom TorchServe inference logic."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "DDpXiVHnBPJ-"
      },
      "outputs": [],
      "source": [
        "custom_handler_script = '''\n",
        "# Based on https://github.com/pytorch/serve/blob/master/examples/image_classifier/mnist/mnist_handler.py\n",
        "\n",
        "from torchvision import transforms\n",
        "from ts.torch_handler.image_classifier import ImageClassifier\n",
        "from torch.profiler import ProfilerActivity\n",
        "\n",
        "\n",
        "class Cifar10Classifier(ImageClassifier):\n",
        "    \"\"\"\n",
        "    Cifar10Classifier handler class. This handler extends ImageClassifier class\n",
        "    \"\"\"\n",
        "\n",
        "    label_names = [\n",
        "        \"plane\", \"car\", \"bird\", \"cat\", \"deer\", \"dog\", \"frog\", \"horse\", \"ship\", \"truck\"\n",
        "    ]\n",
        "\n",
        "    image_processing = transforms.Compose([\n",
        "        transforms.ToTensor(),\n",
        "        transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])\n",
        "    ])\n",
        "\n",
        "    def __init__(self):\n",
        "        super(Cifar10Classifier, self).__init__()\n",
        "        self.profiler_args = {\n",
        "            \"activities\": [ProfilerActivity.CPU],\n",
        "            \"record_shapes\": True,\n",
        "        }\n",
        "\n",
        "    def postprocess(self, data):\n",
        "        \"\"\"\n",
        "        Post-process function to convert the predicted class id to label\n",
        "        \"\"\"\n",
        "        predictions = data.argmax(1).tolist()\n",
        "        return [self.label_names[pred] for pred in predictions]\n",
        "\n",
        "'''\n",
        "\n",
        "with open(build_path / \"custom_handler.py\", \"w\") as custom_handler_file:\n",
        "    custom_handler_file.write(custom_handler_script)\n",
        "custom_handler_file.close()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "M2YY73mBByKt"
      },
      "source": [
        "#### Package the model artifacts in a model archive file\n",
        "\n",
        "Use the Torch Model Archiver tool to create a model archive file (.mar).\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "FVzPb0oxETee"
      },
      "outputs": [],
      "source": [
        "build_script = \"\"\"\n",
        "torch-model-archiver -f --model-name cifar10 \\\n",
        "    --version 1.0  \\\n",
        "    --model-file model.py \\\n",
        "    --serialized-file model.pt \\\n",
        "    --handler custom_handler.py \\\n",
        "    --export-path .\n",
        "\"\"\"\n",
        "\n",
        "with open(build_path / \"build.sh\", \"w\") as build_file:\n",
        "    build_file.write(build_script)\n",
        "build_file.close()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "d09sIVOaCFQE"
      },
      "outputs": [],
      "source": [
        "! cd {str(build_path)} && chmod +x ./build.sh && ./build.sh"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0zljbquGFtHP"
      },
      "source": [
        "#### Upload `mar` to bucket\n",
        "\n",
        "Store the .mar file to Cloud bucket."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "NPejvg6uFzBG"
      },
      "outputs": [],
      "source": [
        "!gsutil cp {str(build_path)}/cifar10.mar {BUILD_URI}/model.mar"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Yi3I6JotDw3k"
      },
      "source": [
        "### Register model in Vertex AI Model Registry\n",
        "\n",
        "Register the model as a Model resource in Vertex AI Model Registry."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9kyOITop38H7"
      },
      "outputs": [],
      "source": [
        "registered_model = vertex_ai.Model.upload(\n",
        "    display_name=MODEL_NAME,\n",
        "    serving_container_image_uri=DEPLOY_IMAGE_URI,\n",
        "    artifact_uri=BUILD_URI,\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ko3IfZgbEG1J"
      },
      "source": [
        "### Deploy model for predictions\n",
        "\n",
        "Create a Vertex AI endpoint and deploy the registered model for predictions."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "MSukFOf04IMv"
      },
      "outputs": [],
      "source": [
        "endpoint = registered_model.deploy(\n",
        "    deployed_model_display_name=MODEL_NAME,\n",
        "    machine_type=DEPLOY_COMPUTE,\n",
        "    accelerator_type=None,\n",
        "    accelerator_count=0,\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "AzHtXzv0E2sg"
      },
      "source": [
        "## Make online predictions\n",
        "\n",
        "Sample an image from the `CIFAR10` dataset for getting online predictions."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "kdRjvCOX4hVn"
      },
      "outputs": [],
      "source": [
        "test_dataset = datasets.CIFAR10(\n",
        "    root=tests_path / \"data\",\n",
        "    transform=transforms.ToTensor(),\n",
        "    train=False,\n",
        "    download=True,\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "VFuZC39WbIJL"
      },
      "source": [
        "Visualize the sampled image."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "fDWLHAuPWo_7"
      },
      "outputs": [],
      "source": [
        "pil_image = plot_image_sample(test_dataset)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "TeuZLWpRbLyT"
      },
      "source": [
        "Run a prediction request and get the predicted class."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "kPqDTh184fM0"
      },
      "outputs": [],
      "source": [
        "predictions = predict_from_image(pil_image, endpoint)\n",
        "\n",
        "for pred in predictions:\n",
        "    print(\"Predicted class:\", pred)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "TpV-iwP9qw9c"
      },
      "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": "sx_vKniMq9ZX"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "\n",
        "delete_endpoint = False\n",
        "delete_model = False\n",
        "delete_ray_cluster = False\n",
        "delete_bucket = False\n",
        "delete_tutorial = False\n",
        "\n",
        "# Delete endpoint resource\n",
        "if delete_endpoint or os.getenv(\"IS_TESTING\"):\n",
        "    endpoint.delete(force=True)\n",
        "\n",
        "# Delete model resource\n",
        "if delete_model or os.getenv(\"IS_TESTING\"):\n",
        "    registered_model.delete()\n",
        "\n",
        "# Delete ray on vertex cluster\n",
        "if delete_ray_cluster or os.getenv(\"IS_TESTING\"):\n",
        "    vertex_ray.delete_ray_cluster(ray_cluster.cluster_resource_name)\n",
        "\n",
        "# Delete tutorial folder\n",
        "if delete_tutorial or os.getenv(\"IS_TESTING\"):\n",
        "    shutil.rmtree(tutorial_path)\n",
        "\n",
        "# Delete Cloud Storage objects that were created\n",
        "if delete_bucket or os.getenv(\"IS_TESTING\"):\n",
        "    ! gsutil -q -m rm -r $BUCKET_URI"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "name": "get_started_with_pytorch_rov.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}