feature: Add huggingface-llm 0.6.0 dlc images

src/sagemaker/huggingface/__init__.py

@@ -16,5 +16,6 @@
 from sagemaker.huggingface.estimator import HuggingFace  # noqa: F401
 from sagemaker.huggingface.model import HuggingFaceModel, HuggingFacePredictor  # noqa: F401
 from sagemaker.huggingface.processing import HuggingFaceProcessor  # noqa:F401
+from sagemaker.huggingface.llm_utils import get_huggingface_llm_image_uri  # noqa: F401
 
 from sagemaker.huggingface.training_compiler.config import TrainingCompilerConfig  # noqa: F401

src/sagemaker/huggingface/llm_utils.py

@@ -0,0 +1,56 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You
+# may not use this file except in compliance with the License. A copy of
+# the License is located at
+#
+#     http://aws.amazon.com/apache2.0/
+#
+# or in the "license" file accompanying this file. This file is
+# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
+# ANY KIND, either express or implied. See the License for the specific
+# language governing permissions and limitations under the License.
+"""Functions for generating ECR image URIs for pre-built SageMaker Docker images."""
+from __future__ import absolute_import
+
+from typing import Optional
+
+from sagemaker import image_uris
+from sagemaker.session import Session
+
+
+def get_huggingface_llm_image_uri(
+    backend: str,
+    session: Optional[Session] = None,
+    region: Optional[str] = None,
+    version: Optional[str] = None,
+) -> str:
+    """Retrieves the image URI for inference.
+
+    Args:
+        backend (str): The backend to use. Valid values include "huggingface" and "lmi".
+        session (Session): The SageMaker Session to use. (Default: None).
+        region (str): The AWS region to use for image URI. (default: None).
+        version (str): The framework version for which to retrieve an
+            image URI. If no version is set, defaults to latest version. (default: None).
+
+    Returns:
+        str: The image URI string.
+    """
+
+    if region is None:
+        if session is None:
+            region = Session().boto_session.region_name
+        else:
+            region = session.boto_session.region_name
+    if backend == "huggingface":
+        return image_uris.retrieve(
+            "huggingface-llm",
+            region=region,
+            version=version,
+            image_scope="inference",
+        )
+    if backend == "lmi":
+        version = version or "0.22.1"
+        return image_uris.retrieve(framework="djl-deepspeed", region=region, version=version)
+    raise ValueError("Unsupported backend: %s" % backend)

src/sagemaker/image_uri_config/huggingface-llm.json

@@ -0,0 +1,51 @@
+{
+    "inference": {
+        "processors": ["gpu"],
+        "version_aliases": {
+            "0.6": "0.6.0"
+        },
+        "versions": {
+            "0.6.0": {
+                "py_versions": ["py39"],
+                "registries": {
+                    "af-south-1": "626614931356",
+                    "ap-east-1": "871362719292",
+                    "ap-northeast-1": "763104351884",
+                    "ap-northeast-2": "763104351884",
+                    "ap-northeast-3": "364406365360",
+                    "ap-south-1": "763104351884",
+                    "ap-south-2": "772153158452",
+                    "ap-southeast-1": "763104351884",
+                    "ap-southeast-2": "763104351884",
+                    "ap-southeast-3": "907027046896",
+                    "ap-southeast-4": "457447274322",
+                    "ca-central-1": "763104351884",
+                    "cn-north-1": "727897471807",
+                    "cn-northwest-1": "727897471807",
+                    "eu-central-1": "763104351884",
+                    "eu-central-2": "380420809688",
+                    "eu-north-1": "763104351884",
+                    "eu-west-1": "763104351884",
+                    "eu-west-2": "763104351884",
+                    "eu-west-3": "763104351884",
+                    "eu-south-1": "692866216735",
+                    "eu-south-2": "503227376785",
+                    "me-south-1": "217643126080",
+                    "me-central-1": "914824155844",
+                    "sa-east-1": "763104351884",
+                    "us-east-1": "763104351884",
+                    "us-east-2": "763104351884",
+                    "us-gov-east-1": "446045086412",
+                    "us-gov-west-1": "442386744353",
+                    "us-iso-east-1": "886529160074",
+                    "us-isob-east-1": "094389454867",
+                    "us-west-1": "763104351884",
+                    "us-west-2": "763104351884"
+                },
+                "tag_prefix": "2.0.0-tgi0.6.0",
+                "repository": "huggingface-pytorch-tgi-inference",
+                "container_version": {"gpu": "cu118-ubuntu20.04"}
+            }
+        }
+    }
+}

src/sagemaker/image_uris.py

