sagemaker.core.image_retriever.image_retriever

sagemaker.core.image_retriever.image_retriever#

Classes

ImageRetriever()

class sagemaker.core.image_retriever.image_retriever.ImageRetriever[source]#

Bases: object

static retrieve(framework: str, region: str, version: str | None = None, py_version: str | None = None, instance_type: str | None = None, accelerator_type: str | None = None, image_scope: str | None = None, container_version: str = None, distributed: bool = False, smp: bool = False, base_framework_version: str | None = None, training_compiler_config: TrainingCompilerConfig = None, model_id: str | None = None, model_version: str | None = None, sdk_version: str | None = None, inference_tool: str | None = None, serverless_inference_config: ServerlessInferenceConfig = None)[source]#

Retrieves the ECR URI for the Docker image matching the given arguments.

Parameters:

  • framework (str) – The name of the framework or algorithm.

  • region (str) – The AWS region.

  • version (Optional[str]) – The framework or algorithm version. This is required if there is more than one supported version for the given framework or algorithm.

  • py_version (Optional[str]) – The Python version. This is required if there is more than one supported Python version for the given framework version.

  • instance_type (Optional[str]) – The SageMaker instance type. For supported types, see https://aws.amazon.com/sagemaker/pricing. This is required if there are different images for different processor types.

  • accelerator_type (Optional[str]) – Elastic Inference accelerator type. For more, see https://docs.aws.amazon.com/sagemaker/latest/dg/ei.html.

  • image_scope (Optional[str]) – The image type, i.e. what it is used for. Valid values: “training”, “inference”, “inference_graviton”, “eia”. If accelerator_type is set, image_scope is ignored.

  • container_version (Optional[str]) – the version of docker image. Ideally the value of parameter should be created inside the framework. For custom use, see the list of supported container versions: aws/deep-learning-containers (default: None).

  • distributed (bool) – If the image is for running distributed training.

  • smp (bool) – If using the SMP library for distributed training.

  • base_framework_version (Optional[str]) – The base version number of PyTorch or Tensorflow. (default: None).

  • training_compiler_config (TrainingCompilerConfig) – A configuration class for the SageMaker Training Compiler (default: None).

  • model_id (Optional[str]) – The JumpStart model ID for which to retrieve the image URI (default: None).

  • model_version (Optional[str]) – The version of the JumpStart model for which to retrieve the image URI (default: None).

  • hub_arn (Optional[str]) – The arn of the SageMaker Hub for which to retrieve model details from. (Default: None).

  • tolerate_vulnerable_model (bool) – True if vulnerable versions of model specifications should be tolerated without an exception raised. If False, raises an exception if the script used by this version of the model has dependencies with known security vulnerabilities. (Default: False).

  • tolerate_deprecated_model (bool) – True if deprecated versions of model specifications should be tolerated without an exception raised. If False, raises an exception if the version of the model is deprecated. (Default: False).

  • sdk_version (Optional[str]) – the version of python-sdk that will be used in the image retrieval. (default: None).

  • inference_tool (Optional[str]) – the tool that will be used to aid in the inference. Valid values: “neuron, neuronx, None” (default: None).

  • serverless_inference_config (sagemaker.serve.serverless.ServerlessInferenceConfig) – Specifies configuration related to serverless endpoint. Instance type is not provided in serverless inference. So this is used to determine processor type.

  • sagemaker_session (sagemaker.core.helper.Session) – A SageMaker Session object, used for SageMaker interactions. If not specified, one is created using the default AWS configuration chain. (Default: sagemaker.jumpstart.constants.DEFAULT_JUMPSTART_SAGEMAKER_SESSION).

  • config_name (Optional[str]) – Name of the JumpStart Model config to apply. (Default: None).

  • model_type (JumpStartModelType) – The type of the model, can be open weights model or proprietary model. (Default: JumpStartModelType.OPEN_WEIGHTS).

Returns:

The ECR URI for the corresponding SageMaker Docker image.

Return type:

str

Raises:

  • NotImplementedError – If the scope is not supported.

  • ValueError – If the combination of arguments specified is not supported or any PipelineVariable object is passed in.

  • VulnerableJumpStartModelError – If any of the dependencies required by the script have known security vulnerabilities.

  • DeprecatedJumpStartModelError – If the version of the model is deprecated.

static retrieve_base_python_image_uri(region: str, py_version: str = '310') → str[source]#

Retrieves the image URI for base python image.

Parameters:

  • region (str) – The AWS region to use for image URI.

  • py_version (str) – The python version to use for the image.

  • 310 (Default to)

Returns:

The image URI string.

Return type:

str

