From 40749af5684b83d217da0823b6e6db478c7d8e2a Mon Sep 17 00:00:00 2001 From: Sara Robinson Date: Thu, 16 Feb 2023 10:22:26 -0800 Subject: [PATCH] chore: fix tensorboard cleanup in system tests PiperOrigin-RevId: 510180069 --- tests/system/aiplatform/e2e_base.py | 10 +--------- tests/system/aiplatform/test_autologging.py | 6 +++--- 2 files changed, 4 insertions(+), 12 deletions(-) diff --git a/tests/system/aiplatform/e2e_base.py b/tests/system/aiplatform/e2e_base.py index 10fb76560f..9ff3eb5ee8 100644 --- a/tests/system/aiplatform/e2e_base.py +++ b/tests/system/aiplatform/e2e_base.py @@ -179,6 +179,7 @@ def tear_down_resources(self, shared_state: Dict[str, Any]): key=lambda r: 1 if isinstance(r, aiplatform.Endpoint) or isinstance(r, aiplatform.MatchingEngineIndexEndpoint) + or isinstance(r, aiplatform.Experiment) else 2 ) @@ -201,12 +202,3 @@ def tear_down_resources(self, shared_state: Dict[str, Any]): resource.delete() except exceptions.GoogleAPIError as e: logging.error(f"Could not delete resource: {resource} due to: {e}") - - # When an Experiment has a backing_tensorboard, the Experiment needs to be deleted first - # This is used by the autologging tests - if "tensorboard" in shared_state: - for resource in shared_state["tensorboard"]: - try: - resource.delete() - except exceptions.GoogleAPIError as e: - logging.error(f"Could not delete resource: {resource} due to: {e}") diff --git a/tests/system/aiplatform/test_autologging.py b/tests/system/aiplatform/test_autologging.py index bfeea298d0..bae63f8440 100644 --- a/tests/system/aiplatform/test_autologging.py +++ b/tests/system/aiplatform/test_autologging.py @@ -152,11 +152,11 @@ def test_autologging_with_autorun_creation(self, shared_state): experiment_tensorboard=self._backing_tensorboard, ) - shared_state["tensorboard"] = [self._backing_tensorboard] + shared_state["resources"] = [self._backing_tensorboard] - shared_state["resources"] = [ + shared_state["resources"].append( aiplatform.metadata.metadata._experiment_tracker.experiment - ] + ) aiplatform.autolog()