Merge branch 'master' into feat/jumpstart-construct-payload-utility

Author: evakravi

CHANGELOG.md

@@ -1,5 +1,28 @@
 # Changelog
 
+## v2.193.0 (2023-10-18)
+
+### Features
+
+ * jumpstart model artifact instance type variants
+ * jumpstart instance specific hyperparameters
+ * Feature Processor event based triggers (#1132)
+ * Support job checkpoint in remote function
+ * jumpstart model package arn instance type variants
+
+### Bug Fixes and Other Changes
+
+ * Fix hyperlinks in feature_processor.scheduler parameter descriptions
+ * add image_uris_unit_test pytest mark
+ * bump apache-airflow to `v2.7.2`
+ * clone distribution in validate_distribution
+ * fix flaky Inference Recommender integration tests
+
+### Documentation Changes
+
+ * Update PipelineModel.register documentation
+ * specify that input_shape in no longer required for torch 2.0 mod…
+
 ## v2.192.1 (2023-10-13)
 
 ### Bug Fixes and Other Changes

VERSION

@@ -1 +1 @@
-2.192.2.dev0
+2.193.1.dev0

doc/amazon_sagemaker_featurestore.rst

@@ -230,9 +230,11 @@ The following code from the fraud detection example shows a minimal
        enable_online_store=True
    )
 
-Creating a feature group takes time as the data is loaded. You will need
-to wait until it is created before you can use it. You can check status
-using the following method.
+Creating a feature group takes time as the data is loaded. You will
+need to wait until it is created before you can use it. You can
+check status using the following method. Note that it can take
+approximately 10-15 minutes to provision an online ``FeatureGroup``
+with the ``InMemory`` ``StorageType``.
 
 .. code:: python
 
