Fix typo and merge with master branch

CHANGELOG.md

@@ -1,5 +1,28 @@
 # Changelog
 
+## v2.218.0 (2024-05-01)
+
+### Features
+
+ * set default allow_pickle param to False
+
+### Bug Fixes and Other Changes
+
+ * properly close files in lineage queries and tests
+
+## v2.217.0 (2024-04-24)
+
+### Features
+
+ * support session tag chaining for training job
+
+### Bug Fixes and Other Changes
+
+ * Add Triton v24.03 URI
+ * mainline alt config parsing
+ * Fix tox installs
+ * Add PT 2.2 Graviton Inference DLC
+
 ## v2.216.1 (2024-04-22)
 
 ### Bug Fixes and Other Changes

VERSION

@@ -1 +1 @@
-2.216.2.dev0
+2.218.1.dev0

src/sagemaker/base_deserializers.py

@@ -196,14 +196,14 @@ class NumpyDeserializer(SimpleBaseDeserializer):
     single array.
     """
 
-    def __init__(self, dtype=None, accept="application/x-npy", allow_pickle=True):
+    def __init__(self, dtype=None, accept="application/x-npy", allow_pickle=False):
         """Initialize a ``NumpyDeserializer`` instance.
 
         Args:
             dtype (str): The dtype of the data (default: None).
             accept (union[str, tuple[str]]): The MIME type (or tuple of allowable MIME types) that
                 is expected from the inference endpoint (default: "application/x-npy").
-            allow_pickle (bool): Allow loading pickled object arrays (default: True).
+            allow_pickle (bool): Allow loading pickled object arrays (default: False).
         """
         super(NumpyDeserializer, self).__init__(accept=accept)
         self.dtype = dtype
@@ -227,10 +227,21 @@ def deserialize(self, stream, content_type):
             if content_type == "application/json":
                 return np.array(json.load(codecs.getreader("utf-8")(stream)), dtype=self.dtype)
             if content_type == "application/x-npy":
-                return np.load(io.BytesIO(stream.read()), allow_pickle=self.allow_pickle)
+                try:
+                    return np.load(io.BytesIO(stream.read()), allow_pickle=self.allow_pickle)
+                except ValueError as ve:
+                    raise ValueError(
+                        "Please set the param allow_pickle=True \
+                        to deserialize pickle objects in NumpyDeserializer"
+                    ).with_traceback(ve.__traceback__)
             if content_type == "application/x-npz":
                 try:
                     return np.load(io.BytesIO(stream.read()), allow_pickle=self.allow_pickle)
+                except ValueError as ve:
+                    raise ValueError(
+                        "Please set the param allow_pickle=True \
+                        to deserialize pickle objectsin NumpyDeserializer"
+                    ).with_traceback(ve.__traceback__)
                 finally:
                     stream.close()
         finally:

src/sagemaker/estimator.py

@@ -181,6 +181,7 @@ def __init__(
         container_arguments: Optional[List[str]] = None,
         disable_output_compression: bool = False,
         enable_remote_debug: Optional[Union[bool, PipelineVariable]] = None,
+        enable_session_tag_chaining: Optional[Union[bool, PipelineVariable]] = None,
         **kwargs,
     ):
         """Initialize an ``EstimatorBase`` instance.
@@ -544,7 +545,9 @@ def __init__(
             enable_infra_check (bool or PipelineVariable): Optional.
                 Specifies whether it is running Sagemaker built-in infra check jobs.
             enable_remote_debug (bool or PipelineVariable): Optional.
-                Specifies whether RemoteDebug is enabled for the training job
+                Specifies whether RemoteDebug is enabled for the training job.
+            enable_session_tag_chaining (bool or PipelineVariable): Optional.
+                Specifies whether SessionTagChaining is enabled for the training job.
         """
         instance_count = renamed_kwargs(
             "train_instance_count", "instance_count", instance_count, kwargs
@@ -785,6 +788,8 @@ def __init__(
 
         self._enable_remote_debug = enable_remote_debug
 
+        self._enable_session_tag_chaining = enable_session_tag_chaining
+
     @abstractmethod
     def training_image_uri(self):
         """Return the Docker image to use for training.
@@ -2318,6 +2323,14 @@ def get_remote_debug_config(self):
             else {"EnableRemoteDebug": self._enable_remote_debug}
         )
 
