HorovodPytorch/01_TrainPyTorchModel.ipynb

{
    "cells": [
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": "# Train PyTorch Model Distributed on Batch AI\nIn this notebook we will train a PyTorch model ([ResNet50](https://arxiv.org/abs/1512.03385)) in a distributed fashion using [Horovod](https://github.com/uber/horovod) on the Imagenet dataset. This tutorial will take you through the following steps:\n * [Create Experiment](#experiment)\n * [Upload Training Scripts](#training_scripts)\n * [Submit and Monitor Job](#job)\n * [Clean Up Resources](#clean_up)"
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": "import sys\nsys.path.append(\"../common\") \n\nimport json\nfrom dotenv import get_key\nimport os\nfrom utils import write_json_to_file, dotenv_for"
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": "Set the USE_FAKE to True if you want to use fake data rather than the Imagenet dataset. This is often a good way to debug your models as well as checking what IO overhead is."
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {
                "tags": [
                    "parameters"
                ]
            },
            "outputs": [],
            "source": "# Variables for Batch AI - change as necessary\ndotenv_path = dotenv_for()\nGROUP_NAME             = get_key(dotenv_path, 'GROUP_NAME')\nFILE_SHARE_NAME        = get_key(dotenv_path, 'FILE_SHARE_NAME')\nWORKSPACE              = get_key(dotenv_path, 'WORKSPACE')\nNUM_NODES              = int(get_key(dotenv_path, 'NUM_NODES'))\nCLUSTER_NAME           = get_key(dotenv_path, 'CLUSTER_NAME')\nGPU_TYPE               = get_key(dotenv_path, 'GPU_TYPE')\nPROCESSES_PER_NODE     = int(get_key(dotenv_path, 'PROCESSES_PER_NODE'))\nSTORAGE_ACCOUNT_NAME   = get_key(dotenv_path, 'STORAGE_ACCOUNT_NAME')\n\nEXPERIMENT             = f\"distributed_pytorch_{GPU_TYPE}\"\nUSE_FAKE               = False\nDOCKERHUB              = os.getenv('DOCKER_REPOSITORY', \"masalvar\")  #\"<YOUR DOCKERHUB>\""
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": "FAKE='-x FAKE=True' if USE_FAKE else ''\nTOTAL_PROCESSES = PROCESSES_PER_NODE * NUM_NODES"
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": "<a id='experiment'></a>\n# Create Experiment\nNext we create our experiment."
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": "!az batchai experiment create -n $EXPERIMENT -g $GROUP_NAME -w $WORKSPACE"
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": "<a id='training_scripts'></a>\n# Upload Training Scripts\nWe need to upload our training scripts and associated files"
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": "json_data = !az storage account keys list -n $STORAGE_ACCOUNT_NAME -g $GROUP_NAME\nstorage_account_key = json.loads(''.join([i for i in json_data if 'WARNING' not in i]))[0]['value']"
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {
                "tags": [
                    "stripout"
                ]
            },
            "outputs": [],
            "source": "%env AZURE_STORAGE_ACCOUNT $STORAGE_ACCOUNT_NAME\n%env AZURE_STORAGE_KEY=$storage_account_key"
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": "Upload our training scripts"
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": "!az storage file upload --share-name $FILE_SHARE_NAME --source src/imagenet_pytorch_horovod.py --path scripts\n!az storage file upload --share-name $FILE_SHARE_NAME --source ../common/timer.py --path scripts"
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": "Let's check our cluster we created earlier"
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": "!az batchai cluster list -w $WORKSPACE -o table"
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": "<a id='job'></a>\n## Submit and Monitor Job\nBelow we specify the job we wish to execute.  "
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": "jobs_dict = {\n  \"$schema\": \"https://raw.githubusercontent.com/Azure/BatchAI/master/schemas/2017-09-01-preview/job.json\",\n  \"properties\": {\n    \"nodeCount\": NUM_NODES,\n    \"customToolkitSettings\": {\n      \"commandLine\": f\"echo $AZ_BATCH_HOST_LIST; \\\n    cat $AZ_BATCHAI_MPI_HOST_FILE; \\\n    mpirun -np {TOTAL_PROCESSES} --hostfile $AZ_BATCHAI_MPI_HOST_FILE \\\n    -bind-to none -map-by slot \\\n    -x NCCL_DEBUG=INFO -x LD_LIBRARY_PATH \\\n    -mca btl_tcp_if_include eth0 \\\n    -x NCCL_SOCKET_IFNAME=eth0 \\\n    -mca btl ^openib \\\n    -x NCCL_IB_DISABLE=1 \\\n    -x DISTRIBUTED=True \\\n    -x AZ_BATCHAI_INPUT_TRAIN \\\n    -x AZ_BATCHAI_INPUT_TEST \\\n    --allow-run-as-root \\\n      {FAKE} \\\n      python -u $AZ_BATCHAI_INPUT_SCRIPTS/imagenet_pytorch_horovod.py\"\n    },\n    \"stdOutErrPathPrefix\": \"$AZ_BATCHAI_MOUNT_ROOT/extfs\",\n    \"inputDirectories\": [{\n        \"id\": \"SCRIPTS\",\n        \"path\": \"$AZ_BATCHAI_MOUNT_ROOT/extfs/scripts\"\n      },\n      {\n        \"id\": \"TRAIN\",\n        \"path\": \"$AZ_BATCHAI_MOUNT_ROOT/nfs/imagenet\",\n      },\n      {\n        \"id\": \"TEST\",\n        \"path\": \"$AZ_BATCHAI_MOUNT_ROOT/nfs/imagenet\",\n      },\n    ],\n    \"outputDirectories\": [{\n        \"id\": \"MODEL\",\n        \"pathPrefix\": \"$AZ_BATCHAI_MOUNT_ROOT/extfs\",\n        \"pathSuffix\": \"Models\"\n    }],\n    \"containerSettings\": {\n      \"imageSourceRegistry\": {\n        \"image\": f\"{DOCKERHUB}/caia-horovod-pytorch\"\n      }\n    }\n  }\n}"
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": "write_json_to_file(jobs_dict, 'job.json')"
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": "JOB_NAME='pytorch-horovod-{}'.format(NUM_NODES*PROCESSES_PER_NODE)"
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": "We now submit the job to Batch AI"
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {
                "tags": [
                    "stripout"
                ]
            },
            "outputs": [],
            "source": "!az batchai job create -n $JOB_NAME --cluster $CLUSTER_NAME -w $WORKSPACE -e $EXPERIMENT -f job.json"
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": "With the command below we can check the status of the job"
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": "!az batchai job list -w $WORKSPACE -e $EXPERIMENT -o table"
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": "To view the files that the job has generated use the command below"
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {
                "tags": [
                    "stripout"
                ]
            },
            "outputs": [],
            "source": "!az batchai job file list -w $WORKSPACE -e $EXPERIMENT --j $JOB_NAME --output-directory-id stdouterr"
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": "We are also able to stream the stdout and stderr that our job produces. This is great to check the progress of our job as well as debug issues."
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {
                "tags": [
                    "stripout"
                ]
            },
            "outputs": [],
            "source": "!az batchai job file stream -w $WORKSPACE -e $EXPERIMENT --j $JOB_NAME --output-directory-id stdouterr -f stdout.txt"
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {
                "tags": [
                    "stripout"
                ]
            },
            "outputs": [],
            "source": "!az batchai job file stream -w $WORKSPACE -e $EXPERIMENT --j $JOB_NAME --output-directory-id stdouterr -f stderr.txt"
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": "We can either wait for the job to complete or delete it with the command below."
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": "!az batchai job delete -w $WORKSPACE -e $EXPERIMENT --name $JOB_NAME -y"
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": "<a id='clean_up'></a>\n## Clean Up Resources\nNext we wish to tidy up the resource we created.  \nFirst we reset the default values we set earlier."
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": "!az configure --defaults group=''\n!az configure --defaults location=''"
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": " Next we delete the cluster"
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": "!az batchai cluster delete -w $WORKSPACE --name $CLUSTER_NAME -g $GROUP_NAME -y"
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": "Once the cluster is deleted you will not incur any cost for the computation but you can still retain your experiments and workspace. If you wish to delete those as well execute the commands below."
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": "!az batchai experiment delete -w $WORKSPACE --name $EXPERIMENT -g $GROUP_NAME -y"
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": "!az batchai workspace delete -n $WORKSPACE -g $GROUP_NAME -y"
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": "Finally we can delete the group and we will have deleted everything created for this tutorial."
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": "!az group delete --name $GROUP_NAME -y"
        }
    ],
    "metadata": {
        "jupytext": {
            "text_representation": {
                "extension": ".py",
                "format_name": "light",
                "format_version": "1.3",
                "jupytext_version": "0.8.6"
            }
        },
        "kernelspec": {
            "display_name": "Python 3",
            "language": "python",
            "name": "python3"
        }
    },
    "nbformat": 4,
    "nbformat_minor": 2
}