Merge branch 'master' into master

Author: Qingzi-Lan

CHANGELOG.md

@@ -1,5 +1,15 @@
 # Changelog
 
+## v2.181.0 (2023-08-28)
+
+### Features
+
+ * StabilityAI DLC Image URIs
+
+### Bug Fixes and Other Changes
+
+ * temporarily skip kmeans notebook
+
 ## v2.180.0 (2023-08-24)
 
 ### Features

VERSION

@@ -1 +1 @@
-2.180.1.dev0
+2.181.1.dev0

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/remote_function/client.py

@@ -83,6 +83,8 @@ def remote(
     volume_size: int = 30,
     encrypt_inter_container_traffic: bool = None,
     spark_config: SparkConfig = None,
+    use_spot_instances=False,
+    max_wait_time_in_seconds=None,
 ):
     """Decorator for running the annotated function as a SageMaker training job.
 
@@ -255,6 +257,14 @@ def remote(
           Spark image. If ``spark_config`` is specified, a SageMaker Spark image uri
           will be used for training. Note that ``image_uri`` can not be specified at the
           same time otherwise a ``ValueError`` is thrown. Defaults to ``None``.
+
+        use_spot_instances (bool): Specifies whether to use SageMaker Managed Spot instances for
+          training. If enabled then the ``max_wait_time_in_seconds`` arg should also be set.
+          Defaults to ``False``.
+
+        max_wait_time_in_seconds (int): Timeout in seconds waiting for spot training job.
+          After this amount of time Amazon SageMaker will stop waiting for managed spot training
+          job to complete. Defaults to ``None``.
     """
 
     def _remote(func):
@@ -284,6 +294,8 @@ def _remote(func):
             volume_size=volume_size,
             encrypt_inter_container_traffic=encrypt_inter_container_traffic,
             spark_config=spark_config,
+            use_spot_instances=use_spot_instances,
+            max_wait_time_in_seconds=max_wait_time_in_seconds,
         )
 
         @functools.wraps(func)
@@ -492,6 +504,8 @@ def __init__(
         volume_size: int = 30,
         encrypt_inter_container_traffic: bool = None,
         spark_config: SparkConfig = None,
+        use_spot_instances=False,
+        max_wait_time_in_seconds=None,
     ):
         """Constructor for RemoteExecutor
 
@@ -670,6 +684,14 @@ def __init__(
               Spark image. If ``spark_config`` is specified, a SageMaker Spark image uri
               will be used for training. Note that ``image_uri`` can not be specified at the
               same time otherwise a ``ValueError`` is thrown. Defaults to ``None``.
+
+            use_spot_instances (bool): Specifies whether to use SageMaker Managed Spot instances for
+              training. If enabled then the ``max_wait_time_in_seconds`` arg should also be set.
+              Defaults to ``False``.
+
+            max_wait_time_in_seconds (int): Timeout in seconds waiting for spot training job.
+              After this amount of time Amazon SageMaker will stop waiting for managed spot training
+              job to complete. Defaults to ``None``.
         """
         self.max_parallel_jobs = max_parallel_jobs
 
@@ -707,6 +729,8 @@ def __init__(
             volume_size=volume_size,
             encrypt_inter_container_traffic=encrypt_inter_container_traffic,
             spark_config=spark_config,
+            use_spot_instances=use_spot_instances,
+            max_wait_time_in_seconds=max_wait_time_in_seconds,
         )
 
         self._state_condition = threading.Condition()

src/sagemaker/remote_function/job.py

