Skip to content

Commit 2a5955f

Browse files
vertex-sdk-botcopybara-github
vertex-sdk-bot
authored and
copybara-github
committed
fix: Fix experiments failure when backing tensorboard has been deleted.
PiperOrigin-RevId: 592351876
1 parent e8c5057 commit 2a5955f

5 files changed

Lines changed: 106 additions & 19 deletions

File tree

‎google/cloud/aiplatform/initializer.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -229,6 +229,7 @@ def init(
229229
experiment=experiment,
230230
description=experiment_description,
231231
backing_tensorboard=experiment_tensorboard,
232+
credentials=credentials,
232233
)
233234

234235
def get_encryption_spec(

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

Lines changed: 11 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -502,7 +502,7 @@ def _lookup_backing_tensorboard(self) -> Optional[tensorboard_resource.Tensorboa
502502
"""Returns backing tensorboard if one is set.
503503
504504
Returns:
505-
Tensorboard resource if one exists.
505+
Tensorboard resource if one exists, otherwise returns None.
506506
"""
507507
tensorboard_resource_name = self._metadata_context.metadata.get(
508508
constants._BACKING_TENSORBOARD_RESOURCE_KEY
@@ -516,10 +516,16 @@ def _lookup_backing_tensorboard(self) -> Optional[tensorboard_resource.Tensorboa
516516
)
517517

518518
if tensorboard_resource_name:
519-
return tensorboard_resource.Tensorboard(
520-
tensorboard_resource_name,
521-
credentials=self._metadata_context.credentials,
522-
)
519+
try:
520+
return tensorboard_resource.Tensorboard(
521+
tensorboard_resource_name,
522+
credentials=self._metadata_context.credentials,
523+
)
524+
except exceptions.NotFound:
525+
self._metadata_context.update(
526+
metadata={constants._BACKING_TENSORBOARD_RESOURCE_KEY: None}
527+
)
528+
return None
523529

524530
def get_backing_tensorboard_resource(
525531
self,

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

Lines changed: 23 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -229,6 +229,25 @@ def reset(self):
229229
self._experiment = None
230230
self._experiment_run = None
231231

232+
def _get_global_tensorboard(
233+
self,
234+
credentials: Optional[auth_credentials.Credentials] = None,
235+
) -> Optional[tensorboard_resource.Tensorboard]:
236+
"""Helper method to get the global TensorBoard instance.
237+
238+
Returns:
239+
tensorboard_resource.Tensorboard: the global TensorBoard instance.
240+
"""
241+
if self._global_tensorboard:
242+
try:
243+
return tensorboard_resource.Tensorboard(
244+
self._global_tensorboard.resource_name,
245+
credentials=credentials,
246+
)
247+
except exceptions.NotFound:
248+
self._global_tensorboard = None
249+
return None
250+
232251
@property
233252
def experiment_name(self) -> Optional[str]:
234253
"""Return the currently set experiment name, if experiment is not set, return None"""
@@ -269,6 +288,7 @@ def set_experiment(
269288
backing_tensorboard: Optional[
270289
Union[str, tensorboard_resource.Tensorboard, bool]
271290
] = None,
291+
credentials: Optional[auth_credentials.Credentials] = None,
272292
):
273293
"""Set the experiment. Will retrieve the Experiment if it exists or create one with the provided name.
274294
@@ -286,6 +306,8 @@ def set_experiment(
286306
287307
To disable using a backign tensorboard, set `backing_tensorboard` to `False`.
288308
To maintain this behavior, set `experiment_tensorboard` to `False` in subsequent calls to aiplatform.init().
309+
credentials (auth_credentials.Credentials):
310+
Optional. Custom credentials used to set this Tensorboard resource.
289311
"""
290312
self.reset()
291313

@@ -299,7 +321,7 @@ def set_experiment(
299321
backing_tb = None
300322
else:
301323
backing_tb = (
302-
self._global_tensorboard or _get_or_create_default_tensorboard()
324+
self._get_global_tensorboard(credentials) or _get_or_create_default_tensorboard()
303325
)
304326

305327
current_backing_tb = experiment.backing_tensorboard_resource_name

‎google/cloud/aiplatform/tensorboard/uploader_tracker.py‎

Lines changed: 13 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -236,24 +236,24 @@ def _create_uploader(
236236
project, location, tensorboard_id
237237
)
238238
else:
239-
if _experiment_tracker._global_tensorboard:
239+
credentials = None if not _experiment_tracker.experiment else _experiment_tracker.experiment._metadata_context.credentials
240+
if _experiment_tracker._get_global_tensorboard(credentials):
240241
tensorboard_resource_name = (
241-
_experiment_tracker._global_tensorboard.resource_name
242+
_experiment_tracker._get_global_tensorboard(credentials).resource_name
242243
)
243-
else:
244-
if _experiment_tracker._experiment:
245-
if _experiment_tracker._experiment._lookup_backing_tensorboard():
246-
tensorboard_resource_name = (
247-
_experiment_tracker._experiment._lookup_backing_tensorboard().resource_name
248-
)
249-
else:
250-
raise ValueError(
251-
f"No TensorBoard associated with experiment {initializer.global_config.experiment_name}. Please provide tensorboard_id in the argument."
252-
)
244+
elif _experiment_tracker._experiment:
245+
if _experiment_tracker._experiment._lookup_backing_tensorboard():
246+
tensorboard_resource_name = (
247+
_experiment_tracker._experiment._lookup_backing_tensorboard().resource_name
248+
)
253249
else:
254250
raise ValueError(
255-
"No TensorBoard found. Please provide tensorboard_id in the argument."
251+
f"No TensorBoard associated with experiment {initializer.global_config.experiment_name}. Please provide tensorboard_id in the argument."
256252
)
253+
else:
254+
raise ValueError(
255+
"No TensorBoard found. Please provide tensorboard_id in the argument."
256+
)
257257

258258
api_client = initializer.global_config.create_client(
259259
client_class=TensorboardClientWithOverride,

‎tests/system/aiplatform/test_experiments.py‎

Lines changed: 58 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -618,3 +618,61 @@ def test_init_associates_global_tensorboard_to_experiment(self, shared_state):
618618
)
619619
== tensorboard.resource_name
620620
)
621+
622+
def test_get_backing_tensorboard_resource_returns_tensorboard(self, shared_state):
623+
tensorboard = aiplatform.Tensorboard.create(
624+
project=e2e_base._PROJECT,
625+
location=e2e_base._LOCATION,
626+
display_name=self._make_display_name("")[:64],
627+
)
628+
shared_state["resources"] = [tensorboard]
629+
aiplatform.init(
630+
project=e2e_base._PROJECT,
631+
location=e2e_base._LOCATION,
632+
experiment=self._experiment_name,
633+
experiment_tensorboard=tensorboard,
634+
)
635+
experiment = aiplatform.Experiment(
636+
self._experiment_name,
637+
project=e2e_base._PROJECT,
638+
location=e2e_base._LOCATION,
639+
)
640+
641+
assert (
642+
experiment.get_backing_tensorboard_resource().resource_name
643+
== tensorboard.resource_name
644+
)
645+
646+
def test_get_backing_tensorboard_resource_returns_none(self):
647+
new_experiment_name = "new-experiment"
648+
aiplatform.init(
649+
project=e2e_base._PROJECT,
650+
location=e2e_base._LOCATION,
651+
experiment=new_experiment_name,
652+
experiment_tensorboard=False,
653+
)
654+
experiment = aiplatform.Experiment(
655+
new_experiment_name,
656+
project=e2e_base._PROJECT,
657+
location=e2e_base._LOCATION,
658+
)
659+
660+
assert experiment.get_backing_tensorboard_resource() is None
661+
662+
def test_delete_backing_tensorboard_experiment_run_success(self):
663+
aiplatform.init(
664+
project=e2e_base._PROJECT,
665+
location=e2e_base._LOCATION,
666+
experiment=self._experiment_name,
667+
)
668+
experiment = aiplatform.Experiment(
669+
self._experiment_name,
670+
project=e2e_base._PROJECT,
671+
location=e2e_base._LOCATION,
672+
)
673+
experiment.get_backing_tensorboard_resource().delete()
674+
run = aiplatform.start_run(_RUN)
675+
aiplatform.end_run()
676+
677+
assert experiment.get_backing_tensorboard_resource() is None
678+
assert run.name == _RUN

0 commit comments

Comments
 (0)