Open in Colab

Distributed Vizier

This documentation shows how to perform distributed optimization over multiple clients.

Installation and reference imports

!pip install google-vizier[jax]
import multiprocessing

from vizier import service
from vizier.service import clients
from vizier.service import pyvizier as vz
from vizier.service import servers

Regular setup

We setup a regular study configuration below.

study_config = vz.StudyConfig()
study_config.search_space.root.add_float_param('x', 0.0, 1.0)
study_config.metric_information.append(vz.MetricInformation(name='metric', goal=vz.ObjectiveMetricGoal.MAXIMIZE))
study_config.algorithm = 'DEFAULT'


def evaluate(x: float) -> float:
  return 2*x - x**2

Server creation

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.

server = servers.DefaultVizierServer()  # Ideally created on a separate process such as a server machine.

Client parallelization

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.

clients.environment_variables.server_endpoint = server.endpoint  # Server address.
study_client = clients.Study.from_study_config(study_config, owner='owner', study_id = 'example_study_id')  # Now connects to the explicitly created server.
another_study_client = clients.Study.from_resource_name(study_client.resource_name)  # Another way to fork clients.

Distributed suggestions

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.

def tuning_loop(client_id: str):
  for i in range(10):
    suggestions = study_client.suggest(count=1, client_id=client_id)
    for suggestion in suggestions:
      objective = evaluate(suggestion.parameters['x'])
      final_measurement = vz.Measurement({'metric': objective})
      suggestion.complete(final_measurement)

For example, we may perform a threadpool and construct multiple clients to parallelize evaluations on a single machine.

NUM_CLIENTS = 10
NUM_TRIALS_PER_CLIENT = 50

pool = multiprocessing.pool.ThreadPool(NUM_CLIENTS)
pool.map(tuning_loop, range(NUM_CLIENTS))