Amazon SageMaker Operators¶

Amazon SageMaker is a fully managed machine learning service. With Amazon SageMaker, data scientists and developers can quickly build and train machine learning models, and then deploy them into a production-ready hosted environment.

Airflow provides operators to create and interact with SageMaker Jobs.

Prerequisite Tasks¶

To use these operators, you must do a few things:

Create necessary resources using AWS Console or AWS CLI.
Install API libraries via pip.
pip install 'apache-airflow[amazon]'
Detailed information is available Installation
Setup Connection.

Manage Amazon SageMaker Jobs¶

Create an Amazon SageMaker Processing Job¶

To create an Amazon Sagemaker processing job to sanitize your dataset you can use SageMakerProcessingOperator.

airflow/providers/amazon/aws/example_dags/example_sagemaker.py[source]

preprocess_raw_data = SageMakerProcessingOperator(
    task_id='preprocess_raw_data',
    config=SAGEMAKER_PROCESSING_JOB_CONFIG,
    do_xcom_push=False,
)

Create an Amazon SageMaker Training Job¶

To create an Amazon Sagemaker training job you can use SageMakerTrainingOperator.

airflow/providers/amazon/aws/example_dags/example_sagemaker.py[source]

train_model = SageMakerTrainingOperator(
    task_id='train_model',
    config=TRAINING_CONFIG,
    # Waits by default, setting as False to demonstrate the Sensor below.
    wait_for_completion=False,
    do_xcom_push=False,
)

Create an Amazon SageMaker Model¶

To create an Amazon Sagemaker model you can use SageMakerModelOperator.

airflow/providers/amazon/aws/example_dags/example_sagemaker.py[source]

create_model = SageMakerModelOperator(
    task_id='create_model',
    config=MODEL_CONFIG,
    do_xcom_push=False,
)

Start a Hyperparameter Tuning Job¶

To start a hyperparameter tuning job for an Amazon Sagemaker model you can use SageMakerTuningOperator.

airflow/providers/amazon/aws/example_dags/example_sagemaker.py[source]

tune_model = SageMakerTuningOperator(
    task_id='tune_model',
    config=TUNING_CONFIG,
    # Waits by default, setting as False to demonstrate the Sensor below.
    wait_for_completion=False,
    do_xcom_push=False,
)

Delete an Amazon SageMaker Model¶

To delete an Amazon Sagemaker model you can use SageMakerDeleteModelOperator.

airflow/providers/amazon/aws/example_dags/example_sagemaker.py[source]

delete_model = SageMakerDeleteModelOperator(
    task_id='delete_model',
    config={'ModelName': MODEL_NAME},
    trigger_rule='all_done',
)

Create an Amazon SageMaker Transform Job¶

To create an Amazon Sagemaker transform job you can use SageMakerTransformOperator.

airflow/providers/amazon/aws/example_dags/example_sagemaker.py[source]

test_model = SageMakerTransformOperator(
    task_id='test_model',
    config=TRANSFORM_CONFIG,
    # Waits by default, setting as False to demonstrate the Sensor below.
    wait_for_completion=False,
    do_xcom_push=False,
)

Create an Amazon SageMaker Endpoint Config Job¶

To create an Amazon Sagemaker endpoint config job you can use SageMakerEndpointConfigOperator.

airflow/providers/amazon/aws/example_dags/example_sagemaker_endpoint.py[source]

configure_endpoint = SageMakerEndpointConfigOperator(
    task_id='configure_endpoint',
    config=ENDPOINT_CONFIG_CONFIG,
    do_xcom_push=False,
)

Create an Amazon SageMaker Endpoint Job¶

To create an Amazon Sagemaker endpoint you can use SageMakerEndpointOperator.

airflow/providers/amazon/aws/example_dags/example_sagemaker_endpoint.py[source]

deploy_endpoint = SageMakerEndpointOperator(
    task_id='deploy_endpoint',
    config=DEPLOY_ENDPOINT_CONFIG,
    # Waits by default, <setting as False to demonstrate the Sensor below.
    wait_for_completion=False,
    do_xcom_push=False,
)

Amazon SageMaker Sensors¶

Amazon SageMaker Training Sensor¶

To check the state of an Amazon Sagemaker training job until it reaches a terminal state you can use SageMakerTrainingSensor.

airflow/providers/amazon/aws/example_dags/example_sagemaker.py[source]

await_training = SageMakerTrainingSensor(
    task_id='await_training',
    job_name=TRAINING_JOB_NAME,
)

Amazon SageMaker Transform Sensor¶

To check the state of an Amazon Sagemaker transform job until it reaches a terminal state you can use SageMakerTransformOperator.

airflow/providers/amazon/aws/example_dags/example_sagemaker.py[source]

await_transform = SageMakerTransformSensor(
    task_id='await_transform',
    job_name=TRANSFORM_JOB_NAME,
)

Amazon SageMaker Tuning Sensor¶

To check the state of an Amazon Sagemaker hyperparameter tuning job until it reaches a terminal state you can use SageMakerTuningSensor.

airflow/providers/amazon/aws/example_dags/example_sagemaker.py[source]

await_tune = SageMakerTuningSensor(
    task_id='await_tuning',
    job_name=TUNING_JOB_NAME,
)

Amazon SageMaker Endpoint Sensor¶

To check the state of an Amazon Sagemaker hyperparameter tuning job until it reaches a terminal state you can use SageMakerEndpointSensor.

airflow/providers/amazon/aws/example_dags/example_sagemaker_endpoint.py[source]

await_endpoint = SageMakerEndpointSensor(
    task_id='await_endpoint',
    endpoint_name=ENDPOINT_NAME,
    do_xcom_push=False,
)

Reference¶

For further information, look at: