# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.
from __future__ import annotations

import json
import logging
import subprocess
from datetime import datetime
from tempfile import NamedTemporaryFile
from textwrap import dedent

import boto3
from botocore.exceptions import ClientError

from airflow.providers.amazon.aws.operators.s3 import (
    S3CreateBucketOperator,
    S3CreateObjectOperator,
    S3DeleteBucketOperator,
)
from airflow.providers.amazon.aws.operators.sagemaker import (
    SageMakerAutoMLOperator,
    SageMakerCreateExperimentOperator,
    SageMakerDeleteModelOperator,
    SageMakerModelOperator,
    SageMakerProcessingOperator,
    SageMakerRegisterModelVersionOperator,
    SageMakerTrainingOperator,
    SageMakerTransformOperator,
    SageMakerTuningOperator,
)
from airflow.providers.amazon.aws.sensors.sagemaker import (
    SageMakerAutoMLSensor,
    SageMakerProcessingSensor,
    SageMakerTrainingSensor,
    SageMakerTransformSensor,
    SageMakerTuningSensor,
)

from tests_common.test_utils.version_compat import AIRFLOW_V_3_0_PLUS

if AIRFLOW_V_3_0_PLUS:
    from airflow.sdk import DAG, chain, task
else:
    # Airflow 2 path
    from airflow.decorators import task  # type: ignore[attr-defined,no-redef]
    from airflow.models.baseoperator import chain  # type: ignore[attr-defined,no-redef]
    from airflow.models.dag import DAG  # type: ignore[attr-defined,no-redef,assignment]

try:
    from airflow.sdk import TriggerRule
except ImportError:
    # Compatibility for Airflow < 3.1
    from airflow.utils.trigger_rule import TriggerRule  # type: ignore[no-redef,attr-defined]

from system.amazon.aws.utils import ENV_ID_KEY, SystemTestContextBuilder, prune_logs

logger = logging.getLogger(__name__)

DAG_ID = "example_sagemaker"

# Externally fetched variables:
ROLE_ARN_KEY = "ROLE_ARN"

sys_test_context_task = SystemTestContextBuilder().add_variable(ROLE_ARN_KEY).build()

# The URI of a Docker image for handling KNN model training.
# To find the URI of a free Amazon-provided image that can be used, substitute your
# desired region in the following link and find the URI under "Registry Path".
# https://docs.aws.amazon.com/sagemaker/latest/dg/ecr-us-east-1.html#knn-us-east-1.title
# This URI should be in the format of {12-digits}.dkr.ecr.{region}.amazonaws.com/knn
KNN_IMAGES_BY_REGION = {
    "us-east-1": "382416733822.dkr.ecr.us-east-1.amazonaws.com/knn:1",
    "us-west-2": "174872318107.dkr.ecr.us-west-2.amazonaws.com/knn:1",
}

SAMPLE_SIZE = 600

# This script will be the entrypoint for the docker image which will handle preprocessing the raw data
PREPROCESS_SCRIPT_TEMPLATE = dedent("""
    import numpy as np
    import pandas as pd
    import json
    import os

    def main():
        # Load the dataset from {input_path}/input.csv, split it into train/test
        # subsets, and write them to {output_path}/ for the Processing Operator.

        data = pd.read_csv('{input_path}/input.csv')

        # Split into test and train data
        data_train, data_test = np.split(
            data.sample(frac=1, random_state=np.random.RandomState()), [int(0.7 * len(data))]
        )

        # Remove the "answers" from the test set
        data_test.drop(['class'], axis=1, inplace=True)

        # Write the splits to disk
        data_train.to_csv('{output_path}/train.csv', index=False, header=False)
        data_test.to_csv('{output_path}/test.csv', index=False, header=False)

        # Write evaluation metrics for output_files_to_xcom testing
        evaluation_dir = '{output_path}/evaluation'
        os.makedirs(evaluation_dir, exist_ok=True)
        metrics = {{"metrics": {{"mse": {{"value": 5.332878}}, "rmse": {{"value": 2.309302}}}}}}
        with open(os.path.join(evaluation_dir, 'evaluation.json'), 'w') as f:
            json.dump(metrics, f)

        print('Preprocessing Done.')

    if __name__ == "__main__":
        main()
""")


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 = {
        "Version": "2012-10-17",
        "Statement": [
            {
                "Sid": "Allow access to the system test execution role",
                "Effect": "Allow",
                "Principal": {"AWS": execution_role_arn},
                "Action": "ecr:*",
            }
        ],
    }

    client = boto3.client("ecr")
    repo = client.create_repository(repositoryName=repo_name)["repository"]
    client.set_repository_policy(repositoryName=repo["repositoryName"], policyText=json.dumps(access_policy))

    return repo["repositoryUri"]


