Merge branch 'master' into dw-tlv-3x-updates

Author: secastillo

src/sagemaker/image_uri_config/stabilityai.json

@@ -0,0 +1,48 @@
+{
+    "inference": {
+        "processors": ["gpu"],
+        "version_aliases": {
+            "0.1": "0.1.0"
+        },
+        "versions": {
+            "0.1.0": {
+                "py_versions": ["py310"],
+                "registries": {
+                    "af-south-1": "626614931356",
+                    "il-central-1": "780543022126",
+                    "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",
+                    "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-west-1": "763104351884",
+                    "us-west-2": "763104351884"
+                },
+                "tag_prefix": "2.0.1-sgm0.1.0",
+                "repository": "stabilityai-pytorch-inference",
+                "container_version": {
+                    "gpu": "cu118-ubuntu20.04-sagemaker"
+                }
+            }
+        }
+    }
+}

src/sagemaker/image_uris.py

@@ -39,6 +39,7 @@
 TRAINIUM_ALLOWED_FRAMEWORKS = "pytorch"
 INFERENCE_GRAVITON = "inference_graviton"
 DATA_WRANGLER_FRAMEWORK = "data-wrangler"
+STABILITYAI_FRAMEWORK = "stabilityai"
 
 
 @override_pipeline_parameter_var
@@ -476,7 +477,11 @@ def _validate_version_and_set_if_needed(version, config, framework):
 
         return available_versions[0]
 
-    if version is None and framework in [DATA_WRANGLER_FRAMEWORK, HUGGING_FACE_LLM_FRAMEWORK]:
+    if version is None and framework in [
+        DATA_WRANGLER_FRAMEWORK,
+        HUGGING_FACE_LLM_FRAMEWORK,
+        STABILITYAI_FRAMEWORK,
+    ]:
         version = _get_latest_versions(available_versions)
 
     _validate_arg(version, available_versions + aliased_versions, "{} version".format(framework))

src/sagemaker/stabilityai/__init__.py

@@ -0,0 +1,16 @@
+# 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.
+"""StabilityAI module."""
+from __future__ import absolute_import
+
+from sagemaker.stabilityai.stability_utils import get_stabilityai_image_uri  # noqa: F401

src/sagemaker/stabilityai/stability_utils.py

@@ -0,0 +1,50 @@
+# 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.
+"""Utility functions."""
+
+from __future__ import absolute_import
+
+from typing import Optional
+
+from sagemaker import image_uris
+from sagemaker.session import Session
+
+
+def get_stabilityai_image_uri(
+    session: Optional[Session] = None,
+    region: Optional[str] = None,
+    version: Optional[str] = None,
+    image_scope: Optional[str] = "inference",
+) -> str:
+    """Very basic utility function to fetch image URI of StabilityAI images.
+
+    Args:
+        session (Session): SageMaker session.
+        region (str): AWS region of image URI.
+        version (str): Framework version. Latest version used if not specified.
+        image_scope (str): Image type. e.g. inference, training
+    Returns:
+        Image URI string.
+    """
+
+    if region is None:
+        if session is None:
+            region = Session().boto_session.region_name
+        else:
+            region = session.boto_session.region_name
+    return image_uris.retrieve(
+        framework="stabilityai",
+        region=region,
+        version=version,
+        image_scope=image_scope,
+    )

tests/scripts/run-notebook-test.sh

@@ -123,12 +123,14 @@ echo "set SAGEMAKER_ROLE_ARN=$SAGEMAKER_ROLE_ARN"
 --platformIdentifier notebook-al2-v2 \
 --consider-skips-failures \
 ./amazon-sagemaker-examples/sagemaker_processing/spark_distributed_data_processing/sagemaker-spark-processing.ipynb \
-./amazon-sagemaker-examples/advanced_functionality/kmeans_bring_your_own_model/kmeans_bring_your_own_model.ipynb \
 ./amazon-sagemaker-examples/advanced_functionality/tensorflow_iris_byom/tensorflow_BYOM_iris.ipynb \
 ./amazon-sagemaker-examples/sagemaker-python-sdk/1P_kmeans_highlevel/kmeans_mnist.ipynb \
 ./amazon-sagemaker-examples/sagemaker-python-sdk/scikit_learn_randomforest/Sklearn_on_SageMaker_end2end.ipynb \
 ./amazon-sagemaker-examples/sagemaker-pipelines/tabular/abalone_build_train_deploy/sagemaker-pipelines-preprocess-train-evaluate-batch-transform.ipynb \
 
+# Skipping test until fix in example notebook to move to new conda environment
+#./amazon-sagemaker-examples/advanced_functionality/kmeans_bring_your_own_model/kmeans_bring_your_own_model.ipynb \
+
 # Skipping test until fix in example notebook to install docker-compose is complete
 #./amazon-sagemaker-examples/sagemaker-python-sdk/tensorflow_moving_from_framework_mode_to_script_mode/tensorflow_moving_from_framework_mode_to_script_mode.ipynb \
 

tests/unit/sagemaker/image_uris/expected_uris.py

@@ -91,6 +91,11 @@ def huggingface_llm_framework_uri(
     return IMAGE_URI_FORMAT.format(account, region, domain, repo, tag)
 
 
+def stabilityai_framework_uri(repo, account, 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_stabilityai.py

@@ -0,0 +1,59 @@
+# 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.stabilityai import get_stabilityai_image_uri
+from tests.unit.sagemaker.image_uris import expected_uris
+
+ACCOUNTS = {
+    "af-south-1": "626614931356",
+    "il-central-1": "780543022126",
+    "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",
+    "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",
+}
+SAI_VERSIONS = ["0.1.0"]
+SAI_VERSIONS_MAPPING = {"0.1.0": "2.0.1-sgm0.1.0-gpu-py310-cu118-ubuntu20.04-sagemaker"}
+
+
+@pytest.mark.parametrize("version", SAI_VERSIONS)
+def test_stabilityai_image_uris(version):
+    for region in ACCOUNTS.keys():
+        result = get_stabilityai_image_uri(region=region, version=version)
+        expected = expected_uris.stabilityai_framework_uri(
+            "stabilityai-pytorch-inference",
+            ACCOUNTS[region],
+            SAI_VERSIONS_MAPPING[version],
+            region=region,
+        )
+        assert expected == result