fix: SDK defaults for volume size, JS Estimator image uri region, Predictor str method (#3870)

Author: evakravi

src/sagemaker/base_predictor.py

@@ -47,7 +47,7 @@
     NumpySerializer,
 )
 from sagemaker.session import production_variant, Session
-from sagemaker.utils import name_from_base
+from sagemaker.utils import name_from_base, stringify_object
 
 from sagemaker.model_monitor.model_monitoring import DEFAULT_REPOSITORY_NAME
 
@@ -75,6 +75,10 @@ def content_type(self) -> str:
     def accept(self) -> Tuple[str]:
         """The content type(s) that are expected from the inference server."""
 
+    def __str__(self) -> str:
+        """Overriding str(*) method to make more human-readable."""
+        return stringify_object(self)
+
 
 class Predictor(PredictorBase):
     """Make prediction requests to an Amazon SageMaker endpoint."""

src/sagemaker/estimator.py

@@ -63,6 +63,7 @@
 )
 from sagemaker.inputs import TrainingInput, FileSystemInput
 from sagemaker.instance_group import InstanceGroup
+from sagemaker.utils import instance_supports_kms
 from sagemaker.job import _Job
 from sagemaker.jumpstart.utils import (
     add_jumpstart_tags,
@@ -95,6 +96,7 @@
 )
 from sagemaker.workflow import is_pipeline_variable
 from sagemaker.workflow.entities import PipelineVariable
+from sagemaker.workflow.parameters import ParameterString
 from sagemaker.workflow.pipeline_context import PipelineSession, runnable_by_pipeline
 
 logger = logging.getLogger(__name__)
