airflow.providers.amazon.aws.hooks.sagemaker

Module Contents

Classes

LogState

Enum-style class holding all possible states of CloudWatch log streams.

SageMakerHook

Interact with Amazon SageMaker.

Functions

argmin(arr, f)

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

secondary_training_status_changed(...)

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

secondary_training_status_message(job_description, ...)

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

Attributes

Position

class airflow.providers.amazon.aws.hooks.sagemaker.LogState[source]

Enum-style class holding all possible states of CloudWatch log streams. https://sagemaker.readthedocs.io/en/stable/session.html#sagemaker.session.LogState

STARTING = 1[source]
WAIT_IN_PROGRESS = 2[source]
TAILING = 3[source]
JOB_COMPLETE = 4[source]
COMPLETE = 5[source]
airflow.providers.amazon.aws.hooks.sagemaker.Position[source]
airflow.providers.amazon.aws.hooks.sagemaker.argmin(arr, f)[source]

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

airflow.providers.amazon.aws.hooks.sagemaker.secondary_training_status_changed(current_job_description, prev_job_description)[source]

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.

Return type

bool

airflow.providers.amazon.aws.hooks.sagemaker.secondary_training_status_message(job_description, prev_description)[source]

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

Parameters
  • job_description (dict[str, list[Any]]) – Returned response from DescribeTrainingJob call

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

Returns

Job status string to be printed.

Return type

str

class airflow.providers.amazon.aws.hooks.sagemaker.SageMakerHook(*args, **kwargs)[source]

Bases: airflow.providers.amazon.aws.hooks.base_aws.AwsBaseHook

Interact with Amazon SageMaker. Provide thick wrapper around boto3.client("sagemaker").

Additional arguments (such as aws_conn_id) may be specified and are passed down to the underlying AwsBaseHook.

non_terminal_states[source]
endpoint_non_terminal_states[source]
pipeline_non_terminal_states[source]
failed_states[source]
tar_and_s3_upload(path, key, bucket)[source]

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

Return type

None

configure_s3_resources(config)[source]

Extract the S3 operations from the configuration and execute them.

Parameters

config (dict) – config of SageMaker operation

check_s3_url(s3url)[source]

Check if an S3 URL exists

Parameters

s3url (str) – S3 url

check_training_config(training_config)[source]

Check if a training configuration is valid

Parameters

training_config (dict) – training_config

Returns

None

Return type

None

check_tuning_config(tuning_config)[source]

Check if a tuning configuration is valid

Parameters

tuning_config (dict) – tuning_config

Returns

None

Return type

None

multi_stream_iter(log_group, streams, positions=None)[source]

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 – 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).

Return type

Generator

create_training_job(config, wait_for_completion=True, print_log=True, check_interval=30, max_ingestion_time=None)[source]

Starts a model training job. After training completes, Amazon SageMaker saves the resulting model artifacts to an Amazon S3 location that you specify.

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 | None) – 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(config, wait_for_completion=True, check_interval=30, max_ingestion_time=None)[source]

Starts a hyperparameter tuning job. A hyperparameter tuning job finds the best version of a model by running many training jobs on your dataset using the algorithm you choose and values for hyperparameters within ranges that you specify. It then chooses the hyperparameter values that result in a model that performs the best, as measured by an objective metric that you choose.

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 | None) – 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(config, wait_for_completion=True, check_interval=30, max_ingestion_time=None)[source]

Starts a transform job. A transform job uses a trained model to get inferences on a dataset and saves these results to an Amazon S3 location that you specify.

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 | None) – 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_processing_job(config, wait_for_completion=True, check_interval=30, max_ingestion_time=None)[source]

Use Amazon SageMaker Processing to analyze data and evaluate machine learning models on Amazon SageMaker. With Processing, you can use a simplified, managed experience on SageMaker to run your data processing workloads, such as feature engineering, data validation, model evaluation, and model interpretation.

Parameters
  • config (dict) – the config for processing 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 | None) – 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(config)[source]

Creates a model in Amazon SageMaker. In the request, you name the model and describe a primary container. For the primary container, you specify the Docker image that contains inference code, artifacts (from prior training), and a custom environment map that the inference code uses when you deploy the model for predictions.

Parameters

config (dict) – the config for model

Returns

A response to model creation

create_endpoint_config(config)[source]

Creates an endpoint configuration that Amazon SageMaker hosting services uses to deploy models. In the configuration, you identify one or more models, created using the CreateModel API, to deploy and the resources that you want Amazon SageMaker to provision.

Parameters

config (dict) – the config for endpoint-config