@@ -480,7 +482,9 @@ Feature Store `DatasetBuilder API Reference
 .. rubric:: Delete a feature group
    :name: bCe9CA61b78
 
-You can delete a feature group with the ``delete`` function.
+You can delete a feature group with the ``delete`` function. Note that it
+can take approximately 10-15 minutes to delete an online ``FeatureGroup``
+with the ``InMemory`` ``StorageType``.
 
 .. code:: python
 

requirements/extras/test_requirements.txt

@@ -12,7 +12,7 @@ awslogs==0.14.0
 black==22.3.0
 stopit==1.1.2
 # Update tox.ini to have correct version of airflow constraints file
-apache-airflow==2.7.1
+apache-airflow==2.7.2
 apache-airflow-providers-amazon==7.2.1
 attrs>=23.1.0,<24
 fabric==2.6.0

src/sagemaker/feature_store/feature_processor/feature_scheduler.py

@@ -281,13 +281,13 @@ def schedule(
     Args:
         pipeline_name (str): The SageMaker Pipeline name that will be scheduled.
         schedule_expression (str): The expression that defines when the schedule runs. It supports
-            at expression, rate expression and cron expression. See https://docs.aws.amazon.com/
-            scheduler/latest/APIReference/API_CreateSchedule.html#scheduler-CreateSchedule-request
-            -ScheduleExpression for more details.
+            at expression, rate expression and cron expression. See '''https://docs.aws.amazon.com\
+                /scheduler/latest/APIReference/API_CreateSchedule.html#scheduler-CreateSchedule-\
+                request-ScheduleExpression''' for more details.
         state (str): Specifies whether the schedule is enabled or disabled. Valid values are
-            ENABLED and DISABLED. See https://docs.aws.amazon.com/scheduler/latest/APIReference/
-            API_CreateSchedule.html#scheduler-CreateSchedule-request-State for more details.
-            If not specified, it will default to ENABLED.
+            ENABLED and DISABLED. See '''https://docs.aws.amazon.com/scheduler/latest/APIReference\
+                /API_CreateSchedule.html#scheduler-CreateSchedule-request-State'''
+            for more details. If not specified, it will default to ENABLED.
         start_date (Optional[datetime]): The date, in UTC, after which the schedule can begin
             invoking its target. Depending on the schedule’s recurrence expression, invocations
             might occur on, or after, the StartDate you specify.

src/sagemaker/jumpstart/artifacts/model_uris.py

@@ -28,12 +28,68 @@
     verify_model_region_and_return_specs,
 )
 from sagemaker.session import Session
+from sagemaker.jumpstart.types import JumpStartModelSpecs
+
+
+def _retrieve_hosting_prepacked_artifact_key(
+    model_specs: JumpStartModelSpecs, instance_type: str
+) -> str:
+    """Returns instance specific hosting prepacked artifact key or default one as fallback."""
+    instance_specific_prepacked_hosting_artifact_key: Optional[str] = (
+        model_specs.hosting_instance_type_variants.get_instance_specific_prepacked_artifact_key(
+            instance_type=instance_type
+        )
+        if instance_type
+        and getattr(model_specs, "hosting_instance_type_variants", None) is not None
+        else None
+    )
+
+    default_prepacked_hosting_artifact_key: Optional[str] = getattr(
+        model_specs, "hosting_prepacked_artifact_key"
+    )
+
+    return (
+        instance_specific_prepacked_hosting_artifact_key or default_prepacked_hosting_artifact_key
+    )
+
+
+def _retrieve_hosting_artifact_key(model_specs: JumpStartModelSpecs, instance_type: str) -> str:
+    """Returns instance specific hosting artifact key or default one as fallback."""
+    instance_specific_hosting_artifact_key: Optional[str] = (
+        model_specs.hosting_instance_type_variants.get_instance_specific_artifact_key(
+            instance_type=instance_type
+        )
+        if instance_type
+        and getattr(model_specs, "hosting_instance_type_variants", None) is not None
+        else None
+    )
+
+    default_hosting_artifact_key: str = model_specs.hosting_artifact_key
+
+    return instance_specific_hosting_artifact_key or default_hosting_artifact_key
+
+
+def _retrieve_training_artifact_key(model_specs: JumpStartModelSpecs, instance_type: str) -> str:
+    """Returns instance specific training artifact key or default one as fallback."""
+    instance_specific_training_artifact_key: Optional[str] = (
+        model_specs.training_instance_type_variants.get_instance_specific_artifact_key(
+            instance_type=instance_type
+        )
+        if instance_type
+        and getattr(model_specs, "training_instance_type_variants", None) is not None
+        else None
+    )
+
+    default_training_artifact_key: str = model_specs.training_artifact_key
+
+    return instance_specific_training_artifact_key or default_training_artifact_key
 
 
 def _retrieve_model_uri(
     model_id: str,
     model_version: str,
     model_scope: Optional[str] = None,
+    instance_type: Optional[str] = None,
     region: Optional[str] = None,
     tolerate_vulnerable_model: bool = False,
     tolerate_deprecated_model: bool = False,
@@ -50,6 +106,7 @@ def _retrieve_model_uri(
             artifact S3 URI.
         model_scope (str): The model type, i.e. what it is used for.
             Valid values: "training" and "inference".
+        instance_type (str): The ML compute instance type for the specified scope. (Default: None).
         region (str): Region for which to retrieve model S3 URI. (Default: None).
         tolerate_vulnerable_model (bool): True if vulnerable versions of model
             specifications should be tolerated (exception not raised). If False, raises an
@@ -84,14 +141,21 @@ def _retrieve_model_uri(
         sagemaker_session=sagemaker_session,
     )
 
+    model_artifact_key: str
+
     if model_scope == JumpStartScriptScope.INFERENCE:
+
+        is_prepacked = not model_specs.use_inference_script_uri()
+
         model_artifact_key = (
-            getattr(model_specs, "hosting_prepacked_artifact_key", None)
-            or model_specs.hosting_artifact_key
+            _retrieve_hosting_prepacked_artifact_key(model_specs, instance_type)
+            if is_prepacked
+            else _retrieve_hosting_artifact_key(model_specs, instance_type)
         )
 
     elif model_scope == JumpStartScriptScope.TRAINING:
-        model_artifact_key = model_specs.training_artifact_key
+
+        model_artifact_key = _retrieve_training_artifact_key(model_specs, instance_type)
 
     bucket = os.environ.get(
         ENV_VARIABLE_JUMPSTART_MODEL_ARTIFACT_BUCKET_OVERRIDE

src/sagemaker/jumpstart/factory/estimator.py

@@ -475,6 +475,7 @@ def _add_model_uri_to_kwargs(kwargs: JumpStartEstimatorInitKwargs) -> JumpStartE
             tolerate_deprecated_model=kwargs.tolerate_deprecated_model,
             tolerate_vulnerable_model=kwargs.tolerate_vulnerable_model,
             sagemaker_session=kwargs.sagemaker_session,
+            instance_type=kwargs.instance_type,
         )
 
         if (

src/sagemaker/jumpstart/factory/model.py

@@ -215,6 +215,7 @@ def _add_model_data_to_kwargs(kwargs: JumpStartModelInitKwargs) -> JumpStartMode
         tolerate_deprecated_model=kwargs.tolerate_deprecated_model,
         tolerate_vulnerable_model=kwargs.tolerate_vulnerable_model,
         sagemaker_session=kwargs.sagemaker_session,
+        instance_type=kwargs.instance_type,
     )
 
     if isinstance(model_data, str) and model_data.startswith("s3://") and model_data.endswith("/"):

src/sagemaker/jumpstart/types.py

@@ -405,6 +405,62 @@ def to_json(self) -> Dict[str, Any]:
         json_obj = {att: getattr(self, att) for att in self.__slots__ if hasattr(self, att)}
         return json_obj
 
+    def get_instance_specific_prepacked_artifact_key(self, instance_type: str) -> Optional[str]:
+        """Returns instance specific model artifact key.
+
+        Returns None if a model, instance type tuple does not have specific
+        artifact key.
+        """
+
+        return self._get_instance_specific_property(
+            instance_type=instance_type, property_name="prepacked_artifact_key"
+        )
+
+    def get_instance_specific_artifact_key(self, instance_type: str) -> Optional[str]:
+        """Returns instance specific model artifact key.
+
+        Returns None if a model, instance type tuple does not have specific
+        artifact key.
+        """
+
+        return self._get_instance_specific_property(
+            instance_type=instance_type, property_name="artifact_key"
+        )
+
+    def _get_instance_specific_property(
+        self, instance_type: str, property_name: str
+    ) -> Optional[str]:
+        """Returns instance specific property.
+
+        If a value exists for both the instance family and instance type,
+        the instance type value is chosen.
+
+        Returns None if a (model, instance type, property name) tuple does not have
+        specific prepacked artifact key.
+        """
+
+        if self.variants is None:
+            return None
+
+        instance_specific_property: Optional[str] = (
+            self.variants.get(instance_type, {}).get("properties", {}).get(property_name, None)
+        )
+
+        if instance_specific_property:
+            return instance_specific_property
+
+        instance_type_family = get_instance_type_family(instance_type)
+
+        instance_family_property: Optional[str] = (
+            self.variants.get(instance_type_family, {})
+            .get("properties", {})
+            .get(property_name, None)
+            if instance_type_family not in {"", None}
+            else None
+        )
+
+        return instance_family_property
+
     def get_instance_specific_hyperparameters(
         self, instance_type: str
     ) -> List[JumpStartHyperparameter]:

src/sagemaker/model.py

@@ -623,14 +623,7 @@ def prepare_container_def(
         )
         deploy_env = copy.deepcopy(self.env)
         if self.source_dir or self.dependencies or self.entry_point or self.git_config:
-            is_repack = (
-                self.source_dir
-                and self.entry_point
-                and not (
-                    (self.key_prefix and issubclass(type(self), FrameworkModel)) or self.git_config
-                )
-            )
-            self._upload_code(deploy_key_prefix, repack=is_repack)
+            self._upload_code(deploy_key_prefix, repack=self.is_repack())
             deploy_env.update(self._script_mode_env_vars())
 
         return sagemaker.container_def(
@@ -640,6 +633,14 @@ def prepare_container_def(
             image_config=self.image_config,
         )
 
+    def is_repack(self) -> bool:
+        """Whether the source code needs to be repacked before uploading to S3.
+
+        Returns:
+            bool: if the source need to be repacked or not
+        """
+        return self.source_dir and self.entry_point and not self.git_config
+
     def _upload_code(self, key_prefix: str, repack: bool = False) -> None:
         """Uploads code to S3 to be used with script mode with SageMaker inference.
 
@@ -1775,6 +1776,14 @@ def __init__(
             **kwargs,
         )
 
+    def is_repack(self) -> bool:
+        """Whether the source code needs to be repacked before uploading to S3.
+
+        Returns:
+            bool: if the source need to be repacked or not
+        """
+        return self.source_dir and self.entry_point and not (self.key_prefix or self.git_config)
+
 
 # works for MODEL_PACKAGE_ARN with or without version info.
 MODEL_PACKAGE_ARN_PATTERN = r"arn:aws:sagemaker:(.*?):(.*?):model-package/(.*?)(?:/(\d+))?$"

src/sagemaker/model_uris.py

@@ -30,6 +30,7 @@ def retrieve(
     model_id: Optional[str] = None,
     model_version: Optional[str] = None,
     model_scope: Optional[str] = None,
+    instance_type: Optional[str] = None,
     tolerate_vulnerable_model: bool = False,
     tolerate_deprecated_model: bool = False,
     sagemaker_session: Session = DEFAULT_JUMPSTART_SAGEMAKER_SESSION,
@@ -44,6 +45,7 @@ def retrieve(
             the model artifact S3 URI.
         model_scope (str): The model type.
             Valid values: "training" and "inference".
+        instance_type (str): The ML compute instance type for the specified scope. (Default: None).
         tolerate_vulnerable_model (bool): ``True`` if vulnerable versions of model
             specifications should be tolerated without raising an exception. If ``False``, raises an
             exception if the script used by this version of the model has dependencies with known
@@ -71,11 +73,12 @@ def retrieve(
         )
 
     return artifacts._retrieve_model_uri(
-        model_id,
-        model_version,  # type: ignore
-        model_scope,
-        region,
-        tolerate_vulnerable_model,
-        tolerate_deprecated_model,
+        model_id=model_id,
+        model_version=model_version,  # type: ignore
+        model_scope=model_scope,
+        instance_type=instance_type,
+        region=region,
+        tolerate_vulnerable_model=tolerate_vulnerable_model,
+        tolerate_deprecated_model=tolerate_deprecated_model,
         sagemaker_session=sagemaker_session,
     )

tests/conftest.py

@@ -14,6 +14,7 @@
 
 import json
 import os
+import pathlib
 
 import boto3
 import pytest
@@ -86,6 +87,16 @@
 PYTORCH_RENEWED_GPU = "ml.g4dn.xlarge"
 
 
+image_uris_unit_tests_dir = pathlib.Path("tests/unit/sagemaker/image_uris")
+
+
+def pytest_collection_modifyitems(config, items):
+    for item in items:
+        testmod = pathlib.Path(item.fspath)
+        if config.rootdir / image_uris_unit_tests_dir in testmod.parents:
+            item.add_marker(pytest.mark.image_uris_unit_test)
+
+
 def pytest_addoption(parser):
     parser.addoption("--sagemaker-client-config", action="store", default=None)
     parser.addoption("--sagemaker-runtime-config", action="store", default=None)