01_CreateResources.ipynb

{
    "cells": [
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": "# Create Azure and Batch AI Resources\nIn this notebook we will create the necessary resources to train a ResNet50 model([ResNet50](https://arxiv.org/abs/1512.03385)) in a distributed fashion using [Horovod](https://github.com/uber/horovod) on the ImageNet dataset. If you plan on using fake data then the sections marked optional can be skipped. This notebook will take you through the following steps:\n * [Create Azure Resources](#azure_resources)\n * [Create Fileserver(NFS)](#create_fileshare)\n * [Upload Data to Blob (Optional)](#upload_data)\n * [Configure Batch AI Cluster](#configure_cluster)"
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": "import sys\nsys.path.append(\"common\") \n\nfrom dotenv import set_key\nimport os\nimport json\nfrom utils import get_password, dotenv_for\nfrom pathlib import Path"
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": "Below are the variables that describe our experiment. By default we are using the NC24rs_v3 (Standard_NC24rs_v3) VMs which have V100 GPUs and Infiniband. By default we are using 2 nodes with each node having 4 GPUs, this equates to 8 GPUs. Feel free to increase the number of nodes but be aware what limitations your subscription may have.\n\nSet 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\nID                     = \"dtdemo\"\nGROUP_NAME             = f\"batch{ID}rg\"\nSTORAGE_ACCOUNT_NAME   = f\"batch{ID}st\"\nFILE_SHARE_NAME        = f\"batch{ID}share\"\nSELECTED_SUBSCRIPTION  = \"<YOUR_SUBSCRIPTION>\"\nWORKSPACE              = \"workspace\"\nNUM_NODES              = 2\nCLUSTER_NAME           = \"msv100\"\nVM_SIZE                = \"Standard_NC24rs_v3\"\nGPU_TYPE               = \"V100\"\nPROCESSES_PER_NODE     = 4\nLOCATION               = \"eastus\"\nNFS_NAME               = f\"batch{ID}nfs\"\nUSERNAME               = \"batchai_user\"\nUSE_FAKE               = False\nDOCKERHUB              = os.getenv('DOCKER_REPOSITORY', \"masalvar\")\nDATA                   = Path(\"/data\")\nCONTAINER_NAME         = f\"batch{ID}container\"\nDOCKER_PWD             = \"<YOUR_DOCKER_PWD>\"\n\ndotenv_path = dotenv_for()\nset_key(dotenv_path, 'DOCKER_PWD', DOCKER_PWD)\nset_key(dotenv_path, 'GROUP_NAME', GROUP_NAME)\nset_key(dotenv_path, 'FILE_SHARE_NAME', FILE_SHARE_NAME)\nset_key(dotenv_path, 'WORKSPACE', WORKSPACE)\nset_key(dotenv_path, 'NUM_NODES', str(NUM_NODES))\nset_key(dotenv_path, 'CLUSTER_NAME', CLUSTER_NAME)\nset_key(dotenv_path, 'GPU_TYPE', GPU_TYPE)\nset_key(dotenv_path, 'PROCESSES_PER_NODE', str(PROCESSES_PER_NODE))\nset_key(dotenv_path, 'STORAGE_ACCOUNT_NAME', STORAGE_ACCOUNT_NAME)"
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": "<a id='azure_resources'></a>\n## Create Azure Resources\nFirst we need to log in to our Azure account. "
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {
                "tags": [
                    "stripout"
                ]
            },
            "outputs": [],
            "source": "!az login -o table"
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": "If you have more than one Azure account you will need to select it with the command below. If you only have one account you can skip this step."
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": "!az account set --subscription \"$SELECTED_SUBSCRIPTION\""
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {
                "tags": [
                    "stripout"
                ]
            },
            "outputs": [],
            "source": "!az account list -o table"
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": "Next we create the group that will hold all our Azure resources."
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": "!az group create -n $GROUP_NAME -l $LOCATION -o table"
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": "We will create the storage account that will store our fileshare where all the outputs from the jobs will be stored."
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": "json_data = !az storage account create -l $LOCATION -n $STORAGE_ACCOUNT_NAME -g $GROUP_NAME --sku Standard_LRS\nprint('Storage account {} provisioning state: {}'.format(STORAGE_ACCOUNT_NAME, \n                                                         json.loads(''.join(json_data))['provisioningState']))"
        },
        {
            "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": {},
            "outputs": [],
            "source": "!az storage share create --account-name $STORAGE_ACCOUNT_NAME \\\n--account-key $storage_account_key --name $FILE_SHARE_NAME"
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": "!az storage directory create --share-name $FILE_SHARE_NAME  --name scripts \\\n--account-name $STORAGE_ACCOUNT_NAME --account-key $storage_account_key"
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": "Here we are setting some defaults so we don't have to keep adding them to every command"
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": "!az configure --defaults location=$LOCATION\n!az configure --defaults group=$GROUP_NAME"
        },
        {
            "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": "#### Create Workspace\nBatch AI has the concept of workspaces and experiments. Below we will create the workspace for our work."
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {
                "tags": [
                    "stripout"
                ]
            },
            "outputs": [],
            "source": "!az batchai workspace create -n $WORKSPACE -g $GROUP_NAME"
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": "<a id='upload_data'></a>\n## Upload Data to Blob (Optional)\nIn this section we will create a blob container and upload the imagenet data we prepared locally in the previous notebook.\n\n**You only need to run this section if you want to use real data. If USE_FAKE is set to False the commands below won't be executed.**\n"
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": "if USE_FAKE is False:\n    !az storage container create --account-name {STORAGE_ACCOUNT_NAME} \\\n                                 --account-key {storage_account_key} \\\n                                 --name {CONTAINER_NAME}"
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {
                "tags": [
                    "stripout"
                ]
            },
            "outputs": [],
            "source": "if USE_FAKE is False:\n    # Should take about 20 minutes\n    !azcopy --source {DATA/\"train.tar.gz\"} \\\n    --destination https://{STORAGE_ACCOUNT_NAME}.blob.core.windows.net/{CONTAINER_NAME}/train.tar.gz \\\n    --dest-key {storage_account_key} --quiet"
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {
                "tags": [
                    "stripout"
                ]
            },
            "outputs": [],
            "source": "if USE_FAKE is False:\n    !azcopy --source {DATA/\"validation.tar.gz\"} \\\n    --destination https://{STORAGE_ACCOUNT_NAME}.blob.core.windows.net/{CONTAINER_NAME}/validation.tar.gz \\\n    --dest-key {storage_account_key} --quiet"
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": "<a id='create_fileshare'></a>\n## Create Fileserver\nIn this example we will store the data on an NFS fileshare. It is possible to use many storage solutions with Batch AI. NFS offers the best tradeoff between performance and ease of use. The best performance is achieved by loading the data locally but this can be cumbersome since it requires that the data is download by the all the nodes which with the ImageNet dataset can take hours. If you are using fake data we won't be using the fileserver but we will create one so that if you want to run the real ImageNet data later the server is ready."
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {
                "tags": [
                    "stripout"
                ]
            },
            "outputs": [],
            "source": "!az batchai file-server create -n $NFS_NAME --disk-count 4 --disk-size 250 -w $WORKSPACE \\\n-s Standard_DS4_v2 -u $USERNAME -p {get_password(dotenv_for())} -g $GROUP_NAME --storage-sku Premium_LRS"
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": "!az batchai file-server list -o table -w $WORKSPACE -g $GROUP_NAME"
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": "json_data = !az batchai file-server list -w $WORKSPACE -g $GROUP_NAME\nnfs_ip=json.loads(''.join([i for i in json_data if 'WARNING' not in i]))[0]['mountSettings']['fileServerPublicIp']"
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": "After we have created the NFS share we need to copy the data to it. To do this we write the script below which will be executed on the fileserver. It installs a tool called azcopy and then downloads and extracts the data to the appropriate directory."
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": "nodeprep_script = f\"\"\"\n#!/usr/bin/env bash\nwget https://gist.githubusercontent.com/msalvaris/073c28a9993d58498957294d20d74202/raw/87a78275879f7c9bb8d6fb9de8a2d2996bb66c24/install_azcopy\nchmod 777 install_azcopy\nsudo ./install_azcopy\n\nmkdir -p /data/imagenet\n\nazcopy --source https://{STORAGE_ACCOUNT_NAME}.blob.core.windows.net/{CONTAINER_NAME}/validation.tar.gz \\\n        --destination  /data/imagenet/validation.tar.gz\\\n        --source-key {storage_account_key}\\\n        --quiet\n\n\nazcopy --source https://{STORAGE_ACCOUNT_NAME}.blob.core.windows.net/{CONTAINER_NAME}/train.tar.gz \\\n        --destination  /data/imagenet/train.tar.gz\\\n        --source-key {storage_account_key}\\\n        --quiet\n\ncd /data/imagenet\ntar -xzf train.tar.gz\ntar -xzf validation.tar.gz\n\"\"\""
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": "with open('nodeprep.sh', 'w') as f:\n    f.write(nodeprep_script)"
        },
        {
            "cell_type": "markdown",
            "metadata": {
                "lines_to_next_cell": 2
            },
            "source": "Next we will copy the file over and run it on the NFS VM. This will install azcopy and download and prepare the data"
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": "if USE_FAKE:\n    raise Warning(\"You should not be running this section if you simply want to use fake data\")"
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {
                "tags": [
                    "stripout"
                ]
            },
            "outputs": [],
            "source": "if USE_FAKE is False:\n    !sshpass -p {get_password(dotenv_for())} scp -o \"StrictHostKeyChecking=no\" nodeprep.sh $USERNAME@{nfs_ip}:~/"
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {
                "tags": [
                    "stripout"
                ]
            },
            "outputs": [],
            "source": "if USE_FAKE is False:\n    !sshpass -p {get_password(dotenv_for())} ssh -o \"StrictHostKeyChecking=no\" $USERNAME@{nfs_ip} \"sudo chmod 777 ~/nodeprep.sh && ./nodeprep.sh\""
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": "<a id='configure_cluster'></a>\n## Configure Batch AI Cluster\nWe then upload the scripts we wish to execute onto the fileshare. The fileshare will later be mounted by Batch AI. An alternative to uploading the scripts would be to embedd them inside the Docker image."
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": "!az storage file upload --share-name $FILE_SHARE_NAME --source HorovodPytorch/cluster_config/docker.service --path scripts\n!az storage file upload --share-name $FILE_SHARE_NAME --source HorovodPytorch/cluster_config/nodeprep.sh --path scripts"
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": "Below it the command to create the cluster. "
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {
                "tags": [
                    "stripout"
                ]
            },
            "outputs": [],
            "source": "!az batchai cluster create \\\n    -w $WORKSPACE \\\n    --name $CLUSTER_NAME \\\n    --image UbuntuLTS \\\n    --vm-size $VM_SIZE \\\n    --min $NUM_NODES --max $NUM_NODES \\\n    --afs-name $FILE_SHARE_NAME \\\n    --afs-mount-path extfs \\\n    --user-name $USERNAME \\\n    --password {get_password(dotenv_for())} \\\n    --storage-account-name $STORAGE_ACCOUNT_NAME \\\n    --storage-account-key $storage_account_key \\\n    --nfs $NFS_NAME \\\n    --nfs-mount-path nfs \\\n    --config-file HorovodPytorch/cluster_config/cluster.json"
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": "Let's check that the cluster was created succesfully."
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {
                "tags": [
                    "stripout"
                ]
            },
            "outputs": [],
            "source": "!az batchai cluster show -n $CLUSTER_NAME -w $WORKSPACE"
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": "!az batchai cluster list -w $WORKSPACE -o table"
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": "!az batchai cluster node list -c $CLUSTER_NAME -w $WORKSPACE -o table"
        }
    ],
    "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
}