colab_codelab.ipynb

{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "qrzr4aCx_2ue"
      },
      "outputs": [],
      "source": [
        "# Copyright 2021 DeepMind Technologies Limited\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",
        "#      http://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": "JUsdeEk8CGcf"
      },
      "source": [
        "# XManager Codelab Notebook\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Laa1Anzh0H8L"
      },
      "source": [
        "This notebook will take you through running an XManager experiment on Google Cloud Platform (GCP).\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "T8oYMjJECXmE"
      },
      "source": [
        "## Install XManager"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "OvCRFe9_CZ7c"
      },
      "source": [
        "Remember to restart the runtime after running the cell below for the first time. ",
        "Avoid using run all, as this colab has async functions that will not work correctly."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "KrMsvX3vjq_x"
      },
      "outputs": [],
      "source": [
        "!pip install git+https://github.com/deepmind/xmanager\n",
        "!git clone https://github.com/deepmind/xmanager xmanager_repo"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8jWVaFeDDlgp"
      },
      "source": [
        "## Some utilities"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "C_rpcBUj_AcE"
      },
      "outputs": [],
      "source": [
        "# Dependencies used\n",
        "from IPython.display import display\n",
        "import ipywidgets\n",
        "import os\n",
        "import itertools\n",
        "import asyncio\n",
        "import nest_asyncio\n",
        "\n",
        "# Display text box for updating environmental variables\n",
        "def display_env_variable_box(env_variable):\n",
        "  def variable_changed(change):\n",
        "      os.environ[env_variable] = change.new\n",
        "\n",
        "  ENV_VARIABLE_WIDGET = ipywidgets.Text(\n",
        "      description=env_variable + ':',\n",
        "      style={'description_width': 'initial'},\n",
        "      layout=ipywidgets.Layout(width='50%'),\n",
        "  )\n",
        "  ENV_VARIABLE_WIDGET.observe(variable_changed, names='value')\n",
        "\n",
        "  display(ENV_VARIABLE_WIDGET)\n",
        "\n",
        "# Allows running async functions in Colab\n",
        "nest_asyncio.apply()\n",
        "\n",
        "run_async_function = lambda fn: asyncio.get_event_loop().run_until_complete(fn())"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WBEfpkwqCkeV"
      },
      "source": [
        "## Setup GCP"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "RXfCT5il1RVK"
      },
      "source": [
        "### 1. [Create](https://console.cloud.google.com/) a GCP project if one does not already exist and enter its name below"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "HYGaFHj4BhC-"
      },
      "outputs": [],
      "source": [
        "display_env_variable_box('GOOGLE_CLOUD_PROJECT')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "55t4TlhfOuNl"
      },
      "outputs": [],
      "source": [
        "!gcloud config set project \"${GOOGLE_CLOUD_PROJECT}\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "h1XgOtRJ1Wtm"
      },
      "source": [
        "### 2. Authenticate"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "z_N-TglIAq54"
      },
      "outputs": [],
      "source": [
        "!gcloud auth login"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Na7Qe1nKBUCk"
      },
      "outputs": [],
      "source": [
        "!gcloud auth application-default login"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_7jXxwCX2IjB"
      },
      "source": [
        "### 3. Create a Google Cloud Storage Bucket if one does not already exist and enter its name in the box below"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "MzPM0raR5KgP"
      },
      "outputs": [],
      "source": [
        "display_env_variable_box('GOOGLE_CLOUD_BUCKET_NAME')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "jBZb8A0fDS7a"
      },
      "source": [
        "## Launching an experiment"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "fTbKiDPq73ni"
      },
      "source": [
        "An experiment can be broken down into 5 steps:\n",
        "\n",
        "1. Creating the experiment.\n",
        "2. Defining the executable specification.\n",
        "3. Defining the execution environment.\n",
        "4. Creating the jobs and experiment units.\n",
        "5. Defining the hyperparameters."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Y0jhekyW79NH"
      },
      "source": [
        "### 1. Creating the experiment"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_rC4jkLV5dsb"
      },
      "outputs": [],
      "source": [
        "from xmanager import xm\n",
        "from xmanager import xm_local\n",
        "# This code block sets FLAGS to use default values to avoid an absl.flags.UnparsedFlagAccessError.\n",
        "# Normally XManager flags are set via the command-line with `xmanager train.py -- --key=value`\n",
        "\n",
        "from absl import flags\n",
        "flags.FLAGS([''])\n",
        "flags.FLAGS.xm_wrap_late_bindings = True"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ALx50GrH8Cus"
      },
      "source": [
        "Experiments are the core of XManager. An experiment typically involves running a computation (e.g., training a model in JAX or TensorFlow) in different hyperparameter configurations. It can have associated metadata (name, description, notes, tags, links, etc.). Experiments are made up of various experiment units, including work units which do the computation(s) in question and auxiliary units which perform other functions like TensorBoard."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DcF4km_i9noh"
      },
      "source": [
        "Give the experiment a name. The `create_experiment` method will also create a unique integer id for the experiment and save this experiment to a database."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6YLC2xNj9Of9"
      },
      "outputs": [],
      "source": [
        "async def create_experiment_demo():\n",
        "  async with xm_local.create_experiment(experiment_title='my-experiment') as experiment:\n",
        "    print(f'Local Experiment created with experiment_id={experiment.experiment_id}')\n",
        "\n",
        "run_async_function(create_experiment_demo)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "YN18pmaC9fLA"
      },
      "source": [
        "### 2. Defining the executable specification"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "fOhs6j3v9tMB"
      },
      "source": [
        "Define the job that will run in the experiment. A `PythonContainer` is an example of an executable specification. This executable specification tells XManager to package everything inside the `PythonContainer.path` as a container and use `PythonContainer.entrypoint` as the main module. Because we cloned XManager to `~/xmanager_repo` in an early step, we can use one of the examples, `/content/xmanager_repo/examples/cifar10_torch` as the path.\n",
        "\n",
        "We also need to declare where the executable should be staged. This step will upload the executable specification to the correct storage option that is best suited for the execution environment. For example, if the execution environment is Vertex AI, the executable must be stored in Google Container Registry. The `Vertex.Spec()` specification will upload the specification to Google Container Registry, where it will be accessible by Vertex AI."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "TtonbZZ33PoW"
      },
      "source": [
        "```python\n",
        "[executable] = experiment.package([\n",
        "  xm.python_container(\n",
        "    executor_spec=xm_local.Vertex.Spec(),\n",
        "    path=os.path.expanduser('/content/xmanager_repo/examples/cifar10_torch'),\n",
        "    entrypoint=xm.ModuleName('cifar10'),\n",
        "  )\n",
        "])\n",
        "```"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "VK9t8Xw4-Xza"
      },
      "source": [
        "### 3. Defining the execution environment"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "485vWvlO-a5r"
      },
      "source": [
        "Declare where the job will run and what compute requirements are necessary to run one job. To run on AI Vertex, we must use the `xm_local.Vertex` executor. Each job should use 1 NVidia T4 GPU, so we must pass in a `xm.JobRequirements` to the executor."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0D7tqVt_3YXP"
      },
      "source": [
        "```python\n",
        "executor = xm_local.Vertex(xm.JobRequirements(T4=1))\n",
        "```"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xLCfqZg8-9j5"
      },
      "source": [
        "### 4. Launching the jobs"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "v2kPjxXj_Kge"
      },
      "source": [
        "Finally, we can create an experiment and jobs to the experiment. A **job** is an unit of execution. A job contains an executable representing \"what to run\" and an executor reprenting \"how to run it\". Jobs can be reused.\n",
        "\n",
        "To add a single job to the experiment, create a `xm.Job` object that combine the executable, compute requirements, and custom arguments hyperparameters, and the job to the experiment."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "87ikHuuy3cOH"
      },
      "source": [
        "```python\n",
        "experiment.add(xm.Job(\n",
        "    executable=executable,\n",
        "    executor=executor,\n",
        "    args={'batch_size': 64, 'learning_rate': 0.01},\n",
        "))\n",
        "```"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "n8vggs0C_ff5"
      },
      "source": [
        "#### Job Groups\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "sDWKF90Z_mK7"
      },
      "source": [
        "A **job group** is a collection of jobs. Job groups can be added to experiments in the same way. The jobs in a job groups should be connected in some way. Some examples include: a client and server, multiple workers in a distributed experiment, a TPU session and a coordinator, a trainer and a evaluator, etc."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Ptwtv7cN_tLX"
      },
      "source": [
        "```python\n",
        "async with xm_local.create_experiment(experiment_title='cifar10') as experiment:\n",
        "    await experiment.add(xm.JobGroup(\n",
        "        client=xm.Job(\n",
        "            executable=executable,\n",
        "            executor=executor,\n",
        "        ),\n",
        "        server=xm.Job(\n",
        "            executable=executable,\n",
        "            executor=executor,\n",
        "        ),\n",
        "    ))\n",
        "```"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "E6QqXqyj_xTK"
      },
      "source": [
        "### 5. Defining the hyperparameters"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "APGQHBT3_5PP"
      },
      "source": [
        "In research, it is often required to run the experimental setup multiple times with different hyperparameter values. This is called **hyperparameter optimization**. The simplest form of hyperparameter optimization is called *grid search* or *parameter sweep*, which is an exhaustive search through all possible Cartesian products of hyperparameter values. Grid search trials can be constructed using `itertools`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "rEtzBvpX_8Zy"
      },
      "outputs": [],
      "source": [
        "inputs = {\n",
        "    'batch_size': [64, 128],\n",
        "    'learning_rate': [0.01, 0.001],\n",
        "}\n",
        "hyperparameters = list(dict(zip(inputs, x)) for x in itertools.product(*inputs.values()))\n",
        "\n",
        "from pprint import pprint\n",
        "pprint(hyperparameters)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7-3Esf-_AAnv"
      },
      "source": [
        "To perform the grid search, loop over all the hyperparameters, passing a different hyperparameter configuration to the `args` parameter of each job. Add each job to the experiment."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-0DfojJBAE8v"
      },
      "source": [
        "```python\n",
        "async with xm_local.create_experiment(experiment_title='cifar10') as experiment:\n",
        "    for hparams in trials:\n",
        "        experiment.add(xm.Job(\n",
        "            executable=executable,\n",
        "            executor=executor,\n",
        "            args=hparams,\n",
        "        ))\n",
        "```"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "crLkWI86APJz"
      },
      "source": [
        "### Experiment Units\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0J2u_XtxARZm"
      },
      "source": [
        "The return result of `experiment.add` is an awaitable experiment unit.\n",
        "\n",
        "An **experiment unit** is created by XManager for every instance the experiment launches a set launched jobs. Jobs inside an experiment unit must be created together or should be destroyed together.\n",
        "\n",
        "The common type of experiment unit is the **work unit**. A work unit contains the main task of training in an ML experiment, and they are part of the hyperparameter sweep. A helpful way of thinking of work units is that each work unit represents a trial. Just as an experiment can trial different hyperparameter sets, an XManager experiment can create a work unit for each trial. Work units are enumerated and assigned a `work_unit_id` starting from 0.\n",
        "\n",
        "There are other experiment unit types which represent differnt roles. For example, an **auxiliary unit** represents a job that runs alongside all the trials such as a Tensorboard or an hyperparameter optimizer job."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2ZVLBdJTAced"
      },
      "source": [
        "### Job Generator"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "x8uOjNOFAf1z"
      },
      "source": [
        "It's possible to modify the args of a job based on the attributes of the work unit. A good example of wanting to do this is when you want to create a pass a different log directory argument to each work unit. Another example is if you need to adjust the args with the correct DNS names based on the `work_unit_id`.\n",
        "\n",
        "A job generator allows you modify the arguments of a job before they are added to the work unit."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-frD6agIAk5q"
      },
      "source": [
        "```python\n",
        "async generate(work_unit, hparams):\n",
        "    print(work_unit.work_unit_id, hparams)\n",
        "    # hparams.update({'log_dir', f'/tmp/{work_unit.work_unit_id}'})\n",
        "    work_unit.add(xm.Job(\n",
        "        executable=executable,\n",
        "        executor=executor,\n",
        "        args=hparams,\n",
        "    ))\n",
        "\n",
        "async with xm_local.create_experiment(experiment_title='cifar10') as experiment:\n",
        "    for hparams in trials:\n",
        "        work_unit = await experiment.add(generate, hparams)\n",
        "```"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "U_rI_3ccArED"
      },
      "source": [
        "### Tracking job status\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Gu2_yEZ7AtoC"
      },
      "source": [
        "You can list all of your previous experiments."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "NFDASr8JAvLb"
      },
      "outputs": [],
      "source": [
        "[e.experiment_id for e in xm_local.list_experiments()]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Z5mXwH1qAzZq"
      },
      "source": [
        "Some execution environments allow you to track the status of jobs in an experiment. Vertex AI is one of the execution environments that supports job-tracking."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "beIesla1Bf4y"
      },
      "source": [
        "```python\n",
        "# TODO: Use experiment.work_units instead of private member.\n",
        "for i, unit in enumerate(experiment._experiment_units):\n",
        "    print(f'[{i}] Completed: {unit.get_status().is_completed}, Failed: {unit.get_status().is_failed}')\n",
        "```"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "VmJL_mm6A7XL"
      },
      "source": [
        "# End to end"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xTsY3lUhA-RU"
      },
      "source": [
        "Combining everything above into a single code-block, the launch script looks like this:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "dbjyU6P1OX7b"
      },
      "outputs": [],
      "source": [
        "async def launch_experiment():\n",
        "  async with xm_local.create_experiment(experiment_title='cifar10') as experiment:\n",
        "    [executable] = experiment.package([\n",
        "        xm.python_container(\n",
        "            executor_spec=xm_local.Vertex.Spec(),\n",
        "            path=os.path.expanduser('/content/xmanager_repo/examples/cifar10_torch'),\n",
        "            entrypoint=xm.ModuleName('cifar10'),\n",
        "        )\n",
        "    ])\n",
        "\n",
        "    batch_sizes = [64, 128]\n",
        "    learning_rates = [0.01, 0.001]\n",
        "    trials = list(\n",
        "        dict([('batch_size', bs), ('learning_rate', lr)])\n",
        "        for (bs, lr) in itertools.product(batch_sizes, learning_rates)\n",
        "    )\n",
        "    for hyperparameters in trials:\n",
        "        experiment.add(xm.Job(\n",
        "            executable=executable,\n",
        "            executor=xm_local.Vertex(requirements=xm.JobRequirements(T4=1)),\n",
        "            args=hyperparameters,\n",
        "        ))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "pAz-xSrTBDfh"
      },
      "outputs": [],
      "source": [
        "# Since Docker can't be used in Colab, the image will be built using the CloudBuild API.\n",
        "\n",
        "# Make sure required APIs are enabled for the project\n",
        "run_async_function(launch_experiment)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tUraiepOtiHU"
      },
      "source": [
        "## Revoke credentials"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "cx70AqbmtlYh"
      },
      "outputs": [],
      "source": [
        "!gcloud auth revoke\n",
        "!gcloud auth application-default revoke\n",
        "!gcloud config unset project"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [],
      "name": "XManager Codelab.ipynb",
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}