Skip to content

Commit

Permalink
chore: fix tensorboard cleanup in system tests
Browse files Browse the repository at this point in the history
PiperOrigin-RevId: 510180069
  • Loading branch information
sararob authored and copybara-github committed Feb 16, 2023
1 parent ef350be commit 40749af
Show file tree
Hide file tree
Showing 2 changed files with 4 additions and 12 deletions.
10 changes: 1 addition & 9 deletions tests/system/aiplatform/e2e_base.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
)

Expand All @@ -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}")
6 changes: 3 additions & 3 deletions tests/system/aiplatform/test_autologging.py
Original file line number Diff line number Diff line change
Expand Up @@ -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()

Expand Down

0 comments on commit 40749af

Please sign in to comment.