@@ -599,10 +601,33 @@ def __init__(
         self.output_kms_key = resolve_value_from_config(
             output_kms_key, TRAINING_JOB_KMS_KEY_ID_PATH, sagemaker_session=self.sagemaker_session
         )
-        self.volume_kms_key = resolve_value_from_config(
-            volume_kms_key,
-            TRAINING_JOB_VOLUME_KMS_KEY_ID_PATH,
-            sagemaker_session=self.sagemaker_session,
+        if instance_type is None or isinstance(instance_type, str):
+            instance_type_for_volume_kms = instance_type
+        elif isinstance(instance_type, ParameterString):
+            instance_type_for_volume_kms = instance_type.default_value
+        else:
+            raise ValueError(f"Bad value for instance type: '{instance_type}'")
+
+        # KMS can only be attached to supported instances
+        use_volume_kms_config = (
+            (instance_type_for_volume_kms and instance_supports_kms(instance_type_for_volume_kms))
+            or instance_groups is not None
+            and any(
+                [
+                    instance_supports_kms(instance_group.instance_type)
+                    for instance_group in instance_groups
+                ]
+            )
+        )
+
+        self.volume_kms_key = (
+            resolve_value_from_config(
+                volume_kms_key,
+                TRAINING_JOB_VOLUME_KMS_KEY_ID_PATH,
+                sagemaker_session=self.sagemaker_session,
+            )
+            if use_volume_kms_config
+            else volume_kms_key
         )
 
         # VPC configurations

src/sagemaker/instance_types.py

@@ -119,22 +119,3 @@ def retrieve(
         tolerate_vulnerable_model,
         tolerate_deprecated_model,
     )
-
-
-def volume_size_supported(instance_type: str) -> bool:
-    """Returns True if SageMaker allows volume_size to be used for the instance type.
-
-    Raises:
-        ValueError: If the instance type is improperly formatted.
-    """
-    try:
-        parts: List[str] = instance_type.split(".")
-        if len(parts) != 3 or parts[0] != "ml":
-            raise ValueError("Instance type must have 2 periods and start with 'ml'.")
-
-        # Any instance type with a "d" in the instance family (i.e. c5d, p4d, etc) + g5
-        # does not support attaching an EBS volume.
-        family = parts[1]
-        return "d" not in family and not family.startswith("g5")
-    except Exception as e:
-        raise ValueError(f"Failed to parse instance type '{instance_type}': {str(e)}")

src/sagemaker/jumpstart/artifacts/kwargs.py

@@ -14,7 +14,7 @@
 from __future__ import absolute_import
 from copy import deepcopy
 from typing import Optional
-from sagemaker.instance_types import volume_size_supported
+from sagemaker.utils import volume_size_supported
 from sagemaker.jumpstart.constants import (
     JUMPSTART_DEFAULT_REGION_NAME,
 )

src/sagemaker/jumpstart/estimator.py

@@ -36,7 +36,7 @@
     is_valid_model_id,
     resolve_model_intelligent_default_field,
 )
-from sagemaker.jumpstart.utils import stringify_object
+from sagemaker.utils import stringify_object
 from sagemaker.model_monitor.data_capture_config import DataCaptureConfig
 from sagemaker.predictor import PredictorBase
 

src/sagemaker/jumpstart/factory/estimator.py

@@ -425,7 +425,7 @@ def _add_image_uri_to_kwargs(kwargs: JumpStartEstimatorInitKwargs) -> JumpStartE
     """Sets image uri in kwargs based on default or override, returns full kwargs."""
 
     kwargs.image_uri = kwargs.image_uri or image_uris.retrieve(
-        region=None,
+        region=kwargs.region,
         framework=None,
         image_scope=JumpStartScriptScope.TRAINING,
         model_id=kwargs.model_id,

src/sagemaker/jumpstart/model.py

@@ -28,7 +28,7 @@
     get_init_kwargs,
 )
 from sagemaker.jumpstart.utils import is_valid_model_id
-from sagemaker.jumpstart.utils import stringify_object
+from sagemaker.utils import stringify_object
 from sagemaker.model import Model
 from sagemaker.model_monitor.data_capture_config import DataCaptureConfig
 from sagemaker.predictor import PredictorBase

src/sagemaker/jumpstart/utils.py

@@ -530,12 +530,6 @@ def resolve_estimator_intelligent_default_field(
     return field_val
 
 
-def stringify_object(obj: Any) -> str:
-    """Returns string representation of object, returning only non-None fields."""
-    non_none_atts = {key: value for key, value in obj.__dict__.items() if value is not None}
-    return f"{type(obj).__name__}: {str(non_none_atts)}"
-
-
 def is_valid_model_id(
     model_id: Optional[str],
     region: Optional[str] = None,

src/sagemaker/pipeline.py

@@ -34,6 +34,7 @@
 from sagemaker.transformer import Transformer
 from sagemaker.workflow.entities import PipelineVariable
 from sagemaker.workflow.pipeline_context import runnable_by_pipeline
+from sagemaker.utils import instance_supports_kms
 
 
 class PipelineModel(object):
@@ -235,8 +236,12 @@ def deploy(
             container_startup_health_check_timeout=container_startup_health_check_timeout,
         )
         self.endpoint_name = endpoint_name or self.name
-        kms_key = resolve_value_from_config(
-            kms_key, ENDPOINT_CONFIG_KMS_KEY_ID_PATH, sagemaker_session=self.sagemaker_session
+        kms_key = (
+            resolve_value_from_config(
+                kms_key, ENDPOINT_CONFIG_KMS_KEY_ID_PATH, sagemaker_session=self.sagemaker_session
+            )
+            if instance_supports_kms(instance_type)
+            else kms_key
         )
 
         data_capture_config_dict = None

src/sagemaker/session.py

@@ -29,6 +29,7 @@
 import botocore.config
 from botocore.exceptions import ClientError
 import six
+from sagemaker.utils import instance_supports_kms
 
 import sagemaker.logs
 from sagemaker import vpc_utils, s3_utils
@@ -811,9 +812,17 @@ def train(  # noqa: C901
         inferred_output_config = update_nested_dictionary_with_values_from_config(
             output_config, TRAINING_JOB_OUTPUT_DATA_CONFIG_PATH, sagemaker_session=self
         )
+        customer_supplied_kms_key = "VolumeKmsKeyId" in resource_config
         inferred_resource_config = update_nested_dictionary_with_values_from_config(
             resource_config, TRAINING_JOB_RESOURCE_CONFIG_PATH, sagemaker_session=self
         )
+        if (
+            not customer_supplied_kms_key
+            and "InstanceType" in inferred_resource_config
+            and not instance_supports_kms(inferred_resource_config["InstanceType"])
+            and "VolumeKmsKeyId" in inferred_resource_config
+        ):
+            del inferred_resource_config["VolumeKmsKeyId"]
         train_request = self._get_train_request(
             input_mode=input_mode,
             input_config=input_config,
@@ -3750,8 +3759,12 @@ def create_endpoint_config(
         )
         if tags is not None:
             request["Tags"] = tags
-        kms_key = resolve_value_from_config(
-            kms_key, ENDPOINT_CONFIG_KMS_KEY_ID_PATH, sagemaker_session=self
+        kms_key = (
+            resolve_value_from_config(
+                kms_key, ENDPOINT_CONFIG_KMS_KEY_ID_PATH, sagemaker_session=self
+            )
+            if instance_supports_kms(instance_type)
+            else kms_key
         )
         if kms_key is not None:
             request["KmsKeyId"] = kms_key
@@ -3844,7 +3857,16 @@ def create_endpoint_config_from_existing(
 
         if new_kms_key is not None or existing_endpoint_config_desc.get("KmsKeyId") is not None:
             request["KmsKeyId"] = new_kms_key or existing_endpoint_config_desc.get("KmsKeyId")
-        if KMS_KEY_ID not in request:
+
+        supports_kms = any(
+            [
+                instance_supports_kms(production_variant["InstanceType"])
+                for production_variant in production_variants
+                if "InstanceType" in production_variant
+            ]
+        )
+
+        if KMS_KEY_ID not in request and supports_kms:
             kms_key_from_config = resolve_value_from_config(
                 config_path=ENDPOINT_CONFIG_KMS_KEY_ID_PATH, sagemaker_session=self
             )
@@ -4465,15 +4487,28 @@ def endpoint_from_production_variants(
         Returns:
             str: The name of the created ``Endpoint``.
         """
+
+        supports_kms = any(
+            [
+                instance_supports_kms(production_variant["InstanceType"])
+                for production_variant in production_variants
+                if "InstanceType" in production_variant
+            ]
+        )
+
         update_list_of_dicts_with_values_from_config(
             production_variants,
             ENDPOINT_CONFIG_PRODUCTION_VARIANTS_PATH,
             required_key_paths=["CoreDumpConfig.DestinationS3Uri"],
             sagemaker_session=self,
         )
         config_options = {"EndpointConfigName": name, "ProductionVariants": production_variants}
-        kms_key = resolve_value_from_config(
-            kms_key, ENDPOINT_CONFIG_KMS_KEY_ID_PATH, sagemaker_session=self
+        kms_key = (
+            resolve_value_from_config(
+                kms_key, ENDPOINT_CONFIG_KMS_KEY_ID_PATH, sagemaker_session=self
+            )
+            if supports_kms
+            else kms_key
         )
         tags = _append_project_tags(tags)
         tags = self._append_sagemaker_config_tags(