HorovodPytorch/00_CreateImageAndTest.ipynb

{
    "cells": [
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": "# Create Docker Image for PyTorch\nIn this notebook we will create the Docker image for our PyTorch script to run in. We will go through the process of creating the image and testing it locally to make sure it runs before submitting it to the cluster. It is often recommended you do this rather than debugging on the cluster since debugging on a cluster can be much more difficult and time consuming.\n \n**You will need to be running everything on a GPU enabled VM to run this notebook.** "
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": "import sys\nsys.path.append(\"../common\") \n\nfrom dotenv import get_key\nimport os\nfrom utils import dotenv_for\nimport docker"
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": "We will use fake data here since we don't want to have to download the data etc. Using fake data is often a good way to debug your models as well as checking what IO overhead is. Here we are setting the number of processes (NUM_PROCESSES) to 2 since the VM we are testing on has 2 GPUs. If you are running on a machine with 1 GPU set NUM_PROCESSES to 1."
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {
                "tags": [
                    "parameters"
                ]
            },
            "outputs": [],
            "source": "dotenv_path = dotenv_for()\nUSE_FAKE               = True\nDOCKERHUB              = os.getenv('DOCKER_REPOSITORY', \"masalvar\")\nNUM_PROCESSES          = 2\nDOCKER_PWD             = get_key(dotenv_path, 'DOCKER_PWD')"
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": "dc = docker.from_env()"
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": "image, log_iter = dc.images.build(path='Docker', \n                          tag='{}/caia-horovod-pytorch'.format(DOCKERHUB))"
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": "container_labels = {'containerName': 'pytorchgpu'}\nenvironment ={\n    \"DISTRIBUTED\":True,\n    \"PYTHONPATH\":'/workspace/common/',\n}\n\nvolumes = {\n    os.getenv('EXT_PWD'): {\n                                'bind': '/workspace', \n                                'mode': 'rw'\n                               }\n}\n\nif USE_FAKE:\n    environment['FAKE'] = True\nelse:\n    environment['FAKE'] = False\n    volumes[os.getenv('EXT_DATA')]={'bind': '/mnt/input', 'mode': 'rw'}\n    environment['AZ_BATCHAI_INPUT_TRAIN'] = '/mnt/input/train'\n    environment['AZ_BATCHAI_INPUT_TEST'] = '/mnt/input/validation'"
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": "cmd=f'mpirun -np {NUM_PROCESSES} -H localhost:{NUM_PROCESSES} '\\\n     'python -u /workspace/HorovodPytorch/src/imagenet_pytorch_horovod.py'\ncontainer = dc.containers.run(image.tags[0], \n                              command=cmd,\n                              detach=True, \n                              labels=container_labels,\n                              runtime='nvidia',\n                              volumes=volumes,\n                              environment=environment,\n                              shm_size='8G',\n                              privileged=True)"
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": "With the code below we are simply monitoring what is happening in the container. Feel free to stop the notebook when you are happy that everything is working."
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {
                "tags": [
                    "stripout"
                ]
            },
            "outputs": [],
            "source": "for line in container.logs(stderr=True, stream=True):\n    print(line.decode(\"utf-8\"),end =\"\")"
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": "container.reload() # Refresh state\nif container.status is 'running':\n    container.kill()"
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {
                "tags": [
                    "stripout"
                ]
            },
            "outputs": [],
            "source": "for line in dc.images.push(image.tags[0], \n                           stream=True,\n                           auth_config={\"username\": DOCKERHUB,\n                                        \"password\": DOCKER_PWD}):\n    print(line)"
        }
    ],
    "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
}