This is an automated email from the ASF dual-hosted git repository.
ferruzzi 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 225a439af18 Add Stop AutoML Job to the Sagemaker system test to clean
up. (#49325)
225a439af18 is described below
commit 225a439af186067aec104a66ee417b191cc60be9
Author: D. Ferruzzi <[email protected]>
AuthorDate: Wed Apr 16 09:13:49 2025 -0700
Add Stop AutoML Job to the Sagemaker system test to clean up. (#49325)
---
.../tests/system/amazon/aws/example_sagemaker.py | 53 ++++++++++++++++++++++
1 file changed, 53 insertions(+)
diff --git a/providers/amazon/tests/system/amazon/aws/example_sagemaker.py
b/providers/amazon/tests/system/amazon/aws/example_sagemaker.py
index d8ad58812f8..7e4cb4e781b 100644
--- a/providers/amazon/tests/system/amazon/aws/example_sagemaker.py
+++ b/providers/amazon/tests/system/amazon/aws/example_sagemaker.py
@@ -24,6 +24,7 @@ from tempfile import NamedTemporaryFile
from textwrap import dedent
import boto3
+from botocore.exceptions import ClientError
from airflow.decorators import task
from airflow.models.baseoperator import chain
@@ -107,6 +108,39 @@ PREPROCESS_SCRIPT_TEMPLATE = dedent("""
""")
+def _install_aws_cli_if_needed():
+ """
+ Check if the AWS CLI tool is installed and install it if needed.
+
+ The AmazonLinux image has flip-flopped a couple of times on whether this
is included in the base image
+ or not, so to future-proof this we are going to check if it's installed
and install if necessary.
+ """
+ check = subprocess.Popen(
+ "aws --version",
+ shell=True,
+ stdout=subprocess.PIPE,
+ stderr=subprocess.PIPE,
+ )
+ _, stderr = check.communicate()
+
+ if check.returncode == 0:
+ logger.info("AWS CLI tool is installed.")
+ return
+
+ if "aws: not found" in str(stderr):
+ logger.info("AWS CLI tool not found; installing.")
+ subprocess.Popen(
+ """
+ curl
"https://awscli.amazonaws.com/awscli-exe-linux-x86_64.zip"; -o "awscliv2.zip"
+ unzip awscliv2.zip
+ sudo ./aws/install
+ """,
+ shell=True,
+ stdout=subprocess.PIPE,
+ stderr=subprocess.PIPE,
+ ).communicate()
+
+
def _create_ecr_repository(repo_name):
execution_role_arn = boto3.client("sts").get_caller_identity()["Arn"]
access_policy = {
@@ -383,6 +417,7 @@ def set_up(env_id, role_arn):
preprocess_script = PREPROCESS_SCRIPT_TEMPLATE.format(
input_path=processing_local_input_path,
output_path=processing_local_output_path
)
+ _install_aws_cli_if_needed()
_build_and_upload_docker_image(preprocess_script, ecr_repository_uri)
ti = get_current_context()["ti"]
@@ -409,6 +444,23 @@ def set_up(env_id, role_arn):
ti.xcom_push(key="transform_job_name", value=transform_job_name)
+@task(trigger_rule=TriggerRule.ALL_DONE)
+def stop_automl_job(job_name: str):
+ try:
+ logger.info("Stopping AutoML job: %s", job_name)
+ boto3.client("sagemaker").stop_auto_ml_job(AutoMLJobName=job_name)
+ except ClientError as e:
+ # If the job has already completed, boto will raise a
ValidationException.
+ # In this case, consider that a successful result.
+ if (
+ e.response["Error"]["Code"] == "ValidationException"
+ and "already reached a terminal state" in
e.response["Error"]["Message"]
+ ):
+ logger.info("AutoML job %s already completed.", job_name)
+ else:
+ raise e
+
+
@task(trigger_rule=TriggerRule.ALL_DONE)
def delete_ecr_repository(repository_name):
client = boto3.client("ecr")
@@ -648,6 +700,7 @@ with DAG(
delete_model_group(test_setup["model_package_group_name"],
register_model.output),
delete_model,
delete_bucket,
+ stop_automl_job(test_setup["auto_ml_job_name"]),
delete_experiments(
[
test_setup["experiment_name"],