def _build_and_upload_docker_image(preprocess_script, repository_uri):
    """
    We need a Docker image with the following requirements:
      - Has numpy, pandas, requests, and boto3 installed
      - Has our data preprocessing script mounted and set as the entry point
    """
    with NamedTemporaryFile(mode="w+t") as preprocessing_script, NamedTemporaryFile(mode="w+t") as dockerfile:
        preprocessing_script.write(preprocess_script)
        preprocessing_script.flush()

        dockerfile.write(
            f"""
            FROM public.ecr.aws/amazonlinux/amazonlinux
            COPY {preprocessing_script.name.split("/")[2]} /preprocessing.py
            RUN yum install python3 pip -y
            RUN pip3 install boto3 pandas requests
            CMD [ "python3", "/preprocessing.py"]
            """
        )
        dockerfile.flush()

        ecr_region = repository_uri.split(".")[3]

        docker_build_and_push_commands = f"""
            # login to public ecr repo containing amazonlinux image (public login is always on us east 1)
            aws ecr-public get-login-password --region us-east-1 |
            docker login --username AWS --password-stdin public.ecr.aws &&
            docker build --platform=linux/amd64 -f {dockerfile.name} -t {repository_uri} /tmp &&

            # login again, this time to the private repo we created to hold that specific image
            aws ecr get-login-password --region {ecr_region} |
            docker login --username AWS --password-stdin {repository_uri} &&
            docker push {repository_uri}
            """
        logger.info("building and uploading docker image for preprocessing...")
        docker_build = subprocess.Popen(
            docker_build_and_push_commands,
            shell=True,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
        )
        _, stderr = docker_build.communicate()
        if docker_build.returncode != 0:
            raise RuntimeError(
                "Failed to prepare docker image for the preprocessing job.\n"
                "The following error happened while executing the sequence of bash commands:\n"
                f"{stderr.decode()}"
            )


def generate_data() -> str:
    """generates a very simple csv dataset with headers"""
    content = "class,x,y\n"  # headers
    for i in range(SAMPLE_SIZE):
        content += f"{i % 100},{i},{SAMPLE_SIZE - i}\n"
    return content


