{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NF8O2wVCztXx"
      },
      "source": [
        "[![Open in Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/google/vizier/blob/main/docs/advanced_topics/pyglove/vizier_as_backend.ipynb)\n",
        "\n",
        "# OSS Vizier as a Backend\n",
        "We demonstrate how OSS Vizier can be used as a distributed backend for PyGlove-based tuning tasks.\n",
        "\n",
        "This assumes the user is already familiar with PyGlove primitives."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NVufIkimznpa"
      },
      "source": [
        "## Installation and reference imports"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "rD_Okl7wzos6"
      },
      "outputs": [],
      "source": [
        "!pip install google-vizier\n",
        "!pip install pyglove"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6DvcXueRzq8C"
      },
      "outputs": [],
      "source": [
        "import multiprocessing\n",
        "import multiprocessing.pool\n",
        "import os\n",
        "\n",
        "import pyglove as pg\n",
        "from vizier import pyglove as pg_vizier\n",
        "from vizier.service import servers"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Ji8Gq-5zrQ2I"
      },
      "source": [
        "## Preliminaries\n",
        "In the original PyGlove setting, one can normally perform evolutionary computation, for example:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "DX4VYRZXraUK"
      },
      "outputs": [],
      "source": [
        "search_space = pg.Dict(x=pg.floatv(0.0, 1.0), y=pg.floatv(0.0, 1.0))\n",
        "algorithm = pg.evolution.regularized_evolution()\n",
        "num_trials = 100\n",
        "\n",
        "\n",
        "def evaluator(value: pg.Dict):\n",
        "  return value.x**2 - value.y**2\n",
        "\n",
        "\n",
        "for value, feedback in pg.sample(\n",
        "    search_space,\n",
        "    algorithm=algorithm,\n",
        "    num_examples=num_trials,\n",
        "    name='basic_run',\n",
        "):\n",
        "  reward = evaluator(value)\n",
        "  feedback(reward=reward)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "OYB-1EcKsOrm"
      },
      "source": [
        "However, in many real-world scenarios, the evaluator may be much more expensive. For example, in neural architecture search applications, `evaluator` may be the result of an entire neural network training pipeline.\n",
        "\n",
        "This leads to the need for a **backend**, in order to:\n",
        "\n",
        "1. Distribute the evaluations over multiple workers.\n",
        "2. Store the valuable results reliably and handle worker faults."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MLtbodhW0UuL"
      },
      "source": [
        "## Initializing the OSS Vizier backend\n",
        "The main initializer to call is `vizier.pyglove.init(...)`, **which should only be called once per \u003cins\u003eprocess\u003c/ins\u003e (not thread).** This function will edit global Python variables for determining values such as:\n",
        "\n",
        "1. Prefix for study names.\n",
        "2. Endpoint of the `VizierService` for storing data and handling requests.\n",
        "3. Port for the `PythiaService` for computing suggestions.\n",
        "\n",
        "In the local case, this can be called as-is:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "y4v-vphkpoZH"
      },
      "outputs": [],
      "source": [
        "pg_vizier.init('my_study')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zIbLEGi6prpm"
      },
      "source": [
        "**Alternatively**, if using a remote server, the endpoint can be specified as well:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ADpQi3ulp3Mq"
      },
      "outputs": [],
      "source": [
        "server = servers.DefaultVizierServer()  # Normally hosted on a remote machine.\n",
        "pg_vizier.init('my_study', vizier_endpoint=server.endpoint)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QkoYBGNmqUv7"
      },
      "source": [
        "## Parallelization\n",
        "Due to the OSS Vizier backend, all workers may conveniently use exactly the same evaluation loop to work on a study:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "w22LaZeHuzSG"
      },
      "outputs": [],
      "source": [
        "NUM_WORKERS = 10\n",
        "\n",
        "\n",
        "def work_fn(worker_id):\n",
        "  print(f\"Worker ID: {worker_id}\")\n",
        "  for value, feedback in pg.sample(\n",
        "      search_space,\n",
        "      algorithm=algorithm,\n",
        "      num_examples=num_trials // NUM_WORKERS,\n",
        "      name=\"worker_run\",\n",
        "  ):\n",
        "    reward = evaluator(value)\n",
        "    feedback(reward=reward)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "iZLIvUXUmCm_"
      },
      "source": [
        "There are three common forms of parallelization over the evaluation computation:\n",
        "\n",
        "1.   Multiple threads, single process.\n",
        "2.   Multiple processes, single machine.\n",
        "3.   Multiple machines.\n",
        "\n",
        "Each of these cases defines the \"worker\", which can be a thread, process or machine respectively. We demonstrate examples of every type of parallelization below."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "epVC14e4tK51"
      },
      "source": [
        "### Multiple threads, single process"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "cDYi_qhE0Q96"
      },
      "outputs": [],
      "source": [
        "with multiprocessing.pool.ThreadPool(num_workers) as pool:\n",
        "  pool.map(work_fn, range(NUM_WORKERS))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "x8dfn8LSt316"
      },
      "source": [
        "### Multiple processes, single machine"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_1b6VEECtxn-"
      },
      "outputs": [],
      "source": [
        "processes = []\n",
        "for worker_id in range(NUM_WORKERS):\n",
        "  p = multiprocessing.Process(target=work_fn, args=(worker_id,))\n",
        "  p.start()\n",
        "  processes.append(p)\n",
        "\n",
        "for p in processes:\n",
        "  p.join()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "O-ivTF1lu32-"
      },
      "source": [
        "### Multiple machines"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "gwKIMbdDulY6"
      },
      "outputs": [],
      "source": [
        "# Server Machine\n",
        "server = servers.DefaultVizierServer()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "T8zbk5IlxTtt"
      },
      "outputs": [],
      "source": [
        "# Worker Machine\n",
        "worker_id = os.uname()[1]\n",
        "pg_vizier.init('my_study', vizier_endpoint=server.endpoint)\n",
        "work_fn(worker_id)"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "last_runtime": {
        "build_target": "//ads/thresholds/kumamon/colab:notebook",
        "kind": "shared"
      },
      "name": "OSS Vizier as a Backend",
      "private_outputs": true,
      "provenance": [
        {
          "file_id": "1UChBMO67WN6pXMB6Am9rjv9_imRC9nHa",
          "timestamp": 1674077608900
        }
      ]
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}