This is an automated email from the ASF dual-hosted git repository.
onikolas pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new f648129fe2 Fix sagemaker experiments cleanup (#41287)
f648129fe2 is described below
commit f648129fe24958b35dbf3e8bedbeea0e99890675
Author: Niko Oliveira <[email protected]>
AuthorDate: Tue Aug 6 11:50:28 2024 -0700
Fix sagemaker experiments cleanup (#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.
---
.../providers/amazon/aws/example_sagemaker.py | 31 +++++++++++++---------
.../amazon/aws/example_sagemaker_pipeline.py | 2 ++
2 files changed, 21 insertions(+), 12 deletions(-)
diff --git a/tests/system/providers/amazon/aws/example_sagemaker.py
b/tests/system/providers/amazon/aws/example_sagemaker.py
index 8c41b9b543..cd291b7971 100644
--- 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 @@ with DAG(
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
index 48bf92fa7f..3f32511702 100644
--- a/tests/system/providers/amazon/aws/example_sagemaker_pipeline.py
+++ b/tests/system/providers/amazon/aws/example_sagemaker_pipeline.py
@@ -31,6 +31,7 @@ from airflow.providers.amazon.aws.sensors.sagemaker import (
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 @@ with DAG(
await_pipeline2,
# TEST TEARDOWN
delete_pipeline(pipeline_name),
+ delete_experiments([pipeline_name]),
)
from tests.system.utils.watcher import watcher
