notebooks/official/training/sdk_pytorch_torchrun_custom_container_training_imagenet.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": [
        "# Vertex AI SDK: Using PyTorch torchrun to simplify multi-node training with custom containers\n",
        "<table align=\"left\">\n",
        "\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/training/sdk_pytorch_torchrun_custom_container_training_imagenet.ipynb\">\n",
        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Colab logo\"><br> Run 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%2Ftraining%2Fsdk_pytorch_torchrun_custom_container_training_imagenet.ipynb\">\n",
        "      <img width=\"32px\" src=\"https://cloud.google.com/ml-engine/images/colab-enterprise-logo-32px.png\" alt=\"Google Cloud Colab Enterprise logo\"><br> Open in Colab Enterprise\n",
        "    </a>\n",
        "  </td>   \n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://console.cloud.google.com/vertex-ai/workbench/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/notebooks/official/training/sdk_pytorch_torchrun_custom_container_training_imagenet.ipynb\">\n",
        "      <img src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" alt=\"Vertex AI logo\"><br>\n",
        "      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/training/sdk_pytorch_torchrun_custom_container_training_imagenet.ipynb\">\n",
        "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\"><br>\n",
        "      View on GitHub\n",
        "    </a>\n",
        "  </td>\n",
        "</table>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "c119189aa978"
      },
      "source": [
        "## Overview\n",
        "\n",
        "This tutorial uses the Tiny ImageNet dataset to run multi-node distributed training on Vertex AI with Torchrun. It runs distributed training on multiple nodes with GPUs.\n",
        "\n",
        "Learn more about [Distributed training](https://cloud.google.com/vertex-ai/docs/training/distributed-training)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "26f217af913b"
      },
      "source": [
        "### Objective\n",
        "\n",
        "In this tutorial, you learn how to train an Imagenet model using PyTorch's Torchrun on multiple nodes.\n",
        "\n",
        "This tutorial uses the following Google Cloud ML services:\n",
        "\n",
        "- Vertex AI Training(Custom Python Package Training) \n",
        "\n",
        "The steps performed include:\n",
        "\n",
        "    * Create a shell script to start an ETCD cluster on the master node\n",
        "    * Create a training script using code from PyTorch Elastic's Github repository\n",
        "    * Create containers that download the data, and start an ETCD cluster on the host\n",
        "    * Train the model using multiple nodes with GPUs"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "c6a298ceea8e"
      },
      "source": [
        "### Dataset\n",
        "\n",
        "For the sake of training time, the Tiny ImageNet dataset is used in this tutorial: https://image-net.org/data/tiny-imagenet-200.zip\n",
        "\n",
        "This dataset consists of many small (~2KB) images. To avoid network bottlenecks with the large volume of network transfers from Cloud Storage to the GPUs, download this dataset to the containers\n",
        "\n",
        "The training code is based on this PyTorch Torchrun example for ImageNet: https://github.com/pytorch/elastic/blob/master/examples/imagenet/main.py"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tvgnzT1CKxrO"
      },
      "source": [
        "### Costs \n",
        "\n",
        "This tutorial uses billable components of Google Cloud:\n",
        "\n",
        "* Vertex AI Training w/ GPUs\n",
        "* Vertex AI TensorBoard\n",
        "* Cloud Storage\n",
        "\n",
        "Learn about [Vertex AI\n",
        "pricing](https://cloud.google.com/vertex-ai/pricing) and [Cloud Storage\n",
        "pricing](https://cloud.google.com/storage/pricing), and use the [Pricing\n",
        "Calculator](https://cloud.google.com/products/calculator/)\n",
        "to generate a cost estimate based on your projected usage."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "f0316df526f8"
      },
      "source": [
        "## Get started"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "i7EUnXsZhAGF"
      },
      "source": [
        "### Install Vertex AI SDK for Python and other required packages"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "2b4ef9b72d43"
      },
      "outputs": [],
      "source": [
        "# Install the packages\n",
        "! pip3 install --upgrade --quiet google-cloud-aiplatform \\\n",
        "                                 python-etcd"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "58707a750154"
      },
      "source": [
        "### Restart runtime (Colab only)\n",
        "\n",
        "To use the newly installed packages, you must restart the runtime on Google Colab."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "f200f10a1da3"
      },
      "outputs": [],
      "source": [
        "import sys\n",
        "\n",
        "if \"google.colab\" in sys.modules:\n",
        "\n",
        "    import IPython\n",
        "\n",
        "    app = IPython.Application.instance()\n",
        "    app.kernel.do_shutdown(True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ee775571c2b5"
      },
      "source": [
        "<div class=\"alert alert-block alert-warning\">\n",
        "<b>⚠️ The kernel is going to restart. Wait until it's finished before continuing to the next step. ⚠️</b>\n",
        "</div>\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "92e68cfc3a90"
      },
      "source": [
        "### Authenticate your notebook environment (Colab only)\n",
        "\n",
        "Authenticate your environment on Google Colab.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "46604f70e831"
      },
      "outputs": [],
      "source": [
        "import sys\n",
        "\n",
        "if \"google.colab\" in sys.modules:\n",
        "\n",
        "    from google.colab import auth\n",
        "\n",
        "    auth.authenticate_user()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BF1j6f9HApxa"
      },
      "source": [
        "### Set Google Cloud project information and initialize Vertex AI SDK for Python\n",
        "\n",
        "To get started using Vertex AI, you must have an existing Google Cloud project and [enable the Vertex AI API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com). Learn more about [setting up a project and a development environment](https://cloud.google.com/vertex-ai/docs/start/cloud-environment)."
      ]
    },
    {
      "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.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "MzGDU7TWdts_"
      },
      "outputs": [],
      "source": [
        "PROJECT_ID = \"[your-project-id]\"  # @param {type:\"string\"}\n",
        "LOCATION = \"us-central1\"  # @param {type:\"string\"}\n",
        "\n",
        "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 {LOCATION} -p {PROJECT_ID} {BUCKET_URI}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "37f325ab5cf8"
      },
      "source": [
        "### Initialize Vertex AI SDK for Python"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "PyQmSRbKA8r-"
      },
      "outputs": [],
      "source": [
        "from google.cloud import aiplatform\n",
        "\n",
        "aiplatform.init(project=PROJECT_ID, location=LOCATION, staging_bucket=BUCKET_URI)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "a7a84e2e4c4e"
      },
      "source": [
        "### Service Account\n",
        "\n",
        "You use a service account to create the Vertex AI Training job. If you do not want to use your project's Compute Engine service account, set SERVICE_ACCOUNT to another service account ID."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "e0c9c4f84849"
      },
      "outputs": [],
      "source": [
        "SERVICE_ACCOUNT = \"[your-service-account]\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2d9440bb3017"
      },
      "source": [
        "If you do not provide a service account, run the code below to get the Compute Engine service account"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "304a9ea0b6d0"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "\n",
        "# The Google Cloud Notebook product has specific requirements\n",
        "IS_GOOGLE_CLOUD_NOTEBOOK = os.path.exists(\"/opt/deeplearning/metadata/env_version\")\n",
        "\n",
        "if (\n",
        "    SERVICE_ACCOUNT == \"\"\n",
        "    or SERVICE_ACCOUNT is None\n",
        "    or SERVICE_ACCOUNT == \"[your-service-account]\"\n",
        "):\n",
        "    # Get your service account from gcloud\n",
        "    if IS_GOOGLE_CLOUD_NOTEBOOK:\n",
        "        shell_output = !gcloud auth list 2>/dev/null\n",
        "        SERVICE_ACCOUNT = shell_output[2].replace(\"*\", \"\").strip()\n",
        "\n",
        "    if not IS_GOOGLE_CLOUD_NOTEBOOK:\n",
        "        shell_output = ! gcloud projects describe  $PROJECT_ID\n",
        "        project_number = shell_output[-1].split(\":\")[1].strip().replace(\"'\", \"\")\n",
        "        SERVICE_ACCOUNT = f\"{project_number}-compute@developer.gserviceaccount.com\"\n",
        "\n",
        "    print(\"Service Account:\", SERVICE_ACCOUNT)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gar_enable_api"
      },
      "source": [
        "### Enable Artifact Registry API\n",
        "\n",
        "First, you must enable the Artifact Registry API service for your project.\n",
        "\n",
        "Learn more about [Enabling service](https://cloud.google.com/artifact-registry/docs/enable-service)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "gar_enable_api"
      },
      "outputs": [],
      "source": [
        "! gcloud services enable artifactregistry.googleapis.com\n",
        "\n",
        "if os.getenv(\"IS_TESTING\"):\n",
        "    ! sudo apt-get update --yes && sudo apt-get --only-upgrade --yes install google-cloud-sdk-cloud-run-proxy google-cloud-sdk-harbourbridge google-cloud-sdk-cbt google-cloud-sdk-gke-gcloud-auth-plugin google-cloud-sdk-kpt google-cloud-sdk-local-extract google-cloud-sdk-minikube google-cloud-sdk-app-engine-java google-cloud-sdk-app-engine-go google-cloud-sdk-app-engine-python google-cloud-sdk-spanner-emulator google-cloud-sdk-bigtable-emulator google-cloud-sdk-nomos google-cloud-sdk-package-go-module google-cloud-sdk-firestore-emulator kubectl google-cloud-sdk-datastore-emulator google-cloud-sdk-app-engine-python-extras google-cloud-sdk-cloud-build-local google-cloud-sdk-kubectl-oidc google-cloud-sdk-anthos-auth google-cloud-sdk-app-engine-grpc google-cloud-sdk-pubsub-emulator google-cloud-sdk-datalab google-cloud-sdk-skaffold google-cloud-sdk google-cloud-sdk-terraform-tools google-cloud-sdk-config-connector\n",
        "    ! gcloud components update --quiet"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gar_create_repo"
      },
      "source": [
        "### Create a private Docker repository\n",
        "\n",
        "Your first step is to create your own Docker repository in Artifact Registry.\n",
        "\n",
        "1. Run the `gcloud artifacts repositories create` command to create a new Docker repository with your region with the description \"docker repository\".\n",
        "\n",
        "2. Run the `gcloud artifacts repositories list` command to verify that your repository was created."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ea7cf85d87d2"
      },
      "outputs": [],
      "source": [
        "REPOSITORY = \"torchrun-imagenet-repo\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "gar_create_repo"
      },
      "outputs": [],
      "source": [
        "! gcloud artifacts repositories create {REPOSITORY} --repository-format=docker --location={LOCATION} --description=\"Docker repository\"\n",
        "\n",
        "! gcloud artifacts repositories list"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gar_auth"
      },
      "source": [
        "### Configure authentication to your private repo\n",
        "\n",
        "Before you push or pull container images, configure Docker to use the `gcloud` command-line tool to authenticate requests to `Artifact Registry` for your region."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "gar_auth"
      },
      "outputs": [],
      "source": [
        "! gcloud auth configure-docker {LOCATION}-docker.pkg.dev --quiet"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8f3ea1210749"
      },
      "source": [
        "## Vertex AI Training with GPUs"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2f6e7336d1a0"
      },
      "source": [
        "### Create files for the host container"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "fd9228f35579"
      },
      "outputs": [],
      "source": [
        "%mkdir -p trainer\n",
        "%cat /dev/null > trainer/__init__.py"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "70886930375a"
      },
      "source": [
        "#### Create the Dockerfile\n",
        "Installs necessary libraries, and downloads the tiny ImageNet data for training"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "b7da2e37e767"
      },
      "outputs": [],
      "source": [
        "%%writefile trainer/Dockerfile\n",
        "FROM gcr.io/deeplearning-platform-release/pytorch-gpu.1-13:m102\n",
        "\n",
        "RUN curl https://packages.cloud.google.com/apt/doc/apt-key.gpg | apt-key add - && \\\n",
        "    # Install reduction server plugin on GPU containers. google-fast-socket is\n",
        "    # previously installed in GPU dlenv containers only and it is not compatible\n",
        "    # with google-reduction-server.\n",
        "    if dpkg -s google-fast-socket; then \\\n",
        "      apt remove -y google-fast-socket && \\\n",
        "      apt install -y google-reduction-server; \\\n",
        "    fi\n",
        "\n",
        "RUN rm -f /etc/apt/sources.list.d/cuda.list && \\\n",
        "    rm -f /etc/apt/sources.list.d/nvidia-ml.list\n",
        "\n",
        "RUN apt-key del 7fa2af80 && \\\n",
        "    apt-key adv --fetch-keys http://developer.download.nvidia.com/compute/cuda/repos/ubuntu1804/x86_64/3bf863cc.pub && \\\n",
        "    apt-key adv --fetch-keys http://developer.download.nvidia.com/compute/machine-learning/repos/ubuntu1804/x86_64/7fa2af80.pub\n",
        "\n",
        "RUN apt-get update --allow-releaseinfo-change\n",
        "RUN apt-get update -y && \\\n",
        "    apt-get install -y curl gnupg telnet nano net-tools iputils-ping\n",
        "\n",
        "# Set ETCD version\n",
        "ARG ETCD_VER=v2.3.0\n",
        "# Choose either URL\n",
        "ARG GOOGLE_URL=https://storage.googleapis.com/etcd\n",
        "ARG GITHUB_URL=https://github.com/etcd-io/etcd/releases/download\n",
        "# Set ETCD URL to download from\n",
        "ARG DOWNLOAD_URL=$GOOGLE_URL\n",
        "\n",
        "# Install ETCD\n",
        "RUN mkdir -p /tmp/etcd-download-test && \\\n",
        "    curl -L ${DOWNLOAD_URL}/${ETCD_VER}/etcd-${ETCD_VER}-linux-amd64.tar.gz -o /tmp/etcd-${ETCD_VER}-linux-amd64.tar.gz && \\\n",
        "    tar xzvf /tmp/etcd-${ETCD_VER}-linux-amd64.tar.gz -C /tmp/etcd-download-test --strip-components=1 && \\\n",
        "    rm -f /tmp/etcd-${ETCD_VER}-linux-amd64.tar.gz\n",
        "\n",
        "# Copy training application code\n",
        "COPY . /trainer\n",
        "\n",
        "WORKDIR /trainer\n",
        "\n",
        "# Install dependencies\n",
        "RUN pip install -r requirements.txt\n",
        "\n",
        "RUN chmod 777 main.sh\n",
        "\n",
        "# Download data to the container\n",
        "RUN wget -q -P /trainer/data https://image-net.org/data/tiny-imagenet-200.zip\n",
        "RUN unzip -q /trainer/data/tiny-imagenet-200.zip\n",
        "RUN rm /trainer/data/tiny-imagenet-200.zip\n",
        "\n",
        "CMD [\"/bin/bash\", \"main.sh\"]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "cbcfa70f5d54"
      },
      "outputs": [],
      "source": [
        "%%writefile trainer/requirements.txt\n",
        "torch==1.13.0\n",
        "torchvision==0.14.0\n",
        "tensorboard==2.5.0\n",
        "protobuf==3.20.*\n",
        "python-etcd\n",
        "python-json-logger"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "127c2963c3b8"
      },
      "source": [
        "#### Create the main.sh file \n",
        "Starts the ETCD server on the host, saves the host IP to Cloud Storage (for the workers), and calls torchrun"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "77dba1b6bf95"
      },
      "outputs": [],
      "source": [
        "%%writefile trainer/main.sh\n",
        "#!/bin/bash\n",
        "# Copyright 2022 Google Inc. All Rights Reserved.\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",
        "#            http://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.\n",
        "\n",
        "# Provision the prerequisites for running a job on TPU VMs in GKE \n",
        "# using a Vertex AI Pipeline \n",
        "# USAGE:  ./install.sh PROJECT_ID GKE_CLUSTER NAME_PREFIX [ZONE=us-central1-b]\n",
        "# ./install.sh your-project-id gke-tpu-cluster gke-tpu us-central1-b\n",
        "# Set up a global error handler\n",
        "err_handler() {\n",
        "    echo \"Error on line: $1\"\n",
        "    echo \"Caused by: $2\"\n",
        "    echo \"That returned exit status: $3\"\n",
        "    echo \"Aborting...\"\n",
        "    exit $3\n",
        "}\n",
        "\n",
        "trap 'err_handler \"$LINENO\" \"$BASH_COMMAND\" \"$?\"' ERR\n",
        "\n",
        "setup_etcd() {\n",
        "    HOST_IP=$1\n",
        "    # Start a local instane of ETCD v2 \n",
        "    export ETCD_ENABLE_V2=true\n",
        "    export ETCDCTL_API=2\n",
        "\n",
        "    /tmp/etcd-download-test/etcd --name s1 --data-dir /tmp/etcd-download-test/s1  \\\n",
        "    --listen-client-urls http://0.0.0.0:2379 --advertise-client-urls http://$HOST_IP:2379 \\\n",
        "    --listen-peer-urls http://0.0.0.0:2380 --initial-advertise-peer-urls http://$HOST_IP:2380 \\\n",
        "    --initial-cluster s1=http://$HOST_IP:2380 --initial-cluster-token tkn \\\n",
        "    --initial-cluster-state new &> /tmp/etcd-download-test/node.log &\n",
        "\n",
        "    /tmp/etcd-download-test/etcd --version\n",
        "    /tmp/etcd-download-test/etcdctl --version\n",
        "}\n",
        "\n",
        "\n",
        "# Process and print passed in variables\n",
        "while getopts e:a:b:d:t:w:v:u:i:p:n:r:c: option\n",
        "do \n",
        "    case \"${option}\"\n",
        "        in\n",
        "        e)epochs=${OPTARG};;\n",
        "        a)arch=${OPTARG};;\n",
        "        b)batchsize=${OPTARG};;\n",
        "        d)distbackend=${OPTARG};;\n",
        "        t)data=${OPTARG};;\n",
        "        w)workers=${OPTARG};;\n",
        "        v)env=${OPTARG};;\n",
        "        u)rdvzbackend=${OPTARG};;\n",
        "        i)rdvzid=${OPTARG};;\n",
        "        p)endpoint=${OPTARG};;\n",
        "        n)nnodes=${OPTARG};;\n",
        "        r)nprocpernode=${OPTARG};;\n",
        "        c)ischief=${OPTARG};;\n",
        "    esac\n",
        "done\n",
        "\n",
        "echo \"epochs : $epochs\"\n",
        "echo \"arch : $arch\"\n",
        "echo \"batchsize : $batchsize\"\n",
        "echo \"distbackend : $distbackend\"\n",
        "echo \"data : $data\"\n",
        "echo \"workers : $workers\"\n",
        "echo \"env : $env\"\n",
        "echo \"rdvzbackend : $rdvzbackend\"\n",
        "echo \"rdvzid : $rdvzid\"\n",
        "echo \"endpoint : $endpoint\"\n",
        "echo \"nnodes : $nnodes\"\n",
        "echo \"nprocpernode : $nprocpernode\"\n",
        "echo \"ischief : $ischief\"\n",
        "\n",
        "# parse cluster config\n",
        "IFS=' ' read -a conf <<< $(python parse_cluster_config.py)\n",
        "WORKERPOOL_TYPE=\"${conf[0]}\"\n",
        "\n",
        "echo \"WORKERPOOL_TYPE=${WORKERPOOL_TYPE}\"\n",
        "echo \"CLUSTER_SPEC=${CLUSTER_SPEC}\"\n",
        "\n",
        "gcsfilepath=\"${env//\\/gcs\\//gs://}\"\n",
        "\n",
        "if [ \"$WORKERPOOL_TYPE\" == \"workerpool0\" ] || [ \"$WORKERPOOL_TYPE\" == \"chief\" ]; then\n",
        "    HOST_IP=$(hostname -i)\n",
        "    echo \"HOST_IP=\"$HOST_IP\n",
        "    echo \"Writing host IP address to \"$gcsfilepath\n",
        "    echo $HOST_IP| gsutil cp - $gcsfilepath\n",
        "    setup_etcd $HOST_IP\n",
        "else\n",
        "    echo \"Wait 60s for the host server to come online\"\n",
        "    sleep 60\n",
        "    echo \"reading host IP address from \"$gcsfilepath\n",
        "    HOST_IP=$(gsutil cat $gcsfilepath)\n",
        "    echo \"HOST_IP=\"$HOST_IP\n",
        "fi\n",
        "\n",
        "env=\"env://\"\n",
        "ping -c 1 $HOST_IP\n",
        "\n",
        "set -x\n",
        "\n",
        "torchrun --rdzv_backend $rdvzbackend --rdzv_id $rdvzid --rdzv_endpoint $HOST_IP:2379 \\\n",
        "--nnodes $nnodes --nproc_per_node $nprocpernode --master_addr $HOST_IP --master_port 2379 \\\n",
        "main.py --epochs $epochs --arch $arch --batch-size $batchsize --dist-backend $distbackend \\\n",
        "--data $data \\\n",
        "--env $env \\\n",
        "--hostip $HOST_IP \\\n",
        "--hostipport 2379 \\\n",
        "--workers $workers \\\n",
        "--ischief $ischief"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "3294e388117a"
      },
      "outputs": [],
      "source": [
        "%%writefile trainer/parse_cluster_config.py\n",
        "import os\n",
        "import json\n",
        "\n",
        "cluster_config_str = os.environ.get('CLUSTER_SPEC')\n",
        "cluster_config_dict  = json.loads(cluster_config_str)\n",
        "workerpool_type = cluster_config_dict['task']['type']\n",
        "\n",
        "print(workerpool_type)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "10fdcf6db6bb"
      },
      "source": [
        "#### Create the main.py file \n",
        "Main trainer for the ImageNet training job"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "24793c94ebff"
      },
      "outputs": [],
      "source": [
        "%%writefile trainer/main.py\n",
        "#!/usr/bin/env python3\n",
        "\n",
        "# Copyright (c) Facebook, Inc. and its affiliates.\n",
        "# All rights reserved.\n",
        "#\n",
        "# This source code is licensed under the BSD-style license found in the\n",
        "# LICENSE file in the root directory of this source tree.\n",
        "\n",
        "r\"\"\"\n",
        "Source: `pytorch imagenet example <https://github.com/pytorch/examples/blob/master/imagenet/main.py>`_ # noqa B950\n",
        "Modified and simplified to make the original pytorch example compatible with\n",
        "torchelastic.distributed.launch.\n",
        "Changes:\n",
        "1. Removed ``rank``, ``gpu``, ``multiprocessing-distributed``, ``dist_url`` options.\n",
        "   These are obsolete parameters when using ``torchelastic.distributed.launch``.\n",
        "2. Removed ``seed``, ``evaluate``, ``pretrained`` options for simplicity.\n",
        "3. Removed ``resume``, ``start-epoch`` options.\n",
        "   Loads the most recent checkpoint by default.\n",
        "4. ``batch-size`` is now per GPU (worker) batch size rather than for all GPUs.\n",
        "5. Defaults ``workers`` (num data loader workers) to ``0``.\n",
        "Usage\n",
        "::\n",
        " >>> python -m torchelastic.distributed.launch\n",
        "        --nnodes=$NUM_NODES\n",
        "        --nproc_per_node=$WORKERS_PER_NODE\n",
        "        --rdzv_id=$JOB_ID\n",
        "        --rdzv_backend=etcd\n",
        "        --rdzv_endpoint=$ETCD_HOST:$ETCD_PORT\n",
        "        main.py\n",
        "        --arch resnet18\n",
        "        --epochs 20\n",
        "        --batch-size 32\n",
        "        <DATA_DIR>\n",
        "\"\"\"\n",
        "\n",
        "import traceback\n",
        "import argparse\n",
        "import io\n",
        "import os\n",
        "import shutil\n",
        "import time\n",
        "from contextlib import contextmanager\n",
        "from datetime import timedelta\n",
        "from typing import List, Tuple\n",
        "\n",
        "import numpy\n",
        "import torch\n",
        "import torch.backends.cudnn as cudnn\n",
        "import torch.distributed as dist\n",
        "import torch.nn as nn\n",
        "import torch.nn.parallel\n",
        "import torch.optim\n",
        "import torch.utils.data\n",
        "import torch.utils.data.distributed\n",
        "import torchvision.datasets as datasets\n",
        "import torchvision.models as models\n",
        "import torchvision.transforms as transforms\n",
        "from torch.distributed.elastic.utils.data import ElasticDistributedSampler\n",
        "from torch.nn.parallel import DistributedDataParallel\n",
        "from torch.optim import SGD\n",
        "from torch.utils.data import DataLoader\n",
        "\n",
        "\n",
        "model_names = sorted(\n",
        "    name\n",
        "    for name in models.__dict__\n",
        "    if name.islower() and not name.startswith(\"__\") and callable(models.__dict__[name])\n",
        ")\n",
        "\n",
        "parser = argparse.ArgumentParser(description=\"PyTorch Elastic ImageNet Training\")\n",
        "parser.add_argument(\"--data\", metavar=\"DIR\", help=\"path to dataset\")\n",
        "parser.add_argument(\n",
        "    \"-a\",\n",
        "    \"--arch\",\n",
        "    metavar=\"ARCH\",\n",
        "    default=\"resnet18\",\n",
        "    choices=model_names,\n",
        "    help=\"model architecture: \" + \" | \".join(model_names) + \" (default: resnet18)\",\n",
        ")\n",
        "parser.add_argument(\n",
        "    \"-j\",\n",
        "    \"--workers\",\n",
        "    default=0,\n",
        "    type=int,\n",
        "    metavar=\"N\",\n",
        "    help=\"number of data loading workers\",\n",
        ")\n",
        "parser.add_argument(\n",
        "    \"--epochs\", default=90, type=int, metavar=\"N\", help=\"number of total epochs to run\"\n",
        ")\n",
        "parser.add_argument(\n",
        "    \"-b\",\n",
        "    \"--batch-size\",\n",
        "    default=32,\n",
        "    type=int,\n",
        "    metavar=\"N\",\n",
        "    help=\"mini-batch size (default: 32), per worker (GPU)\",\n",
        ")\n",
        "parser.add_argument(\n",
        "    \"--lr\",\n",
        "    \"--learning-rate\",\n",
        "    default=0.1,\n",
        "    type=float,\n",
        "    metavar=\"LR\",\n",
        "    help=\"initial learning rate\",\n",
        "    dest=\"lr\",\n",
        ")\n",
        "parser.add_argument(\"--momentum\", default=0.9, type=float, metavar=\"M\", help=\"momentum\")\n",
        "parser.add_argument(\n",
        "    \"--wd\",\n",
        "    \"--weight-decay\",\n",
        "    default=1e-4,\n",
        "    type=float,\n",
        "    metavar=\"W\",\n",
        "    help=\"weight decay (default: 1e-4)\",\n",
        "    dest=\"weight_decay\",\n",
        ")\n",
        "parser.add_argument(\n",
        "    \"-p\",\n",
        "    \"--print-freq\",\n",
        "    default=10,\n",
        "    type=int,\n",
        "    metavar=\"N\",\n",
        "    help=\"print frequency (default: 10)\",\n",
        ")\n",
        "parser.add_argument(\n",
        "    \"--dist-backend\",\n",
        "    default=\"nccl\",\n",
        "    choices=[\"nccl\", \"gloo\"],\n",
        "    type=str,\n",
        "    help=\"distributed backend\",\n",
        ")\n",
        "parser.add_argument(\n",
        "    \"--checkpoint-file\",\n",
        "    default=\"/tmp/checkpoint.pth.tar\",\n",
        "    type=str,\n",
        "    help=\"checkpoint file path, to load and save to\",\n",
        ")\n",
        "parser.add_argument(\n",
        "    \"--env\",\n",
        "    default=\"env://\",\n",
        "    type=str,\n",
        "    help=\"setting for init_method for torch.distributed.init_process_group. Leave default unless you want to pass a shared gcs path\",\n",
        ")\n",
        "parser.add_argument(\n",
        "    \"--hostip\",\n",
        "    default=\"localhost\",\n",
        "    type=str,\n",
        "    help=\"setting for etcd host ip\",\n",
        ")\n",
        "parser.add_argument(\n",
        "    \"--hostipport\",\n",
        "    default=2379,\n",
        "    type=int,\n",
        "    help=\"setting for etcd host ip port\",\n",
        ")\n",
        "parser.add_argument(\n",
        "    \"--ischief\",\n",
        "    default=\"n\", \n",
        "    type=str,\n",
        "    help='is this cheif or worker')\n",
        "\n",
        "def main():\n",
        "    args = parser.parse_args()\n",
        "    print(args)\n",
        "    device_id = int(os.environ[\"LOCAL_RANK\"])\n",
        "    torch.cuda.set_device(device_id)\n",
        "    print(f\"=> set cuda device = {device_id}\")\n",
        "\n",
        "    LOCAL_RANK=int(os.environ[\"LOCAL_RANK\"])\n",
        "    RANK=int(os.environ[\"RANK\"])\n",
        "    WORLD_SIZE=int(os.environ[\"WORLD_SIZE\"])\n",
        "\n",
        "    print (f\"LOCAL_RANK={os.environ['LOCAL_RANK']} RANK={os.environ['RANK']} WORLD_SIZE={os.environ['WORLD_SIZE']}\")\n",
        "    print (f\"args env= {args.env}\")\n",
        "    \n",
        "    print (f\"Host address: {os.environ['MASTER_ADDR']}:{os.environ['MASTER_PORT']}\")\n",
        "    os.environ['MASTER_ADDR']=args.hostip\n",
        "    print (f\"Updated IPv4 host address: {os.environ['MASTER_ADDR']}:{os.environ['MASTER_PORT']}\")\n",
        "    \n",
        "    print ('Initialize process group')\n",
        "    if args.env == \"env://\":\n",
        "        dist.init_process_group(\n",
        "            backend=args.dist_backend, init_method=f\"{args.env}\", timeout=timedelta(seconds=120)\n",
        "        )\n",
        "    else:\n",
        "        if args.ischief.lower() == 'y':\n",
        "            print ('Setting store')\n",
        "            #STORE = dist.FileStore(args.env, WORLD_SIZE)\n",
        "            STORE = dist.TCPStore(host_name=args.hostip, port=args.hostipport, world_size=WORLD_SIZE, is_master=True, timeout=timedelta(seconds=30))\n",
        "            print (f'Store set = {STORE}')        \n",
        "            dist.init_process_group(\n",
        "                backend=args.dist_backend, store=STORE, timeout=timedelta(seconds=30),\n",
        "                rank=RANK, world_size=WORLD_SIZE\n",
        "            )\n",
        "\n",
        "        dist.init_process_group(\n",
        "            backend=args.dist_backend, init_method=f\"tcp://{args.hostip}:{args.hostipport}\", timeout=timedelta(seconds=120), rank=RANK, world_size=WORLD_SIZE\n",
        "        )\n",
        "    print ('Process initialized')\n",
        "\n",
        "    model, criterion, optimizer = initialize_model(\n",
        "        args.arch, args.lr, args.momentum, args.weight_decay, device_id\n",
        "    )\n",
        "\n",
        "    train_loader, val_loader = initialize_data_loader(\n",
        "        args.data, args.batch_size, args.workers\n",
        "    )\n",
        "\n",
        "    # resume from checkpoint if one exists;\n",
        "    state = load_checkpoint(\n",
        "        args.checkpoint_file, device_id, args.arch, model, optimizer\n",
        "    )\n",
        "\n",
        "    start_epoch = state.epoch + 1\n",
        "    print(f\"=> start_epoch: {start_epoch}, best_acc1: {state.best_acc1}\")\n",
        "\n",
        "    print_freq = args.print_freq\n",
        "    for epoch in range(start_epoch, args.epochs):\n",
        "        state.epoch = epoch\n",
        "        train_loader.batch_sampler.sampler.set_epoch(epoch)\n",
        "        adjust_learning_rate(optimizer, epoch, args.lr)\n",
        "\n",
        "        # train for one epoch\n",
        "        train(train_loader, model, criterion, optimizer, epoch, device_id, print_freq)\n",
        "\n",
        "        # evaluate on validation set\n",
        "        acc1 = validate(val_loader, model, criterion, device_id, print_freq)\n",
        "\n",
        "        # remember best acc@1 and save checkpoint\n",
        "        is_best = acc1 > state.best_acc1\n",
        "        state.best_acc1 = max(acc1, state.best_acc1)\n",
        "\n",
        "        if device_id == 0:\n",
        "            save_checkpoint(state, is_best, args.checkpoint_file)\n",
        "\n",
        "\n",
        "class State:\n",
        "    \"\"\"\n",
        "    Container for objects that we want to checkpoint. Represents the\n",
        "    current \"state\" of the worker. This object is mutable.\n",
        "    \"\"\"\n",
        "\n",
        "    def __init__(self, arch, model, optimizer):\n",
        "        self.epoch = -1\n",
        "        self.best_acc1 = 0\n",
        "        self.arch = arch\n",
        "        self.model = model\n",
        "        self.optimizer = optimizer\n",
        "\n",
        "    def capture_snapshot(self):\n",
        "        \"\"\"\n",
        "        Essentially a ``serialize()`` function, returns the state as an\n",
        "        object compatible with ``torch.save()``. The following should work\n",
        "        ::\n",
        "        snapshot = state_0.capture_snapshot()\n",
        "        state_1.apply_snapshot(snapshot)\n",
        "        assert state_0 == state_1\n",
        "        \"\"\"\n",
        "        return {\n",
        "            \"epoch\": self.epoch,\n",
        "            \"best_acc1\": self.best_acc1,\n",
        "            \"arch\": self.arch,\n",
        "            \"state_dict\": self.model.state_dict(),\n",
        "            \"optimizer\": self.optimizer.state_dict(),\n",
        "        }\n",
        "\n",
        "    def apply_snapshot(self, obj, device_id):\n",
        "        \"\"\"\n",
        "        The complimentary function of ``capture_snapshot()``. Applies the\n",
        "        snapshot object that was returned by ``capture_snapshot()``.\n",
        "        This function mutates this state object.\n",
        "        \"\"\"\n",
        "\n",
        "        self.epoch = obj[\"epoch\"]\n",
        "        self.best_acc1 = obj[\"best_acc1\"]\n",
        "        self.state_dict = obj[\"state_dict\"]\n",
        "        self.model.load_state_dict(obj[\"state_dict\"])\n",
        "        self.optimizer.load_state_dict(obj[\"optimizer\"])\n",
        "\n",
        "    def save(self, f):\n",
        "        torch.save(self.capture_snapshot(), f)\n",
        "\n",
        "    def load(self, f, device_id):\n",
        "        # Map model to be loaded to specified single gpu.\n",
        "        snapshot = torch.load(f, map_location=f\"cuda:{device_id}\")\n",
        "        self.apply_snapshot(snapshot, device_id)\n",
        "\n",
        "\n",
        "def initialize_model(\n",
        "    arch: str, lr: float, momentum: float, weight_decay: float, device_id: int\n",
        "):\n",
        "    print(f\"=> creating model: {arch}\")\n",
        "    model = models.__dict__[arch]()\n",
        "    # For multiprocessing distributed, DistributedDataParallel constructor\n",
        "    # should always set the single device scope, otherwise,\n",
        "    # DistributedDataParallel will use all available devices.\n",
        "    model.cuda(device_id)\n",
        "    cudnn.benchmark = True\n",
        "    model = DistributedDataParallel(model, device_ids=[device_id])\n",
        "    # define loss function (criterion) and optimizer\n",
        "    criterion = nn.CrossEntropyLoss().cuda(device_id)\n",
        "    optimizer = SGD(\n",
        "        model.parameters(), lr, momentum=momentum, weight_decay=weight_decay\n",
        "    )\n",
        "    return model, criterion, optimizer\n",
        "\n",
        "\n",
        "def initialize_data_loader(\n",
        "    data_dir, batch_size, num_data_workers\n",
        ") -> Tuple[DataLoader, DataLoader]:\n",
        "    traindir = os.path.join(data_dir, \"train\")\n",
        "    valdir = os.path.join(data_dir, \"val\")\n",
        "    normalize = transforms.Normalize(\n",
        "        mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]\n",
        "    )\n",
        "    train_dataset = datasets.ImageFolder(\n",
        "        traindir,\n",
        "        transforms.Compose(\n",
        "            [\n",
        "                transforms.RandomResizedCrop(224),\n",
        "                transforms.RandomHorizontalFlip(),\n",
        "                transforms.ToTensor(),\n",
        "                normalize,\n",
        "            ]\n",
        "        ),\n",
        "    )\n",
        "    train_sampler = ElasticDistributedSampler(train_dataset)\n",
        "    train_loader = DataLoader(\n",
        "        train_dataset,\n",
        "        batch_size=batch_size,\n",
        "        num_workers=num_data_workers,\n",
        "        pin_memory=True,\n",
        "        sampler=train_sampler,\n",
        "    )\n",
        "    val_loader = DataLoader(\n",
        "        datasets.ImageFolder(\n",
        "            valdir,\n",
        "            transforms.Compose(\n",
        "                [\n",
        "                    transforms.Resize(256),\n",
        "                    transforms.CenterCrop(224),\n",
        "                    transforms.ToTensor(),\n",
        "                    normalize,\n",
        "                ]\n",
        "            ),\n",
        "        ),\n",
        "        batch_size=batch_size,\n",
        "        shuffle=False,\n",
        "        num_workers=num_data_workers,\n",
        "        pin_memory=True,\n",
        "    )\n",
        "    return train_loader, val_loader\n",
        "\n",
        "\n",
        "def load_checkpoint(\n",
        "    checkpoint_file: str,\n",
        "    device_id: int,\n",
        "    arch: str,\n",
        "    model: DistributedDataParallel,\n",
        "    optimizer,  # SGD\n",
        ") -> State:\n",
        "    \"\"\"\n",
        "    Loads a local checkpoint (if any). Otherwise, checks to see if any of\n",
        "    the neighbors have a non-zero state. If so, restore the state\n",
        "    from the rank that has the most up-to-date checkpoint.\n",
        "    .. note:: when your job has access to a globally visible persistent storage\n",
        "              (e.g. nfs mount, S3) you can simply have all workers load\n",
        "              from the most recent checkpoint from such storage. Since this\n",
        "              example is expected to run on vanilla hosts (with no shared\n",
        "              storage) the checkpoints are written to local disk, hence\n",
        "              we have the extra logic to broadcast the checkpoint from a\n",
        "              surviving node.\n",
        "    \"\"\"\n",
        "\n",
        "    state = State(arch, model, optimizer)\n",
        "\n",
        "    if os.path.isfile(checkpoint_file):\n",
        "        print(f\"=> loading checkpoint file: {checkpoint_file}\")\n",
        "        state.load(checkpoint_file, device_id)\n",
        "        print(f\"=> loaded checkpoint file: {checkpoint_file}\")\n",
        "\n",
        "    # logic below is unnecessary when the checkpoint is visible on all nodes!\n",
        "    # create a temporary cpu pg to broadcast most up-to-date checkpoint\n",
        "    with tmp_process_group(backend=\"gloo\") as pg:\n",
        "        rank = dist.get_rank(group=pg)\n",
        "\n",
        "        # get rank that has the largest state.epoch\n",
        "        epochs = torch.zeros(dist.get_world_size(), dtype=torch.int32)\n",
        "        epochs[rank] = state.epoch\n",
        "        dist.all_reduce(epochs, op=dist.ReduceOp.SUM, group=pg)\n",
        "        t_max_epoch, t_max_rank = torch.max(epochs, dim=0)\n",
        "        max_epoch = t_max_epoch.item()\n",
        "        max_rank = t_max_rank.item()\n",
        "\n",
        "        # max_epoch == -1 means no one has checkpointed return base state\n",
        "        if max_epoch == -1:\n",
        "            print(f\"=> no workers have checkpoints, starting from epoch 0\")\n",
        "            return state\n",
        "\n",
        "        # broadcast the state from max_rank (which has the most up-to-date state)\n",
        "        # pickle the snapshot, convert it into a byte-blob tensor\n",
        "        # then broadcast it, unpickle it and apply the snapshot\n",
        "        print(f\"=> using checkpoint from rank: {max_rank}, max_epoch: {max_epoch}\")\n",
        "\n",
        "        with io.BytesIO() as f:\n",
        "            torch.save(state.capture_snapshot(), f)\n",
        "            raw_blob = numpy.frombuffer(f.getvalue(), dtype=numpy.uint8)\n",
        "\n",
        "        blob_len = torch.tensor(len(raw_blob))\n",
        "        dist.broadcast(blob_len, src=max_rank, group=pg)\n",
        "        print(f\"=> checkpoint broadcast size is: {blob_len}\")\n",
        "\n",
        "        if rank != max_rank:\n",
        "            # pyre-fixme[6]: For 1st param expected `Union[List[int], Size,\n",
        "            #  typing.Tuple[int, ...]]` but got `Union[bool, float, int]`.\n",
        "            blob = torch.zeros(blob_len.item(), dtype=torch.uint8)\n",
        "        else:\n",
        "            blob = torch.as_tensor(raw_blob, dtype=torch.uint8)\n",
        "\n",
        "        dist.broadcast(blob, src=max_rank, group=pg)\n",
        "        print(f\"=> done broadcasting checkpoint\")\n",
        "\n",
        "        if rank != max_rank:\n",
        "            with io.BytesIO(blob.numpy()) as f:\n",
        "                snapshot = torch.load(f)\n",
        "            state.apply_snapshot(snapshot, device_id)\n",
        "\n",
        "        # wait till everyone has loaded the checkpoint\n",
        "        dist.barrier(group=pg)\n",
        "\n",
        "    print(f\"=> done restoring from previous checkpoint\")\n",
        "    return state\n",
        "\n",
        "\n",
        "@contextmanager\n",
        "def tmp_process_group(backend):\n",
        "    cpu_pg = dist.new_group(backend=backend)\n",
        "    try:\n",
        "        yield cpu_pg\n",
        "    finally:\n",
        "        dist.destroy_process_group(cpu_pg)\n",
        "\n",
        "\n",
        "def save_checkpoint(state: State, is_best: bool, filename: str):\n",
        "    checkpoint_dir = os.path.dirname(filename)\n",
        "    os.makedirs(checkpoint_dir, exist_ok=True)\n",
        "\n",
        "    # save to tmp, then commit by moving the file in case the job\n",
        "    # gets interrupted while writing the checkpoint\n",
        "    tmp_filename = filename + \".tmp\"\n",
        "    torch.save(state.capture_snapshot(), tmp_filename)\n",
        "    os.rename(tmp_filename, filename)\n",
        "    print(f\"=> saved checkpoint for epoch {state.epoch} at {filename}\")\n",
        "    if is_best:\n",
        "        best = os.path.join(checkpoint_dir, \"model_best.pth.tar\")\n",
        "        print(f\"=> best model found at epoch {state.epoch} saving to {best}\")\n",
        "        shutil.copyfile(filename, best)\n",
        "\n",
        "\n",
        "def train(\n",
        "    train_loader: DataLoader,\n",
        "    model: DistributedDataParallel,\n",
        "    criterion,  # nn.CrossEntropyLoss\n",
        "    optimizer,  # SGD,\n",
        "    epoch: int,\n",
        "    device_id: int,\n",
        "    print_freq: int,\n",
        "):\n",
        "    batch_time = AverageMeter(\"Time\", \":6.3f\")\n",
        "    data_time = AverageMeter(\"Data\", \":6.3f\")\n",
        "    losses = AverageMeter(\"Loss\", \":.4e\")\n",
        "    top1 = AverageMeter(\"Acc@1\", \":6.2f\")\n",
        "    top5 = AverageMeter(\"Acc@5\", \":6.2f\")\n",
        "    progress = ProgressMeter(\n",
        "        len(train_loader),\n",
        "        [batch_time, data_time, losses, top1, top5],\n",
        "        prefix=\"Epoch: [{}]\".format(epoch+1),\n",
        "    )\n",
        "\n",
        "    # switch to train mode\n",
        "    model.train()\n",
        "\n",
        "    end = time.time()\n",
        "    for i, (images, target) in enumerate(train_loader):\n",
        "        # measure data loading time\n",
        "        data_time.update(time.time() - end)\n",
        "\n",
        "        images = images.cuda(device_id, non_blocking=True)\n",
        "        target = target.cuda(device_id, non_blocking=True)\n",
        "\n",
        "        # compute output\n",
        "        output = model(images)\n",
        "        loss = criterion(output, target)\n",
        "\n",
        "        # measure accuracy and record loss\n",
        "        acc1, acc5 = accuracy(output, target, topk=(1, 5))\n",
        "        losses.update(loss.item(), images.size(0))\n",
        "        top1.update(acc1[0], images.size(0))\n",
        "        top5.update(acc5[0], images.size(0))\n",
        "\n",
        "        # compute gradient and do SGD step\n",
        "        optimizer.zero_grad()\n",
        "        loss.backward()\n",
        "        optimizer.step()\n",
        "\n",
        "        # measure elapsed time\n",
        "        batch_time.update(time.time() - end)\n",
        "        end = time.time()\n",
        "\n",
        "        if i % print_freq == 0:\n",
        "            progress.display(i)\n",
        "\n",
        "\n",
        "def validate(\n",
        "    val_loader: DataLoader,\n",
        "    model: DistributedDataParallel,\n",
        "    criterion,  # nn.CrossEntropyLoss\n",
        "    device_id: int,\n",
        "    print_freq: int,\n",
        "):\n",
        "    batch_time = AverageMeter(\"Time\", \":6.3f\")\n",
        "    losses = AverageMeter(\"Loss\", \":.4e\")\n",
        "    top1 = AverageMeter(\"Acc@1\", \":6.2f\")\n",
        "    top5 = AverageMeter(\"Acc@5\", \":6.2f\")\n",
        "    progress = ProgressMeter(\n",
        "        len(val_loader), [batch_time, losses, top1, top5], prefix=\"Test: \"\n",
        "    )\n",
        "\n",
        "    # switch to evaluate mode\n",
        "    model.eval()\n",
        "\n",
        "    with torch.no_grad():\n",
        "        end = time.time()\n",
        "        for i, (images, target) in enumerate(val_loader):\n",
        "            if device_id is not None:\n",
        "                images = images.cuda(device_id, non_blocking=True)\n",
        "            target = target.cuda(device_id, non_blocking=True)\n",
        "\n",
        "            # compute output\n",
        "            output = model(images)\n",
        "            loss = criterion(output, target)\n",
        "\n",
        "            # measure accuracy and record loss\n",
        "            acc1, acc5 = accuracy(output, target, topk=(1, 5))\n",
        "            losses.update(loss.item(), images.size(0))\n",
        "            top1.update(acc1[0], images.size(0))\n",
        "            top5.update(acc5[0], images.size(0))\n",
        "\n",
        "            # measure elapsed time\n",
        "            batch_time.update(time.time() - end)\n",
        "            end = time.time()\n",
        "\n",
        "            if i % print_freq == 0:\n",
        "                progress.display(i)\n",
        "\n",
        "        # TODO: this should also be done with the ProgressMeter\n",
        "        print(\n",
        "            \" * Acc@1 {top1.avg:.3f} Acc@5 {top5.avg:.3f}\".format(top1=top1, top5=top5)\n",
        "        )\n",
        "\n",
        "    return top1.avg\n",
        "\n",
        "\n",
        "class AverageMeter(object):\n",
        "    \"\"\"Computes and stores the average and current value\"\"\"\n",
        "\n",
        "    def __init__(self, name: str, fmt: str = \":f\"):\n",
        "        self.name = name\n",
        "        self.fmt = fmt\n",
        "        self.reset()\n",
        "\n",
        "    def reset(self) -> None:\n",
        "        self.val = 0\n",
        "        self.avg = 0\n",
        "        self.sum = 0\n",
        "        self.count = 0\n",
        "\n",
        "    def update(self, val, n=1) -> None:\n",
        "        self.val = val\n",
        "        self.sum += val * n\n",
        "        self.count += n\n",
        "        self.avg = self.sum / self.count\n",
        "\n",
        "    def __str__(self):\n",
        "        fmtstr = \"{name} {val\" + self.fmt + \"} ({avg\" + self.fmt + \"})\"\n",
        "        return fmtstr.format(**self.__dict__)\n",
        "\n",
        "\n",
        "class ProgressMeter(object):\n",
        "    def __init__(self, num_batches: int, meters: List[AverageMeter], prefix: str = \"\"):\n",
        "        self.batch_fmtstr = self._get_batch_fmtstr(num_batches)\n",
        "        self.meters = meters\n",
        "        self.prefix = prefix\n",
        "\n",
        "    def display(self, batch: int) -> None:\n",
        "        entries = [self.prefix + self.batch_fmtstr.format(batch)]\n",
        "        entries += [str(meter) for meter in self.meters]\n",
        "        print(\"\\t\".join(entries))\n",
        "\n",
        "    def _get_batch_fmtstr(self, num_batches: int) -> str:\n",
        "        num_digits = len(str(num_batches // 1))\n",
        "        fmt = \"{:\" + str(num_digits) + \"d}\"\n",
        "        return \"[\" + fmt + \"/\" + fmt.format(num_batches) + \"]\"\n",
        "\n",
        "\n",
        "def adjust_learning_rate(optimizer, epoch: int, lr: float) -> None:\n",
        "    \"\"\"\n",
        "    Sets the learning rate to the initial LR decayed by 10 every 30 epochs\n",
        "    \"\"\"\n",
        "    learning_rate = lr * (0.1 ** (epoch // 30))\n",
        "    for param_group in optimizer.param_groups:\n",
        "        param_group[\"lr\"] = learning_rate\n",
        "\n",
        "\n",
        "def accuracy(output, target, topk=(1,)):\n",
        "    \"\"\"\n",
        "    Computes the accuracy over the k top predictions for the specified values of k\n",
        "    \"\"\"\n",
        "    with torch.no_grad():\n",
        "        maxk = max(topk)\n",
        "        batch_size = target.size(0)\n",
        "\n",
        "        _, pred = output.topk(maxk, 1, True, True)\n",
        "        pred = pred.t()\n",
        "        correct = pred.eq(target.view(1, -1).expand_as(pred))\n",
        "\n",
        "        res = []\n",
        "        for k in topk:\n",
        "            correct_k = correct[:k].reshape(1, -1).view(-1).float().sum(0, keepdim=True)\n",
        "            res.append(correct_k.mul_(100.0 / batch_size))\n",
        "        return res\n",
        "\n",
        "\n",
        "if __name__ == \"__main__\":\n",
        "    try:\n",
        "        main()\n",
        "    except Exception as e:\n",
        "        trace_str = ''.join(traceback.format_tb(e.__traceback__))\n",
        "        print(trace_str)\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "93002a20a2a6"
      },
      "source": [
        "### Build custom container"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "8ea6fc98b2a9"
      },
      "outputs": [],
      "source": [
        "CONTENT_NAME = \"pytorch-torchrun-imagenet-multi-node\"\n",
        "CONTAINER_NAME = CONTENT_NAME + \"-gpu\"\n",
        "TAG = \"latest\"\n",
        "\n",
        "custom_container_host_image_uri = f\"{LOCATION}-docker.pkg.dev/{PROJECT_ID}/{REPOSITORY}/{CONTAINER_NAME}:{TAG}\"  # noqa"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ee1a0a06d0b4"
      },
      "outputs": [],
      "source": [
        "!gcloud builds submit \\\n",
        "   --region $LOCATION \\\n",
        "   --tag $custom_container_host_image_uri \\\n",
        "   --timeout \"2h\" \\\n",
        "   --machine-type=e2-highcpu-32 \\\n",
        "   trainer"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4a87f56e869a"
      },
      "source": [
        "### Run training on Vertex AI using `torchrun` with ETCD on host"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "63ca30b33176"
      },
      "outputs": [],
      "source": [
        "from datetime import datetime\n",
        "\n",
        "BUCKET_NAME = BUCKET_URI.replace(\"gs://\", \"\")\n",
        "TIMESTAMP = datetime.now().strftime(\"%Y%m%d%H%M%S\")\n",
        "PRIMARY_COMPUTE = \"n1-highmem-16\"\n",
        "TRAIN_COMPUTE = \"n1-highmem-16\"\n",
        "NUM_CPUS = 14  # Set to a few less than max CPUs per instance for paralle data loading\n",
        "TRAIN_GPU = \"NVIDIA_TESLA_T4\"\n",
        "TRAIN_NGPU = 1\n",
        "BATCH_SIZE = 512\n",
        "REPLICAS = 2\n",
        "EPOCHS = 5\n",
        "ARCH = \"resnet18\"\n",
        "BACKEND = \"nccl\"  # gloo for CPU only, nccl for GPUs\n",
        "TRAIN_DATA_LOCATION = (\n",
        "    \"/trainer/tiny-imagenet-200\"  # Data location of filed downloaded in Dockerfile\n",
        ")\n",
        "\n",
        "display_name = (\n",
        "    CONTAINER_NAME\n",
        "    + \"-LOCAL-ETCD-\"\n",
        "    + f\"{REPLICAS}workers-{TRAIN_NGPU}{TRAIN_GPU}-{BATCH_SIZE}batch-\"\n",
        "    + TIMESTAMP\n",
        ")\n",
        "gcs_output_uri_prefix = f\"{BUCKET_URI}/{display_name}\"\n",
        "\n",
        "RDZV_BACKEND = \"etcd-v2\"\n",
        "RDZV_BACKEND_STORE = f\"/gcs/{BUCKET_NAME}/sharedfile-{display_name}\"\n",
        "RDZV_ENDPOINT = \"localhost:2379\"\n",
        "\n",
        "# Use letters for each parameter to be processed in the shell script\n",
        "\"\"\"\n",
        "e)epochs=${OPTARG};;\n",
        "a)arch=${OPTARG};;\n",
        "b)batchsize=${OPTARG};;\n",
        "d)distbackend=${OPTARG};;\n",
        "t)data=${OPTARG};;\n",
        "w)workers=${OPTARG};;\n",
        "v)env=${OPTARG};;\n",
        "u)rdvzbackend=${OPTARG};;\n",
        "i)rdvzid=${OPTARG};;\n",
        "p)endpoint=${OPTARG};;\n",
        "n)nnodes=${OPTARG};;\n",
        "r)nprocpernode=${OPTARG};;\n",
        "c)ischief=${OPTARG};;\n",
        "\"\"\"\n",
        "\n",
        "CONTAINER_SPEC = {\n",
        "    \"image_uri\": custom_container_host_image_uri,\n",
        "    \"command\": [\n",
        "        \"/bin/bash\",\n",
        "        \"main.sh\",\n",
        "        f\"-e {EPOCHS}\",\n",
        "        f\"-a {ARCH}\",\n",
        "        f\"-b {BATCH_SIZE}\",\n",
        "        f\"-d {BACKEND}\",\n",
        "        f\"-t {TRAIN_DATA_LOCATION}\",\n",
        "        f\"-w {NUM_CPUS}\",\n",
        "        f\"-v {RDZV_BACKEND_STORE}\",\n",
        "        f\"-u {RDZV_BACKEND}\",\n",
        "        f\"-i {display_name}\",\n",
        "        f\"-p {RDZV_ENDPOINT}\",\n",
        "        f\"-n {REPLICAS+1}\",\n",
        "        f\"-r {TRAIN_NGPU}\",\n",
        "        \"-c y\",\n",
        "    ],\n",
        "}\n",
        "\n",
        "CONTAINER_WORKER_SPEC = {\n",
        "    \"image_uri\": custom_container_host_image_uri,\n",
        "    \"command\": [\n",
        "        \"/bin/bash\",\n",
        "        \"main.sh\",\n",
        "        f\"-e {EPOCHS}\",\n",
        "        f\"-a {ARCH}\",\n",
        "        f\"-b {BATCH_SIZE}\",\n",
        "        f\"-d {BACKEND}\",\n",
        "        f\"-t {TRAIN_DATA_LOCATION}\",\n",
        "        f\"-w {NUM_CPUS}\",\n",
        "        f\"-v {RDZV_BACKEND_STORE}\",\n",
        "        f\"-u {RDZV_BACKEND}\",\n",
        "        f\"-i {display_name}\",\n",
        "        f\"-p {RDZV_ENDPOINT}\",\n",
        "        f\"-n {REPLICAS+1}\",\n",
        "        f\"-r {TRAIN_NGPU}\",\n",
        "        \"-c n\",\n",
        "    ],\n",
        "}\n",
        "\n",
        "PRIMARY_WORKER_POOL = {\n",
        "    \"replica_count\": 1,\n",
        "    \"machine_spec\": {\n",
        "        \"machine_type\": PRIMARY_COMPUTE,\n",
        "        \"accelerator_count\": TRAIN_NGPU,\n",
        "        \"accelerator_type\": TRAIN_GPU,\n",
        "    },\n",
        "    \"container_spec\": CONTAINER_SPEC,\n",
        "}\n",
        "\n",
        "WORKER_POOL_SPECS = [PRIMARY_WORKER_POOL]\n",
        "\n",
        "TRAIN_WORKER_POOL = {\n",
        "    \"replica_count\": REPLICAS,\n",
        "    \"machine_spec\": {\n",
        "        \"machine_type\": TRAIN_COMPUTE,\n",
        "        \"accelerator_count\": TRAIN_NGPU,\n",
        "        \"accelerator_type\": TRAIN_GPU,\n",
        "    },\n",
        "    \"container_spec\": CONTAINER_WORKER_SPEC,\n",
        "}\n",
        "\n",
        "WORKER_POOL_SPECS.append(TRAIN_WORKER_POOL)\n",
        "\n",
        "job = aiplatform.CustomJob(\n",
        "    display_name=display_name,\n",
        "    base_output_dir=gcs_output_uri_prefix,\n",
        "    worker_pool_specs=WORKER_POOL_SPECS,\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "2b4684fde1a6"
      },
      "outputs": [],
      "source": [
        "if not os.getenv(\"IS_TESTING\"):\n",
        "    job.run(\n",
        "        sync=True\n",
        "        # comment out the line below to turn off interactive debug\n",
        "        ,\n",
        "        enable_web_access=True,\n",
        "        service_account=SERVICE_ACCOUNT,\n",
        "    )"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "796de99ae13a"
      },
      "source": [
        "### Run training on Vertex AI using `torchrun` with ETCD on host and reduction server"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6815ecb07ad9"
      },
      "outputs": [],
      "source": [
        "from datetime import datetime\n",
        "\n",
        "BUCKET_NAME = BUCKET_URI.replace(\"gs://\", \"\")\n",
        "TIMESTAMP = datetime.now().strftime(\"%Y%m%d%H%M%S\")\n",
        "PRIMARY_COMPUTE = \"n1-highmem-16\"\n",
        "TRAIN_COMPUTE = \"n1-highmem-16\"\n",
        "REDUCTION_COMPUTE = \"n1-highcpu-16\"\n",
        "NUM_CPUS = 14  # Set to a few less than max CPUs per instance for paralle data loading\n",
        "TRAIN_GPU = \"NVIDIA_TESLA_T4\"\n",
        "TRAIN_NGPU = 1\n",
        "BATCH_SIZE = 512\n",
        "REPLICAS = 2\n",
        "EPOCHS = 5\n",
        "ARCH = \"resnet18\"\n",
        "BACKEND = \"nccl\"  # gloo for CPU only, nccl for GPUs\n",
        "TRAIN_DATA_LOCATION = (\n",
        "    \"/trainer/tiny-imagenet-200\"  # Data location of filed downloaded in Dockerfile\n",
        ")\n",
        "\n",
        "\n",
        "display_name = (\n",
        "    CONTAINER_NAME\n",
        "    + \"-LOCAL-ETCD-reduc-server-\"\n",
        "    + f\"{REPLICAS}workers-{TRAIN_NGPU}{TRAIN_GPU}-{BATCH_SIZE}batch-\"\n",
        "    + TIMESTAMP\n",
        ")\n",
        "gcs_output_uri_prefix = f\"{BUCKET_URI}/{display_name}\"\n",
        "\n",
        "RDZV_BACKEND = \"etcd-v2\"\n",
        "RDZV_BACKEND_STORE = f\"/gcs/{BUCKET_NAME}/sharedfile-{display_name}\"\n",
        "RDZV_ENDPOINT = \"localhost:2379\"\n",
        "\n",
        "\n",
        "# Use letters for each parameter to be processed in the shell script\n",
        "\"\"\"\n",
        "e)epochs=${OPTARG};;\n",
        "a)arch=${OPTARG};;\n",
        "b)batchsize=${OPTARG};;\n",
        "d)distbackend=${OPTARG};;\n",
        "t)data=${OPTARG};;\n",
        "w)workers=${OPTARG};;\n",
        "v)env=${OPTARG};;\n",
        "u)rdvzbackend=${OPTARG};;\n",
        "i)rdvzid=${OPTARG};;\n",
        "p)endpoint=${OPTARG};;\n",
        "n)nnodes=${OPTARG};;\n",
        "r)nprocpernode=${OPTARG};;\n",
        "c)ischief=${OPTARG};;\n",
        "\"\"\"\n",
        "\n",
        "CONTAINER_SPEC = {\n",
        "    \"image_uri\": custom_container_host_image_uri,\n",
        "    \"command\": [\n",
        "        \"/bin/bash\",\n",
        "        \"main.sh\",\n",
        "        f\"-e {EPOCHS}\",\n",
        "        f\"-a {ARCH}\",\n",
        "        f\"-b {BATCH_SIZE}\",\n",
        "        f\"-d {BACKEND}\",\n",
        "        f\"-t {TRAIN_DATA_LOCATION}\",\n",
        "        f\"-w {NUM_CPUS}\",\n",
        "        f\"-v {RDZV_BACKEND_STORE}\",\n",
        "        f\"-u {RDZV_BACKEND}\",\n",
        "        f\"-i {display_name}\",\n",
        "        f\"-p {RDZV_ENDPOINT}\",\n",
        "        f\"-n {REPLICAS+1}\",\n",
        "        f\"-r {TRAIN_NGPU}\",\n",
        "        \"-c y\",\n",
        "    ],\n",
        "}\n",
        "\n",
        "CONTAINER_WORKER_SPEC = {\n",
        "    \"image_uri\": custom_container_host_image_uri,\n",
        "    \"command\": [\n",
        "        \"/bin/bash\",\n",
        "        \"main.sh\",\n",
        "        f\"-e {EPOCHS}\",\n",
        "        f\"-a {ARCH}\",\n",
        "        f\"-b {BATCH_SIZE}\",\n",
        "        f\"-d {BACKEND}\",\n",
        "        f\"-t {TRAIN_DATA_LOCATION}\",\n",
        "        f\"-w {NUM_CPUS}\",\n",
        "        f\"-v {RDZV_BACKEND_STORE}\",\n",
        "        f\"-u {RDZV_BACKEND}\",\n",
        "        f\"-i {display_name}\",\n",
        "        f\"-p {RDZV_ENDPOINT}\",\n",
        "        f\"-n {REPLICAS+1}\",\n",
        "        f\"-r {TRAIN_NGPU}\",\n",
        "        \"-c n\",\n",
        "    ],\n",
        "}\n",
        "\n",
        "PRIMARY_WORKER_POOL = {\n",
        "    \"replica_count\": 1,\n",
        "    \"machine_spec\": {\n",
        "        \"machine_type\": PRIMARY_COMPUTE,\n",
        "        \"accelerator_count\": TRAIN_NGPU,\n",
        "        \"accelerator_type\": TRAIN_GPU,\n",
        "    },\n",
        "    \"container_spec\": CONTAINER_SPEC,\n",
        "}\n",
        "\n",
        "WORKER_POOL_SPECS = [PRIMARY_WORKER_POOL]\n",
        "\n",
        "TRAIN_WORKER_POOL = {\n",
        "    \"replica_count\": REPLICAS,\n",
        "    \"machine_spec\": {\n",
        "        \"machine_type\": TRAIN_COMPUTE,\n",
        "        \"accelerator_count\": TRAIN_NGPU,\n",
        "        \"accelerator_type\": TRAIN_GPU,\n",
        "    },\n",
        "    \"container_spec\": CONTAINER_WORKER_SPEC,\n",
        "}\n",
        "\n",
        "WORKER_POOL_SPECS.append(TRAIN_WORKER_POOL)\n",
        "\n",
        "# Add Reduction Server worker pool\n",
        "REDUCTION_SERVER_REPLICAS = 3\n",
        "REDUCTION_SERVER_IMAGE_URI = (\n",
        "    \"us-docker.pkg.dev/vertex-ai-restricted/training/reductionserver:latest\"\n",
        ")\n",
        "\n",
        "CONTAINER_REDUCTION_SPEC = {\"image_uri\": REDUCTION_SERVER_IMAGE_URI}\n",
        "\n",
        "REDUCTION_WORKER_POOL = {\n",
        "    \"replica_count\": REDUCTION_SERVER_REPLICAS,\n",
        "    \"machine_spec\": {\n",
        "        \"machine_type\": REDUCTION_COMPUTE,\n",
        "    },\n",
        "    \"container_spec\": CONTAINER_REDUCTION_SPEC,\n",
        "}\n",
        "\n",
        "WORKER_POOL_SPECS.append(REDUCTION_WORKER_POOL)\n",
        "\n",
        "job = aiplatform.CustomJob(\n",
        "    display_name=display_name,\n",
        "    base_output_dir=gcs_output_uri_prefix,\n",
        "    worker_pool_specs=WORKER_POOL_SPECS,\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "d4591d0a861a"
      },
      "outputs": [],
      "source": [
        "if not os.getenv(\"IS_TESTING\"):\n",
        "    job.run(\n",
        "        sync=True\n",
        "        # comment out the line below to turn off interactive debug\n",
        "        ,\n",
        "        enable_web_access=True,\n",
        "        service_account=SERVICE_ACCOUNT,\n",
        "    )"
      ]
    },
    {
      "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": [
        "try:\n",
        "    job.delete()\n",
        "except Exception as e:\n",
        "    print(e)\n",
        "    \n",
        "gar_images = ! gcloud artifacts docker images list $LOCATION-docker.pkg.dev/$PROJECT_ID/$REPOSITORY \\\n",
        "      --filter=\"package~\"$(echo $CONTAINER_NAME | sed 's/:.*//') \\\n",
        "      --format=\"get(package)\"\n",
        "\n",
        "delete_image = True\n",
        "try:\n",
        "    if delete_image or os.getenv(\"IS_TESTING\"):\n",
        "        for image in gar_images:\n",
        "            # delete only if image name starts with valid region\n",
        "            if image.startswith(f'{LOCATION}-docker.pkg.dev'):\n",
        "                print(f\"Deleting image {image} including all tags\")\n",
        "                ! gcloud artifacts docker images delete $image --delete-tags --quiet\n",
        "except Exception as e:\n",
        "    print(e)\n",
        "\n",
        "# Delete Cloud Storage objects that were created\n",
        "delete_bucket = False\n",
        "if delete_bucket:\n",
        "    ! gsutil -m rm -r $BUCKET_URI"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [],
      "name": "sdk_pytorch_torchrun_custom_container_training_imagenet.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}