+    def get_session_chaining_config(self):
+        """dict: Return the configuration of SessionChaining"""
+        return (
+            None
+            if self._enable_session_tag_chaining is None
+            else {"EnableSessionTagChaining": self._enable_session_tag_chaining}
+        )
+
     def enable_remote_debug(self):
         """Enable remote debug for a training job."""
         self._update_remote_debug(True)
@@ -2574,6 +2587,9 @@ def _get_train_args(cls, estimator, inputs, experiment_config):
         if estimator.get_remote_debug_config() is not None:
             train_args["remote_debug_config"] = estimator.get_remote_debug_config()
 
+        if estimator.get_session_chaining_config() is not None:
+            train_args["session_chaining_config"] = estimator.get_session_chaining_config()
+
         return train_args
 
     @classmethod
@@ -2766,6 +2782,7 @@ def __init__(
         disable_output_compression: bool = False,
         enable_infra_check: Optional[Union[bool, PipelineVariable]] = None,
         enable_remote_debug: Optional[Union[bool, PipelineVariable]] = None,
+        enable_session_tag_chaining: Optional[Union[bool, PipelineVariable]] = None,
         **kwargs,
     ):
         """Initialize an ``Estimator`` instance.
@@ -3129,6 +3146,8 @@ def __init__(
                 Specifies whether it is running Sagemaker built-in infra check jobs.
             enable_remote_debug (bool or PipelineVariable): Optional.
                 Specifies whether RemoteDebug is enabled for the training job
+            enable_session_tag_chaining (bool or PipelineVariable): Optional.
+                 Specifies whether SessionTagChaining is enabled for the training job
         """
         self.image_uri = image_uri
         self._hyperparameters = hyperparameters.copy() if hyperparameters else {}
@@ -3181,6 +3200,7 @@ def __init__(
             container_arguments=container_arguments,
             disable_output_compression=disable_output_compression,
             enable_remote_debug=enable_remote_debug,
+            enable_session_tag_chaining=enable_session_tag_chaining,
             **kwargs,
         )
 

src/sagemaker/image_uri_config/sagemaker-tritonserver.json

@@ -7,7 +7,7 @@
 		"inference"
 	],
 	"versions": {
-		"23.12": {
+		"24.03": {
 			"registries": {
 				"af-south-1": "626614931356",
 				"il-central-1": "780543022126",
@@ -37,7 +37,7 @@
 				"ca-west-1": "204538143572"
 			},
 			"repository": "sagemaker-tritonserver",
-			"tag_prefix": "23.12-py3"
+			"tag_prefix": "24.03-py3"
 		},
 		"24.01": {
 			"registries": {
@@ -70,6 +70,38 @@
 			},
 			"repository": "sagemaker-tritonserver",
 			"tag_prefix": "24.01-py3"
+		},
+		"23.12": {
+			"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-southeast-1": "763104351884",
+				"ap-southeast-2": "763104351884",
+				"ap-southeast-3": "907027046896",
+				"ca-central-1": "763104351884",
+				"cn-north-1": "727897471807",
+				"cn-northwest-1": "727897471807",
+				"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",
+				"ca-west-1": "204538143572"
+			},
+			"repository": "sagemaker-tritonserver",
+			"tag_prefix": "23.12-py3"
 		}
 	}
 }

src/sagemaker/jumpstart/estimator.py

