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))