@task
def set_up(env_id, role_arn):
    bucket_name = f"{env_id}-sagemaker-example"
    ecr_repository_name = f"{env_id}-repo"
    model_name = f"{env_id}-KNN-model"
    processing_job_name = f"{env_id}-processing"
    training_job_name = f"{env_id}-train"
    transform_job_name = f"{env_id}-transform"
    tuning_job_name = f"{env_id}-tune"
    model_package_group_name = f"{env_id}-group"
    auto_ml_job_name = f"{env_id}-automl"
    experiment_name = f"{env_id}-experiment"

    input_data_S3_key = f"{env_id}/processed-input-data"
    prediction_output_s3_key = f"{env_id}/transform"
    processing_local_input_path = "/opt/ml/processing/input"
    processing_local_output_path = "/opt/ml/processing/output"
    raw_data_s3_key = f"{env_id}/preprocessing/input.csv"
    training_output_s3_key = f"{env_id}/results"

    ecr_repository_uri = _create_ecr_repository(ecr_repository_name)
    region = boto3.session.Session().region_name
    try:
        knn_image_uri = KNN_IMAGES_BY_REGION[region]
    except KeyError:
        raise KeyError(
            f"Region name {region} does not have a known KNN "
            f"Image URI.  Please add the region and URI following "
            f"the directions at the top of the system testfile "
        )

    resource_config = {
        "InstanceCount": 1,
        "InstanceType": "ml.m5.large",
        "VolumeSizeInGB": 1,
    }
    input_data_uri = f"s3://{bucket_name}/{raw_data_s3_key}"
    processing_config = {
        "ProcessingJobName": processing_job_name,
        "ProcessingInputs": [
            {
                "InputName": "input",
                "AppManaged": False,
                "S3Input": {
                    "S3Uri": input_data_uri,
                    "LocalPath": processing_local_input_path,
                    "S3DataType": "S3Prefix",
                    "S3InputMode": "File",
                    "S3DataDistributionType": "FullyReplicated",
                    "S3CompressionType": "None",
                },
            },
        ],
        "ProcessingOutputConfig": {
            "Outputs": [
                {
                    "OutputName": "output",
                    "S3Output": {
                        "S3Uri": f"s3://{bucket_name}/{input_data_S3_key}",
                        "LocalPath": processing_local_output_path,
                        "S3UploadMode": "EndOfJob",
                    },
                    "AppManaged": False,
                },
                {
                    "OutputName": "evaluation",
                    "S3Output": {
                        "S3Uri": f"s3://{bucket_name}/{env_id}/evaluation",
                        "LocalPath": f"{processing_local_output_path}/evaluation",
                        "S3UploadMode": "EndOfJob",
                    },
                    "AppManaged": False,
                },
            ]
        },
        "ProcessingResources": {
            "ClusterConfig": resource_config,
        },
        "StoppingCondition": {"MaxRuntimeInSeconds": 600},
        "AppSpecification": {
            "ImageUri": ecr_repository_uri,
        },
        "RoleArn": role_arn,
    }

    output_files_to_xcom_config = [
        {
            "result_name": "EvaluationReport",
            "output_name": "evaluation",
            "file_name": "evaluation.json",
        }
    ]

    training_data_source = {
        "CompressionType": "None",
        "ContentType": "text/csv",
        "DataSource": {
            "S3DataSource": {
                "S3DataDistributionType": "FullyReplicated",
                "S3DataType": "S3Prefix",
                "S3Uri": f"s3://{bucket_name}/{input_data_S3_key}/train.csv",
            }
        },
    }
    training_config = {
        "AlgorithmSpecification": {
            "TrainingImage": knn_image_uri,
            "TrainingInputMode": "File",
        },
        "HyperParameters": {
            "predictor_type": "classifier",
            "feature_dim": "2",
            "k": "3",
            "sample_size": str(SAMPLE_SIZE),
        },
        "InputDataConfig": [
            {
                "ChannelName": "train",
                **training_data_source,
            }
        ],
        "OutputDataConfig": {"S3OutputPath": f"s3://{bucket_name}/{training_output_s3_key}/"},
        "ExperimentConfig": {"ExperimentName": experiment_name},
        "ResourceConfig": resource_config,
        "RoleArn": role_arn,
        "StoppingCondition": {"MaxRuntimeInSeconds": 600},
        "TrainingJobName": training_job_name,
    }
    model_trained_weights = (
        f"s3://{bucket_name}/{training_output_s3_key}/{training_job_name}/output/model.tar.gz"
    )
    model_config = {
        "ExecutionRoleArn": role_arn,
        "ModelName": model_name,
        "PrimaryContainer": {
            "Mode": "SingleModel",
            "Image": knn_image_uri,
            "ModelDataUrl": model_trained_weights,
        },
    }
    tuning_config = {
        "HyperParameterTuningJobName": tuning_job_name,
        "HyperParameterTuningJobConfig": {
            "Strategy": "Bayesian",
            "HyperParameterTuningJobObjective": {
                "MetricName": "test:accuracy",
                "Type": "Maximize",
            },
            "ResourceLimits": {
                "MaxNumberOfTrainingJobs": 10,
                "MaxParallelTrainingJobs": 10,
            },
            "ParameterRanges": {
                "CategoricalParameterRanges": [],
                "IntegerParameterRanges": [
                    # Set the min and max values of the hyperparameters you want to tune.
                    {
                        "Name": "k",
                        "MinValue": "1",
                        "MaxValue": str(SAMPLE_SIZE),
                    },
                    {
                        "Name": "sample_size",
                        "MinValue": "1",
                        "MaxValue": str(SAMPLE_SIZE),
                    },
                ],
            },
        },
        "TrainingJobDefinition": {
            "StaticHyperParameters": {
                "predictor_type": "classifier",
                "feature_dim": "2",
            },
            "AlgorithmSpecification": {"TrainingImage": knn_image_uri, "TrainingInputMode": "File"},
            "InputDataConfig": [
                {
                    "ChannelName": "train",
                    **training_data_source,
                },
                {
                    "ChannelName": "test",
                    **training_data_source,
                },
            ],
            "OutputDataConfig": {"S3OutputPath": f"s3://{bucket_name}/{training_output_s3_key}"},
            "ResourceConfig": resource_config,
            "RoleArn": role_arn,
            "StoppingCondition": {"MaxRuntimeInSeconds": 600},
        },
    }
    transform_config = {
        "TransformJobName": transform_job_name,
        "TransformInput": {
            "DataSource": {
                "S3DataSource": {
                    "S3DataType": "S3Prefix",
                    "S3Uri": f"s3://{bucket_name}/{input_data_S3_key}/test.csv",
                }
            },
            "SplitType": "Line",
            "ContentType": "text/csv",
        },
        "TransformOutput": {"S3OutputPath": f"s3://{bucket_name}/{prediction_output_s3_key}"},
        "TransformResources": {
            "InstanceCount": 1,
            "InstanceType": "ml.m5.large",
        },
        "ModelName": model_name,
    }

    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)

    from tests_common.test_utils.version_compat import AIRFLOW_V_3_0_PLUS

    if AIRFLOW_V_3_0_PLUS:
        from airflow.sdk import get_current_context
    else:
        from airflow.operators.python import get_current_context  # type: ignore[attr-defined,no-redef]

    ti = get_current_context()["ti"]
    ti.xcom_push(key="docker_image", value=ecr_repository_uri)
    ti.xcom_push(key="bucket_name", value=bucket_name)
    ti.xcom_push(key="raw_data_s3_key", value=raw_data_s3_key)
    ti.xcom_push(key="ecr_repository_name", value=ecr_repository_name)
    ti.xcom_push(key="processing_config", value=processing_config)
    ti.xcom_push(key="output_files_to_xcom_config", value=output_files_to_xcom_config)
    ti.xcom_push(key="processing_job_name", value=processing_job_name)
    ti.xcom_push(key="input_data_uri", value=input_data_uri)
    ti.xcom_push(key="output_data_uri", value=f"s3://{bucket_name}/{training_output_s3_key}")
    ti.xcom_push(key="training_config", value=training_config)
    ti.xcom_push(key="training_job_name", value=training_job_name)
    ti.xcom_push(key="model_package_group_name", value=model_package_group_name)
    ti.xcom_push(key="auto_ml_job_name", value=auto_ml_job_name)
    ti.xcom_push(key="experiment_name", value=experiment_name)
    ti.xcom_push(key="model_config", value=model_config)
    ti.xcom_push(key="model_name", value=model_name)
    ti.xcom_push(key="inference_code_image", value=knn_image_uri)
    ti.xcom_push(key="model_trained_weights", value=model_trained_weights)
    ti.xcom_push(key="tuning_config", value=tuning_config)
    ti.xcom_push(key="tuning_job_name", value=tuning_job_name)
    ti.xcom_push(key="transform_config", value=transform_config)
    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")

    # All images must be removed from the repo before it can be deleted.
    image_ids = client.list_images(repositoryName=repository_name)["imageIds"]
    client.batch_delete_image(
        repositoryName=repository_name,
        imageIds=[{"imageDigest": image["imageDigest"]} for image in image_ids],
    )
    client.delete_repository(repositoryName=repository_name)


