Skip to content

Commit e0d923e

Browse files
xingyousongcopybara-github
authored andcommitted
Separate out the distributed tuning doc, and make sure client_id is unique.
User didn't realize `client_id` was supposed to be unique: #1215 (comment) PiperOrigin-RevId: 712563313
1 parent 33e91e1 commit e0d923e

File tree

7 files changed

+221
-58
lines changed

7 files changed

+221
-58
lines changed

demos/run_vizier_client.py

Lines changed: 9 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -36,7 +36,6 @@
3636
from absl import app
3737
from absl import flags
3838
from absl import logging
39-
4039
from vizier import service
4140
from vizier.service import clients
4241
from vizier.service import pyvizier as vz
@@ -72,6 +71,12 @@
7271
' and API.'
7372
),
7473
)
74+
flags.DEFINE_string(
75+
'client_id',
76+
clients.UNUSED_CLIENT_ID,
77+
'The client id to use for the study. NOTE: For distributed cases, this'
78+
' needs to be unique for every client.',
79+
)
7580

7681
FLAGS = flags.FLAGS
7782

@@ -138,7 +143,9 @@ def main(argv: Sequence[str]) -> None:
138143

139144
# Evaluate the suggestion(s) and report the results to Vizier.
140145
for _ in range(FLAGS.max_num_iterations):
141-
trials = study.suggest(count=FLAGS.suggestion_count)
146+
trials = study.suggest(
147+
count=FLAGS.suggestion_count, client_id=FLAGS.client_id
148+
)
142149
for trial in trials:
143150
materialized_trial = trial.materialize()
144151
measurement = evaluate_trial(materialized_trial)

docs/guides/index.rst

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@ For Users
88
:maxdepth: 1
99

1010
user/running_vizier
11+
user/distributed
1112
user/search_spaces
1213
user/converters
1314
Switching to Vertex <https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/vizier/conversions_vertex_vizier_and_open_source_vizier.ipynb>

docs/guides/user/distributed.ipynb

