Merge branch 'master' into transform_ami_version_v2

Author: pintaoz-aws

sagemaker-core/src/sagemaker/core/helper/pipeline_variable.py

@@ -80,3 +80,7 @@ def __get_pydantic_core_schema__(cls, source_type, handler):
 
 # This is a type that could be either string or pipeline variable
 StrPipeVar = Union[str, PipelineVariable]
+# This is a type that could be either integer or pipeline variable
+IntPipeVar = Union[int, PipelineVariable]
+# This is a type that could be either boolean or pipeline variable
+BoolPipeVar = Union[bool, PipelineVariable]

sagemaker-core/src/sagemaker/core/image_uri_config/djl-neuronx.json

@@ -18,6 +18,7 @@
                 "cn-northwest-1": "727897471807",
                 "eu-central-1": "763104351884",
                 "eu-central-2": "380420809688",
+                "eu-isoe-west-1": "371248457586",
                 "eu-west-1": "763104351884",
                 "eu-west-3": "763104351884",
                 "mx-central-1": "637423239942",
@@ -26,6 +27,10 @@
                 "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-isof-east-1": "303241398832",
+                "us-isof-south-1": "454834333376",
                 "us-west-2": "763104351884",
                 "ca-west-1": "204538143572",
                 "eu-north-1": "763104351884",

sagemaker-core/src/sagemaker/core/jumpstart/region_config.json

@@ -119,6 +119,10 @@
         "gated_content_bucket": "jumpstart-private-cache-prod-eu-west-3",
         "neo_content_bucket": "sagemaker-sd-models-prod-eu-west-3"
     },
+    "eu-isoe-west-1": {
+        "content_bucket": "jumpstart-cache-prod-eu-isoe-west-1",
+        "gated_content_bucket": "jumpstart-private-cache-prod-eu-isoe-west-1"
+    },
     "il-central-1": {
         "content_bucket": "jumpstart-cache-prod-il-central-1",
         "gated_content_bucket": "jumpstart-private-cache-prod-il-central-1"
@@ -158,6 +162,22 @@
         "content_bucket": "jumpstart-cache-prod-us-gov-west-1",
         "gated_content_bucket": "jumpstart-private-cache-prod-us-gov-west-1"
     },
+    "us-iso-east-1": {
+        "content_bucket": "jumpstart-cache-prod-us-iso-east-1",
+        "gated_content_bucket": "jumpstart-private-cache-prod-us-iso-east-1"
+    },
+    "us-isob-east-1": {
+        "content_bucket": "jumpstart-cache-prod-us-isob-east-1",
+        "gated_content_bucket": "jumpstart-private-cache-prod-us-isob-east-1"
+    },
+    "us-isof-east-1": {
+        "content_bucket": "jumpstart-cache-prod-us-isof-east-1",
+        "gated_content_bucket": "jumpstart-private-cache-prod-us-isof-east-1"
+    },
+    "us-isof-south-1": {
+        "content_bucket": "jumpstart-cache-prod-us-isof-south-1",
+        "gated_content_bucket": "jumpstart-private-cache-prod-us-isof-south-1"
+    },
     "us-west-1": {
         "content_bucket": "jumpstart-cache-prod-us-west-1",
         "gated_content_bucket": "jumpstart-private-cache-prod-us-west-1",

sagemaker-core/src/sagemaker/core/shapes/shapes.py

@@ -16,7 +16,7 @@
 from pydantic import BaseModel, ConfigDict, Field
 from typing import List, Dict, Optional, Any, Union
 from sagemaker.core.utils.utils import Unassigned
-from sagemaker.core.helper.pipeline_variable import StrPipeVar
+from sagemaker.core.helper.pipeline_variable import StrPipeVar, IntPipeVar, BoolPipeVar
 
 # Suppress Pydantic warnings about field names shadowing parent attributes
 warnings.filterwarnings("ignore", message=".*shadows an attribute.*")
@@ -1324,10 +1324,10 @@ class ResourceConfig(Base):
     """
 
     instance_type: Optional[StrPipeVar] = Unassigned()
-    instance_count: Optional[int] = Unassigned()
-    volume_size_in_gb: Optional[int] = Unassigned()
+    instance_count: Optional[IntPipeVar] = Unassigned()
+    volume_size_in_gb: Optional[IntPipeVar] = Unassigned()
     volume_kms_key_id: Optional[StrPipeVar] = Unassigned()
-    keep_alive_period_in_seconds: Optional[int] = Unassigned()
+    keep_alive_period_in_seconds: Optional[IntPipeVar] = Unassigned()
     capacity_reservation_ids: Optional[List[StrPipeVar]] = Unassigned()
     instance_groups: Optional[List[InstanceGroup]] = Unassigned()
     capacity_schedules_config: Optional[CapacitySchedulesConfig] = Unassigned()

sagemaker-core/src/sagemaker/core/training/configs.py

@@ -25,7 +25,7 @@
 from pydantic import BaseModel, model_validator, ConfigDict
 
 import sagemaker.core.shapes as shapes
-from sagemaker.core.helper.pipeline_variable import StrPipeVar
+from sagemaker.core.helper.pipeline_variable import StrPipeVar, IntPipeVar, BoolPipeVar
 
 # TODO: Can we add custom logic to some of these to set better defaults?
 from sagemaker.core.shapes import (
@@ -158,23 +158,23 @@ class Compute(shapes.ResourceConfig):
         instance_type (Optional[StrPipeVar]):
             The ML compute instance type. For information about available instance types,
             see https://aws.amazon.com/sagemaker/pricing/.
-        instance_count (Optional[int]): The number of ML compute instances to use. For distributed
+        instance_count (Optional[IntPipeVar]): The number of ML compute instances to use. For distributed
             training, provide a value greater than 1.
-        volume_size_in_gb (Optional[int]):
+        volume_size_in_gb (Optional[IntPipeVar]):
             The size of the ML storage volume that you want to provision.  ML storage volumes store
             model artifacts and incremental states. Training algorithms might also use the ML
             storage volume for scratch space. Default: 30
         volume_kms_key_id (Optional[StrPipeVar]):
             The Amazon Web Services KMS key that SageMaker uses to encrypt data on the storage
             volume attached to the ML compute instance(s) that run the training job.
-        keep_alive_period_in_seconds (Optional[int]):
+        keep_alive_period_in_seconds (Optional[IntPipeVar]):
             The duration of time in seconds to retain configured resources in a warm pool for
             subsequent training jobs.
         instance_groups (Optional[List[InstanceGroup]]):
             A list of instance groups for heterogeneous clusters to be used in the training job.
         training_plan_arn (Optional[StrPipeVar]):
             The Amazon Resource Name (ARN) of the training plan to use for this resource configuration.
-        enable_managed_spot_training (Optional[bool]):
+        enable_managed_spot_training (Optional[BoolPipeVar]):
             To train models using managed spot training, choose True. Managed spot training
             provides a fully managed and scalable infrastructure for training machine learning
             models. this option is useful when training jobs can be interrupted and when there