notebooks/official/training/get_started_with_vertex_distributed_training.ipynb

{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "copyright"
      },
      "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": "title:generic,gcp"
      },
      "source": [
        "# Get started with Vertex AI distributed training\n",
        "\n",
        "<table align=\"left\">\n",
        "      <td>\n",
        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/training/get_started_with_vertex_distributed_training.ipynb\">\n",
        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Colab logo\"> Run in Colab\n",
        "    </a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/training/get_started_with_vertex_distributed_training.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/official/training/get_started_with_vertex_distributed_training.ipynb\">\n",
        "      <img src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" alt=\"Vertex AI logo\">\n",
        "Open in Vertex AI Workbench\n",
        "    </a>\n",
        "  </td>\n",
        "</table>\n",
        "<br/><br/><br/>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "overview:mlops"
      },
      "source": [
        "## Overview\n",
        "\n",
        "This tutorial demonstrates how to use the Vertex AI Python client library to do distrbuted training of a TensorFlow model.\n",
        "\n",
        "*Note:* There are incompatibilities between Colab and Docker and the Docker section may not work until resolved by the platform.\n",
        "\n",
        "Learn more about [Vertex AI distributed training](https://cloud.google.com/vertex-ai/docs/training/distributed-training)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "objective:mlops,stage2,get_started_vertex_distributed_training"
      },
      "source": [
        "### Objective\n",
        "\n",
        "In this tutorial, you learn how to use `Vertex AI distributed training` when training with `Vertex AI`.\n",
        "\n",
        "This tutorial uses the following Google Cloud ML services:\n",
        "\n",
        "- `Vertex AI distributed training`\n",
        "- `Vertex AI Reduction Server`\n",
        "\n",
        "The steps performed include:\n",
        "\n",
        "- `MirroredStrategy`: Train on a single VM with multiple GPUs.\n",
        "- `MultiWorkerMirroredStrategy`: Train on multiple VMs with automatic setup of replicas.\n",
        "- `MultiWorkerMirroredStrategy`: Train on multiple VMs with fine grain control of replicas.\n",
        "- `ReductionServer`: Train on multiple VMS and sync updates across VMS with `Vertex AI Reduction Server`.\n",
        "- `TPUTraining`: Train with multiple Cloud TPUs."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "recommendation:mlops,stage2,vertex,distributed_training"
      },
      "source": [
        "### Recommendations\n",
        "\n",
        "When doing E2E MLOps on Google Cloud, the following are best practices for when to use Vertex AI distributed training:\n",
        "\n",
        "**Single VM / Single Device (OneDeviceStrategy)**\n",
        "\n",
        "You are experimenting and the total training data and number of model parameters is small.\n",
        "\n",
        "If the number of model parameters is very small, you may not get much benefit from a GPU and may consider using the VM's CPU.\n",
        "\n",
        "**Single VM / Multiple Compute Devices (MirroredStrategy)**\n",
        "\n",
        "The number of model parameters is very large, but the total training data is small.\n",
        "\n",
        "**Multiple VM / Multiple Compute Devices (MultiWorkerMirroredStrategy)**\n",
        "\n",
        "The number of model parameters is very large and the total training data is very large.\n",
        "\n",
        "**ReductionServer**\n",
        "\n",
        "While training across a large number of VMs and the model parameters updates to sync is very large."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dataset:custom,boston,lrg"
      },
      "source": [
        "### Dataset\n",
        "\n",
        "The dataset used for this tutorial is the [Boston Housing Prices dataset](https://www.cs.toronto.edu/~delve/data/boston/bostonDetail.html). The version of the dataset you use in this tutorial is built into TensorFlow. The trained model predicts the median price of a house in units of 1K USD."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "d10166df7141"
      },
      "source": [
        "### Costs\n",
        " \n",
        "This tutorial uses billable components of Google Cloud:\n",
        "\n",
        "Vertex AI\n",
        "Cloud Storage\n",
        "\n",
        "Learn about [Vertex AI pricing](https://cloud.google.com/vertex-ai/pricing) and [Cloud Storage pricing](https://cloud.google.com/storage/pricing), and use the [Pricing Calculator](https://cloud.google.com/products/calculator/),\n",
        "        to generate a cost estimate based on your projected usage."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "XkYpRvOQyVYb"
      },
      "source": [
        "## Installation\n",
        "\n",
        "Install the packages required for executing this notebook."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "xs_Kt8RcyXTC"
      },
      "outputs": [],
      "source": [
        "! pip3 install --upgrade --quiet google-cloud-aiplatform"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "58707a750154"
      },
      "source": [
        "### Colab only: Uncomment the following cell to restart the kernel."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "f200f10a1da3"
      },
      "outputs": [],
      "source": [
        "# Automatically restart kernel after installs so that your environment can access the new packages\n",
        "# import IPython\n",
        "\n",
        "# app = IPython.Application.instance()\n",
        "# app.kernel.do_shutdown(True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BF1j6f9HApxa"
      },
      "source": [
        "## Before you begin\n",
        "\n",
        "### Set up your Google Cloud project\n",
        "\n",
        "**The following steps are required, regardless of your notebook environment.**\n",
        "\n",
        "1. [Select or create a Google Cloud project](https://console.cloud.google.com/cloud-resource-manager). When you first create an account, you get a $300 free credit towards your compute/storage costs.\n",
        "\n",
        "2. [Make sure that billing is enabled for your project](https://cloud.google.com/billing/docs/how-to/modify-project).\n",
        "\n",
        "3. [Enable the Vertex AI API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com).\n",
        "\n",
        "4. If you are running this notebook locally, install the [Cloud SDK](https://cloud.google.com/sdk)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WReHDGG5g0XY"
      },
      "source": [
        "#### Set your project ID\n",
        "\n",
        "**If you don't know your project ID**, try the following:\n",
        "* Run `gcloud config list`.\n",
        "* Run `gcloud projects list`.\n",
        "* See the support page: [Locate the project ID](https://support.google.com/googleapi/answer/7014113)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "oM1iC_MfAts1"
      },
      "outputs": [],
      "source": [
        "PROJECT_ID = \"[your-project-id]\"  # @param {type:\"string\"}\n",
        "\n",
        "# Set the project id\n",
        "! gcloud config set project {PROJECT_ID}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "region"
      },
      "source": [
        "#### Set the region\n",
        "\n",
        "**Optional**: Update the 'REGION' variable to specify the region that you want to use. Learn more about [Vertex AI regions](https://cloud.google.com/vertex-ai/docs/general/locations)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "nsN5NJKSu-GU"
      },
      "outputs": [],
      "source": [
        "REGION = \"us-central1\"  # @param {type: \"string\"}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "06571eb4063b"
      },
      "source": [
        "#### UUID\n",
        "\n",
        "If you are in a live tutorial session, you might be using a shared test account or project. To avoid name collisions between users on resources created, you create a uuid 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": "4e166d927e36"
      },
      "outputs": [],
      "source": [
        "import random\n",
        "import string\n",
        "\n",
        "\n",
        "# Generate a uuid of a specifed length(default=8)\n",
        "def generate_uuid(length: int = 8) -> str:\n",
        "    return \"\".join(random.choices(string.ascii_lowercase + string.digits, k=length))\n",
        "\n",
        "\n",
        "UUID = generate_uuid()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "sBCra4QMA2wR"
      },
      "source": [
        "### Authenticate your Google Cloud account\n",
        "\n",
        "To authenticate your Google Cloud account, follow the instructions for your Jupyter environment:\n",
        "\n",
        "**1. Vertex AI Workbench**\n",
        "<br>You are already authenticated.\n",
        "\n",
        "**2. Local JupyterLab instance**\n",
        "<br>Uncomment and run the following code:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "254614fa0c46"
      },
      "outputs": [],
      "source": [
        "# ! gcloud auth login"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ef21552ccea8"
      },
      "source": [
        "**3. Colab**\n",
        "<br>Uncomment and run the following code:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "603adbbf0532"
      },
      "outputs": [],
      "source": [
        "# from google.colab import auth\n",
        "\n",
        "# auth.authenticate_user()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "c13224697bfb"
      },
      "source": [
        "**4. Service account or other**\n",
        "* See how to grant Cloud Storage permissions to your service account at https://cloud.google.com/storage/docs/gsutil/commands/iam#ch-examples."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zgPO1eR3CYjk"
      },
      "source": [
        "### Create a Cloud Storage bucket\n",
        "\n",
        "Create a storage bucket to store intermediate artifacts such as datasets."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "MzGDU7TWdts_"
      },
      "outputs": [],
      "source": [
        "BUCKET_URI = f\"gs://your-bucket-name-{PROJECT_ID}-unique\"  # @param {type:\"string\"}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "create_bucket"
      },
      "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": "setup_vars"
      },
      "source": [
        "### Set up variables\n",
        "\n",
        "Next, set up some variables used throughout the tutorial.\n",
        "### Import libraries and define constants"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "import_aip:mbsdk"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "\n",
        "import google.cloud.aiplatform as aip"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "init_aip:mbsdk"
      },
      "source": [
        "### Initialize Vertex AI SDK for Python\n",
        "\n",
        "Initialize the Vertex AI SDK for Python for your project and corresponding bucket."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "wbvYPSTDulvS"
      },
      "outputs": [],
      "source": [
        "aip.init(project=PROJECT_ID, staging_bucket=BUCKET_URI)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "accelerators:training,prediction,ngpu,mbsdk"
      },
      "source": [
        "#### Set hardware accelerators\n",
        "\n",
        "You can set hardware accelerators for training and prediction.\n",
        "\n",
        "Set the variables `TRAIN_GPU/TRAIN_NGPU` and `DEPLOY_GPU/DEPLOY_NGPU` to use a container image supporting a GPU and the number of GPUs allocated to the virtual machine (VM) instance. For example, to use a GPU container image with 4 Nvidia Telsa K80 GPUs allocated to each VM, you would specify:\n",
        "\n",
        "    (aip.AcceleratorType.NVIDIA_TESLA_K80, 4)\n",
        "\n",
        "\n",
        "Otherwise specify `(None, None)` to use a container image to run on a CPU.\n",
        "\n",
        "Learn more about [hardware accelerator support for your region](https://cloud.google.com/vertex-ai/docs/general/locations#accelerators).\n",
        "\n",
        "*Note*: TF releases before 2.3 for GPU support will fail to load the custom model in this tutorial. It is a known issue and fixed in TF 2.3. This is caused by static graph ops that are generated in the serving function. If you encounter this issue on your own custom models, use a container image for TF 2.3 with GPU support."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "PryARdnoulvT"
      },
      "outputs": [],
      "source": [
        "if os.getenv(\"IS_TESTING_TRAIN_GPU\"):\n",
        "    TRAIN_GPU, TRAIN_NGPU = (\n",
        "        aip.gapic.AcceleratorType.NVIDIA_TESLA_K80,\n",
        "        int(os.getenv(\"IS_TESTING_TRAIN_GPU\")),\n",
        "    )\n",
        "else:\n",
        "    TRAIN_GPU, TRAIN_NGPU = (aip.gapic.AcceleratorType.NVIDIA_TESLA_K80, 4)\n",
        "\n",
        "if os.getenv(\"IS_TESTING_DEPLOY_GPU\"):\n",
        "    DEPLOY_GPU, DEPLOY_NGPU = (\n",
        "        aip.gapic.AcceleratorType.NVIDIA_TESLA_K80,\n",
        "        int(os.getenv(\"IS_TESTING_DEPLOY_GPU\")),\n",
        "    )\n",
        "else:\n",
        "    DEPLOY_GPU, DEPLOY_NGPU = (None, None)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "container:training,prediction"
      },
      "source": [
        "#### Set pre-built containers\n",
        "\n",
        "Set the pre-built Docker container image for training and prediction.\n",
        "\n",
        "\n",
        "For the latest list, see [Pre-built containers for training](https://cloud.google.com/ai-platform-unified/docs/training/pre-built-containers).\n",
        "\n",
        "\n",
        "For the latest list, see [Pre-built containers for prediction](https://cloud.google.com/ai-platform-unified/docs/predictions/pre-built-containers)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "LhhUFw2nulvT"
      },
      "outputs": [],
      "source": [
        "if os.getenv(\"IS_TESTING_TF\"):\n",
        "    TF = os.getenv(\"IS_TESTING_TF\")\n",
        "else:\n",
        "    TF = \"2.5\".replace(\".\", \"-\")\n",
        "\n",
        "if TF[0] == \"2\":\n",
        "    if TRAIN_GPU:\n",
        "        TRAIN_VERSION = \"tf-gpu.{}\".format(TF)\n",
        "    else:\n",
        "        TRAIN_VERSION = \"tf-cpu.{}\".format(TF)\n",
        "    if DEPLOY_GPU:\n",
        "        DEPLOY_VERSION = \"tf2-gpu.{}\".format(TF)\n",
        "    else:\n",
        "        DEPLOY_VERSION = \"tf2-cpu.{}\".format(TF)\n",
        "else:\n",
        "    if TRAIN_GPU:\n",
        "        TRAIN_VERSION = \"tf-gpu.{}\".format(TF)\n",
        "    else:\n",
        "        TRAIN_VERSION = \"tf-cpu.{}\".format(TF)\n",
        "    if DEPLOY_GPU:\n",
        "        DEPLOY_VERSION = \"tf-gpu.{}\".format(TF)\n",
        "    else:\n",
        "        DEPLOY_VERSION = \"tf-cpu.{}\".format(TF)\n",
        "\n",
        "TRAIN_IMAGE = \"{}-docker.pkg.dev/vertex-ai/training/{}:latest\".format(\n",
        "    REGION.split(\"-\")[0], TRAIN_VERSION\n",
        ")\n",
        "DEPLOY_IMAGE = \"{}-docker.pkg.dev/vertex-ai/prediction/{}:latest\".format(\n",
        "    REGION.split(\"-\")[0], DEPLOY_VERSION\n",
        ")\n",
        "\n",
        "print(\"Training:\", TRAIN_IMAGE, TRAIN_GPU, TRAIN_NGPU)\n",
        "print(\"Deployment:\", DEPLOY_IMAGE, DEPLOY_GPU, DEPLOY_NGPU)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "machine:training"
      },
      "source": [
        "#### Set machine type\n",
        "\n",
        "Next, set the machine type to use for training.\n",
        "\n",
        "- Set the variable `TRAIN_COMPUTE` to configure  the compute resources for the VMs you will use for for training.\n",
        " - `machine type`\n",
        "     - `n1-standard`: 3.75GB of memory per vCPU.\n",
        "     - `n1-highmem`: 6.5GB of memory per vCPU\n",
        "     - `n1-highcpu`: 0.9 GB of memory per vCPU\n",
        " - `vCPUs`: number of \\[2, 4, 8, 16, 32, 64, 96 \\]\n",
        "\n",
        "*Note: The following is not supported for training:*\n",
        "\n",
        " - `standard`: 2 vCPUs\n",
        " - `highcpu`: 2, 4 and 8 vCPUs\n",
        "\n",
        "*Note: You may also use n2 and e2 machine types for training and deployment, but they do not support GPUs*."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "vytMaukeulvT"
      },
      "outputs": [],
      "source": [
        "if os.getenv(\"IS_TESTING_TRAIN_MACHINE\"):\n",
        "    MACHINE_TYPE = os.getenv(\"IS_TESTING_TRAIN_MACHINE\")\n",
        "else:\n",
        "    MACHINE_TYPE = \"n1-standard\"\n",
        "\n",
        "VCPU = \"4\"\n",
        "TRAIN_COMPUTE = MACHINE_TYPE + \"-\" + VCPU\n",
        "print(\"Train machine type\", TRAIN_COMPUTE)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "mirrored_intro"
      },
      "source": [
        "## Mirrored Strategy\n",
        "\n",
        "When training on a single VM, one can either train was a single compute device or with multiple compute devices on the same VM. With Vertex AI distributed training you can specify both the number of compute devices for the VM instance and type of compute devices: CPU, GPU.\n",
        "\n",
        "Vertex AI distributed training supports `tf.distribute.MirroredStrategy' for TensorFlow models. To enable training across multiple compute devices on the same VM, you do the following additional steps in your Python training script:\n",
        "\n",
        "1. Set the tf.distribute.MirrorStrategy\n",
        "2. Compile the model within the scope of tf.distribute.MirrorStrategy. *Note:* Tells MirroredStrategy which variables to mirror across your compute devices.\n",
        "3. Increase the batch size for each compute device to num_devices * batch size.\n",
        "\n",
        "During transitions, the distribution of batches will be synchronized as well as the updates to the model parameters."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "create_custom_pp_training_job:mbsdk"
      },
      "source": [
        "### Create and run custom training job\n",
        "\n",
        "\n",
        "To train a custom model, you perform two steps: 1) create a custom training job, and 2) run the job.\n",
        "\n",
        "#### Create custom training job\n",
        "\n",
        "A custom training job is created with the `CustomTrainingJob` class, with the following parameters:\n",
        "\n",
        "- `display_name`: The human readable name for the custom training job.\n",
        "- `container_uri`: The training container image.\n",
        "\n",
        "- `python_package_gcs_uri`: The location of the Python training package as a tarball.\n",
        "- `python_module_name`: The relative path to the training script in the Python package.\n",
        "- `model_serving_container_uri`: The container image for deploying the model.\n",
        "\n",
        "*Note:* There is no requirements parameter. You specify any requirements in the `setup.py` script in your Python package."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "mhw34XoOulvU"
      },
      "outputs": [],
      "source": [
        "DISPLAY_NAME = \"boston_\" + UUID\n",
        "\n",
        "job = aip.CustomPythonPackageTrainingJob(\n",
        "    display_name=DISPLAY_NAME,\n",
        "    python_package_gcs_uri=f\"{BUCKET_URI}/trainer_boston.tar.gz\",\n",
        "    python_module_name=\"trainer.task\",\n",
        "    container_uri=TRAIN_IMAGE,\n",
        "    model_serving_container_image_uri=DEPLOY_IMAGE,\n",
        "    project=PROJECT_ID,\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "examine_training_package"
      },
      "source": [
        "### Examine the training package\n",
        "\n",
        "#### Package layout\n",
        "\n",
        "Before you start the training, you will look at how a Python package is assembled for a custom training job. When unarchived, the package contains the following directory/file layout.\n",
        "\n",
        "- PKG-INFO\n",
        "- README.md\n",
        "- setup.cfg\n",
        "- setup.py\n",
        "- trainer\n",
        "  - \\_\\_init\\_\\_.py\n",
        "  - task.py\n",
        "\n",
        "The files `setup.cfg` and `setup.py` are the instructions for installing the package into the operating environment of the Docker image.\n",
        "\n",
        "The file `trainer/task.py` is the Python script for executing the custom training job. *Note*, when we referred to it in the worker pool specification, we replace the directory slash with a dot (`trainer.task`) and dropped the file suffix (`.py`).\n",
        "\n",
        "#### Package Assembly\n",
        "\n",
        "In the following cells, you will assemble the training package."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "IAaZpZyyulvU"
      },
      "outputs": [],
      "source": [
        "# Make folder for Python training script\n",
        "! rm -rf custom\n",
        "! mkdir custom\n",
        "\n",
        "# Add package information\n",
        "! touch custom/README.md\n",
        "\n",
        "setup_cfg = \"[egg_info]\\n\\ntag_build =\\n\\ntag_date = 0\"\n",
        "! echo \"$setup_cfg\" > custom/setup.cfg\n",
        "\n",
        "setup_py = \"import setuptools\\n\\nsetuptools.setup(\\n\\n    install_requires=[\\n\\n        'tensorflow==2.5.0',\\n\\n        'tensorflow_datasets==1.3.0',\\n\\n    ],\\n\\n    packages=setuptools.find_packages())\"\n",
        "! echo \"$setup_py\" > custom/setup.py\n",
        "\n",
        "pkg_info = \"Metadata-Version: 1.0\\n\\nName: Boston Housing cloud\\n\\nVersion: 0.0.0\\n\\nSummary: Demostration training script\\n\\nHome-page: www.google.com\\n\\nAuthor: Google\\n\\nAuthor-email: aferlitsch@google.com\\n\\nLicense: Public\\n\\nDescription: Demo\\n\\nPlatform: Vertex\"\n",
        "! echo \"$pkg_info\" > custom/PKG-INFO\n",
        "\n",
        "# Make the training subfolder\n",
        "! mkdir custom/trainer\n",
        "! touch custom/trainer/__init__.py"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "taskpy_contents:mirrored,boston"
      },
      "source": [
        "#### Task.py contents\n",
        "\n",
        "In the next cell, you write the contents of the training script task.py. I won't go into detail, it's just there for you to browse. In summary:\n",
        "\n",
        "- Get the directory where to save the model artifacts from the command line (`--model_dir`), and if not specified, then from the environment variable `AIP_MODEL_DIR`.\n",
        "- Loads Boston Housing dataset from TF.Keras builtin datasets\n",
        "- Builds a simple deep neural network model using TF.Keras model API.\n",
        "- Compiles the model (`compile()`).\n",
        "- Sets a training distribution strategy according to the argument `args.distribute`.\n",
        "- Trains the model (`fit()`) with epochs specified by `args.epochs`.\n",
        "- Saves the trained model (`save(args.model_dir)`) to the specified model directory.\n",
        "- Saves the maximum value for each feature `f.write(str(params))` to the specified parameters file."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "zKzddzl6ulvV"
      },
      "outputs": [],
      "source": [
        "%%writefile custom/trainer/task.py\n",
        "# Single, Mirrored and MultiWorker Distributed Training\n",
        "\n",
        "import tensorflow_datasets as tfds\n",
        "import tensorflow as tf\n",
        "from tensorflow.python.client import device_lib\n",
        "import numpy as np\n",
        "import argparse\n",
        "import os\n",
        "import sys\n",
        "import logging\n",
        "\n",
        "parser = argparse.ArgumentParser()\n",
        "parser.add_argument('--model-dir', dest='model_dir',\n",
        "                    default=os.getenv('AIP_MODEL_DIR'), type=str, help='Model dir.')\n",
        "parser.add_argument('--lr', dest='lr',\n",
        "                    default=0.001, type=float,\n",
        "                    help='Learning rate.')\n",
        "parser.add_argument('--epochs', dest='epochs',\n",
        "                    default=10, type=int,\n",
        "                    help='Number of epochs.')\n",
        "parser.add_argument('--steps', dest='steps',\n",
        "                    default=100, type=int,\n",
        "                    help='Number of steps per epoch.')\n",
        "parser.add_argument('--batch_size', dest='batch_size',\n",
        "                    default=16, type=int,\n",
        "                    help='Size of a batch.')\n",
        "parser.add_argument('--distribute', dest='distribute', type=str, default='single',\n",
        "                    help='distributed training strategy')\n",
        "parser.add_argument('--param-file', dest='param_file',\n",
        "                    default='/tmp/param.txt', type=str,\n",
        "                    help='Output file for parameters')\n",
        "args = parser.parse_args()\n",
        "\n",
        "logging.info('DEVICES'  + str(device_lib.list_local_devices()))\n",
        "\n",
        "# Single Machine, single compute device\n",
        "if args.distribute == 'single':\n",
        "    if tf.test.is_gpu_available():\n",
        "        strategy = tf.distribute.OneDeviceStrategy(device=\"/gpu:0\")\n",
        "    else:\n",
        "        strategy = tf.distribute.OneDeviceStrategy(device=\"/cpu:0\")\n",
        "    logging.info(\"Single device training\")\n",
        "# Single Machine, multiple compute device\n",
        "elif args.distribute == 'mirrored':\n",
        "    strategy = tf.distribute.MirroredStrategy()\n",
        "    logging.info(\"Mirrored Strategy distributed training\")\n",
        "# Multi Machine, multiple compute device\n",
        "elif args.distribute == 'multiworker':\n",
        "    strategy = tf.distribute.MultiWorkerMirroredStrategy()\n",
        "    logging.info(\"Multi-worker Strategy distributed training\")\n",
        "    logging.info('TF_CONFIG = {}'.format(os.environ.get('TF_CONFIG', 'Not found')))\n",
        "    # Single Machine, multiple TPU devices\n",
        "elif args.distribute == 'tpu':\n",
        "    cluster_resolver = tf.distribute.cluster_resolver.TPUClusterResolver(tpu=\"local\")\n",
        "    tf.config.experimental_connect_to_cluster(cluster_resolver)\n",
        "    tf.tpu.experimental.initialize_tpu_system(cluster_resolver)\n",
        "    strategy = tf.distribute.TPUStrategy(cluster_resolver)\n",
        "    print(\"All devices: \", tf.config.list_logical_devices('TPU'))\n",
        "\n",
        "logging.info('num_replicas_in_sync = {}'.format(strategy.num_replicas_in_sync))\n",
        "\n",
        "def _is_chief(task_type, task_id):\n",
        "    ''' Check for primary if multiworker training\n",
        "    '''\n",
        "    return (task_type == 'chief') or (task_type == 'worker' and task_id == 0) or task_type is None\n",
        "\n",
        "\n",
        "def get_data():\n",
        "    # Scaling Boston Housing data features\n",
        "    def scale(feature):\n",
        "        max = np.max(feature)\n",
        "        feature = (feature / max).astype(np.float)\n",
        "        return feature, max\n",
        "\n",
        "    (x_train, y_train), (x_test, y_test) = tf.keras.datasets.boston_housing.load_data(\n",
        "        path=\"boston_housing.npz\", test_split=0.2, seed=113\n",
        "    )\n",
        "\n",
        "    params = []\n",
        "    for _ in range(13):\n",
        "        x_train[_], max = scale(x_train[_])\n",
        "        x_test[_], _ = scale(x_test[_])\n",
        "    params.append(max)\n",
        "\n",
        "    # store the normalization (max) value for each feature\n",
        "    with tf.io.gfile.GFile(args.param_file, 'w') as f:\n",
        "        f.write(str(params))\n",
        "    return (x_train, y_train), (x_test, y_test)\n",
        "\n",
        "def get_model():\n",
        "    model = tf.keras.Sequential([\n",
        "          tf.keras.layers.Dense(128, activation='relu', input_shape=(13,)),\n",
        "          tf.keras.layers.Dense(128, activation='relu'),\n",
        "          tf.keras.layers.Dense(1, activation='linear')\n",
        "    ])\n",
        "\n",
        "    model.compile(\n",
        "          loss='mse',\n",
        "          optimizer=tf.keras.optimizers.RMSprop(learning_rate=args.lr)\n",
        "    )\n",
        "    return model\n",
        "\n",
        "def train(model, x_train, y_train):\n",
        "    NUM_WORKERS = strategy.num_replicas_in_sync\n",
        "    # Here the batch size scales up by number of workers since\n",
        "    # `tf.data.Dataset.batch` expects the global batch size.\n",
        "    GLOBAL_BATCH_SIZE = args.batch_size * NUM_WORKERS\n",
        "\n",
        "    model.fit(x_train, y_train, epochs=args.epochs, batch_size=GLOBAL_BATCH_SIZE)\n",
        "\n",
        "    if args.distribute == 'multiworker':\n",
        "        task_type, task_id = (strategy.cluster_resolver.task_type,\n",
        "                              strategy.cluster_resolver.task_id)\n",
        "    else:\n",
        "        task_type, task_id = None, None\n",
        "\n",
        "    if args.distribute==\"tpu\":\n",
        "        save_locally = tf.saved_model.SaveOptions(experimental_io_device='/job:localhost')\n",
        "        model.save(args.model_dir, options=save_locally)\n",
        "    # single, mirrored or primary for multiworker\n",
        "    elif _is_chief(task_type, task_id):\n",
        "        model.save(args.model_dir)\n",
        "    # non-primary workers for multi-workers\n",
        "    else:\n",
        "        # each worker saves their model instance to a unique temp location\n",
        "        worker_dir = args.model_dir + '/workertemp_' + str(task_id)\n",
        "        tf.io.gfile.makedirs(worker_dir)\n",
        "        model.save(worker_dir)\n",
        "\n",
        "with strategy.scope():\n",
        "    # Creation of dataset, and model building/compiling need to be within\n",
        "    # `strategy.scope()`.\n",
        "    model = get_model()\n",
        "\n",
        "(x_train, y_train), (x_test, y_test) = get_data()\n",
        "\n",
        "train(model, x_train, y_train)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tarball_training_script"
      },
      "source": [
        "#### Store training script on your Cloud Storage bucket\n",
        "\n",
        "Next, you package the training folder into a compressed tar ball, and then store it in your Cloud Storage bucket."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "LFUHioqTulvV"
      },
      "outputs": [],
      "source": [
        "! rm -f custom.tar custom.tar.gz\n",
        "! tar cvf custom.tar custom\n",
        "! gzip custom.tar\n",
        "! gsutil cp custom.tar.gz $BUCKET_URI/trainer_boston.tar.gz"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "run_custom_pp_training_job:mirrored"
      },
      "source": [
        "#### Run the custom Python package training job\n",
        "\n",
        "Next, you run the custom job to start the training job by invoking the method `run()`. The parameters are the same as when running a CustomTrainingJob."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "LnUX0UkvulvV"
      },
      "outputs": [],
      "source": [
        "MODEL_DIR = BUCKET_URI\n",
        "\n",
        "CMDARGS = [\"--epochs=5\", \"--batch_size=16\", \"--distribute=mirrored\"]\n",
        "\n",
        "model = job.run(\n",
        "    model_display_name=\"boston_\" + UUID,\n",
        "    args=CMDARGS,\n",
        "    replica_count=1,\n",
        "    machine_type=TRAIN_COMPUTE,\n",
        "    accelerator_type=TRAIN_GPU.name,\n",
        "    accelerator_count=TRAIN_NGPU,\n",
        "    base_output_dir=MODEL_DIR,\n",
        "    sync=True,\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "delete_job"
      },
      "source": [
        "### Delete a custom training job\n",
        "\n",
        "After a training job is completed, you can delete the training job with the method `delete()`.  Prior to completion, a training job can be canceled with the method `cancel()`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "iUWHFpPoulvW"
      },
      "outputs": [],
      "source": [
        "job.delete()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "model_delete:mbsdk"
      },
      "source": [
        "#### Delete the model\n",
        "\n",
        "The method 'delete()' will delete the model."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "-0gqCUTEulvW"
      },
      "outputs": [],
      "source": [
        "model.delete()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "multiworker_intro"
      },
      "source": [
        "## Multi-Worker Mirrored Strategy\n",
        "\n",
        "With Vertex AI distributed training you can train with multiple VM instances\n",
        "\n",
        "Vertex AI distributed training supports `tf.distribute.MultiWorkerMirroredStrategy' for TensorFlow and PyTorch models. To enable training across multiple VMS, you do the following additional steps in your Python training script:\n",
        "\n",
        "1. All the additional steps for MirroredStrategy, except that MultiWorkerStrategy is set in place of MirroredStrategy.\n",
        "2. Setup the worker pools.\n",
        "3. Alter the saving of the model so that the non-primary workers save their model instance to a unique temporary directory each.\n",
        "\n",
        "*Note:* You do not need to construct the TF_CONFIG environment variable. It is automatically constructed by Vertex AI distributed training.\n",
        "\n",
        "Learn more about [distributed training](https://cloud.google.com/vertex-ai/docs/training/distributed-training)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "worker_pools"
      },
      "source": [
        "### Worker pools\n",
        "\n",
        "If you run a distributed training job with Vertex AI, you specify multiple machines (nodes) in a training cluster. The training service allocates the resources for the machine types you specify. Your running job on a given node is called a replica. A group of replicas with the same configuration is called a worker pool.\n",
        "\n",
        "Each replica in the training cluster is given a single role or task in distributed training. For example:\n",
        "\n",
        "- **Primary replica**: Exactly one replica is designated the primary replica. This task manages the others and reports status for the job as a whole.\n",
        "\n",
        "- **Worker(s)**: One or more replicas may be designated as workers. These replicas do their portion of the work as you designate in your job configuration.\n",
        "\n",
        "- Parameter server(s): If supported by your ML framework, one or more replicas may be designated as parameter servers. These replicas store model parameters and coordinate shared model state between the workers.\n",
        "\n",
        "Evaluator(s): If supported by your ML framework, one or more replicas may be designated as evaluators. These replicas can be used to evaluate your model. If you are using TensorFlow, note that TensorFlow generally expects that you use no more than one evaluator.\n",
        "\n",
        "To configure a distributed training job, define your list of worker pools (workerPoolSpecs[]), designating one WorkerPoolSpec for each type of task:\n",
        "\n",
        "*Note:* The worker pool is order dependent (0..3):\n",
        "\n",
        "**workerPoolSpecs[0]**: Primary, chief, scheduler, or \"master\"\n",
        "\n",
        "**workerPoolSpecs[1]**: Secondary, replicas, workers\n",
        "\n",
        "**workerPoolSpecs[2]**: Parameter servers, Reduction Server\n",
        "\n",
        "**workerPoolSpecs[2]**: Evaluators"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "multiworker_methods"
      },
      "source": [
        "### Distributed training options for Multi-Worker Mirrored Strategy\n",
        "\n",
        "How you setup the worker pools is dependent on the Vertex AI method you use for training.\n",
        "\n",
        "**CustomTrainingJob** / **CustomContainerTrainingJob** / **CustomPythonPackageTrainingJob**\n",
        "\n",
        "The `replica_count` includes the primary and secondary (replica_count-1), and share the same machine type and accelerators.\n",
        "\n",
        "You cannot specify a parameter server or evaluation node.\n",
        "\n",
        "**CustomJob**\n",
        "\n",
        "You specify a `worker_pool_spec`, where you can specify detailed settings for each of the four worker pools."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "aXvPN8P6ulvX"
      },
      "source": [
        "### Create and run custom training job\n",
        "\n",
        "\n",
        "To train a custom model, you perform two steps: 1) create a custom training job, and 2) run the job.\n",
        "\n",
        "#### Create custom training job\n",
        "\n",
        "A custom training job is created with the `CustomTrainingJob` class, with the following parameters:\n",
        "\n",
        "- `display_name`: The human readable name for the custom training job.\n",
        "- `container_uri`: The training container image.\n",
        "\n",
        "- `python_package_gcs_uri`: The location of the Python training package as a tarball.\n",
        "- `python_module_name`: The relative path to the training script in the Python package.\n",
        "- `model_serving_container_uri`: The container image for deploying the model.\n",
        "\n",
        "*Note:* There is no requirements parameter. You specify any requirements in the `setup.py` script in your Python package."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "kYcFsVSEulvX"
      },
      "outputs": [],
      "source": [
        "DISPLAY_NAME = \"boston_\" + UUID\n",
        "\n",
        "job = aip.CustomPythonPackageTrainingJob(\n",
        "    display_name=DISPLAY_NAME,\n",
        "    python_package_gcs_uri=f\"{BUCKET_URI}/trainer_boston.tar.gz\",\n",
        "    python_module_name=\"trainer.task\",\n",
        "    container_uri=TRAIN_IMAGE,\n",
        "    model_serving_container_image_uri=DEPLOY_IMAGE,\n",
        "    project=PROJECT_ID,\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "run_custom_pp_training_job:multiworker"
      },
      "source": [
        "#### Run the custom Python package training job\n",
        "\n",
        "Next, you run the custom job to start the training job by invoking the method `run()`. The parameters are the same as when running a CustomTrainingJob."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "GHRxPU32ulvX"
      },
      "outputs": [],
      "source": [
        "MODEL_DIR = BUCKET_URI\n",
        "\n",
        "CMDARGS = [\"--epochs=5\", \"--batch_size=16\", \"--distribute=multiworker\"]\n",
        "\n",
        "try:\n",
        "    model = job.run(\n",
        "        model_display_name=\"boston_\" + UUID,\n",
        "        args=CMDARGS,\n",
        "        replica_count=4,\n",
        "        machine_type=TRAIN_COMPUTE,\n",
        "        accelerator_type=TRAIN_GPU.name,\n",
        "        accelerator_count=TRAIN_NGPU,\n",
        "        base_output_dir=MODEL_DIR,\n",
        "        sync=True,\n",
        "    )\n",
        "except Exception as e:\n",
        "    # may fail duing model.save() -- seems to be some issue when merging checkpoints from the workers\n",
        "    print(e)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "92D_hbuVulvX"
      },
      "source": [
        "### Delete a custom training job\n",
        "\n",
        "After a training job is completed, you can delete the training job with the method `delete()`.  Prior to completion, a training job can be canceled with the method `cancel()`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "CqrfWkB3ulvX"
      },
      "outputs": [],
      "source": [
        "job.delete()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "customjob_intro:multiworker"
      },
      "source": [
        "### Multiworker distributed training with CustomJob\n",
        "\n",
        "Multiworker distributed training with `CustomJob` has the advantages of fine detail control of the primary replica and optionally specifying worker pools for parameter server and evaluators. Creating a `CustomJob` includes the following steps:\n",
        "\n",
        "\n",
        "1. Specify individual details for each worker pool.\n",
        "2. Embed training package into Docker image."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "create_docker_container:training"
      },
      "source": [
        "### Create a Docker file\n",
        "\n",
        "To use your own custom training container, you build a Docker file and embed into the container your training scripts."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "write_docker_file:training,multiworker"
      },
      "source": [
        "#### Write the Docker file contents\n",
        "\n",
        "Your first step in containerizing your code is to create a Docker file. In your Docker you’ll include all the commands needed to run your container image. It’ll install all the libraries you’re using and set up the entry point for your training code.\n",
        "\n",
        "1. Install a pre-defined container image from TensorFlow repository for deep learning images.\n",
        "2. Copies in the Python training code, to be shown subsequently.\n",
        "3. Sets the entry into the Python training script as `trainer/task.py`. Note, the `.py` is dropped in the ENTRYPOINT command, as it is implied."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "pGI2viDAulvY"
      },
      "outputs": [],
      "source": [
        "%%writefile custom/Dockerfile\n",
        "\n",
        "FROM gcr.io/deeplearning-platform-release/tf2-gpu.2-5\n",
        "\n",
        "WORKDIR /\n",
        "\n",
        "# Copies the trainer code to the docker image.\n",
        "COPY trainer /trainer\n",
        "\n",
        "# Sets up the entry point to invoke the trainer.\n",
        "ENTRYPOINT [\"python\", \"-m\", \"trainer.task\"]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "name_container:training"
      },
      "source": [
        "#### Build the container locally\n",
        "\n",
        "Next, you will provide a name for your customer container that you will use when you submit it to the Google Container Registry."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7P8cdlFtulvY"
      },
      "outputs": [],
      "source": [
        "TRAIN_IMAGE = \"gcr.io/\" + PROJECT_ID + \"/boston:v1\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "build_container:training"
      },
      "source": [
        "Next, build the container."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "jmw5cakNulvY"
      },
      "outputs": [],
      "source": [
        "import sys\n",
        "\n",
        "IS_COLAB = \"google.colab\" in sys.modules\n",
        "if not IS_COLAB:\n",
        "    ! docker build custom -t $TRAIN_IMAGE\n",
        "else:\n",
        "    # install docker daemon\n",
        "    ! apt-get -qq install docker.io"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "test_container:training"
      },
      "source": [
        "#### Test the container locally\n",
        "\n",
        "Run the container within your notebook instance to ensure it’s working correctly. You will run it for 5 epochs."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "jJGLjU-TulvZ"
      },
      "outputs": [],
      "source": [
        "if not IS_COLAB:\n",
        "    ! docker run $TRAIN_IMAGE --epochs=5 --model-dir=./"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "register_container:training"
      },
      "source": [
        "#### Register the custom container\n",
        "\n",
        "When you’ve finished running the container locally, push it to Google Container Registry."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "GAXGjae7ulvZ"
      },
      "outputs": [],
      "source": [
        "if not IS_COLAB:\n",
        "    ! docker push $TRAIN_IMAGE"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "f50e9c553fb7"
      },
      "source": [
        "*Executes in Colab*"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "a7e8c98f1e56"
      },
      "outputs": [],
      "source": [
        "%%bash -s $IS_COLAB $TRAIN_IMAGE\n",
        "if [ $1 == \"False\" ]; then\n",
        "  exit 0\n",
        "fi\n",
        "set -x\n",
        "dockerd -b none --iptables=0 -l warn &\n",
        "for i in $(seq 5); do [ ! -S \"/var/run/docker.sock\" ] && sleep 2 || break; done\n",
        "docker build custom -t $2\n",
        "docker run $2 --epochs=5 --model-dir=./\n",
        "docker push $2\n",
        "kill $(jobs -p)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "worker_pool_primary"
      },
      "source": [
        "#### Primary worker pool\n",
        "\n",
        "The primary worker pool (index 0) coordinates the work done by all the other replicas. Set the replicaCount to 1. Since the worker is coordinating and not training, use a general purpose CPU, instead of a GPU.\n",
        "\n",
        "Learn more about [Machine Types for Training](https://cloud.google.com/vertex-ai/docs/training/configure-compute#machine-types)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "CEAnXBzCulvZ"
      },
      "outputs": [],
      "source": [
        "PRIMARY_COMPUTE = \"n2-highcpu-64\"\n",
        "\n",
        "MODEL_DIR = BUCKET_URI\n",
        "\n",
        "CMDARGS = [\n",
        "    \"--model-dir=\" + MODEL_DIR,\n",
        "    \"--epochs=5\",\n",
        "    \"--batch_size=16\",\n",
        "    \"--distribute=multiworker\",\n",
        "]\n",
        "\n",
        "CONTAINER_SPEC = {\"image_uri\": TRAIN_IMAGE, \"command\": \"trainer.task\", \"args\": CMDARGS}\n",
        "\n",
        "PRIMARY_WORKER_POOL = {\n",
        "    \"replica_count\": 1,\n",
        "    \"machine_spec\": {\"machine_type\": PRIMARY_COMPUTE, \"accelerator_count\": 0},\n",
        "    \"container_spec\": CONTAINER_SPEC,\n",
        "}\n",
        "\n",
        "WORKER_POOL_SPECS = [PRIMARY_WORKER_POOL]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "worker_pool_training"
      },
      "source": [
        "#### Training worker pool\n",
        "\n",
        "The secondary worker pool (index 1) performs model training. Each of the replicas will have an instance of the your training package installed on it.\n",
        "\n",
        "Each replica may have one (single device training) or multiple (mirrored) compute devices for training."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6dchPSfNulvZ"
      },
      "outputs": [],
      "source": [
        "TRAIN_WORKER_POOL = {\n",
        "    \"replica_count\": 4,\n",
        "    \"machine_spec\": {\n",
        "        \"machine_type\": TRAIN_COMPUTE,\n",
        "        \"accelerator_count\": TRAIN_NGPU,\n",
        "        \"accelerator_type\": TRAIN_GPU,\n",
        "    },\n",
        "    \"container_spec\": CONTAINER_SPEC,\n",
        "}\n",
        "\n",
        "WORKER_POOL_SPECS.append(TRAIN_WORKER_POOL)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "custom_job:worker_pool"
      },
      "source": [
        "### Create CustomJob with worker pool specifications\n",
        "\n",
        "Next, you create a `CustomJob` for the multi-worker distributed training job:\n",
        "\n",
        "-`display_name`: The display name for the custom job.\n",
        "\n",
        "-`worker_pool_specs`: The detailed specifications for each worker pool."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "m2VgmqEOulva"
      },
      "outputs": [],
      "source": [
        "DISPLAY_NAME = \"boston_\" + UUID\n",
        "\n",
        "job = aip.CustomJob(display_name=DISPLAY_NAME, worker_pool_specs=WORKER_POOL_SPECS)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "run_custom_job:multiworker"
      },
      "source": [
        "### Run the CustomJob\n",
        "\n",
        "Next, you run the custom job."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "hg8vnI_Wulva"
      },
      "outputs": [],
      "source": [
        "try:\n",
        "    job.run(sync=True)\n",
        "except Exception as e:\n",
        "    # may fail in multi-worker to find startup script\n",
        "    print(e)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WT76Sc-culva"
      },
      "source": [
        "### Delete a custom training job\n",
        "\n",
        "After a training job is completed, you can delete the training job with the method `delete()`.  Prior to completion, a training job can be canceled with the method `cancel()`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "I_IxVfuDulva"
      },
      "outputs": [],
      "source": [
        "job.delete()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "reduction_server_intro"
      },
      "source": [
        "## Reduction Server\n",
        "\n",
        "To speed up training of large models, many engineering teams are adopting distributed training using scale-out clusters of ML accelerators. However, distributed training at scale brings its own set of challenges. Specifically, limited network bandwidth between nodes makes optimizing performance of distributed training inherently difficult, particularly for large cluster configurations.\n",
        "\n",
        "Vertex AI Reduction Server optimizes bandwidth and latency of multi-node distributed training on NVIDIA GPUs for synchronous data parallel algorithms. Synchronous data parallelism is the foundation of many widely adopted distributed training frameworks, including TensorFlow’s MultiWorkerMirroredStrategy, Horovod, and PyTorch Distributed. By optimizing bandwidth usage and latency of the all-reduce collective operation used by these frameworks, Reduction Server can decrease both the time and cost of large training jobs.\n",
        "\n",
        "Learn more about [Optimizing training performance using Vertex Reduction Server](https://cloud.google.com/blog/topics/developers-practitioners/optimize-training-performance-reduction-server-vertex-ai)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "worker_pool_reduction_server"
      },
      "outputs": [],
      "source": [
        "reduction_server_count = 1\n",
        "reduction_server_machine_type = \"n1-highcpu-16\"\n",
        "reduction_server_image_uri = (\n",
        "    \"us-docker.pkg.dev/vertex-ai-restricted/training/reductionserver:latest\"\n",
        ")\n",
        "\n",
        "PARAMETER_POOL = {\n",
        "    \"replica_count\": reduction_server_count,\n",
        "    \"machine_spec\": {\n",
        "        \"machine_type\": reduction_server_machine_type,\n",
        "    },\n",
        "    \"container_spec\": {\"image_uri\": reduction_server_image_uri},\n",
        "}\n",
        "WORKER_POOL_SPECS.append(PARAMETER_POOL)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "L8Av8ATVulvb"
      },
      "source": [
        "### Create CustomJob with worker pool specifications\n",
        "\n",
        "Next, you create a `CustomJob` for the multi-worker distributed training job:\n",
        "\n",
        "-`display_name`: The display name for the custom job.\n",
        "\n",
        "-`worker_pool_specs`: The detailed specifications for each worker pool."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "TUWEP1Lmulvb"
      },
      "outputs": [],
      "source": [
        "DISPLAY_NAME = \"boston_\" + UUID\n",
        "\n",
        "job = aip.CustomJob(display_name=DISPLAY_NAME, worker_pool_specs=WORKER_POOL_SPECS)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_95FH8jeulvb"
      },
      "source": [
        "### Run the CustomJob\n",
        "\n",
        "Next, you run the custom job."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "IEbrY05Gulvb"
      },
      "outputs": [],
      "source": [
        "try:\n",
        "    job.run(sync=True)\n",
        "except Exception as e:\n",
        "    # may fail in multi-worker to find startup script\n",
        "    print(e)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8R2Bnmwmulvb"
      },
      "source": [
        "### Delete a custom training job\n",
        "\n",
        "After a training job is completed, you can delete the training job with the method `delete()`.  Prior to completion, a training job can be canceled with the method `cancel()`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "s1geVE3Lulvb"
      },
      "outputs": [],
      "source": [
        "job.delete()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tpu_intro"
      },
      "source": [
        "## Cloud TPU Training\n",
        "\n",
        "To further speed up trainig, your organization can utilize Google's Cloud Tensor Processing Units (TPU) pods.\n",
        "\n",
        "Cloud TPU is the custom-designed machine learning ASIC that powers Google products like Translate, Photos, Search, Assistant, and Gmail. Cloud TPU is designed to run cutting-edge machine learning models with AI services on Google Cloud. And its custom high-speed network offers over 100 petaflops of performance in a single pod.\n",
        "\n",
        "Learn more about [Cloud TPU](https://cloud.google.com/tpu)\n",
        "\n",
        "*Note*: TPU VM Training is currently an opt-in feature. Your GCP project must first be added to the feature allowlist. Please email your project information(project id/number) to vertex-ai-tpu-vm-training-support@google.com for the allowlist. You will receive an email as soon as your project is ready."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "docker_write:tpu"
      },
      "source": [
        "### Write Docker file for TPU training\n",
        "\n",
        "Currently, there is no pre-built Vertex AI Docker image for training with TPUs. No problems, you can make your own, as follows:\n",
        "\n",
        "1. Create a vanilla Python 3 image (e.g., `python3:8`).\n",
        "2. Get and install the TPU library (`libtpu.so`).\n",
        "3. Copy in your training package"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "nQVPtknpulvb"
      },
      "outputs": [],
      "source": [
        "%%writefile custom/Dockerfile\n",
        "FROM python:3.8\n",
        "\n",
        "WORKDIR /\n",
        "\n",
        "# Copies the trainer code to the docker image.\n",
        "COPY trainer /trainer\n",
        "\n",
        "RUN pip3 install tensorflow-datasets\n",
        "\n",
        "# Install TPU Tensorflow and dependencies.\n",
        "# libtpu.so must be under the '/lib' directory.\n",
        "RUN wget https://storage.googleapis.com/cloud-tpu-tpuvm-artifacts/libtpu/20210525/libtpu.so -O /lib/libtpu.so\n",
        "RUN chmod 777 /lib/libtpu.so\n",
        "\n",
        "RUN wget https://storage.googleapis.com/cloud-tpu-tpuvm-artifacts/tensorflow/20210525/tf_nightly-2.6.0-cp38-cp38-linux_x86_64.whl\n",
        "RUN pip3 install tf_nightly-2.6.0-cp38-cp38-linux_x86_64.whl\n",
        "RUN rm tf_nightly-2.6.0-cp38-cp38-linux_x86_64.whl\n",
        "# Sets up the entry point to invoke the trainer.\n",
        "ENTRYPOINT [\"python\", \"-m\", \"trainer.task\"]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "docker_push:tpu"
      },
      "source": [
        "### Build and push the Docker image to the Artifact Registry"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "J_d_zEXUulvc"
      },
      "outputs": [],
      "source": [
        "TRAIN_IMAGE = \"gcr.io/\" + PROJECT_ID + \"/tpu-train:latest\"\n",
        "\n",
        "os.chdir(\"custom\")\n",
        "! docker build --quiet --tag={TRAIN_IMAGE} .\n",
        "! docker push {TRAIN_IMAGE}\n",
        "os.chdir(\"..\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "worker_pool_tpu"
      },
      "source": [
        "### TPU worker specification pool\n",
        "\n",
        "Next, you create the worker specification pool. For TPUs, you do:\n",
        "\n",
        "- Create only one worker pool (Primary).\n",
        "- Set the machine type to `cloud-tpu`.\n",
        "- Set the accelerator type to a `TPU`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "d514eU7lulvc"
      },
      "outputs": [],
      "source": [
        "# Use TPU Accelerators. Temporarily using numeric codes, until types are added to the SDK\n",
        "#   6 = TPU_V2\n",
        "#   7 = TPU_V3\n",
        "TRAIN_TPU, TRAIN_NTPU = (7, 8)\n",
        "TRAIN_COMPUTE = \"cloud-tpu\"\n",
        "\n",
        "\n",
        "if not TRAIN_NTPU or TRAIN_NTPU < 2:\n",
        "    TRAIN_STRATEGY = \"single\"\n",
        "else:\n",
        "    TRAIN_STRATEGY = \"tpu\"\n",
        "print(TRAIN_STRATEGY)\n",
        "\n",
        "EPOCHS = 20\n",
        "STEPS = 10000\n",
        "\n",
        "TRAINER_ARGS = [\n",
        "    \"--epochs=\" + str(EPOCHS),\n",
        "    \"--steps=\" + str(STEPS),\n",
        "    \"--distribute=\" + TRAIN_STRATEGY,\n",
        "]\n",
        "\n",
        "\n",
        "WORKER_POOL_SPECS = [\n",
        "    {\n",
        "        \"container_spec\": {\n",
        "            \"args\": TRAINER_ARGS,\n",
        "            \"image_uri\": TRAIN_IMAGE,\n",
        "        },\n",
        "        \"replica_count\": 1,\n",
        "        \"machine_spec\": {\n",
        "            \"machine_type\": TRAIN_COMPUTE,\n",
        "            \"accelerator_type\": TRAIN_TPU,\n",
        "            \"accelerator_count\": TRAIN_NTPU,\n",
        "        },\n",
        "    }\n",
        "]\n",
        "\n",
        "print(WORKER_POOL_SPECS[0])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "RruSqNfrulvc"
      },
      "source": [
        "### Create CustomJob with worker pool specifications\n",
        "\n",
        "Next, you create a `CustomJob` for the multi-worker distributed training job:\n",
        "\n",
        "-`display_name`: The display name for the custom job.\n",
        "\n",
        "-`worker_pool_specs`: The detailed specifications for each worker pool."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "2QvSqbbHulvc"
      },
      "outputs": [],
      "source": [
        "DISPLAY_NAME = \"boston_\" + UUID\n",
        "\n",
        "job = aip.CustomJob(display_name=DISPLAY_NAME, worker_pool_specs=WORKER_POOL_SPECS)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Iw4L3UIfulvd"
      },
      "source": [
        "### Run the CustomJob\n",
        "\n",
        "Next, you run the custom job."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "zmqCNS78ulvd"
      },
      "outputs": [],
      "source": [
        "try:\n",
        "    job.run(sync=True)\n",
        "except Exception as e:\n",
        "    # may fail in multi-worker to find startup script\n",
        "    print(e)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gWZoH9QKulvd"
      },
      "source": [
        "### Delete a custom training job\n",
        "\n",
        "After a training job is completed, you can delete the training job with the method `delete()`.  Prior to completion, a training job can be canceled with the method `cancel()`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Lt8BJ4iBulvd"
      },
      "outputs": [],
      "source": [
        "job.delete()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cleanup"
      },
      "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:\n",
        "\n",
        "\n",
        "- Cloud Storage Bucket"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "U98Wzc01ulvd"
      },
      "outputs": [],
      "source": [
        "# Set this to true only if you'd like to delete your bucket\n",
        "delete_bucket = False\n",
        "\n",
        "if delete_bucket or os.getenv(\"IS_TESTING\"):\n",
        "    ! gsutil rm -r $BUCKET_URI"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "name": "get_started_with_vertex_distributed_training.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}