Lines changed: 204 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,204 @@
1+
{
2+
"cells": [
3+
{
4+
"cell_type": "markdown",
5+
"metadata": {
6+
"id": "Xgqk7eHswDpB"
7+
},
8+
"source": [
9+
"[![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",
10+
"\n",
11+
"# Distributed Vizier\n",
12+
"This documentation shows how to perform distributed optimization over multiple clients."
13+
]
14+
},
15+
{
16+
"cell_type": "markdown",
17+
"metadata": {
18+
"id": "O5RnMytPR8Aw"
19+
},
20+
"source": [
21+
"## Installation and reference imports"
22+
]
23+
},
24+
{
25+
"cell_type": "code",
26+
"execution_count": null,
27+
"metadata": {
28+
"id": "kSG8XlxLvCJO"
29+
},
30+
"outputs": [],
31+
"source": [
32+
"!pip install google-vizier[jax]"
33+
]
34+
},
35+
{
36+
"cell_type": "code",
37+
"execution_count": null,
38+
"metadata": {
39+
"id": "fzYr0bPYSHfQ"
40+
},
41+
"outputs": [],
42+
"source": [
43+
"import multiprocessing\n",
44+
"\n",
45+
"from vizier import service\n",
46+
"from vizier.service import clients\n",
47+
"from vizier.service import pyvizier as vz\n",
48+
"from vizier.service import servers"
49+
]
50+
},
51+
{
52+
"cell_type": "markdown",
53+
"metadata": {
54+
"id": "qJ1kRiHaKOVt"
55+
},
56+
"source": [
57+
"## Regular setup\n",
58+
"We setup a regular study configuration below."
59+
]
60+
},
61+
{
62+
"cell_type": "code",
63+
"execution_count": null,
64+
"metadata": {
65+
"id": "zX2G3_pcKYdG"
66+
},
67+
"outputs": [],
68+
"source": [
69+
"study_config = vz.StudyConfig()\n",
70+
"study_config.search_space.root.add_float_param('x', 0.0, 1.0)\n",
71+
"study_config.metric_information.append(vz.MetricInformation(name='metric', goal=vz.ObjectiveMetricGoal.MAXIMIZE))\n",
72+
"study_config.algorithm = 'DEFAULT'\n",
73+
"\n",
74+
"\n",
75+
"def evaluate(x: float) -\u003e float:\n",
76+
" return 2*x - x**2"
77+
]
78+
},
79+
{
80+
"cell_type": "markdown",
81+
"metadata": {
82+
"id": "w3m48cPsXcxD"
83+
},
84+
"source": [
85+
"## Server creation\n",
86+
"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."
87+
]
88+
},
89+
{
90+
"cell_type": "code",
91+
"execution_count": null,
92+
"metadata": {
93+
"id": "V6ef6OfMXdpz"
94+
},
95+
"outputs": [],
96+
"source": [
97+
"server = servers.DefaultVizierServer() # Ideally created on a separate process such as a server machine."
98+
]
99+
},
100+
{
101+
"cell_type": "markdown",
102+
"metadata": {
103+
"id": "ktExEiS0xlH_"
104+
},
105+
"source": [
106+
"## Client parallelization\n",
107+
"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`."
108+
]
109+
},
110+
{
111+
"cell_type": "code",
112+
"execution_count": null,
113+
"metadata": {
114+
"id": "EQR1_u-VxEwn"
115+
},
116+
"outputs": [],
117+
"source": [
118+
"clients.environment_variables.server_endpoint = server.endpoint # Server address.\n",
119+
"study_client = clients.Study.from_study_config(study_config, owner='owner', study_id = 'example_study_id') # Now connects to the explicitly created server.\n",
120+
"another_study_client = clients.Study.from_resource_name(study_client.resource_name) # Another way to fork clients."
121+
]
122+
},
123+
{
124+
"cell_type": "markdown",
125+
"metadata": {
126+
"id": "Vh3eNsrAdaMJ"
127+
},
128+
"source": [
129+
"## Distributed suggestions\n",
130+
"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."
131+
]
132+
},
133+
{
134+
"cell_type": "code",
135+
"execution_count": null,
136+
"metadata": {
137+
"id": "BnFKc7FadkJV"
138+
},
139+
"outputs": [],
140+
"source": [
141+
"def tuning_loop(client_id: str):\n",
142+
" for i in range(10):\n",
143+
" suggestions = study_client.suggest(count=1, client_id=client_id)\n",
144+
" for suggestion in suggestions:\n",
145+
" objective = evaluate(suggestion.parameters['x'])\n",
146+
" final_measurement = vz.Measurement({'metric': objective})\n",
147+
" suggestion.complete(final_measurement)"
148+
]
149+
},
150+
{
151+
"cell_type": "markdown",
152+
"metadata": {
153+
"id": "NVGcVEzb0Gxe"
154+
},
155+
"source": [
156+
"For example, we may perform a threadpool and construct multiple clients to parallelize evaluations on a single machine."
157+
]
158+
},
159+
{
160+
"cell_type": "code",
161+
"execution_count": null,
162+
"metadata": {
163+
"id": "R0pcPViUz9zC"
164+
},
165+
"outputs": [],
166+
"source": [
167+
"NUM_CLIENTS = 10\n",
168+
"NUM_TRIALS_PER_CLIENT = 50\n",
169+
"\n",
170+
"pool = multiprocessing.pool.ThreadPool(NUM_CLIENTS)\n",
171+
"pool.map(tuning_loop, range(NUM_CLIENTS))"
172+
]
173+
}
174+
],
175+
"metadata": {
176+
"colab": {
177+
"last_runtime": {
178+
"build_target": "//ads/thresholds/kumamon/colab:notebook",
179+
"kind": "shared"
180+
},
181+
"name": "Distributed.ipynb",
182+
"private_outputs": true,
183+
"provenance": [
184+
{
185+
"file_id": "/piper/depot/http://github.com/google/vizier/tree/main/vizier/docs/guides/user/running_vizier.ipynb",
186+
"timestamp": 1673247218127
187+
},
188+
{
189+
"file_id": "1q87rsDDUJLHci3o9Gv-sU0g7H3O3lAbU",
190+
"timestamp": 1659555396142
191+
}
192+
]
193+
},
194+
"kernelspec": {
195+
"display_name": "Python 3",
196+
"name": "python3"
197+
},
198+
"language_info": {
199+
"name": "python"
200+
}
201+
},
202+
"nbformat": 4,
203+
"nbformat_minor": 0
204+
}

docs/guides/user/running_vizier.ipynb

