airflow.providers.amazon.aws.hooks.sagemaker
¶
Module Contents¶
Classes¶
Enum-style class holding all possible states of CloudWatch log streams. |
|
Interact with Amazon SageMaker. |
Functions¶
|
Given callable |
Check if training job's secondary status message has changed. |
|
|
Format string containing start time and the secondary training job status message. |
Attributes¶
- 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
- airflow.providers.amazon.aws.hooks.sagemaker.argmin(arr, f)[source]¶
Given callable
f
, find index inarr
to minimizef(arr[i])
.None is returned if
arr
is empty.
- airflow.providers.amazon.aws.hooks.sagemaker.secondary_training_status_changed(current_job_description, prev_job_description)[source]¶
Check if training job’s secondary status message has changed.
- Parameters
- Returns
Whether the secondary status message of a training job changed or not.
- Return type
- airflow.providers.amazon.aws.hooks.sagemaker.secondary_training_status_message(job_description, prev_description)[source]¶
Format string containing start time and the secondary training job status message.
- 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.- configure_s3_resources(config)[source]¶
Extract the S3 operations from the configuration and execute them.
- Parameters
config (dict) – config of SageMaker operation
- check_training_config(training_config)[source]¶
Check if a training configuration is valid.
- Parameters
training_config (dict) – training_config
- check_tuning_config(tuning_config)[source]¶
Check if a tuning configuration is valid.
- Parameters
tuning_config (dict) – tuning_config
- multi_stream_iter(log_group, streams, positions=None)[source]¶
Iterate over the available events.
The 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
- 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]¶
Start 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]¶
Start 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]¶
Start 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 models.
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]¶
Create 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.
See also
- Parameters
config (dict) – the config for model
- Returns
A response to model creation
- create_endpoint_config(config)[source]¶
Create an endpoint configuration 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.
See also
- 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]¶
Create an endpoint from configuration.
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.
See also
- 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]¶
Deploy the config in the request and switch to using the new endpoint.
Resources provisioned for the endpoint using the previous EndpointConfig are deleted (there is no availability loss).
See also
- 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]¶
Get 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]¶
Get the associated training job info and print CloudWatch logs.
- 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
- check_training_status_with_log(job_name, non_terminal_states, failed_states, wait_for_completion, check_interval, max_ingestion_time=None)[source]¶
Display 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.
- list_training_jobs(name_contains=None, max_results=None, **kwargs)[source]¶
Call 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 that boto3 expects these in CamelCase, for example:
list_training_jobs(name_contains="myjob", StatusEquals="Failed")
- Parameters
- Returns
results of the list_training_jobs request
- Return type
- list_transform_jobs(name_contains=None, max_results=None, **kwargs)[source]¶
Call 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 that boto3 expects these in CamelCase, for example:
list_transform_jobs(name_contains="myjob", StatusEquals="Failed")
- Parameters
- Returns
results of the list_transform_jobs request.
- Return type
- list_processing_jobs(**kwargs)[source]¶
Call boto3’s list_processing_jobs.
All arguments should be provided via kwargs. Note that boto3 expects these in CamelCase, for example:
list_processing_jobs(NameContains="myjob", StatusEquals="Failed")
- count_processing_jobs_by_name(processing_job_name, job_name_suffix=None, throttle_retry_delay=2, retries=3)[source]¶
Get the number of processing jobs found with the provided name prefix.
- Parameters
- Returns
The number of processing jobs that start with the provided prefix.
- Return type
- delete_model(model_name)[source]¶
Delete a SageMaker model.
See also
- Parameters
model_name (str) – name of the model
- describe_pipeline_exec(pipeline_exec_arn, verbose=False)[source]¶
Get info about a SageMaker pipeline execution.
See also
- start_pipeline(pipeline_name, display_name='airflow-triggered-execution', pipeline_params=None)[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.
- Returns
the ARN of the pipeline execution launched.
- Return type
- stop_pipeline(pipeline_exec_arn, 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.
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
- create_model_package_group(package_group_name, package_group_desc='')[source]¶
Create a Model Package Group if it does not already exist.
- Parameters
- Returns
True if the model package group was created, False if it already existed.
- Return type
- 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]¶
Create an auto ML job to predict the given column.
The learning input is based on data provided through S3 , and the 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
- async describe_training_job_async(job_name)[source]¶
Return the training job info associated with the name.
- Parameters
job_name (str) – the name of the training job
- async describe_training_job_with_log_async(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.
- Parameters
job_name (str) – name of the job to check status
positions (dict[str, Any]) – A list of pairs of (timestamp, skip) which represents the last record read from each stream.
stream_names (list[str]) – A list of the log stream names. The position of the stream in this list is the stream number.
instance_count (int) – Count of the instance created for the job initially
state (int) – log state
last_description (dict[str, Any]) – Latest description of the training job
last_describe_job_call (float) – previous job called time