@@ -112,6 +112,7 @@ def __init__(
         disable_output_compression: Optional[bool] = None,
         enable_remote_debug: Optional[Union[bool, PipelineVariable]] = None,
         config_name: Optional[str] = None,
+        enable_session_tag_chaining: Optional[Union[bool, PipelineVariable]] = None,
     ):
         """Initializes a ``JumpStartEstimator``.
 
@@ -505,6 +506,8 @@ def __init__(
                 Specifies whether RemoteDebug is enabled for the training job
             config_name (Optional[str]):
                 Name of the training configuration to apply to the Estimator. (Default: None).
+            enable_session_tag_chaining (bool or PipelineVariable): Optional.
+                Specifies whether SessionTagChaining is enabled for the training job
 
         Raises:
             ValueError: If the model ID is not recognized by JumpStart.
@@ -584,6 +587,7 @@ def _validate_model_id_and_get_type_hook():
             enable_infra_check=enable_infra_check,
             enable_remote_debug=enable_remote_debug,
             config_name=config_name,
+            enable_session_tag_chaining=enable_session_tag_chaining,
         )
 
         self.model_id = estimator_init_kwargs.model_id

src/sagemaker/jumpstart/factory/estimator.py

@@ -131,6 +131,7 @@ def get_init_kwargs(
     enable_infra_check: Optional[Union[bool, PipelineVariable]] = None,
     enable_remote_debug: Optional[Union[bool, PipelineVariable]] = None,
     config_name: Optional[str] = None,
+    enable_session_tag_chaining: Optional[Union[bool, PipelineVariable]] = None,
 ) -> JumpStartEstimatorInitKwargs:
     """Returns kwargs required to instantiate `sagemaker.estimator.Estimator` object."""
 
@@ -190,6 +191,7 @@ def get_init_kwargs(
         enable_infra_check=enable_infra_check,
         enable_remote_debug=enable_remote_debug,
         config_name=config_name,
+        enable_session_tag_chaining=enable_session_tag_chaining,
     )
 
     estimator_init_kwargs = _add_model_version_to_kwargs(estimator_init_kwargs)

src/sagemaker/jumpstart/session_utils.py

@@ -219,7 +219,7 @@ def get_model_info_from_training_job(
         model_id,
         inferred_model_version,
         inference_config_name,
-        trainig_config_name,
+        training_config_name,
     ) = get_jumpstart_model_info_from_resource_arn(training_job_arn, sagemaker_session)
 
     model_version = inferred_model_version or None
@@ -231,4 +231,4 @@ def get_model_info_from_training_job(
             "for this training job."
         )
 
-    return model_id, model_version, inference_config_name, trainig_config_name
+    return model_id, model_version, inference_config_name, training_config_name

src/sagemaker/jumpstart/types.py

@@ -1078,7 +1078,7 @@ class JumpStartMetadataConfig(JumpStartDataHolderType):
         "resolved_metadata_config",
         "config_name",
         "default_inference_config",
-        "default_incremental_trainig_config",
+        "default_incremental_training_config",
         "supported_inference_configs",
         "supported_incremental_training_configs",
     ]
@@ -1114,7 +1114,7 @@ def __init__(
         self.resolved_metadata_config: Optional[Dict[str, Any]] = None
         self.config_name: Optional[str] = config_name
         self.default_inference_config: Optional[str] = config.get("default_inference_config")
-        self.default_incremental_trainig_config: Optional[str] = config.get(
+        self.default_incremental_training_config: Optional[str] = config.get(
             "default_incremental_training_config"
         )
         self.supported_inference_configs: Optional[List[str]] = config.get(
@@ -1775,6 +1775,7 @@ class JumpStartEstimatorInitKwargs(JumpStartKwargs):
         "enable_infra_check",
         "enable_remote_debug",
         "config_name",
+        "enable_session_tag_chaining",
     ]
 
     SERIALIZATION_EXCLUSION_SET = {
@@ -1844,6 +1845,7 @@ def __init__(
         enable_infra_check: Optional[Union[bool, PipelineVariable]] = None,
         enable_remote_debug: Optional[Union[bool, PipelineVariable]] = None,
         config_name: Optional[str] = None,
+        enable_session_tag_chaining: Optional[Union[bool, PipelineVariable]] = None,
     ) -> None:
         """Instantiates JumpStartEstimatorInitKwargs object."""
 
@@ -1904,6 +1906,7 @@ def __init__(
         self.enable_infra_check = enable_infra_check
         self.enable_remote_debug = enable_remote_debug
         self.config_name = config_name
+        self.enable_session_tag_chaining = enable_session_tag_chaining
 
 
 class JumpStartEstimatorFitKwargs(JumpStartKwargs):

src/sagemaker/lineage/query.py

@@ -335,8 +335,8 @@ def _get_legend_line(self, component_name):
 
     def _add_legend(self, path):
         """Embed legend to html file generated by pyvis."""
-        f = open(path, "r")
-        content = self.BeautifulSoup(f, "html.parser")
+        with open(path, "r") as f:
+            content = self.BeautifulSoup(f, "html.parser")
 
         legend = """
             <div style="display: inline-block; font-size: 1vw; font-family: verdana;

