notebooks/community/prediction/custom_prediction_routines/SDK_Pytorch_Custom_Predict.ipynb

{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ur8xi4C7S06n"
      },
      "outputs": [],
      "source": [
        "# Copyright 2022 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": [
        "<table align=\"left\">\n",
        "  <td>\n",
        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/prediction/custom_prediction_routines/SDK_Pytorch_Custom_Predict.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/vertex-ai-samples/main/notebooks/community/prediction/custom_prediction_routines/SDK_Pytorch_Custom_Predict.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>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tvgnzT1CKxrO"
      },
      "source": [
        "## Overview\n",
        "\n",
        "This tutorial demonstrates how to use Vertex AI SDK to build a custom container that uses the Custom Prediction Routine model server to serve a PyTorch model on Vertex AI Predictions.\n",
        "\n",
        "\n",
        "\n",
        "### Dataset\n",
        "\n",
        "This tutorial uses R.A. Fisher's Iris dataset, a small dataset that is popular for trying out machine learning techniques. Each instance has four numerical features, which are different measurements of a flower, and a target label that\n",
        "marks it as one of three types of iris: Iris setosa, Iris versicolour, or Iris virginica.\n",
        "\n",
        "This tutorial uses [iris dataset](https://archive.ics.uci.edu/ml/datasets/iris).\n",
        "\n",
        "### Objective\n",
        "\n",
        "The goal is to:\n",
        "- Train a model that uses a flower's measurements as input to predict what type of iris it is.\n",
        "- Save the model.\n",
        "- Build a custom PyTorch serving container with custom preprocessing using the Custom Prediction Routine feature in the Vertex AI SDK.\n",
        "- Test the built container locally.\n",
        "- Upload and deploy custom container to Vertex Prediction.\n",
        "\n",
        "This tutorial focuses more on deploying this model with Vertex AI than on\n",
        "the design of the model itself.\n",
        "\n",
        "### Costs \n",
        "\n",
        "This tutorial uses billable components of Google Cloud:\n",
        "\n",
        "* Vertex AI\n",
        "\n",
        "Learn about [Vertex AI\n",
        "pricing](https://cloud.google.com/vertex-ai/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": "ze4-nDLfK4pw"
      },
      "source": [
        "### Set up your local development environment\n",
        "\n",
        "**If you are using Vertex AI Workbench Notebooks**, your environment already meets\n",
        "all the requirements to run this notebook. You can skip this step."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gCuSR8GkAgzl"
      },
      "source": [
        "**Otherwise**, make sure your environment meets this notebook's requirements.\n",
        "You need the following:\n",
        "\n",
        "* Docker\n",
        "* Git\n",
        "* Google Cloud SDK (gcloud)\n",
        "* Python 3\n",
        "* virtualenv\n",
        "* Jupyter notebook running in a virtual environment with Python 3\n",
        "\n",
        "The Google Cloud guide to [Setting up a Python development\n",
        "environment](https://cloud.google.com/python/setup) and the [Jupyter\n",
        "installation guide](https://jupyter.org/install) provide detailed instructions\n",
        "for meeting these requirements. The following steps provide a condensed set of\n",
        "instructions:\n",
        "\n",
        "1. [Install and initialize the Cloud SDK.](https://cloud.google.com/sdk/docs/)\n",
        "\n",
        "1. [Install Python 3.](https://cloud.google.com/python/setup#installing_python)\n",
        "\n",
        "1. [Install\n",
        "   virtualenv](https://cloud.google.com/python/setup#installing_and_using_virtualenv)\n",
        "   and create a virtual environment that uses Python 3. Activate the virtual environment.\n",
        "\n",
        "1. To install Jupyter, run `pip install jupyter` on the\n",
        "command-line in a terminal shell.\n",
        "\n",
        "1. To launch Jupyter, run `jupyter notebook` on the command-line in a terminal shell.\n",
        "\n",
        "1. Open this notebook in the Jupyter Notebook Dashboard."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "i7EUnXsZhAGF"
      },
      "source": [
        "### Install additional packages\n",
        "\n",
        "Install additional package dependencies not installed in your notebook environment, such as NumPy, Scikit-learn, FastAPI, Uvicorn, and joblib."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "747f59abb3a5"
      },
      "outputs": [],
      "source": [
        "%%writefile requirements.txt\n",
        "fastapi\n",
        "uvicorn==0.17.6\n",
        "pandas\n",
        "torch==1.11.0\n",
        "google-cloud-storage>=1.26.0,<2.0.0dev\n",
        "google-cloud-aiplatform[prediction]>=1.16.0"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "356f45ac7942"
      },
      "source": [
        "**The model you deploy will have a different set of dependencies pre-installed than your notebook environment has. You should not assume that because things work in the notebook, they will work in the model. Instead, we will be very explicit about the dependencies for the model by listing them in requirements.txt and then use `pip install` to install the exact same dependencies in the notebook. Please note, of course, that there is a chance that we miss a dependency in requirements.txt that already exists in the notebook. If that's the case, things will run in the notebook, but not in the model. To guard against that, we will test the model locally before deploying to the cloud.**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "wyy5Lbnzg5fi"
      },
      "outputs": [],
      "source": [
        "# Install the same dependencies used in the serving container in the notebook\n",
        "# environment.\n",
        "%pip install -U --user -r requirements.txt"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hhq5zEbGg0XX"
      },
      "source": [
        "### Restart the kernel\n",
        "\n",
        "After you install the additional packages, you need to restart the notebook kernel so it can find the packages."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "EzrelQZ22IZj"
      },
      "outputs": [],
      "source": [
        "# Automatically restart kernel after installs\n",
        "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",
      "metadata": {
        "id": "lWEdiXsJg0XY"
      },
      "source": [
        "## Before you begin"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "64a0a87c5b66"
      },
      "source": [
        "### Set up logging in this sample"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "348ca5392459"
      },
      "outputs": [],
      "source": [
        "import logging\n",
        "\n",
        "logging.basicConfig(level=logging.INFO)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BF1j6f9HApxa"
      },
      "source": [
        "### 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",
        "1. [Make sure that billing is enabled for your project](https://cloud.google.com/billing/docs/how-to/modify-project).\n",
        "\n",
        "1. [Enable the Vertex AI API and Compute Engine API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com,compute_component).\n",
        "\n",
        "1. If you are running this notebook locally, you will need to install the [Cloud SDK](https://cloud.google.com/sdk).\n",
        "\n",
        "1. Enter your project ID in the cell below. Then run the cell to make sure the\n",
        "Cloud SDK uses the right project for all the commands in this notebook.\n",
        "\n",
        "**Note**: Jupyter runs lines prefixed with `!` or `%` as shell commands, and it interpolates Python variables with `$` or `{}` into these commands."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WReHDGG5g0XY"
      },
      "source": [
        "#### Set your project ID\n",
        "\n",
        "**If you don't know your project ID**, you may be able to get your project ID using `gcloud`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "oM1iC_MfAts1"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "\n",
        "PROJECT_ID = \"\"\n",
        "\n",
        "# Get your Google Cloud project ID from gcloud\n",
        "if not os.getenv(\"IS_TESTING\"):\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": "markdown",
      "metadata": {
        "id": "qJYoRfYng0XZ"
      },
      "source": [
        "Otherwise, set your project ID here."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "riG_qUokg0XZ"
      },
      "outputs": [],
      "source": [
        "if PROJECT_ID == \"\" or PROJECT_ID is None:\n",
        "    PROJECT_ID = \"[your-project-id]\"  # @param {type:\"string\"}"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "set_gcloud_project_id"
      },
      "outputs": [],
      "source": [
        "! gcloud config set project $PROJECT_ID"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "region"
      },
      "source": [
        "#### Region\n",
        "\n",
        "You can also change the `REGION` variable, which is used for operations\n",
        "throughout the rest of this notebook.  Below are regions supported for Vertex AI. We recommend that you choose the region closest to you.\n",
        "\n",
        "- Americas: `us-central1`\n",
        "- Europe: `europe-west4`\n",
        "- Asia Pacific: `asia-east1`\n",
        "\n",
        "You may not use a multi-regional bucket for training with Vertex AI. Not all regions provide support for all Vertex AI services.\n",
        "\n",
        "Learn more about [Vertex AI regions](https://cloud.google.com/vertex-ai/docs/general/locations)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "region"
      },
      "outputs": [],
      "source": [
        "REGION = \"[your-region]\"  # @param {type: \"string\"}\n",
        "\n",
        "if REGION == \"[your-region]\":\n",
        "    REGION = \"us-central1\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "XoEqT2Y4DJmf"
      },
      "source": [
        "### Configure project and resource names"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UYxy9q6vVNIw"
      },
      "source": [
        "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 it onto the name of resources you create in this tutorial."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "AZ6rEERzVPdb"
      },
      "outputs": [],
      "source": [
        "from datetime import datetime\n",
        "\n",
        "TIMESTAMP = datetime.now().strftime(\"%Y%m%d%H%M%S\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vE0Zo2mxVVMT"
      },
      "source": [
        "Configure GCP resource names."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "MzGDU7TWdts_"
      },
      "outputs": [],
      "source": [
        "MODEL_ARTIFACT_DIR = \"pytorch-cpr-model-sdk\"  # @param {type:\"string\"}\n",
        "REPOSITORY = \"custom-container-prediction-sdk\"  # @param {type:\"string\"}\n",
        "IMAGE = \"pytorch-cpr-server-sdk\"  # @param {type:\"string\"}\n",
        "MODEL_DISPLAY_NAME = \"pytorch-cpr-model-sdk\"  # @param {type:\"string\"}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ca1a915d641d"
      },
      "source": [
        "`REGION` - Used for operations\n",
        "throughout the rest of this notebook. Make sure to [choose a region where Cloud\n",
        "Vertex AI services are\n",
        "available](https://cloud.google.com/vertex-ai/docs/general/locations#feature-availability). You may\n",
        "not use a Multi-Regional Storage bucket for prediction with Vertex AI.\n",
        "\n",
        "`MODEL_ARTIFACT_DIR` - Folder directory path to your model artifacts within a Cloud Storage bucket, for example: \"my-models/fraud-detection/trial-4\"\n",
        "\n",
        "`REPOSITORY` - Name of the Artifact Repository to create or use.\n",
        "\n",
        "`IMAGE` - Name of the container image that will be pushed.\n",
        "\n",
        "`MODEL_DISPLAY_NAME` - Display name of Vertex AI Model resource."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zgPO1eR3CYjk"
      },
      "source": [
        "### Create a Cloud Storage bucket\n",
        "\n",
        "**The following steps are required, regardless of your notebook environment.**\n",
        "\n",
        "To update your model artifacts without re-building the container, you must upload your model\n",
        "artifacts and any custom code to Cloud Storage.\n",
        "\n",
        "Set the name of your Cloud Storage bucket below. It must be unique across all\n",
        "Cloud Storage buckets."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "MzGDU7TWdts_"
      },
      "outputs": [],
      "source": [
        "BUCKET_NAME = \"[your-bucket-name]\"  # @param {type:\"string\"}\n",
        "BUCKET_URI = f\"gs://{BUCKET_NAME}\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "cf221059d072"
      },
      "outputs": [],
      "source": [
        "if BUCKET_NAME == \"\" or BUCKET_NAME is None or BUCKET_NAME == \"[your-bucket-name]\":\n",
        "    BUCKET_NAME = PROJECT_ID + \"-aip-\" + TIMESTAMP\n",
        "    BUCKET_URI = f\"gs://{BUCKET_NAME}\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-EcIXiGsCePi"
      },
      "source": [
        "**Only 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": "ucvCsknMCims"
      },
      "source": [
        "Finally, validate access to your Cloud Storage bucket by examining its contents:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "vhOb7YnwClBb"
      },
      "outputs": [],
      "source": [
        "! gsutil ls -al $BUCKET_URI"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3c2d091d9e73"
      },
      "source": [
        "### Set up directories"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3c2d091d9e73"
      },
      "source": [
        "Decide the directory to put your all custom code."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "3c2d091d9e73"
      },
      "outputs": [],
      "source": [
        "USER_SRC_DIR = \"src_dir_pytorch\"  # @param {type:\"string\"}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cee1ac13b767"
      },
      "source": [
        "Decide the directory to put your trained models."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "104037d0dbf5"
      },
      "outputs": [],
      "source": [
        "LOCAL_MODEL_ARTIFACTS_DIR = \"model_artifacts\"  # @param {type:\"string\"}"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6e74556ea0b4"
      },
      "outputs": [],
      "source": [
        "%mkdir $USER_SRC_DIR\n",
        "%mkdir $LOCAL_MODEL_ARTIFACTS_DIR"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "c389f200eed5"
      },
      "source": [
        "### Download iris data\n",
        "In this example, we want to build a classifier for the simple [iris dataset](https://archive.ics.uci.edu/ml/datasets/iris). So first, we download the data csv file locally."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "950c429a65ec"
      },
      "outputs": [],
      "source": [
        "DATA_DIR = \"data_iris\"\n",
        "\n",
        "%mkdir $DATA_DIR\n",
        "\n",
        "LOCAL_DATA_FILE = f\"{DATA_DIR}/iris.csv\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "3759597371ab"
      },
      "outputs": [],
      "source": [
        "from urllib.request import urlretrieve\n",
        "\n",
        "urlretrieve(\n",
        "    \"https://archive.ics.uci.edu/ml/machine-learning-databases/iris/iris.data\",\n",
        "    LOCAL_DATA_FILE,\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "29fcb6559390"
      },
      "source": [
        "### Build a PyTorch NN Classifier"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4b816cd52f4b"
      },
      "source": [
        "Make sure that pytorch package is [installed](https://pytorch.org/get-started/locally/)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "43e47249f736"
      },
      "outputs": [],
      "source": [
        "import torch\n",
        "from torch.autograd import Variable\n",
        "\n",
        "print(\"PyTorch Version: {}\".format(torch.__version__))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3c693178ec3b"
      },
      "source": [
        "#### Step 1. Load data\n",
        "\n",
        "In this step, we are going to:\n",
        "1. Load the data to Pandas Dataframe.\n",
        "1. Convert the class feature (species) from string to a numeric indicator.\n",
        "1. Split the Dataframe into input feature (xtrain) and target feature (ytrain)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ecac2569ceec"
      },
      "outputs": [],
      "source": [
        "import pandas as pd\n",
        "\n",
        "CLASS_VOCAB = [\"setosa\", \"versicolor\", \"virginica\"]\n",
        "\n",
        "datatrain = pd.read_csv(\n",
        "    LOCAL_DATA_FILE,\n",
        "    names=[\"sepal_length\", \"sepal_width\", \"petal_length\", \"petal_width\", \"species\"],\n",
        ")\n",
        "\n",
        "# change string value to numeric\n",
        "datatrain.loc[datatrain[\"species\"] == \"Iris-setosa\", \"species\"] = 0\n",
        "datatrain.loc[datatrain[\"species\"] == \"Iris-versicolor\", \"species\"] = 1\n",
        "datatrain.loc[datatrain[\"species\"] == \"Iris-virginica\", \"species\"] = 2\n",
        "datatrain = datatrain.apply(pd.to_numeric)\n",
        "\n",
        "# change dataframe to array\n",
        "datatrain_array = datatrain.values\n",
        "\n",
        "# split x and y (feature and target)\n",
        "xtrain = datatrain_array[:, :4]\n",
        "ytrain = datatrain_array[:, 4]\n",
        "\n",
        "input_features = xtrain.shape[1]\n",
        "num_classes = len(CLASS_VOCAB)\n",
        "\n",
        "print(\"Records loaded: {}\".format(len(xtrain)))\n",
        "print(\"Number of input features: {}\".format(input_features))\n",
        "print(\"Number of classes: {}\".format(num_classes))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cc140a358557"
      },
      "source": [
        "#### Step 2. Set model parameters\n",
        "You can try different values for **hidden_units** or **learning_rate**."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "296a7986f8f2"
      },
      "outputs": [],
      "source": [
        "HIDDEN_UNITS = 10\n",
        "LEARNING_RATE = 0.1"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dc58f3f86d9f"
      },
      "source": [
        "#### Step 3. Define the PyTorch NN model\n",
        "Here, we build a a neural network with one hidden layer, and a Softmax output layer for classification."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "0a6e2ebdb237"
      },
      "outputs": [],
      "source": [
        "model = torch.nn.Sequential(\n",
        "    torch.nn.Linear(input_features, HIDDEN_UNITS),\n",
        "    torch.nn.Sigmoid(),\n",
        "    torch.nn.Linear(HIDDEN_UNITS, num_classes),\n",
        "    torch.nn.Softmax(),\n",
        ")\n",
        "\n",
        "loss_metric = torch.nn.CrossEntropyLoss()\n",
        "optimizer = torch.optim.SGD(model.parameters(), lr=LEARNING_RATE)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "60c65751452b"
      },
      "source": [
        "#### Step 4. Train the model\n",
        "We are going to train the model for **num_epoch** epochs."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "2a38c4a8f6cd"
      },
      "outputs": [],
      "source": [
        "NUM_EPOCHS = 10000\n",
        "\n",
        "for epoch in range(NUM_EPOCHS):\n",
        "\n",
        "    x = Variable(torch.Tensor(xtrain).float())\n",
        "    y = Variable(torch.Tensor(ytrain).long())\n",
        "    optimizer.zero_grad()\n",
        "    y_pred = model(x)\n",
        "    loss = loss_metric(y_pred, y)\n",
        "    loss.backward()\n",
        "    optimizer.step()\n",
        "    if (epoch) % 1000 == 0:\n",
        "        print(\n",
        "            \"Epoch [{}/{}] Loss: {}\".format(\n",
        "                epoch + 1, NUM_EPOCHS, round(loss.item(), 3)\n",
        "            )\n",
        "        )\n",
        "\n",
        "print(\"Epoch [{}/{}] Loss: {}\".format(epoch + 1, NUM_EPOCHS, round(loss.item(), 3)))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "a22f07e6f56b"
      },
      "source": [
        "#### Step 5. Save the model"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "48172aba64eb"
      },
      "outputs": [],
      "source": [
        "LOCAL_MODEL_FILE = f\"{LOCAL_MODEL_ARTIFACTS_DIR}/model.pt\"\n",
        "\n",
        "torch.save(model, LOCAL_MODEL_FILE)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7d8ac3463cbb"
      },
      "source": [
        "#### Step 6. Test the loaded model for predictions"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "c97cdbb9f792"
      },
      "outputs": [],
      "source": [
        "iris_classifier = torch.load(LOCAL_MODEL_FILE)\n",
        "\n",
        "\n",
        "def predict_class(instances):\n",
        "    instances = torch.Tensor(instances)\n",
        "    output = iris_classifier(instances)\n",
        "    _, predicted = torch.max(output, 1)\n",
        "    return predicted\n",
        "\n",
        "\n",
        "predicted = predict_class(xtrain[0:5])\n",
        "print([CLASS_VOCAB[class_index] for class_index in predicted])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3849066a33bd"
      },
      "source": [
        "### Upload model artifacts and custom code to Cloud Storage\n",
        "\n",
        "Before you can deploy your model for serving, Vertex AI needs access to the following files in Cloud Storage:\n",
        "\n",
        "* `model.pt` (model artifact)\n",
        "\n",
        "Run the following commands to upload your files:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ca67ee52d4d9"
      },
      "outputs": [],
      "source": [
        "!gsutil cp {LOCAL_MODEL_ARTIFACTS_DIR}/* {BUCKET_URI}/{MODEL_ARTIFACT_DIR}/\n",
        "!gsutil ls {BUCKET_URI}/{MODEL_ARTIFACT_DIR}/"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "b2b675059515"
      },
      "source": [
        "## Build a custom serving container using the CPR model server"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "480a1d88ecdb"
      },
      "source": [
        "Now that the model and processor has been trained and saved, it's time to build the custom serving container. Typically building a serving container requires writing model server code. However, with the Custom Prediction Routine feature, Vertex AI Prediction provides a model server that can be used out of the box.\n",
        "\n",
        "A custom serving container contains the following 3 pieces of code:\n",
        "1. [Model Server](https://github.com/googleapis/python-aiplatform/blob/main/google/cloud/aiplatform/prediction/model_server.py)\n",
        "    * HTTP server that hosts the model.\n",
        "    * Responsible for setting up routes/ports/etc.\n",
        "1. [Request Handler](https://github.com/googleapis/python-aiplatform/blob/main/google/cloud/aiplatform/prediction/handler.py)\n",
        "    * Responsible for webserver aspects of handling a request, such as deserializing the request body, serializing the response, setting response headers, etc.\n",
        "    * In this example, we will use the default Handler, `google.cloud.aiplatform.prediction.handler.PredictionHandler` provided in the SDK.\n",
        "1. [Predictor](https://github.com/googleapis/python-aiplatform/blob/main/google/cloud/aiplatform/prediction/predictor.py)\n",
        "    * Responsible for the ML logic for processing a prediction request.\n",
        "\n",
        "Each of these three pieces can be customized based on the requirements of the custom container. In this example, we will only be implementing the `Predictor`.\n",
        "\n",
        "\n",
        "A [`Predictor`](https://github.com/googleapis/python-aiplatform/blob/main/google/cloud/aiplatform/prediction/predictor.py) must implement the following interface:\n",
        "\n",
        "```\n",
        "from abc import ABC, abstractmethod\n",
        "from typing import Any\n",
        "\n",
        "\n",
        "class Predictor(ABC):\n",
        "    \"\"\"Interface of the Predictor class for Custom Prediction Routines.\n",
        "    The Predictor is responsible for the ML logic for processing a prediction request.\n",
        "    Specifically, the Predictor must define:\n",
        "    (1) How to load all model artifacts used during prediction into memory.\n",
        "    (2) The logic that should be executed at predict time.\n",
        "    When using the default PredictionHandler, the Predictor will be invoked as follows:\n",
        "      predictor.postprocess(predictor.predict(predictor.preprocess(prediction_input)))\n",
        "    \"\"\"\n",
        "\n",
        "    def __init__(self):\n",
        "        return\n",
        "\n",
        "    @abstractmethod\n",
        "    def load(self, artifacts_uri: str) -> None:\n",
        "        \"\"\"Loads the model artifact.\n",
        "        Args:\n",
        "            artifacts_uri (str):\n",
        "                Required. The value of the environment variable AIP_STORAGE_URI.\n",
        "        \"\"\"\n",
        "        pass\n",
        "\n",
        "    def preprocess(self, prediction_input: Any) -> Any:\n",
        "        \"\"\"Preprocesses the prediction input before doing the prediction.\n",
        "        Args:\n",
        "            prediction_input (Any):\n",
        "                Required. The prediction input that needs to be preprocessed.\n",
        "        Returns:\n",
        "            The preprocessed prediction input.\n",
        "        \"\"\"\n",
        "        return prediction_input\n",
        "\n",
        "    @abstractmethod\n",
        "    def predict(self, instances: Any) -> Any:\n",
        "        \"\"\"Performs prediction.\n",
        "        Args:\n",
        "            instances (Any):\n",
        "                Required. The instance(s) used for performing prediction.\n",
        "        Returns:\n",
        "            Prediction results.\n",
        "        \"\"\"\n",
        "        pass\n",
        "\n",
        "    def postprocess(self, prediction_results: Any) -> Any:\n",
        "        \"\"\"Postprocesses the prediction results.\n",
        "        Args:\n",
        "            prediction_results (Any):\n",
        "                Required. The prediction results.\n",
        "        Returns:\n",
        "            The postprocessed prediction results.\n",
        "        \"\"\"\n",
        "        return prediction_results\n",
        "```"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-KODNAqK5k8R"
      },
      "source": [
        "First, implement a custom `Predictor` that loads in both the preprocesor and the model. The preprocessor and the model will then be used at `predict` time.\n",
        "\n",
        "Custom Prediction Routine supports a way to run the containers locally for testing your images. You can pass either a GCS path or a local path while testing your images locally.\n",
        "- You need to set up the credentials if you pass a GCS path. \n",
        "- You need to support loading your models remotely and locally in your `Predictor` if you want to testing by passing a local path.\n",
        "\n",
        "Vertex SDK provides a [function `download_model_artifacts`](https://github.com/googleapis/python-aiplatform/blob/main/google/cloud/aiplatform/utils/prediction_utils.py) to help you download model artifacts from either GCS paths or local paths. See the example in the `load` function below."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "fkh8DT7CCgQO"
      },
      "outputs": [],
      "source": [
        "%%writefile $USER_SRC_DIR/predictor.py\n",
        "\n",
        "import pandas as pd\n",
        "import pickle\n",
        "import torch\n",
        "from typing import Dict\n",
        "\n",
        "from google.cloud.aiplatform.prediction.predictor import Predictor\n",
        "from google.cloud.aiplatform.utils import prediction_utils\n",
        "\n",
        "\n",
        "class CustomPyTorchPredictor(Predictor):\n",
        "    \n",
        "    def __init__(self):\n",
        "        self._class_names = [\"setosa\", \"versicolor\", \"virginica\"]\n",
        "    \n",
        "    def load(self, artifacts_uri: str):\n",
        "        \"\"\"Loads the model artifacts.\"\"\"\n",
        "        prediction_utils.download_model_artifacts(artifacts_uri)\n",
        "\n",
        "        self._model = torch.load(\"model.pt\")\n",
        "\n",
        "    def preprocess(self, prediction_input: Dict) -> torch.Tensor:\n",
        "        instances = prediction_input[\"instances\"]\n",
        "        data = pd.DataFrame(instances).values\n",
        "        return torch.Tensor(data)\n",
        "\n",
        "    @torch.inference_mode()\n",
        "    def predict(self, instances: torch.Tensor) -> torch.Tensor:\n",
        "        \"\"\"Performs prediction.\"\"\"\n",
        "        outputs = self._model(instances)\n",
        "        _ , predicted = torch.max(outputs, 1)\n",
        "        return predicted\n",
        "\n",
        "    def postprocess(self, prediction_results: torch.Tensor) -> Dict:\n",
        "        return {\"predictions\": [self._class_names[class_num] for class_num in prediction_results]}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "waA_Jl5Q6E0W"
      },
      "source": [
        "To build a custom container, we also need to write an entrypoint of the image that starts the model server. However, with the Custom Prediction Routine feature, you don't need to write the entrypoint anymore. Vertex SDK will populate the entrypoint with the custom predictor you provide."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "04ddb6d57ceb"
      },
      "source": [
        "Write the dependencies to the source directory which will be installed in the image."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "a94b4589fb20"
      },
      "outputs": [],
      "source": [
        "!cp requirements.txt $USER_SRC_DIR/requirements.txt"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "51e149fdec1b"
      },
      "source": [
        "## Build and push container to Artifact Registry"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3bdb9a7768a5"
      },
      "source": [
        "### Build your custom container\n",
        "\n",
        "To build a custom image, a Dockerfile is necessary where you need to implement what the image looks like. However, with the Custom Prediction Routine features, Vertex SDK also generates Dockerfile and build images for you.\n",
        "\n",
        "Using `python:3.7` as a base image by default."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "3d3a6b9ed22b"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "\n",
        "from google.cloud.aiplatform.prediction import LocalModel\n",
        "from src_dir_pytorch.predictor import \\\n",
        "    CustomPyTorchPredictor  # Update this path as the variable $USER_SRC_DIR to import the custom predictor.\n",
        "\n",
        "local_model = LocalModel.build_cpr_model(\n",
        "    USER_SRC_DIR,\n",
        "    f\"{REGION}-docker.pkg.dev/{PROJECT_ID}/{REPOSITORY}/{IMAGE}\",\n",
        "    predictor=CustomPyTorchPredictor,\n",
        "    requirements_path=os.path.join(USER_SRC_DIR, \"requirements.txt\"),\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "04c988201499"
      },
      "source": [
        "You can check out the serving container spec of the built image."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "f1e7d639b9cc"
      },
      "outputs": [],
      "source": [
        "local_model.get_serving_container_spec()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8b62ddf1def3"
      },
      "source": [
        "### Store test instances\n",
        "\n",
        "To learn more about formatting input instances in JSON, [read the documentation.](https://cloud.google.com/vertex-ai/docs/predictions/online-predictions-custom-models#request-body-details)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "caf53dbc65d9"
      },
      "outputs": [],
      "source": [
        "INPUT_FILE = \"instances.json\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "b6605e9e6186"
      },
      "outputs": [],
      "source": [
        "%%writefile $INPUT_FILE\n",
        "{\n",
        "    \"instances\": [\n",
        "        [6.7, 3.1, 4.7, 1.5],\n",
        "        [4.6, 3.1, 1.5, 0.2]\n",
        "    ]\n",
        "}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "147a555f6c93"
      },
      "source": [
        "### Run and test the container locally"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "a72cf39f6d53"
      },
      "source": [
        "Custom Prediction Routines provide two ways to download model artifacts for testing images locally.\n",
        "1. A local path.\n",
        "2. A GCS path.\n",
        "\n",
        "To use a local path, your `Predictor` needs to support both ways to load models so it can be tested locally and deployed to Vertex AI Prediction service. SDK provides a method `download_model_artifacts` to support these two ways in [prediction_utils.py](https://github.com/googleapis/python-aiplatform/blob/main/google/cloud/aiplatform/utils/prediction_utils.py) which you can call in the `load` function in your `Predictor`.\n",
        "\n",
        "If you want to test images locally with a GCS path, you need to follow the instructions to set up the credentials."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8phJWBhe7u9u"
      },
      "source": [
        "#### Set up credentials\n",
        "\n",
        "Setting up credentials is only required to run the custom serving container locally with GCS paths. Credentials set up is required to execute the `Predictor`'s `load` function, which downloads the model artifacts from Google Cloud Storage.\n",
        "\n",
        "To access Google Cloud Storage in your project, you'll need to set up credentials by using one of the following:\n",
        "\n",
        "1. User account\n",
        "\n",
        "1. Service account\n",
        "\n",
        "You can learn more about each of the above [here](https://cloud.google.com/docs/authentication#principals)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "eac44f62b898"
      },
      "source": [
        "**Option 1. Use Google user credentials**\n",
        "\n",
        "First authorize Google auth library to access the Cloud Platform with Google user credentials. This step involves an interactive prompt which requires user inputs. If you are using a non-interactive shell, you should open a terminal to run this command. See [here](https://jupyterlab.readthedocs.io/en/stable/user/terminal.html) for how to open a terminal in a Vertex Workbench notebook environment.\n",
        "\n",
        "Note that if you are running the command below on a machine without a window manager/browser, it will prompt you to run the command `gcloud auth application-default login --remote-bootstrap=\"...\"`. Only machines that can launch web browsers, e.g. laptops or workstations, are allowed to run this command. After completing the authentication steps using your browser, you'll need to copy and paste the verification code back into the original terminal to continue the login flow. If you are running this command on a machine with web browser support, there is no need to use a separate machine.\n",
        "\n",
        "This command will print the credential file as `Credentials saved to file: [CREDENTIALS_FILE]`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "867c773995d6"
      },
      "outputs": [],
      "source": [
        "!gcloud auth application-default login"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0fe3f1ff17e5"
      },
      "source": [
        "Specify the credential file, which is a json file, as the path."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "14d8e7e0691b"
      },
      "outputs": [],
      "source": [
        "CREDENTIALS_FILE = \"[CREDENTIALS_FILE]\"  # @param {type:\"string\"}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8a544c84b96f"
      },
      "source": [
        "Next, authorize gcloud to access the Cloud Platform with Google user credentials. The user credentials need to have `setIamPolicy` permission in the project. See [here](https://cloud.google.com/resource-manager/docs/access-control-proj) for more details.\n",
        "\n",
        "Similar to the previous step, this also requires user inputs. Instructions for the previous step apply here as well."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "e63e7d6a5e93"
      },
      "outputs": [],
      "source": [
        "!gcloud auth login"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "d8685d6f9aef"
      },
      "source": [
        "Grant roles to the user account. Use the email address that you used in the gcloud auth step, e.g., `myemail@gmail.com`. Since we will be using this user account to download the model artifacts from Google Cloud Storage, we grant the user account **Storage Object Viewer** role. See [IAM roles for Cloud Storage](https://cloud.google.com/storage/docs/access-control/iam-roles) for more details."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "dad0f966c10b"
      },
      "outputs": [],
      "source": [
        "USER_ACCOUNT = \"[USER_ACCOUNT]\"  # @param {type:\"string\"}\n",
        "\n",
        "!gcloud projects add-iam-policy-binding $PROJECT_ID \\\n",
        "    --member=user:$USER_ACCOUNT --role=roles/storage.objectViewer"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ca95b2d79b57"
      },
      "source": [
        "Initialize the Vertex AI SDK with the project ID, which is required for using user credentials."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7eb690cb1182"
      },
      "outputs": [],
      "source": [
        "from google.cloud import aiplatform\n",
        "\n",
        "aiplatform.init(project=PROJECT_ID, location=REGION)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "a6c474e860f3"
      },
      "source": [
        "**Option 2. Use Google Service Account credentials**\n",
        "\n",
        "Skip this if you've already completed Option 1 above. First enable the IAM API if it's not already enabled."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "mERTlLb6dluB"
      },
      "outputs": [],
      "source": [
        "!gcloud services enable iam.googleapis.com"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "b4c72d351bba"
      },
      "source": [
        "Specify service account name which should be unique in your project."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "fece5ffc5938"
      },
      "outputs": [],
      "source": [
        "SERVICE_ACCOUNT = \"[SERVICE_ACCOUNT]\"  # @param {type:\"string\"}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "42a08c49ceb3"
      },
      "source": [
        "Create the service account."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "398ac116171f"
      },
      "outputs": [],
      "source": [
        "!gcloud iam service-accounts create $SERVICE_ACCOUNT"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "d47ad5c4d7d2"
      },
      "source": [
        "Authorize gcloud to access the Cloud Platform with Google user credentials. The user credentials need to have `setIamPolicy` permission in the project. See [here](https://cloud.google.com/resource-manager/docs/access-control-proj) for more details.\n",
        "\n",
        "This step involves an interactive prompt which requires user inputs. If you are using a non-interactive shell, you should open a terminal to run this command. See [here](https://jupyterlab.readthedocs.io/en/stable/user/terminal.html) for how to open a terminal in a Vertex Workbench notebook environment.\n",
        "\n",
        "Note that if you are running the command below on a machine without a window manager/browser, it will prompt you to run the command `gcloud auth application-default login --remote-bootstrap=\"...\"`. Only machines that can launch web browsers, e.g. laptops or workstations, are allowed to run this command. After completing the authentication steps using your browser, you'll need to copy and paste the verification code back into the original terminal to continue the login flow. If you are running this command on a machine with web browser support, there is no need to use a separate machine."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "cf5458c115f5"
      },
      "outputs": [],
      "source": [
        "!gcloud auth login"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "af9c4cf647e6"
      },
      "source": [
        "Grant roles to the service account. Since we will be using this service account to download the model artifacts from Google Cloud Storage, we grant the service account **Storage Object Viewer** role. See [IAM roles for Cloud Storage](https://cloud.google.com/storage/docs/access-control/iam-roles) for more details."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "a90517374617"
      },
      "outputs": [],
      "source": [
        "!gcloud projects add-iam-policy-binding $PROJECT_ID \\\n",
        "    --member=serviceAccount:{SERVICE_ACCOUNT}@{PROJECT_ID}.iam.gserviceaccount.com \\\n",
        "    --role=roles/storage.objectViewer"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7417cf95a651"
      },
      "source": [
        "Specify the service account key file name."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "2968d51dae50"
      },
      "outputs": [],
      "source": [
        "CREDENTIALS_FILE = \"./credentials.json\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6fb51f953e19"
      },
      "source": [
        "Generate the service account key file. This will be used while running the custom serving container locally."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "356a0aeba296"
      },
      "outputs": [],
      "source": [
        "!gcloud iam service-accounts keys create $CREDENTIALS_FILE \\\n",
        "    --iam-account=\"{SERVICE_ACCOUNT}@{PROJECT_ID}.iam.gserviceaccount.com\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pJe8TCr4ec4F"
      },
      "source": [
        "#### Run and send requests to the container locally\n",
        "\n",
        "With the Custom Prediction Routine feature, it is easy to test the container locally.\n",
        "\n",
        "In this example, the container executes a prediction request and a health check."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6a8cc38a6cce"
      },
      "source": [
        "**Option 1. Pass local paths**\n",
        "\n",
        "This option requires the `load` function in the `Predictor` supports loading model artifacts from both of GCS paths and local paths."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "62ed2d334d0f"
      },
      "outputs": [],
      "source": [
        "with local_model.deploy_to_local_endpoint(\n",
        "    artifact_uri=f\"{LOCAL_MODEL_ARTIFACTS_DIR}\",\n",
        ") as local_endpoint:\n",
        "    predict_response = local_endpoint.predict(\n",
        "        request_file=INPUT_FILE,\n",
        "        headers={\"Content-Type\": \"application/json\"},\n",
        "    )\n",
        "\n",
        "    health_check_response = local_endpoint.run_health_check()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "07485db39fd0"
      },
      "source": [
        "Print out the predict response and its content."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ce629eea32fd"
      },
      "outputs": [],
      "source": [
        "predict_response, predict_response.content"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "e9d1b7982414"
      },
      "source": [
        "Print out the health check response and its content."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "56986f93438e"
      },
      "outputs": [],
      "source": [
        "health_check_response, health_check_response.content"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "a29fcbbe0188"
      },
      "source": [
        "Also print out all the container logs."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "90470050560d"
      },
      "outputs": [],
      "source": [
        "local_endpoint.print_container_logs(show_all=True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "39f41bb37659"
      },
      "source": [
        "**Option 2. Pass GCS paths**\n",
        "\n",
        "If you want to pass GCS paths, you need to have the credentials set up in the previous step and pass the path to the credentials while running the container. The service account should have the **Storage Object Viewer** permission."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "62ed2d334d0f"
      },
      "outputs": [],
      "source": [
        "with local_model.deploy_to_local_endpoint(\n",
        "    artifact_uri=f\"{BUCKET_URI}/{MODEL_ARTIFACT_DIR}\",\n",
        "    credential_path=CREDENTIALS_FILE,\n",
        ") as local_endpoint:\n",
        "    predict_response = local_endpoint.predict(\n",
        "        request_file=INPUT_FILE,\n",
        "        headers={\"Content-Type\": \"application/json\"},\n",
        "    )\n",
        "\n",
        "    health_check_response = local_endpoint.run_health_check()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cdbc0dfeeed0"
      },
      "source": [
        "Print out the predict response and its content."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ce629eea32fd"
      },
      "outputs": [],
      "source": [
        "predict_response, predict_response.content"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "193440c77fab"
      },
      "source": [
        "Print out the health check response and its content."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "56986f93438e"
      },
      "outputs": [],
      "source": [
        "health_check_response, health_check_response.content"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "a29fcbbe0188"
      },
      "source": [
        "Also print out all the container logs."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "bd28ac20dfb5"
      },
      "outputs": [],
      "source": [
        "local_endpoint.print_container_logs(show_all=True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "212b2935ea12"
      },
      "source": [
        "### Push the container to artifact registry\n",
        "\n",
        "Configure Docker to access Artifact Registry. Then push your container image to your Artifact Registry repository."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "wSFXCj3LdluJ"
      },
      "outputs": [],
      "source": [
        "!gcloud services list"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ABE9UpwSdluK"
      },
      "source": [
        "If `artifactregistry.googleapis.com` is not enabled in your project, enable the API before proceeding."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "qDhLoQMydluK"
      },
      "outputs": [],
      "source": [
        "!gcloud services enable artifactregistry.googleapis.com"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "09ffe2434e3d"
      },
      "outputs": [],
      "source": [
        "!gcloud artifacts repositories create {REPOSITORY} \\\n",
        "    --repository-format=docker \\\n",
        "    --location=$REGION"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "293437024749"
      },
      "outputs": [],
      "source": [
        "!gcloud auth configure-docker {REGION}-docker.pkg.dev --quiet"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "89af2542b0a2"
      },
      "source": [
        "Use SDK to push the image."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "1dd7448f4703"
      },
      "outputs": [],
      "source": [
        "local_model.push_image()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "b438bfa2129f"
      },
      "source": [
        "## Deploy to Vertex AI"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4ae19df6a33e"
      },
      "source": [
        "### Upload the custom container model"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "8d682d8388ec"
      },
      "outputs": [],
      "source": [
        "from google.cloud import aiplatform"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "574fb82d3eed"
      },
      "outputs": [],
      "source": [
        "aiplatform.init(project=PROJECT_ID, location=REGION)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "148292367a46"
      },
      "source": [
        "Use the LocalModel instance to upload the model. It will populate the container spec automatically for you."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "2738154345d5"
      },
      "outputs": [],
      "source": [
        "model = aiplatform.Model.upload(\n",
        "    local_model=local_model,\n",
        "    display_name=MODEL_DISPLAY_NAME,\n",
        "    artifact_uri=f\"{BUCKET_URI}/{MODEL_ARTIFACT_DIR}\",\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bd1b85afc7df"
      },
      "source": [
        "### Deploy the model on Vertex AI\n",
        "After this step completes, the model is deployed and ready for online prediction."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "62cf66498a28"
      },
      "outputs": [],
      "source": [
        "endpoint = model.deploy(machine_type=\"n1-standard-4\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6883e7b07143"
      },
      "source": [
        "## Send predictions\n",
        "\n",
        "### Using Python SDK"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "d69ed411c2d3"
      },
      "outputs": [],
      "source": [
        "endpoint.predict(instances=[[6.7, 3.1, 4.7, 1.5], [4.6, 3.1, 1.5, 0.2]])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "370d22f53427"
      },
      "source": [
        "### Using REST"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ba55bc560d58"
      },
      "outputs": [],
      "source": [
        "ENDPOINT_ID = endpoint.name"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "95c562b4e98b"
      },
      "outputs": [],
      "source": [
        "! curl \\\n",
        "-H \"Authorization: Bearer $(gcloud auth print-access-token)\" \\\n",
        "-H \"Content-Type: application/json\" \\\n",
        "-d @instances.json \\\n",
        "https://{REGION}-aiplatform.googleapis.com/v1/projects/{PROJECT_ID}/locations/{REGION}/endpoints/{ENDPOINT_ID}:predict"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "fa71174a7dd0"
      },
      "source": [
        "### Using gcloud CLI"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "23b8e807b02c"
      },
      "outputs": [],
      "source": [
        "!gcloud ai endpoints predict $ENDPOINT_ID \\\n",
        "  --region=$REGION \\\n",
        "  --json-request=instances.json"
      ]
    },
    {
      "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": [
        "# Undeploy model and delete endpoint\n",
        "endpoint.delete(force=True)\n",
        "\n",
        "# Delete the model resource\n",
        "model.delete()\n",
        "\n",
        "# Delete the container image from Artifact Registry\n",
        "!gcloud artifacts docker images delete \\\n",
        "    --quiet \\\n",
        "    --delete-tags \\\n",
        "    {REGION}-docker.pkg.dev/{PROJECT_ID}/{REPOSITORY}/{IMAGE}\n",
        "delete_bucket = False\n",
        "\n",
        "if delete_bucket or os.getenv(\"IS_TESTING\"):\n",
        "    ! gsutil rm -r $BUCKET_URI"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [],
      "name": "SDK_Pytorch_Custom_Predict.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}