training/heterogeneous-clusters/pt.grpc.sagemaker/hetero-pytorch-mnist.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "41408d9e",
   "metadata": {},
   "source": [
    "# PyTorch's example to demonstrate Amazon SageMaker Heterogeneous Cluster for model training\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "300aa8d3",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "This notebook's CI test result for us-west-2 is as follows. CI test results in other regions can be found at the end of the notebook. \n",
    "\n",
    "![This us-west-2 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/us-west-2/training|heterogeneous-clusters|pt.grpc.sagemaker|hetero-pytorch-mnist.ipynb)\n",
    "\n",
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "84a83bc5",
   "metadata": {},
   "source": [
    "### Description\n",
    "Heterogeneous clusters enable launching training jobs that use multiple instance types in a single job. This capability can improve your training cost and speed by running different parts of the model training on the most suitable instance type. This use case typically happens in computer vision DL training, where training is bottleneck on CPU resources needed for data augmentation, leaving the expensive GPU underutilized. Heterogeneous clusters enable you to add more CPU resources to fully utilize GPUs, thus increase training speed and cost-efficiency. For more details, you can find the documentation of this feature [here](https://docs.aws.amazon.com/sagemaker/latest/dg/train-heterogeneous-cluster.html).\n",
    "\n",
    "This notebook demonstrates how to use Heterogeneous Cluster feature of SageMaker Training with PyTorch.  The notebook works on Python 3 (_PyTorch 1.12 Python 3.8 CPU Optimized_) image of SageMaker Studio Notebook instance, and runs on _ml.t3.medium_ instance type.\n",
    "\n",
    "The notebook covers:\n",
    "- Setting up SageMaker Studio Notebook \n",
    "- Setting up the Training environment \n",
    "- Submit a Training job\n",
    "- Monitor and visualize the CloudWatch metrics\n",
    "- Comparing time-to-train and cost-to-train\n",
    "- Considerations\n",
    "- Conclusion \n",
    "\n",
    "In this sample notebook, we have taken the PyTorch model based on this [official MNIST example](https://github.com/pytorch/examples/tree/main/MNIST). We modified the training code to be heavy on data pre-processing. We are going to train this model in both Homogeneous and Heterogeneous Cluster modes. The flag to train on any of these modes can be set using `IS_HETERO = False or True` in section **B.2 Configure environment variables**. \n",
    "\n",
    "Homogeneous Training Job - In this baseline we observe an ml.p3.2xlarge with an under-utilized GPU due to a CPU bottleneck. \n",
    "<img src=images/homogeneous-cluster-diagram.png alt=\"homogeneous-training job\" />  \n",
    "\n",
    "Heterogeneous Training Job - Where we add ml.c5.9xlarge instance for extra CPU cores, to allow increased GPU usage of ml.p3.2xlarge instance, and improve cost-efficiency. Both the jobs runs the training code, train data set, pre-processing, and other relevant parameters.\n",
    "<img src=images/heterogeneous-cluster-diagram.png alt=\"heterogeneous-training job\" />\n",
    "\n",
    "In homogeneous cluster training job, the data pre-processing and Deep Neural Network (DNN) training code runs on the same instance. However, in heterogeneous cluster training job, the data pre-processing code runs on the CPU nodes (here by referred as **data_group or data group**), whereas the Deep Neural Network (DNN) training code runs on the GPU nodes (here referred as **dnn_group or dnn group**). The inter-node communication between the data and dnn groups is handled by generic implementation of [gRPC client-server interface](https://grpc.io/docs/languages/python/basics/).  \n",
    "\n",
    "The script (`launcher.py`) has the logic to detect (using SageMaker environment variables) whether the node it is running on belongs to data_group or dnn_group. If it is data_group, it spawns a separate process by executing `train_data.py`. This script runs grpc-server service for extracting processed training batches using [Protocol Buffers](https://developers.google.com/protocol-buffers/docs/overview). The gRPC server running on the data_group listens on a specific port (ex. 6000). In the code (`train_data.py`) documentation, we have chosen an implementation that keeps the data loading logic intact  where data batches are entered into a shared queue. The `get_samples` function of the `DataFeedService` pulls batches from the same queue and sends them to the client in the form of a continuous data stream. While fetching the data, the main entrypoint script `launcher.py` listens on port 16000 for a shutdown request coming from gRPC client i.e. data group. The `train_data.py` waits for shutdown action from the parent process. \n",
    "\n",
    "If the node belongs to dnn_group, the main training script (`launcher.py`) spawns a separate set of processes by executing `train_dnn.py`. The script runs gRPC client code and DNN component of the training job. It consumes the processed training data from the gRPC server. We have defined an iterable PyTorch dataset, RemoteDataset, that opens a connection to the gRPC server, and reads from a stream of data batches. Once the model is trained with all the batches of training data, the gRPC client exits, and the parent process`launcher.py` sends a shutdown request on port 16000. This indicates the gRPC server to shutdown, and signals ends of the training job. \n",
    "\n",
    "Here is how the workflow looks like:\n",
    "\n",
    "<img src=images/pytorch-heterogeneous-workflow.png width=600px>\n",
    "\n",
    "This example notebook runs a training job on 2 instances, 1 in each node group. The data_group uses ml.c5.9xlarge whereas dnn_group uses ml.p3.2xlarge.\n",
    "\n",
    "This notebook refers following files and folders:\n",
    "\n",
    "- Folders: \n",
    "  - `code`: this has the training (data pre-processing and dnn) scripts, and grpc client-server start and shutdown scripts\n",
    "  - `images`: contains images referred in notebook\n",
    "- Files: \n",
    "  - `launcher.py`: entry point training script. This script is executed on all the nodes irrespective of which group it belongs to. This is a parent process that makes a decision on where to spawn a data pre-processing or dnn component of the training job. The script runs on all the nodes as entry point. It also handles the shutdown logic for gRPC server. \n",
    "  - `train_data.py`, `dataset_feed_pb2.py`, `dataset_feed_pb2_grpc.py`: these scripts run on the data_group nodes and responsible for setting up grpc-server, start and shutdown.\n",
    "  - `train_dnn.py`: this script runs dnn code on the training data set. It fetches preprocessed data from the data_group node as a stream using gRPC client-server communication. It also sends a shutdown request after all the iterations on the preprocessed training data set. \n",
    "  - `requirement.txt`: defines package required for gRPC \n",
    "  - `train.py`: this script is the entry point script for SageMaker homogeneous cluster training. This script is picked up when you choose IS_HETERO = False. This uses a local dataset and runs both data pre-processing and a dnn component on the same node. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cb913740",
   "metadata": {},
   "source": [
    "### Security groups update if running in private VPC\n",
    "This section is relevant if you plan to [run in a private VPC](https://docs.aws.amazon.com/sagemaker/latest/dg/train-vpc.html) (passing `subnets` and `security_group_ids` parameters when defining an Estimator).  \n",
    "SageMaker documentation recommends you [add](https://docs.aws.amazon.com/sagemaker/latest/dg/train-vpc.html#train-vpc-vpc) a rule for your security group that allows inbound connections between members of the same security group, for all TCP communication. This will also cover for the gRPC related traffic between instances:\n",
    "- the data_group instances will listen on port 6000 for connections from all nodes. This stream is not encrypted. You can change the code to encrypted the connection if needed.\n",
    "- the data_group intances listen on port 16000 for a shutdown signal from all nodes."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d698424a",
   "metadata": {},
   "source": [
    "### A. Setting up SageMaker Studio notebook\n",
    "\n",
    "#### Step 1 - Upgrade SageMaker SDK and dependent packages \n",
    "Heterogeneous Clusters for Amazon SageMaker model training was [announced](https://aws.amazon.com/about-aws/whats-new/2022/07/announcing-heterogeneous-clusters-amazon-sagemaker-model-training) on 07/08/2022. As a first step, ensure you have updated SageMaker SDK, PyTorch, and Boto3 client that enables this feature."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5d4c9ff1",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%bash\n",
    "python3 -m pip install --upgrade boto3 botocore awscli sagemaker"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "55a00d60",
   "metadata": {},
   "source": [
    "#### Step 2 - Restart the notebook kernel "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1d664b18",
   "metadata": {},
   "outputs": [],
   "source": [
    "# import IPython\n",
    "# IPython.Application.instance().kernel.do_shutdown(True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "afda813d",
   "metadata": {},
   "source": [
    "#### Step 3 - Validate SageMaker Python SDK and PyTorch versions\n",
    "Ensure the output of the cell below reflects:\n",
    "\n",
    "- SageMaker Python SDK version 2.98.0 or above, \n",
    "- boto3 1.24 or above \n",
    "- botocore 1.27 or above \n",
    "- PyTorch 1.10 or above "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c2af1c49",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip show sagemaker torch boto3 botocore |egrep 'Name|Version|---'"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5e2742a5",
   "metadata": {},
   "source": [
    "--------------\n",
    "### B. Setting up the Training environment\n",
    "\n",
    "#### Step 1 - Import SageMaker components and set up the IAM role and Amazon S3 bucket"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f159a269",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import json\n",
    "import datetime\n",
    "import os\n",
    "\n",
    "import sagemaker\n",
    "from sagemaker.pytorch import PyTorch\n",
    "from sagemaker import get_execution_role\n",
    "from sagemaker.instance_group import InstanceGroup\n",
    "\n",
    "\n",
    "sess = sagemaker.Session()\n",
    "\n",
    "role = get_execution_role()\n",
    "\n",
    "output_path = \"s3://\" + sess.default_bucket() + \"/DEMO-MNIST\"\n",
    "print(role)\n",
    "print(output_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1f7216a6",
   "metadata": {},
   "source": [
    "#### Step 2 - Configure environment variables \n",
    "This step defines whether you want to run training job in heterogeneous cluster mode or not. Also, defines instance groups, multiple nodes in group, and hyperparameter values. For baselining, run a homogeneous cluster training job by setting `IS_HETERO = False`. This will let both the data pre-processing and DNN code run on the same node i.e. `ml.p3.2xlarge`. \n",
    "\n",
    "\n",
    "Test configuration (if running training on p3.2xl or g5.2xl as dnn_group instance type, and c5.2xl as data_group instance type: (training duration: 7-8 mins)  \n",
    "`num-data-workers: 4`  \n",
    "`grpc-workers: 4`   \n",
    "`num-dnn-workers: 4`  \n",
    "`pin-memory\": True`   \n",
    "`iterations : 100`   \n",
    "\n",
    "Performance configuration (if running training on p3.2xl as dnn_group instance type, and c5.9xl as data_group instance type OR training in homogeneous cluster mode i.e. g5.8xl): (training duration - 30 mins)  \n",
    "`num-data-workers: 32`  \n",
    "`grpc-workers: 2`   \n",
    "`num-dnn-workers: 2`  \n",
    "`pin-memory\": True`   \n",
    "`iterations : 4800`\n",
    "\n",
    "Performance configuration (if running training on p3.2xl in homogeneous cluster mode):   \n",
    "`num-data-workers: 8`  \n",
    "`grpc-workers: 2`   \n",
    "`num-dnn-workers: 2`  \n",
    "`pin-memory\": True`   \n",
    "`iterations : 2400`\n",
    "\n",
    "Note: This PyTorch example has not been tested with multiple instances in an instance group. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ec14fd57",
   "metadata": {},
   "outputs": [],
   "source": [
    "IS_CLOUD_JOB = True\n",
    "IS_HETERO = True  # if set to false, uses homogeneous cluster\n",
    "PT_DATA_MODE = \"service\" if IS_HETERO else \"local\"  # local | service\n",
    "IS_DNN_DISTRIBUTION = False  # Distributed Training with DNN nodes not tested, set it to False\n",
    "\n",
    "data_group = InstanceGroup(\n",
    "    \"data_group\", \"ml.c5.9xlarge\", 1\n",
    ")  # 36 vCPU #change the instance type if IS_HETERO=True\n",
    "dnn_group = InstanceGroup(\n",
    "    \"dnn_group\", \"ml.p3.2xlarge\", 1\n",
    ")  # 8 vCPU #change the instance type if IS_HETERO=True\n",
    "\n",
    "kwargs = dict()\n",
    "kwargs[\"hyperparameters\"] = {\n",
    "    \"batch-size\": 8192,\n",
    "    \"num-data-workers\": 4,  # This number drives the avg. step time. More workers help parallel pre-processing of data. Recommendation: Total no. of cpu 'n' = 'num-data-wokers'+'grpc-workers'+ 2 (reserved)\n",
    "    \"grpc-workers\": 4,  # No. of workers serving pre-processed data to DNN group (gRPC client). see above formula.\n",
    "    \"num-dnn-workers\": 4,  # Modify this no. to be less than the cpu core of your training instances in dnn group\n",
    "    \"pin-memory\": True,  # Pin to GPU memory\n",
    "    \"iterations\": 100,  # No. of iterations in an epoch (must be multiple of 10).\n",
    "    \"region\": sess.boto_region_name,\n",
    "}\n",
    "\n",
    "if IS_HETERO:\n",
    "    kwargs[\"instance_groups\"] = [data_group, dnn_group]\n",
    "    entry_point = \"launcher.py\"\n",
    "else:\n",
    "    kwargs[\"instance_type\"] = (\n",
    "        \"ml.p3.2xlarge\" if IS_CLOUD_JOB else \"local\"\n",
    "    )  # change the instance type if IS_HETERO=False\n",
    "    kwargs[\"instance_count\"] = 1\n",
    "    entry_point = \"train.py\"\n",
    "\n",
    "if IS_DNN_DISTRIBUTION:\n",
    "    processes_per_host_dict = {\n",
    "        \"ml.g5.xlarge\": 1,\n",
    "        \"ml.g5.12xlarge\": 4,\n",
    "        \"ml.p3.8xlarge\": 4,\n",
    "        \"ml.p4d.24xlarge\": 8,\n",
    "    }\n",
    "    kwargs[\"distribution\"] = {\n",
    "        \"mpi\": {\n",
    "            \"enabled\": True,\n",
    "            \"processes_per_host\": processes_per_host_dict[dnn_instance_type],\n",
    "            \"custom_mpi_options\": \"--NCCL_DEBUG INFO\",\n",
    "        },\n",
    "    }\n",
    "    if IS_HETERO:\n",
    "        kwargs[\"distribution\"][\"instance_groups\"] = [dnn_group]\n",
    "\n",
    "    print(f\"distribution={kwargs['distribution']}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "08bb18c8",
   "metadata": {},
   "source": [
    "#### Step 3: Set up the Estimator\n",
    "In order to use SageMaker to fit our algorithm, we'll create `Estimator` that defines how to use the container to train. This includes the configuration we need to invoke SageMaker training."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "76b1d36d",
   "metadata": {},
   "outputs": [],
   "source": [
    "estimator = PyTorch(\n",
    "    framework_version=\"1.11.0\",  # 1.10.0 or later\n",
    "    py_version=\"py38\",  # Python v3.8\n",
    "    role=role,\n",
    "    entry_point=entry_point,\n",
    "    source_dir=\"code\",\n",
    "    volume_size=10,\n",
    "    max_run=4800,\n",
    "    disable_profiler=True,\n",
    "    debugger_hook_config=False,\n",
    "    **kwargs,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fedab49f",
   "metadata": {},
   "source": [
    "#### Step 4: Download the MNIST Data and Upload it to S3 bucket\n",
    "\n",
    "This is an optional step for now. The training job downloads the data on its run directly from MNIST website to the data_group node (grpc server). "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5781cd67",
   "metadata": {},
   "outputs": [],
   "source": [
    "import logging\n",
    "import boto3\n",
    "from botocore.exceptions import ClientError\n",
    "\n",
    "# Download training and testing data from a public S3 bucket\n",
    "\n",
    "\n",
    "def download_from_s3(data_dir=\"./data\", train=True):\n",
    "    \"\"\"Download MNIST dataset and convert it to numpy array\n",
    "\n",
    "    Args:\n",
    "        data_dir (str): directory to save the data\n",
    "        train (bool): download training set\n",
    "\n",
    "    Returns:\n",
    "        None\n",
    "    \"\"\"\n",
    "\n",
    "    if not os.path.exists(data_dir):\n",
    "        os.makedirs(data_dir)\n",
    "\n",
    "    if train:\n",
    "        images_file = \"train-images-idx3-ubyte.gz\"\n",
    "        labels_file = \"train-labels-idx1-ubyte.gz\"\n",
    "    else:\n",
    "        images_file = \"t10k-images-idx3-ubyte.gz\"\n",
    "        labels_file = \"t10k-labels-idx1-ubyte.gz\"\n",
    "\n",
    "    # download objects\n",
    "    s3 = boto3.client(\"s3\")\n",
    "    bucket = f\"sagemaker-example-files-prod-{sess.boto_region_name}\"\n",
    "    for obj in [images_file, labels_file]:\n",
    "        key = os.path.join(\"datasets/image/MNIST\", obj)\n",
    "        dest = os.path.join(data_dir, obj)\n",
    "        if not os.path.exists(dest):\n",
    "            s3.download_file(bucket, key, dest)\n",
    "    return\n",
    "\n",
    "\n",
    "download_from_s3(\"./data\", True)\n",
    "download_from_s3(\"./data\", False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8d1e1880",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Upload to the default bucket\n",
    "\n",
    "prefix = \"DEMO-MNIST\"\n",
    "bucket = sess.default_bucket()\n",
    "loc = sess.upload_data(path=\"./data\", bucket=bucket, key_prefix=prefix)\n",
    "\n",
    "channels = {\"training\": loc, \"testing\": loc}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c0b9f285",
   "metadata": {},
   "source": [
    "## C. Submit the training job\n",
    "\n",
    "The job runs for the predefined iterations. DNN instance group sends a shutdown request to data group after done with the training. You can see the following entries in the CloudWatch logs of dnn instance. A job with 4800 iterations finishes in 29 mins in a Heterogeneous cluster composed of 1x ml.c5.9xlarge as data node and 1x ml.p3.2xlarge as DNN node.\n",
    "\n",
    "Log excerpt from algo-1 (DNN instance)\n",
    "```\n",
    "4780: avg step time: 0.19709917231025106\n",
    "INFO:__main__:4780: avg step time: 0.19709917231025106\n",
    "4790: avg step time: 0.19694106239373696\n",
    "INFO:__main__:4790: avg step time: 0.19694106239373696\n",
    "4800: avg step time: 0.196784295383125\n",
    "Saving the model\n",
    "INFO:__main__:4800: avg step time: 0.196784295383125\n",
    "INFO:__main__:Saving the model\n",
    "Training job completed!\n",
    "INFO:__main__:Training job completed!\n",
    "Process train_dnn.py closed with returncode=0\n",
    "Shutting downdata service dispatcher via: [algo-2:16000]\n",
    "shutdown request sent to algo-2:16000\n",
    "2022-08-16 01:15:05,555 sagemaker-training-toolkit INFO     Waiting for the process to finish and give a return code.\n",
    "2022-08-16 01:15:05,555 sagemaker-training-toolkit INFO     Done waiting for a return code. Received 0 from exiting process.\n",
    "2022-08-16 01:15:05,556 sagemaker-training-toolkit INFO     Reporting training SUCCESS\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "607143e2",
   "metadata": {},
   "outputs": [],
   "source": [
    "estimator.fit(\n",
    "    inputs=channels,\n",
    "    job_name=\"pt-hetero\"\n",
    "    + \"-\"\n",
    "    + \"H-\"\n",
    "    + str(IS_HETERO)[0]\n",
    "    + \"-\"\n",
    "    + datetime.datetime.utcnow().strftime(\"%Y%m%dT%H%M%SZ\"),\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2c866a60",
   "metadata": {},
   "source": [
    "## D. Monitoring Instance Metrics for GPU and CPU utilization\n",
    "\n",
    "Click on **View instance metrics** from the **Training jobs** node in **Amazon SageMaker Console**. In the run above, all 30 vCPU of Data node (algo-1) is approx. 100% utilized, and the GPU utilization is at 100% at frequent intervals in the DNN node (algo-2). To rescale the CloudWatch Metrics to 100% on CPU utilization for algo-1 and algo-2, use CloudWatch \"Add Math\" feature and average it out by no. of cores on those instance types.\n",
    "\n",
    "<img src=images/heterogeneous-instance-metrics.png width=900px>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a768b324",
   "metadata": {},
   "source": [
    "## E. Comparing time-to-train and cost-to-train\n",
    "\n",
    "Let's continue with the above example i.e. train a heavy data pre-processing (CPU intensive) model (MNIST) requiring only 1 GPU. We start with ml.p3.2xlarge (1xV100 GPU, 8x vCPU) in homogeneous cluster mode to get the baseline performance numbers. Due to the no. of CPU cores, we could not go beyond 8 data loader/workers for data pre-processing. The avg. step cost was `7.6 cents` and avg. step time is `1.19 seconds`. \n",
    "\n",
    "Our objective is to reduce the cost and speed up the model training time. The first choice here is to scale up the instance type in the same family. If we leverage the next instance type (4 GPU) in the P3 family, the GPUs would have gone underutilized. In this case, we needed more vCPU to GPU ratio. Assuming we haven't had any instance type in another instance family or the model is incompatible with the CPU/GPU architectures of other instance families, we are constrained to use ml.p3.2xlarge. The only way then to have more vCPUs to GPU ratio is to use SageMaker feature, Heterogeneous Cluster, which enables customers to offload data pre-processing logic to CPU only instance types example ml.c5. In the next test, we offloaded CPU intensive work i.e. data preprocessing to ml.c5.9xlarge (36 vCPU) and continued using ml.p3.2xlarge for DNN. The avg. step cost was `1.9 cents` and avg. step time is `0.18 seconds`. \n",
    "\n",
    "In summary, we reduced the training cost by 4.75 times, and the avg. step reduced by 6.5 times. This was possible because with higher cpu count, we could use 32 data loader workers (compared to 8 with p3.2xl) to preprocess the data, and kept GPU close to 100% utilized at frequent intervals. Note: These numbers are just taken as a sample, you have to do benchmarking with your own model and dataset to come up with the exact price-performance benefits. \n",
    "\n",
    "## F. Considerations\n",
    "This PyTorch example implementation of gRPC client-server supports a single instance in each instance group. You may like to extend the logic to support multiple instances in the instance group. \n",
    "\n",
    "## G. Conclusion\n",
    "In this notebook, we demonstrated how to leverage heterogeneous cluster feature of SageMaker Training to achieve better price performance. To get started you can copy this example project, and only change the `train_dnn.py` script.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a5536db5",
   "metadata": {},
   "source": [
    "## Notebook CI Test Results\n",
    "\n",
    "This notebook was tested in multiple regions. The test results are as follows, except for us-west-2 which is shown at the top of the notebook.\n",
    "\n",
    "![This us-east-1 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/us-east-1/training|heterogeneous-clusters|pt.grpc.sagemaker|hetero-pytorch-mnist.ipynb)\n",
    "\n",
    "![This us-east-2 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/us-east-2/training|heterogeneous-clusters|pt.grpc.sagemaker|hetero-pytorch-mnist.ipynb)\n",
    "\n",
    "![This us-west-1 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/us-west-1/training|heterogeneous-clusters|pt.grpc.sagemaker|hetero-pytorch-mnist.ipynb)\n",
    "\n",
    "![This ca-central-1 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/ca-central-1/training|heterogeneous-clusters|pt.grpc.sagemaker|hetero-pytorch-mnist.ipynb)\n",
    "\n",
    "![This sa-east-1 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/sa-east-1/training|heterogeneous-clusters|pt.grpc.sagemaker|hetero-pytorch-mnist.ipynb)\n",
    "\n",
    "![This eu-west-1 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/eu-west-1/training|heterogeneous-clusters|pt.grpc.sagemaker|hetero-pytorch-mnist.ipynb)\n",
    "\n",
    "![This eu-west-2 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/eu-west-2/training|heterogeneous-clusters|pt.grpc.sagemaker|hetero-pytorch-mnist.ipynb)\n",
    "\n",
    "![This eu-west-3 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/eu-west-3/training|heterogeneous-clusters|pt.grpc.sagemaker|hetero-pytorch-mnist.ipynb)\n",
    "\n",
    "![This eu-central-1 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/eu-central-1/training|heterogeneous-clusters|pt.grpc.sagemaker|hetero-pytorch-mnist.ipynb)\n",
    "\n",
    "![This eu-north-1 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/eu-north-1/training|heterogeneous-clusters|pt.grpc.sagemaker|hetero-pytorch-mnist.ipynb)\n",
    "\n",
    "![This ap-southeast-1 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/ap-southeast-1/training|heterogeneous-clusters|pt.grpc.sagemaker|hetero-pytorch-mnist.ipynb)\n",
    "\n",
    "![This ap-southeast-2 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/ap-southeast-2/training|heterogeneous-clusters|pt.grpc.sagemaker|hetero-pytorch-mnist.ipynb)\n",
    "\n",
    "![This ap-northeast-1 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/ap-northeast-1/training|heterogeneous-clusters|pt.grpc.sagemaker|hetero-pytorch-mnist.ipynb)\n",
    "\n",
    "![This ap-northeast-2 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/ap-northeast-2/training|heterogeneous-clusters|pt.grpc.sagemaker|hetero-pytorch-mnist.ipynb)\n",
    "\n",
    "![This ap-south-1 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/ap-south-1/training|heterogeneous-clusters|pt.grpc.sagemaker|hetero-pytorch-mnist.ipynb)\n"
   ]
  }
 ],
 "metadata": {
  "availableInstances": [
   {
    "_defaultOrder": 0,
    "_isFastLaunch": true,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 4,
    "name": "ml.t3.medium",
    "vcpuNum": 2
   },
   {
    "_defaultOrder": 1,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 8,
    "name": "ml.t3.large",
    "vcpuNum": 2
   },
   {
    "_defaultOrder": 2,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 16,
    "name": "ml.t3.xlarge",
    "vcpuNum": 4
   },
   {
    "_defaultOrder": 3,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 32,
    "name": "ml.t3.2xlarge",
    "vcpuNum": 8
   },
   {
    "_defaultOrder": 4,
    "_isFastLaunch": true,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 8,
    "name": "ml.m5.large",
    "vcpuNum": 2
   },
   {
    "_defaultOrder": 5,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 16,
    "name": "ml.m5.xlarge",
    "vcpuNum": 4
   },
   {
    "_defaultOrder": 6,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 32,
    "name": "ml.m5.2xlarge",
    "vcpuNum": 8
   },
   {
    "_defaultOrder": 7,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 64,
    "name": "ml.m5.4xlarge",
    "vcpuNum": 16
   },
   {
    "_defaultOrder": 8,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 128,
    "name": "ml.m5.8xlarge",
    "vcpuNum": 32
   },
   {
    "_defaultOrder": 9,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 192,
    "name": "ml.m5.12xlarge",
    "vcpuNum": 48
   },
   {
    "_defaultOrder": 10,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 256,
    "name": "ml.m5.16xlarge",
    "vcpuNum": 64
   },
   {
    "_defaultOrder": 11,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 384,
    "name": "ml.m5.24xlarge",
    "vcpuNum": 96
   },
   {
    "_defaultOrder": 12,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 8,
    "name": "ml.m5d.large",
    "vcpuNum": 2
   },
   {
    "_defaultOrder": 13,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 16,
    "name": "ml.m5d.xlarge",
    "vcpuNum": 4
   },
   {
    "_defaultOrder": 14,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 32,
    "name": "ml.m5d.2xlarge",
    "vcpuNum": 8
   },
   {
    "_defaultOrder": 15,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 64,
    "name": "ml.m5d.4xlarge",
    "vcpuNum": 16
   },
   {
    "_defaultOrder": 16,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 128,
    "name": "ml.m5d.8xlarge",
    "vcpuNum": 32
   },
   {
    "_defaultOrder": 17,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 192,
    "name": "ml.m5d.12xlarge",
    "vcpuNum": 48
   },
   {
    "_defaultOrder": 18,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 256,
    "name": "ml.m5d.16xlarge",
    "vcpuNum": 64
   },
   {
    "_defaultOrder": 19,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 384,
    "name": "ml.m5d.24xlarge",
    "vcpuNum": 96
   },
   {
    "_defaultOrder": 20,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": true,
    "memoryGiB": 0,
    "name": "ml.geospatial.interactive",
    "supportedImageNames": [
     "sagemaker-geospatial-v1-0"
    ],
    "vcpuNum": 0
   },
   {
    "_defaultOrder": 21,
    "_isFastLaunch": true,
    "category": "Compute optimized",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 4,
    "name": "ml.c5.large",
    "vcpuNum": 2
   },
   {
    "_defaultOrder": 22,
    "_isFastLaunch": false,
    "category": "Compute optimized",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 8,
    "name": "ml.c5.xlarge",
    "vcpuNum": 4
   },
   {
    "_defaultOrder": 23,
    "_isFastLaunch": false,
    "category": "Compute optimized",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 16,
    "name": "ml.c5.2xlarge",
    "vcpuNum": 8
   },
   {
    "_defaultOrder": 24,
    "_isFastLaunch": false,
    "category": "Compute optimized",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 32,
    "name": "ml.c5.4xlarge",
    "vcpuNum": 16
   },
   {
    "_defaultOrder": 25,
    "_isFastLaunch": false,
    "category": "Compute optimized",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 72,
    "name": "ml.c5.9xlarge",
    "vcpuNum": 36
   },
   {
    "_defaultOrder": 26,
    "_isFastLaunch": false,
    "category": "Compute optimized",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 96,
    "name": "ml.c5.12xlarge",
    "vcpuNum": 48
   },
   {
    "_defaultOrder": 27,
    "_isFastLaunch": false,
    "category": "Compute optimized",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 144,
    "name": "ml.c5.18xlarge",
    "vcpuNum": 72
   },
   {
    "_defaultOrder": 28,
    "_isFastLaunch": false,
    "category": "Compute optimized",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 192,
    "name": "ml.c5.24xlarge",
    "vcpuNum": 96
   },
   {
    "_defaultOrder": 29,
    "_isFastLaunch": true,
    "category": "Accelerated computing",
    "gpuNum": 1,
    "hideHardwareSpecs": false,
    "memoryGiB": 16,
    "name": "ml.g4dn.xlarge",
    "vcpuNum": 4
   },
   {
    "_defaultOrder": 30,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 1,
    "hideHardwareSpecs": false,
    "memoryGiB": 32,
    "name": "ml.g4dn.2xlarge",
    "vcpuNum": 8
   },
   {
    "_defaultOrder": 31,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 1,
    "hideHardwareSpecs": false,
    "memoryGiB": 64,
    "name": "ml.g4dn.4xlarge",
    "vcpuNum": 16
   },
   {
    "_defaultOrder": 32,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 1,
    "hideHardwareSpecs": false,
    "memoryGiB": 128,
    "name": "ml.g4dn.8xlarge",
    "vcpuNum": 32
   },
   {
    "_defaultOrder": 33,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 4,
    "hideHardwareSpecs": false,
    "memoryGiB": 192,
    "name": "ml.g4dn.12xlarge",
    "vcpuNum": 48
   },
   {
    "_defaultOrder": 34,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 1,
    "hideHardwareSpecs": false,
    "memoryGiB": 256,
    "name": "ml.g4dn.16xlarge",
    "vcpuNum": 64
   },
   {
    "_defaultOrder": 35,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 1,
    "hideHardwareSpecs": false,
    "memoryGiB": 61,
    "name": "ml.p3.2xlarge",
    "vcpuNum": 8
   },
   {
    "_defaultOrder": 36,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 4,
    "hideHardwareSpecs": false,
    "memoryGiB": 244,
    "name": "ml.p3.8xlarge",
    "vcpuNum": 32
   },
   {
    "_defaultOrder": 37,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 8,
    "hideHardwareSpecs": false,
    "memoryGiB": 488,
    "name": "ml.p3.16xlarge",
    "vcpuNum": 64
   },
   {
    "_defaultOrder": 38,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 8,
    "hideHardwareSpecs": false,
    "memoryGiB": 768,
    "name": "ml.p3dn.24xlarge",
    "vcpuNum": 96
   },
   {
    "_defaultOrder": 39,
    "_isFastLaunch": false,
    "category": "Memory Optimized",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 16,
    "name": "ml.r5.large",
    "vcpuNum": 2
   },
   {
    "_defaultOrder": 40,
    "_isFastLaunch": false,
    "category": "Memory Optimized",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 32,
    "name": "ml.r5.xlarge",
    "vcpuNum": 4
   },
   {
    "_defaultOrder": 41,
    "_isFastLaunch": false,
    "category": "Memory Optimized",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 64,
    "name": "ml.r5.2xlarge",
    "vcpuNum": 8
   },
   {
    "_defaultOrder": 42,
    "_isFastLaunch": false,
    "category": "Memory Optimized",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 128,
    "name": "ml.r5.4xlarge",
    "vcpuNum": 16
   },
   {
    "_defaultOrder": 43,
    "_isFastLaunch": false,
    "category": "Memory Optimized",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 256,
    "name": "ml.r5.8xlarge",
    "vcpuNum": 32
   },
   {
    "_defaultOrder": 44,
    "_isFastLaunch": false,
    "category": "Memory Optimized",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 384,
    "name": "ml.r5.12xlarge",
    "vcpuNum": 48
   },
   {
    "_defaultOrder": 45,
    "_isFastLaunch": false,
    "category": "Memory Optimized",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 512,
    "name": "ml.r5.16xlarge",
    "vcpuNum": 64
   },
   {
    "_defaultOrder": 46,
    "_isFastLaunch": false,
    "category": "Memory Optimized",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 768,
    "name": "ml.r5.24xlarge",
    "vcpuNum": 96
   },
   {
    "_defaultOrder": 47,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 1,
    "hideHardwareSpecs": false,
    "memoryGiB": 16,
    "name": "ml.g5.xlarge",
    "vcpuNum": 4
   },
   {
    "_defaultOrder": 48,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 1,
    "hideHardwareSpecs": false,
    "memoryGiB": 32,
    "name": "ml.g5.2xlarge",
    "vcpuNum": 8
   },
   {
    "_defaultOrder": 49,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 1,
    "hideHardwareSpecs": false,
    "memoryGiB": 64,
    "name": "ml.g5.4xlarge",
    "vcpuNum": 16
   },
   {
    "_defaultOrder": 50,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 1,
    "hideHardwareSpecs": false,
    "memoryGiB": 128,
    "name": "ml.g5.8xlarge",
    "vcpuNum": 32
   },
   {
    "_defaultOrder": 51,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 1,
    "hideHardwareSpecs": false,
    "memoryGiB": 256,
    "name": "ml.g5.16xlarge",
    "vcpuNum": 64
   },
   {
    "_defaultOrder": 52,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 4,
    "hideHardwareSpecs": false,
    "memoryGiB": 192,
    "name": "ml.g5.12xlarge",
    "vcpuNum": 48
   },
   {
    "_defaultOrder": 53,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 4,
    "hideHardwareSpecs": false,
    "memoryGiB": 384,
    "name": "ml.g5.24xlarge",
    "vcpuNum": 96
   },
   {
    "_defaultOrder": 54,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 8,
    "hideHardwareSpecs": false,
    "memoryGiB": 768,
    "name": "ml.g5.48xlarge",
    "vcpuNum": 192
   },
   {
    "_defaultOrder": 55,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 8,
    "hideHardwareSpecs": false,
    "memoryGiB": 1152,
    "name": "ml.p4d.24xlarge",
    "vcpuNum": 96
   },
   {
    "_defaultOrder": 56,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 8,
    "hideHardwareSpecs": false,
    "memoryGiB": 1152,
    "name": "ml.p4de.24xlarge",
    "vcpuNum": 96
   }
  ],
  "kernelspec": {
   "display_name": "Python 3 (PyTorch 1.13 Python 3.9 CPU Optimized)",
   "language": "python",
   "name": "python3__SAGEMAKER_INTERNAL__arn:aws:sagemaker:us-west-2:236514542706:image/pytorch-1.13-cpu-py39"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.16"
  },
  "vscode": {
   "interpreter": {
    "hash": "77c0de85c2cb739aa5100af7b92fb9d2075368f0e653f4148499a56c989df5f7"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}