Returns

A response to endpoint config creation

create_endpoint(config, wait_for_completion=True, check_interval=30, max_ingestion_time=None)[source]

When you create a serverless endpoint, SageMaker provisions and manages the compute resources for you. Then, you can make inference requests to the endpoint and receive model predictions in response. SageMaker scales the compute resources up and down as needed to handle your request traffic.

Requires an Endpoint Config.

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 | None) – 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(config, wait_for_completion=True, check_interval=30, max_ingestion_time=None)[source]

Deploys the new EndpointConfig specified in the request, switches to using newly created endpoint, and then deletes resources provisioned for the endpoint using the previous EndpointConfig (there is no availability loss).

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 | None) – 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(name)[source]

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(job_name, positions, stream_names, instance_count, state, last_description, last_describe_job_call)[source]

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

describe_tuning_job(name)[source]

Return the tuning job info associated with the name

Parameters

name (str) – the name of the tuning job

Returns

A dict contains all the tuning job info

Return type

dict

describe_model(name)[source]

Return the SageMaker model info associated with the name

Parameters

name (str) – the name of the SageMaker model

Returns

A dict contains all the model info

Return type

dict

describe_transform_job(name)[source]

Return the transform job info associated with the name

Parameters

name (str) – the name of the transform job

Returns

A dict contains all the transform job info

Return type

dict

describe_processing_job(name)[source]

Return the processing job info associated with the name

Parameters

name (str) – the name of the processing job

Returns

A dict contains all the processing job info

Return type

dict

describe_endpoint_config(name)[source]

Return the endpoint config info associated with the name

Parameters

name (str) – the name of the endpoint config

Returns

A dict contains all the endpoint config info

Return type

dict

describe_endpoint(name)[source]

Returns the description of an endpoint.

Parameters

name (str) – the name of the endpoint

Returns

A dict contains all the endpoint info

Return type

dict

check_status(job_name, key, describe_function, check_interval, max_ingestion_time=None, non_terminal_states=None)[source]

Check status of a SageMaker resource

Parameters
  • job_name (str) – name of the resource to check status, can be a job but also pipeline for instance.

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

  • describe_function (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 resource

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

  • non_terminal_states (set | None) – the set of nonterminal states

Returns

response of describe call after resource is done

Return type

dict

check_training_status_with_log(job_name, non_terminal_states, failed_states, wait_for_completion, check_interval, max_ingestion_time=None)[source]

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 | None) – 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

list_training_jobs(name_contains=None, max_results=None, **kwargs)[source]

This method wraps boto3’s list_training_jobs. The training job name and max results are configurable via arguments. Other arguments are not, and should be provided via kwargs. Note boto3 expects these in CamelCase format, for example:

list_training_jobs(name_contains="myjob", StatusEquals="Failed")
Parameters
  • name_contains (str | None) – (optional) partial name to match

  • max_results (int | None) – (optional) maximum number of results to return. None returns infinite results

  • kwargs – (optional) kwargs to boto3’s list_training_jobs method

Returns

results of the list_training_jobs request

Return type

list[dict]

list_transform_jobs(name_contains=None, max_results=None, **kwargs)[source]

This method wraps boto3’s list_transform_jobs. The transform job name and max results are configurable via arguments. Other arguments are not, and should be provided via kwargs. Note boto3 expects these in CamelCase format, for example:

list_transform_jobs(name_contains="myjob", StatusEquals="Failed")
Parameters
  • name_contains (str | None) – (optional) partial name to match

  • max_results (int | None) – (optional) maximum number of results to return. None returns infinite results

  • kwargs – (optional) kwargs to boto3’s list_transform_jobs method

Returns

results of the list_transform_jobs request

Return type

list[dict]

list_processing_jobs(**kwargs)[source]

This method wraps boto3’s list_processing_jobs. All arguments should be provided via kwargs. Note boto3 expects these in CamelCase format, for example:

list_processing_jobs(NameContains="myjob", StatusEquals="Failed")
Parameters

kwargs – (optional) kwargs to boto3’s list_training_jobs method

Returns

results of the list_processing_jobs request

Return type

list[dict]

count_processing_jobs_by_name(processing_job_name, job_name_suffix=None, throttle_retry_delay=2, retries=3)[source]

Returns the number of processing jobs found with the provided name prefix. :param processing_job_name: The prefix to look for. :param job_name_suffix: The optional suffix which may be appended to deduplicate an existing job name. :param throttle_retry_delay: Seconds to wait if a ThrottlingException is hit. :param retries: The max number of times to retry. :returns: The number of processing jobs that start with the provided prefix.