static retrieve_hugging_face_uri(region: str, version: str | None = None, py_version: str | None = None, instance_type: str | None = None, accelerator_type: str | None = None, image_scope: str | None = None, container_version: str | None = None, distributed: bool = False, base_framework_version: str | None = None, training_compiler_config: TrainingCompilerConfig | None = None, sdk_version: str | None = None, inference_tool: str | None = None, serverless_inference_config: ServerlessInferenceConfig | None = None)[source]#

Retrieves the ECR URI for the Docker image of HuggingFace models.

Parameters:

  • region (str) – The AWS region.

  • version (Optional[str]) – The framework or algorithm version. This is required if there is more than one supported version for the given framework or algorithm.

  • py_version (Optional[str]) – The Python version. This is required if there is more than one supported Python version for the given framework version.

  • instance_type (Optional[str]) – The SageMaker instance type. For supported types, see https://aws.amazon.com/sagemaker/pricing. This is required if there are different images for different processor types.

  • accelerator_type (Optional[str]) – Elastic Inference accelerator type. For more, see https://docs.aws.amazon.com/sagemaker/latest/dg/ei.html.

  • image_scope (Optional[str]) – The image type, i.e. what it is used for. Valid values: “training”, “inference”, “inference_graviton”, “eia”. If accelerator_type is set, image_scope is ignored.

  • container_version (Optional[str]) – the version of docker image. Ideally the value of parameter should be created inside the framework. For custom use, see the list of supported container versions: aws/deep-learning-containers (default: None).

  • distributed (bool) – If the image is for running distributed training.

  • base_framework_version (Optional[str]) – The base version number of PyTorch or Tensorflow. (default: None).

  • training_compiler_config (TrainingCompilerConfig) – A configuration class for the SageMaker Training Compiler (default: None).

  • sdk_version (Optional[str]) – the version of python-sdk that will be used in the image retrieval. (default: None).

  • inference_tool (Optional[str]) – the tool that will be used to aid in the inference. Valid values: “neuron, neuronx, None” (default: None).

  • serverless_inference_config (sagemaker.serve.serverless.ServerlessInferenceConfig) – Specifies configuration related to serverless endpoint. Instance type is not provided in serverless inference. So this is used to determine processor type..

Returns:

The ECR URI for the corresponding SageMaker Docker image.

Return type:

str

static retrieve_jumpstart_uri()[source]#

Retrieves the container image URI for JumpStart models.

TODO: We are changing the Jumpstart utilities and moving from S3 based metadata to API based metadata. Implement this function after utility changes are finished.

static retrieve_pytorch_uri(region: str, version: str | None = None, py_version: str | None = None, instance_type: str | None = None, accelerator_type: str | None = None, image_scope: str | None = None, container_version: str | None = None, distributed: bool = False, smp: bool = False, training_compiler_config: TrainingCompilerConfig | None = None, sdk_version: str | None = None, inference_tool: str | None = None, serverless_inference_config: ServerlessInferenceConfig | None = None)[source]#

Retrieves the ECR URI for the Docker image of PyTorch models.

Parameters:

  • region (str) – The AWS region.

  • version (Optional[str]) – The framework or algorithm version. This is required if there is more than one supported version for the given framework or algorithm.

  • py_version (Optional[str]) – The Python version. This is required if there is more than one supported Python version for the given framework version.

  • instance_type (Optional[str]) – The SageMaker instance type. For supported types, see https://aws.amazon.com/sagemaker/pricing. This is required if there are different images for different processor types.

  • accelerator_type (Optional[str]) – Elastic Inference accelerator type. For more, see https://docs.aws.amazon.com/sagemaker/latest/dg/ei.html.

  • image_scope (Optional[str]) – The image type, i.e. what it is used for. Valid values: “training”, “inference”, “inference_graviton”, “eia”. If accelerator_type is set, image_scope is ignored.

  • container_version (Optional[str]) – the version of docker image. Ideally the value of parameter should be created inside the framework. For custom use, see the list of supported container versions: aws/deep-learning-containers (default: None).

  • distributed (bool) – If the image is for running distributed training.

  • smp (bool) – If using the SMP library for distributed training.

  • training_compiler_config (TrainingCompilerConfig) – A configuration class for the SageMaker Training Compiler (default: None).

  • sdk_version (Optional[str]) – the version of python-sdk that will be used in the image retrieval. (default: None).

  • inference_tool (Optional[str]) – the tool that will be used to aid in the inference. Valid values: “neuron, neuronx, None” (default: None).

  • serverless_inference_config (sagemaker.serve.serverless.ServerlessInferenceConfig) – Specifies configuration related to serverless endpoint. Instance type is not provided in serverless inference. So this is used to determine processor type.

Returns:

The ECR URI for the corresponding SageMaker Docker image.

Return type:

str

Contents