airflow.contrib.hooks.sagemaker_hook

Module Contents

class airflow.contrib.hooks.sagemaker_hook.LogState

Bases: object

STARTING = 1

WAIT_IN_PROGRESS = 2

TAILING = 3

JOB_COMPLETE = 4

COMPLETE = 5

airflow.contrib.hooks.sagemaker_hook.Position

airflow.contrib.hooks.sagemaker_hook.argmin(arr, f)

Return the index, i, in arr that minimizes f(arr[i])

airflow.contrib.hooks.sagemaker_hook.secondary_training_status_changed(current_job_description, prev_job_description)

Returns true if training job's secondary status message has changed.

Parameters

• current_job_description (dict) – Current job description, returned from DescribeTrainingJob call.

• prev_job_description (dict) – Previous job description, returned from DescribeTrainingJob call.

Returns

Whether the secondary status message of a training job changed or not.

airflow.contrib.hooks.sagemaker_hook.secondary_training_status_message(job_description, prev_description)

Returns a string contains start time and the secondary training job status message.

Parameters

• job_description (dict) – Returned response from DescribeTrainingJob call

• prev_description (dict) – Previous job description from DescribeTrainingJob call

Returns

Job status string to be printed.

class airflow.contrib.hooks.sagemaker_hook.SageMakerHook(*args, **kwargs)

Bases: airflow.contrib.hooks.aws_hook.AwsHook

Interact with Amazon SageMaker.

non_terminal_states

endpoint_non_terminal_states

failed_states

tar_and_s3_upload(self, path, key, bucket)

Tar the local file or directory and upload to s3

Parameters

• path (str) – local file or directory

• key (str) – s3 key

• bucket (str) – s3 bucket

Returns

None

configure_s3_resources(self, config)

Extract the S3 operations from the configuration and execute them.

Parameters

config (dict) – config of SageMaker operation

Return type

dict

check_s3_url(self, s3url)

Check if an S3 URL exists

Parameters

s3url (str) – S3 url

Return type

bool

check_training_config(self, training_config)

Check if a training configuration is valid

Parameters

training_config (dict) – training_config

Returns

None

check_tuning_config(self, tuning_config)

Check if a tuning configuration is valid

Parameters

tuning_config (dict) – tuning_config

Returns

None

get_conn(self)

Establish an AWS connection for SageMaker

Return type

SageMaker.Client

get_log_conn(self)

This method is deprecated. Please use airflow.contrib.hooks.AwsLogsHook.get_conn() instead.

log_stream(self, log_group, stream_name, start_time=0, skip=0)

This method is deprecated. Please use airflow.contrib.hooks.AwsLogsHook.get_log_events() instead.

multi_stream_iter(self, log_group, streams, positions=None)

Iterate over the available events coming from a set of log streams in a single log group interleaving the events from each stream so they’re yielded in timestamp order.

Parameters

• log_group (str) – The name of the log group.

• streams (list) – A list of the log stream names. The position of the stream in this list is the stream number.

• positions (list) – A list of pairs of (timestamp, skip) which represents the last record read from each stream.

Returns

A tuple of (stream number, cloudwatch log event).

create_training_job(self, config, wait_for_completion=True, print_log=True, check_interval=30, max_ingestion_time=None)

Create a training job

Parameters

• config (dict) – the config for training

• wait_for_completion (bool) – if the program should keep running until job finishes

• check_interval (int) – the time interval in seconds which the operator will check the status of any SageMaker job

• max_ingestion_time (int) – the maximum ingestion time in seconds. Any SageMaker jobs that run longer than this will fail. Setting this to None implies no timeout for any SageMaker job.

Returns

A response to training job creation

create_tuning_job(self, config, wait_for_completion=True, check_interval=30, max_ingestion_time=None)

Create a tuning job

Parameters

• config (dict) – the config for tuning

• wait_for_completion (bool) – if the program should keep running until job finishes

• check_interval (int) – the time interval in seconds which the operator will check the status of any SageMaker job

• max_ingestion_time (int) – the maximum ingestion time in seconds. Any SageMaker jobs that run longer than this will fail. Setting this to None implies no timeout for any SageMaker job.

Returns

A response to tuning job creation

