docs/guides/benchmarks/ray_benchmarks.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wBP-KRCq6G6b"
      },
      "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/benchmarks/ray_benchmarks.ipynb)\n",
        "\n",
        "# Benchmarking with Ray\n",
        "\n",
        "We provide a brief guide below on the Vizier + Ray integration, and how to\n",
        "benchmark with all publicly available algorithms on\n",
        "[Ray Tune](https://docs.ray.io/en/latest/tune/). Notably, Tune integrates with a\n",
        "wide range of additional hyperparameter optimization tools, including Ax,\n",
        "BayesOpt, BOHB, Dragonfly, FLAML, HEBO, Hyperopt, Nevergrad, Optuna, SigOpt,\n",
        "skopt, and ZOOpt.\n",
        "\n",
        "![alt-text](https://docs.ray.io/en/latest/_images/tune_overview.png)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "f76xE1Os5DFD"
      },
      "source": [
        "## Initial Installation"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "BUfxI1EPwXfH"
      },
      "outputs": [],
      "source": [
        "!pip install google-vizier[jax]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "keo91_QSY9pr"
      },
      "outputs": [],
      "source": [
        "!pip install -U \"ray[default]\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "M4H9Q3y85dwC"
      },
      "source": [
        "## Algorithm and Experimenter Factories\n",
        "\n",
        "As mentioned in previous guides, since we want to compare algorithms across multiple benchmarks, we first create a bunch of relevant benchmark experimenters. To do so, we use `SerializableExperimenterFactory` from our [Experimenters API](https://github.com/google/vizier/blob/main/vizier/benchmarks/experimenters/__init__.py) to modularize the construction of multiple benchmark components.\n",
        "\n",
        "For example, here we can create a diverse set of BBOB functions with different dimensions via the `BBOBExperimenterFactory`. Then, we can print out the full serialization of the benchmarks that we have created."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "jDQNz57RYymT"
      },
      "outputs": [],
      "source": [
        "import itertools\n",
        "import numpy as np\n",
        "from vizier.benchmarks import experimenters\n",
        "\n",
        "function_names = [\n",
        "    'Sphere',\n",
        "    'BentCigar',\n",
        "    'Katsuura',\n",
        "]\n",
        "dimensions = [4, 8]\n",
        "product_list = list(itertools.product(function_names, dimensions))\n",
        "\n",
        "experimenter_factories = []\n",
        "for product in product_list:\n",
        "  name, dim = product\n",
        "  bbob_factory = experimenters.BBOBExperimenterFactory(name=name, dim=dim)\n",
        "  experimenter_factory = experimenters.SingleObjectiveExperimenterFactory(\n",
        "      bbob_factory,\n",
        "      shift=np.random.uniform(low=-2, high=2, size=dim),\n",
        "      noise_type='LIGHT_ADDITIVE_GAUSSIAN',\n",
        "  )\n",
        "  experimenter_factories.append(experimenter_factory)\n",
        "  print(experimenter_factory.dump())"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cSfnV6CjZzPb"
      },
      "source": [
        "Next, we need to define our algorithms by installing the relevant packages and importing the relevant algorithms. For simplicity, we only compare against only a subset of the algorithms that Ray supports.\n",
        "\n",
        "**NOTE:** We provide the `VizierSearch` class in our own libaries that can directly use the `Searcher` API in Ray. The imports are given below."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "xkw_F32tZwiN"
      },
      "outputs": [],
      "source": [
        "pip install ax-platform scikit-optimize hyperopt optuna bayesian-optimization"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "MytRchwhau0z"
      },
      "outputs": [],
      "source": [
        "from ray import tune\n",
        "from ray.tune.search.ax import AxSearch\n",
        "from ray.tune.search.bayesopt import BayesOptSearch\n",
        "from ray.tune.search.hyperopt import HyperOptSearch\n",
        "from ray.tune.search.optuna import OptunaSearch\n",
        "from ray.tune.search.skopt import SkOptSearch\n",
        "from vizier import raytune as vzr\n",
        "from vizier._src.raytune.vizier_search import VizierSearch\n",
        "\n",
        "algorithm_factories = {\n",
        "    'ray': lambda: None,\n",
        "    'vizier': VizierSearch,\n",
        "    'ax': AxSearch,\n",
        "    'bayesopt': BayesOptSearch,\n",
        "    'optuna': OptunaSearch,\n",
        "    'hyperopt': HyperOptSearch,\n",
        "    'skopt': SkOptSearch,\n",
        "}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "OdXa6kDM5pOb"
      },
      "source": [
        "## Running RayTune\n",
        "Running RayTune using `ExperimenterFactory` is made easy using our utility libraries which takes in any factory with a `TuneConfig` to run the algorithm on the corresponding benchmark. Let us first run one algorithm on the first benchmark and see the results that we get.\n",
        "\n",
        "**NOTE:** This uses a local Ray instance."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "KRufPFLc7rwT"
      },
      "outputs": [],
      "source": [
        "ALGORITHM_NAME = 'ray'  # @param str\n",
        "experimenter_factory = experimenter_factories[0]\n",
        "factory = algorithm_factories[ALGORITHM_NAME]\n",
        "tune_config = tune.TuneConfig(\n",
        "    search_alg=factory(),\n",
        "    num_samples=4,\n",
        "    max_concurrent_trials=1,\n",
        ")\n",
        "vzr.run_tune.run_tune_from_factory(experimenter_factory, tune_config)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4ZKtke9pdjlU"
      },
      "source": [
        "Now, we repeat our runs for each `ExperimenterFactory` and each algorithm, converting the results into `PlotElements` for easy plotting and comparison."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "s_nuLfYsdtz6"
      },
      "outputs": [],
      "source": [
        "from vizier.benchmarks import analyzers\n",
        "\n",
        "NUM_REPEATS = 3  # @param\n",
        "NUM_ITERATIONS = 50  # @param\n",
        "\n",
        "\n",
        "def results_to_element(results_list):\n",
        "  curves = []\n",
        "  for results in results_list:\n",
        "    raw_ys = np.array(results.get_dataframe()['bbob_eval_before_noise'])\n",
        "    ys = np.minimum.accumulate(raw_ys)\n",
        "    curve = analyzers.ConvergenceCurve(\n",
        "        xs=np.arange(1, len(ys) + 1),\n",
        "        ys=ys.reshape((1, len(ys))),\n",
        "        trend=analyzers.ConvergenceCurve.YTrend.DECREASING,\n",
        "    )\n",
        "    curves.append(curve)\n",
        "  all_curves = analyzers.ConvergenceCurve.align_xs(curves)\n",
        "  ele = analyzers.PlotElement(curve=all_curves[0], yscale='symlog')\n",
        "  return ele\n",
        "\n",
        "\n",
        "all_records = []\n",
        "for experimenter_factory in experimenter_factories:\n",
        "  for algorithm, factory in algorithm_factories.items():\n",
        "    results = []\n",
        "    for _ in range(NUM_REPEATS):\n",
        "      tune_config = tune.TuneConfig(\n",
        "          search_alg=factory(),\n",
        "          num_samples=NUM_ITERATIONS,\n",
        "          max_concurrent_trials=1,\n",
        "      )\n",
        "      results.append(\n",
        "          vzr.run_tune.run_tune_from_factory(experimenter_factory, tune_config)\n",
        "      )\n",
        "    ele = results_to_element(results)\n",
        "    record = analyzers.BenchmarkRecord(\n",
        "        algorithm=algorithm,\n",
        "        experimenter_metadata=experimenter_factory.dump(),\n",
        "        plot_elements={'objective': ele},\n",
        "    )\n",
        "    all_records.append(record)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "JSt2Ypf3UVO1"
      },
      "outputs": [],
      "source": [
        "analyzed_records = analyzers.BenchmarkRecordAnalyzer.add_comparison_metrics(\n",
        "    records=all_records, baseline_algo='ray'\n",
        ")\n",
        "analyzers.plot_from_records(analyzed_records)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZgZemyGUUXUt"
      },
      "source": [
        "## Running Parallelized Ray\n",
        "\n",
        "In the previous example, we are using Ray local instances and running each benchmark in sequential format, which can take minutes. When there are a large number of benchmarks or computationally intensive benchmark runs, using parallelism distributed across each (algorithm, benchmark) tuple is crucial for reasonable benchmarking turnaround. We recommend using the [Ray Jobs API](https://docs.ray.io/en/latest/cluster/running-applications/job-submission/index.html) to distribute work across clusters."
      ]
    }
  ],
  "metadata": {
    "colab": {
      "name": "Ray Benchmarks.ipynb",
      "private_outputs": true,
      "provenance": []
    },
    "gpuClass": "standard",
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}