@task(trigger_rule=TriggerRule.ALL_DONE)
def delete_model_group(group_name, model_version_arn):
    sgmk_client = boto3.client("sagemaker")
    # need to destroy model registered in group first
    sgmk_client.delete_model_package(ModelPackageName=model_version_arn)
    sgmk_client.delete_model_package_group(ModelPackageGroupName=group_name)


@task(trigger_rule=TriggerRule.ALL_DONE)
def delete_experiments(experiment_names):
    sgmk_client = boto3.client("sagemaker")
    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)
def delete_docker_image(image_name):
    docker_build = subprocess.Popen(
        f"docker rmi {image_name}",
        shell=True,
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
    )
    _, stderr = docker_build.communicate()
    if docker_build.returncode != 0:
        logger.error(
            "Failed to delete local docker image. "
            "Run 'docker images' to see if you need to clean it yourself.\nerror message: %s",
            stderr,
        )


with DAG(
    dag_id=DAG_ID,
    schedule="@once",
    start_date=datetime(2021, 1, 1),
    catchup=False,
) as dag:
    test_context = sys_test_context_task()
    env_id = test_context[ENV_ID_KEY]

    test_setup = set_up(
        env_id=env_id,
        role_arn=test_context[ROLE_ARN_KEY],
    )

    create_bucket = S3CreateBucketOperator(
        task_id="create_bucket",
        bucket_name=test_setup["bucket_name"],
    )

    upload_dataset = S3CreateObjectOperator(
        task_id="upload_dataset",
        s3_bucket=test_setup["bucket_name"],
        s3_key=test_setup["raw_data_s3_key"],
        data=generate_data(),
        replace=True,
    )

    # [START howto_operator_sagemaker_auto_ml]
    automl = SageMakerAutoMLOperator(
        task_id="auto_ML",
        job_name=test_setup["auto_ml_job_name"],
        s3_input=test_setup["input_data_uri"],
        target_attribute="class",
        s3_output=test_setup["output_data_uri"],
        role_arn=test_context[ROLE_ARN_KEY],
        time_limit=30,  # will stop the job before it can do anything, but it's not the point here
    )
    # [END howto_operator_sagemaker_auto_ml]
    automl.wait_for_completion = False  # just to be able to test the sensor next

    # [START howto_sensor_sagemaker_auto_ml]
    await_automl = SageMakerAutoMLSensor(job_name=test_setup["auto_ml_job_name"], task_id="await_auto_ML")
    # [END howto_sensor_sagemaker_auto_ml]
    await_automl.poke_interval = 10

    # [START howto_operator_sagemaker_experiment]
    create_experiment = SageMakerCreateExperimentOperator(
        task_id="create_experiment", name=test_setup["experiment_name"]
    )
    # [END howto_operator_sagemaker_experiment]

    # [START howto_operator_sagemaker_processing]
    preprocess_raw_data = SageMakerProcessingOperator(
        task_id="preprocess_raw_data",
        config=test_setup["processing_config"],
    )

    # SageMakerProcessingOperator waits by default, setting as False to test the Sensor below.
    preprocess_raw_data.wait_for_completion = False

    # [END howto_operator_sagemaker_processing]

    # [START howto_sensor_sagemaker_processing]
    await_preprocess = SageMakerProcessingSensor(
        task_id="await_preprocess", job_name=test_setup["processing_job_name"]
    )
    # [END howto_sensor_sagemaker_processing]

    # [START howto_operator_sagemaker_processing_output_files]
    preprocess_with_output_files = SageMakerProcessingOperator(
        task_id="preprocess_with_output_files",
        config=test_setup["processing_config"],
        output_files_to_xcom=test_setup["output_files_to_xcom_config"],
    )
    # [END howto_operator_sagemaker_processing_output_files]

    @task
    def validate_output_files(xcom_result):
        """Verify output files were read from S3 and included in XCom."""
        assert "OutputFiles" in xcom_result, (
            f"OutputFiles missing from XCom. Got keys: {list(xcom_result.keys())}"
        )
        evaluation_report = xcom_result["OutputFiles"]["EvaluationReport"]
        assert "metrics" in evaluation_report, f"metrics key missing. Got: {evaluation_report}"
        assert evaluation_report["metrics"]["mse"]["value"] == 5.332878
        assert evaluation_report["metrics"]["rmse"]["value"] == 2.309302

    validate_pf = validate_output_files(preprocess_with_output_files.output)

    # [START howto_operator_sagemaker_training]
    train_model = SageMakerTrainingOperator(
        task_id="train_model",
        config=test_setup["training_config"],
    )
    # [END howto_operator_sagemaker_training]

    # SageMakerTrainingOperator waits by default, setting as False to test the Sensor below.
    train_model.wait_for_completion = False

    # [START howto_sensor_sagemaker_training]
    await_training = SageMakerTrainingSensor(
        task_id="await_training",
        job_name=test_setup["training_job_name"],
    )
    # [END howto_sensor_sagemaker_training]

    # [START howto_operator_sagemaker_model]
    create_model = SageMakerModelOperator(
        task_id="create_model",
        config=test_setup["model_config"],
    )
    # [END howto_operator_sagemaker_model]

    # [START howto_operator_sagemaker_register]
    register_model = SageMakerRegisterModelVersionOperator(
        task_id="register_model",
        image_uri=test_setup["inference_code_image"],
        model_url=test_setup["model_trained_weights"],
        package_group_name=test_setup["model_package_group_name"],
    )
    # [END howto_operator_sagemaker_register]

    # [START howto_operator_sagemaker_tuning]
    tune_model = SageMakerTuningOperator(
        task_id="tune_model",
        config=test_setup["tuning_config"],
    )
    # [END howto_operator_sagemaker_tuning]

    # SageMakerTuningOperator waits by default, setting as False to test the Sensor below.
    tune_model.wait_for_completion = False

    # [START howto_sensor_sagemaker_tuning]
    await_tuning = SageMakerTuningSensor(
        task_id="await_tuning",
        job_name=test_setup["tuning_job_name"],
    )
    # [END howto_sensor_sagemaker_tuning]

    # [START howto_operator_sagemaker_transform]
    test_model = SageMakerTransformOperator(
        task_id="test_model",
        config=test_setup["transform_config"],
    )
    # [END howto_operator_sagemaker_transform]

    # SageMakerTransformOperator waits by default, setting as False to test the Sensor below.
    test_model.wait_for_completion = False

    # [START howto_sensor_sagemaker_transform]
    await_transform = SageMakerTransformSensor(
        task_id="await_transform",
        job_name=test_setup["transform_job_name"],
    )
    # [END howto_sensor_sagemaker_transform]

    # [START howto_operator_sagemaker_delete_model]
    delete_model = SageMakerDeleteModelOperator(
        task_id="delete_model",
        config={"ModelName": test_setup["model_name"]},
    )
    # [END howto_operator_sagemaker_delete_model]
    delete_model.trigger_rule = TriggerRule.ALL_DONE

    delete_bucket = S3DeleteBucketOperator(
        task_id="delete_bucket",
        trigger_rule=TriggerRule.ALL_DONE,
        bucket_name=test_setup["bucket_name"],
        force_delete=True,
    )

    log_cleanup = prune_logs(
        [
            # Format: ('log group name', 'log stream prefix')
            ("/aws/sagemaker/ProcessingJobs", env_id),
            ("/aws/sagemaker/TrainingJobs", env_id),
            ("/aws/sagemaker/TransformJobs", env_id),
        ]
    )

    chain(
        # TEST SETUP
        test_context,
        test_setup,
        create_bucket,
        upload_dataset,
        # TEST BODY
        automl,
        await_automl,
        create_experiment,
        preprocess_raw_data,
        await_preprocess,
        preprocess_with_output_files,
        validate_pf,
        train_model,
        await_training,
        create_model,
        register_model,
        tune_model,
        await_tuning,
        test_model,
        await_transform,
        # TEST TEARDOWN
        delete_ecr_repository(test_setup["ecr_repository_name"]),
        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"],
                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,
    )

    from tests_common.test_utils.watcher import watcher

    # This test needs watcher in order to properly mark success/failure
    # when "tearDown" task with trigger rule is part of the DAG
    list(dag.tasks) >> watcher()

from tests_common.test_utils.system_tests import get_test_run  # noqa: E402

# Needed to run the example DAG with pytest (see: contributing-docs/testing/system_tests.rst)
test_run = get_test_run(dag)
