AWS Sagemaker Inference API#

Note

To use the Sagemaker agent, you must pip install flytekitplugins-awssagemaker.

Boto#

class flytekitplugins.awssagemaker_inference.BotoAgent#

A general purpose boto3 agent that can be used to call any boto3 method.

async do(task_template, output_prefix, inputs=None, **kwargs)#

This is the method that the agent will run.

Parameters:

  • task_template (TaskTemplate)

  • output_prefix (str)

  • inputs (LiteralMap | None)

Return type:

Resource

class flytekitplugins.awssagemaker_inference.BotoTask(*args, **kwargs)#
get_custom(settings)#

Return additional plugin-specific custom data (if any) as a serializable dictionary.

Parameters:

settings (SerializationSettings)

Return type:

Dict[str, Any]

class flytekitplugins.awssagemaker_inference.SageMakerEndpointAgent#

This agent creates an endpoint.

async create(task_template, inputs=None, **kwargs)#

Return a resource meta that can be used to get the status of the task.

Parameters:

  • task_template (TaskTemplate)

  • inputs (LiteralMap | None)

Return type:

SageMakerEndpointMetadata

async delete(resource_meta, **kwargs)#

Delete the task. This call should be idempotent. It should raise an error if fails to delete the task.

Parameters:

resource_meta (SageMakerEndpointMetadata)

async get(resource_meta, **kwargs)#

Return the status of the task, and return the outputs in some cases. For example, bigquery job can’t write the structured dataset to the output location, so it returns the output literals to the propeller, and the propeller will write the structured dataset to the blob store.

Parameters:

resource_meta (SageMakerEndpointMetadata)

Return type:

Resource

class flytekitplugins.awssagemaker_inference.SageMakerEndpointTask(*args, **kwargs)#
get_custom(settings)#

Return additional plugin-specific custom data (if any) as a serializable dictionary.

Parameters:

settings (SerializationSettings)

Return type:

Dict[str, Any]

class flytekitplugins.awssagemaker_inference.SageMakerEndpointConfigTask(*args, **kwargs)#
class flytekitplugins.awssagemaker_inference.SageMakerInvokeEndpointTask(*args, **kwargs)#
class flytekitplugins.awssagemaker_inference.SageMakerDeleteEndpointConfigTask(*args, **kwargs)#
class flytekitplugins.awssagemaker_inference.SageMakerModelTask(*args, **kwargs)#
class flytekitplugins.awssagemaker_inference.SageMakerDeleteModelTask(*args, **kwargs)#
flytekitplugins.awssagemaker_inference.create_sagemaker_deployment(name, model_config, endpoint_config_config, endpoint_config, images=None, model_input_types=None, endpoint_config_input_types=None, endpoint_input_types=None, region=None, region_at_runtime=False, idempotence_token=True)#

Creates SageMaker model, endpoint config and endpoint.

Parameters:

  • model_config (Dict[str, Any]) – Configuration for the SageMaker model creation API call.

  • endpoint_config_config (Dict[str, Any]) – Configuration for the SageMaker endpoint configuration creation API call.

  • endpoint_config (Dict[str, Any]) – Configuration for the SageMaker endpoint creation API call.

  • images (Dict[str, Any] | None) – A dictionary of images for SageMaker model creation.

  • model_input_types (Dict[str, Type] | None) – Mapping of SageMaker model configuration inputs to their types.

  • endpoint_config_input_types (Dict[str, Type] | None) – Mapping of SageMaker endpoint configuration inputs to their types.

  • endpoint_input_types (Dict[str, Type] | None) – Mapping of SageMaker endpoint inputs to their types.

  • region (str | None) – The region for SageMaker API calls.

  • region_at_runtime (bool) – Set this to True if you want to provide the region at runtime.

  • idempotence_token (bool) – Set this to False if you don’t want the agent to automatically append a token/hash to the deployment names.

  • name (str)

Return type:

ImperativeWorkflow

flytekitplugins.awssagemaker_inference.delete_sagemaker_deployment(name, region=None, region_at_runtime=False)#

Deletes SageMaker model, endpoint config and endpoint.

Parameters:

  • name (str) – The prefix to be added to the task names.

  • region (str | None) – The region to use for SageMaker API calls.

  • region_at_runtime (bool) – Set this to True if you want to provide the region at runtime.

Return type:

ImperativeWorkflow