Vizier Basics

Below, we provide examples of how to:

• Define a problem statement and study configuration.

• Start a client.

• (Optionally) Connect the client to a server.

• Perform a typical tuning loop.

• Use other client APIs.

Installation and reference imports

!pip install google-vizier[jax]

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

Setting up the problem statement

Here we setup the problem statement, which contains information about the search space and the metrics to optimize.

problem = vz.ProblemStatement()
problem.search_space.root.add_float_param('x', 0.0, 1.0)
problem.search_space.root.add_float_param('y', 0.0, 1.0)
problem.metric_information.append(
    vz.MetricInformation(
        name='maximize_metric', goal=vz.ObjectiveMetricGoal.MAXIMIZE))


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

Setting up the study configuration

The study configuration contains additional information, such as the algorithm to use and level of noise that we think the objective will have.

study_config = vz.StudyConfig.from_problem(problem)
study_config.algorithm = 'DEFAULT'

Setting up the client

Starts a study_client, which can be either in local mode (default) or distributed mode.

Local Mode: The client has no endpoint set, and will implicitly create a local Vizier Service which will be shared across other clients in the same Python process. Studies will then be stored locally in a SQL database file located at service.VIZIER_DB_PATH.

study_client = clients.Study.from_study_config(study_config, owner='owner', study_id='example_study_id')
print('Local SQL database file located at: ', service.VIZIER_DB_PATH)

Distributed mode: The service may be explicitly created, wrapped as a server in a separate process to accept requests from all other client processses. Details such as the database_url, port, policy_factory, etc. can be configured in the server’s initializer.

All client processes (on a single machine or over multiple machines) will connect to this server via a globally specified endpoint.

server = servers.DefaultVizierServer()  # Ideally created on a separate process such as a server machine.
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.

Client Parallelization

Regardless of whether the setup is local or distributed, we may simultaneously create multiple clients to work on the same study, useful for parallelizing evaluation workload.

another_study_client = clients.Study.from_resource_name(study_client.resource_name)

Obtaining suggestions

Start requesting suggestions from the server, for evaluating objectives. Suggestions can be made sequentially (count=1) or in batches (count>1).

for i in range(10):
  suggestions = study_client.suggest(count=1)
  for suggestion in suggestions:
    x = suggestion.parameters['x']
    y = suggestion.parameters['y']
    objective = evaluate(x, y)
    print(f'Iteration {i}, suggestion ({x},{y}) led to objective value {objective}.')
    final_measurement = vz.Measurement({'maximize_metric': objective})
    suggestion.complete(final_measurement)

Find optimal trial

Find the best objective so far, with corresponding suggestion value. For multiobjective cases, there may be multiple outputs of optimal_trials(), all corresponding to a Pareto-optimal curve.

for optimal_trial in study_client.optimal_trials():
  optimal_trial = optimal_trial.materialize()
  print("Optimal Trial Suggestion and Objective:", optimal_trial.parameters,
        optimal_trial.final_measurement)

Other client commands

The study_client can also send other requests, such as the following:

study_client.get_trial(1)  # Get the first trial.
study_client.trials()  # Get all trials so far.

# Obtain only the completed trials.
trial_filter = vz.TrialFilter(status=[vz.TrialStatus.COMPLETED])
study_client.trials(trial_filter=trial_filter)