src/sagemaker/session.py

@@ -758,6 +758,7 @@ def train(  # noqa: C901
         environment: Optional[Dict[str, str]] = None,
         retry_strategy=None,
         remote_debug_config=None,
+        session_chaining_config=None,
     ):
         """Create an Amazon SageMaker training job.
 
@@ -877,6 +878,15 @@ def train(  # noqa: C901
                     remote_debug_config = {
                         "EnableRemoteDebug": True,
                     }
+            session_chaining_config(dict): Configuration for SessionChaining. (default: ``None``)
+                The dict can contain 'EnableSessionTagChaining'(bool).
+                For example,
+
+                .. code:: python
+
+                    session_chaining_config = {
+                        "EnableSessionTagChaining": True,
+                    }
             environment (dict[str, str]) : Environment variables to be set for
                 use during training job (default: ``None``)
             retry_strategy(dict): Defines RetryStrategy for InternalServerFailures.
@@ -970,6 +980,7 @@ def train(  # noqa: C901
             profiler_rule_configs=profiler_rule_configs,
             profiler_config=inferred_profiler_config,
             remote_debug_config=remote_debug_config,
+            session_chaining_config=session_chaining_config,
             environment=environment,
             retry_strategy=retry_strategy,
         )
@@ -1013,6 +1024,7 @@ def _get_train_request(  # noqa: C901
         profiler_rule_configs=None,
         profiler_config=None,
         remote_debug_config=None,
+        session_chaining_config=None,
         environment=None,
         retry_strategy=None,
     ):
@@ -1133,6 +1145,15 @@ def _get_train_request(  # noqa: C901
                     remote_debug_config = {
                         "EnableRemoteDebug": True,
                     }
+            session_chaining_config(dict): Configuration for SessionChaining. (default: ``None``)
+                The dict can contain 'EnableSessionTagChaining'(bool).
+                For example,
+
+                .. code:: python
+
+                    session_chaining_config = {
+                        "EnableSessionTagChaining": True,
+                    }
             environment (dict[str, str]) : Environment variables to be set for
                 use during training job (default: ``None``)
             retry_strategy(dict): Defines RetryStrategy for InternalServerFailures.
@@ -1239,6 +1260,9 @@ def _get_train_request(  # noqa: C901
         if remote_debug_config is not None:
             train_request["RemoteDebugConfig"] = remote_debug_config
 
+        if session_chaining_config is not None:
+            train_request["SessionChainingConfig"] = session_chaining_config
+
         if retry_strategy is not None:
             train_request["RetryStrategy"] = retry_strategy
 

tests/data/sip/training.py

@@ -73,7 +73,8 @@ def main():
     )
 
     model_dir = os.environ.get("SM_MODEL_DIR")
-    pkl.dump(bst, open(model_dir + "/model.bin", "wb"))
+    with open(model_dir + "/model.bin", "wb") as f:
+        pkl.dump(bst, f)
 
 
 if __name__ == "__main__":