Fix sagemaker experiments cleanup (apache#41287)

We had some experiment cleanup logic, for cleaning up the experiment created by the create experiment task, but more experiments are created as side-effects of other tasks (such as creating pipelines and tuning jobs). These changes clean up the remaining experiments.
uzhastik · Aug 6, 2024 · f648129 · f648129
1 parent 91ea7fc
commit f648129
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 12 deletions.
diff --git a/tests/system/providers/amazon/aws/example_sagemaker.py b/tests/system/providers/amazon/aws/example_sagemaker.py
@@ -432,18 +432,19 @@ def delete_model_group(group_name, model_version_arn):
 
 
 @task(trigger_rule=TriggerRule.ALL_DONE)
-def delete_experiment(name):
+def delete_experiments(experiment_names):
     sgmk_client = boto3.client("sagemaker")
-    trials = sgmk_client.list_trials(ExperimentName=name)
-    trials_names = [s["TrialName"] for s in trials["TrialSummaries"]]
-    for trial in trials_names:
-        components = sgmk_client.list_trial_components(TrialName=trial)
-        components_names = [s["TrialComponentName"] for s in components["TrialComponentSummaries"]]
-        for component in components_names:
-            sgmk_client.disassociate_trial_component(TrialComponentName=component, TrialName=trial)
-            sgmk_client.delete_trial_component(TrialComponentName=component)
-        sgmk_client.delete_trial(TrialName=trial)
-    sgmk_client.delete_experiment(ExperimentName=name)
+    for experiment in experiment_names:
+        trials = sgmk_client.list_trials(ExperimentName=experiment)
+        trials_names = [s["TrialName"] for s in trials["TrialSummaries"]]
+        for trial in trials_names:
+            components = sgmk_client.list_trial_components(TrialName=trial)
+            components_names = [s["TrialComponentName"] for s in components["TrialComponentSummaries"]]
+            for component in components_names:
+                sgmk_client.disassociate_trial_component(TrialComponentName=component, TrialName=trial)
+                sgmk_client.delete_trial_component(TrialComponentName=component)
+            sgmk_client.delete_trial(TrialName=trial)
+        sgmk_client.delete_experiment(ExperimentName=experiment)
 
 
 @task(trigger_rule=TriggerRule.ALL_DONE)
@@ -637,7 +638,13 @@ def delete_docker_image(image_name):
         delete_model_group(test_setup["model_package_group_name"], register_model.output),
         delete_model,
         delete_bucket,
-        delete_experiment(test_setup["experiment_name"]),
+        delete_experiments(
+            [
+                test_setup["experiment_name"],
+                f"{test_setup['auto_ml_job_name']}-aws-auto-ml-job",
+                f"{test_setup['tuning_job_name']}-aws-tuning-job",
+            ]
+        ),
         delete_docker_image(test_setup["docker_image"]),
         log_cleanup,
     )

diff --git a/tests/system/providers/amazon/aws/example_sagemaker_pipeline.py b/tests/system/providers/amazon/aws/example_sagemaker_pipeline.py
@@ -31,6 +31,7 @@
     SageMakerPipelineSensor,
 )
 from airflow.utils.trigger_rule import TriggerRule
+from tests.system.providers.amazon.aws.example_sagemaker import delete_experiments
 from tests.system.providers.amazon.aws.utils import ENV_ID_KEY, SystemTestContextBuilder
 
 DAG_ID = "example_sagemaker_pipeline"
@@ -112,6 +113,7 @@ def delete_pipeline(name: str):
         await_pipeline2,
         # TEST TEARDOWN
         delete_pipeline(pipeline_name),
+        delete_experiments([pipeline_name]),
     )
 
     from tests.system.utils.watcher import watcher