@@ -32,6 +32,7 @@
 
 ECR_URI_TEMPLATE = "{registry}.dkr.{hostname}/{repository}"
 HUGGING_FACE_FRAMEWORK = "huggingface"
+HUGGING_FACE_LLM_FRAMEWORK = "huggingface-llm"
 XGBOOST_FRAMEWORK = "xgboost"
 SKLEARN_FRAMEWORK = "sklearn"
 TRAINIUM_ALLOWED_FRAMEWORKS = "pytorch"
@@ -462,7 +463,7 @@ def _validate_version_and_set_if_needed(version, config, framework):
 
         return available_versions[0]
 
-    if version is None and framework in [DATA_WRANGLER_FRAMEWORK]:
+    if version is None and framework in [DATA_WRANGLER_FRAMEWORK, HUGGING_FACE_LLM_FRAMEWORK]:
         version = _get_latest_versions(available_versions)
 
     _validate_arg(version, available_versions + aliased_versions, "{} version".format(framework))

tests/unit/sagemaker/image_uris/expected_uris.py

@@ -80,6 +80,17 @@ def djl_framework_uri(repo, account, djl_version, primary_framework, region=REGI
     return IMAGE_URI_FORMAT.format(account, region, domain, repo, tag)
 
 
+def huggingface_llm_framework_uri(
+    repo,
+    account,
+    version,
+    tag,
+    region=REGION,
+):
+    domain = ALTERNATE_DOMAINS.get(region, DOMAIN)
+    return IMAGE_URI_FORMAT.format(account, region, domain, repo, tag)
+
+
 def base_python_uri(repo, account, region=REGION):
     domain = ALTERNATE_DOMAINS.get(region, DOMAIN)
     tag = "1.0"

tests/unit/sagemaker/image_uris/test_huggingface_llm.py

@@ -0,0 +1,75 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You
+# may not use this file except in compliance with the License. A copy of
+# the License is located at
+#
+#     http://aws.amazon.com/apache2.0/
+#
+# or in the "license" file accompanying this file. This file is
+# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
+# ANY KIND, either express or implied. See the License for the specific
+# language governing permissions and limitations under the License.
+from __future__ import absolute_import
+
+import pytest
+
+from sagemaker.huggingface import get_huggingface_llm_image_uri
+from tests.unit.sagemaker.image_uris import expected_uris
+
+ACCOUNTS = {
+    "af-south-1": "626614931356",
+    "ap-east-1": "871362719292",
+    "ap-northeast-1": "763104351884",
+    "ap-northeast-2": "763104351884",
+    "ap-northeast-3": "364406365360",
+    "ap-south-1": "763104351884",
+    "ap-southeast-1": "763104351884",
+    "ap-southeast-2": "763104351884",
+    "ap-southeast-3": "907027046896",
+    "ca-central-1": "763104351884",
+    "cn-north-1": "727897471807",
+    "cn-northwest-1": "727897471807",
+    "eu-central-1": "763104351884",
+    "eu-north-1": "763104351884",
+    "eu-west-1": "763104351884",
+    "eu-west-2": "763104351884",
+    "eu-west-3": "763104351884",
+    "eu-south-1": "692866216735",
+    "me-south-1": "217643126080",
+    "sa-east-1": "763104351884",
+    "us-east-1": "763104351884",
+    "us-east-2": "763104351884",
+    "us-west-1": "763104351884",
+    "us-west-2": "763104351884",
+}
+HF_VERSIONS = ["0.6.0"]
+LMI_VERSIONS = ["0.22.1"]
+HF_VERSIONS_MAPPING = {"0.6.0": "2.0.0-tgi0.6.0-gpu-py39-cu118-ubuntu20.04"}
+LMI_VERSIONS_MAPPING = {"0.22.1": "deepspeed0.8.3-cu118"}
+
+
+@pytest.mark.parametrize("version", HF_VERSIONS)
+def test_huggingface(version):
+    for region in ACCOUNTS.keys():
+        uri = get_huggingface_llm_image_uri("huggingface", region=region, version=version)
+
+        expected = expected_uris.huggingface_llm_framework_uri(
+            "huggingface-pytorch-tgi-inference",
+            ACCOUNTS[region],
+            version,
+            HF_VERSIONS_MAPPING[version],
+            region=region,
+        )
+        assert expected == uri
+
+
+@pytest.mark.parametrize("version", LMI_VERSIONS)
+def test_lmi(version):
+    for region in ACCOUNTS.keys():
+        uri = get_huggingface_llm_image_uri("lmi", region=region, version=version)
+
+        expected = expected_uris.djl_framework_uri(
+            "djl-inference", ACCOUNTS[region], version, LMI_VERSIONS_MAPPING[version], region=region
+        )
+        assert expected == uri