astronomer.providers.amazon.aws.triggers.sagemaker

Module Contents

Classes

SagemakerProcessingTrigger

SagemakerProcessingTrigger is fired as deferred class with params to run the task in triggerer.

SagemakerTrigger

SagemakerTrigger is common trigger for both transform and training sagemaker job and it is

SagemakerTrainingWithLogTrigger

SagemakerTrainingWithLogTrigger is fired as deferred class with params to run the task in triggerer.
class astronomer.providers.amazon.aws.triggers.sagemaker.SagemakerProcessingTrigger(job_name, poll_interval, end_time, aws_conn_id='aws_default')[source]

Bases: airflow.triggers.base.BaseTrigger

SagemakerProcessingTrigger is fired as deferred class with params to run the task in triggerer.

Parameters:

  • job_name (str) – name of the job to check status

  • poll_interval (float) – polling period in seconds to check for the status

  • aws_conn_id (str) – AWS connection ID for sagemaker

  • end_time (Optional[float]) – the end time in seconds. Any SageMaker jobs that run longer than this will fail.

NON_TERMINAL_STATES = ['InProgress', 'Stopping']
TERMINAL_STATE = ['Failed']
serialize()[source]

Serializes SagemakerProcessingTrigger arguments and classpath.

async run()[source]

Makes async connection to sagemaker async hook and gets job status for a job submitted by the operator. Trigger returns a failure event if any error and success in state return the success event.

class astronomer.providers.amazon.aws.triggers.sagemaker.SagemakerTrigger(job_name, job_type, response_key, poke_interval, end_time=None, aws_conn_id='aws_default')[source]

Bases: airflow.triggers.base.BaseTrigger

SagemakerTrigger is common trigger for both transform and training sagemaker job and it is

fired as deferred class with params to run the task in triggerer.

Parameters:

  • job_name (str) – name of the job to check status

  • job_type (str) – Type of the sagemaker job whether it is Transform or Training

  • response_key (str) – The key which needs to be look in the response.

  • poke_interval (float) – polling period in seconds to check for the status

  • end_time (Optional[float]) – Time in seconds to wait for a job run to reach a terminal status.

  • aws_conn_id (str) – AWS connection ID for sagemaker

NON_TERMINAL_STATES = ['InProgress', 'Stopping', 'Stopped']
TERMINAL_STATE = ['Failed']
serialize()[source]

Serializes SagemakerTrigger arguments and classpath.

async run()[source]

Makes async connection to sagemaker async hook and gets job status for a job submitted by the operator. Trigger returns a failure event if any error and success in state return the success event.

async static get_job_status(hook, job_name, job_type)[source]

Based on the job type the SageMakerHookAsync connect to sagemaker related function and get the response of the job and return it

class astronomer.providers.amazon.aws.triggers.sagemaker.SagemakerTrainingWithLogTrigger(job_name, instance_count, status, poke_interval, end_time=None, aws_conn_id='aws_default')[source]

Bases: airflow.triggers.base.BaseTrigger

SagemakerTrainingWithLogTrigger is fired as deferred class with params to run the task in triggerer.

Parameters:

  • job_name (str) – name of the job to check status

  • instance_count (int) – count of the instance created for running the training job

  • status (str) – The status of the training job created.

  • poke_interval (float) – polling period in seconds to check for the status

  • end_time (Optional[float]) – Time in seconds to wait for a job run to reach a terminal status.

  • aws_conn_id (str) – AWS connection ID for sagemaker

NON_TERMINAL_STATES = ['InProgress', 'Stopping', 'Stopped']
TERMINAL_STATE = ['Failed']
serialize()[source]

Serializes SagemakerTrainingWithLogTrigger arguments and classpath.

async run()[source]

Makes async connection to sagemaker async hook and gets job status for a job submitted by the operator. Trigger returns a failure event if any error and success in state return the success event.