@@ -191,6 +191,8 @@ def __init__(
         volume_size: int = 30,
         encrypt_inter_container_traffic: bool = None,
         spark_config: SparkConfig = None,
+        use_spot_instances=False,
+        max_wait_time_in_seconds=None,
     ):
         """Initialize a _JobSettings instance which configures the remote job.
 
@@ -353,6 +355,14 @@ def __init__(
               Spark image. If ``spark_config`` is specified, a SageMaker Spark image uri
               will be used for training. Note that ``image_uri`` can not be specified at the
               same time otherwise a ``ValueError`` is thrown. Defaults to ``None``.
+
+            use_spot_instances (bool): Specifies whether to use SageMaker Managed Spot instances for
+              training. If enabled then the ``max_wait`` arg should also be set.
+              Defaults to ``False``.
+
+            max_wait_time_in_seconds (int): Timeout in seconds waiting for spot training job.
+              After this amount of time Amazon SageMaker will stop waiting for managed spot
+              training job to complete. Defaults to ``None``.
         """
         self.sagemaker_session = sagemaker_session or Session()
         self.environment_variables = resolve_value_from_config(
@@ -439,6 +449,8 @@ def __init__(
         self.max_retry_attempts = max_retry_attempts
         self.keep_alive_period_in_seconds = keep_alive_period_in_seconds
         self.spark_config = spark_config
+        self.use_spot_instances = use_spot_instances
+        self.max_wait_time_in_seconds = max_wait_time_in_seconds
         self.job_conda_env = resolve_value_from_config(
             direct_input=job_conda_env,
             config_path=REMOTE_FUNCTION_JOB_CONDA_ENV,
@@ -648,12 +660,16 @@ def start(job_settings: _JobSettings, func, func_args, func_kwargs, run_info=Non
 
         stored_function.save(func, *func_args, **func_kwargs)
 
+        stopping_condition = {
+            "MaxRuntimeInSeconds": job_settings.max_runtime_in_seconds,
+        }
+        if job_settings.max_wait_time_in_seconds is not None:
+            stopping_condition["MaxWaitTimeInSeconds"] = job_settings.max_wait_time_in_seconds
+
         request_dict = dict(
             TrainingJobName=job_name,
             RoleArn=job_settings.role,
-            StoppingCondition={
-                "MaxRuntimeInSeconds": job_settings.max_runtime_in_seconds,
-            },
+            StoppingCondition=stopping_condition,
             RetryStrategy={"MaximumRetryAttempts": job_settings.max_retry_attempts},
         )
 
@@ -742,6 +758,8 @@ def start(job_settings: _JobSettings, func, func_args, func_kwargs, run_info=Non
         if job_settings.vpc_config:
             request_dict["VpcConfig"] = job_settings.vpc_config
 
+        request_dict["EnableManagedSpotTraining"] = job_settings.use_spot_instances
+
         request_dict["Environment"] = job_settings.environment_variables
 
         extended_request = _extend_spark_config_to_request(request_dict, job_settings, s3_base_uri)

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/integ/sagemaker/remote_function/test_decorator.py

@@ -601,6 +601,24 @@ def get_file_content(file_names):
         assert "line 2: bws: command not found" in str(e)
 
 
+def test_decorator_with_spot_instances(
+    sagemaker_session, dummy_container_without_error, cpu_instance_type
+):
+    @remote(
+        role=ROLE,
+        image_uri=dummy_container_without_error,
+        instance_type=cpu_instance_type,
+        sagemaker_session=sagemaker_session,
+        use_spot_instances=True,
+        max_wait_time_in_seconds=48 * 60 * 60,
+    )
+    def divide(x, y):
+        return x / y
+
+    assert divide(10, 2) == 5
+    assert divide(20, 2) == 10
+
+
 @pytest.mark.skip
 def test_decorator_with_spark_job(sagemaker_session, cpu_instance_type):
     @remote(

tests/integ/sagemaker/remote_function/test_executor.py

@@ -195,6 +195,46 @@ def cube(x):
             assert metric_summary.avg == 550
 
 
+def test_executor_submit_using_spot_instances(
+    sagemaker_session, dummy_container_without_error, cpu_instance_type
+):
+    def square_on_spot_instance(x):
+        return x * x
+
+    def cube_on_spot_instance(x):
+        return x * x * x
+
+    with RemoteExecutor(
+        max_parallel_jobs=1,
+        role=ROLE,
+        image_uri=dummy_container_without_error,
+        instance_type=cpu_instance_type,
+        sagemaker_session=sagemaker_session,
+        use_spot_instances=True,
+        max_wait_time_in_seconds=48 * 60 * 60,
+    ) as e:
+        future_1 = e.submit(square_on_spot_instance, 10)
+        future_2 = e.submit(cube_on_spot_instance, 10)
+
+    assert future_1.result() == 100
+    assert future_2.result() == 1000
+
+    assert get_future(future_1._job.job_name, sagemaker_session).result() == 100
+    assert get_future(future_2._job.job_name, sagemaker_session).result() == 1000
+
+    describe_job_1 = next(
+        list_futures(job_name_prefix="square-on-spot-instance", sagemaker_session=sagemaker_session)
+    )._job.describe()
+    assert describe_job_1["EnableManagedSpotTraining"] is True
+    assert describe_job_1["StoppingCondition"]["MaxWaitTimeInSeconds"] == 172800
+
+    describe_job_2 = next(
+        list_futures(job_name_prefix="cube-on-spot-instance", sagemaker_session=sagemaker_session)
+    )._job.describe()
+    assert describe_job_2["EnableManagedSpotTraining"] is True
+    assert describe_job_2["StoppingCondition"]["MaxWaitTimeInSeconds"] == 172800
+
+
 def test_executor_map_with_run(sagemaker_session, dummy_container_without_error, cpu_instance_type):
     def square(x):
         with load_run() as run:

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/feature_store/feature_processor/test_feature_scheduler.py

@@ -873,6 +873,8 @@ def test_remote_decorator_fields_consistency(get_execution_role, session):
         "volume_kms_key",
         "vpc_config",
         "tags",
+        "use_spot_instances",
+        "max_wait_time_in_seconds",
     }
 
     job_settings = _JobSettings(

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"