Lines changed: 2 additions & 51 deletions
Original file line numberDiff line numberDiff line change
@@ -80,9 +80,7 @@
8080
"problem = vz.ProblemStatement()\n",
8181
"problem.search_space.root.add_float_param('x', 0.0, 1.0)\n",
8282
"problem.search_space.root.add_float_param('y', 0.0, 1.0)\n",
83-
"problem.metric_information.append(\n",
84-
" vz.MetricInformation(\n",
85-
" name='maximize_metric', goal=vz.ObjectiveMetricGoal.MAXIMIZE))\n",
83+
"problem.metric_information.append(vz.MetricInformation(name='maximize_metric', goal=vz.ObjectiveMetricGoal.MAXIMIZE))\n",
8684
"\n",
8785
"\n",
8886
"def evaluate(x: float, y: float) -\u003e float:\n",
@@ -118,9 +116,7 @@
118116
},
119117
"source": [
120118
"## Setting up the client\n",
121-
"Starts a `study_client`, which can be either in **local mode (default)** or **distributed mode.**\n",
122-
"\n",
123-
"**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`."
119+
"Starts a `study_client`, which 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`."
124120
]
125121
},
126122
{
@@ -135,51 +131,6 @@
135131
"print('Local SQL database file located at: ', service.VIZIER_DB_PATH)"
136132
]
137133
},
138-
{
139-
"cell_type": "markdown",
140-
"metadata": {
141-
"id": "w3m48cPsXcxD"
142-
},
143-
"source": [
144-
"**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.\n",
145-
"\n",
146-
"All client processes (on a single machine or over multiple machines) will connect to this server via a globally specified `endpoint`."
147-
]
148-
},
149-
{
150-
"cell_type": "code",
151-
"execution_count": null,
152-
"metadata": {
153-
"id": "V6ef6OfMXdpz"
154-
},
155-
"outputs": [],
156-
"source": [
157-
"server = servers.DefaultVizierServer() # Ideally created on a separate process such as a server machine.\n",
158-
"clients.environment_variables.server_endpoint = server.endpoint # Server address.\n",
159-
"study_client = clients.Study.from_study_config(study_config, owner='owner', study_id = 'example_study_id') # Now connects to the explicitly created server."
160-
]
161-
},
162-
{
163-
"cell_type": "markdown",
164-
"metadata": {
165-
"id": "Z0Ycmc-exzqm"
166-
},
167-
"source": [
168-
"## Client Parallelization\n",
169-
"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."
170-
]
171-
},
172-
{
173-
"cell_type": "code",
174-
"execution_count": null,
175-
"metadata": {
176-
"id": "VlfFb5t3yILl"
177-
},
178-
"outputs": [],
179-
"source": [
180-
"another_study_client = clients.Study.from_resource_name(study_client.resource_name)"
181-
]
182-
},
183134
{
184135
"cell_type": "markdown",
185136
"metadata": {

vizier/_src/benchmarks/experimenters/infeasible_experimenter_test.py

Lines changed: 0 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -38,9 +38,6 @@ def test_consistency(self):
3838

3939
for t in trials:
4040
self.assertEqual(t.infeasible, trials[0].infeasible)
41-
self.assertEqual(
42-
t.final_measurement_or_die, trials[0].final_measurement_or_die
43-
)
4441

4542

4643
class ParamRegionInfeasibleExperimenterTest(absltest.TestCase):

vizier/_src/service/clients.py

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,8 @@
3232
# Redeclared since clients.py is the user-facing client API.
3333
environment_variables = vizier_client.environment_variables
3434

35+
UNUSED_CLIENT_ID = constants.UNUSED_CLIENT_ID
36+
3537

3638
@attr.define
3739
class Trial(client_abc.TrialInterface):
@@ -204,7 +206,7 @@ def materialize_state(self) -> vz.StudyState:
204206

205207
@classmethod
206208
def from_resource_name(cls: Type['Study'], name: str) -> 'Study':
207-
client = vizier_client.VizierClient(name, constants.UNUSED_CLIENT_ID)
209+
client = vizier_client.VizierClient(name, UNUSED_CLIENT_ID)
208210
try:
209211
_ = client.get_study_config() # Make sure study exists.
210212
except Exception as err:
@@ -251,7 +253,7 @@ def from_study_config(
251253
# expected to provide a client_id in the suggest() call.
252254
client = vizier_client.create_or_load_study(
253255
owner_id=owner,
254-
client_id=constants.UNUSED_CLIENT_ID,
256+
client_id=UNUSED_CLIENT_ID,
255257
study_id=study_id,
256258
study_config=config,
257259
)

vizier/service/clients/__init__.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,3 +21,4 @@
2121
from vizier._src.service.clients import Study
2222
from vizier._src.service.clients import Trial
2323
from vizier._src.service.clients import TrialIterable
24+
from vizier._src.service.clients import UNUSED_CLIENT_ID

0 commit comments

Comments
 (0)