Skip to content

Commit a25c669

Browse files
vertex-sdk-botcopybara-github
vertex-sdk-bot
authored and
copybara-github
committed
feat: Add option to not use default tensorboard
PiperOrigin-RevId: 581402305
1 parent 45d599b commit a25c669

File tree

4 files changed

+109
-21
lines changed

4 files changed

+109
-21
lines changed

‎google/cloud/aiplatform/initializer.py‎

Lines changed: 10 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -112,7 +112,7 @@ def init(
112112
experiment: Optional[str] = None,
113113
experiment_description: Optional[str] = None,
114114
experiment_tensorboard: Optional[
115-
Union[str, tensorboard_resource.Tensorboard]
115+
Union[str, tensorboard_resource.Tensorboard, bool]
116116
] = None,
117117
staging_bucket: Optional[str] = None,
118118
credentials: Optional[auth_credentials.Credentials] = None,
@@ -128,7 +128,7 @@ def init(
128128
set defaults to us-central-1.
129129
experiment (str): Optional. The experiment name.
130130
experiment_description (str): Optional. The description of the experiment.
131-
experiment_tensorboard (Union[str, tensorboard_resource.Tensorboard]):
131+
experiment_tensorboard (Union[str, tensorboard_resource.Tensorboard, bool]):
132132
Optional. The Vertex AI TensorBoard instance, Tensorboard resource name,
133133
or Tensorboard resource ID to use as a backing Tensorboard for the provided
134134
experiment.
@@ -141,6 +141,13 @@ def init(
141141
Any subsequent calls to aiplatform.init() with `experiment` and without
142142
`experiment_tensorboard` will automatically assign the global Tensorboard
143143
to the `experiment`.
144+
145+
If `experiment_tensorboard` is ommitted or set to `True` or `None` the global
146+
Tensorboard will be assigned to the `experiment`. If a global Tensorboard is
147+
not set, the default Tensorboard instance will be used, and created if it deos not exist.
148+
149+
To disable creating and using Tensorboard with `experiment`, set `experiment_tensorboard` to False.
150+
Any subsequent calls to aiplatform.init() should include this setting as well.
144151
staging_bucket (str): The default staging bucket to use to stage artifacts
145152
when making API calls. In the form gs://...
146153
credentials (google.auth.credentials.Credentials): The default custom
@@ -177,7 +184,7 @@ def init(
177184
"Experiment needs to be set in `init` in order to add experiment descriptions."
178185
)
179186

180-
if experiment_tensorboard:
187+
if experiment_tensorboard and not isinstance(experiment_tensorboard, bool):
181188
metadata._experiment_tracker.set_tensorboard(
182189
tensorboard=experiment_tensorboard,
183190
project=project,

‎google/cloud/aiplatform/metadata/metadata.py‎

Lines changed: 17 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -267,7 +267,7 @@ def set_experiment(
267267
*,
268268
description: Optional[str] = None,
269269
backing_tensorboard: Optional[
270-
Union[str, tensorboard_resource.Tensorboard]
270+
Union[str, tensorboard_resource.Tensorboard, bool]
271271
] = None,
272272
):
273273
"""Set the experiment. Will retrieve the Experiment if it exists or create one with the provided name.
@@ -277,21 +277,30 @@ def set_experiment(
277277
Required. Name of the experiment to set.
278278
description (str):
279279
Optional. Description of an experiment.
280-
backing_tensorboard Union[str, aiplatform.Tensorboard]:
280+
backing_tensorboard Union[str, aiplatform.Tensorboard, bool]:
281281
Optional. If provided, assigns tensorboard as backing tensorboard to support time series metrics
282282
logging.
283+
284+
If ommitted, or set to `True` or `None`, the global tensorboard is used.
285+
If no global tensorboard is set, the default tensorboard will be used, and created if it does not exist.
286+
287+
To disable using a backign tensorboard, set `backing_tensorboard` to `False`.
288+
To maintain this behavior, set `experiment_tensorboard` to `False` in subsequent calls to aiplatform.init().
283289
"""
284290
self.reset()
285291

286292
experiment = experiment_resources.Experiment.get_or_create(
287293
experiment_name=experiment, description=description
288294
)
289295

290-
backing_tb = (
291-
backing_tensorboard
292-
or self._global_tensorboard
293-
or _get_or_create_default_tensorboard()
294-
)
296+
if backing_tensorboard and not isinstance(backing_tensorboard, bool):
297+
backing_tb = backing_tensorboard
298+
elif isinstance(backing_tensorboard, bool) and not backing_tensorboard:
299+
backing_tb = None
300+
else:
301+
backing_tb = (
302+
self._global_tensorboard or _get_or_create_default_tensorboard()
303+
)
295304

296305
current_backing_tb = experiment.backing_tensorboard_resource_name
297306

@@ -322,7 +331,7 @@ def set_tensorboard(
322331
credentials (auth_credentials.Credentials):
323332
Optional. Custom credentials used to set this Tensorboard resource.
324333
"""
325-
if isinstance(tensorboard, str):
334+
if tensorboard and isinstance(tensorboard, str):
326335
tensorboard = tensorboard_resource.Tensorboard(
327336
tensorboard,
328337
project=project,

‎tests/unit/aiplatform/test_initializer.py‎

Lines changed: 31 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -114,7 +114,9 @@ def test_init_service_account_sets_service_account(self):
114114
def test_init_experiment_sets_experiment(self, set_experiment_mock):
115115
initializer.global_config.init(experiment=_TEST_EXPERIMENT)
116116
set_experiment_mock.assert_called_once_with(
117-
experiment=_TEST_EXPERIMENT, description=None, backing_tensorboard=None
117+
experiment=_TEST_EXPERIMENT,
118+
description=None,
119+
backing_tensorboard=None,
118120
)
119121

120122
@patch.object(_experiment_tracker, "set_experiment")
@@ -162,25 +164,44 @@ def test_init_with_experiment_tensorboard_resource_sets_global_tensorboard(
162164
credentials=None,
163165
)
164166

165-
@patch.object(_experiment_tracker, "set_tensorboard")
166167
@patch.object(_experiment_tracker, "set_experiment")
167168
def test_init_experiment_without_tensorboard_uses_global_tensorboard(
168-
self,
169-
set_tensorboard_mock,
170-
set_experiment_mock,
169+
self, set_experiment_mock
171170
):
172171

173-
initializer.global_config.init(experiment_tensorboard=_TEST_TENSORBOARD_NAME)
172+
initializer.global_config.tensorboard = _TEST_TENSORBOARD_NAME
174173

175174
initializer.global_config.init(
176175
experiment=_TEST_EXPERIMENT,
177176
)
178177

179178
set_experiment_mock.assert_called_once_with(
180-
tensorboard=_TEST_TENSORBOARD_NAME,
181-
project=None,
182-
location=None,
183-
credentials=None,
179+
experiment=_TEST_EXPERIMENT,
180+
description=None,
181+
backing_tensorboard=None,
182+
)
183+
184+
assert initializer.global_config.tensorboard == _TEST_TENSORBOARD_NAME
185+
186+
@patch.object(_experiment_tracker, "set_tensorboard")
187+
@patch.object(_experiment_tracker, "set_experiment")
188+
def test_init_experiment_tensorboard_false_does_not_set_tensorboard(
189+
self, set_experiment_mock, set_tensorboard_mock
190+
):
191+
192+
initializer.global_config.tensorboard = _TEST_TENSORBOARD_NAME
193+
194+
initializer.global_config.init(
195+
experiment=_TEST_EXPERIMENT,
196+
experiment_tensorboard=False,
197+
)
198+
199+
set_tensorboard_mock.assert_not_called()
200+
201+
set_experiment_mock.assert_called_once_with(
202+
experiment=_TEST_EXPERIMENT,
203+
description=None,
204+
backing_tensorboard=False,
184205
)
185206

186207
def test_init_experiment_description_fail_without_experiment(self):

‎tests/unit/aiplatform/test_metadata.py‎

Lines changed: 51 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1140,6 +1140,57 @@ def test_init_experiment_with_default_tensorboard(
11401140
)
11411141
assign_backing_tensorboard_mock.assert_called_once()
11421142

1143+
@pytest.mark.usefixtures(
1144+
"get_metadata_store_mock",
1145+
"get_experiment_run_run_mock",
1146+
)
1147+
def test_init_experiment_tensorboard_false_doesNotSet_backing_tensorboard(
1148+
self, list_default_tensorboard_mock, assign_backing_tensorboard_mock
1149+
):
1150+
aiplatform.init(
1151+
project=_TEST_PROJECT,
1152+
location=_TEST_LOCATION,
1153+
experiment=_TEST_EXPERIMENT,
1154+
experiment_tensorboard=False,
1155+
)
1156+
1157+
list_default_tensorboard_mock.assert_not_called()
1158+
assign_backing_tensorboard_mock.assert_not_called()
1159+
1160+
@pytest.mark.usefixtures(
1161+
"get_metadata_store_mock",
1162+
"get_experiment_run_run_mock",
1163+
)
1164+
def test_init_experiment_tensorboard_true_sets_backing_tensorboard(
1165+
self, list_default_tensorboard_mock, assign_backing_tensorboard_mock
1166+
):
1167+
aiplatform.init(
1168+
project=_TEST_PROJECT,
1169+
location=_TEST_LOCATION,
1170+
experiment=_TEST_EXPERIMENT,
1171+
experiment_tensorboard=True,
1172+
)
1173+
1174+
list_default_tensorboard_mock.assert_called()
1175+
assign_backing_tensorboard_mock.assert_called()
1176+
1177+
@pytest.mark.usefixtures(
1178+
"get_metadata_store_mock",
1179+
"get_experiment_run_run_mock",
1180+
)
1181+
def test_init_experiment_tensorboard_none_sets_backing_tensorboard(
1182+
self, list_default_tensorboard_mock, assign_backing_tensorboard_mock
1183+
):
1184+
aiplatform.init(
1185+
project=_TEST_PROJECT,
1186+
location=_TEST_LOCATION,
1187+
experiment=_TEST_EXPERIMENT,
1188+
experiment_tensorboard=None,
1189+
)
1190+
1191+
list_default_tensorboard_mock.assert_called()
1192+
assign_backing_tensorboard_mock.assert_called()
1193+
11431194
@pytest.mark.usefixtures("get_metadata_store_mock")
11441195
def test_create_experiment(self, create_experiment_context_mock):
11451196
exp = aiplatform.Experiment.create(

0 commit comments

Comments
 (0)