delete_model(model_name)[source]

Delete SageMaker model.

Parameters

model_name (str) – name of the model

describe_pipeline_exec(pipeline_exec_arn, verbose=False)[source]

Get info about a SageMaker pipeline execution.

Parameters
  • pipeline_exec_arn (str) – arn of the pipeline execution

  • verbose (bool) – Whether to log details about the steps status in the pipeline execution

start_pipeline(pipeline_name, display_name='airflow-triggered-execution', pipeline_params=None, wait_for_completion=False, check_interval=30, verbose=True)[source]

Start a new execution for a SageMaker pipeline.

Parameters
  • pipeline_name (str) – Name of the pipeline to start (this is _not_ the ARN).

  • display_name (str) – The name this pipeline execution will have in the UI. Doesn’t need to be unique.

  • pipeline_params (dict | None) – Optional parameters for the pipeline. All parameters supplied need to already be present in the pipeline definition.

  • wait_for_completion (bool) – Will only return once the pipeline is complete if true.

  • check_interval (int) – How long to wait between checks for pipeline status when waiting for completion.

  • verbose (bool) – Whether to print steps details when waiting for completion. Defaults to true, consider turning off for pipelines that have thousands of steps.

Returns

the ARN of the pipeline execution launched.

Return type

str

stop_pipeline(pipeline_exec_arn, wait_for_completion=False, check_interval=10, verbose=True, fail_if_not_running=False)[source]

Stop SageMaker pipeline execution

Parameters
  • pipeline_exec_arn (str) – Amazon Resource Name (ARN) of the pipeline execution. It’s the ARN of the pipeline itself followed by “/execution/” and an id.

  • wait_for_completion (bool) – Whether to wait for the pipeline to reach a final state. (i.e. either ‘Stopped’ or ‘Failed’)

  • check_interval (int) – How long to wait between checks for pipeline status when waiting for completion.

  • verbose (bool) – Whether to print steps details when waiting for completion. Defaults to true, consider turning off for pipelines that have thousands of steps.

  • fail_if_not_running (bool) – This method will raise an exception if the pipeline we’re trying to stop is not in an “Executing” state when the call is sent (which would mean that the pipeline is already either stopping or stopped). Note that setting this to True will raise an error if the pipeline finished successfully before it was stopped.

Returns

Status of the pipeline execution after the operation. One of ‘Executing’|’Stopping’|’Stopped’|’Failed’|’Succeeded’.

Return type

str

create_model_package_group(package_group_name, package_group_desc='')[source]

Creates a Model Package Group if it does not already exist

Parameters
  • package_group_name (str) – Name of the model package group to create if not already present.

  • package_group_desc (str) – Description of the model package group, if it was to be created (optional).

Returns

True if the model package group was created, False if it already existed.

Return type

bool

create_auto_ml_job(job_name, s3_input, target_attribute, s3_output, role_arn, compressed_input=False, time_limit=None, autodeploy_endpoint_name=None, extras=None, wait_for_completion=True, check_interval=30)[source]

Creates an auto ML job, learning to predict the given column from the data provided through S3. The learning output is written to the specified S3 location.

Parameters
  • job_name (str) – Name of the job to create, needs to be unique within the account.

  • s3_input (str) – The S3 location (folder or file) where to fetch the data. By default, it expects csv with headers.

  • target_attribute (str) – The name of the column containing the values to predict.

  • s3_output (str) – The S3 folder where to write the model artifacts. Must be 128 characters or fewer.

  • role_arn (str) – The ARN or the IAM role to use when interacting with S3. Must have read access to the input, and write access to the output folder.

  • compressed_input (bool) – Set to True if the input is gzipped.

  • time_limit (int | None) – The maximum amount of time in seconds to spend training the model(s).

  • autodeploy_endpoint_name (str | None) – If specified, the best model will be deployed to an endpoint with that name. No deployment made otherwise.

  • extras (dict | None) – Use this dictionary to set any variable input variable for job creation that is not offered through the parameters of this function. The format is described in: https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/sagemaker.html#SageMaker.Client.create_auto_ml_job

  • wait_for_completion (bool) – Whether to wait for the job to finish before returning. Defaults to True.

  • check_interval (int) – Interval in seconds between 2 status checks when waiting for completion.

Returns

Only if waiting for completion, a dictionary detailing the best model. The structure is that of the “BestCandidate” key in: https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/sagemaker.html#SageMaker.Client.describe_auto_ml_job

Return type

dict | None

Was this entry helpful?