create_transform_job(self, config, wait_for_completion=True, check_interval=30, max_ingestion_time=None)

Create a transform job

Parameters

• config (dict) – the config for transform job

• wait_for_completion (bool) – if the program should keep running until job finishes

• check_interval (int) – the time interval in seconds which the operator will check the status of any SageMaker job

• max_ingestion_time (int) – the maximum ingestion time in seconds. Any SageMaker jobs that run longer than this will fail. Setting this to None implies no timeout for any SageMaker job.

Returns

A response to transform job creation

create_model(self, config)

Create a model job

Parameters

config (dict) – the config for model

Returns

A response to model creation

create_endpoint_config(self, config)

Create an endpoint config

Parameters

config (dict) – the config for endpoint-config

Returns

A response to endpoint config creation

create_endpoint(self, config, wait_for_completion=True, check_interval=30, max_ingestion_time=None)

Create an endpoint

Parameters

• config (dict) – the config for endpoint

• wait_for_completion (bool) – if the program should keep running until job finishes

• check_interval (int) – the time interval in seconds which the operator will check the status of any SageMaker job

• max_ingestion_time (int) – the maximum ingestion time in seconds. Any SageMaker jobs that run longer than this will fail. Setting this to None implies no timeout for any SageMaker job.

Returns

A response to endpoint creation

update_endpoint(self, config, wait_for_completion=True, check_interval=30, max_ingestion_time=None)

Update an endpoint

Parameters

• config (dict) – the config for endpoint

• wait_for_completion (bool) – if the program should keep running until job finishes

• check_interval (int) – the time interval in seconds which the operator will check the status of any SageMaker job

• max_ingestion_time (int) – the maximum ingestion time in seconds. Any SageMaker jobs that run longer than this will fail. Setting this to None implies no timeout for any SageMaker job.

Returns

A response to endpoint update

describe_training_job(self, name)

Return the training job info associated with the name

Parameters

name (str) – the name of the training job

Returns

A dict contains all the training job info

describe_training_job_with_log(self, job_name, positions, stream_names, instance_count, state, last_description, last_describe_job_call)

Return the training job info associated with job_name and print CloudWatch logs

describe_tuning_job(self, name)

Return the tuning job info associated with the name

Parameters

name (string) – the name of the tuning job

Returns

A dict contains all the tuning job info

describe_model(self, name)

Return the SageMaker model info associated with the name

Parameters

name (string) – the name of the SageMaker model

Returns

A dict contains all the model info

describe_transform_job(self, name)

Return the transform job info associated with the name

Parameters

name (string) – the name of the transform job

Returns

A dict contains all the transform job info

describe_endpoint_config(self, name)

Return the endpoint config info associated with the name

Parameters

name (string) – the name of the endpoint config

Returns

A dict contains all the endpoint config info

describe_endpoint(self, name)

Parameters

name (string) – the name of the endpoint

Returns

A dict contains all the endpoint info

check_status(self, job_name, key, describe_function, check_interval, max_ingestion_time, non_terminal_states=None)

Check status of a SageMaker job

Parameters

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

• key (str) – the key of the response dict that points to the state

• describe_function (python callable) – the function used to retrieve the status

• args – the arguments for the function

• check_interval (int) – the time interval in seconds which the operator will check the status of any SageMaker job

• max_ingestion_time (int) – the maximum ingestion time in seconds. Any SageMaker jobs that run longer than this will fail. Setting this to None implies no timeout for any SageMaker job.

• non_terminal_states (set) – the set of nonterminal states

Returns

response of describe call after job is done

check_training_status_with_log(self, job_name, non_terminal_states, failed_states, wait_for_completion, check_interval, max_ingestion_time)

Display the logs for a given training job, optionally tailing them until the job is complete.

Parameters

• job_name (str) – name of the training job to check status and display logs for

• non_terminal_states (set) – the set of non_terminal states

• failed_states (set) – the set of failed states

• wait_for_completion (bool) – Whether to keep looking for new log entries until the job completes

• check_interval (int) – The interval in seconds between polling for new log entries and job completion

• max_ingestion_time (int) – the maximum ingestion time in seconds. Any SageMaker jobs that run longer than this will fail. Setting this to None implies no timeout for any SageMaker job.

Returns

None