From c4ac4803b26fd200597a2ec972d9682f697c8c25 Mon Sep 17 00:00:00 2001
From: Haotian An <33510317+Captainia@users.noreply.github.com>
Date: Tue, 23 Apr 2024 12:03:54 -0400
Subject: [PATCH 01/18] fix: mainline alt config parsing (#4602)

* fix: parsing

* fix: commit tests

* fix: types

* updated

* fix
---
 src/sagemaker/jumpstart/types.py             | 72 +++++++++++-----
 tests/unit/sagemaker/jumpstart/constants.py  | 32 ++++---
 tests/unit/sagemaker/jumpstart/test_types.py | 90 ++++++++++++++++----
 tests/unit/sagemaker/jumpstart/test_utils.py | 60 ++++++-------
 4 files changed, 174 insertions(+), 80 deletions(-)

diff --git a/src/sagemaker/jumpstart/types.py b/src/sagemaker/jumpstart/types.py
index 8e53bf6f83..05c6a00961 100644
--- a/src/sagemaker/jumpstart/types.py
+++ b/src/sagemaker/jumpstart/types.py
@@ -744,12 +744,12 @@ def _get_regional_property(
 
 
 class JumpStartBenchmarkStat(JumpStartDataHolderType):
-    """Data class JumpStart benchmark stats."""
+    """Data class JumpStart benchmark stat."""
 
     __slots__ = ["name", "value", "unit"]
 
     def __init__(self, spec: Dict[str, Any]):
-        """Initializes a JumpStartBenchmarkStat object
+        """Initializes a JumpStartBenchmarkStat object.
 
         Args:
             spec (Dict[str, Any]): Dictionary representation of benchmark stat.
@@ -858,7 +858,7 @@ class JumpStartMetadataBaseFields(JumpStartDataHolderType):
         "model_subscription_link",
     ]
 
-    def __init__(self, fields: Optional[Dict[str, Any]]):
+    def __init__(self, fields: Dict[str, Any]):
         """Initializes a JumpStartMetadataFields object.
 
         Args:
@@ -877,7 +877,7 @@ def from_json(self, json_obj: Dict[str, Any]) -> None:
         self.version: str = json_obj.get("version")
         self.min_sdk_version: str = json_obj.get("min_sdk_version")
         self.incremental_training_supported: bool = bool(
-            json_obj.get("incremental_training_supported")
+            json_obj.get("incremental_training_supported", False)
         )
         self.hosting_ecr_specs: Optional[JumpStartECRSpecs] = (
             JumpStartECRSpecs(json_obj["hosting_ecr_specs"])
@@ -1038,7 +1038,7 @@ class JumpStartConfigComponent(JumpStartMetadataBaseFields):
 
     __slots__ = slots + JumpStartMetadataBaseFields.__slots__
 
-    def __init__(  # pylint: disable=super-init-not-called
+    def __init__(
         self,
         component_name: str,
         component: Optional[Dict[str, Any]],
@@ -1049,7 +1049,10 @@ def __init__(  # pylint: disable=super-init-not-called
             component_name (str): Name of the component.
             component (Dict[str, Any]):
                 Dictionary representation of the config component.
+        Raises:
+            ValueError: If the component field is invalid.
         """
+        super().__init__(component)
         self.component_name = component_name
         self.from_json(component)
 
@@ -1080,7 +1083,7 @@ def __init__(
         self,
         base_fields: Dict[str, Any],
         config_components: Dict[str, JumpStartConfigComponent],
-        benchmark_metrics: Dict[str, JumpStartBenchmarkStat],
+        benchmark_metrics: Dict[str, List[JumpStartBenchmarkStat]],
     ):
         """Initializes a JumpStartMetadataConfig object from its json representation.
 
@@ -1089,12 +1092,12 @@ def __init__(
                 The default base fields that are used to construct the final resolved config.
             config_components (Dict[str, JumpStartConfigComponent]):
                 The list of components that are used to construct the resolved config.
-            benchmark_metrics (Dict[str, JumpStartBenchmarkStat]):
+            benchmark_metrics (Dict[str, List[JumpStartBenchmarkStat]]):
                 The dictionary of benchmark metrics with name being the key.
         """
         self.base_fields = base_fields
         self.config_components: Dict[str, JumpStartConfigComponent] = config_components
-        self.benchmark_metrics: Dict[str, JumpStartBenchmarkStat] = benchmark_metrics
+        self.benchmark_metrics: Dict[str, List[JumpStartBenchmarkStat]] = benchmark_metrics
         self.resolved_metadata_config: Optional[Dict[str, Any]] = None
 
     def to_json(self) -> Dict[str, Any]:
@@ -1104,7 +1107,7 @@ def to_json(self) -> Dict[str, Any]:
 
     @property
     def resolved_config(self) -> Dict[str, Any]:
-        """Returns the final config that is resolved from the list of components.
+        """Returns the final config that is resolved from the components map.
 
         Construct the final config by applying the list of configs from list index,
         and apply to the base default fields in the current model specs.
@@ -1139,7 +1142,7 @@ def __init__(
 
         Args:
             configs (Dict[str, JumpStartMetadataConfig]):
-                List of configs that the current model has.
+                The map of JumpStartMetadataConfig object, with config name being the key.
             config_rankings (JumpStartConfigRanking):
                 Config ranking class represents the ranking of the configs in the model.
             scope (JumpStartScriptScope):
@@ -1158,19 +1161,30 @@ def get_top_config_from_ranking(
         self,
         ranking_name: str = JumpStartConfigRankingName.DEFAULT,
         instance_type: Optional[str] = None,
-    ) -> JumpStartMetadataConfig:
-        """Gets the best the config based on config ranking."""
+    ) -> Optional[JumpStartMetadataConfig]:
+        """Gets the best the config based on config ranking.
+
+        Args:
+            ranking_name (str):
+                The ranking name that config priority is based on.
+            instance_type (Optional[str]):
+                The instance type which the config selection is based on.
+
+        Raises:
+            ValueError: If the config exists but missing config ranking.
+            NotImplementedError: If the scope is unrecognized.
+        """
         if self.configs and (
             not self.config_rankings or not self.config_rankings.get(ranking_name)
         ):
-            raise ValueError("Config exists but missing config ranking.")
+            raise ValueError(f"Config exists but missing config ranking {ranking_name}.")
 
         if self.scope == JumpStartScriptScope.INFERENCE:
             instance_type_attribute = "supported_inference_instance_types"
         elif self.scope == JumpStartScriptScope.TRAINING:
             instance_type_attribute = "supported_training_instance_types"
         else:
-            raise ValueError(f"Unknown script scope {self.scope}")
+            raise NotImplementedError(f"Unknown script scope {self.scope}")
 
         rankings = self.config_rankings.get(ranking_name)
         for config_name in rankings.rankings:
@@ -1198,12 +1212,13 @@ class JumpStartModelSpecs(JumpStartMetadataBaseFields):
 
     __slots__ = JumpStartMetadataBaseFields.__slots__ + slots
 
-    def __init__(self, spec: Dict[str, Any]):  # pylint: disable=super-init-not-called
+    def __init__(self, spec: Dict[str, Any]):
         """Initializes a JumpStartModelSpecs object from its json representation.
 
         Args:
             spec (Dict[str, Any]): Dictionary representation of spec.
         """
+        super().__init__(spec)
         self.from_json(spec)
         if self.inference_configs and self.inference_configs.get_top_config_from_ranking():
             super().from_json(self.inference_configs.get_top_config_from_ranking().resolved_config)
@@ -1245,8 +1260,8 @@ def from_json(self, json_obj: Dict[str, Any]) -> None:
                     ),
                     (
                         {
-                            stat_name: JumpStartBenchmarkStat(stat)
-                            for stat_name, stat in config.get("benchmark_metrics").items()
+                            stat_name: [JumpStartBenchmarkStat(stat) for stat in stats]
+                            for stat_name, stats in config.get("benchmark_metrics").items()
                         }
                         if config and config.get("benchmark_metrics")
                         else None
@@ -1297,8 +1312,8 @@ def from_json(self, json_obj: Dict[str, Any]) -> None:
                         ),
                         (
                             {
-                                stat_name: JumpStartBenchmarkStat(stat)
-                                for stat_name, stat in config.get("benchmark_metrics").items()
+                                stat_name: [JumpStartBenchmarkStat(stat) for stat in stats]
+                                for stat_name, stats in config.get("benchmark_metrics").items()
                             }
                             if config and config.get("benchmark_metrics")
                             else None
@@ -1330,13 +1345,26 @@ def set_config(
             config_name (str): Name of the config.
             scope (JumpStartScriptScope, optional):
                 Scope of the config. Defaults to JumpStartScriptScope.INFERENCE.
+
+        Raises:
+            ValueError: If the scope is not supported, or cannot find config name.
         """
         if scope == JumpStartScriptScope.INFERENCE:
-            super().from_json(self.inference_configs.configs[config_name].resolved_config)
+            metadata_configs = self.inference_configs
         elif scope == JumpStartScriptScope.TRAINING and self.training_supported:
-            super().from_json(self.training_configs.configs[config_name].resolved_config)
+            metadata_configs = self.training_configs
         else:
-            raise ValueError(f"Unknown Jumpstart Script scope {scope}.")
+            raise ValueError(f"Unknown Jumpstart script scope {scope}.")
+
+        config_object = metadata_configs.configs.get(config_name)
+        if not config_object:
+            error_msg = f"Cannot find Jumpstart config name {config_name}. "
+            config_names = list(metadata_configs.configs.keys())
+            if config_names:
+                error_msg += f"List of config names that is supported by the model: {config_names}"
+            raise ValueError(error_msg)
+
+        super().from_json(config_object.resolved_config)
 
     def supports_prepacked_inference(self) -> bool:
         """Returns True if the model has a prepacked inference artifact."""
diff --git a/tests/unit/sagemaker/jumpstart/constants.py b/tests/unit/sagemaker/jumpstart/constants.py
index 2b6856b1f3..f165a513a9 100644
--- a/tests/unit/sagemaker/jumpstart/constants.py
+++ b/tests/unit/sagemaker/jumpstart/constants.py
@@ -6270,6 +6270,10 @@
             "framework_version": "1.5.0",
             "py_version": "py3",
         },
+        "default_inference_instance_type": "ml.p2.xlarge",
+        "supported_inference_instance_type": ["ml.p2.xlarge", "ml.p3.xlarge"],
+        "default_training_instance_type": "ml.p2.xlarge",
+        "supported_training_instance_type": ["ml.p2.xlarge", "ml.p3.xlarge"],
         "hosting_artifact_key": "pytorch-infer/infer-pytorch-eqa-bert-base-cased.tar.gz",
         "hosting_script_key": "source-directory-tarballs/pytorch/inference/eqa/v1.0.0/sourcedir.tar.gz",
         "inference_vulnerable": False,
@@ -7658,25 +7662,25 @@
     "inference_configs": {
         "neuron-inference": {
             "benchmark_metrics": {
-                "ml.inf2.2xlarge": {"name": "Latency", "value": "100", "unit": "Tokens/S"}
+                "ml.inf2.2xlarge": [{"name": "Latency", "value": "100", "unit": "Tokens/S"}]
             },
-            "component_names": ["neuron-base"],
+            "component_names": ["neuron-inference"],
         },
         "neuron-inference-budget": {
             "benchmark_metrics": {
-                "ml.inf2.2xlarge": {"name": "Latency", "value": "100", "unit": "Tokens/S"}
+                "ml.inf2.2xlarge": [{"name": "Latency", "value": "100", "unit": "Tokens/S"}]
             },
             "component_names": ["neuron-base"],
         },
         "gpu-inference-budget": {
             "benchmark_metrics": {
-                "ml.p3.2xlarge": {"name": "Latency", "value": "100", "unit": "Tokens/S"}
+                "ml.p3.2xlarge": [{"name": "Latency", "value": "100", "unit": "Tokens/S"}]
             },
             "component_names": ["gpu-inference-budget"],
         },
         "gpu-inference": {
             "benchmark_metrics": {
-                "ml.p3.2xlarge": {"name": "Latency", "value": "100", "unit": "Tokens/S"}
+                "ml.p3.2xlarge": [{"name": "Latency", "value": "100", "unit": "Tokens/S"}]
             },
             "component_names": ["gpu-inference"],
         },
@@ -7686,7 +7690,13 @@
             "supported_inference_instance_types": ["ml.inf2.xlarge", "ml.inf2.2xlarge"]
         },
         "neuron-inference": {
+            "default_inference_instance_type": "ml.inf2.xlarge",
             "supported_inference_instance_types": ["ml.inf2.xlarge", "ml.inf2.2xlarge"],
+            "hosting_ecr_specs": {
+                "framework": "huggingface-llm-neuronx",
+                "framework_version": "0.0.17",
+                "py_version": "py310",
+            },
             "hosting_artifact_key": "artifacts/meta-textgeneration-llama-2-7b/neuron-inference/model/",
             "hosting_instance_type_variants": {
                 "regional_aliases": {
@@ -7738,27 +7748,27 @@
     "training_configs": {
         "neuron-training": {
             "benchmark_metrics": {
-                "ml.tr1n1.2xlarge": {"name": "Latency", "value": "100", "unit": "Tokens/S"},
-                "ml.tr1n1.4xlarge": {"name": "Latency", "value": "50", "unit": "Tokens/S"},
+                "ml.tr1n1.2xlarge": [{"name": "Latency", "value": "100", "unit": "Tokens/S"}],
+                "ml.tr1n1.4xlarge": [{"name": "Latency", "value": "50", "unit": "Tokens/S"}],
             },
             "component_names": ["neuron-training"],
         },
         "neuron-training-budget": {
             "benchmark_metrics": {
-                "ml.tr1n1.2xlarge": {"name": "Latency", "value": "100", "unit": "Tokens/S"},
-                "ml.tr1n1.4xlarge": {"name": "Latency", "value": "50", "unit": "Tokens/S"},
+                "ml.tr1n1.2xlarge": [{"name": "Latency", "value": "100", "unit": "Tokens/S"}],
+                "ml.tr1n1.4xlarge": [{"name": "Latency", "value": "50", "unit": "Tokens/S"}],
             },
             "component_names": ["neuron-training-budget"],
         },
         "gpu-training": {
             "benchmark_metrics": {
-                "ml.p3.2xlarge": {"name": "Latency", "value": "200", "unit": "Tokens/S"},
+                "ml.p3.2xlarge": [{"name": "Latency", "value": "200", "unit": "Tokens/S"}],
             },
             "component_names": ["gpu-training"],
         },
         "gpu-training-budget": {
             "benchmark_metrics": {
-                "ml.p3.2xlarge": {"name": "Latency", "value": "100", "unit": "Tokens/S"}
+                "ml.p3.2xlarge": [{"name": "Latency", "value": "100", "unit": "Tokens/S"}]
             },
             "component_names": ["gpu-training-budget"],
         },
diff --git a/tests/unit/sagemaker/jumpstart/test_types.py b/tests/unit/sagemaker/jumpstart/test_types.py
index 3048bbc320..5ca01c3c52 100644
--- a/tests/unit/sagemaker/jumpstart/test_types.py
+++ b/tests/unit/sagemaker/jumpstart/test_types.py
@@ -12,6 +12,7 @@
 # language governing permissions and limitations under the License.
 from __future__ import absolute_import
 import copy
+import pytest
 from sagemaker.jumpstart.enums import JumpStartScriptScope
 from sagemaker.jumpstart.types import (
     JumpStartBenchmarkStat,
@@ -934,9 +935,9 @@ def test_inference_configs_parsing():
     assert specs1.incremental_training_supported
     assert specs1.hosting_ecr_specs == JumpStartECRSpecs(
         {
-            "framework": "pytorch",
-            "framework_version": "1.5.0",
-            "py_version": "py3",
+            "framework": "huggingface-llm-neuronx",
+            "framework_version": "0.0.17",
+            "py_version": "py310",
         }
     )
     assert specs1.training_ecr_specs == JumpStartECRSpecs(
@@ -946,7 +947,10 @@ def test_inference_configs_parsing():
             "py_version": "py3",
         }
     )
-    assert specs1.hosting_artifact_key == "pytorch-infer/infer-pytorch-ic-mobilenet-v2.tar.gz"
+    assert (
+        specs1.hosting_artifact_key
+        == "artifacts/meta-textgeneration-llama-2-7b/neuron-inference/model/"
+    )
     assert specs1.training_artifact_key == "pytorch-training/train-pytorch-ic-mobilenet-v2.tar.gz"
     assert (
         specs1.hosting_script_key
@@ -1019,16 +1023,58 @@ def test_inference_configs_parsing():
     config = specs1.inference_configs.get_top_config_from_ranking()
 
     assert config.benchmark_metrics == {
-        "ml.inf2.2xlarge": JumpStartBenchmarkStat(
-            {"name": "Latency", "value": "100", "unit": "Tokens/S"}
-        )
+        "ml.inf2.2xlarge": [
+            JumpStartBenchmarkStat({"name": "Latency", "value": "100", "unit": "Tokens/S"})
+        ]
     }
     assert len(config.config_components) == 1
-    assert config.config_components["neuron-base"] == JumpStartConfigComponent(
-        "neuron-base",
-        {"supported_inference_instance_types": ["ml.inf2.xlarge", "ml.inf2.2xlarge"]},
+    assert config.config_components["neuron-inference"] == JumpStartConfigComponent(
+        "neuron-inference",
+        {
+            "default_inference_instance_type": "ml.inf2.xlarge",
+            "supported_inference_instance_types": ["ml.inf2.xlarge", "ml.inf2.2xlarge"],
+            "hosting_ecr_specs": {
+                "framework": "huggingface-llm-neuronx",
+                "framework_version": "0.0.17",
+                "py_version": "py310",
+            },
+            "hosting_artifact_key": "artifacts/meta-textgeneration-llama-2-7b/neuron-inference/model/",
+            "hosting_instance_type_variants": {
+                "regional_aliases": {
+                    "us-west-2": {
+                        "neuron-ecr-uri": "763104351884.dkr.ecr.us-west-2.amazonaws.com/"
+                        "huggingface-pytorch-hosting:2.0.0-transformers4.28.1-gpu-py310-cu118-ubuntu20.04"
+                    }
+                },
+                "variants": {"inf2": {"regional_properties": {"image_uri": "$neuron-ecr-uri"}}},
+            },
+        },
     )
-    assert list(config.config_components.keys()) == ["neuron-base"]
+    assert list(config.config_components.keys()) == ["neuron-inference"]
+
+
+def test_set_inference_configs():
+    spec = {**BASE_SPEC, **INFERENCE_CONFIGS, **INFERENCE_CONFIG_RANKINGS}
+    specs1 = JumpStartModelSpecs(spec)
+
+    assert list(specs1.inference_config_components.keys()) == [
+        "neuron-base",
+        "neuron-inference",
+        "neuron-budget",
+        "gpu-inference",
+        "gpu-inference-budget",
+    ]
+
+    with pytest.raises(ValueError) as error:
+        specs1.set_config("invalid_name")
+    assert "Cannot find Jumpstart config name invalid_name."
+    "List of config names that is supported by the model: "
+    "['neuron-inference', 'neuron-inference-budget', "
+    "'gpu-inference-budget', 'gpu-inference']" in str(error.value)
+
+    assert specs1.supported_inference_instance_types == ["ml.inf2.xlarge", "ml.inf2.2xlarge"]
+    specs1.set_config("gpu-inference")
+    assert specs1.supported_inference_instance_types == ["ml.p2.xlarge", "ml.p3.2xlarge"]
 
 
 def test_training_configs_parsing():
@@ -1133,12 +1179,12 @@ def test_training_configs_parsing():
     config = specs1.training_configs.get_top_config_from_ranking()
 
     assert config.benchmark_metrics == {
-        "ml.tr1n1.2xlarge": JumpStartBenchmarkStat(
-            {"name": "Latency", "value": "100", "unit": "Tokens/S"}
-        ),
-        "ml.tr1n1.4xlarge": JumpStartBenchmarkStat(
-            {"name": "Latency", "value": "50", "unit": "Tokens/S"}
-        ),
+        "ml.tr1n1.2xlarge": [
+            JumpStartBenchmarkStat({"name": "Latency", "value": "100", "unit": "Tokens/S"})
+        ],
+        "ml.tr1n1.4xlarge": [
+            JumpStartBenchmarkStat({"name": "Latency", "value": "50", "unit": "Tokens/S"})
+        ],
     }
     assert len(config.config_components) == 1
     assert config.config_components["neuron-training"] == JumpStartConfigComponent(
@@ -1192,3 +1238,13 @@ def test_set_training_config():
         specs1.training_artifact_key
         == "artifacts/meta-textgeneration-llama-2-7b/gpu-training-budget/model/"
     )
+
+    with pytest.raises(ValueError) as error:
+        specs1.set_config("invalid_name", scope=JumpStartScriptScope.TRAINING)
+    assert "Cannot find Jumpstart config name invalid_name."
+    "List of config names that is supported by the model: "
+    "['neuron-training', 'neuron-training-budget', "
+    "'gpu-training-budget', 'gpu-training']" in str(error.value)
+
+    with pytest.raises(ValueError) as error:
+        specs1.set_config("invalid_name", scope="unknown scope")
diff --git a/tests/unit/sagemaker/jumpstart/test_utils.py b/tests/unit/sagemaker/jumpstart/test_utils.py
index e7a7d522c3..c1ea8abcb8 100644
--- a/tests/unit/sagemaker/jumpstart/test_utils.py
+++ b/tests/unit/sagemaker/jumpstart/test_utils.py
@@ -1598,24 +1598,24 @@ def test_get_jumpstart_benchmark_stats_full_list(
             "mock-region", "mock-model", "mock-model-version", config_names=None
         ) == {
             "neuron-inference": {
-                "ml.inf2.2xlarge": JumpStartBenchmarkStat(
-                    {"name": "Latency", "value": "100", "unit": "Tokens/S"}
-                )
+                "ml.inf2.2xlarge": [
+                    JumpStartBenchmarkStat({"name": "Latency", "value": "100", "unit": "Tokens/S"})
+                ]
             },
             "neuron-inference-budget": {
-                "ml.inf2.2xlarge": JumpStartBenchmarkStat(
-                    {"name": "Latency", "value": "100", "unit": "Tokens/S"}
-                )
+                "ml.inf2.2xlarge": [
+                    JumpStartBenchmarkStat({"name": "Latency", "value": "100", "unit": "Tokens/S"})
+                ]
             },
             "gpu-inference-budget": {
-                "ml.p3.2xlarge": JumpStartBenchmarkStat(
-                    {"name": "Latency", "value": "100", "unit": "Tokens/S"}
-                )
+                "ml.p3.2xlarge": [
+                    JumpStartBenchmarkStat({"name": "Latency", "value": "100", "unit": "Tokens/S"})
+                ]
             },
             "gpu-inference": {
-                "ml.p3.2xlarge": JumpStartBenchmarkStat(
-                    {"name": "Latency", "value": "100", "unit": "Tokens/S"}
-                )
+                "ml.p3.2xlarge": [
+                    JumpStartBenchmarkStat({"name": "Latency", "value": "100", "unit": "Tokens/S"})
+                ]
             },
         }
 
@@ -1633,14 +1633,14 @@ def test_get_jumpstart_benchmark_stats_partial_list(
             config_names=["neuron-inference-budget", "gpu-inference-budget"],
         ) == {
             "neuron-inference-budget": {
-                "ml.inf2.2xlarge": JumpStartBenchmarkStat(
-                    {"name": "Latency", "value": "100", "unit": "Tokens/S"}
-                )
+                "ml.inf2.2xlarge": [
+                    JumpStartBenchmarkStat({"name": "Latency", "value": "100", "unit": "Tokens/S"})
+                ]
             },
             "gpu-inference-budget": {
-                "ml.p3.2xlarge": JumpStartBenchmarkStat(
-                    {"name": "Latency", "value": "100", "unit": "Tokens/S"}
-                )
+                "ml.p3.2xlarge": [
+                    JumpStartBenchmarkStat({"name": "Latency", "value": "100", "unit": "Tokens/S"})
+                ]
             },
         }
 
@@ -1658,9 +1658,9 @@ def test_get_jumpstart_benchmark_stats_single_stat(
             config_names=["neuron-inference-budget"],
         ) == {
             "neuron-inference-budget": {
-                "ml.inf2.2xlarge": JumpStartBenchmarkStat(
-                    {"name": "Latency", "value": "100", "unit": "Tokens/S"}
-                )
+                "ml.inf2.2xlarge": [
+                    JumpStartBenchmarkStat({"name": "Latency", "value": "100", "unit": "Tokens/S"})
+                ]
             }
         }
 
@@ -1695,16 +1695,16 @@ def test_get_jumpstart_benchmark_stats_training(
             config_names=["neuron-training", "gpu-training-budget"],
         ) == {
             "neuron-training": {
-                "ml.tr1n1.2xlarge": JumpStartBenchmarkStat(
-                    {"name": "Latency", "value": "100", "unit": "Tokens/S"}
-                ),
-                "ml.tr1n1.4xlarge": JumpStartBenchmarkStat(
-                    {"name": "Latency", "value": "50", "unit": "Tokens/S"}
-                ),
+                "ml.tr1n1.2xlarge": [
+                    JumpStartBenchmarkStat({"name": "Latency", "value": "100", "unit": "Tokens/S"})
+                ],
+                "ml.tr1n1.4xlarge": [
+                    JumpStartBenchmarkStat({"name": "Latency", "value": "50", "unit": "Tokens/S"})
+                ],
             },
             "gpu-training-budget": {
-                "ml.p3.2xlarge": JumpStartBenchmarkStat(
-                    {"name": "Latency", "value": "100", "unit": "Tokens/S"}
-                )
+                "ml.p3.2xlarge": [
+                    JumpStartBenchmarkStat({"name": "Latency", "value": "100", "unit": "Tokens/S"})
+                ]
             },
         }

From 30c9bf670cc5acb735376046ee6aeb71d628b0f9 Mon Sep 17 00:00:00 2001
From: Nikhil Kulkarni <knikhil29@gmail.com>
Date: Tue, 23 Apr 2024 09:30:46 -0700
Subject: [PATCH 02/18] Add Triton v24.03 URI (#4605)

Co-authored-by: Nikhil Kulkarni <nikhilsk@amazon.com>
---
 .../sagemaker-tritonserver.json               | 36 +++++++++++++++++--
 1 file changed, 34 insertions(+), 2 deletions(-)

diff --git a/src/sagemaker/image_uri_config/sagemaker-tritonserver.json b/src/sagemaker/image_uri_config/sagemaker-tritonserver.json
index 82397d913e..b2257ce803 100644
--- a/src/sagemaker/image_uri_config/sagemaker-tritonserver.json
+++ b/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"
 		}
 	}
 }
\ No newline at end of file

From fe32d799fa991ab0cca8caedcb07a379e5b6acef Mon Sep 17 00:00:00 2001
From: jessicazhu3 <106775307+jessicazhu3@users.noreply.github.com>
Date: Wed, 24 Apr 2024 11:37:29 -0700
Subject: [PATCH 03/18] feature: support session tag chaining for training job
 (#4596)

* feature: support session tag chaining for training job

* fix: resolve typo

* fix: resolve typo and build failure

* fix: resolve typo and unit test failure

---------

Co-authored-by: Jessica Zhu <jessicazhu3@106775307+jessicazhu3@users.noreply.github.com>
---
 src/sagemaker/estimator.py                   | 22 +++++++++++-
 src/sagemaker/jumpstart/estimator.py         |  4 +++
 src/sagemaker/jumpstart/factory/estimator.py |  2 ++
 src/sagemaker/jumpstart/types.py             |  3 ++
 src/sagemaker/session.py                     | 24 ++++++++++++++
 tests/unit/test_estimator.py                 | 35 ++++++++++++++++++++
 tests/unit/test_session.py                   |  3 ++
 7 files changed, 92 insertions(+), 1 deletion(-)

diff --git a/src/sagemaker/estimator.py b/src/sagemaker/estimator.py
index 066846564e..58a5fabc2f 100644
--- a/src/sagemaker/estimator.py
+++ b/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,
         )
 
diff --git a/src/sagemaker/jumpstart/estimator.py b/src/sagemaker/jumpstart/estimator.py
index 88927ae931..bade834cc6 100644
--- a/src/sagemaker/jumpstart/estimator.py
+++ b/src/sagemaker/jumpstart/estimator.py
@@ -109,6 +109,7 @@ def __init__(
         container_arguments: Optional[List[str]] = None,
         disable_output_compression: Optional[bool] = None,
         enable_remote_debug: Optional[Union[bool, PipelineVariable]] = None,
+        enable_session_tag_chaining: Optional[Union[bool, PipelineVariable]] = None,
     ):
         """Initializes a ``JumpStartEstimator``.
 
@@ -500,6 +501,8 @@ def __init__(
                 to Amazon S3 without compression after training finishes.
             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
 
         Raises:
             ValueError: If the model ID is not recognized by JumpStart.
@@ -578,6 +581,7 @@ def _validate_model_id_and_get_type_hook():
             disable_output_compression=disable_output_compression,
             enable_infra_check=enable_infra_check,
             enable_remote_debug=enable_remote_debug,
+            enable_session_tag_chaining=enable_session_tag_chaining,
         )
 
         self.model_id = estimator_init_kwargs.model_id
diff --git a/src/sagemaker/jumpstart/factory/estimator.py b/src/sagemaker/jumpstart/factory/estimator.py
index 875ec9d003..387a4a843c 100644
--- a/src/sagemaker/jumpstart/factory/estimator.py
+++ b/src/sagemaker/jumpstart/factory/estimator.py
@@ -130,6 +130,7 @@ def get_init_kwargs(
     disable_output_compression: Optional[bool] = None,
     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,
 ) -> JumpStartEstimatorInitKwargs:
     """Returns kwargs required to instantiate `sagemaker.estimator.Estimator` object."""
 
@@ -188,6 +189,7 @@ def get_init_kwargs(
         disable_output_compression=disable_output_compression,
         enable_infra_check=enable_infra_check,
         enable_remote_debug=enable_remote_debug,
+        enable_session_tag_chaining=enable_session_tag_chaining,
     )
 
     estimator_init_kwargs = _add_model_version_to_kwargs(estimator_init_kwargs)
diff --git a/src/sagemaker/jumpstart/types.py b/src/sagemaker/jumpstart/types.py
index 05c6a00961..dae879494e 100644
--- a/src/sagemaker/jumpstart/types.py
+++ b/src/sagemaker/jumpstart/types.py
@@ -1751,6 +1751,7 @@ class JumpStartEstimatorInitKwargs(JumpStartKwargs):
         "disable_output_compression",
         "enable_infra_check",
         "enable_remote_debug",
+        "enable_session_tag_chaining",
     ]
 
     SERIALIZATION_EXCLUSION_SET = {
@@ -1818,6 +1819,7 @@ def __init__(
         disable_output_compression: Optional[bool] = None,
         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,
     ) -> None:
         """Instantiates JumpStartEstimatorInitKwargs object."""
 
@@ -1877,6 +1879,7 @@ def __init__(
         self.disable_output_compression = disable_output_compression
         self.enable_infra_check = enable_infra_check
         self.enable_remote_debug = enable_remote_debug
+        self.enable_session_tag_chaining = enable_session_tag_chaining
 
 
 class JumpStartEstimatorFitKwargs(JumpStartKwargs):
diff --git a/src/sagemaker/session.py b/src/sagemaker/session.py
index 9e593706c1..5ea3d5f8a1 100644
--- a/src/sagemaker/session.py
+++ b/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
 
diff --git a/tests/unit/test_estimator.py b/tests/unit/test_estimator.py
index 382c48fde6..fd45601801 100644
--- a/tests/unit/test_estimator.py
+++ b/tests/unit/test_estimator.py
@@ -2089,6 +2089,41 @@ def test_framework_disable_remote_debug(sagemaker_session):
     assert len(args) == 2
 
 
+def test_framework_with_session_chaining_config(sagemaker_session):
+    f = DummyFramework(
+        entry_point=SCRIPT_PATH,
+        role=ROLE,
+        sagemaker_session=sagemaker_session,
+        instance_groups=[
+            InstanceGroup("group1", "ml.c4.xlarge", 1),
+            InstanceGroup("group2", "ml.m4.xlarge", 2),
+        ],
+        enable_session_tag_chaining=True,
+    )
+    f.fit("s3://mydata")
+    sagemaker_session.train.assert_called_once()
+    _, args = sagemaker_session.train.call_args
+    assert args["session_chaining_config"]["EnableSessionTagChaining"]
+    assert f.get_session_chaining_config()["EnableSessionTagChaining"]
+
+
+def test_framework_without_session_chaining_config(sagemaker_session):
+    f = DummyFramework(
+        entry_point=SCRIPT_PATH,
+        role=ROLE,
+        sagemaker_session=sagemaker_session,
+        instance_groups=[
+            InstanceGroup("group1", "ml.c4.xlarge", 1),
+            InstanceGroup("group2", "ml.m4.xlarge", 2),
+        ],
+    )
+    f.fit("s3://mydata")
+    sagemaker_session.train.assert_called_once()
+    _, args = sagemaker_session.train.call_args
+    assert args.get("SessionTagChaining") is None
+    assert f.get_remote_debug_config() is None
+
+
 @patch("time.strftime", return_value=TIMESTAMP)
 def test_custom_code_bucket(time, sagemaker_session):
     code_bucket = "codebucket"
diff --git a/tests/unit/test_session.py b/tests/unit/test_session.py
index 19f9d0ae3d..944f22acff 100644
--- a/tests/unit/test_session.py
+++ b/tests/unit/test_session.py
@@ -2197,6 +2197,7 @@ def test_train_pack_to_request_with_optional_params(sagemaker_session):
     CONTAINER_ENTRY_POINT = ["bin/bash", "test.sh"]
     CONTAINER_ARGUMENTS = ["--arg1", "value1", "--arg2", "value2"]
     remote_debug_config = {"EnableRemoteDebug": True}
+    session_chaining_config = {"EnableSessionTagChaining": True}
 
     sagemaker_session.train(
         image_uri=IMAGE,
@@ -2222,6 +2223,7 @@ def test_train_pack_to_request_with_optional_params(sagemaker_session):
         container_entry_point=CONTAINER_ENTRY_POINT,
         container_arguments=CONTAINER_ARGUMENTS,
         remote_debug_config=remote_debug_config,
+        session_chaining_config=session_chaining_config,
     )
 
     _, _, actual_train_args = sagemaker_session.sagemaker_client.method_calls[0]
@@ -2245,6 +2247,7 @@ def test_train_pack_to_request_with_optional_params(sagemaker_session):
     )
     assert actual_train_args["AlgorithmSpecification"]["ContainerArguments"] == CONTAINER_ARGUMENTS
     assert actual_train_args["RemoteDebugConfig"]["EnableRemoteDebug"]
+    assert actual_train_args["SessionChainingConfig"]["EnableSessionTagChaining"]
 
 
 def test_create_transform_job_with_sagemaker_config_injection(sagemaker_session):

From 8984d9243fd37cc56639889d1c21ded68d868c0b Mon Sep 17 00:00:00 2001
From: ci <ci>
Date: Wed, 24 Apr 2024 21:35:04 +0000
Subject: [PATCH 04/18] prepare release v2.217.0

---
 CHANGELOG.md | 13 +++++++++++++
 VERSION      |  2 +-
 2 files changed, 14 insertions(+), 1 deletion(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 546c3438a0..880e5df8c5 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,5 +1,18 @@
 # Changelog
 
+## 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
diff --git a/VERSION b/VERSION
index 9558cc93a5..b236067a9e 100644
--- a/VERSION
+++ b/VERSION
@@ -1 +1 @@
-2.216.2.dev0
+2.217.0

From ed390dddffa95049e39ab45f75f276015cf12ff6 Mon Sep 17 00:00:00 2001
From: ci <ci>
Date: Wed, 24 Apr 2024 21:35:06 +0000
Subject: [PATCH 05/18] update development version to v2.217.1.dev0

---
 VERSION | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/VERSION b/VERSION
index b236067a9e..70303736d8 100644
--- a/VERSION
+++ b/VERSION
@@ -1 +1 @@
-2.217.0
+2.217.1.dev0

From 2a52478e7fd185f197e944568173c58bd2495d78 Mon Sep 17 00:00:00 2001
From: Justin <justinm088@hotmail.com>
Date: Thu, 25 Apr 2024 11:30:56 -0500
Subject: [PATCH 06/18] fix: properly close files in lineage queries and tests
 (#4587)

Closes #4458
---
 src/sagemaker/lineage/query.py                       |  4 ++--
 tests/data/sip/training.py                           |  3 ++-
 .../sagemaker/lineage/test_lineage_visualize.py      |  4 ++--
 tests/integ/sagemaker/workflow/test_workflow.py      |  3 ++-
 tests/integ/test_sagemaker_config.py                 |  6 ++++--
 tests/unit/sagemaker/local/test_local_image.py       | 12 ++++++++----
 tests/unit/sagemaker/serializers/test_serializers.py |  6 ++++--
 7 files changed, 24 insertions(+), 14 deletions(-)

diff --git a/src/sagemaker/lineage/query.py b/src/sagemaker/lineage/query.py
index 182f117913..3e2003674b 100644
--- a/src/sagemaker/lineage/query.py
+++ b/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;
diff --git a/tests/data/sip/training.py b/tests/data/sip/training.py
index 9b643b2c35..f13b8b3533 100644
--- a/tests/data/sip/training.py
+++ b/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__":
diff --git a/tests/integ/sagemaker/lineage/test_lineage_visualize.py b/tests/integ/sagemaker/lineage/test_lineage_visualize.py
index 4b9e816623..2fdf735d21 100644
--- a/tests/integ/sagemaker/lineage/test_lineage_visualize.py
+++ b/tests/integ/sagemaker/lineage/test_lineage_visualize.py
@@ -142,8 +142,8 @@ def test_graph_visualize(sagemaker_session, extract_data_from_html):
         lq_result.visualize(path="testGraph.html")
 
         # check generated graph info
-        fo = open("testGraph.html", "r")
-        lines = fo.readlines()
+        with open("testGraph.html", "r") as fo:
+            lines = fo.readlines()
         for line in lines:
             if "nodes = " in line:
                 node = line
diff --git a/tests/integ/sagemaker/workflow/test_workflow.py b/tests/integ/sagemaker/workflow/test_workflow.py
index 81871d13f4..2643a3b88e 100644
--- a/tests/integ/sagemaker/workflow/test_workflow.py
+++ b/tests/integ/sagemaker/workflow/test_workflow.py
@@ -527,7 +527,8 @@ def test_one_step_ingestion_pipeline(
 
     temp_flow_path = "./ingestion.flow"
     with cleanup_feature_group(feature_group):
-        json.dump(ingestion_only_flow, open(temp_flow_path, "w"))
+        with open(temp_flow_path, "w") as f:
+            json.dump(ingestion_only_flow, f)
 
         data_wrangler_processor = DataWranglerProcessor(
             role=role,
diff --git a/tests/integ/test_sagemaker_config.py b/tests/integ/test_sagemaker_config.py
index 45efb1f6ab..8cc8d50053 100644
--- a/tests/integ/test_sagemaker_config.py
+++ b/tests/integ/test_sagemaker_config.py
@@ -66,7 +66,8 @@ def expected_merged_config():
     expected_merged_config_file_path = os.path.join(
         CONFIG_DATA_DIR, "expected_output_config_after_merge.yaml"
     )
-    return yaml.safe_load(open(expected_merged_config_file_path, "r").read())
+    with open(expected_merged_config_file_path, "r") as f:
+        return yaml.safe_load(f.read())
 
 
 @pytest.fixture(scope="module")
@@ -171,7 +172,8 @@ def test_config_download_from_s3_and_merge(
         CONFIG_DATA_DIR, "sample_additional_config_for_merge.yaml"
     )
 
-    config_file_1_as_yaml = open(config_file_1_local_path, "r").read()
+    with open(config_file_1_local_path, "r") as f:
+        config_file_1_as_yaml = f.read()
     s3_uri_config_1 = os.path.join(s3_uri_prefix, "config_1.yaml")
 
     # Upload S3 files in case they dont already exist
diff --git a/tests/unit/sagemaker/local/test_local_image.py b/tests/unit/sagemaker/local/test_local_image.py
index 08c55fa0b4..3142fa6dfa 100644
--- a/tests/unit/sagemaker/local/test_local_image.py
+++ b/tests/unit/sagemaker/local/test_local_image.py
@@ -221,9 +221,12 @@ def test_write_config_file(LocalSession, tmpdir):
     assert os.path.exists(resource_config_file)
     assert os.path.exists(input_data_config_file)
 
-    hyperparameters_data = json.load(open(hyperparameters_file))
-    resource_config_data = json.load(open(resource_config_file))
-    input_data_config_data = json.load(open(input_data_config_file))
+    with open(hyperparameters_file) as f:
+        hyperparameters_data = json.load(f)
+    with open(resource_config_file) as f:
+        resource_config_data = json.load(f)
+    with open(input_data_config_file) as f:
+        input_data_config_data = json.load(f)
 
     # Validate HyperParameters
     for k, v in HYPERPARAMETERS.items():
@@ -280,7 +283,8 @@ def test_write_config_files_input_content_type(LocalSession, tmpdir):
     sagemaker_container.write_config_files(host, HYPERPARAMETERS, input_data_config)
 
     assert os.path.exists(input_data_config_file)
-    parsed_input_config = json.load(open(input_data_config_file))
+    with open(input_data_config_file) as f:
+        parsed_input_config = json.load(f)
     # Validate Input Data Config
     for channel in input_data_config:
         assert channel["ChannelName"] in parsed_input_config
diff --git a/tests/unit/sagemaker/serializers/test_serializers.py b/tests/unit/sagemaker/serializers/test_serializers.py
index 6b70c600ca..cb69e9e4ac 100644
--- a/tests/unit/sagemaker/serializers/test_serializers.py
+++ b/tests/unit/sagemaker/serializers/test_serializers.py
@@ -345,7 +345,8 @@ def test_data_serializer_raw(data_serializer):
         input_image = image.read()
     input_image_data = data_serializer.serialize(input_image)
     validation_image_file_path = os.path.join(DATA_DIR, "", "cuteCat.raw")
-    validation_image_data = open(validation_image_file_path, "rb").read()
+    with open(validation_image_file_path, "rb") as f:
+        validation_image_data = f.read()
     assert input_image_data == validation_image_data
 
 
@@ -353,5 +354,6 @@ def test_data_serializer_file_like(data_serializer):
     input_image_file_path = os.path.join(DATA_DIR, "", "cuteCat.jpg")
     validation_image_file_path = os.path.join(DATA_DIR, "", "cuteCat.raw")
     input_image_data = data_serializer.serialize(input_image_file_path)
-    validation_image_data = open(validation_image_file_path, "rb").read()
+    with open(validation_image_file_path, "rb") as f:
+        validation_image_data = f.read()
     assert input_image_data == validation_image_data

From 72e0c9712aec6fbb82fb40fda091dfc2a42c70a0 Mon Sep 17 00:00:00 2001
From: ASHWIN KRISHNA <38850354+akrishna1995@users.noreply.github.com>
Date: Mon, 29 Apr 2024 11:34:39 -0700
Subject: [PATCH 07/18] feature: set default allow_pickle param to False
 (#4557)

* breaking: set default allow_pickle param to False

* breaking: fix unit tests and linting

NumpyDeserializer will not allow deserialization
unless allow_pickle flag is set to True explicitly

* fix: black-check

---------

Co-authored-by: Ashwin Krishna <ashwikri@amazon.com>
---
 src/sagemaker/base_deserializers.py             | 17 ++++++++++++++---
 .../deserializers/test_deserializers.py         |  3 ++-
 2 files changed, 16 insertions(+), 4 deletions(-)

diff --git a/src/sagemaker/base_deserializers.py b/src/sagemaker/base_deserializers.py
index 7162e5274d..a152f0144d 100644
--- a/src/sagemaker/base_deserializers.py
+++ b/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:
diff --git a/tests/unit/sagemaker/deserializers/test_deserializers.py b/tests/unit/sagemaker/deserializers/test_deserializers.py
index b8ede11ba5..cb1923a094 100644
--- a/tests/unit/sagemaker/deserializers/test_deserializers.py
+++ b/tests/unit/sagemaker/deserializers/test_deserializers.py
@@ -142,7 +142,8 @@ def test_numpy_deserializer_from_npy(numpy_deserializer):
     assert np.array_equal(array, result)
 
 
-def test_numpy_deserializer_from_npy_object_array(numpy_deserializer):
+def test_numpy_deserializer_from_npy_object_array():
+    numpy_deserializer = NumpyDeserializer(allow_pickle=True)
     array = np.array([{"a": "", "b": ""}, {"c": "", "d": ""}])
     stream = io.BytesIO()
     np.save(stream, array)

From b17d332a5e4542d57d2039d08b124edc6042f9fb Mon Sep 17 00:00:00 2001
From: Haotian An <33510317+Captainia@users.noreply.github.com>
Date: Tue, 30 Apr 2024 17:58:41 -0400
Subject: [PATCH 08/18] Fix:invalid component error with new metadata (#4634)

* fix: invalid component name

* tests

* format

* fix vulnerable model integ tests llama 2

* updated

* fix: training dataset location
---
 src/sagemaker/jumpstart/estimator.py                      | 7 ++++++-
 src/sagemaker/jumpstart/types.py                          | 5 ++---
 tests/integ/sagemaker/jumpstart/constants.py              | 1 +
 .../jumpstart/estimator/test_jumpstart_estimator.py       | 3 ++-
 .../unit/sagemaker/jumpstart/estimator/test_estimator.py  | 4 ++++
 tests/unit/sagemaker/jumpstart/test_types.py              | 8 ++++++++
 6 files changed, 23 insertions(+), 5 deletions(-)

diff --git a/src/sagemaker/jumpstart/estimator.py b/src/sagemaker/jumpstart/estimator.py
index bade834cc6..f53d109dc8 100644
--- a/src/sagemaker/jumpstart/estimator.py
+++ b/src/sagemaker/jumpstart/estimator.py
@@ -734,7 +734,12 @@ def attach(
 
         model_version = model_version or "*"
 
-        additional_kwargs = {"model_id": model_id, "model_version": model_version}
+        additional_kwargs = {
+            "model_id": model_id,
+            "model_version": model_version,
+            "tolerate_vulnerable_model": True,  # model is already trained
+            "tolerate_deprecated_model": True,  # model is already trained
+        }
 
         model_specs = verify_model_region_and_return_specs(
             model_id=model_id,
diff --git a/src/sagemaker/jumpstart/types.py b/src/sagemaker/jumpstart/types.py
index dae879494e..05c38da266 100644
--- a/src/sagemaker/jumpstart/types.py
+++ b/src/sagemaker/jumpstart/types.py
@@ -1064,9 +1064,8 @@ def from_json(self, json_obj: Dict[str, Any]) -> None:
                 Dictionary representation of the config component.
         """
         for field in json_obj.keys():
-            if field not in self.__slots__:
-                raise ValueError(f"Invalid component field: {field}")
-            setattr(self, field, json_obj[field])
+            if field in self.__slots__:
+                setattr(self, field, json_obj[field])
 
 
 class JumpStartMetadataConfig(JumpStartDataHolderType):
diff --git a/tests/integ/sagemaker/jumpstart/constants.py b/tests/integ/sagemaker/jumpstart/constants.py
index f5ffbf7a3a..b839866b1f 100644
--- a/tests/integ/sagemaker/jumpstart/constants.py
+++ b/tests/integ/sagemaker/jumpstart/constants.py
@@ -48,6 +48,7 @@ def _to_s3_path(filename: str, s3_prefix: Optional[str]) -> str:
     ("meta-textgeneration-llama-2-7b", "*"): ("training-datasets/sec_amazon/"),
     ("meta-textgeneration-llama-2-7b", "2.*"): ("training-datasets/sec_amazon/"),
     ("meta-textgeneration-llama-2-7b", "3.*"): ("training-datasets/sec_amazon/"),
+    ("meta-textgeneration-llama-2-7b", "4.*"): ("training-datasets/sec_amazon/"),
     ("meta-textgenerationneuron-llama-2-7b", "*"): ("training-datasets/sec_amazon/"),
 }
 
diff --git a/tests/integ/sagemaker/jumpstart/estimator/test_jumpstart_estimator.py b/tests/integ/sagemaker/jumpstart/estimator/test_jumpstart_estimator.py
index a839a293c5..0da64ecf05 100644
--- a/tests/integ/sagemaker/jumpstart/estimator/test_jumpstart_estimator.py
+++ b/tests/integ/sagemaker/jumpstart/estimator/test_jumpstart_estimator.py
@@ -140,7 +140,7 @@ def test_gated_model_training_v1(setup):
 def test_gated_model_training_v2(setup):
 
     model_id = "meta-textgeneration-llama-2-7b"
-    model_version = "3.*"  # model artifacts retrieved from jumpstart-private-cache-* buckets
+    model_version = "4.*"  # model artifacts retrieved from jumpstart-private-cache-* buckets
 
     estimator = JumpStartEstimator(
         model_id=model_id,
@@ -150,6 +150,7 @@ def test_gated_model_training_v2(setup):
         tags=[{"Key": JUMPSTART_TAG, "Value": os.environ[ENV_VAR_JUMPSTART_SDK_TEST_SUITE_ID]}],
         environment={"accept_eula": "true"},
         max_run=259200,  # avoid exceeding resource limits
+        tolerate_vulnerable_model=True,  # tolerate old version of model
     )
 
     # uses ml.g5.12xlarge instance
diff --git a/tests/unit/sagemaker/jumpstart/estimator/test_estimator.py b/tests/unit/sagemaker/jumpstart/estimator/test_estimator.py
index ce5f15b287..36d8b11fab 100644
--- a/tests/unit/sagemaker/jumpstart/estimator/test_estimator.py
+++ b/tests/unit/sagemaker/jumpstart/estimator/test_estimator.py
@@ -1010,6 +1010,8 @@ def test_jumpstart_estimator_attach_eula_model(
                 "model_id": "gemma-model",
                 "model_version": "*",
                 "environment": {"accept_eula": "true"},
+                "tolerate_vulnerable_model": True,
+                "tolerate_deprecated_model": True,
             },
         )
 
@@ -1053,6 +1055,8 @@ def test_jumpstart_estimator_attach_no_model_id_happy_case(
             additional_kwargs={
                 "model_id": "js-trainable-model-prepacked",
                 "model_version": "1.0.0",
+                "tolerate_vulnerable_model": True,
+                "tolerate_deprecated_model": True,
             },
         )
 
diff --git a/tests/unit/sagemaker/jumpstart/test_types.py b/tests/unit/sagemaker/jumpstart/test_types.py
index 5ca01c3c52..b2758c73ef 100644
--- a/tests/unit/sagemaker/jumpstart/test_types.py
+++ b/tests/unit/sagemaker/jumpstart/test_types.py
@@ -1052,6 +1052,14 @@ def test_inference_configs_parsing():
     )
     assert list(config.config_components.keys()) == ["neuron-inference"]
 
+    spec = {
+        **BASE_SPEC,
+        **INFERENCE_CONFIGS,
+        **INFERENCE_CONFIG_RANKINGS,
+        "unrecognized-field": "blah",  # New fields in base metadata fields should be ignored
+    }
+    specs1 = JumpStartModelSpecs(spec)
+
 
 def test_set_inference_configs():
     spec = {**BASE_SPEC, **INFERENCE_CONFIGS, **INFERENCE_CONFIG_RANKINGS}

From 15094ee208ec2b84f9ca7a53bd1afb291406b8e3 Mon Sep 17 00:00:00 2001
From: ci <ci>
Date: Wed, 1 May 2024 21:14:30 +0000
Subject: [PATCH 09/18] prepare release v2.218.0

---
 CHANGELOG.md | 10 ++++++++++
 VERSION      |  2 +-
 2 files changed, 11 insertions(+), 1 deletion(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 880e5df8c5..99416fe44a 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,5 +1,15 @@
 # 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
diff --git a/VERSION b/VERSION
index 70303736d8..45aef98018 100644
--- a/VERSION
+++ b/VERSION
@@ -1 +1 @@
-2.217.1.dev0
+2.218.0

From 7c49f5d43d31e1543cb01cb59e72735c0cb901de Mon Sep 17 00:00:00 2001
From: ci <ci>
Date: Wed, 1 May 2024 21:14:32 +0000
Subject: [PATCH 10/18] update development version to v2.218.1.dev0

---
 VERSION | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/VERSION b/VERSION
index 45aef98018..c611a0a1ab 100644
--- a/VERSION
+++ b/VERSION
@@ -1 +1 @@
-2.218.0
+2.218.1.dev0

From 45e31921994cb11612e7c44e262b48d8ff5f4d9c Mon Sep 17 00:00:00 2001
From: Haotian An <33510317+Captainia@users.noreply.github.com>
Date: Thu, 2 May 2024 13:03:32 -0400
Subject: [PATCH 11/18] chore: update skipped flaky tests (#4644)

* Update skipped flaky tests

* flake8

* format

* format
---
 src/sagemaker/jumpstart/notebook_utils.py     |  9 ++--
 src/sagemaker/jumpstart/payload_utils.py      |  6 ++-
 .../jumpstart/test_notebook_utils.py          | 52 +++++++++----------
 3 files changed, 35 insertions(+), 32 deletions(-)

diff --git a/src/sagemaker/jumpstart/notebook_utils.py b/src/sagemaker/jumpstart/notebook_utils.py
index 732493ce3b..83613cd71b 100644
--- a/src/sagemaker/jumpstart/notebook_utils.py
+++ b/src/sagemaker/jumpstart/notebook_utils.py
@@ -329,9 +329,12 @@ def list_jumpstart_models(  # pylint: disable=redefined-builtin
         return sorted(list(model_id_version_dict.keys()))
 
     if not list_old_models:
-        model_id_version_dict = {
-            model_id: set([max(versions)]) for model_id, versions in model_id_version_dict.items()
-        }
+        for model_id, versions in model_id_version_dict.items():
+            try:
+                model_id_version_dict.update({model_id: set([max(versions)])})
+            except TypeError:
+                versions = [str(v) for v in versions]
+                model_id_version_dict.update({model_id: set([max(versions)])})
 
     model_id_version_set: Set[Tuple[str, str]] = set()
     for model_id in model_id_version_dict:
diff --git a/src/sagemaker/jumpstart/payload_utils.py b/src/sagemaker/jumpstart/payload_utils.py
index 595f801598..e4d31e9c83 100644
--- a/src/sagemaker/jumpstart/payload_utils.py
+++ b/src/sagemaker/jumpstart/payload_utils.py
@@ -23,7 +23,7 @@
 from sagemaker.jumpstart.constants import (
     DEFAULT_JUMPSTART_SAGEMAKER_SESSION,
 )
-from sagemaker.jumpstart.enums import MIMEType
+from sagemaker.jumpstart.enums import JumpStartModelType, MIMEType
 from sagemaker.jumpstart.types import JumpStartSerializablePayload
 from sagemaker.jumpstart.utils import (
     get_jumpstart_content_bucket,
@@ -61,6 +61,7 @@ def _construct_payload(
     tolerate_vulnerable_model: bool = False,
     tolerate_deprecated_model: bool = False,
     sagemaker_session: Session = DEFAULT_JUMPSTART_SAGEMAKER_SESSION,
+    model_type: JumpStartModelType = JumpStartModelType.OPEN_WEIGHTS,
 ) -> Optional[JumpStartSerializablePayload]:
     """Returns example payload from prompt.
 
@@ -83,6 +84,8 @@ def _construct_payload(
             object, used for SageMaker interactions. If not
             specified, one is created using the default AWS configuration
             chain. (Default: sagemaker.jumpstart.constants.DEFAULT_JUMPSTART_SAGEMAKER_SESSION).
+        model_type (JumpStartModelType): The type of the model, can be open weights model or
+            proprietary model. (Default: JumpStartModelType.OPEN_WEIGHTS).
     Returns:
         Optional[JumpStartSerializablePayload]: serializable payload with prompt, or None if
             this feature is unavailable for the specified model.
@@ -94,6 +97,7 @@ def _construct_payload(
         tolerate_vulnerable_model=tolerate_vulnerable_model,
         tolerate_deprecated_model=tolerate_deprecated_model,
         sagemaker_session=sagemaker_session,
+        model_type=model_type,
     )
     if payloads is None or len(payloads) == 0:
         return None
diff --git a/tests/unit/sagemaker/jumpstart/test_notebook_utils.py b/tests/unit/sagemaker/jumpstart/test_notebook_utils.py
index c00d271ef1..6544c59019 100644
--- a/tests/unit/sagemaker/jumpstart/test_notebook_utils.py
+++ b/tests/unit/sagemaker/jumpstart/test_notebook_utils.py
@@ -3,7 +3,6 @@
 
 from unittest import TestCase
 from unittest.mock import Mock, patch
-import datetime
 
 import pytest
 from sagemaker.jumpstart.constants import (
@@ -17,7 +16,6 @@
     get_prototype_manifest,
     get_prototype_model_spec,
 )
-from tests.unit.sagemaker.jumpstart.constants import BASE_PROPRIETARY_MANIFEST
 from sagemaker.jumpstart.enums import JumpStartModelType
 from sagemaker.jumpstart.notebook_utils import (
     _generate_jumpstart_model_versions,
@@ -227,10 +225,6 @@ def test_list_jumpstart_models_simple_case(
         patched_get_manifest.assert_called()
         patched_get_model_specs.assert_not_called()
 
-    @pytest.mark.skipif(
-        datetime.datetime.now() < datetime.datetime(year=2024, month=5, day=1),
-        reason="Contact JumpStart team to fix flaky test.",
-    )
     @patch("sagemaker.jumpstart.accessors.JumpStartModelsAccessor._get_manifest")
     @patch("sagemaker.jumpstart.notebook_utils.DEFAULT_JUMPSTART_SAGEMAKER_SESSION.read_s3_file")
     def test_list_jumpstart_models_script_filter(
@@ -246,23 +240,25 @@ def test_list_jumpstart_models_script_filter(
         manifest_length = len(get_prototype_manifest())
         vals = [True, False]
         for val in vals:
-            kwargs = {"filter": f"training_supported == {val}"}
+            kwargs = {"filter": And(f"training_supported == {val}", "model_type is open_weights")}
             list_jumpstart_models(**kwargs)
             assert patched_read_s3_file.call_count == manifest_length
-            patched_get_manifest.assert_called_once()
+            assert patched_get_manifest.call_count == 2
 
             patched_get_manifest.reset_mock()
             patched_read_s3_file.reset_mock()
 
-            kwargs = {"filter": f"training_supported != {val}"}
+            kwargs = {"filter": And(f"training_supported != {val}", "model_type is open_weights")}
             list_jumpstart_models(**kwargs)
             assert patched_read_s3_file.call_count == manifest_length
             assert patched_get_manifest.call_count == 2
 
             patched_get_manifest.reset_mock()
             patched_read_s3_file.reset_mock()
-
-        kwargs = {"filter": f"training_supported in {vals}", "list_versions": True}
+        kwargs = {
+            "filter": And(f"training_supported != {val}", "model_type is open_weights"),
+            "list_versions": True,
+        }
         assert list_jumpstart_models(**kwargs) == [
             ("catboost-classification-model", "1.0.0"),
             ("huggingface-spc-bert-base-cased", "1.0.0"),
@@ -279,7 +275,7 @@ def test_list_jumpstart_models_script_filter(
         patched_get_manifest.reset_mock()
         patched_read_s3_file.reset_mock()
 
-        kwargs = {"filter": f"training_supported not in {vals}"}
+        kwargs = {"filter": And(f"training_supported not in {vals}", "model_type is open_weights")}
         models = list_jumpstart_models(**kwargs)
         assert [] == models
         assert patched_read_s3_file.call_count == manifest_length
@@ -518,10 +514,6 @@ def get_manifest_more_versions(region: str = JUMPSTART_DEFAULT_REGION_NAME):
             list_old_models=False, list_versions=True
         ) == list_jumpstart_models(list_versions=True)
 
-    @pytest.mark.skipif(
-        datetime.datetime.now() < datetime.datetime(year=2024, month=5, day=1),
-        reason="Contact JumpStart team to fix flaky test.",
-    )
     @patch("sagemaker.jumpstart.accessors.JumpStartModelsAccessor._get_manifest")
     @patch("sagemaker.jumpstart.notebook_utils.DEFAULT_JUMPSTART_SAGEMAKER_SESSION.read_s3_file")
     def test_list_jumpstart_models_vulnerable_models(
@@ -547,12 +539,15 @@ def vulnerable_training_model_spec(bucket, key, *args, **kwargs):
         patched_read_s3_file.side_effect = vulnerable_inference_model_spec
 
         num_specs = len(PROTOTYPICAL_MODEL_SPECS_DICT)
-        num_prop_specs = len(BASE_PROPRIETARY_MANIFEST)
         assert [] == list_jumpstart_models(
-            And("inference_vulnerable is false", "training_vulnerable is false")
+            And(
+                "inference_vulnerable is false",
+                "training_vulnerable is false",
+                "model_type is open_weights",
+            )
         )
 
-        assert patched_read_s3_file.call_count == num_specs + num_prop_specs
+        assert patched_read_s3_file.call_count == num_specs
         assert patched_get_manifest.call_count == 2
 
         patched_get_manifest.reset_mock()
@@ -561,10 +556,14 @@ def vulnerable_training_model_spec(bucket, key, *args, **kwargs):
         patched_read_s3_file.side_effect = vulnerable_training_model_spec
 
         assert [] == list_jumpstart_models(
-            And("inference_vulnerable is false", "training_vulnerable is false")
+            And(
+                "inference_vulnerable is false",
+                "training_vulnerable is false",
+                "model_type is open_weights",
+            )
         )
 
-        assert patched_read_s3_file.call_count == num_specs + num_prop_specs
+        assert patched_read_s3_file.call_count == num_specs
         assert patched_get_manifest.call_count == 2
 
         patched_get_manifest.reset_mock()
@@ -574,10 +573,6 @@ def vulnerable_training_model_spec(bucket, key, *args, **kwargs):
 
         assert patched_read_s3_file.call_count == 0
 
-    @pytest.mark.skipif(
-        datetime.datetime.now() < datetime.datetime(year=2024, month=5, day=1),
-        reason="Contact JumpStart team to fix flaky test.",
-    )
     @patch("sagemaker.jumpstart.accessors.JumpStartModelsAccessor._get_manifest")
     @patch("sagemaker.jumpstart.notebook_utils.DEFAULT_JUMPSTART_SAGEMAKER_SESSION.read_s3_file")
     def test_list_jumpstart_models_deprecated_models(
@@ -598,10 +593,11 @@ def deprecated_model_spec(bucket, key, *args, **kwargs) -> str:
         patched_read_s3_file.side_effect = deprecated_model_spec
 
         num_specs = len(PROTOTYPICAL_MODEL_SPECS_DICT)
-        num_prop_specs = len(BASE_PROPRIETARY_MANIFEST)
-        assert [] == list_jumpstart_models("deprecated equals false")
+        assert [] == list_jumpstart_models(
+            And("deprecated equals false", "model_type is open_weights")
+        )
 
-        assert patched_read_s3_file.call_count == num_specs + num_prop_specs
+        assert patched_read_s3_file.call_count == num_specs
         assert patched_get_manifest.call_count == 2
 
         patched_get_manifest.reset_mock()

From c751dbd9050b9efb58eca4bd33d81e54f02e4a81 Mon Sep 17 00:00:00 2001
From: Haixin Wang <98612668+haixiw@users.noreply.github.com>
Date: Thu, 2 May 2024 14:04:18 -0700
Subject: [PATCH 12/18] chore: release tgi 2.0.1 (#4642)

* chore: release tgi 2.0.1

* minor fix

---------

Co-authored-by: Zhaoqi <52220743+zhaoqizqwang@users.noreply.github.com>
---
 .../image_uri_config/huggingface-llm.json     | 49 ++++++++++++++++++-
 .../image_uris/test_huggingface_llm.py        |  1 +
 2 files changed, 49 insertions(+), 1 deletion(-)

diff --git a/src/sagemaker/image_uri_config/huggingface-llm.json b/src/sagemaker/image_uri_config/huggingface-llm.json
index 10073338e7..d357367e6e 100644
--- a/src/sagemaker/image_uri_config/huggingface-llm.json
+++ b/src/sagemaker/image_uri_config/huggingface-llm.json
@@ -12,7 +12,7 @@
             "1.2": "1.2.0",
             "1.3": "1.3.3",
             "1.4": "1.4.5",
-            "2.0": "2.0.0"
+            "2.0": "2.0.1"
         },
         "versions": {
             "0.6.0": {
@@ -578,6 +578,53 @@
                 "container_version": {
                     "gpu": "cu121-ubuntu22.04"
                 }
+            },
+            "2.0.1": {
+                "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",
+                    "cn-north-1": "727897471807",
+                    "cn-northwest-1": "727897471807",
+                    "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-gov-east-1": "446045086412",
+                    "us-gov-west-1": "442386744353",
+                    "us-iso-east-1": "886529160074",
+                    "us-isob-east-1": "094389454867",
+                    "us-west-1": "763104351884",
+                    "us-west-2": "763104351884",
+                    "ca-west-1": "204538143572"
+                },
+                "tag_prefix": "2.1.1-tgi2.0.1",
+                "repository": "huggingface-pytorch-tgi-inference",
+                "container_version": {
+                    "gpu": "cu121-ubuntu22.04"
+                }
             }
         }
     }
diff --git a/tests/unit/sagemaker/image_uris/test_huggingface_llm.py b/tests/unit/sagemaker/image_uris/test_huggingface_llm.py
index 582e5cf82d..2ef981a109 100644
--- a/tests/unit/sagemaker/image_uris/test_huggingface_llm.py
+++ b/tests/unit/sagemaker/image_uris/test_huggingface_llm.py
@@ -32,6 +32,7 @@
         "1.4.2": "2.1.1-tgi1.4.2-gpu-py310-cu121-ubuntu22.04",
         "1.4.5": "2.1.1-tgi1.4.5-gpu-py310-cu121-ubuntu22.04",
         "2.0.0": "2.1.1-tgi2.0.0-gpu-py310-cu121-ubuntu22.04",
+        "2.0.1": "2.1.1-tgi2.0.1-gpu-py310-cu121-ubuntu22.04",
     },
     "inf2": {
         "0.0.16": "1.13.1-optimum0.0.16-neuronx-py310-ubuntu22.04",

From 0f7e6780fea2ae9ba304d2223e3b932f6c7d7ef8 Mon Sep 17 00:00:00 2001
From: Kalyani Nikure <110067132+knikure@users.noreply.github.com>
Date: Fri, 3 May 2024 11:22:32 -0700
Subject: [PATCH 13/18] fix: Fix UserAgent logging in Python SDK (#4647)

---
 src/sagemaker/session.py      | 29 +++++++++++----
 src/sagemaker/user_agent.py   | 45 ++++------------------
 tests/unit/test_session.py    | 70 +++++++++++++----------------------
 tests/unit/test_user_agent.py | 64 +++++++++-----------------------
 4 files changed, 71 insertions(+), 137 deletions(-)

diff --git a/src/sagemaker/session.py b/src/sagemaker/session.py
index 5ea3d5f8a1..bf2a736871 100644
--- a/src/sagemaker/session.py
+++ b/src/sagemaker/session.py
@@ -121,7 +121,7 @@
 from sagemaker.deprecations import deprecated_class
 from sagemaker.enums import EndpointType
 from sagemaker.inputs import ShuffleConfig, TrainingInput, BatchDataCaptureConfig
-from sagemaker.user_agent import prepend_user_agent
+from sagemaker.user_agent import get_user_agent_extra_suffix
 from sagemaker.utils import (
     name_from_image,
     secondary_training_status_changed,
@@ -285,6 +285,7 @@ def _initialize(
         Creates or uses a boto_session, sagemaker_client and sagemaker_runtime_client.
         Sets the region_name.
         """
+
         self.boto_session = boto_session or boto3.DEFAULT_SESSION or boto3.Session()
 
         self._region_name = self.boto_session.region_name
@@ -293,19 +294,30 @@ def _initialize(
                 "Must setup local AWS configuration with a region supported by SageMaker."
             )
 
-        self.sagemaker_client = sagemaker_client or self.boto_session.client("sagemaker")
-        prepend_user_agent(self.sagemaker_client)
+        # Make use of user_agent_extra field of the botocore_config object
+        # to append SageMaker Python SDK specific user_agent suffix
+        # to the current User-Agent header value from boto3
+        # This config will also make sure that user_agent never fails to log the User-Agent string
+        # even if boto User-Agent header format is updated in the future
+        # Ref: https://botocore.amazonaws.com/v1/documentation/api/latest/reference/config.html
+        botocore_config = botocore.config.Config(user_agent_extra=get_user_agent_extra_suffix())
+
+        # Create sagemaker_client with the botocore_config object
+        # This config is customized to append SageMaker Python SDK specific user_agent suffix
+        self.sagemaker_client = sagemaker_client or self.boto_session.client(
+            "sagemaker", config=botocore_config
+        )
 
         if sagemaker_runtime_client is not None:
             self.sagemaker_runtime_client = sagemaker_runtime_client
         else:
-            config = botocore.config.Config(read_timeout=80)
+            config = botocore.config.Config(
+                read_timeout=80, user_agent_extra=get_user_agent_extra_suffix()
+            )
             self.sagemaker_runtime_client = self.boto_session.client(
                 "runtime.sagemaker", config=config
             )
 
-        prepend_user_agent(self.sagemaker_runtime_client)
-
         if sagemaker_featurestore_runtime_client:
             self.sagemaker_featurestore_runtime_client = sagemaker_featurestore_runtime_client
         else:
@@ -316,8 +328,9 @@ def _initialize(
         if sagemaker_metrics_client:
             self.sagemaker_metrics_client = sagemaker_metrics_client
         else:
-            self.sagemaker_metrics_client = self.boto_session.client("sagemaker-metrics")
-        prepend_user_agent(self.sagemaker_metrics_client)
+            self.sagemaker_metrics_client = self.boto_session.client(
+                "sagemaker-metrics", config=botocore_config
+            )
 
         self.s3_client = self.boto_session.client("s3", region_name=self.boto_region_name)
         self.s3_resource = self.boto_session.resource("s3", region_name=self.boto_region_name)
diff --git a/src/sagemaker/user_agent.py b/src/sagemaker/user_agent.py
index 8af89696c2..c1b2bcac07 100644
--- a/src/sagemaker/user_agent.py
+++ b/src/sagemaker/user_agent.py
@@ -13,8 +13,6 @@
 """Placeholder docstring"""
 from __future__ import absolute_import
 
-import platform
-import sys
 import json
 import os
 
@@ -28,12 +26,6 @@
 STUDIO_METADATA_FILE = "/opt/ml/metadata/resource-metadata.json"
 
 SDK_VERSION = importlib_metadata.version("sagemaker")
-OS_NAME = platform.system() or "UnresolvedOS"
-OS_VERSION = platform.release() or "UnresolvedOSVersion"
-OS_NAME_VERSION = "{}/{}".format(OS_NAME, OS_VERSION)
-PYTHON_VERSION = "Python/{}.{}.{}".format(
-    sys.version_info.major, sys.version_info.minor, sys.version_info.micro
-)
 
 
 def process_notebook_metadata_file():
@@ -63,45 +55,24 @@ def process_studio_metadata_file():
     return None
 
 
-def determine_prefix(user_agent=""):
-    """Determines the prefix for the user agent string.
+def get_user_agent_extra_suffix():
+    """Get the user agent extra suffix string specific to SageMaker Python SDK
 
-    Args:
-        user_agent (str): The user agent string to prepend the prefix to.
+    Adhers to new boto recommended User-Agent 2.0 header format
 
     Returns:
-        str: The user agent string with the prefix prepended.
+        str: The user agent extra suffix string to be appended
     """
-    prefix = "{}/{}".format(SDK_PREFIX, SDK_VERSION)
-
-    if PYTHON_VERSION not in user_agent:
-        prefix = "{} {}".format(prefix, PYTHON_VERSION)
-
-    if OS_NAME_VERSION not in user_agent:
-        prefix = "{} {}".format(prefix, OS_NAME_VERSION)
+    suffix = "lib/{}#{}".format(SDK_PREFIX, SDK_VERSION)
 
     # Get the notebook instance type and prepend it to the user agent string if exists
     notebook_instance_type = process_notebook_metadata_file()
     if notebook_instance_type:
-        prefix = "{} {}/{}".format(prefix, NOTEBOOK_PREFIX, notebook_instance_type)
+        suffix = "{} md/{}#{}".format(suffix, NOTEBOOK_PREFIX, notebook_instance_type)
 
     # Get the studio app type and prepend it to the user agent string if exists
     studio_app_type = process_studio_metadata_file()
     if studio_app_type:
-        prefix = "{} {}/{}".format(prefix, STUDIO_PREFIX, studio_app_type)
-
-    return prefix
-
-
-def prepend_user_agent(client):
-    """Prepends the user agent string with the SageMaker Python SDK version.
-
-    Args:
-        client (botocore.client.BaseClient): The client to prepend the user agent string for.
-    """
-    prefix = determine_prefix(client._client_config.user_agent)
+        suffix = "{} md/{}#{}".format(suffix, STUDIO_PREFIX, studio_app_type)
 
-    if client._client_config.user_agent is None:
-        client._client_config.user_agent = prefix
-    else:
-        client._client_config.user_agent = "{} {}".format(prefix, client._client_config.user_agent)
+    return suffix
diff --git a/tests/unit/test_session.py b/tests/unit/test_session.py
index 944f22acff..f7dede1ce9 100644
--- a/tests/unit/test_session.py
+++ b/tests/unit/test_session.py
@@ -43,8 +43,6 @@
 from sagemaker.utils import update_list_of_dicts_with_values_from_config
 from sagemaker.user_agent import (
     SDK_PREFIX,
-    STUDIO_PREFIX,
-    NOTEBOOK_PREFIX,
 )
 from sagemaker.compute_resource_requirements.resource_requirements import ResourceRequirements
 from tests.unit import (
@@ -87,15 +85,20 @@
     limits={},
 )
 
+SDK_DEFAULT_SUFFIX = f"lib/{SDK_PREFIX}#2.218.0"
+NOTEBOOK_SUFFIX = f"{SDK_DEFAULT_SUFFIX} md/AWS-SageMaker-Notebook-Instance#instance_type"
+STUDIO_SUFFIX = f"{SDK_DEFAULT_SUFFIX} md/AWS-SageMaker-Studio#app_type"
 
-@pytest.fixture()
-def boto_session():
-    boto_mock = Mock(name="boto_session", region_name=REGION)
 
+@pytest.fixture
+def boto_session(request):
+    boto_user_agent = "Boto3/1.33.9 md/Botocore#1.33.9 ua/2.0 os/linux#linux-ver md/arch#x86_64 lang/python#3.10.6"
+    user_agent_suffix = getattr(request, "param", "")
+    boto_mock = Mock(name="boto_session", region_name=REGION)
     client_mock = Mock()
-    client_mock._client_config.user_agent = (
-        "Boto3/1.9.69 Python/3.6.5 Linux/4.14.77-70.82.amzn1.x86_64 Botocore/1.12.69 Resource"
-    )
+    user_agent = f"{boto_user_agent} {SDK_DEFAULT_SUFFIX} {user_agent_suffix}"
+    with patch("sagemaker.user_agent.get_user_agent_extra_suffix", return_value=user_agent_suffix):
+        client_mock._client_config.user_agent = user_agent
     boto_mock.client.return_value = client_mock
     return boto_mock
 
@@ -887,65 +890,42 @@ def test_delete_model(boto_session):
     boto_session.client().delete_model.assert_called_with(ModelName=model_name)
 
 
+@pytest.mark.parametrize("boto_session", [""], indirect=True)
 def test_user_agent_injected(boto_session):
-    assert SDK_PREFIX not in boto_session.client("sagemaker")._client_config.user_agent
-
     sess = Session(boto_session)
-
+    expected_user_agent_suffix = "lib/AWS-SageMaker-Python-SDK#2.218.0"
     for client in [
         sess.sagemaker_client,
         sess.sagemaker_runtime_client,
         sess.sagemaker_metrics_client,
     ]:
-        assert SDK_PREFIX in client._client_config.user_agent
-        assert NOTEBOOK_PREFIX not in client._client_config.user_agent
-        assert STUDIO_PREFIX not in client._client_config.user_agent
+        assert expected_user_agent_suffix in client._client_config.user_agent
 
 
-@patch("sagemaker.user_agent.process_notebook_metadata_file", return_value="ml.t3.medium")
-def test_user_agent_injected_with_nbi(
-    mock_process_notebook_metadata_file,
-    boto_session,
-):
-    assert SDK_PREFIX not in boto_session.client("sagemaker")._client_config.user_agent
-
-    sess = Session(
-        boto_session=boto_session,
+@pytest.mark.parametrize("boto_session", [f"{NOTEBOOK_SUFFIX}"], indirect=True)
+def test_user_agent_with_notebook_instance_type(boto_session):
+    sess = Session(boto_session)
+    expected_user_agent_suffix = (
+        "lib/AWS-SageMaker-Python-SDK#2.218.0 md/AWS-SageMaker-Notebook-Instance#instance_type"
     )
-
     for client in [
         sess.sagemaker_client,
         sess.sagemaker_runtime_client,
         sess.sagemaker_metrics_client,
     ]:
-        mock_process_notebook_metadata_file.assert_called()
-
-        assert SDK_PREFIX in client._client_config.user_agent
-        assert NOTEBOOK_PREFIX in client._client_config.user_agent
-        assert STUDIO_PREFIX not in client._client_config.user_agent
+        assert expected_user_agent_suffix in client._client_config.user_agent
 
 
-@patch("sagemaker.user_agent.process_studio_metadata_file", return_value="dymmy-app-type")
-def test_user_agent_injected_with_studio_app_type(
-    mock_process_studio_metadata_file,
-    boto_session,
-):
-    assert SDK_PREFIX not in boto_session.client("sagemaker")._client_config.user_agent
-
-    sess = Session(
-        boto_session=boto_session,
-    )
-
+@pytest.mark.parametrize("boto_session", [f"{STUDIO_SUFFIX}"], indirect=True)
+def test_user_agent_with_studio_app_type(boto_session):
+    sess = Session(boto_session)
+    expected_user_agent = "lib/AWS-SageMaker-Python-SDK#2.218.0 md/AWS-SageMaker-Studio#app_type"
     for client in [
         sess.sagemaker_client,
         sess.sagemaker_runtime_client,
         sess.sagemaker_metrics_client,
     ]:
-        mock_process_studio_metadata_file.assert_called()
-
-        assert SDK_PREFIX in client._client_config.user_agent
-        assert NOTEBOOK_PREFIX not in client._client_config.user_agent
-        assert STUDIO_PREFIX in client._client_config.user_agent
+        assert expected_user_agent in client._client_config.user_agent
 
 
 def test_training_input_all_defaults():
diff --git a/tests/unit/test_user_agent.py b/tests/unit/test_user_agent.py
index c116fef951..fb46988e7b 100644
--- a/tests/unit/test_user_agent.py
+++ b/tests/unit/test_user_agent.py
@@ -13,20 +13,17 @@
 from __future__ import absolute_import
 
 import json
-from mock import MagicMock, patch, mock_open
+from mock import patch, mock_open
 
 
 from sagemaker.user_agent import (
     SDK_PREFIX,
     SDK_VERSION,
-    PYTHON_VERSION,
-    OS_NAME_VERSION,
     NOTEBOOK_PREFIX,
     STUDIO_PREFIX,
     process_notebook_metadata_file,
     process_studio_metadata_file,
-    determine_prefix,
-    prepend_user_agent,
+    get_user_agent_extra_suffix,
 )
 
 
@@ -60,45 +57,18 @@ def test_process_studio_metadata_file_not_exists(tmp_path):
         assert process_studio_metadata_file() is None
 
 
-# Test determine_prefix function
-def test_determine_prefix_notebook_instance_type(monkeypatch):
-    monkeypatch.setattr(
-        "sagemaker.user_agent.process_notebook_metadata_file", lambda: "instance_type"
-    )
-    assert (
-        determine_prefix()
-        == f"{SDK_PREFIX}/{SDK_VERSION} {PYTHON_VERSION} {OS_NAME_VERSION} {NOTEBOOK_PREFIX}/instance_type"
-    )
-
-
-def test_determine_prefix_studio_app_type(monkeypatch):
-    monkeypatch.setattr(
-        "sagemaker.user_agent.process_studio_metadata_file", lambda: "studio_app_type"
-    )
-    assert (
-        determine_prefix()
-        == f"{SDK_PREFIX}/{SDK_VERSION} {PYTHON_VERSION} {OS_NAME_VERSION} {STUDIO_PREFIX}/studio_app_type"
-    )
-
-
-def test_determine_prefix_no_metadata(monkeypatch):
-    monkeypatch.setattr("sagemaker.user_agent.process_notebook_metadata_file", lambda: None)
-    monkeypatch.setattr("sagemaker.user_agent.process_studio_metadata_file", lambda: None)
-    assert determine_prefix() == f"{SDK_PREFIX}/{SDK_VERSION} {PYTHON_VERSION} {OS_NAME_VERSION}"
-
-
-# Test prepend_user_agent function
-def test_prepend_user_agent_existing_user_agent(monkeypatch):
-    client = MagicMock()
-    client._client_config.user_agent = "existing_user_agent"
-    monkeypatch.setattr("sagemaker.user_agent.determine_prefix", lambda _: "prefix")
-    prepend_user_agent(client)
-    assert client._client_config.user_agent == "prefix existing_user_agent"
-
-
-def test_prepend_user_agent_no_user_agent(monkeypatch):
-    client = MagicMock()
-    client._client_config.user_agent = None
-    monkeypatch.setattr("sagemaker.user_agent.determine_prefix", lambda _: "prefix")
-    prepend_user_agent(client)
-    assert client._client_config.user_agent == "prefix"
+# Test get_user_agent_extra_suffix function
+def test_get_user_agent_extra_suffix():
+    assert get_user_agent_extra_suffix() == f"lib/{SDK_PREFIX}#{SDK_VERSION}"
+
+    with patch("sagemaker.user_agent.process_notebook_metadata_file", return_value="instance_type"):
+        assert (
+            get_user_agent_extra_suffix()
+            == f"lib/{SDK_PREFIX}#{SDK_VERSION} md/{NOTEBOOK_PREFIX}#instance_type"
+        )
+
+    with patch("sagemaker.user_agent.process_studio_metadata_file", return_value="studio_type"):
+        assert (
+            get_user_agent_extra_suffix()
+            == f"lib/{SDK_PREFIX}#{SDK_VERSION} md/{STUDIO_PREFIX}#studio_type"
+        )

From fa1a8bf5dc91e9fa64fb3cd3c699824cee33886a Mon Sep 17 00:00:00 2001
From: ci <ci>
Date: Fri, 3 May 2024 20:28:25 +0000
Subject: [PATCH 14/18] prepare release v2.218.1

---
 CHANGELOG.md | 8 ++++++++
 VERSION      | 2 +-
 2 files changed, 9 insertions(+), 1 deletion(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 99416fe44a..38092bf59e 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,5 +1,13 @@
 # Changelog
 
+## v2.218.1 (2024-05-03)
+
+### Bug Fixes and Other Changes
+
+ * Fix UserAgent logging in Python SDK
+ * chore: release tgi 2.0.1
+ * chore: update skipped flaky tests
+
 ## v2.218.0 (2024-05-01)
 
 ### Features
diff --git a/VERSION b/VERSION
index c611a0a1ab..a80e33fcf7 100644
--- a/VERSION
+++ b/VERSION
@@ -1 +1 @@
-2.218.1.dev0
+2.218.1

From 0075fb3bea06fd9eac36bccf2e9f99be802b9aa1 Mon Sep 17 00:00:00 2001
From: ci <ci>
Date: Fri, 3 May 2024 20:28:27 +0000
Subject: [PATCH 15/18] update development version to v2.218.2.dev0

---
 VERSION | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/VERSION b/VERSION
index a80e33fcf7..b298acdcc9 100644
--- a/VERSION
+++ b/VERSION
@@ -1 +1 @@
-2.218.1
+2.218.2.dev0

From d79e53e780213a3a0c5af35309e074ef56fbd867 Mon Sep 17 00:00:00 2001
From: Haotian An <haotiaa@amazon.com>
Date: Mon, 6 May 2024 15:15:43 +0000
Subject: [PATCH 16/18] formatting

---
 src/sagemaker/jumpstart/types.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/sagemaker/jumpstart/types.py b/src/sagemaker/jumpstart/types.py
index ab2eeed7f0..7a51d075ae 100644
--- a/src/sagemaker/jumpstart/types.py
+++ b/src/sagemaker/jumpstart/types.py
@@ -1097,7 +1097,7 @@ def __init__(
             config (Dict[str, Any]):
                 Dictionary representation of the config.
             base_fields (Dict[str, Any]):
-                The default base fields that are used to construct the final resolved config.
+                The default base fields that are used to construct the resolved config.
             config_components (Dict[str, JumpStartConfigComponent]):
                 The list of components that are used to construct the resolved config.
         """

From 03cdee7745de6b2089e30a2dc2496454df1f92d9 Mon Sep 17 00:00:00 2001
From: Haotian An <haotiaa@amazon.com>
Date: Mon, 6 May 2024 16:13:08 +0000
Subject: [PATCH 17/18] Revert "formatting"

This reverts commit d79e53e780213a3a0c5af35309e074ef56fbd867.
---
 src/sagemaker/jumpstart/types.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/sagemaker/jumpstart/types.py b/src/sagemaker/jumpstart/types.py
index 7a51d075ae..ab2eeed7f0 100644
--- a/src/sagemaker/jumpstart/types.py
+++ b/src/sagemaker/jumpstart/types.py
@@ -1097,7 +1097,7 @@ def __init__(
             config (Dict[str, Any]):
                 Dictionary representation of the config.
             base_fields (Dict[str, Any]):
-                The default base fields that are used to construct the resolved config.
+                The default base fields that are used to construct the final resolved config.
             config_components (Dict[str, JumpStartConfigComponent]):
                 The list of components that are used to construct the resolved config.
         """

From 3b098dba80f5e1702f21c6c44f665b263eefb6c4 Mon Sep 17 00:00:00 2001
From: Haotian An <haotiaa@amazon.com>
Date: Mon, 6 May 2024 20:21:16 +0000
Subject: [PATCH 18/18] placeholder commit

---
 src/sagemaker/jumpstart/types.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/sagemaker/jumpstart/types.py b/src/sagemaker/jumpstart/types.py
index ab2eeed7f0..7a51d075ae 100644
--- a/src/sagemaker/jumpstart/types.py
+++ b/src/sagemaker/jumpstart/types.py
@@ -1097,7 +1097,7 @@ def __init__(
             config (Dict[str, Any]):
                 Dictionary representation of the config.
             base_fields (Dict[str, Any]):
-                The default base fields that are used to construct the final resolved config.
+                The default base fields that are used to construct the resolved config.
             config_components (Dict[str, JumpStartConfigComponent]):
                 The list of components that are used to construct the resolved config.
         """