{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Xgqk7eHswDpB"
      },
      "source": [
        "[![Open in Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/google/vizier/blob/main/docs/guides/user/distributed.ipynb)\n",
        "\n",
        "# Distributed Vizier\n",
        "This documentation shows how to perform distributed optimization over multiple clients."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "O5RnMytPR8Aw"
      },
      "source": [
        "## Installation and reference imports"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "kSG8XlxLvCJO"
      },
      "outputs": [],
      "source": [
        "!pip install google-vizier[jax]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "fzYr0bPYSHfQ"
      },
      "outputs": [],
      "source": [
        "import multiprocessing\n",
        "\n",
        "from vizier import service\n",
        "from vizier.service import clients\n",
        "from vizier.service import pyvizier as vz\n",
        "from vizier.service import servers"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qJ1kRiHaKOVt"
      },
      "source": [
        "## Regular setup\n",
        "We setup a regular study configuration below."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "zX2G3_pcKYdG"
      },
      "outputs": [],
      "source": [
        "study_config = vz.StudyConfig()\n",
        "study_config.search_space.root.add_float_param('x', 0.0, 1.0)\n",
        "study_config.metric_information.append(vz.MetricInformation(name='metric', goal=vz.ObjectiveMetricGoal.MAXIMIZE))\n",
        "study_config.algorithm = 'DEFAULT'\n",
        "\n",
        "\n",
        "def evaluate(x: float) -\u003e float:\n",
        "  return 2*x - x**2"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "w3m48cPsXcxD"
      },
      "source": [
        "## Server creation\n",
        "Unlike the single-client case, in the distributed case, we require a single explicit server to accept requests from all other client processses. Details such as the `host`, `port`, `database_url`, `policy_factory`, etc. can be configured in the server's initializer."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "V6ef6OfMXdpz"
      },
      "outputs": [],
      "source": [
        "server = servers.DefaultVizierServer()  # Ideally created on a separate process such as a server machine."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ktExEiS0xlH_"
      },
      "source": [
        "## Client parallelization\n",
        "We may simultaneously create multiple clients to work on the same study, useful for parallelizing evaluation workload. All client processes (on a single machine or over multiple machines) will connect to this server via a globally specified `endpoint`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "EQR1_u-VxEwn"
      },
      "outputs": [],
      "source": [
        "clients.environment_variables.server_endpoint = server.endpoint  # Server address.\n",
        "study_client = clients.Study.from_study_config(study_config, owner='owner', study_id = 'example_study_id')  # Now connects to the explicitly created server.\n",
        "another_study_client = clients.Study.from_resource_name(study_client.resource_name)  # Another way to fork clients."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Vh3eNsrAdaMJ"
      },
      "source": [
        "## Distributed suggestions\n",
        "We may now distribute our workflow, with each worker/client using the same loop below. Each client requires a unique `client_id` however, to ensure the server can identify client workers and distribute workloads properly."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "BnFKc7FadkJV"
      },
      "outputs": [],
      "source": [
        "def tuning_loop(client_id: str):\n",
        "  for i in range(10):\n",
        "    suggestions = study_client.suggest(count=1, client_id=client_id)\n",
        "    for suggestion in suggestions:\n",
        "      objective = evaluate(suggestion.parameters['x'])\n",
        "      final_measurement = vz.Measurement({'metric': objective})\n",
        "      suggestion.complete(final_measurement)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NVGcVEzb0Gxe"
      },
      "source": [
        "For example, we may perform a threadpool and construct multiple clients to parallelize evaluations on a single machine."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "R0pcPViUz9zC"
      },
      "outputs": [],
      "source": [
        "NUM_CLIENTS = 10\n",
        "NUM_TRIALS_PER_CLIENT = 50\n",
        "\n",
        "pool = multiprocessing.pool.ThreadPool(NUM_CLIENTS)\n",
        "pool.map(tuning_loop, range(NUM_CLIENTS))"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "last_runtime": {
        "build_target": "//ads/thresholds/kumamon/colab:notebook",
        "kind": "shared"
      },
      "name": "Distributed.ipynb",
      "private_outputs": true,
      "provenance": [
        {
          "file_id": "/piper/depot/http://github.com/google/vizier/tree/main/vizier/docs/guides/user/running_vizier.ipynb",
          "timestamp": 1673247218127
        },
        {
          "file_id": "1q87rsDDUJLHci3o9Gv-sU0g7H3O3lAbU",
          "timestamp": 1659555396142
        }
      ]
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}