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/58] 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/58] 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/58] 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/58] 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/58] 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/58] 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/58] 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/58] 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/58] 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/58] 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/58] 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/58] 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/58] 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/58] 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/58] 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 49e09c3bbc9976b35d86c499884378a6b9cf5285 Mon Sep 17 00:00:00 2001
From: Keerthan Vasist <kvasist@amazon.com>
Date: Thu, 2 May 2024 10:38:24 -0700
Subject: [PATCH 16/58] feature: allow choosing js payload by alias in private
 method

---
 src/sagemaker/jumpstart/payload_utils.py      |  5 ++-
 .../sagemaker/jumpstart/test_payload_utils.py | 34 ++++++++++++++++---
 2 files changed, 34 insertions(+), 5 deletions(-)

diff --git a/src/sagemaker/jumpstart/payload_utils.py b/src/sagemaker/jumpstart/payload_utils.py
index e4d31e9c83..9c6716dc64 100644
--- a/src/sagemaker/jumpstart/payload_utils.py
+++ b/src/sagemaker/jumpstart/payload_utils.py
@@ -62,6 +62,7 @@ def _construct_payload(
     tolerate_deprecated_model: bool = False,
     sagemaker_session: Session = DEFAULT_JUMPSTART_SAGEMAKER_SESSION,
     model_type: JumpStartModelType = JumpStartModelType.OPEN_WEIGHTS,
+    alias: Optional[str] = None,
 ) -> Optional[JumpStartSerializablePayload]:
     """Returns example payload from prompt.
 
@@ -102,7 +103,9 @@ def _construct_payload(
     if payloads is None or len(payloads) == 0:
         return None
 
-    payload_to_use: JumpStartSerializablePayload = list(payloads.values())[0]
+    payload_to_use: JumpStartSerializablePayload = (
+        payloads[alias] if alias else list(payloads.values())[0]
+    )
 
     prompt_key: Optional[str] = payload_to_use.prompt_key
     if prompt_key is None:
diff --git a/tests/unit/sagemaker/jumpstart/test_payload_utils.py b/tests/unit/sagemaker/jumpstart/test_payload_utils.py
index afc955e2f3..3c339c9b95 100644
--- a/tests/unit/sagemaker/jumpstart/test_payload_utils.py
+++ b/tests/unit/sagemaker/jumpstart/test_payload_utils.py
@@ -32,10 +32,36 @@ def test_construct_payload(self, patched_get_model_specs):
         region = "us-west-2"
 
         constructed_payload_body = _construct_payload(
-            prompt="kobebryant",
-            model_id=model_id,
-            model_version="*",
-            region=region,
+            prompt="kobebryant", model_id=model_id, model_version="*", region=region
+        ).body
+
+        self.assertEqual(
+            {
+                "hello": {"prompt": "kobebryant"},
+                "seed": 43,
+            },
+            constructed_payload_body,
+        )
+
+        # Unsupported model
+        self.assertIsNone(
+            _construct_payload(
+                prompt="blah",
+                model_id="default_payloads",
+                model_version="*",
+                region=region,
+            )
+        )
+
+    @patch("sagemaker.jumpstart.accessors.JumpStartModelsAccessor.get_model_specs")
+    def test_construct_payload_with_specific_alias(self, patched_get_model_specs):
+        patched_get_model_specs.side_effect = get_special_model_spec
+
+        model_id = "prompt-key"
+        region = "us-west-2"
+
+        constructed_payload_body = _construct_payload(
+            prompt="kobebryant", model_id=model_id, model_version="*", region=region, alias="Dog"
         ).body
 
         self.assertEqual(

From 3ae20b7476100c629dd036306b506f48dda0bb5b Mon Sep 17 00:00:00 2001
From: SuhitK <kodgule.suhit@gmail.com>
Date: Tue, 7 May 2024 17:42:59 -0700
Subject: [PATCH 17/58] Updates for SMP v2.3.1 (#4660)

Co-authored-by: Suhit Kodgule <skodgule@amazon.com>
---
 src/sagemaker/fw_utils.py                     |  2 +-
 .../image_uri_config/pytorch-smp.json         | 29 +++++++++++++++++--
 2 files changed, 28 insertions(+), 3 deletions(-)

diff --git a/src/sagemaker/fw_utils.py b/src/sagemaker/fw_utils.py
index cf9291a139..33018becdd 100644
--- a/src/sagemaker/fw_utils.py
+++ b/src/sagemaker/fw_utils.py
@@ -142,7 +142,6 @@
         "2.1.0",
         "2.1.2",
         "2.2.0",
-        "2.3.0",
     ],
 }
 
@@ -170,6 +169,7 @@
     "2.1.2",
     "2.2.0",
     "2.3.0",
+    "2.3.1",
 ]
 
 TRAINIUM_SUPPORTED_DISTRIBUTION_STRATEGIES = ["torch_distributed"]
diff --git a/src/sagemaker/image_uri_config/pytorch-smp.json b/src/sagemaker/image_uri_config/pytorch-smp.json
index faf7d6a14a..518da5f15d 100644
--- a/src/sagemaker/image_uri_config/pytorch-smp.json
+++ b/src/sagemaker/image_uri_config/pytorch-smp.json
@@ -6,8 +6,8 @@
         "version_aliases": {
             "2.0": "2.0.1",
             "2.1": "2.1.2",
-            "2.2": "2.3.0",
-            "2.2.0": "2.3.0"
+            "2.2": "2.3.1",
+            "2.2.0": "2.3.1"
         },
         "versions": {
             "2.0.1": {
@@ -109,6 +109,31 @@
                     "us-west-2": "658645717510"
                 },
                 "repository": "smdistributed-modelparallel"
+            },
+            "2.3.1": {
+                "py_versions": [
+                    "py310"
+                ],
+                "registries": {
+                    "ap-northeast-1": "658645717510",
+                    "ap-northeast-2": "658645717510",
+                    "ap-northeast-3": "658645717510",
+                    "ap-south-1": "658645717510",
+                    "ap-southeast-1": "658645717510",
+                    "ap-southeast-2": "658645717510",
+                    "ca-central-1": "658645717510",
+                    "eu-central-1": "658645717510",
+                    "eu-north-1": "658645717510",
+                    "eu-west-1": "658645717510",
+                    "eu-west-2": "658645717510",
+                    "eu-west-3": "658645717510",
+                    "sa-east-1": "658645717510",
+                    "us-east-1": "658645717510",
+                    "us-east-2": "658645717510",
+                    "us-west-1": "658645717510",
+                    "us-west-2": "658645717510"
+                },
+                "repository": "smdistributed-modelparallel"
             }
         }
     }

From b1823cb16eb3483de6a716facaa22102dd132904 Mon Sep 17 00:00:00 2001
From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com>
Date: Tue, 7 May 2024 22:30:38 -0700
Subject: [PATCH 18/58] chore(deps): bump jinja2 from 3.1.3 to 3.1.4 in /doc
 (#4655)

Bumps [jinja2](https://github.com/pallets/jinja) from 3.1.3 to 3.1.4.
- [Release notes](https://github.com/pallets/jinja/releases)
- [Changelog](https://github.com/pallets/jinja/blob/main/CHANGES.rst)
- [Commits](https://github.com/pallets/jinja/compare/3.1.3...3.1.4)

---
updated-dependencies:
- dependency-name: jinja2
  dependency-type: direct:production
...

Signed-off-by: dependabot[bot] <support@github.com>
Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>
---
 doc/requirements.txt | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/doc/requirements.txt b/doc/requirements.txt
index a65e0e4050..8193dfa22a 100644
--- a/doc/requirements.txt
+++ b/doc/requirements.txt
@@ -2,6 +2,6 @@ sphinx==5.1.1
 sphinx-rtd-theme==0.5.0
 docutils==0.15.2
 packaging==20.9
-jinja2==3.1.3
+jinja2==3.1.4
 schema==0.7.5
 accelerate>=0.24.1,<=0.27.0

From 2058912738f7a48459afe167ce619f8ebd281c06 Mon Sep 17 00:00:00 2001
From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com>
Date: Tue, 7 May 2024 22:31:00 -0700
Subject: [PATCH 19/58] chore(deps): bump tqdm from 4.66.2 to 4.66.3 in
 /tests/data/serve_resources/mlflow/pytorch (#4650)

Bumps [tqdm](https://github.com/tqdm/tqdm) from 4.66.2 to 4.66.3.
- [Release notes](https://github.com/tqdm/tqdm/releases)
- [Commits](https://github.com/tqdm/tqdm/compare/v4.66.2...v4.66.3)

---
updated-dependencies:
- dependency-name: tqdm
  dependency-type: direct:production
...

Signed-off-by: dependabot[bot] <support@github.com>
Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>
---
 tests/data/serve_resources/mlflow/pytorch/requirements.txt | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/data/serve_resources/mlflow/pytorch/requirements.txt b/tests/data/serve_resources/mlflow/pytorch/requirements.txt
index 9848949b0f..a10494de69 100644
--- a/tests/data/serve_resources/mlflow/pytorch/requirements.txt
+++ b/tests/data/serve_resources/mlflow/pytorch/requirements.txt
@@ -13,4 +13,4 @@ pyyaml==6.0.1
 requests==2.31.0
 torch==2.0.1
 torchvision==0.15.2
-tqdm==4.66.2
+tqdm==4.66.3

From 533f30a00de1e86cd4d2d5d2a0768c4cd3e9add4 Mon Sep 17 00:00:00 2001
From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com>
Date: Tue, 7 May 2024 22:31:20 -0700
Subject: [PATCH 20/58] chore(deps): bump jinja2 from 3.1.3 to 3.1.4 in
 /requirements/extras (#4654)

Bumps [jinja2](https://github.com/pallets/jinja) from 3.1.3 to 3.1.4.
- [Release notes](https://github.com/pallets/jinja/releases)
- [Changelog](https://github.com/pallets/jinja/blob/main/CHANGES.rst)
- [Commits](https://github.com/pallets/jinja/compare/3.1.3...3.1.4)

---
updated-dependencies:
- dependency-name: jinja2
  dependency-type: direct:production
...

Signed-off-by: dependabot[bot] <support@github.com>
Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>
---
 requirements/extras/test_requirements.txt | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/requirements/extras/test_requirements.txt b/requirements/extras/test_requirements.txt
index 43da930636..2713f88c04 100644
--- a/requirements/extras/test_requirements.txt
+++ b/requirements/extras/test_requirements.txt
@@ -18,7 +18,7 @@ attrs>=23.1.0,<24
 fabric==2.6.0
 requests==2.31.0
 sagemaker-experiments==0.1.35
-Jinja2==3.1.3
+Jinja2==3.1.4
 pyvis==0.2.1
 pandas>=1.3.5,<1.5
 scikit-learn==1.3.0

From 80d24f307e491a2c47121310b1a8966185115f50 Mon Sep 17 00:00:00 2001
From: ci <ci>
Date: Wed, 8 May 2024 20:32:46 +0000
Subject: [PATCH 21/58] prepare release v2.219.0

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

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 38092bf59e..b41a169cdc 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,5 +1,18 @@
 # Changelog
 
+## v2.219.0 (2024-05-08)
+
+### Features
+
+ * allow choosing js payload by alias in private method
+
+### Bug Fixes and Other Changes
+
+ * chore(deps): bump jinja2 from 3.1.3 to 3.1.4 in /requirements/extras
+ * chore(deps): bump tqdm from 4.66.2 to 4.66.3 in /tests/data/serve_resources/mlflow/pytorch
+ * chore(deps): bump jinja2 from 3.1.3 to 3.1.4 in /doc
+ * Updates for SMP v2.3.1
+
 ## v2.218.1 (2024-05-03)
 
 ### Bug Fixes and Other Changes
diff --git a/VERSION b/VERSION
index b298acdcc9..8411e80b25 100644
--- a/VERSION
+++ b/VERSION
@@ -1 +1 @@
-2.218.2.dev0
+2.219.0

From 2a902cd116fe0ecdeb9f0849f94a7eb49a7cf755 Mon Sep 17 00:00:00 2001
From: ci <ci>
Date: Wed, 8 May 2024 20:32:48 +0000
Subject: [PATCH 22/58] update development version to v2.219.1.dev0

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

diff --git a/VERSION b/VERSION
index 8411e80b25..6f59b6ff9a 100644
--- a/VERSION
+++ b/VERSION
@@ -1 +1 @@
-2.219.0
+2.219.1.dev0

From 6fbb6f857aded594fa08723845c96ffaddfc9d4e Mon Sep 17 00:00:00 2001
From: Haotian An <33510317+Captainia@users.noreply.github.com>
Date: Wed, 8 May 2024 17:47:24 -0400
Subject: [PATCH 23/58] fix: skip flakey tests pending investigation (#4667)

---
 tests/unit/sagemaker/jumpstart/test_notebook_utils.py | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/tests/unit/sagemaker/jumpstart/test_notebook_utils.py b/tests/unit/sagemaker/jumpstart/test_notebook_utils.py
index 6544c59019..c0a37c5b38 100644
--- a/tests/unit/sagemaker/jumpstart/test_notebook_utils.py
+++ b/tests/unit/sagemaker/jumpstart/test_notebook_utils.py
@@ -1,5 +1,6 @@
 from __future__ import absolute_import
 import json
+import datetime
 
 from unittest import TestCase
 from unittest.mock import Mock, patch
@@ -514,6 +515,10 @@ 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=8, 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(

From d549e7d8273cd5e91372865a67e2a3295a6149b2 Mon Sep 17 00:00:00 2001
From: sagemaker-bot <sagemaker-bot@amazon.com>
Date: Thu, 9 May 2024 14:17:41 +0000
Subject: [PATCH 24/58] change: update image_uri_configs  05-09-2024 07:17:41
 PST

---
 src/sagemaker/image_uri_config/pytorch.json   |  606 ++++---
 .../image_uri_config/tensorflow.json          | 1605 +++++++++--------
 2 files changed, 1139 insertions(+), 1072 deletions(-)

diff --git a/src/sagemaker/image_uri_config/pytorch.json b/src/sagemaker/image_uri_config/pytorch.json
index f068c68149..f04f9c6b69 100644
--- a/src/sagemaker/image_uri_config/pytorch.json
+++ b/src/sagemaker/image_uri_config/pytorch.json
@@ -23,17 +23,17 @@
                     "ap-southeast-3": "907027046896",
                     "ap-southeast-4": "457447274322",
                     "ca-central-1": "763104351884",
+                    "ca-west-1": "204538143572",
                     "cn-north-1": "727897471807",
                     "cn-northwest-1": "727897471807",
                     "eu-central-1": "763104351884",
-                    "eu-north-1": "763104351884",
                     "eu-central-2": "380420809688",
-                    "eu-west-1": "763104351884",
+                    "eu-north-1": "763104351884",
                     "eu-south-2": "503227376785",
+                    "eu-west-1": "763104351884",
                     "us-east-1": "763104351884",
                     "us-east-2": "763104351884",
-                    "us-west-2": "763104351884",
-                    "ca-west-1": "204538143572"
+                    "us-west-2": "763104351884"
                 },
                 "repository": "pytorch-inference-eia"
             },
@@ -47,13 +47,13 @@
                     "ap-south-2": "772153158452",
                     "ap-southeast-3": "907027046896",
                     "ap-southeast-4": "457447274322",
+                    "ca-west-1": "204538143572",
                     "eu-central-2": "380420809688",
-                    "eu-west-1": "763104351884",
                     "eu-south-2": "503227376785",
+                    "eu-west-1": "763104351884",
                     "us-east-1": "763104351884",
                     "us-east-2": "763104351884",
-                    "us-west-2": "763104351884",
-                    "ca-west-1": "204538143572"
+                    "us-west-2": "763104351884"
                 },
                 "repository": "pytorch-inference-eia"
             }
@@ -81,7 +81,8 @@
             "1.12": "1.12.1",
             "1.13": "1.13.1",
             "2.0": "2.0.1",
-            "2.1": "2.1.0"
+            "2.1": "2.1.0",
+            "2.2": "2.2.0"
         },
         "versions": {
             "0.4.0": {
@@ -193,7 +194,6 @@
                 ],
                 "registries": {
                     "af-south-1": "626614931356",
-                    "il-central-1": "780543022126",
                     "ap-east-1": "871362719292",
                     "ap-northeast-1": "763104351884",
                     "ap-northeast-2": "763104351884",
@@ -205,18 +205,20 @@
                     "ap-southeast-3": "907027046896",
                     "ap-southeast-4": "457447274322",
                     "ca-central-1": "763104351884",
+                    "ca-west-1": "204538143572",
                     "cn-north-1": "727897471807",
                     "cn-northwest-1": "727897471807",
                     "eu-central-1": "763104351884",
                     "eu-central-2": "380420809688",
                     "eu-north-1": "763104351884",
+                    "eu-south-1": "692866216735",
+                    "eu-south-2": "503227376785",
                     "eu-west-1": "763104351884",
                     "eu-west-2": "763104351884",
                     "eu-west-3": "763104351884",
-                    "eu-south-1": "692866216735",
-                    "eu-south-2": "503227376785",
-                    "me-south-1": "217643126080",
+                    "il-central-1": "780543022126",
                     "me-central-1": "914824155844",
+                    "me-south-1": "217643126080",
                     "sa-east-1": "763104351884",
                     "us-east-1": "763104351884",
                     "us-east-2": "763104351884",
@@ -225,8 +227,7 @@
                     "us-iso-east-1": "886529160074",
                     "us-isob-east-1": "094389454867",
                     "us-west-1": "763104351884",
-                    "us-west-2": "763104351884",
-                    "ca-west-1": "204538143572"
+                    "us-west-2": "763104351884"
                 },
                 "repository": "pytorch-inference"
             },
@@ -237,7 +238,6 @@
                 ],
                 "registries": {
                     "af-south-1": "626614931356",
-                    "il-central-1": "780543022126",
                     "ap-east-1": "871362719292",
                     "ap-northeast-1": "763104351884",
                     "ap-northeast-2": "763104351884",
@@ -249,18 +249,20 @@
                     "ap-southeast-3": "907027046896",
                     "ap-southeast-4": "457447274322",
                     "ca-central-1": "763104351884",
+                    "ca-west-1": "204538143572",
                     "cn-north-1": "727897471807",
                     "cn-northwest-1": "727897471807",
                     "eu-central-1": "763104351884",
                     "eu-central-2": "380420809688",
                     "eu-north-1": "763104351884",
+                    "eu-south-1": "692866216735",
+                    "eu-south-2": "503227376785",
                     "eu-west-1": "763104351884",
                     "eu-west-2": "763104351884",
                     "eu-west-3": "763104351884",
-                    "eu-south-1": "692866216735",
-                    "eu-south-2": "503227376785",
-                    "me-south-1": "217643126080",
+                    "il-central-1": "780543022126",
                     "me-central-1": "914824155844",
+                    "me-south-1": "217643126080",
                     "sa-east-1": "763104351884",
                     "us-east-1": "763104351884",
                     "us-east-2": "763104351884",
@@ -269,18 +271,17 @@
                     "us-iso-east-1": "886529160074",
                     "us-isob-east-1": "094389454867",
                     "us-west-1": "763104351884",
-                    "us-west-2": "763104351884",
-                    "ca-west-1": "204538143572"
+                    "us-west-2": "763104351884"
                 },
                 "repository": "pytorch-inference"
             },
             "1.4.0": {
                 "py_versions": [
-                    "py3"
+                    "py3",
+                    "py36"
                 ],
                 "registries": {
                     "af-south-1": "626614931356",
-                    "il-central-1": "780543022126",
                     "ap-east-1": "871362719292",
                     "ap-northeast-1": "763104351884",
                     "ap-northeast-2": "763104351884",
@@ -292,18 +293,20 @@
                     "ap-southeast-3": "907027046896",
                     "ap-southeast-4": "457447274322",
                     "ca-central-1": "763104351884",
+                    "ca-west-1": "204538143572",
                     "cn-north-1": "727897471807",
                     "cn-northwest-1": "727897471807",
                     "eu-central-1": "763104351884",
                     "eu-central-2": "380420809688",
                     "eu-north-1": "763104351884",
+                    "eu-south-1": "692866216735",
+                    "eu-south-2": "503227376785",
                     "eu-west-1": "763104351884",
                     "eu-west-2": "763104351884",
                     "eu-west-3": "763104351884",
-                    "eu-south-1": "692866216735",
-                    "eu-south-2": "503227376785",
-                    "me-south-1": "217643126080",
+                    "il-central-1": "780543022126",
                     "me-central-1": "914824155844",
+                    "me-south-1": "217643126080",
                     "sa-east-1": "763104351884",
                     "us-east-1": "763104351884",
                     "us-east-2": "763104351884",
@@ -312,18 +315,17 @@
                     "us-iso-east-1": "886529160074",
                     "us-isob-east-1": "094389454867",
                     "us-west-1": "763104351884",
-                    "us-west-2": "763104351884",
-                    "ca-west-1": "204538143572"
+                    "us-west-2": "763104351884"
                 },
                 "repository": "pytorch-inference"
             },
             "1.5.0": {
                 "py_versions": [
-                    "py3"
+                    "py3",
+                    "py36"
                 ],
                 "registries": {
                     "af-south-1": "626614931356",
-                    "il-central-1": "780543022126",
                     "ap-east-1": "871362719292",
                     "ap-northeast-1": "763104351884",
                     "ap-northeast-2": "763104351884",
@@ -335,18 +337,20 @@
                     "ap-southeast-3": "907027046896",
                     "ap-southeast-4": "457447274322",
                     "ca-central-1": "763104351884",
+                    "ca-west-1": "204538143572",
                     "cn-north-1": "727897471807",
                     "cn-northwest-1": "727897471807",
                     "eu-central-1": "763104351884",
                     "eu-central-2": "380420809688",
                     "eu-north-1": "763104351884",
+                    "eu-south-1": "692866216735",
+                    "eu-south-2": "503227376785",
                     "eu-west-1": "763104351884",
                     "eu-west-2": "763104351884",
                     "eu-west-3": "763104351884",
-                    "eu-south-1": "692866216735",
-                    "eu-south-2": "503227376785",
-                    "me-south-1": "217643126080",
+                    "il-central-1": "780543022126",
                     "me-central-1": "914824155844",
+                    "me-south-1": "217643126080",
                     "sa-east-1": "763104351884",
                     "us-east-1": "763104351884",
                     "us-east-2": "763104351884",
@@ -355,8 +359,7 @@
                     "us-iso-east-1": "886529160074",
                     "us-isob-east-1": "094389454867",
                     "us-west-1": "763104351884",
-                    "us-west-2": "763104351884",
-                    "ca-west-1": "204538143572"
+                    "us-west-2": "763104351884"
                 },
                 "repository": "pytorch-inference"
             },
@@ -367,7 +370,6 @@
                 ],
                 "registries": {
                     "af-south-1": "626614931356",
-                    "il-central-1": "780543022126",
                     "ap-east-1": "871362719292",
                     "ap-northeast-1": "763104351884",
                     "ap-northeast-2": "763104351884",
@@ -379,18 +381,20 @@
                     "ap-southeast-3": "907027046896",
                     "ap-southeast-4": "457447274322",
                     "ca-central-1": "763104351884",
+                    "ca-west-1": "204538143572",
                     "cn-north-1": "727897471807",
                     "cn-northwest-1": "727897471807",
                     "eu-central-1": "763104351884",
                     "eu-central-2": "380420809688",
                     "eu-north-1": "763104351884",
+                    "eu-south-1": "692866216735",
+                    "eu-south-2": "503227376785",
                     "eu-west-1": "763104351884",
                     "eu-west-2": "763104351884",
                     "eu-west-3": "763104351884",
-                    "eu-south-1": "692866216735",
-                    "eu-south-2": "503227376785",
-                    "me-south-1": "217643126080",
+                    "il-central-1": "780543022126",
                     "me-central-1": "914824155844",
+                    "me-south-1": "217643126080",
                     "sa-east-1": "763104351884",
                     "us-east-1": "763104351884",
                     "us-east-2": "763104351884",
@@ -399,8 +403,7 @@
                     "us-iso-east-1": "886529160074",
                     "us-isob-east-1": "094389454867",
                     "us-west-1": "763104351884",
-                    "us-west-2": "763104351884",
-                    "ca-west-1": "204538143572"
+                    "us-west-2": "763104351884"
                 },
                 "repository": "pytorch-inference"
             },
@@ -411,7 +414,6 @@
                 ],
                 "registries": {
                     "af-south-1": "626614931356",
-                    "il-central-1": "780543022126",
                     "ap-east-1": "871362719292",
                     "ap-northeast-1": "763104351884",
                     "ap-northeast-2": "763104351884",
@@ -423,18 +425,20 @@
                     "ap-southeast-3": "907027046896",
                     "ap-southeast-4": "457447274322",
                     "ca-central-1": "763104351884",
+                    "ca-west-1": "204538143572",
                     "cn-north-1": "727897471807",
                     "cn-northwest-1": "727897471807",
                     "eu-central-1": "763104351884",
                     "eu-central-2": "380420809688",
                     "eu-north-1": "763104351884",
+                    "eu-south-1": "692866216735",
+                    "eu-south-2": "503227376785",
                     "eu-west-1": "763104351884",
                     "eu-west-2": "763104351884",
                     "eu-west-3": "763104351884",
-                    "eu-south-1": "692866216735",
-                    "eu-south-2": "503227376785",
-                    "me-south-1": "217643126080",
+                    "il-central-1": "780543022126",
                     "me-central-1": "914824155844",
+                    "me-south-1": "217643126080",
                     "sa-east-1": "763104351884",
                     "us-east-1": "763104351884",
                     "us-east-2": "763104351884",
@@ -443,8 +447,7 @@
                     "us-iso-east-1": "886529160074",
                     "us-isob-east-1": "094389454867",
                     "us-west-1": "763104351884",
-                    "us-west-2": "763104351884",
-                    "ca-west-1": "204538143572"
+                    "us-west-2": "763104351884"
                 },
                 "repository": "pytorch-inference"
             },
@@ -455,7 +458,6 @@
                 ],
                 "registries": {
                     "af-south-1": "626614931356",
-                    "il-central-1": "780543022126",
                     "ap-east-1": "871362719292",
                     "ap-northeast-1": "763104351884",
                     "ap-northeast-2": "763104351884",
@@ -467,18 +469,20 @@
                     "ap-southeast-3": "907027046896",
                     "ap-southeast-4": "457447274322",
                     "ca-central-1": "763104351884",
+                    "ca-west-1": "204538143572",
                     "cn-north-1": "727897471807",
                     "cn-northwest-1": "727897471807",
                     "eu-central-1": "763104351884",
                     "eu-central-2": "380420809688",
                     "eu-north-1": "763104351884",
+                    "eu-south-1": "692866216735",
+                    "eu-south-2": "503227376785",
                     "eu-west-1": "763104351884",
                     "eu-west-2": "763104351884",
                     "eu-west-3": "763104351884",
-                    "eu-south-1": "692866216735",
-                    "eu-south-2": "503227376785",
-                    "me-south-1": "217643126080",
+                    "il-central-1": "780543022126",
                     "me-central-1": "914824155844",
+                    "me-south-1": "217643126080",
                     "sa-east-1": "763104351884",
                     "us-east-1": "763104351884",
                     "us-east-2": "763104351884",
@@ -487,8 +491,7 @@
                     "us-iso-east-1": "886529160074",
                     "us-isob-east-1": "094389454867",
                     "us-west-1": "763104351884",
-                    "us-west-2": "763104351884",
-                    "ca-west-1": "204538143572"
+                    "us-west-2": "763104351884"
                 },
                 "repository": "pytorch-inference"
             },
@@ -499,7 +502,6 @@
                 ],
                 "registries": {
                     "af-south-1": "626614931356",
-                    "il-central-1": "780543022126",
                     "ap-east-1": "871362719292",
                     "ap-northeast-1": "763104351884",
                     "ap-northeast-2": "763104351884",
@@ -511,18 +513,20 @@
                     "ap-southeast-3": "907027046896",
                     "ap-southeast-4": "457447274322",
                     "ca-central-1": "763104351884",
+                    "ca-west-1": "204538143572",
                     "cn-north-1": "727897471807",
                     "cn-northwest-1": "727897471807",
                     "eu-central-1": "763104351884",
                     "eu-central-2": "380420809688",
                     "eu-north-1": "763104351884",
+                    "eu-south-1": "692866216735",
+                    "eu-south-2": "503227376785",
                     "eu-west-1": "763104351884",
                     "eu-west-2": "763104351884",
                     "eu-west-3": "763104351884",
-                    "eu-south-1": "692866216735",
-                    "eu-south-2": "503227376785",
-                    "me-south-1": "217643126080",
+                    "il-central-1": "780543022126",
                     "me-central-1": "914824155844",
+                    "me-south-1": "217643126080",
                     "sa-east-1": "763104351884",
                     "us-east-1": "763104351884",
                     "us-east-2": "763104351884",
@@ -531,8 +535,7 @@
                     "us-iso-east-1": "886529160074",
                     "us-isob-east-1": "094389454867",
                     "us-west-1": "763104351884",
-                    "us-west-2": "763104351884",
-                    "ca-west-1": "204538143572"
+                    "us-west-2": "763104351884"
                 },
                 "repository": "pytorch-inference"
             },
@@ -542,7 +545,6 @@
                 ],
                 "registries": {
                     "af-south-1": "626614931356",
-                    "il-central-1": "780543022126",
                     "ap-east-1": "871362719292",
                     "ap-northeast-1": "763104351884",
                     "ap-northeast-2": "763104351884",
@@ -554,18 +556,20 @@
                     "ap-southeast-3": "907027046896",
                     "ap-southeast-4": "457447274322",
                     "ca-central-1": "763104351884",
+                    "ca-west-1": "204538143572",
                     "cn-north-1": "727897471807",
                     "cn-northwest-1": "727897471807",
                     "eu-central-1": "763104351884",
                     "eu-central-2": "380420809688",
                     "eu-north-1": "763104351884",
+                    "eu-south-1": "692866216735",
+                    "eu-south-2": "503227376785",
                     "eu-west-1": "763104351884",
                     "eu-west-2": "763104351884",
                     "eu-west-3": "763104351884",
-                    "eu-south-1": "692866216735",
-                    "eu-south-2": "503227376785",
-                    "me-south-1": "217643126080",
+                    "il-central-1": "780543022126",
                     "me-central-1": "914824155844",
+                    "me-south-1": "217643126080",
                     "sa-east-1": "763104351884",
                     "us-east-1": "763104351884",
                     "us-east-2": "763104351884",
@@ -574,8 +578,7 @@
                     "us-iso-east-1": "886529160074",
                     "us-isob-east-1": "094389454867",
                     "us-west-1": "763104351884",
-                    "us-west-2": "763104351884",
-                    "ca-west-1": "204538143572"
+                    "us-west-2": "763104351884"
                 },
                 "repository": "pytorch-inference"
             },
@@ -585,7 +588,6 @@
                 ],
                 "registries": {
                     "af-south-1": "626614931356",
-                    "il-central-1": "780543022126",
                     "ap-east-1": "871362719292",
                     "ap-northeast-1": "763104351884",
                     "ap-northeast-2": "763104351884",
@@ -597,18 +599,20 @@
                     "ap-southeast-3": "907027046896",
                     "ap-southeast-4": "457447274322",
                     "ca-central-1": "763104351884",
+                    "ca-west-1": "204538143572",
                     "cn-north-1": "727897471807",
                     "cn-northwest-1": "727897471807",
                     "eu-central-1": "763104351884",
                     "eu-central-2": "380420809688",
                     "eu-north-1": "763104351884",
+                    "eu-south-1": "692866216735",
+                    "eu-south-2": "503227376785",
                     "eu-west-1": "763104351884",
                     "eu-west-2": "763104351884",
                     "eu-west-3": "763104351884",
-                    "eu-south-1": "692866216735",
-                    "eu-south-2": "503227376785",
-                    "me-south-1": "217643126080",
+                    "il-central-1": "780543022126",
                     "me-central-1": "914824155844",
+                    "me-south-1": "217643126080",
                     "sa-east-1": "763104351884",
                     "us-east-1": "763104351884",
                     "us-east-2": "763104351884",
@@ -617,8 +621,7 @@
                     "us-iso-east-1": "886529160074",
                     "us-isob-east-1": "094389454867",
                     "us-west-1": "763104351884",
-                    "us-west-2": "763104351884",
-                    "ca-west-1": "204538143572"
+                    "us-west-2": "763104351884"
                 },
                 "repository": "pytorch-inference"
             },
@@ -628,7 +631,6 @@
                 ],
                 "registries": {
                     "af-south-1": "626614931356",
-                    "il-central-1": "780543022126",
                     "ap-east-1": "871362719292",
                     "ap-northeast-1": "763104351884",
                     "ap-northeast-2": "763104351884",
@@ -640,18 +642,20 @@
                     "ap-southeast-3": "907027046896",
                     "ap-southeast-4": "457447274322",
                     "ca-central-1": "763104351884",
+                    "ca-west-1": "204538143572",
                     "cn-north-1": "727897471807",
                     "cn-northwest-1": "727897471807",
                     "eu-central-1": "763104351884",
                     "eu-central-2": "380420809688",
                     "eu-north-1": "763104351884",
+                    "eu-south-1": "692866216735",
+                    "eu-south-2": "503227376785",
                     "eu-west-1": "763104351884",
                     "eu-west-2": "763104351884",
                     "eu-west-3": "763104351884",
-                    "eu-south-1": "692866216735",
-                    "eu-south-2": "503227376785",
-                    "me-south-1": "217643126080",
+                    "il-central-1": "780543022126",
                     "me-central-1": "914824155844",
+                    "me-south-1": "217643126080",
                     "sa-east-1": "763104351884",
                     "us-east-1": "763104351884",
                     "us-east-2": "763104351884",
@@ -660,8 +664,7 @@
                     "us-iso-east-1": "886529160074",
                     "us-isob-east-1": "094389454867",
                     "us-west-1": "763104351884",
-                    "us-west-2": "763104351884",
-                    "ca-west-1": "204538143572"
+                    "us-west-2": "763104351884"
                 },
                 "repository": "pytorch-inference"
             },
@@ -671,7 +674,6 @@
                 ],
                 "registries": {
                     "af-south-1": "626614931356",
-                    "il-central-1": "780543022126",
                     "ap-east-1": "871362719292",
                     "ap-northeast-1": "763104351884",
                     "ap-northeast-2": "763104351884",
@@ -683,18 +685,20 @@
                     "ap-southeast-3": "907027046896",
                     "ap-southeast-4": "457447274322",
                     "ca-central-1": "763104351884",
+                    "ca-west-1": "204538143572",
                     "cn-north-1": "727897471807",
                     "cn-northwest-1": "727897471807",
                     "eu-central-1": "763104351884",
                     "eu-central-2": "380420809688",
                     "eu-north-1": "763104351884",
+                    "eu-south-1": "692866216735",
+                    "eu-south-2": "503227376785",
                     "eu-west-1": "763104351884",
                     "eu-west-2": "763104351884",
                     "eu-west-3": "763104351884",
-                    "eu-south-1": "692866216735",
-                    "eu-south-2": "503227376785",
-                    "me-south-1": "217643126080",
+                    "il-central-1": "780543022126",
                     "me-central-1": "914824155844",
+                    "me-south-1": "217643126080",
                     "sa-east-1": "763104351884",
                     "us-east-1": "763104351884",
                     "us-east-2": "763104351884",
@@ -703,8 +707,7 @@
                     "us-iso-east-1": "886529160074",
                     "us-isob-east-1": "094389454867",
                     "us-west-1": "763104351884",
-                    "us-west-2": "763104351884",
-                    "ca-west-1": "204538143572"
+                    "us-west-2": "763104351884"
                 },
                 "repository": "pytorch-inference"
             },
@@ -714,7 +717,6 @@
                 ],
                 "registries": {
                     "af-south-1": "626614931356",
-                    "il-central-1": "780543022126",
                     "ap-east-1": "871362719292",
                     "ap-northeast-1": "763104351884",
                     "ap-northeast-2": "763104351884",
@@ -726,18 +728,20 @@
                     "ap-southeast-3": "907027046896",
                     "ap-southeast-4": "457447274322",
                     "ca-central-1": "763104351884",
+                    "ca-west-1": "204538143572",
                     "cn-north-1": "727897471807",
                     "cn-northwest-1": "727897471807",
                     "eu-central-1": "763104351884",
                     "eu-central-2": "380420809688",
                     "eu-north-1": "763104351884",
+                    "eu-south-1": "692866216735",
+                    "eu-south-2": "503227376785",
                     "eu-west-1": "763104351884",
                     "eu-west-2": "763104351884",
                     "eu-west-3": "763104351884",
-                    "eu-south-1": "692866216735",
-                    "eu-south-2": "503227376785",
-                    "me-south-1": "217643126080",
+                    "il-central-1": "780543022126",
                     "me-central-1": "914824155844",
+                    "me-south-1": "217643126080",
                     "sa-east-1": "763104351884",
                     "us-east-1": "763104351884",
                     "us-east-2": "763104351884",
@@ -746,8 +750,7 @@
                     "us-iso-east-1": "886529160074",
                     "us-isob-east-1": "094389454867",
                     "us-west-1": "763104351884",
-                    "us-west-2": "763104351884",
-                    "ca-west-1": "204538143572"
+                    "us-west-2": "763104351884"
                 },
                 "repository": "pytorch-inference"
             },
@@ -757,7 +760,6 @@
                 ],
                 "registries": {
                     "af-south-1": "626614931356",
-                    "il-central-1": "780543022126",
                     "ap-east-1": "871362719292",
                     "ap-northeast-1": "763104351884",
                     "ap-northeast-2": "763104351884",
@@ -769,18 +771,20 @@
                     "ap-southeast-3": "907027046896",
                     "ap-southeast-4": "457447274322",
                     "ca-central-1": "763104351884",
+                    "ca-west-1": "204538143572",
                     "cn-north-1": "727897471807",
                     "cn-northwest-1": "727897471807",
                     "eu-central-1": "763104351884",
                     "eu-central-2": "380420809688",
                     "eu-north-1": "763104351884",
+                    "eu-south-1": "692866216735",
+                    "eu-south-2": "503227376785",
                     "eu-west-1": "763104351884",
                     "eu-west-2": "763104351884",
                     "eu-west-3": "763104351884",
-                    "eu-south-1": "692866216735",
-                    "eu-south-2": "503227376785",
-                    "me-south-1": "217643126080",
+                    "il-central-1": "780543022126",
                     "me-central-1": "914824155844",
+                    "me-south-1": "217643126080",
                     "sa-east-1": "763104351884",
                     "us-east-1": "763104351884",
                     "us-east-2": "763104351884",
@@ -789,8 +793,7 @@
                     "us-iso-east-1": "886529160074",
                     "us-isob-east-1": "094389454867",
                     "us-west-1": "763104351884",
-                    "us-west-2": "763104351884",
-                    "ca-west-1": "204538143572"
+                    "us-west-2": "763104351884"
                 },
                 "repository": "pytorch-inference"
             },
@@ -800,7 +803,6 @@
                 ],
                 "registries": {
                     "af-south-1": "626614931356",
-                    "il-central-1": "780543022126",
                     "ap-east-1": "871362719292",
                     "ap-northeast-1": "763104351884",
                     "ap-northeast-2": "763104351884",
@@ -812,16 +814,19 @@
                     "ap-southeast-3": "907027046896",
                     "ap-southeast-4": "457447274322",
                     "ca-central-1": "763104351884",
+                    "ca-west-1": "204538143572",
                     "cn-north-1": "727897471807",
                     "cn-northwest-1": "727897471807",
                     "eu-central-1": "763104351884",
                     "eu-central-2": "380420809688",
                     "eu-north-1": "763104351884",
+                    "eu-south-1": "692866216735",
+                    "eu-south-2": "503227376785",
                     "eu-west-1": "763104351884",
                     "eu-west-2": "763104351884",
                     "eu-west-3": "763104351884",
-                    "eu-south-1": "692866216735",
-                    "eu-south-2": "503227376785",
+                    "il-central-1": "780543022126",
+                    "me-central-1": "914824155844",
                     "me-south-1": "217643126080",
                     "sa-east-1": "763104351884",
                     "us-east-1": "763104351884",
@@ -831,8 +836,7 @@
                     "us-iso-east-1": "886529160074",
                     "us-isob-east-1": "094389454867",
                     "us-west-1": "763104351884",
-                    "us-west-2": "763104351884",
-                    "ca-west-1": "204538143572"
+                    "us-west-2": "763104351884"
                 },
                 "repository": "pytorch-inference"
             },
@@ -842,7 +846,6 @@
                 ],
                 "registries": {
                     "af-south-1": "626614931356",
-                    "il-central-1": "780543022126",
                     "ap-east-1": "871362719292",
                     "ap-northeast-1": "763104351884",
                     "ap-northeast-2": "763104351884",
@@ -854,16 +857,19 @@
                     "ap-southeast-3": "907027046896",
                     "ap-southeast-4": "457447274322",
                     "ca-central-1": "763104351884",
+                    "ca-west-1": "204538143572",
                     "cn-north-1": "727897471807",
                     "cn-northwest-1": "727897471807",
                     "eu-central-1": "763104351884",
                     "eu-central-2": "380420809688",
                     "eu-north-1": "763104351884",
+                    "eu-south-1": "692866216735",
+                    "eu-south-2": "503227376785",
                     "eu-west-1": "763104351884",
                     "eu-west-2": "763104351884",
                     "eu-west-3": "763104351884",
-                    "eu-south-1": "692866216735",
-                    "eu-south-2": "503227376785",
+                    "il-central-1": "780543022126",
+                    "me-central-1": "914824155844",
                     "me-south-1": "217643126080",
                     "sa-east-1": "763104351884",
                     "us-east-1": "763104351884",
@@ -873,8 +879,7 @@
                     "us-iso-east-1": "886529160074",
                     "us-isob-east-1": "094389454867",
                     "us-west-1": "763104351884",
-                    "us-west-2": "763104351884",
-                    "ca-west-1": "204538143572"
+                    "us-west-2": "763104351884"
                 },
                 "repository": "pytorch-inference"
             },
@@ -884,7 +889,6 @@
                 ],
                 "registries": {
                     "af-south-1": "626614931356",
-                    "il-central-1": "780543022126",
                     "ap-east-1": "871362719292",
                     "ap-northeast-1": "763104351884",
                     "ap-northeast-2": "763104351884",
@@ -896,16 +900,19 @@
                     "ap-southeast-3": "907027046896",
                     "ap-southeast-4": "457447274322",
                     "ca-central-1": "763104351884",
+                    "ca-west-1": "204538143572",
                     "cn-north-1": "727897471807",
                     "cn-northwest-1": "727897471807",
                     "eu-central-1": "763104351884",
                     "eu-central-2": "380420809688",
                     "eu-north-1": "763104351884",
+                    "eu-south-1": "692866216735",
+                    "eu-south-2": "503227376785",
                     "eu-west-1": "763104351884",
                     "eu-west-2": "763104351884",
                     "eu-west-3": "763104351884",
-                    "eu-south-1": "692866216735",
-                    "eu-south-2": "503227376785",
+                    "il-central-1": "780543022126",
+                    "me-central-1": "914824155844",
                     "me-south-1": "217643126080",
                     "sa-east-1": "763104351884",
                     "us-east-1": "763104351884",
@@ -915,8 +922,7 @@
                     "us-iso-east-1": "886529160074",
                     "us-isob-east-1": "094389454867",
                     "us-west-1": "763104351884",
-                    "us-west-2": "763104351884",
-                    "ca-west-1": "204538143572"
+                    "us-west-2": "763104351884"
                 },
                 "repository": "pytorch-inference"
             },
@@ -926,7 +932,6 @@
                 ],
                 "registries": {
                     "af-south-1": "626614931356",
-                    "il-central-1": "780543022126",
                     "ap-east-1": "871362719292",
                     "ap-northeast-1": "763104351884",
                     "ap-northeast-2": "763104351884",
@@ -938,16 +943,19 @@
                     "ap-southeast-3": "907027046896",
                     "ap-southeast-4": "457447274322",
                     "ca-central-1": "763104351884",
+                    "ca-west-1": "204538143572",
                     "cn-north-1": "727897471807",
                     "cn-northwest-1": "727897471807",
                     "eu-central-1": "763104351884",
                     "eu-central-2": "380420809688",
                     "eu-north-1": "763104351884",
+                    "eu-south-1": "692866216735",
+                    "eu-south-2": "503227376785",
                     "eu-west-1": "763104351884",
                     "eu-west-2": "763104351884",
                     "eu-west-3": "763104351884",
-                    "eu-south-1": "692866216735",
-                    "eu-south-2": "503227376785",
+                    "il-central-1": "780543022126",
+                    "me-central-1": "914824155844",
                     "me-south-1": "217643126080",
                     "sa-east-1": "763104351884",
                     "us-east-1": "763104351884",
@@ -957,8 +965,7 @@
                     "us-iso-east-1": "886529160074",
                     "us-isob-east-1": "094389454867",
                     "us-west-1": "763104351884",
-                    "us-west-2": "763104351884",
-                    "ca-west-1": "204538143572"
+                    "us-west-2": "763104351884"
                 },
                 "repository": "pytorch-inference"
             },
@@ -968,7 +975,6 @@
                 ],
                 "registries": {
                     "af-south-1": "626614931356",
-                    "il-central-1": "780543022126",
                     "ap-east-1": "871362719292",
                     "ap-northeast-1": "763104351884",
                     "ap-northeast-2": "763104351884",
@@ -980,16 +986,19 @@
                     "ap-southeast-3": "907027046896",
                     "ap-southeast-4": "457447274322",
                     "ca-central-1": "763104351884",
+                    "ca-west-1": "204538143572",
                     "cn-north-1": "727897471807",
                     "cn-northwest-1": "727897471807",
                     "eu-central-1": "763104351884",
                     "eu-central-2": "380420809688",
                     "eu-north-1": "763104351884",
+                    "eu-south-1": "692866216735",
+                    "eu-south-2": "503227376785",
                     "eu-west-1": "763104351884",
                     "eu-west-2": "763104351884",
                     "eu-west-3": "763104351884",
-                    "eu-south-1": "692866216735",
-                    "eu-south-2": "503227376785",
+                    "il-central-1": "780543022126",
+                    "me-central-1": "914824155844",
                     "me-south-1": "217643126080",
                     "sa-east-1": "763104351884",
                     "us-east-1": "763104351884",
@@ -999,8 +1008,7 @@
                     "us-iso-east-1": "886529160074",
                     "us-isob-east-1": "094389454867",
                     "us-west-1": "763104351884",
-                    "us-west-2": "763104351884",
-                    "ca-west-1": "204538143572"
+                    "us-west-2": "763104351884"
                 },
                 "repository": "pytorch-inference"
             },
@@ -1010,7 +1018,6 @@
                 ],
                 "registries": {
                     "af-south-1": "626614931356",
-                    "il-central-1": "780543022126",
                     "ap-east-1": "871362719292",
                     "ap-northeast-1": "763104351884",
                     "ap-northeast-2": "763104351884",
@@ -1022,16 +1029,19 @@
                     "ap-southeast-3": "907027046896",
                     "ap-southeast-4": "457447274322",
                     "ca-central-1": "763104351884",
+                    "ca-west-1": "204538143572",
                     "cn-north-1": "727897471807",
                     "cn-northwest-1": "727897471807",
                     "eu-central-1": "763104351884",
                     "eu-central-2": "380420809688",
                     "eu-north-1": "763104351884",
+                    "eu-south-1": "692866216735",
+                    "eu-south-2": "503227376785",
                     "eu-west-1": "763104351884",
                     "eu-west-2": "763104351884",
                     "eu-west-3": "763104351884",
-                    "eu-south-1": "692866216735",
-                    "eu-south-2": "503227376785",
+                    "il-central-1": "780543022126",
+                    "me-central-1": "914824155844",
                     "me-south-1": "217643126080",
                     "sa-east-1": "763104351884",
                     "us-east-1": "763104351884",
@@ -1041,8 +1051,7 @@
                     "us-iso-east-1": "886529160074",
                     "us-isob-east-1": "094389454867",
                     "us-west-1": "763104351884",
-                    "us-west-2": "763104351884",
-                    "ca-west-1": "204538143572"
+                    "us-west-2": "763104351884"
                 },
                 "repository": "pytorch-inference"
             }
@@ -1060,12 +1069,14 @@
         },
         "versions": {
             "1.12.1": {
+                "container_version": {
+                    "cpu": "ubuntu20.04"
+                },
                 "py_versions": [
                     "py38"
                 ],
                 "registries": {
                     "af-south-1": "626614931356",
-                    "il-central-1": "780543022126",
                     "ap-east-1": "871362719292",
                     "ap-northeast-1": "763104351884",
                     "ap-northeast-2": "763104351884",
@@ -1077,16 +1088,19 @@
                     "ap-southeast-3": "907027046896",
                     "ap-southeast-4": "457447274322",
                     "ca-central-1": "763104351884",
+                    "ca-west-1": "204538143572",
                     "cn-north-1": "727897471807",
                     "cn-northwest-1": "727897471807",
                     "eu-central-1": "763104351884",
                     "eu-central-2": "380420809688",
                     "eu-north-1": "763104351884",
+                    "eu-south-1": "692866216735",
+                    "eu-south-2": "503227376785",
                     "eu-west-1": "763104351884",
                     "eu-west-2": "763104351884",
                     "eu-west-3": "763104351884",
-                    "eu-south-1": "692866216735",
-                    "eu-south-2": "503227376785",
+                    "il-central-1": "780543022126",
+                    "me-central-1": "914824155844",
                     "me-south-1": "217643126080",
                     "sa-east-1": "763104351884",
                     "us-east-1": "763104351884",
@@ -1096,21 +1110,19 @@
                     "us-iso-east-1": "886529160074",
                     "us-isob-east-1": "094389454867",
                     "us-west-1": "763104351884",
-                    "us-west-2": "763104351884",
-                    "ca-west-1": "204538143572"
+                    "us-west-2": "763104351884"
                 },
-                "repository": "pytorch-inference-graviton",
-                "container_version": {
-                    "cpu": "ubuntu20.04"
-                }
+                "repository": "pytorch-inference-graviton"
             },
             "2.0.0": {
+                "container_version": {
+                    "cpu": "ubuntu20.04"
+                },
                 "py_versions": [
                     "py310"
                 ],
                 "registries": {
                     "af-south-1": "626614931356",
-                    "il-central-1": "780543022126",
                     "ap-east-1": "871362719292",
                     "ap-northeast-1": "763104351884",
                     "ap-northeast-2": "763104351884",
@@ -1122,34 +1134,39 @@
                     "ap-southeast-3": "907027046896",
                     "ap-southeast-4": "457447274322",
                     "ca-central-1": "763104351884",
+                    "ca-west-1": "204538143572",
+                    "cn-north-1": "727897471807",
+                    "cn-northwest-1": "727897471807",
                     "eu-central-1": "763104351884",
                     "eu-central-2": "380420809688",
                     "eu-north-1": "763104351884",
+                    "eu-south-1": "692866216735",
+                    "eu-south-2": "503227376785",
                     "eu-west-1": "763104351884",
                     "eu-west-2": "763104351884",
                     "eu-west-3": "763104351884",
-                    "eu-south-1": "692866216735",
-                    "eu-south-2": "503227376785",
+                    "il-central-1": "780543022126",
+                    "me-central-1": "914824155844",
                     "me-south-1": "217643126080",
                     "sa-east-1": "763104351884",
                     "us-east-1": "763104351884",
                     "us-east-2": "763104351884",
+                    "us-gov-east-1": "446045086412",
+                    "us-gov-west-1": "442386744353",
                     "us-west-1": "763104351884",
-                    "us-west-2": "763104351884",
-                    "ca-west-1": "204538143572"
+                    "us-west-2": "763104351884"
                 },
-                "repository": "pytorch-inference-graviton",
-                "container_version": {
-                    "cpu": "ubuntu20.04"
-                }
+                "repository": "pytorch-inference-graviton"
             },
             "2.0.1": {
+                "container_version": {
+                    "cpu": "ubuntu20.04"
+                },
                 "py_versions": [
                     "py310"
                 ],
                 "registries": {
                     "af-south-1": "626614931356",
-                    "il-central-1": "780543022126",
                     "ap-east-1": "871362719292",
                     "ap-northeast-1": "763104351884",
                     "ap-northeast-2": "763104351884",
@@ -1161,34 +1178,39 @@
                     "ap-southeast-3": "907027046896",
                     "ap-southeast-4": "457447274322",
                     "ca-central-1": "763104351884",
+                    "ca-west-1": "204538143572",
+                    "cn-north-1": "727897471807",
+                    "cn-northwest-1": "727897471807",
                     "eu-central-1": "763104351884",
                     "eu-central-2": "380420809688",
                     "eu-north-1": "763104351884",
+                    "eu-south-1": "692866216735",
+                    "eu-south-2": "503227376785",
                     "eu-west-1": "763104351884",
                     "eu-west-2": "763104351884",
                     "eu-west-3": "763104351884",
-                    "eu-south-1": "692866216735",
-                    "eu-south-2": "503227376785",
+                    "il-central-1": "780543022126",
+                    "me-central-1": "914824155844",
                     "me-south-1": "217643126080",
                     "sa-east-1": "763104351884",
                     "us-east-1": "763104351884",
                     "us-east-2": "763104351884",
+                    "us-gov-east-1": "446045086412",
+                    "us-gov-west-1": "442386744353",
                     "us-west-1": "763104351884",
-                    "us-west-2": "763104351884",
-                    "ca-west-1": "204538143572"
+                    "us-west-2": "763104351884"
                 },
-                "repository": "pytorch-inference-graviton",
-                "container_version": {
-                    "cpu": "ubuntu20.04"
-                }
+                "repository": "pytorch-inference-graviton"
             },
             "2.1.0": {
+                "container_version": {
+                    "cpu": "ubuntu20.04"
+                },
                 "py_versions": [
                     "py310"
                 ],
                 "registries": {
                     "af-south-1": "626614931356",
-                    "il-central-1": "780543022126",
                     "ap-east-1": "871362719292",
                     "ap-northeast-1": "763104351884",
                     "ap-northeast-2": "763104351884",
@@ -1200,34 +1222,39 @@
                     "ap-southeast-3": "907027046896",
                     "ap-southeast-4": "457447274322",
                     "ca-central-1": "763104351884",
+                    "ca-west-1": "204538143572",
+                    "cn-north-1": "727897471807",
+                    "cn-northwest-1": "727897471807",
                     "eu-central-1": "763104351884",
                     "eu-central-2": "380420809688",
                     "eu-north-1": "763104351884",
+                    "eu-south-1": "692866216735",
+                    "eu-south-2": "503227376785",
                     "eu-west-1": "763104351884",
                     "eu-west-2": "763104351884",
                     "eu-west-3": "763104351884",
-                    "eu-south-1": "692866216735",
-                    "eu-south-2": "503227376785",
+                    "il-central-1": "780543022126",
+                    "me-central-1": "914824155844",
                     "me-south-1": "217643126080",
                     "sa-east-1": "763104351884",
                     "us-east-1": "763104351884",
                     "us-east-2": "763104351884",
+                    "us-gov-east-1": "446045086412",
+                    "us-gov-west-1": "442386744353",
                     "us-west-1": "763104351884",
-                    "us-west-2": "763104351884",
-                    "ca-west-1": "204538143572"
+                    "us-west-2": "763104351884"
                 },
-                "repository": "pytorch-inference-graviton",
-                "container_version": {
-                    "cpu": "ubuntu20.04"
-                }
+                "repository": "pytorch-inference-graviton"
             },
             "2.2.1": {
+                "container_version": {
+                    "cpu": "ubuntu20.04"
+                },
                 "py_versions": [
                     "py310"
                 ],
                 "registries": {
                     "af-south-1": "626614931356",
-                    "il-central-1": "780543022126",
                     "ap-east-1": "871362719292",
                     "ap-northeast-1": "763104351884",
                     "ap-northeast-2": "763104351884",
@@ -1239,16 +1266,19 @@
                     "ap-southeast-3": "907027046896",
                     "ap-southeast-4": "457447274322",
                     "ca-central-1": "763104351884",
+                    "ca-west-1": "204538143572",
                     "cn-north-1": "727897471807",
                     "cn-northwest-1": "727897471807",
                     "eu-central-1": "763104351884",
                     "eu-central-2": "380420809688",
                     "eu-north-1": "763104351884",
+                    "eu-south-1": "692866216735",
+                    "eu-south-2": "503227376785",
                     "eu-west-1": "763104351884",
                     "eu-west-2": "763104351884",
                     "eu-west-3": "763104351884",
-                    "eu-south-1": "692866216735",
-                    "eu-south-2": "503227376785",
+                    "il-central-1": "780543022126",
+                    "me-central-1": "914824155844",
                     "me-south-1": "217643126080",
                     "sa-east-1": "763104351884",
                     "us-east-1": "763104351884",
@@ -1256,13 +1286,9 @@
                     "us-gov-east-1": "446045086412",
                     "us-gov-west-1": "442386744353",
                     "us-west-1": "763104351884",
-                    "us-west-2": "763104351884",
-                    "ca-west-1": "204538143572"
+                    "us-west-2": "763104351884"
                 },
-                "repository": "pytorch-inference-graviton",
-                "container_version": {
-                    "cpu": "ubuntu20.04"
-                }
+                "repository": "pytorch-inference-graviton"
             }
         }
     },
@@ -1401,7 +1427,6 @@
                 ],
                 "registries": {
                     "af-south-1": "626614931356",
-                    "il-central-1": "780543022126",
                     "ap-east-1": "871362719292",
                     "ap-northeast-1": "763104351884",
                     "ap-northeast-2": "763104351884",
@@ -1413,18 +1438,20 @@
                     "ap-southeast-3": "907027046896",
                     "ap-southeast-4": "457447274322",
                     "ca-central-1": "763104351884",
+                    "ca-west-1": "204538143572",
                     "cn-north-1": "727897471807",
                     "cn-northwest-1": "727897471807",
                     "eu-central-1": "763104351884",
                     "eu-central-2": "380420809688",
                     "eu-north-1": "763104351884",
+                    "eu-south-1": "692866216735",
+                    "eu-south-2": "503227376785",
                     "eu-west-1": "763104351884",
                     "eu-west-2": "763104351884",
                     "eu-west-3": "763104351884",
-                    "eu-south-1": "692866216735",
-                    "eu-south-2": "503227376785",
-                    "me-south-1": "217643126080",
+                    "il-central-1": "780543022126",
                     "me-central-1": "914824155844",
+                    "me-south-1": "217643126080",
                     "sa-east-1": "763104351884",
                     "us-east-1": "763104351884",
                     "us-east-2": "763104351884",
@@ -1433,8 +1460,7 @@
                     "us-iso-east-1": "886529160074",
                     "us-isob-east-1": "094389454867",
                     "us-west-1": "763104351884",
-                    "us-west-2": "763104351884",
-                    "ca-west-1": "204538143572"
+                    "us-west-2": "763104351884"
                 },
                 "repository": "pytorch-training"
             },
@@ -1445,7 +1471,6 @@
                 ],
                 "registries": {
                     "af-south-1": "626614931356",
-                    "il-central-1": "780543022126",
                     "ap-east-1": "871362719292",
                     "ap-northeast-1": "763104351884",
                     "ap-northeast-2": "763104351884",
@@ -1457,18 +1482,20 @@
                     "ap-southeast-3": "907027046896",
                     "ap-southeast-4": "457447274322",
                     "ca-central-1": "763104351884",
+                    "ca-west-1": "204538143572",
                     "cn-north-1": "727897471807",
                     "cn-northwest-1": "727897471807",
                     "eu-central-1": "763104351884",
                     "eu-central-2": "380420809688",
                     "eu-north-1": "763104351884",
+                    "eu-south-1": "692866216735",
+                    "eu-south-2": "503227376785",
                     "eu-west-1": "763104351884",
                     "eu-west-2": "763104351884",
                     "eu-west-3": "763104351884",
-                    "eu-south-1": "692866216735",
-                    "eu-south-2": "503227376785",
-                    "me-south-1": "217643126080",
+                    "il-central-1": "780543022126",
                     "me-central-1": "914824155844",
+                    "me-south-1": "217643126080",
                     "sa-east-1": "763104351884",
                     "us-east-1": "763104351884",
                     "us-east-2": "763104351884",
@@ -1477,19 +1504,18 @@
                     "us-iso-east-1": "886529160074",
                     "us-isob-east-1": "094389454867",
                     "us-west-1": "763104351884",
-                    "us-west-2": "763104351884",
-                    "ca-west-1": "204538143572"
+                    "us-west-2": "763104351884"
                 },
                 "repository": "pytorch-training"
             },
             "1.4.0": {
                 "py_versions": [
                     "py2",
-                    "py3"
+                    "py3",
+                    "py36"
                 ],
                 "registries": {
                     "af-south-1": "626614931356",
-                    "il-central-1": "780543022126",
                     "ap-east-1": "871362719292",
                     "ap-northeast-1": "763104351884",
                     "ap-northeast-2": "763104351884",
@@ -1501,18 +1527,20 @@
                     "ap-southeast-3": "907027046896",
                     "ap-southeast-4": "457447274322",
                     "ca-central-1": "763104351884",
+                    "ca-west-1": "204538143572",
                     "cn-north-1": "727897471807",
                     "cn-northwest-1": "727897471807",
                     "eu-central-1": "763104351884",
                     "eu-central-2": "380420809688",
                     "eu-north-1": "763104351884",
+                    "eu-south-1": "692866216735",
+                    "eu-south-2": "503227376785",
                     "eu-west-1": "763104351884",
                     "eu-west-2": "763104351884",
                     "eu-west-3": "763104351884",
-                    "eu-south-1": "692866216735",
-                    "eu-south-2": "503227376785",
-                    "me-south-1": "217643126080",
+                    "il-central-1": "780543022126",
                     "me-central-1": "914824155844",
+                    "me-south-1": "217643126080",
                     "sa-east-1": "763104351884",
                     "us-east-1": "763104351884",
                     "us-east-2": "763104351884",
@@ -1521,18 +1549,17 @@
                     "us-iso-east-1": "886529160074",
                     "us-isob-east-1": "094389454867",
                     "us-west-1": "763104351884",
-                    "us-west-2": "763104351884",
-                    "ca-west-1": "204538143572"
+                    "us-west-2": "763104351884"
                 },
                 "repository": "pytorch-training"
             },
             "1.5.0": {
                 "py_versions": [
-                    "py3"
+                    "py3",
+                    "py36"
                 ],
                 "registries": {
                     "af-south-1": "626614931356",
-                    "il-central-1": "780543022126",
                     "ap-east-1": "871362719292",
                     "ap-northeast-1": "763104351884",
                     "ap-northeast-2": "763104351884",
@@ -1544,18 +1571,20 @@
                     "ap-southeast-3": "907027046896",
                     "ap-southeast-4": "457447274322",
                     "ca-central-1": "763104351884",
+                    "ca-west-1": "204538143572",
                     "cn-north-1": "727897471807",
                     "cn-northwest-1": "727897471807",
                     "eu-central-1": "763104351884",
                     "eu-central-2": "380420809688",
                     "eu-north-1": "763104351884",
+                    "eu-south-1": "692866216735",
+                    "eu-south-2": "503227376785",
                     "eu-west-1": "763104351884",
                     "eu-west-2": "763104351884",
                     "eu-west-3": "763104351884",
-                    "eu-south-1": "692866216735",
-                    "eu-south-2": "503227376785",
-                    "me-south-1": "217643126080",
+                    "il-central-1": "780543022126",
                     "me-central-1": "914824155844",
+                    "me-south-1": "217643126080",
                     "sa-east-1": "763104351884",
                     "us-east-1": "763104351884",
                     "us-east-2": "763104351884",
@@ -1564,8 +1593,7 @@
                     "us-iso-east-1": "886529160074",
                     "us-isob-east-1": "094389454867",
                     "us-west-1": "763104351884",
-                    "us-west-2": "763104351884",
-                    "ca-west-1": "204538143572"
+                    "us-west-2": "763104351884"
                 },
                 "repository": "pytorch-training"
             },
@@ -1576,7 +1604,6 @@
                 ],
                 "registries": {
                     "af-south-1": "626614931356",
-                    "il-central-1": "780543022126",
                     "ap-east-1": "871362719292",
                     "ap-northeast-1": "763104351884",
                     "ap-northeast-2": "763104351884",
@@ -1588,18 +1615,20 @@
                     "ap-southeast-3": "907027046896",
                     "ap-southeast-4": "457447274322",
                     "ca-central-1": "763104351884",
+                    "ca-west-1": "204538143572",
                     "cn-north-1": "727897471807",
                     "cn-northwest-1": "727897471807",
                     "eu-central-1": "763104351884",
                     "eu-central-2": "380420809688",
                     "eu-north-1": "763104351884",
+                    "eu-south-1": "692866216735",
+                    "eu-south-2": "503227376785",
                     "eu-west-1": "763104351884",
                     "eu-west-2": "763104351884",
                     "eu-west-3": "763104351884",
-                    "eu-south-1": "692866216735",
-                    "eu-south-2": "503227376785",
-                    "me-south-1": "217643126080",
+                    "il-central-1": "780543022126",
                     "me-central-1": "914824155844",
+                    "me-south-1": "217643126080",
                     "sa-east-1": "763104351884",
                     "us-east-1": "763104351884",
                     "us-east-2": "763104351884",
@@ -1608,8 +1637,7 @@
                     "us-iso-east-1": "886529160074",
                     "us-isob-east-1": "094389454867",
                     "us-west-1": "763104351884",
-                    "us-west-2": "763104351884",
-                    "ca-west-1": "204538143572"
+                    "us-west-2": "763104351884"
                 },
                 "repository": "pytorch-training"
             },
@@ -1620,7 +1648,6 @@
                 ],
                 "registries": {
                     "af-south-1": "626614931356",
-                    "il-central-1": "780543022126",
                     "ap-east-1": "871362719292",
                     "ap-northeast-1": "763104351884",
                     "ap-northeast-2": "763104351884",
@@ -1632,18 +1659,20 @@
                     "ap-southeast-3": "907027046896",
                     "ap-southeast-4": "457447274322",
                     "ca-central-1": "763104351884",
+                    "ca-west-1": "204538143572",
                     "cn-north-1": "727897471807",
                     "cn-northwest-1": "727897471807",
                     "eu-central-1": "763104351884",
                     "eu-central-2": "380420809688",
                     "eu-north-1": "763104351884",
+                    "eu-south-1": "692866216735",
+                    "eu-south-2": "503227376785",
                     "eu-west-1": "763104351884",
                     "eu-west-2": "763104351884",
                     "eu-west-3": "763104351884",
-                    "eu-south-1": "692866216735",
-                    "eu-south-2": "503227376785",
-                    "me-south-1": "217643126080",
+                    "il-central-1": "780543022126",
                     "me-central-1": "914824155844",
+                    "me-south-1": "217643126080",
                     "sa-east-1": "763104351884",
                     "us-east-1": "763104351884",
                     "us-east-2": "763104351884",
@@ -1652,8 +1681,7 @@
                     "us-iso-east-1": "886529160074",
                     "us-isob-east-1": "094389454867",
                     "us-west-1": "763104351884",
-                    "us-west-2": "763104351884",
-                    "ca-west-1": "204538143572"
+                    "us-west-2": "763104351884"
                 },
                 "repository": "pytorch-training"
             },
@@ -1664,7 +1692,6 @@
                 ],
                 "registries": {
                     "af-south-1": "626614931356",
-                    "il-central-1": "780543022126",
                     "ap-east-1": "871362719292",
                     "ap-northeast-1": "763104351884",
                     "ap-northeast-2": "763104351884",
@@ -1676,18 +1703,20 @@
                     "ap-southeast-3": "907027046896",
                     "ap-southeast-4": "457447274322",
                     "ca-central-1": "763104351884",
+                    "ca-west-1": "204538143572",
                     "cn-north-1": "727897471807",
                     "cn-northwest-1": "727897471807",
                     "eu-central-1": "763104351884",
                     "eu-central-2": "380420809688",
                     "eu-north-1": "763104351884",
+                    "eu-south-1": "692866216735",
+                    "eu-south-2": "503227376785",
                     "eu-west-1": "763104351884",
                     "eu-west-2": "763104351884",
                     "eu-west-3": "763104351884",
-                    "eu-south-1": "692866216735",
-                    "eu-south-2": "503227376785",
-                    "me-south-1": "217643126080",
+                    "il-central-1": "780543022126",
                     "me-central-1": "914824155844",
+                    "me-south-1": "217643126080",
                     "sa-east-1": "763104351884",
                     "us-east-1": "763104351884",
                     "us-east-2": "763104351884",
@@ -1696,8 +1725,7 @@
                     "us-iso-east-1": "886529160074",
                     "us-isob-east-1": "094389454867",
                     "us-west-1": "763104351884",
-                    "us-west-2": "763104351884",
-                    "ca-west-1": "204538143572"
+                    "us-west-2": "763104351884"
                 },
                 "repository": "pytorch-training"
             },
@@ -1708,7 +1736,6 @@
                 ],
                 "registries": {
                     "af-south-1": "626614931356",
-                    "il-central-1": "780543022126",
                     "ap-east-1": "871362719292",
                     "ap-northeast-1": "763104351884",
                     "ap-northeast-2": "763104351884",
@@ -1720,18 +1747,20 @@
                     "ap-southeast-3": "907027046896",
                     "ap-southeast-4": "457447274322",
                     "ca-central-1": "763104351884",
+                    "ca-west-1": "204538143572",
                     "cn-north-1": "727897471807",
                     "cn-northwest-1": "727897471807",
                     "eu-central-1": "763104351884",
                     "eu-central-2": "380420809688",
                     "eu-north-1": "763104351884",
+                    "eu-south-1": "692866216735",
+                    "eu-south-2": "503227376785",
                     "eu-west-1": "763104351884",
                     "eu-west-2": "763104351884",
                     "eu-west-3": "763104351884",
-                    "eu-south-1": "692866216735",
-                    "eu-south-2": "503227376785",
-                    "me-south-1": "217643126080",
+                    "il-central-1": "780543022126",
                     "me-central-1": "914824155844",
+                    "me-south-1": "217643126080",
                     "sa-east-1": "763104351884",
                     "us-east-1": "763104351884",
                     "us-east-2": "763104351884",
@@ -1740,8 +1769,7 @@
                     "us-iso-east-1": "886529160074",
                     "us-isob-east-1": "094389454867",
                     "us-west-1": "763104351884",
-                    "us-west-2": "763104351884",
-                    "ca-west-1": "204538143572"
+                    "us-west-2": "763104351884"
                 },
                 "repository": "pytorch-training"
             },
@@ -1751,7 +1779,6 @@
                 ],
                 "registries": {
                     "af-south-1": "626614931356",
-                    "il-central-1": "780543022126",
                     "ap-east-1": "871362719292",
                     "ap-northeast-1": "763104351884",
                     "ap-northeast-2": "763104351884",
@@ -1763,18 +1790,20 @@
                     "ap-southeast-3": "907027046896",
                     "ap-southeast-4": "457447274322",
                     "ca-central-1": "763104351884",
+                    "ca-west-1": "204538143572",
                     "cn-north-1": "727897471807",
                     "cn-northwest-1": "727897471807",
                     "eu-central-1": "763104351884",
                     "eu-central-2": "380420809688",
                     "eu-north-1": "763104351884",
+                    "eu-south-1": "692866216735",
+                    "eu-south-2": "503227376785",
                     "eu-west-1": "763104351884",
                     "eu-west-2": "763104351884",
                     "eu-west-3": "763104351884",
-                    "eu-south-1": "692866216735",
-                    "eu-south-2": "503227376785",
-                    "me-south-1": "217643126080",
+                    "il-central-1": "780543022126",
                     "me-central-1": "914824155844",
+                    "me-south-1": "217643126080",
                     "sa-east-1": "763104351884",
                     "us-east-1": "763104351884",
                     "us-east-2": "763104351884",
@@ -1783,8 +1812,7 @@
                     "us-iso-east-1": "886529160074",
                     "us-isob-east-1": "094389454867",
                     "us-west-1": "763104351884",
-                    "us-west-2": "763104351884",
-                    "ca-west-1": "204538143572"
+                    "us-west-2": "763104351884"
                 },
                 "repository": "pytorch-training"
             },
@@ -1794,7 +1822,6 @@
                 ],
                 "registries": {
                     "af-south-1": "626614931356",
-                    "il-central-1": "780543022126",
                     "ap-east-1": "871362719292",
                     "ap-northeast-1": "763104351884",
                     "ap-northeast-2": "763104351884",
@@ -1806,18 +1833,20 @@
                     "ap-southeast-3": "907027046896",
                     "ap-southeast-4": "457447274322",
                     "ca-central-1": "763104351884",
+                    "ca-west-1": "204538143572",
                     "cn-north-1": "727897471807",
                     "cn-northwest-1": "727897471807",
                     "eu-central-1": "763104351884",
                     "eu-central-2": "380420809688",
                     "eu-north-1": "763104351884",
+                    "eu-south-1": "692866216735",
+                    "eu-south-2": "503227376785",
                     "eu-west-1": "763104351884",
                     "eu-west-2": "763104351884",
                     "eu-west-3": "763104351884",
-                    "eu-south-1": "692866216735",
-                    "eu-south-2": "503227376785",
-                    "me-south-1": "217643126080",
+                    "il-central-1": "780543022126",
                     "me-central-1": "914824155844",
+                    "me-south-1": "217643126080",
                     "sa-east-1": "763104351884",
                     "us-east-1": "763104351884",
                     "us-east-2": "763104351884",
@@ -1826,8 +1855,7 @@
                     "us-iso-east-1": "886529160074",
                     "us-isob-east-1": "094389454867",
                     "us-west-1": "763104351884",
-                    "us-west-2": "763104351884",
-                    "ca-west-1": "204538143572"
+                    "us-west-2": "763104351884"
                 },
                 "repository": "pytorch-training"
             },
@@ -1837,7 +1865,6 @@
                 ],
                 "registries": {
                     "af-south-1": "626614931356",
-                    "il-central-1": "780543022126",
                     "ap-east-1": "871362719292",
                     "ap-northeast-1": "763104351884",
                     "ap-northeast-2": "763104351884",
@@ -1849,18 +1876,20 @@
                     "ap-southeast-3": "907027046896",
                     "ap-southeast-4": "457447274322",
                     "ca-central-1": "763104351884",
+                    "ca-west-1": "204538143572",
                     "cn-north-1": "727897471807",
                     "cn-northwest-1": "727897471807",
                     "eu-central-1": "763104351884",
                     "eu-central-2": "380420809688",
                     "eu-north-1": "763104351884",
+                    "eu-south-1": "692866216735",
+                    "eu-south-2": "503227376785",
                     "eu-west-1": "763104351884",
                     "eu-west-2": "763104351884",
                     "eu-west-3": "763104351884",
-                    "eu-south-1": "692866216735",
-                    "eu-south-2": "503227376785",
-                    "me-south-1": "217643126080",
+                    "il-central-1": "780543022126",
                     "me-central-1": "914824155844",
+                    "me-south-1": "217643126080",
                     "sa-east-1": "763104351884",
                     "us-east-1": "763104351884",
                     "us-east-2": "763104351884",
@@ -1869,8 +1898,7 @@
                     "us-iso-east-1": "886529160074",
                     "us-isob-east-1": "094389454867",
                     "us-west-1": "763104351884",
-                    "us-west-2": "763104351884",
-                    "ca-west-1": "204538143572"
+                    "us-west-2": "763104351884"
                 },
                 "repository": "pytorch-training"
             },
@@ -1880,7 +1908,6 @@
                 ],
                 "registries": {
                     "af-south-1": "626614931356",
-                    "il-central-1": "780543022126",
                     "ap-east-1": "871362719292",
                     "ap-northeast-1": "763104351884",
                     "ap-northeast-2": "763104351884",
@@ -1892,18 +1919,20 @@
                     "ap-southeast-3": "907027046896",
                     "ap-southeast-4": "457447274322",
                     "ca-central-1": "763104351884",
+                    "ca-west-1": "204538143572",
                     "cn-north-1": "727897471807",
                     "cn-northwest-1": "727897471807",
                     "eu-central-1": "763104351884",
                     "eu-central-2": "380420809688",
                     "eu-north-1": "763104351884",
+                    "eu-south-1": "692866216735",
+                    "eu-south-2": "503227376785",
                     "eu-west-1": "763104351884",
                     "eu-west-2": "763104351884",
                     "eu-west-3": "763104351884",
-                    "eu-south-1": "692866216735",
-                    "eu-south-2": "503227376785",
-                    "me-south-1": "217643126080",
+                    "il-central-1": "780543022126",
                     "me-central-1": "914824155844",
+                    "me-south-1": "217643126080",
                     "sa-east-1": "763104351884",
                     "us-east-1": "763104351884",
                     "us-east-2": "763104351884",
@@ -1912,8 +1941,7 @@
                     "us-iso-east-1": "886529160074",
                     "us-isob-east-1": "094389454867",
                     "us-west-1": "763104351884",
-                    "us-west-2": "763104351884",
-                    "ca-west-1": "204538143572"
+                    "us-west-2": "763104351884"
                 },
                 "repository": "pytorch-training"
             },
@@ -1923,7 +1951,6 @@
                 ],
                 "registries": {
                     "af-south-1": "626614931356",
-                    "il-central-1": "780543022126",
                     "ap-east-1": "871362719292",
                     "ap-northeast-1": "763104351884",
                     "ap-northeast-2": "763104351884",
@@ -1935,18 +1962,20 @@
                     "ap-southeast-3": "907027046896",
                     "ap-southeast-4": "457447274322",
                     "ca-central-1": "763104351884",
+                    "ca-west-1": "204538143572",
                     "cn-north-1": "727897471807",
                     "cn-northwest-1": "727897471807",
                     "eu-central-1": "763104351884",
                     "eu-central-2": "380420809688",
                     "eu-north-1": "763104351884",
+                    "eu-south-1": "692866216735",
+                    "eu-south-2": "503227376785",
                     "eu-west-1": "763104351884",
                     "eu-west-2": "763104351884",
                     "eu-west-3": "763104351884",
-                    "eu-south-1": "692866216735",
-                    "eu-south-2": "503227376785",
-                    "me-south-1": "217643126080",
+                    "il-central-1": "780543022126",
                     "me-central-1": "914824155844",
+                    "me-south-1": "217643126080",
                     "sa-east-1": "763104351884",
                     "us-east-1": "763104351884",
                     "us-east-2": "763104351884",
@@ -1955,8 +1984,7 @@
                     "us-iso-east-1": "886529160074",
                     "us-isob-east-1": "094389454867",
                     "us-west-1": "763104351884",
-                    "us-west-2": "763104351884",
-                    "ca-west-1": "204538143572"
+                    "us-west-2": "763104351884"
                 },
                 "repository": "pytorch-training"
             },
@@ -1966,7 +1994,6 @@
                 ],
                 "registries": {
                     "af-south-1": "626614931356",
-                    "il-central-1": "780543022126",
                     "ap-east-1": "871362719292",
                     "ap-northeast-1": "763104351884",
                     "ap-northeast-2": "763104351884",
@@ -1978,18 +2005,20 @@
                     "ap-southeast-3": "907027046896",
                     "ap-southeast-4": "457447274322",
                     "ca-central-1": "763104351884",
+                    "ca-west-1": "204538143572",
                     "cn-north-1": "727897471807",
                     "cn-northwest-1": "727897471807",
                     "eu-central-1": "763104351884",
                     "eu-central-2": "380420809688",
                     "eu-north-1": "763104351884",
+                    "eu-south-1": "692866216735",
+                    "eu-south-2": "503227376785",
                     "eu-west-1": "763104351884",
                     "eu-west-2": "763104351884",
                     "eu-west-3": "763104351884",
-                    "eu-south-1": "692866216735",
-                    "eu-south-2": "503227376785",
-                    "me-south-1": "217643126080",
+                    "il-central-1": "780543022126",
                     "me-central-1": "914824155844",
+                    "me-south-1": "217643126080",
                     "sa-east-1": "763104351884",
                     "us-east-1": "763104351884",
                     "us-east-2": "763104351884",
@@ -1998,8 +2027,7 @@
                     "us-iso-east-1": "886529160074",
                     "us-isob-east-1": "094389454867",
                     "us-west-1": "763104351884",
-                    "us-west-2": "763104351884",
-                    "ca-west-1": "204538143572"
+                    "us-west-2": "763104351884"
                 },
                 "repository": "pytorch-training"
             },
@@ -2009,7 +2037,6 @@
                 ],
                 "registries": {
                     "af-south-1": "626614931356",
-                    "il-central-1": "780543022126",
                     "ap-east-1": "871362719292",
                     "ap-northeast-1": "763104351884",
                     "ap-northeast-2": "763104351884",
@@ -2021,16 +2048,19 @@
                     "ap-southeast-3": "907027046896",
                     "ap-southeast-4": "457447274322",
                     "ca-central-1": "763104351884",
+                    "ca-west-1": "204538143572",
                     "cn-north-1": "727897471807",
                     "cn-northwest-1": "727897471807",
                     "eu-central-1": "763104351884",
                     "eu-central-2": "380420809688",
                     "eu-north-1": "763104351884",
+                    "eu-south-1": "692866216735",
+                    "eu-south-2": "503227376785",
                     "eu-west-1": "763104351884",
                     "eu-west-2": "763104351884",
                     "eu-west-3": "763104351884",
-                    "eu-south-1": "692866216735",
-                    "eu-south-2": "503227376785",
+                    "il-central-1": "780543022126",
+                    "me-central-1": "914824155844",
                     "me-south-1": "217643126080",
                     "sa-east-1": "763104351884",
                     "us-east-1": "763104351884",
@@ -2040,8 +2070,7 @@
                     "us-iso-east-1": "886529160074",
                     "us-isob-east-1": "094389454867",
                     "us-west-1": "763104351884",
-                    "us-west-2": "763104351884",
-                    "ca-west-1": "204538143572"
+                    "us-west-2": "763104351884"
                 },
                 "repository": "pytorch-training"
             },
@@ -2051,7 +2080,6 @@
                 ],
                 "registries": {
                     "af-south-1": "626614931356",
-                    "il-central-1": "780543022126",
                     "ap-east-1": "871362719292",
                     "ap-northeast-1": "763104351884",
                     "ap-northeast-2": "763104351884",
@@ -2063,16 +2091,19 @@
                     "ap-southeast-3": "907027046896",
                     "ap-southeast-4": "457447274322",
                     "ca-central-1": "763104351884",
+                    "ca-west-1": "204538143572",
                     "cn-north-1": "727897471807",
                     "cn-northwest-1": "727897471807",
                     "eu-central-1": "763104351884",
                     "eu-central-2": "380420809688",
                     "eu-north-1": "763104351884",
+                    "eu-south-1": "692866216735",
+                    "eu-south-2": "503227376785",
                     "eu-west-1": "763104351884",
                     "eu-west-2": "763104351884",
                     "eu-west-3": "763104351884",
-                    "eu-south-1": "692866216735",
-                    "eu-south-2": "503227376785",
+                    "il-central-1": "780543022126",
+                    "me-central-1": "914824155844",
                     "me-south-1": "217643126080",
                     "sa-east-1": "763104351884",
                     "us-east-1": "763104351884",
@@ -2082,8 +2113,7 @@
                     "us-iso-east-1": "886529160074",
                     "us-isob-east-1": "094389454867",
                     "us-west-1": "763104351884",
-                    "us-west-2": "763104351884",
-                    "ca-west-1": "204538143572"
+                    "us-west-2": "763104351884"
                 },
                 "repository": "pytorch-training"
             },
@@ -2093,7 +2123,6 @@
                 ],
                 "registries": {
                     "af-south-1": "626614931356",
-                    "il-central-1": "780543022126",
                     "ap-east-1": "871362719292",
                     "ap-northeast-1": "763104351884",
                     "ap-northeast-2": "763104351884",
@@ -2105,16 +2134,19 @@
                     "ap-southeast-3": "907027046896",
                     "ap-southeast-4": "457447274322",
                     "ca-central-1": "763104351884",
+                    "ca-west-1": "204538143572",
                     "cn-north-1": "727897471807",
                     "cn-northwest-1": "727897471807",
                     "eu-central-1": "763104351884",
                     "eu-central-2": "380420809688",
                     "eu-north-1": "763104351884",
+                    "eu-south-1": "692866216735",
+                    "eu-south-2": "503227376785",
                     "eu-west-1": "763104351884",
                     "eu-west-2": "763104351884",
                     "eu-west-3": "763104351884",
-                    "eu-south-1": "692866216735",
-                    "eu-south-2": "503227376785",
+                    "il-central-1": "780543022126",
+                    "me-central-1": "914824155844",
                     "me-south-1": "217643126080",
                     "sa-east-1": "763104351884",
                     "us-east-1": "763104351884",
@@ -2124,8 +2156,7 @@
                     "us-iso-east-1": "886529160074",
                     "us-isob-east-1": "094389454867",
                     "us-west-1": "763104351884",
-                    "us-west-2": "763104351884",
-                    "ca-west-1": "204538143572"
+                    "us-west-2": "763104351884"
                 },
                 "repository": "pytorch-training"
             },
@@ -2135,7 +2166,6 @@
                 ],
                 "registries": {
                     "af-south-1": "626614931356",
-                    "il-central-1": "780543022126",
                     "ap-east-1": "871362719292",
                     "ap-northeast-1": "763104351884",
                     "ap-northeast-2": "763104351884",
@@ -2147,16 +2177,19 @@
                     "ap-southeast-3": "907027046896",
                     "ap-southeast-4": "457447274322",
                     "ca-central-1": "763104351884",
+                    "ca-west-1": "204538143572",
                     "cn-north-1": "727897471807",
                     "cn-northwest-1": "727897471807",
                     "eu-central-1": "763104351884",
                     "eu-central-2": "380420809688",
                     "eu-north-1": "763104351884",
+                    "eu-south-1": "692866216735",
+                    "eu-south-2": "503227376785",
                     "eu-west-1": "763104351884",
                     "eu-west-2": "763104351884",
                     "eu-west-3": "763104351884",
-                    "eu-south-1": "692866216735",
-                    "eu-south-2": "503227376785",
+                    "il-central-1": "780543022126",
+                    "me-central-1": "914824155844",
                     "me-south-1": "217643126080",
                     "sa-east-1": "763104351884",
                     "us-east-1": "763104351884",
@@ -2166,8 +2199,7 @@
                     "us-iso-east-1": "886529160074",
                     "us-isob-east-1": "094389454867",
                     "us-west-1": "763104351884",
-                    "us-west-2": "763104351884",
-                    "ca-west-1": "204538143572"
+                    "us-west-2": "763104351884"
                 },
                 "repository": "pytorch-training"
             },
@@ -2177,7 +2209,6 @@
                 ],
                 "registries": {
                     "af-south-1": "626614931356",
-                    "il-central-1": "780543022126",
                     "ap-east-1": "871362719292",
                     "ap-northeast-1": "763104351884",
                     "ap-northeast-2": "763104351884",
@@ -2189,16 +2220,19 @@
                     "ap-southeast-3": "907027046896",
                     "ap-southeast-4": "457447274322",
                     "ca-central-1": "763104351884",
+                    "ca-west-1": "204538143572",
                     "cn-north-1": "727897471807",
                     "cn-northwest-1": "727897471807",
                     "eu-central-1": "763104351884",
                     "eu-central-2": "380420809688",
                     "eu-north-1": "763104351884",
+                    "eu-south-1": "692866216735",
+                    "eu-south-2": "503227376785",
                     "eu-west-1": "763104351884",
                     "eu-west-2": "763104351884",
                     "eu-west-3": "763104351884",
-                    "eu-south-1": "692866216735",
-                    "eu-south-2": "503227376785",
+                    "il-central-1": "780543022126",
+                    "me-central-1": "914824155844",
                     "me-south-1": "217643126080",
                     "sa-east-1": "763104351884",
                     "us-east-1": "763104351884",
@@ -2208,8 +2242,7 @@
                     "us-iso-east-1": "886529160074",
                     "us-isob-east-1": "094389454867",
                     "us-west-1": "763104351884",
-                    "us-west-2": "763104351884",
-                    "ca-west-1": "204538143572"
+                    "us-west-2": "763104351884"
                 },
                 "repository": "pytorch-training"
             },
@@ -2219,7 +2252,6 @@
                 ],
                 "registries": {
                     "af-south-1": "626614931356",
-                    "il-central-1": "780543022126",
                     "ap-east-1": "871362719292",
                     "ap-northeast-1": "763104351884",
                     "ap-northeast-2": "763104351884",
@@ -2231,16 +2263,19 @@
                     "ap-southeast-3": "907027046896",
                     "ap-southeast-4": "457447274322",
                     "ca-central-1": "763104351884",
+                    "ca-west-1": "204538143572",
                     "cn-north-1": "727897471807",
                     "cn-northwest-1": "727897471807",
                     "eu-central-1": "763104351884",
                     "eu-central-2": "380420809688",
                     "eu-north-1": "763104351884",
+                    "eu-south-1": "692866216735",
+                    "eu-south-2": "503227376785",
                     "eu-west-1": "763104351884",
                     "eu-west-2": "763104351884",
                     "eu-west-3": "763104351884",
-                    "eu-south-1": "692866216735",
-                    "eu-south-2": "503227376785",
+                    "il-central-1": "780543022126",
+                    "me-central-1": "914824155844",
                     "me-south-1": "217643126080",
                     "sa-east-1": "763104351884",
                     "us-east-1": "763104351884",
@@ -2250,11 +2285,10 @@
                     "us-iso-east-1": "886529160074",
                     "us-isob-east-1": "094389454867",
                     "us-west-1": "763104351884",
-                    "us-west-2": "763104351884",
-                    "ca-west-1": "204538143572"
+                    "us-west-2": "763104351884"
                 },
                 "repository": "pytorch-training"
             }
         }
     }
-}
+}
\ No newline at end of file
diff --git a/src/sagemaker/image_uri_config/tensorflow.json b/src/sagemaker/image_uri_config/tensorflow.json
index 5dc8d35af2..9194dfbe4a 100644
--- a/src/sagemaker/image_uri_config/tensorflow.json
+++ b/src/sagemaker/image_uri_config/tensorflow.json
@@ -140,7 +140,6 @@
             "1.14.0": {
                 "registries": {
                     "af-south-1": "626614931356",
-                    "il-central-1": "780543022126",
                     "ap-east-1": "871362719292",
                     "ap-northeast-1": "763104351884",
                     "ap-northeast-2": "763104351884",
@@ -152,6 +151,7 @@
                     "ap-southeast-3": "907027046896",
                     "ap-southeast-4": "457447274322",
                     "ca-central-1": "763104351884",
+                    "ca-west-1": "204538143572",
                     "cn-north-1": "727897471807",
                     "cn-northwest-1": "727897471807",
                     "eu-central-1": "763104351884",
@@ -162,8 +162,9 @@
                     "eu-west-1": "763104351884",
                     "eu-west-2": "763104351884",
                     "eu-west-3": "763104351884",
-                    "me-south-1": "217643126080",
+                    "il-central-1": "780543022126",
                     "me-central-1": "914824155844",
+                    "me-south-1": "217643126080",
                     "sa-east-1": "763104351884",
                     "us-east-1": "763104351884",
                     "us-east-2": "763104351884",
@@ -172,15 +173,13 @@
                     "us-iso-east-1": "886529160074",
                     "us-isob-east-1": "094389454867",
                     "us-west-1": "763104351884",
-                    "us-west-2": "763104351884",
-                    "ca-west-1": "204538143572"
+                    "us-west-2": "763104351884"
                 },
                 "repository": "tensorflow-inference-eia"
             },
             "1.15.0": {
                 "registries": {
                     "af-south-1": "626614931356",
-                    "il-central-1": "780543022126",
                     "ap-east-1": "871362719292",
                     "ap-northeast-1": "763104351884",
                     "ap-northeast-2": "763104351884",
@@ -192,6 +191,7 @@
                     "ap-southeast-3": "907027046896",
                     "ap-southeast-4": "457447274322",
                     "ca-central-1": "763104351884",
+                    "ca-west-1": "204538143572",
                     "cn-north-1": "727897471807",
                     "cn-northwest-1": "727897471807",
                     "eu-central-1": "763104351884",
@@ -202,8 +202,9 @@
                     "eu-west-1": "763104351884",
                     "eu-west-2": "763104351884",
                     "eu-west-3": "763104351884",
-                    "me-south-1": "217643126080",
+                    "il-central-1": "780543022126",
                     "me-central-1": "914824155844",
+                    "me-south-1": "217643126080",
                     "sa-east-1": "763104351884",
                     "us-east-1": "763104351884",
                     "us-east-2": "763104351884",
@@ -212,15 +213,13 @@
                     "us-iso-east-1": "886529160074",
                     "us-isob-east-1": "094389454867",
                     "us-west-1": "763104351884",
-                    "us-west-2": "763104351884",
-                    "ca-west-1": "204538143572"
+                    "us-west-2": "763104351884"
                 },
                 "repository": "tensorflow-inference-eia"
             },
             "2.0.0": {
                 "registries": {
                     "af-south-1": "626614931356",
-                    "il-central-1": "780543022126",
                     "ap-east-1": "871362719292",
                     "ap-northeast-1": "763104351884",
                     "ap-northeast-2": "763104351884",
@@ -232,6 +231,7 @@
                     "ap-southeast-3": "907027046896",
                     "ap-southeast-4": "457447274322",
                     "ca-central-1": "763104351884",
+                    "ca-west-1": "204538143572",
                     "cn-north-1": "727897471807",
                     "cn-northwest-1": "727897471807",
                     "eu-central-1": "763104351884",
@@ -242,8 +242,9 @@
                     "eu-west-1": "763104351884",
                     "eu-west-2": "763104351884",
                     "eu-west-3": "763104351884",
-                    "me-south-1": "217643126080",
+                    "il-central-1": "780543022126",
                     "me-central-1": "914824155844",
+                    "me-south-1": "217643126080",
                     "sa-east-1": "763104351884",
                     "us-east-1": "763104351884",
                     "us-east-2": "763104351884",
@@ -252,15 +253,13 @@
                     "us-iso-east-1": "886529160074",
                     "us-isob-east-1": "094389454867",
                     "us-west-1": "763104351884",
-                    "us-west-2": "763104351884",
-                    "ca-west-1": "204538143572"
+                    "us-west-2": "763104351884"
                 },
                 "repository": "tensorflow-inference-eia"
             },
             "2.3.0": {
                 "registries": {
                     "af-south-1": "626614931356",
-                    "il-central-1": "780543022126",
                     "ap-east-1": "871362719292",
                     "ap-northeast-1": "763104351884",
                     "ap-northeast-2": "763104351884",
@@ -272,6 +271,7 @@
                     "ap-southeast-3": "907027046896",
                     "ap-southeast-4": "457447274322",
                     "ca-central-1": "763104351884",
+                    "ca-west-1": "204538143572",
                     "cn-north-1": "727897471807",
                     "cn-northwest-1": "727897471807",
                     "eu-central-1": "763104351884",
@@ -282,8 +282,9 @@
                     "eu-west-1": "763104351884",
                     "eu-west-2": "763104351884",
                     "eu-west-3": "763104351884",
-                    "me-south-1": "217643126080",
+                    "il-central-1": "780543022126",
                     "me-central-1": "914824155844",
+                    "me-south-1": "217643126080",
                     "sa-east-1": "763104351884",
                     "us-east-1": "763104351884",
                     "us-east-2": "763104351884",
@@ -292,8 +293,7 @@
                     "us-iso-east-1": "886529160074",
                     "us-isob-east-1": "094389454867",
                     "us-west-1": "763104351884",
-                    "us-west-2": "763104351884",
-                    "ca-west-1": "204538143572"
+                    "us-west-2": "763104351884"
                 },
                 "repository": "tensorflow-inference-eia"
             }
@@ -305,18 +305,18 @@
             "gpu"
         ],
         "version_aliases": {
-            "1.10": "1.10.0",
-            "1.11": "1.11.0",
-            "1.12": "1.12.0",
-            "1.13": "1.13.0",
-            "1.14": "1.14.0",
-            "1.15": "1.15.5",
             "1.4": "1.4.1",
             "1.5": "1.5.0",
             "1.6": "1.6.0",
             "1.7": "1.7.0",
             "1.8": "1.8.0",
             "1.9": "1.9.0",
+            "1.10": "1.10.0",
+            "1.11": "1.11.0",
+            "1.12": "1.12.0",
+            "1.13": "1.13.0",
+            "1.14": "1.14.0",
+            "1.15": "1.15.5",
             "2.0": "2.0.4",
             "2.1": "2.1.3",
             "2.2": "2.2.2",
@@ -334,6 +334,204 @@
             "2.14": "2.14.1"
         },
         "versions": {
+            "1.4.1": {
+                "py_versions": [
+                    "py2"
+                ],
+                "registries": {
+                    "af-south-1": "313743910680",
+                    "ap-east-1": "057415533634",
+                    "ap-northeast-1": "520713654638",
+                    "ap-northeast-2": "520713654638",
+                    "ap-south-1": "520713654638",
+                    "ap-southeast-1": "520713654638",
+                    "ap-southeast-2": "520713654638",
+                    "ca-central-1": "520713654638",
+                    "cn-north-1": "422961961927",
+                    "cn-northwest-1": "423003514399",
+                    "eu-central-1": "520713654638",
+                    "eu-north-1": "520713654638",
+                    "eu-south-1": "048378556238",
+                    "eu-west-1": "520713654638",
+                    "eu-west-2": "520713654638",
+                    "eu-west-3": "520713654638",
+                    "me-south-1": "724002660598",
+                    "sa-east-1": "520713654638",
+                    "us-east-1": "520713654638",
+                    "us-east-2": "520713654638",
+                    "us-gov-west-1": "246785580436",
+                    "us-iso-east-1": "744548109606",
+                    "us-isob-east-1": "453391408702",
+                    "us-west-1": "520713654638",
+                    "us-west-2": "520713654638"
+                },
+                "repository": "sagemaker-tensorflow"
+            },
+            "1.5.0": {
+                "py_versions": [
+                    "py2"
+                ],
+                "registries": {
+                    "af-south-1": "313743910680",
+                    "ap-east-1": "057415533634",
+                    "ap-northeast-1": "520713654638",
+                    "ap-northeast-2": "520713654638",
+                    "ap-south-1": "520713654638",
+                    "ap-southeast-1": "520713654638",
+                    "ap-southeast-2": "520713654638",
+                    "ca-central-1": "520713654638",
+                    "cn-north-1": "422961961927",
+                    "cn-northwest-1": "423003514399",
+                    "eu-central-1": "520713654638",
+                    "eu-north-1": "520713654638",
+                    "eu-south-1": "048378556238",
+                    "eu-west-1": "520713654638",
+                    "eu-west-2": "520713654638",
+                    "eu-west-3": "520713654638",
+                    "me-south-1": "724002660598",
+                    "sa-east-1": "520713654638",
+                    "us-east-1": "520713654638",
+                    "us-east-2": "520713654638",
+                    "us-gov-west-1": "246785580436",
+                    "us-iso-east-1": "744548109606",
+                    "us-isob-east-1": "453391408702",
+                    "us-west-1": "520713654638",
+                    "us-west-2": "520713654638"
+                },
+                "repository": "sagemaker-tensorflow"
+            },
+            "1.6.0": {
+                "py_versions": [
+                    "py2"
+                ],
+                "registries": {
+                    "af-south-1": "313743910680",
+                    "ap-east-1": "057415533634",
+                    "ap-northeast-1": "520713654638",
+                    "ap-northeast-2": "520713654638",
+                    "ap-south-1": "520713654638",
+                    "ap-southeast-1": "520713654638",
+                    "ap-southeast-2": "520713654638",
+                    "ca-central-1": "520713654638",
+                    "cn-north-1": "422961961927",
+                    "cn-northwest-1": "423003514399",
+                    "eu-central-1": "520713654638",
+                    "eu-north-1": "520713654638",
+                    "eu-south-1": "048378556238",
+                    "eu-west-1": "520713654638",
+                    "eu-west-2": "520713654638",
+                    "eu-west-3": "520713654638",
+                    "me-south-1": "724002660598",
+                    "sa-east-1": "520713654638",
+                    "us-east-1": "520713654638",
+                    "us-east-2": "520713654638",
+                    "us-gov-west-1": "246785580436",
+                    "us-iso-east-1": "744548109606",
+                    "us-isob-east-1": "453391408702",
+                    "us-west-1": "520713654638",
+                    "us-west-2": "520713654638"
+                },
+                "repository": "sagemaker-tensorflow"
+            },
+            "1.7.0": {
+                "py_versions": [
+                    "py2"
+                ],
+                "registries": {
+                    "af-south-1": "313743910680",
+                    "ap-east-1": "057415533634",
+                    "ap-northeast-1": "520713654638",
+                    "ap-northeast-2": "520713654638",
+                    "ap-south-1": "520713654638",
+                    "ap-southeast-1": "520713654638",
+                    "ap-southeast-2": "520713654638",
+                    "ca-central-1": "520713654638",
+                    "cn-north-1": "422961961927",
+                    "cn-northwest-1": "423003514399",
+                    "eu-central-1": "520713654638",
+                    "eu-north-1": "520713654638",
+                    "eu-south-1": "048378556238",
+                    "eu-west-1": "520713654638",
+                    "eu-west-2": "520713654638",
+                    "eu-west-3": "520713654638",
+                    "me-south-1": "724002660598",
+                    "sa-east-1": "520713654638",
+                    "us-east-1": "520713654638",
+                    "us-east-2": "520713654638",
+                    "us-gov-west-1": "246785580436",
+                    "us-iso-east-1": "744548109606",
+                    "us-isob-east-1": "453391408702",
+                    "us-west-1": "520713654638",
+                    "us-west-2": "520713654638"
+                },
+                "repository": "sagemaker-tensorflow"
+            },
+            "1.8.0": {
+                "py_versions": [
+                    "py2"
+                ],
+                "registries": {
+                    "af-south-1": "313743910680",
+                    "ap-east-1": "057415533634",
+                    "ap-northeast-1": "520713654638",
+                    "ap-northeast-2": "520713654638",
+                    "ap-south-1": "520713654638",
+                    "ap-southeast-1": "520713654638",
+                    "ap-southeast-2": "520713654638",
+                    "ca-central-1": "520713654638",
+                    "cn-north-1": "422961961927",
+                    "cn-northwest-1": "423003514399",
+                    "eu-central-1": "520713654638",
+                    "eu-north-1": "520713654638",
+                    "eu-south-1": "048378556238",
+                    "eu-west-1": "520713654638",
+                    "eu-west-2": "520713654638",
+                    "eu-west-3": "520713654638",
+                    "me-south-1": "724002660598",
+                    "sa-east-1": "520713654638",
+                    "us-east-1": "520713654638",
+                    "us-east-2": "520713654638",
+                    "us-gov-west-1": "246785580436",
+                    "us-iso-east-1": "744548109606",
+                    "us-isob-east-1": "453391408702",
+                    "us-west-1": "520713654638",
+                    "us-west-2": "520713654638"
+                },
+                "repository": "sagemaker-tensorflow"
+            },
+            "1.9.0": {
+                "py_versions": [
+                    "py2"
+                ],
+                "registries": {
+                    "af-south-1": "313743910680",
+                    "ap-east-1": "057415533634",
+                    "ap-northeast-1": "520713654638",
+                    "ap-northeast-2": "520713654638",
+                    "ap-south-1": "520713654638",
+                    "ap-southeast-1": "520713654638",
+                    "ap-southeast-2": "520713654638",
+                    "ca-central-1": "520713654638",
+                    "cn-north-1": "422961961927",
+                    "cn-northwest-1": "423003514399",
+                    "eu-central-1": "520713654638",
+                    "eu-north-1": "520713654638",
+                    "eu-south-1": "048378556238",
+                    "eu-west-1": "520713654638",
+                    "eu-west-2": "520713654638",
+                    "eu-west-3": "520713654638",
+                    "me-south-1": "724002660598",
+                    "sa-east-1": "520713654638",
+                    "us-east-1": "520713654638",
+                    "us-east-2": "520713654638",
+                    "us-gov-west-1": "246785580436",
+                    "us-iso-east-1": "744548109606",
+                    "us-isob-east-1": "453391408702",
+                    "us-west-1": "520713654638",
+                    "us-west-2": "520713654638"
+                },
+                "repository": "sagemaker-tensorflow"
+            },
             "1.10.0": {
                 "py_versions": [
                     "py2"
@@ -430,7 +628,6 @@
             "1.13.0": {
                 "registries": {
                     "af-south-1": "626614931356",
-                    "il-central-1": "780543022126",
                     "ap-east-1": "871362719292",
                     "ap-northeast-1": "763104351884",
                     "ap-northeast-2": "763104351884",
@@ -442,6 +639,7 @@
                     "ap-southeast-3": "907027046896",
                     "ap-southeast-4": "457447274322",
                     "ca-central-1": "763104351884",
+                    "ca-west-1": "204538143572",
                     "cn-north-1": "727897471807",
                     "cn-northwest-1": "727897471807",
                     "eu-central-1": "763104351884",
@@ -452,8 +650,9 @@
                     "eu-west-1": "763104351884",
                     "eu-west-2": "763104351884",
                     "eu-west-3": "763104351884",
-                    "me-south-1": "217643126080",
+                    "il-central-1": "780543022126",
                     "me-central-1": "914824155844",
+                    "me-south-1": "217643126080",
                     "sa-east-1": "763104351884",
                     "us-east-1": "763104351884",
                     "us-east-2": "763104351884",
@@ -462,15 +661,13 @@
                     "us-iso-east-1": "886529160074",
                     "us-isob-east-1": "094389454867",
                     "us-west-1": "763104351884",
-                    "us-west-2": "763104351884",
-                    "ca-west-1": "204538143572"
+                    "us-west-2": "763104351884"
                 },
                 "repository": "tensorflow-inference"
             },
             "1.14.0": {
                 "registries": {
                     "af-south-1": "626614931356",
-                    "il-central-1": "780543022126",
                     "ap-east-1": "871362719292",
                     "ap-northeast-1": "763104351884",
                     "ap-northeast-2": "763104351884",
@@ -482,6 +679,7 @@
                     "ap-southeast-3": "907027046896",
                     "ap-southeast-4": "457447274322",
                     "ca-central-1": "763104351884",
+                    "ca-west-1": "204538143572",
                     "cn-north-1": "727897471807",
                     "cn-northwest-1": "727897471807",
                     "eu-central-1": "763104351884",
@@ -492,8 +690,9 @@
                     "eu-west-1": "763104351884",
                     "eu-west-2": "763104351884",
                     "eu-west-3": "763104351884",
-                    "me-south-1": "217643126080",
+                    "il-central-1": "780543022126",
                     "me-central-1": "914824155844",
+                    "me-south-1": "217643126080",
                     "sa-east-1": "763104351884",
                     "us-east-1": "763104351884",
                     "us-east-2": "763104351884",
@@ -502,15 +701,13 @@
                     "us-iso-east-1": "886529160074",
                     "us-isob-east-1": "094389454867",
                     "us-west-1": "763104351884",
-                    "us-west-2": "763104351884",
-                    "ca-west-1": "204538143572"
+                    "us-west-2": "763104351884"
                 },
                 "repository": "tensorflow-inference"
             },
             "1.15.0": {
                 "registries": {
                     "af-south-1": "626614931356",
-                    "il-central-1": "780543022126",
                     "ap-east-1": "871362719292",
                     "ap-northeast-1": "763104351884",
                     "ap-northeast-2": "763104351884",
@@ -522,6 +719,7 @@
                     "ap-southeast-3": "907027046896",
                     "ap-southeast-4": "457447274322",
                     "ca-central-1": "763104351884",
+                    "ca-west-1": "204538143572",
                     "cn-north-1": "727897471807",
                     "cn-northwest-1": "727897471807",
                     "eu-central-1": "763104351884",
@@ -532,8 +730,9 @@
                     "eu-west-1": "763104351884",
                     "eu-west-2": "763104351884",
                     "eu-west-3": "763104351884",
-                    "me-south-1": "217643126080",
+                    "il-central-1": "780543022126",
                     "me-central-1": "914824155844",
+                    "me-south-1": "217643126080",
                     "sa-east-1": "763104351884",
                     "us-east-1": "763104351884",
                     "us-east-2": "763104351884",
@@ -542,15 +741,13 @@
                     "us-iso-east-1": "886529160074",
                     "us-isob-east-1": "094389454867",
                     "us-west-1": "763104351884",
-                    "us-west-2": "763104351884",
-                    "ca-west-1": "204538143572"
+                    "us-west-2": "763104351884"
                 },
                 "repository": "tensorflow-inference"
             },
             "1.15.2": {
                 "registries": {
                     "af-south-1": "626614931356",
-                    "il-central-1": "780543022126",
                     "ap-east-1": "871362719292",
                     "ap-northeast-1": "763104351884",
                     "ap-northeast-2": "763104351884",
@@ -562,6 +759,7 @@
                     "ap-southeast-3": "907027046896",
                     "ap-southeast-4": "457447274322",
                     "ca-central-1": "763104351884",
+                    "ca-west-1": "204538143572",
                     "cn-north-1": "727897471807",
                     "cn-northwest-1": "727897471807",
                     "eu-central-1": "763104351884",
@@ -572,8 +770,9 @@
                     "eu-west-1": "763104351884",
                     "eu-west-2": "763104351884",
                     "eu-west-3": "763104351884",
-                    "me-south-1": "217643126080",
+                    "il-central-1": "780543022126",
                     "me-central-1": "914824155844",
+                    "me-south-1": "217643126080",
                     "sa-east-1": "763104351884",
                     "us-east-1": "763104351884",
                     "us-east-2": "763104351884",
@@ -582,15 +781,13 @@
                     "us-iso-east-1": "886529160074",
                     "us-isob-east-1": "094389454867",
                     "us-west-1": "763104351884",
-                    "us-west-2": "763104351884",
-                    "ca-west-1": "204538143572"
+                    "us-west-2": "763104351884"
                 },
                 "repository": "tensorflow-inference"
             },
             "1.15.3": {
                 "registries": {
                     "af-south-1": "626614931356",
-                    "il-central-1": "780543022126",
                     "ap-east-1": "871362719292",
                     "ap-northeast-1": "763104351884",
                     "ap-northeast-2": "763104351884",
@@ -602,6 +799,7 @@
                     "ap-southeast-3": "907027046896",
                     "ap-southeast-4": "457447274322",
                     "ca-central-1": "763104351884",
+                    "ca-west-1": "204538143572",
                     "cn-north-1": "727897471807",
                     "cn-northwest-1": "727897471807",
                     "eu-central-1": "763104351884",
@@ -612,8 +810,9 @@
                     "eu-west-1": "763104351884",
                     "eu-west-2": "763104351884",
                     "eu-west-3": "763104351884",
-                    "me-south-1": "217643126080",
+                    "il-central-1": "780543022126",
                     "me-central-1": "914824155844",
+                    "me-south-1": "217643126080",
                     "sa-east-1": "763104351884",
                     "us-east-1": "763104351884",
                     "us-east-2": "763104351884",
@@ -622,15 +821,13 @@
                     "us-iso-east-1": "886529160074",
                     "us-isob-east-1": "094389454867",
                     "us-west-1": "763104351884",
-                    "us-west-2": "763104351884",
-                    "ca-west-1": "204538143572"
+                    "us-west-2": "763104351884"
                 },
                 "repository": "tensorflow-inference"
             },
             "1.15.4": {
                 "registries": {
                     "af-south-1": "626614931356",
-                    "il-central-1": "780543022126",
                     "ap-east-1": "871362719292",
                     "ap-northeast-1": "763104351884",
                     "ap-northeast-2": "763104351884",
@@ -642,6 +839,7 @@
                     "ap-southeast-3": "907027046896",
                     "ap-southeast-4": "457447274322",
                     "ca-central-1": "763104351884",
+                    "ca-west-1": "204538143572",
                     "cn-north-1": "727897471807",
                     "cn-northwest-1": "727897471807",
                     "eu-central-1": "763104351884",
@@ -652,8 +850,9 @@
                     "eu-west-1": "763104351884",
                     "eu-west-2": "763104351884",
                     "eu-west-3": "763104351884",
-                    "me-south-1": "217643126080",
+                    "il-central-1": "780543022126",
                     "me-central-1": "914824155844",
+                    "me-south-1": "217643126080",
                     "sa-east-1": "763104351884",
                     "us-east-1": "763104351884",
                     "us-east-2": "763104351884",
@@ -662,15 +861,13 @@
                     "us-iso-east-1": "886529160074",
                     "us-isob-east-1": "094389454867",
                     "us-west-1": "763104351884",
-                    "us-west-2": "763104351884",
-                    "ca-west-1": "204538143572"
+                    "us-west-2": "763104351884"
                 },
                 "repository": "tensorflow-inference"
             },
             "1.15.5": {
                 "registries": {
                     "af-south-1": "626614931356",
-                    "il-central-1": "780543022126",
                     "ap-east-1": "871362719292",
                     "ap-northeast-1": "763104351884",
                     "ap-northeast-2": "763104351884",
@@ -682,6 +879,7 @@
                     "ap-southeast-3": "907027046896",
                     "ap-southeast-4": "457447274322",
                     "ca-central-1": "763104351884",
+                    "ca-west-1": "204538143572",
                     "cn-north-1": "727897471807",
                     "cn-northwest-1": "727897471807",
                     "eu-central-1": "763104351884",
@@ -692,8 +890,9 @@
                     "eu-west-1": "763104351884",
                     "eu-west-2": "763104351884",
                     "eu-west-3": "763104351884",
-                    "me-south-1": "217643126080",
+                    "il-central-1": "780543022126",
                     "me-central-1": "914824155844",
+                    "me-south-1": "217643126080",
                     "sa-east-1": "763104351884",
                     "us-east-1": "763104351884",
                     "us-east-2": "763104351884",
@@ -702,213 +901,13 @@
                     "us-iso-east-1": "886529160074",
                     "us-isob-east-1": "094389454867",
                     "us-west-1": "763104351884",
-                    "us-west-2": "763104351884",
-                    "ca-west-1": "204538143572"
+                    "us-west-2": "763104351884"
                 },
                 "repository": "tensorflow-inference"
             },
-            "1.4.1": {
-                "py_versions": [
-                    "py2"
-                ],
-                "registries": {
-                    "af-south-1": "313743910680",
-                    "ap-east-1": "057415533634",
-                    "ap-northeast-1": "520713654638",
-                    "ap-northeast-2": "520713654638",
-                    "ap-south-1": "520713654638",
-                    "ap-southeast-1": "520713654638",
-                    "ap-southeast-2": "520713654638",
-                    "ca-central-1": "520713654638",
-                    "cn-north-1": "422961961927",
-                    "cn-northwest-1": "423003514399",
-                    "eu-central-1": "520713654638",
-                    "eu-north-1": "520713654638",
-                    "eu-south-1": "048378556238",
-                    "eu-west-1": "520713654638",
-                    "eu-west-2": "520713654638",
-                    "eu-west-3": "520713654638",
-                    "me-south-1": "724002660598",
-                    "sa-east-1": "520713654638",
-                    "us-east-1": "520713654638",
-                    "us-east-2": "520713654638",
-                    "us-gov-west-1": "246785580436",
-                    "us-iso-east-1": "744548109606",
-                    "us-isob-east-1": "453391408702",
-                    "us-west-1": "520713654638",
-                    "us-west-2": "520713654638"
-                },
-                "repository": "sagemaker-tensorflow"
-            },
-            "1.5.0": {
-                "py_versions": [
-                    "py2"
-                ],
-                "registries": {
-                    "af-south-1": "313743910680",
-                    "ap-east-1": "057415533634",
-                    "ap-northeast-1": "520713654638",
-                    "ap-northeast-2": "520713654638",
-                    "ap-south-1": "520713654638",
-                    "ap-southeast-1": "520713654638",
-                    "ap-southeast-2": "520713654638",
-                    "ca-central-1": "520713654638",
-                    "cn-north-1": "422961961927",
-                    "cn-northwest-1": "423003514399",
-                    "eu-central-1": "520713654638",
-                    "eu-north-1": "520713654638",
-                    "eu-south-1": "048378556238",
-                    "eu-west-1": "520713654638",
-                    "eu-west-2": "520713654638",
-                    "eu-west-3": "520713654638",
-                    "me-south-1": "724002660598",
-                    "sa-east-1": "520713654638",
-                    "us-east-1": "520713654638",
-                    "us-east-2": "520713654638",
-                    "us-gov-west-1": "246785580436",
-                    "us-iso-east-1": "744548109606",
-                    "us-isob-east-1": "453391408702",
-                    "us-west-1": "520713654638",
-                    "us-west-2": "520713654638"
-                },
-                "repository": "sagemaker-tensorflow"
-            },
-            "1.6.0": {
-                "py_versions": [
-                    "py2"
-                ],
-                "registries": {
-                    "af-south-1": "313743910680",
-                    "ap-east-1": "057415533634",
-                    "ap-northeast-1": "520713654638",
-                    "ap-northeast-2": "520713654638",
-                    "ap-south-1": "520713654638",
-                    "ap-southeast-1": "520713654638",
-                    "ap-southeast-2": "520713654638",
-                    "ca-central-1": "520713654638",
-                    "cn-north-1": "422961961927",
-                    "cn-northwest-1": "423003514399",
-                    "eu-central-1": "520713654638",
-                    "eu-north-1": "520713654638",
-                    "eu-south-1": "048378556238",
-                    "eu-west-1": "520713654638",
-                    "eu-west-2": "520713654638",
-                    "eu-west-3": "520713654638",
-                    "me-south-1": "724002660598",
-                    "sa-east-1": "520713654638",
-                    "us-east-1": "520713654638",
-                    "us-east-2": "520713654638",
-                    "us-gov-west-1": "246785580436",
-                    "us-iso-east-1": "744548109606",
-                    "us-isob-east-1": "453391408702",
-                    "us-west-1": "520713654638",
-                    "us-west-2": "520713654638"
-                },
-                "repository": "sagemaker-tensorflow"
-            },
-            "1.7.0": {
-                "py_versions": [
-                    "py2"
-                ],
-                "registries": {
-                    "af-south-1": "313743910680",
-                    "ap-east-1": "057415533634",
-                    "ap-northeast-1": "520713654638",
-                    "ap-northeast-2": "520713654638",
-                    "ap-south-1": "520713654638",
-                    "ap-southeast-1": "520713654638",
-                    "ap-southeast-2": "520713654638",
-                    "ca-central-1": "520713654638",
-                    "cn-north-1": "422961961927",
-                    "cn-northwest-1": "423003514399",
-                    "eu-central-1": "520713654638",
-                    "eu-north-1": "520713654638",
-                    "eu-south-1": "048378556238",
-                    "eu-west-1": "520713654638",
-                    "eu-west-2": "520713654638",
-                    "eu-west-3": "520713654638",
-                    "me-south-1": "724002660598",
-                    "sa-east-1": "520713654638",
-                    "us-east-1": "520713654638",
-                    "us-east-2": "520713654638",
-                    "us-gov-west-1": "246785580436",
-                    "us-iso-east-1": "744548109606",
-                    "us-isob-east-1": "453391408702",
-                    "us-west-1": "520713654638",
-                    "us-west-2": "520713654638"
-                },
-                "repository": "sagemaker-tensorflow"
-            },
-            "1.8.0": {
-                "py_versions": [
-                    "py2"
-                ],
-                "registries": {
-                    "af-south-1": "313743910680",
-                    "ap-east-1": "057415533634",
-                    "ap-northeast-1": "520713654638",
-                    "ap-northeast-2": "520713654638",
-                    "ap-south-1": "520713654638",
-                    "ap-southeast-1": "520713654638",
-                    "ap-southeast-2": "520713654638",
-                    "ca-central-1": "520713654638",
-                    "cn-north-1": "422961961927",
-                    "cn-northwest-1": "423003514399",
-                    "eu-central-1": "520713654638",
-                    "eu-north-1": "520713654638",
-                    "eu-south-1": "048378556238",
-                    "eu-west-1": "520713654638",
-                    "eu-west-2": "520713654638",
-                    "eu-west-3": "520713654638",
-                    "me-south-1": "724002660598",
-                    "sa-east-1": "520713654638",
-                    "us-east-1": "520713654638",
-                    "us-east-2": "520713654638",
-                    "us-gov-west-1": "246785580436",
-                    "us-iso-east-1": "744548109606",
-                    "us-isob-east-1": "453391408702",
-                    "us-west-1": "520713654638",
-                    "us-west-2": "520713654638"
-                },
-                "repository": "sagemaker-tensorflow"
-            },
-            "1.9.0": {
-                "py_versions": [
-                    "py2"
-                ],
-                "registries": {
-                    "af-south-1": "313743910680",
-                    "ap-east-1": "057415533634",
-                    "ap-northeast-1": "520713654638",
-                    "ap-northeast-2": "520713654638",
-                    "ap-south-1": "520713654638",
-                    "ap-southeast-1": "520713654638",
-                    "ap-southeast-2": "520713654638",
-                    "ca-central-1": "520713654638",
-                    "cn-north-1": "422961961927",
-                    "cn-northwest-1": "423003514399",
-                    "eu-central-1": "520713654638",
-                    "eu-north-1": "520713654638",
-                    "eu-south-1": "048378556238",
-                    "eu-west-1": "520713654638",
-                    "eu-west-2": "520713654638",
-                    "eu-west-3": "520713654638",
-                    "me-south-1": "724002660598",
-                    "sa-east-1": "520713654638",
-                    "us-east-1": "520713654638",
-                    "us-east-2": "520713654638",
-                    "us-gov-west-1": "246785580436",
-                    "us-iso-east-1": "744548109606",
-                    "us-isob-east-1": "453391408702",
-                    "us-west-1": "520713654638",
-                    "us-west-2": "520713654638"
-                },
-                "repository": "sagemaker-tensorflow"
-            },
             "2.0.0": {
                 "registries": {
                     "af-south-1": "626614931356",
-                    "il-central-1": "780543022126",
                     "ap-east-1": "871362719292",
                     "ap-northeast-1": "763104351884",
                     "ap-northeast-2": "763104351884",
@@ -920,6 +919,7 @@
                     "ap-southeast-3": "907027046896",
                     "ap-southeast-4": "457447274322",
                     "ca-central-1": "763104351884",
+                    "ca-west-1": "204538143572",
                     "cn-north-1": "727897471807",
                     "cn-northwest-1": "727897471807",
                     "eu-central-1": "763104351884",
@@ -930,8 +930,9 @@
                     "eu-west-1": "763104351884",
                     "eu-west-2": "763104351884",
                     "eu-west-3": "763104351884",
-                    "me-south-1": "217643126080",
+                    "il-central-1": "780543022126",
                     "me-central-1": "914824155844",
+                    "me-south-1": "217643126080",
                     "sa-east-1": "763104351884",
                     "us-east-1": "763104351884",
                     "us-east-2": "763104351884",
@@ -940,15 +941,13 @@
                     "us-iso-east-1": "886529160074",
                     "us-isob-east-1": "094389454867",
                     "us-west-1": "763104351884",
-                    "us-west-2": "763104351884",
-                    "ca-west-1": "204538143572"
+                    "us-west-2": "763104351884"
                 },
                 "repository": "tensorflow-inference"
             },
             "2.0.1": {
                 "registries": {
                     "af-south-1": "626614931356",
-                    "il-central-1": "780543022126",
                     "ap-east-1": "871362719292",
                     "ap-northeast-1": "763104351884",
                     "ap-northeast-2": "763104351884",
@@ -960,6 +959,7 @@
                     "ap-southeast-3": "907027046896",
                     "ap-southeast-4": "457447274322",
                     "ca-central-1": "763104351884",
+                    "ca-west-1": "204538143572",
                     "cn-north-1": "727897471807",
                     "cn-northwest-1": "727897471807",
                     "eu-central-1": "763104351884",
@@ -970,8 +970,9 @@
                     "eu-west-1": "763104351884",
                     "eu-west-2": "763104351884",
                     "eu-west-3": "763104351884",
-                    "me-south-1": "217643126080",
+                    "il-central-1": "780543022126",
                     "me-central-1": "914824155844",
+                    "me-south-1": "217643126080",
                     "sa-east-1": "763104351884",
                     "us-east-1": "763104351884",
                     "us-east-2": "763104351884",
@@ -980,15 +981,13 @@
                     "us-iso-east-1": "886529160074",
                     "us-isob-east-1": "094389454867",
                     "us-west-1": "763104351884",
-                    "us-west-2": "763104351884",
-                    "ca-west-1": "204538143572"
+                    "us-west-2": "763104351884"
                 },
                 "repository": "tensorflow-inference"
             },
             "2.0.2": {
                 "registries": {
                     "af-south-1": "626614931356",
-                    "il-central-1": "780543022126",
                     "ap-east-1": "871362719292",
                     "ap-northeast-1": "763104351884",
                     "ap-northeast-2": "763104351884",
@@ -1000,6 +999,7 @@
                     "ap-southeast-3": "907027046896",
                     "ap-southeast-4": "457447274322",
                     "ca-central-1": "763104351884",
+                    "ca-west-1": "204538143572",
                     "cn-north-1": "727897471807",
                     "cn-northwest-1": "727897471807",
                     "eu-central-1": "763104351884",
@@ -1010,8 +1010,9 @@
                     "eu-west-1": "763104351884",
                     "eu-west-2": "763104351884",
                     "eu-west-3": "763104351884",
-                    "me-south-1": "217643126080",
+                    "il-central-1": "780543022126",
                     "me-central-1": "914824155844",
+                    "me-south-1": "217643126080",
                     "sa-east-1": "763104351884",
                     "us-east-1": "763104351884",
                     "us-east-2": "763104351884",
@@ -1020,15 +1021,13 @@
                     "us-iso-east-1": "886529160074",
                     "us-isob-east-1": "094389454867",
                     "us-west-1": "763104351884",
-                    "us-west-2": "763104351884",
-                    "ca-west-1": "204538143572"
+                    "us-west-2": "763104351884"
                 },
                 "repository": "tensorflow-inference"
             },
             "2.0.3": {
                 "registries": {
                     "af-south-1": "626614931356",
-                    "il-central-1": "780543022126",
                     "ap-east-1": "871362719292",
                     "ap-northeast-1": "763104351884",
                     "ap-northeast-2": "763104351884",
@@ -1040,6 +1039,7 @@
                     "ap-southeast-3": "907027046896",
                     "ap-southeast-4": "457447274322",
                     "ca-central-1": "763104351884",
+                    "ca-west-1": "204538143572",
                     "cn-north-1": "727897471807",
                     "cn-northwest-1": "727897471807",
                     "eu-central-1": "763104351884",
@@ -1050,8 +1050,9 @@
                     "eu-west-1": "763104351884",
                     "eu-west-2": "763104351884",
                     "eu-west-3": "763104351884",
-                    "me-south-1": "217643126080",
+                    "il-central-1": "780543022126",
                     "me-central-1": "914824155844",
+                    "me-south-1": "217643126080",
                     "sa-east-1": "763104351884",
                     "us-east-1": "763104351884",
                     "us-east-2": "763104351884",
@@ -1060,15 +1061,13 @@
                     "us-iso-east-1": "886529160074",
                     "us-isob-east-1": "094389454867",
                     "us-west-1": "763104351884",
-                    "us-west-2": "763104351884",
-                    "ca-west-1": "204538143572"
+                    "us-west-2": "763104351884"
                 },
                 "repository": "tensorflow-inference"
             },
             "2.0.4": {
                 "registries": {
                     "af-south-1": "626614931356",
-                    "il-central-1": "780543022126",
                     "ap-east-1": "871362719292",
                     "ap-northeast-1": "763104351884",
                     "ap-northeast-2": "763104351884",
@@ -1080,6 +1079,7 @@
                     "ap-southeast-3": "907027046896",
                     "ap-southeast-4": "457447274322",
                     "ca-central-1": "763104351884",
+                    "ca-west-1": "204538143572",
                     "cn-north-1": "727897471807",
                     "cn-northwest-1": "727897471807",
                     "eu-central-1": "763104351884",
@@ -1090,8 +1090,9 @@
                     "eu-west-1": "763104351884",
                     "eu-west-2": "763104351884",
                     "eu-west-3": "763104351884",
-                    "me-south-1": "217643126080",
+                    "il-central-1": "780543022126",
                     "me-central-1": "914824155844",
+                    "me-south-1": "217643126080",
                     "sa-east-1": "763104351884",
                     "us-east-1": "763104351884",
                     "us-east-2": "763104351884",
@@ -1100,15 +1101,13 @@
                     "us-iso-east-1": "886529160074",
                     "us-isob-east-1": "094389454867",
                     "us-west-1": "763104351884",
-                    "us-west-2": "763104351884",
-                    "ca-west-1": "204538143572"
+                    "us-west-2": "763104351884"
                 },
                 "repository": "tensorflow-inference"
             },
             "2.1.0": {
                 "registries": {
                     "af-south-1": "626614931356",
-                    "il-central-1": "780543022126",
                     "ap-east-1": "871362719292",
                     "ap-northeast-1": "763104351884",
                     "ap-northeast-2": "763104351884",
@@ -1120,6 +1119,7 @@
                     "ap-southeast-3": "907027046896",
                     "ap-southeast-4": "457447274322",
                     "ca-central-1": "763104351884",
+                    "ca-west-1": "204538143572",
                     "cn-north-1": "727897471807",
                     "cn-northwest-1": "727897471807",
                     "eu-central-1": "763104351884",
@@ -1130,8 +1130,9 @@
                     "eu-west-1": "763104351884",
                     "eu-west-2": "763104351884",
                     "eu-west-3": "763104351884",
-                    "me-south-1": "217643126080",
+                    "il-central-1": "780543022126",
                     "me-central-1": "914824155844",
+                    "me-south-1": "217643126080",
                     "sa-east-1": "763104351884",
                     "us-east-1": "763104351884",
                     "us-east-2": "763104351884",
@@ -1140,15 +1141,13 @@
                     "us-iso-east-1": "886529160074",
                     "us-isob-east-1": "094389454867",
                     "us-west-1": "763104351884",
-                    "us-west-2": "763104351884",
-                    "ca-west-1": "204538143572"
+                    "us-west-2": "763104351884"
                 },
                 "repository": "tensorflow-inference"
             },
             "2.1.1": {
                 "registries": {
                     "af-south-1": "626614931356",
-                    "il-central-1": "780543022126",
                     "ap-east-1": "871362719292",
                     "ap-northeast-1": "763104351884",
                     "ap-northeast-2": "763104351884",
@@ -1160,6 +1159,7 @@
                     "ap-southeast-3": "907027046896",
                     "ap-southeast-4": "457447274322",
                     "ca-central-1": "763104351884",
+                    "ca-west-1": "204538143572",
                     "cn-north-1": "727897471807",
                     "cn-northwest-1": "727897471807",
                     "eu-central-1": "763104351884",
@@ -1170,8 +1170,9 @@
                     "eu-west-1": "763104351884",
                     "eu-west-2": "763104351884",
                     "eu-west-3": "763104351884",
-                    "me-south-1": "217643126080",
+                    "il-central-1": "780543022126",
                     "me-central-1": "914824155844",
+                    "me-south-1": "217643126080",
                     "sa-east-1": "763104351884",
                     "us-east-1": "763104351884",
                     "us-east-2": "763104351884",
@@ -1180,15 +1181,13 @@
                     "us-iso-east-1": "886529160074",
                     "us-isob-east-1": "094389454867",
                     "us-west-1": "763104351884",
-                    "us-west-2": "763104351884",
-                    "ca-west-1": "204538143572"
+                    "us-west-2": "763104351884"
                 },
                 "repository": "tensorflow-inference"
             },
             "2.1.2": {
                 "registries": {
                     "af-south-1": "626614931356",
-                    "il-central-1": "780543022126",
                     "ap-east-1": "871362719292",
                     "ap-northeast-1": "763104351884",
                     "ap-northeast-2": "763104351884",
@@ -1200,6 +1199,7 @@
                     "ap-southeast-3": "907027046896",
                     "ap-southeast-4": "457447274322",
                     "ca-central-1": "763104351884",
+                    "ca-west-1": "204538143572",
                     "cn-north-1": "727897471807",
                     "cn-northwest-1": "727897471807",
                     "eu-central-1": "763104351884",
@@ -1210,8 +1210,9 @@
                     "eu-west-1": "763104351884",
                     "eu-west-2": "763104351884",
                     "eu-west-3": "763104351884",
-                    "me-south-1": "217643126080",
+                    "il-central-1": "780543022126",
                     "me-central-1": "914824155844",
+                    "me-south-1": "217643126080",
                     "sa-east-1": "763104351884",
                     "us-east-1": "763104351884",
                     "us-east-2": "763104351884",
@@ -1220,15 +1221,13 @@
                     "us-iso-east-1": "886529160074",
                     "us-isob-east-1": "094389454867",
                     "us-west-1": "763104351884",
-                    "us-west-2": "763104351884",
-                    "ca-west-1": "204538143572"
+                    "us-west-2": "763104351884"
                 },
                 "repository": "tensorflow-inference"
             },
             "2.1.3": {
                 "registries": {
                     "af-south-1": "626614931356",
-                    "il-central-1": "780543022126",
                     "ap-east-1": "871362719292",
                     "ap-northeast-1": "763104351884",
                     "ap-northeast-2": "763104351884",
@@ -1240,6 +1239,7 @@
                     "ap-southeast-3": "907027046896",
                     "ap-southeast-4": "457447274322",
                     "ca-central-1": "763104351884",
+                    "ca-west-1": "204538143572",
                     "cn-north-1": "727897471807",
                     "cn-northwest-1": "727897471807",
                     "eu-central-1": "763104351884",
@@ -1250,8 +1250,9 @@
                     "eu-west-1": "763104351884",
                     "eu-west-2": "763104351884",
                     "eu-west-3": "763104351884",
-                    "me-south-1": "217643126080",
+                    "il-central-1": "780543022126",
                     "me-central-1": "914824155844",
+                    "me-south-1": "217643126080",
                     "sa-east-1": "763104351884",
                     "us-east-1": "763104351884",
                     "us-east-2": "763104351884",
@@ -1260,15 +1261,13 @@
                     "us-iso-east-1": "886529160074",
                     "us-isob-east-1": "094389454867",
                     "us-west-1": "763104351884",
-                    "us-west-2": "763104351884",
-                    "ca-west-1": "204538143572"
+                    "us-west-2": "763104351884"
                 },
                 "repository": "tensorflow-inference"
             },
             "2.2.0": {
                 "registries": {
                     "af-south-1": "626614931356",
-                    "il-central-1": "780543022126",
                     "ap-east-1": "871362719292",
                     "ap-northeast-1": "763104351884",
                     "ap-northeast-2": "763104351884",
@@ -1280,6 +1279,7 @@
                     "ap-southeast-3": "907027046896",
                     "ap-southeast-4": "457447274322",
                     "ca-central-1": "763104351884",
+                    "ca-west-1": "204538143572",
                     "cn-north-1": "727897471807",
                     "cn-northwest-1": "727897471807",
                     "eu-central-1": "763104351884",
@@ -1290,8 +1290,9 @@
                     "eu-west-1": "763104351884",
                     "eu-west-2": "763104351884",
                     "eu-west-3": "763104351884",
-                    "me-south-1": "217643126080",
+                    "il-central-1": "780543022126",
                     "me-central-1": "914824155844",
+                    "me-south-1": "217643126080",
                     "sa-east-1": "763104351884",
                     "us-east-1": "763104351884",
                     "us-east-2": "763104351884",
@@ -1300,15 +1301,13 @@
                     "us-iso-east-1": "886529160074",
                     "us-isob-east-1": "094389454867",
                     "us-west-1": "763104351884",
-                    "us-west-2": "763104351884",
-                    "ca-west-1": "204538143572"
+                    "us-west-2": "763104351884"
                 },
                 "repository": "tensorflow-inference"
             },
             "2.2.1": {
                 "registries": {
                     "af-south-1": "626614931356",
-                    "il-central-1": "780543022126",
                     "ap-east-1": "871362719292",
                     "ap-northeast-1": "763104351884",
                     "ap-northeast-2": "763104351884",
@@ -1320,6 +1319,7 @@
                     "ap-southeast-3": "907027046896",
                     "ap-southeast-4": "457447274322",
                     "ca-central-1": "763104351884",
+                    "ca-west-1": "204538143572",
                     "cn-north-1": "727897471807",
                     "cn-northwest-1": "727897471807",
                     "eu-central-1": "763104351884",
@@ -1330,8 +1330,9 @@
                     "eu-west-1": "763104351884",
                     "eu-west-2": "763104351884",
                     "eu-west-3": "763104351884",
-                    "me-south-1": "217643126080",
+                    "il-central-1": "780543022126",
                     "me-central-1": "914824155844",
+                    "me-south-1": "217643126080",
                     "sa-east-1": "763104351884",
                     "us-east-1": "763104351884",
                     "us-east-2": "763104351884",
@@ -1340,15 +1341,13 @@
                     "us-iso-east-1": "886529160074",
                     "us-isob-east-1": "094389454867",
                     "us-west-1": "763104351884",
-                    "us-west-2": "763104351884",
-                    "ca-west-1": "204538143572"
+                    "us-west-2": "763104351884"
                 },
                 "repository": "tensorflow-inference"
             },
             "2.2.2": {
                 "registries": {
                     "af-south-1": "626614931356",
-                    "il-central-1": "780543022126",
                     "ap-east-1": "871362719292",
                     "ap-northeast-1": "763104351884",
                     "ap-northeast-2": "763104351884",
@@ -1360,6 +1359,7 @@
                     "ap-southeast-3": "907027046896",
                     "ap-southeast-4": "457447274322",
                     "ca-central-1": "763104351884",
+                    "ca-west-1": "204538143572",
                     "cn-north-1": "727897471807",
                     "cn-northwest-1": "727897471807",
                     "eu-central-1": "763104351884",
@@ -1370,8 +1370,9 @@
                     "eu-west-1": "763104351884",
                     "eu-west-2": "763104351884",
                     "eu-west-3": "763104351884",
-                    "me-south-1": "217643126080",
+                    "il-central-1": "780543022126",
                     "me-central-1": "914824155844",
+                    "me-south-1": "217643126080",
                     "sa-east-1": "763104351884",
                     "us-east-1": "763104351884",
                     "us-east-2": "763104351884",
@@ -1380,15 +1381,13 @@
                     "us-iso-east-1": "886529160074",
                     "us-isob-east-1": "094389454867",
                     "us-west-1": "763104351884",
-                    "us-west-2": "763104351884",
-                    "ca-west-1": "204538143572"
+                    "us-west-2": "763104351884"
                 },
                 "repository": "tensorflow-inference"
             },
             "2.3.0": {
                 "registries": {
                     "af-south-1": "626614931356",
-                    "il-central-1": "780543022126",
                     "ap-east-1": "871362719292",
                     "ap-northeast-1": "763104351884",
                     "ap-northeast-2": "763104351884",
@@ -1400,6 +1399,7 @@
                     "ap-southeast-3": "907027046896",
                     "ap-southeast-4": "457447274322",
                     "ca-central-1": "763104351884",
+                    "ca-west-1": "204538143572",
                     "cn-north-1": "727897471807",
                     "cn-northwest-1": "727897471807",
                     "eu-central-1": "763104351884",
@@ -1410,8 +1410,9 @@
                     "eu-west-1": "763104351884",
                     "eu-west-2": "763104351884",
                     "eu-west-3": "763104351884",
-                    "me-south-1": "217643126080",
+                    "il-central-1": "780543022126",
                     "me-central-1": "914824155844",
+                    "me-south-1": "217643126080",
                     "sa-east-1": "763104351884",
                     "us-east-1": "763104351884",
                     "us-east-2": "763104351884",
@@ -1420,15 +1421,13 @@
                     "us-iso-east-1": "886529160074",
                     "us-isob-east-1": "094389454867",
                     "us-west-1": "763104351884",
-                    "us-west-2": "763104351884",
-                    "ca-west-1": "204538143572"
+                    "us-west-2": "763104351884"
                 },
                 "repository": "tensorflow-inference"
             },
             "2.3.1": {
                 "registries": {
                     "af-south-1": "626614931356",
-                    "il-central-1": "780543022126",
                     "ap-east-1": "871362719292",
                     "ap-northeast-1": "763104351884",
                     "ap-northeast-2": "763104351884",
@@ -1440,6 +1439,7 @@
                     "ap-southeast-3": "907027046896",
                     "ap-southeast-4": "457447274322",
                     "ca-central-1": "763104351884",
+                    "ca-west-1": "204538143572",
                     "cn-north-1": "727897471807",
                     "cn-northwest-1": "727897471807",
                     "eu-central-1": "763104351884",
@@ -1450,8 +1450,9 @@
                     "eu-west-1": "763104351884",
                     "eu-west-2": "763104351884",
                     "eu-west-3": "763104351884",
-                    "me-south-1": "217643126080",
+                    "il-central-1": "780543022126",
                     "me-central-1": "914824155844",
+                    "me-south-1": "217643126080",
                     "sa-east-1": "763104351884",
                     "us-east-1": "763104351884",
                     "us-east-2": "763104351884",
@@ -1460,15 +1461,13 @@
                     "us-iso-east-1": "886529160074",
                     "us-isob-east-1": "094389454867",
                     "us-west-1": "763104351884",
-                    "us-west-2": "763104351884",
-                    "ca-west-1": "204538143572"
+                    "us-west-2": "763104351884"
                 },
                 "repository": "tensorflow-inference"
             },
             "2.3.2": {
                 "registries": {
                     "af-south-1": "626614931356",
-                    "il-central-1": "780543022126",
                     "ap-east-1": "871362719292",
                     "ap-northeast-1": "763104351884",
                     "ap-northeast-2": "763104351884",
@@ -1480,6 +1479,7 @@
                     "ap-southeast-3": "907027046896",
                     "ap-southeast-4": "457447274322",
                     "ca-central-1": "763104351884",
+                    "ca-west-1": "204538143572",
                     "cn-north-1": "727897471807",
                     "cn-northwest-1": "727897471807",
                     "eu-central-1": "763104351884",
@@ -1490,8 +1490,9 @@
                     "eu-west-1": "763104351884",
                     "eu-west-2": "763104351884",
                     "eu-west-3": "763104351884",
-                    "me-south-1": "217643126080",
+                    "il-central-1": "780543022126",
                     "me-central-1": "914824155844",
+                    "me-south-1": "217643126080",
                     "sa-east-1": "763104351884",
                     "us-east-1": "763104351884",
                     "us-east-2": "763104351884",
@@ -1500,15 +1501,13 @@
                     "us-iso-east-1": "886529160074",
                     "us-isob-east-1": "094389454867",
                     "us-west-1": "763104351884",
-                    "us-west-2": "763104351884",
-                    "ca-west-1": "204538143572"
+                    "us-west-2": "763104351884"
                 },
                 "repository": "tensorflow-inference"
             },
             "2.4.1": {
                 "registries": {
                     "af-south-1": "626614931356",
-                    "il-central-1": "780543022126",
                     "ap-east-1": "871362719292",
                     "ap-northeast-1": "763104351884",
                     "ap-northeast-2": "763104351884",
@@ -1520,6 +1519,7 @@
                     "ap-southeast-3": "907027046896",
                     "ap-southeast-4": "457447274322",
                     "ca-central-1": "763104351884",
+                    "ca-west-1": "204538143572",
                     "cn-north-1": "727897471807",
                     "cn-northwest-1": "727897471807",
                     "eu-central-1": "763104351884",
@@ -1530,8 +1530,9 @@
                     "eu-west-1": "763104351884",
                     "eu-west-2": "763104351884",
                     "eu-west-3": "763104351884",
-                    "me-south-1": "217643126080",
+                    "il-central-1": "780543022126",
                     "me-central-1": "914824155844",
+                    "me-south-1": "217643126080",
                     "sa-east-1": "763104351884",
                     "us-east-1": "763104351884",
                     "us-east-2": "763104351884",
@@ -1540,15 +1541,13 @@
                     "us-iso-east-1": "886529160074",
                     "us-isob-east-1": "094389454867",
                     "us-west-1": "763104351884",
-                    "us-west-2": "763104351884",
-                    "ca-west-1": "204538143572"
+                    "us-west-2": "763104351884"
                 },
                 "repository": "tensorflow-inference"
             },
             "2.4.3": {
                 "registries": {
                     "af-south-1": "626614931356",
-                    "il-central-1": "780543022126",
                     "ap-east-1": "871362719292",
                     "ap-northeast-1": "763104351884",
                     "ap-northeast-2": "763104351884",
@@ -1560,6 +1559,7 @@
                     "ap-southeast-3": "907027046896",
                     "ap-southeast-4": "457447274322",
                     "ca-central-1": "763104351884",
+                    "ca-west-1": "204538143572",
                     "cn-north-1": "727897471807",
                     "cn-northwest-1": "727897471807",
                     "eu-central-1": "763104351884",
@@ -1570,8 +1570,9 @@
                     "eu-west-1": "763104351884",
                     "eu-west-2": "763104351884",
                     "eu-west-3": "763104351884",
-                    "me-south-1": "217643126080",
+                    "il-central-1": "780543022126",
                     "me-central-1": "914824155844",
+                    "me-south-1": "217643126080",
                     "sa-east-1": "763104351884",
                     "us-east-1": "763104351884",
                     "us-east-2": "763104351884",
@@ -1580,15 +1581,13 @@
                     "us-iso-east-1": "886529160074",
                     "us-isob-east-1": "094389454867",
                     "us-west-1": "763104351884",
-                    "us-west-2": "763104351884",
-                    "ca-west-1": "204538143572"
+                    "us-west-2": "763104351884"
                 },
                 "repository": "tensorflow-inference"
             },
             "2.5.1": {
                 "registries": {
                     "af-south-1": "626614931356",
-                    "il-central-1": "780543022126",
                     "ap-east-1": "871362719292",
                     "ap-northeast-1": "763104351884",
                     "ap-northeast-2": "763104351884",
@@ -1600,6 +1599,7 @@
                     "ap-southeast-3": "907027046896",
                     "ap-southeast-4": "457447274322",
                     "ca-central-1": "763104351884",
+                    "ca-west-1": "204538143572",
                     "cn-north-1": "727897471807",
                     "cn-northwest-1": "727897471807",
                     "eu-central-1": "763104351884",
@@ -1610,8 +1610,9 @@
                     "eu-west-1": "763104351884",
                     "eu-west-2": "763104351884",
                     "eu-west-3": "763104351884",
-                    "me-south-1": "217643126080",
+                    "il-central-1": "780543022126",
                     "me-central-1": "914824155844",
+                    "me-south-1": "217643126080",
                     "sa-east-1": "763104351884",
                     "us-east-1": "763104351884",
                     "us-east-2": "763104351884",
@@ -1620,15 +1621,13 @@
                     "us-iso-east-1": "886529160074",
                     "us-isob-east-1": "094389454867",
                     "us-west-1": "763104351884",
-                    "us-west-2": "763104351884",
-                    "ca-west-1": "204538143572"
+                    "us-west-2": "763104351884"
                 },
                 "repository": "tensorflow-inference"
             },
             "2.6.0": {
                 "registries": {
                     "af-south-1": "626614931356",
-                    "il-central-1": "780543022126",
                     "ap-east-1": "871362719292",
                     "ap-northeast-1": "763104351884",
                     "ap-northeast-2": "763104351884",
@@ -1640,6 +1639,7 @@
                     "ap-southeast-3": "907027046896",
                     "ap-southeast-4": "457447274322",
                     "ca-central-1": "763104351884",
+                    "ca-west-1": "204538143572",
                     "cn-north-1": "727897471807",
                     "cn-northwest-1": "727897471807",
                     "eu-central-1": "763104351884",
@@ -1650,8 +1650,9 @@
                     "eu-west-1": "763104351884",
                     "eu-west-2": "763104351884",
                     "eu-west-3": "763104351884",
-                    "me-south-1": "217643126080",
+                    "il-central-1": "780543022126",
                     "me-central-1": "914824155844",
+                    "me-south-1": "217643126080",
                     "sa-east-1": "763104351884",
                     "us-east-1": "763104351884",
                     "us-east-2": "763104351884",
@@ -1660,15 +1661,13 @@
                     "us-iso-east-1": "886529160074",
                     "us-isob-east-1": "094389454867",
                     "us-west-1": "763104351884",
-                    "us-west-2": "763104351884",
-                    "ca-west-1": "204538143572"
+                    "us-west-2": "763104351884"
                 },
                 "repository": "tensorflow-inference"
             },
             "2.6.3": {
                 "registries": {
                     "af-south-1": "626614931356",
-                    "il-central-1": "780543022126",
                     "ap-east-1": "871362719292",
                     "ap-northeast-1": "763104351884",
                     "ap-northeast-2": "763104351884",
@@ -1680,6 +1679,7 @@
                     "ap-southeast-3": "907027046896",
                     "ap-southeast-4": "457447274322",
                     "ca-central-1": "763104351884",
+                    "ca-west-1": "204538143572",
                     "cn-north-1": "727897471807",
                     "cn-northwest-1": "727897471807",
                     "eu-central-1": "763104351884",
@@ -1690,8 +1690,9 @@
                     "eu-west-1": "763104351884",
                     "eu-west-2": "763104351884",
                     "eu-west-3": "763104351884",
-                    "me-south-1": "217643126080",
+                    "il-central-1": "780543022126",
                     "me-central-1": "914824155844",
+                    "me-south-1": "217643126080",
                     "sa-east-1": "763104351884",
                     "us-east-1": "763104351884",
                     "us-east-2": "763104351884",
@@ -1700,15 +1701,13 @@
                     "us-iso-east-1": "886529160074",
                     "us-isob-east-1": "094389454867",
                     "us-west-1": "763104351884",
-                    "us-west-2": "763104351884",
-                    "ca-west-1": "204538143572"
+                    "us-west-2": "763104351884"
                 },
                 "repository": "tensorflow-inference"
             },
             "2.7.0": {
                 "registries": {
                     "af-south-1": "626614931356",
-                    "il-central-1": "780543022126",
                     "ap-east-1": "871362719292",
                     "ap-northeast-1": "763104351884",
                     "ap-northeast-2": "763104351884",
@@ -1720,6 +1719,7 @@
                     "ap-southeast-3": "907027046896",
                     "ap-southeast-4": "457447274322",
                     "ca-central-1": "763104351884",
+                    "ca-west-1": "204538143572",
                     "cn-north-1": "727897471807",
                     "cn-northwest-1": "727897471807",
                     "eu-central-1": "763104351884",
@@ -1730,8 +1730,9 @@
                     "eu-west-1": "763104351884",
                     "eu-west-2": "763104351884",
                     "eu-west-3": "763104351884",
-                    "me-south-1": "217643126080",
+                    "il-central-1": "780543022126",
                     "me-central-1": "914824155844",
+                    "me-south-1": "217643126080",
                     "sa-east-1": "763104351884",
                     "us-east-1": "763104351884",
                     "us-east-2": "763104351884",
@@ -1740,15 +1741,13 @@
                     "us-iso-east-1": "886529160074",
                     "us-isob-east-1": "094389454867",
                     "us-west-1": "763104351884",
-                    "us-west-2": "763104351884",
-                    "ca-west-1": "204538143572"
+                    "us-west-2": "763104351884"
                 },
                 "repository": "tensorflow-inference"
             },
             "2.8.0": {
                 "registries": {
                     "af-south-1": "626614931356",
-                    "il-central-1": "780543022126",
                     "ap-east-1": "871362719292",
                     "ap-northeast-1": "763104351884",
                     "ap-northeast-2": "763104351884",
@@ -1760,6 +1759,7 @@
                     "ap-southeast-3": "907027046896",
                     "ap-southeast-4": "457447274322",
                     "ca-central-1": "763104351884",
+                    "ca-west-1": "204538143572",
                     "cn-north-1": "727897471807",
                     "cn-northwest-1": "727897471807",
                     "eu-central-1": "763104351884",
@@ -1770,8 +1770,9 @@
                     "eu-west-1": "763104351884",
                     "eu-west-2": "763104351884",
                     "eu-west-3": "763104351884",
-                    "me-south-1": "217643126080",
+                    "il-central-1": "780543022126",
                     "me-central-1": "914824155844",
+                    "me-south-1": "217643126080",
                     "sa-east-1": "763104351884",
                     "us-east-1": "763104351884",
                     "us-east-2": "763104351884",
@@ -1780,15 +1781,13 @@
                     "us-iso-east-1": "886529160074",
                     "us-isob-east-1": "094389454867",
                     "us-west-1": "763104351884",
-                    "us-west-2": "763104351884",
-                    "ca-west-1": "204538143572"
+                    "us-west-2": "763104351884"
                 },
                 "repository": "tensorflow-inference"
             },
             "2.8.4": {
                 "registries": {
                     "af-south-1": "626614931356",
-                    "il-central-1": "780543022126",
                     "ap-east-1": "871362719292",
                     "ap-northeast-1": "763104351884",
                     "ap-northeast-2": "763104351884",
@@ -1800,6 +1799,7 @@
                     "ap-southeast-3": "907027046896",
                     "ap-southeast-4": "457447274322",
                     "ca-central-1": "763104351884",
+                    "ca-west-1": "204538143572",
                     "cn-north-1": "727897471807",
                     "cn-northwest-1": "727897471807",
                     "eu-central-1": "763104351884",
@@ -1810,8 +1810,9 @@
                     "eu-west-1": "763104351884",
                     "eu-west-2": "763104351884",
                     "eu-west-3": "763104351884",
-                    "me-south-1": "217643126080",
+                    "il-central-1": "780543022126",
                     "me-central-1": "914824155844",
+                    "me-south-1": "217643126080",
                     "sa-east-1": "763104351884",
                     "us-east-1": "763104351884",
                     "us-east-2": "763104351884",
@@ -1820,15 +1821,13 @@
                     "us-iso-east-1": "886529160074",
                     "us-isob-east-1": "094389454867",
                     "us-west-1": "763104351884",
-                    "us-west-2": "763104351884",
-                    "ca-west-1": "204538143572"
+                    "us-west-2": "763104351884"
                 },
                 "repository": "tensorflow-inference"
             },
             "2.9.2": {
                 "registries": {
                     "af-south-1": "626614931356",
-                    "il-central-1": "780543022126",
                     "ap-east-1": "871362719292",
                     "ap-northeast-1": "763104351884",
                     "ap-northeast-2": "763104351884",
@@ -1840,6 +1839,7 @@
                     "ap-southeast-3": "907027046896",
                     "ap-southeast-4": "457447274322",
                     "ca-central-1": "763104351884",
+                    "ca-west-1": "204538143572",
                     "cn-north-1": "727897471807",
                     "cn-northwest-1": "727897471807",
                     "eu-central-1": "763104351884",
@@ -1850,6 +1850,8 @@
                     "eu-west-1": "763104351884",
                     "eu-west-2": "763104351884",
                     "eu-west-3": "763104351884",
+                    "il-central-1": "780543022126",
+                    "me-central-1": "914824155844",
                     "me-south-1": "217643126080",
                     "sa-east-1": "763104351884",
                     "us-east-1": "763104351884",
@@ -1859,15 +1861,13 @@
                     "us-iso-east-1": "886529160074",
                     "us-isob-east-1": "094389454867",
                     "us-west-1": "763104351884",
-                    "us-west-2": "763104351884",
-                    "ca-west-1": "204538143572"
+                    "us-west-2": "763104351884"
                 },
                 "repository": "tensorflow-inference"
             },
             "2.9.3": {
                 "registries": {
                     "af-south-1": "626614931356",
-                    "il-central-1": "780543022126",
                     "ap-east-1": "871362719292",
                     "ap-northeast-1": "763104351884",
                     "ap-northeast-2": "763104351884",
@@ -1879,6 +1879,7 @@
                     "ap-southeast-3": "907027046896",
                     "ap-southeast-4": "457447274322",
                     "ca-central-1": "763104351884",
+                    "ca-west-1": "204538143572",
                     "cn-north-1": "727897471807",
                     "cn-northwest-1": "727897471807",
                     "eu-central-1": "763104351884",
@@ -1889,6 +1890,8 @@
                     "eu-west-1": "763104351884",
                     "eu-west-2": "763104351884",
                     "eu-west-3": "763104351884",
+                    "il-central-1": "780543022126",
+                    "me-central-1": "914824155844",
                     "me-south-1": "217643126080",
                     "sa-east-1": "763104351884",
                     "us-east-1": "763104351884",
@@ -1898,15 +1901,13 @@
                     "us-iso-east-1": "886529160074",
                     "us-isob-east-1": "094389454867",
                     "us-west-1": "763104351884",
-                    "us-west-2": "763104351884",
-                    "ca-west-1": "204538143572"
+                    "us-west-2": "763104351884"
                 },
                 "repository": "tensorflow-inference"
             },
             "2.10.0": {
                 "registries": {
                     "af-south-1": "626614931356",
-                    "il-central-1": "780543022126",
                     "ap-east-1": "871362719292",
                     "ap-northeast-1": "763104351884",
                     "ap-northeast-2": "763104351884",
@@ -1918,6 +1919,7 @@
                     "ap-southeast-3": "907027046896",
                     "ap-southeast-4": "457447274322",
                     "ca-central-1": "763104351884",
+                    "ca-west-1": "204538143572",
                     "cn-north-1": "727897471807",
                     "cn-northwest-1": "727897471807",
                     "eu-central-1": "763104351884",
@@ -1928,6 +1930,8 @@
                     "eu-west-1": "763104351884",
                     "eu-west-2": "763104351884",
                     "eu-west-3": "763104351884",
+                    "il-central-1": "780543022126",
+                    "me-central-1": "914824155844",
                     "me-south-1": "217643126080",
                     "sa-east-1": "763104351884",
                     "us-east-1": "763104351884",
@@ -1937,15 +1941,13 @@
                     "us-iso-east-1": "886529160074",
                     "us-isob-east-1": "094389454867",
                     "us-west-1": "763104351884",
-                    "us-west-2": "763104351884",
-                    "ca-west-1": "204538143572"
+                    "us-west-2": "763104351884"
                 },
                 "repository": "tensorflow-inference"
             },
             "2.10.1": {
                 "registries": {
                     "af-south-1": "626614931356",
-                    "il-central-1": "780543022126",
                     "ap-east-1": "871362719292",
                     "ap-northeast-1": "763104351884",
                     "ap-northeast-2": "763104351884",
@@ -1957,6 +1959,7 @@
                     "ap-southeast-3": "907027046896",
                     "ap-southeast-4": "457447274322",
                     "ca-central-1": "763104351884",
+                    "ca-west-1": "204538143572",
                     "cn-north-1": "727897471807",
                     "cn-northwest-1": "727897471807",
                     "eu-central-1": "763104351884",
@@ -1967,6 +1970,8 @@
                     "eu-west-1": "763104351884",
                     "eu-west-2": "763104351884",
                     "eu-west-3": "763104351884",
+                    "il-central-1": "780543022126",
+                    "me-central-1": "914824155844",
                     "me-south-1": "217643126080",
                     "sa-east-1": "763104351884",
                     "us-east-1": "763104351884",
@@ -1976,15 +1981,13 @@
                     "us-iso-east-1": "886529160074",
                     "us-isob-east-1": "094389454867",
                     "us-west-1": "763104351884",
-                    "us-west-2": "763104351884",
-                    "ca-west-1": "204538143572"
+                    "us-west-2": "763104351884"
                 },
                 "repository": "tensorflow-inference"
             },
             "2.11.0": {
                 "registries": {
                     "af-south-1": "626614931356",
-                    "il-central-1": "780543022126",
                     "ap-east-1": "871362719292",
                     "ap-northeast-1": "763104351884",
                     "ap-northeast-2": "763104351884",
@@ -1996,6 +1999,7 @@
                     "ap-southeast-3": "907027046896",
                     "ap-southeast-4": "457447274322",
                     "ca-central-1": "763104351884",
+                    "ca-west-1": "204538143572",
                     "cn-north-1": "727897471807",
                     "cn-northwest-1": "727897471807",
                     "eu-central-1": "763104351884",
@@ -2006,6 +2010,8 @@
                     "eu-west-1": "763104351884",
                     "eu-west-2": "763104351884",
                     "eu-west-3": "763104351884",
+                    "il-central-1": "780543022126",
+                    "me-central-1": "914824155844",
                     "me-south-1": "217643126080",
                     "sa-east-1": "763104351884",
                     "us-east-1": "763104351884",
@@ -2015,15 +2021,13 @@
                     "us-iso-east-1": "886529160074",
                     "us-isob-east-1": "094389454867",
                     "us-west-1": "763104351884",
-                    "us-west-2": "763104351884",
-                    "ca-west-1": "204538143572"
+                    "us-west-2": "763104351884"
                 },
                 "repository": "tensorflow-inference"
             },
             "2.11.1": {
                 "registries": {
                     "af-south-1": "626614931356",
-                    "il-central-1": "780543022126",
                     "ap-east-1": "871362719292",
                     "ap-northeast-1": "763104351884",
                     "ap-northeast-2": "763104351884",
@@ -2035,6 +2039,7 @@
                     "ap-southeast-3": "907027046896",
                     "ap-southeast-4": "457447274322",
                     "ca-central-1": "763104351884",
+                    "ca-west-1": "204538143572",
                     "cn-north-1": "727897471807",
                     "cn-northwest-1": "727897471807",
                     "eu-central-1": "763104351884",
@@ -2045,6 +2050,8 @@
                     "eu-west-1": "763104351884",
                     "eu-west-2": "763104351884",
                     "eu-west-3": "763104351884",
+                    "il-central-1": "780543022126",
+                    "me-central-1": "914824155844",
                     "me-south-1": "217643126080",
                     "sa-east-1": "763104351884",
                     "us-east-1": "763104351884",
@@ -2054,15 +2061,13 @@
                     "us-iso-east-1": "886529160074",
                     "us-isob-east-1": "094389454867",
                     "us-west-1": "763104351884",
-                    "us-west-2": "763104351884",
-                    "ca-west-1": "204538143572"
+                    "us-west-2": "763104351884"
                 },
                 "repository": "tensorflow-inference"
             },
             "2.12.1": {
                 "registries": {
                     "af-south-1": "626614931356",
-                    "il-central-1": "780543022126",
                     "ap-east-1": "871362719292",
                     "ap-northeast-1": "763104351884",
                     "ap-northeast-2": "763104351884",
@@ -2074,6 +2079,7 @@
                     "ap-southeast-3": "907027046896",
                     "ap-southeast-4": "457447274322",
                     "ca-central-1": "763104351884",
+                    "ca-west-1": "204538143572",
                     "cn-north-1": "727897471807",
                     "cn-northwest-1": "727897471807",
                     "eu-central-1": "763104351884",
@@ -2084,6 +2090,8 @@
                     "eu-west-1": "763104351884",
                     "eu-west-2": "763104351884",
                     "eu-west-3": "763104351884",
+                    "il-central-1": "780543022126",
+                    "me-central-1": "914824155844",
                     "me-south-1": "217643126080",
                     "sa-east-1": "763104351884",
                     "us-east-1": "763104351884",
@@ -2093,15 +2101,13 @@
                     "us-iso-east-1": "886529160074",
                     "us-isob-east-1": "094389454867",
                     "us-west-1": "763104351884",
-                    "us-west-2": "763104351884",
-                    "ca-west-1": "204538143572"
+                    "us-west-2": "763104351884"
                 },
                 "repository": "tensorflow-inference"
             },
             "2.13.0": {
                 "registries": {
                     "af-south-1": "626614931356",
-                    "il-central-1": "780543022126",
                     "ap-east-1": "871362719292",
                     "ap-northeast-1": "763104351884",
                     "ap-northeast-2": "763104351884",
@@ -2113,6 +2119,7 @@
                     "ap-southeast-3": "907027046896",
                     "ap-southeast-4": "457447274322",
                     "ca-central-1": "763104351884",
+                    "ca-west-1": "204538143572",
                     "cn-north-1": "727897471807",
                     "cn-northwest-1": "727897471807",
                     "eu-central-1": "763104351884",
@@ -2123,17 +2130,18 @@
                     "eu-west-1": "763104351884",
                     "eu-west-2": "763104351884",
                     "eu-west-3": "763104351884",
+                    "il-central-1": "780543022126",
+                    "me-central-1": "914824155844",
                     "me-south-1": "217643126080",
                     "sa-east-1": "763104351884",
                     "us-east-1": "763104351884",
                     "us-east-2": "763104351884",
-                    "us-iso-east-1": "886529160074",
-                    "us-isob-east-1": "094389454867",
                     "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"
+                    "us-west-2": "763104351884"
                 },
                 "repository": "tensorflow-inference"
             },
@@ -2191,12 +2199,14 @@
         },
         "versions": {
             "2.9.1": {
+                "container_version": {
+                    "cpu": "ubuntu20.04"
+                },
                 "py_versions": [
                     "py38"
                 ],
                 "registries": {
                     "af-south-1": "626614931356",
-                    "il-central-1": "780543022126",
                     "ap-east-1": "871362719292",
                     "ap-northeast-1": "763104351884",
                     "ap-northeast-2": "763104351884",
@@ -2208,16 +2218,19 @@
                     "ap-southeast-3": "907027046896",
                     "ap-southeast-4": "457447274322",
                     "ca-central-1": "763104351884",
+                    "ca-west-1": "204538143572",
                     "cn-north-1": "727897471807",
                     "cn-northwest-1": "727897471807",
                     "eu-central-1": "763104351884",
                     "eu-central-2": "380420809688",
                     "eu-north-1": "763104351884",
+                    "eu-south-1": "692866216735",
+                    "eu-south-2": "503227376785",
                     "eu-west-1": "763104351884",
                     "eu-west-2": "763104351884",
                     "eu-west-3": "763104351884",
-                    "eu-south-1": "692866216735",
-                    "eu-south-2": "503227376785",
+                    "il-central-1": "780543022126",
+                    "me-central-1": "914824155844",
                     "me-south-1": "217643126080",
                     "sa-east-1": "763104351884",
                     "us-east-1": "763104351884",
@@ -2227,21 +2240,19 @@
                     "us-iso-east-1": "886529160074",
                     "us-isob-east-1": "094389454867",
                     "us-west-1": "763104351884",
-                    "us-west-2": "763104351884",
-                    "ca-west-1": "204538143572"
+                    "us-west-2": "763104351884"
                 },
-                "repository": "tensorflow-inference-graviton",
-                "container_version": {
-                    "cpu": "ubuntu20.04"
-                }
+                "repository": "tensorflow-inference-graviton"
             },
             "2.12.1": {
+                "container_version": {
+                    "cpu": "ubuntu20.04"
+                },
                 "py_versions": [
                     "py310"
                 ],
                 "registries": {
                     "af-south-1": "626614931356",
-                    "il-central-1": "780543022126",
                     "ap-east-1": "871362719292",
                     "ap-northeast-1": "763104351884",
                     "ap-northeast-2": "763104351884",
@@ -2253,16 +2264,19 @@
                     "ap-southeast-3": "907027046896",
                     "ap-southeast-4": "457447274322",
                     "ca-central-1": "763104351884",
+                    "ca-west-1": "204538143572",
                     "cn-north-1": "727897471807",
                     "cn-northwest-1": "727897471807",
                     "eu-central-1": "763104351884",
                     "eu-central-2": "380420809688",
                     "eu-north-1": "763104351884",
+                    "eu-south-1": "692866216735",
+                    "eu-south-2": "503227376785",
                     "eu-west-1": "763104351884",
                     "eu-west-2": "763104351884",
                     "eu-west-3": "763104351884",
-                    "eu-south-1": "692866216735",
-                    "eu-south-2": "503227376785",
+                    "il-central-1": "780543022126",
+                    "me-central-1": "914824155844",
                     "me-south-1": "217643126080",
                     "sa-east-1": "763104351884",
                     "us-east-1": "763104351884",
@@ -2272,21 +2286,19 @@
                     "us-iso-east-1": "886529160074",
                     "us-isob-east-1": "094389454867",
                     "us-west-1": "763104351884",
-                    "us-west-2": "763104351884",
-                    "ca-west-1": "204538143572"
+                    "us-west-2": "763104351884"
                 },
-                "repository": "tensorflow-inference-graviton",
-                "container_version": {
-                    "cpu": "ubuntu20.04"
-                }
+                "repository": "tensorflow-inference-graviton"
             },
             "2.13.0": {
+                "container_version": {
+                    "cpu": "ubuntu20.04"
+                },
                 "py_versions": [
                     "py310"
                 ],
                 "registries": {
                     "af-south-1": "626614931356",
-                    "il-central-1": "780543022126",
                     "ap-east-1": "871362719292",
                     "ap-northeast-1": "763104351884",
                     "ap-northeast-2": "763104351884",
@@ -2298,16 +2310,19 @@
                     "ap-southeast-3": "907027046896",
                     "ap-southeast-4": "457447274322",
                     "ca-central-1": "763104351884",
+                    "ca-west-1": "204538143572",
                     "cn-north-1": "727897471807",
                     "cn-northwest-1": "727897471807",
                     "eu-central-1": "763104351884",
                     "eu-central-2": "380420809688",
                     "eu-north-1": "763104351884",
+                    "eu-south-1": "692866216735",
+                    "eu-south-2": "503227376785",
                     "eu-west-1": "763104351884",
                     "eu-west-2": "763104351884",
                     "eu-west-3": "763104351884",
-                    "eu-south-1": "692866216735",
-                    "eu-south-2": "503227376785",
+                    "il-central-1": "780543022126",
+                    "me-central-1": "914824155844",
                     "me-south-1": "217643126080",
                     "sa-east-1": "763104351884",
                     "us-east-1": "763104351884",
@@ -2317,21 +2332,19 @@
                     "us-iso-east-1": "886529160074",
                     "us-isob-east-1": "094389454867",
                     "us-west-1": "763104351884",
-                    "us-west-2": "763104351884",
-                    "ca-west-1": "204538143572"
+                    "us-west-2": "763104351884"
                 },
-                "repository": "tensorflow-inference-graviton",
-                "container_version": {
-                    "cpu": "ubuntu20.04"
-                }
+                "repository": "tensorflow-inference-graviton"
             },
             "2.14.1": {
+                "container_version": {
+                    "cpu": "ubuntu20.04"
+                },
                 "py_versions": [
                     "py310"
                 ],
                 "registries": {
                     "af-south-1": "626614931356",
-                    "il-central-1": "780543022126",
                     "ap-east-1": "871362719292",
                     "ap-northeast-1": "763104351884",
                     "ap-northeast-2": "763104351884",
@@ -2343,16 +2356,19 @@
                     "ap-southeast-3": "907027046896",
                     "ap-southeast-4": "457447274322",
                     "ca-central-1": "763104351884",
+                    "ca-west-1": "204538143572",
                     "cn-north-1": "727897471807",
                     "cn-northwest-1": "727897471807",
                     "eu-central-1": "763104351884",
                     "eu-central-2": "380420809688",
                     "eu-north-1": "763104351884",
+                    "eu-south-1": "692866216735",
+                    "eu-south-2": "503227376785",
                     "eu-west-1": "763104351884",
                     "eu-west-2": "763104351884",
                     "eu-west-3": "763104351884",
-                    "eu-south-1": "692866216735",
-                    "eu-south-2": "503227376785",
+                    "il-central-1": "780543022126",
+                    "me-central-1": "914824155844",
                     "me-south-1": "217643126080",
                     "sa-east-1": "763104351884",
                     "us-east-1": "763104351884",
@@ -2362,13 +2378,9 @@
                     "us-iso-east-1": "886529160074",
                     "us-isob-east-1": "094389454867",
                     "us-west-1": "763104351884",
-                    "us-west-2": "763104351884",
-                    "ca-west-1": "204538143572"
+                    "us-west-2": "763104351884"
                 },
-                "repository": "tensorflow-inference-graviton",
-                "container_version": {
-                    "cpu": "ubuntu20.04"
-                }
+                "repository": "tensorflow-inference-graviton"
             }
         }
     },
@@ -2378,18 +2390,18 @@
             "gpu"
         ],
         "version_aliases": {
-            "1.10": "1.10.0",
-            "1.11": "1.11.0",
-            "1.12": "1.12.0",
-            "1.13": "1.13.1",
-            "1.14": "1.14.0",
-            "1.15": "1.15.5",
             "1.4": "1.4.1",
             "1.5": "1.5.0",
             "1.6": "1.6.0",
             "1.7": "1.7.0",
             "1.8": "1.8.0",
             "1.9": "1.9.0",
+            "1.10": "1.10.0",
+            "1.11": "1.11.0",
+            "1.12": "1.12.0",
+            "1.13": "1.13.1",
+            "1.14": "1.14.0",
+            "1.15": "1.15.5",
             "2.0": "2.0.4",
             "2.1": "2.1.3",
             "2.2": "2.2.2",
@@ -2407,6 +2419,204 @@
             "2.14": "2.14.1"
         },
         "versions": {
+            "1.4.1": {
+                "py_versions": [
+                    "py2"
+                ],
+                "registries": {
+                    "af-south-1": "313743910680",
+                    "ap-east-1": "057415533634",
+                    "ap-northeast-1": "520713654638",
+                    "ap-northeast-2": "520713654638",
+                    "ap-south-1": "520713654638",
+                    "ap-southeast-1": "520713654638",
+                    "ap-southeast-2": "520713654638",
+                    "ca-central-1": "520713654638",
+                    "cn-north-1": "422961961927",
+                    "cn-northwest-1": "423003514399",
+                    "eu-central-1": "520713654638",
+                    "eu-north-1": "520713654638",
+                    "eu-south-1": "048378556238",
+                    "eu-west-1": "520713654638",
+                    "eu-west-2": "520713654638",
+                    "eu-west-3": "520713654638",
+                    "me-south-1": "724002660598",
+                    "sa-east-1": "520713654638",
+                    "us-east-1": "520713654638",
+                    "us-east-2": "520713654638",
+                    "us-gov-west-1": "246785580436",
+                    "us-iso-east-1": "744548109606",
+                    "us-isob-east-1": "453391408702",
+                    "us-west-1": "520713654638",
+                    "us-west-2": "520713654638"
+                },
+                "repository": "sagemaker-tensorflow"
+            },
+            "1.5.0": {
+                "py_versions": [
+                    "py2"
+                ],
+                "registries": {
+                    "af-south-1": "313743910680",
+                    "ap-east-1": "057415533634",
+                    "ap-northeast-1": "520713654638",
+                    "ap-northeast-2": "520713654638",
+                    "ap-south-1": "520713654638",
+                    "ap-southeast-1": "520713654638",
+                    "ap-southeast-2": "520713654638",
+                    "ca-central-1": "520713654638",
+                    "cn-north-1": "422961961927",
+                    "cn-northwest-1": "423003514399",
+                    "eu-central-1": "520713654638",
+                    "eu-north-1": "520713654638",
+                    "eu-south-1": "048378556238",
+                    "eu-west-1": "520713654638",
+                    "eu-west-2": "520713654638",
+                    "eu-west-3": "520713654638",
+                    "me-south-1": "724002660598",
+                    "sa-east-1": "520713654638",
+                    "us-east-1": "520713654638",
+                    "us-east-2": "520713654638",
+                    "us-gov-west-1": "246785580436",
+                    "us-iso-east-1": "744548109606",
+                    "us-isob-east-1": "453391408702",
+                    "us-west-1": "520713654638",
+                    "us-west-2": "520713654638"
+                },
+                "repository": "sagemaker-tensorflow"
+            },
+            "1.6.0": {
+                "py_versions": [
+                    "py2"
+                ],
+                "registries": {
+                    "af-south-1": "313743910680",
+                    "ap-east-1": "057415533634",
+                    "ap-northeast-1": "520713654638",
+                    "ap-northeast-2": "520713654638",
+                    "ap-south-1": "520713654638",
+                    "ap-southeast-1": "520713654638",
+                    "ap-southeast-2": "520713654638",
+                    "ca-central-1": "520713654638",
+                    "cn-north-1": "422961961927",
+                    "cn-northwest-1": "423003514399",
+                    "eu-central-1": "520713654638",
+                    "eu-north-1": "520713654638",
+                    "eu-south-1": "048378556238",
+                    "eu-west-1": "520713654638",
+                    "eu-west-2": "520713654638",
+                    "eu-west-3": "520713654638",
+                    "me-south-1": "724002660598",
+                    "sa-east-1": "520713654638",
+                    "us-east-1": "520713654638",
+                    "us-east-2": "520713654638",
+                    "us-gov-west-1": "246785580436",
+                    "us-iso-east-1": "744548109606",
+                    "us-isob-east-1": "453391408702",
+                    "us-west-1": "520713654638",
+                    "us-west-2": "520713654638"
+                },
+                "repository": "sagemaker-tensorflow"
+            },
+            "1.7.0": {
+                "py_versions": [
+                    "py2"
+                ],
+                "registries": {
+                    "af-south-1": "313743910680",
+                    "ap-east-1": "057415533634",
+                    "ap-northeast-1": "520713654638",
+                    "ap-northeast-2": "520713654638",
+                    "ap-south-1": "520713654638",
+                    "ap-southeast-1": "520713654638",
+                    "ap-southeast-2": "520713654638",
+                    "ca-central-1": "520713654638",
+                    "cn-north-1": "422961961927",
+                    "cn-northwest-1": "423003514399",
+                    "eu-central-1": "520713654638",
+                    "eu-north-1": "520713654638",
+                    "eu-south-1": "048378556238",
+                    "eu-west-1": "520713654638",
+                    "eu-west-2": "520713654638",
+                    "eu-west-3": "520713654638",
+                    "me-south-1": "724002660598",
+                    "sa-east-1": "520713654638",
+                    "us-east-1": "520713654638",
+                    "us-east-2": "520713654638",
+                    "us-gov-west-1": "246785580436",
+                    "us-iso-east-1": "744548109606",
+                    "us-isob-east-1": "453391408702",
+                    "us-west-1": "520713654638",
+                    "us-west-2": "520713654638"
+                },
+                "repository": "sagemaker-tensorflow"
+            },
+            "1.8.0": {
+                "py_versions": [
+                    "py2"
+                ],
+                "registries": {
+                    "af-south-1": "313743910680",
+                    "ap-east-1": "057415533634",
+                    "ap-northeast-1": "520713654638",
+                    "ap-northeast-2": "520713654638",
+                    "ap-south-1": "520713654638",
+                    "ap-southeast-1": "520713654638",
+                    "ap-southeast-2": "520713654638",
+                    "ca-central-1": "520713654638",
+                    "cn-north-1": "422961961927",
+                    "cn-northwest-1": "423003514399",
+                    "eu-central-1": "520713654638",
+                    "eu-north-1": "520713654638",
+                    "eu-south-1": "048378556238",
+                    "eu-west-1": "520713654638",
+                    "eu-west-2": "520713654638",
+                    "eu-west-3": "520713654638",
+                    "me-south-1": "724002660598",
+                    "sa-east-1": "520713654638",
+                    "us-east-1": "520713654638",
+                    "us-east-2": "520713654638",
+                    "us-gov-west-1": "246785580436",
+                    "us-iso-east-1": "744548109606",
+                    "us-isob-east-1": "453391408702",
+                    "us-west-1": "520713654638",
+                    "us-west-2": "520713654638"
+                },
+                "repository": "sagemaker-tensorflow"
+            },
+            "1.9.0": {
+                "py_versions": [
+                    "py2"
+                ],
+                "registries": {
+                    "af-south-1": "313743910680",
+                    "ap-east-1": "057415533634",
+                    "ap-northeast-1": "520713654638",
+                    "ap-northeast-2": "520713654638",
+                    "ap-south-1": "520713654638",
+                    "ap-southeast-1": "520713654638",
+                    "ap-southeast-2": "520713654638",
+                    "ca-central-1": "520713654638",
+                    "cn-north-1": "422961961927",
+                    "cn-northwest-1": "423003514399",
+                    "eu-central-1": "520713654638",
+                    "eu-north-1": "520713654638",
+                    "eu-south-1": "048378556238",
+                    "eu-west-1": "520713654638",
+                    "eu-west-2": "520713654638",
+                    "eu-west-3": "520713654638",
+                    "me-south-1": "724002660598",
+                    "sa-east-1": "520713654638",
+                    "us-east-1": "520713654638",
+                    "us-east-2": "520713654638",
+                    "us-gov-west-1": "246785580436",
+                    "us-iso-east-1": "744548109606",
+                    "us-isob-east-1": "453391408702",
+                    "us-west-1": "520713654638",
+                    "us-west-2": "520713654638"
+                },
+                "repository": "sagemaker-tensorflow"
+            },
             "1.10.0": {
                 "py_versions": [
                     "py2"
@@ -2542,7 +2752,6 @@
                 "py3": {
                     "registries": {
                         "af-south-1": "626614931356",
-                        "il-central-1": "780543022126",
                         "ap-east-1": "871362719292",
                         "ap-northeast-1": "763104351884",
                         "ap-northeast-2": "763104351884",
@@ -2554,6 +2763,7 @@
                         "ap-southeast-3": "907027046896",
                         "ap-southeast-4": "457447274322",
                         "ca-central-1": "763104351884",
+                        "ca-west-1": "204538143572",
                         "cn-north-1": "727897471807",
                         "cn-northwest-1": "727897471807",
                         "eu-central-1": "763104351884",
@@ -2564,8 +2774,9 @@
                         "eu-west-1": "763104351884",
                         "eu-west-2": "763104351884",
                         "eu-west-3": "763104351884",
-                        "me-south-1": "217643126080",
+                        "il-central-1": "780543022126",
                         "me-central-1": "914824155844",
+                        "me-south-1": "217643126080",
                         "sa-east-1": "763104351884",
                         "us-east-1": "763104351884",
                         "us-east-2": "763104351884",
@@ -2574,8 +2785,7 @@
                         "us-iso-east-1": "886529160074",
                         "us-isob-east-1": "094389454867",
                         "us-west-1": "763104351884",
-                        "us-west-2": "763104351884",
-                        "ca-west-1": "204538143572"
+                        "us-west-2": "763104351884"
                     },
                     "repository": "tensorflow-training"
                 }
@@ -2587,7 +2797,6 @@
                 ],
                 "registries": {
                     "af-south-1": "626614931356",
-                    "il-central-1": "780543022126",
                     "ap-east-1": "871362719292",
                     "ap-northeast-1": "763104351884",
                     "ap-northeast-2": "763104351884",
@@ -2599,6 +2808,7 @@
                     "ap-southeast-3": "907027046896",
                     "ap-southeast-4": "457447274322",
                     "ca-central-1": "763104351884",
+                    "ca-west-1": "204538143572",
                     "cn-north-1": "727897471807",
                     "cn-northwest-1": "727897471807",
                     "eu-central-1": "763104351884",
@@ -2609,8 +2819,9 @@
                     "eu-west-1": "763104351884",
                     "eu-west-2": "763104351884",
                     "eu-west-3": "763104351884",
-                    "me-south-1": "217643126080",
+                    "il-central-1": "780543022126",
                     "me-central-1": "914824155844",
+                    "me-south-1": "217643126080",
                     "sa-east-1": "763104351884",
                     "us-east-1": "763104351884",
                     "us-east-2": "763104351884",
@@ -2619,8 +2830,7 @@
                     "us-iso-east-1": "886529160074",
                     "us-isob-east-1": "094389454867",
                     "us-west-1": "763104351884",
-                    "us-west-2": "763104351884",
-                    "ca-west-1": "204538143572"
+                    "us-west-2": "763104351884"
                 },
                 "repository": "tensorflow-training"
             },
@@ -2631,7 +2841,6 @@
                 ],
                 "registries": {
                     "af-south-1": "626614931356",
-                    "il-central-1": "780543022126",
                     "ap-east-1": "871362719292",
                     "ap-northeast-1": "763104351884",
                     "ap-northeast-2": "763104351884",
@@ -2643,6 +2852,7 @@
                     "ap-southeast-3": "907027046896",
                     "ap-southeast-4": "457447274322",
                     "ca-central-1": "763104351884",
+                    "ca-west-1": "204538143572",
                     "cn-north-1": "727897471807",
                     "cn-northwest-1": "727897471807",
                     "eu-central-1": "763104351884",
@@ -2653,8 +2863,9 @@
                     "eu-west-1": "763104351884",
                     "eu-west-2": "763104351884",
                     "eu-west-3": "763104351884",
-                    "me-south-1": "217643126080",
+                    "il-central-1": "780543022126",
                     "me-central-1": "914824155844",
+                    "me-south-1": "217643126080",
                     "sa-east-1": "763104351884",
                     "us-east-1": "763104351884",
                     "us-east-2": "763104351884",
@@ -2663,8 +2874,7 @@
                     "us-iso-east-1": "886529160074",
                     "us-isob-east-1": "094389454867",
                     "us-west-1": "763104351884",
-                    "us-west-2": "763104351884",
-                    "ca-west-1": "204538143572"
+                    "us-west-2": "763104351884"
                 },
                 "repository": "tensorflow-training"
             },
@@ -2676,7 +2886,6 @@
                 ],
                 "registries": {
                     "af-south-1": "626614931356",
-                    "il-central-1": "780543022126",
                     "ap-east-1": "871362719292",
                     "ap-northeast-1": "763104351884",
                     "ap-northeast-2": "763104351884",
@@ -2688,6 +2897,7 @@
                     "ap-southeast-3": "907027046896",
                     "ap-southeast-4": "457447274322",
                     "ca-central-1": "763104351884",
+                    "ca-west-1": "204538143572",
                     "cn-north-1": "727897471807",
                     "cn-northwest-1": "727897471807",
                     "eu-central-1": "763104351884",
@@ -2698,8 +2908,9 @@
                     "eu-west-1": "763104351884",
                     "eu-west-2": "763104351884",
                     "eu-west-3": "763104351884",
-                    "me-south-1": "217643126080",
+                    "il-central-1": "780543022126",
                     "me-central-1": "914824155844",
+                    "me-south-1": "217643126080",
                     "sa-east-1": "763104351884",
                     "us-east-1": "763104351884",
                     "us-east-2": "763104351884",
@@ -2708,8 +2919,7 @@
                     "us-iso-east-1": "886529160074",
                     "us-isob-east-1": "094389454867",
                     "us-west-1": "763104351884",
-                    "us-west-2": "763104351884",
-                    "ca-west-1": "204538143572"
+                    "us-west-2": "763104351884"
                 },
                 "repository": "tensorflow-training"
             },
@@ -2721,7 +2931,6 @@
                 ],
                 "registries": {
                     "af-south-1": "626614931356",
-                    "il-central-1": "780543022126",
                     "ap-east-1": "871362719292",
                     "ap-northeast-1": "763104351884",
                     "ap-northeast-2": "763104351884",
@@ -2733,6 +2942,7 @@
                     "ap-southeast-3": "907027046896",
                     "ap-southeast-4": "457447274322",
                     "ca-central-1": "763104351884",
+                    "ca-west-1": "204538143572",
                     "cn-north-1": "727897471807",
                     "cn-northwest-1": "727897471807",
                     "eu-central-1": "763104351884",
@@ -2743,8 +2953,9 @@
                     "eu-west-1": "763104351884",
                     "eu-west-2": "763104351884",
                     "eu-west-3": "763104351884",
-                    "me-south-1": "217643126080",
+                    "il-central-1": "780543022126",
                     "me-central-1": "914824155844",
+                    "me-south-1": "217643126080",
                     "sa-east-1": "763104351884",
                     "us-east-1": "763104351884",
                     "us-east-2": "763104351884",
@@ -2753,8 +2964,7 @@
                     "us-iso-east-1": "886529160074",
                     "us-isob-east-1": "094389454867",
                     "us-west-1": "763104351884",
-                    "us-west-2": "763104351884",
-                    "ca-west-1": "204538143572"
+                    "us-west-2": "763104351884"
                 },
                 "repository": "tensorflow-training"
             },
@@ -2766,7 +2976,6 @@
                 ],
                 "registries": {
                     "af-south-1": "626614931356",
-                    "il-central-1": "780543022126",
                     "ap-east-1": "871362719292",
                     "ap-northeast-1": "763104351884",
                     "ap-northeast-2": "763104351884",
@@ -2778,6 +2987,7 @@
                     "ap-southeast-3": "907027046896",
                     "ap-southeast-4": "457447274322",
                     "ca-central-1": "763104351884",
+                    "ca-west-1": "204538143572",
                     "cn-north-1": "727897471807",
                     "cn-northwest-1": "727897471807",
                     "eu-central-1": "763104351884",
@@ -2788,53 +2998,9 @@
                     "eu-west-1": "763104351884",
                     "eu-west-2": "763104351884",
                     "eu-west-3": "763104351884",
-                    "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"
-                },
-                "repository": "tensorflow-training"
-            },
-            "1.15.5": {
-                "py_versions": [
-                    "py3",
-                    "py36",
-                    "py37"
-                ],
-                "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-south-1": "692866216735",
-                    "eu-south-2": "503227376785",
-                    "eu-west-1": "763104351884",
-                    "eu-west-2": "763104351884",
-                    "eu-west-3": "763104351884",
-                    "me-south-1": "217643126080",
                     "me-central-1": "914824155844",
+                    "me-south-1": "217643126080",
                     "sa-east-1": "763104351884",
                     "us-east-1": "763104351884",
                     "us-east-2": "763104351884",
@@ -2843,208 +3009,54 @@
                     "us-iso-east-1": "886529160074",
                     "us-isob-east-1": "094389454867",
                     "us-west-1": "763104351884",
-                    "us-west-2": "763104351884",
-                    "ca-west-1": "204538143572"
+                    "us-west-2": "763104351884"
                 },
                 "repository": "tensorflow-training"
             },
-            "1.4.1": {
-                "py_versions": [
-                    "py2"
-                ],
-                "registries": {
-                    "af-south-1": "313743910680",
-                    "ap-east-1": "057415533634",
-                    "ap-northeast-1": "520713654638",
-                    "ap-northeast-2": "520713654638",
-                    "ap-south-1": "520713654638",
-                    "ap-southeast-1": "520713654638",
-                    "ap-southeast-2": "520713654638",
-                    "ca-central-1": "520713654638",
-                    "cn-north-1": "422961961927",
-                    "cn-northwest-1": "423003514399",
-                    "eu-central-1": "520713654638",
-                    "eu-north-1": "520713654638",
-                    "eu-south-1": "048378556238",
-                    "eu-west-1": "520713654638",
-                    "eu-west-2": "520713654638",
-                    "eu-west-3": "520713654638",
-                    "me-south-1": "724002660598",
-                    "sa-east-1": "520713654638",
-                    "us-east-1": "520713654638",
-                    "us-east-2": "520713654638",
-                    "us-gov-west-1": "246785580436",
-                    "us-iso-east-1": "744548109606",
-                    "us-isob-east-1": "453391408702",
-                    "us-west-1": "520713654638",
-                    "us-west-2": "520713654638"
-                },
-                "repository": "sagemaker-tensorflow"
-            },
-            "1.5.0": {
-                "py_versions": [
-                    "py2"
-                ],
-                "registries": {
-                    "af-south-1": "313743910680",
-                    "ap-east-1": "057415533634",
-                    "ap-northeast-1": "520713654638",
-                    "ap-northeast-2": "520713654638",
-                    "ap-south-1": "520713654638",
-                    "ap-southeast-1": "520713654638",
-                    "ap-southeast-2": "520713654638",
-                    "ca-central-1": "520713654638",
-                    "cn-north-1": "422961961927",
-                    "cn-northwest-1": "423003514399",
-                    "eu-central-1": "520713654638",
-                    "eu-north-1": "520713654638",
-                    "eu-south-1": "048378556238",
-                    "eu-west-1": "520713654638",
-                    "eu-west-2": "520713654638",
-                    "eu-west-3": "520713654638",
-                    "me-south-1": "724002660598",
-                    "sa-east-1": "520713654638",
-                    "us-east-1": "520713654638",
-                    "us-east-2": "520713654638",
-                    "us-gov-west-1": "246785580436",
-                    "us-iso-east-1": "744548109606",
-                    "us-isob-east-1": "453391408702",
-                    "us-west-1": "520713654638",
-                    "us-west-2": "520713654638"
-                },
-                "repository": "sagemaker-tensorflow"
-            },
-            "1.6.0": {
-                "py_versions": [
-                    "py2"
-                ],
-                "registries": {
-                    "af-south-1": "313743910680",
-                    "ap-east-1": "057415533634",
-                    "ap-northeast-1": "520713654638",
-                    "ap-northeast-2": "520713654638",
-                    "ap-south-1": "520713654638",
-                    "ap-southeast-1": "520713654638",
-                    "ap-southeast-2": "520713654638",
-                    "ca-central-1": "520713654638",
-                    "cn-north-1": "422961961927",
-                    "cn-northwest-1": "423003514399",
-                    "eu-central-1": "520713654638",
-                    "eu-north-1": "520713654638",
-                    "eu-south-1": "048378556238",
-                    "eu-west-1": "520713654638",
-                    "eu-west-2": "520713654638",
-                    "eu-west-3": "520713654638",
-                    "me-south-1": "724002660598",
-                    "sa-east-1": "520713654638",
-                    "us-east-1": "520713654638",
-                    "us-east-2": "520713654638",
-                    "us-gov-west-1": "246785580436",
-                    "us-iso-east-1": "744548109606",
-                    "us-isob-east-1": "453391408702",
-                    "us-west-1": "520713654638",
-                    "us-west-2": "520713654638"
-                },
-                "repository": "sagemaker-tensorflow"
-            },
-            "1.7.0": {
-                "py_versions": [
-                    "py2"
-                ],
-                "registries": {
-                    "af-south-1": "313743910680",
-                    "ap-east-1": "057415533634",
-                    "ap-northeast-1": "520713654638",
-                    "ap-northeast-2": "520713654638",
-                    "ap-south-1": "520713654638",
-                    "ap-southeast-1": "520713654638",
-                    "ap-southeast-2": "520713654638",
-                    "ca-central-1": "520713654638",
-                    "cn-north-1": "422961961927",
-                    "cn-northwest-1": "423003514399",
-                    "eu-central-1": "520713654638",
-                    "eu-north-1": "520713654638",
-                    "eu-south-1": "048378556238",
-                    "eu-west-1": "520713654638",
-                    "eu-west-2": "520713654638",
-                    "eu-west-3": "520713654638",
-                    "me-south-1": "724002660598",
-                    "sa-east-1": "520713654638",
-                    "us-east-1": "520713654638",
-                    "us-east-2": "520713654638",
-                    "us-gov-west-1": "246785580436",
-                    "us-iso-east-1": "744548109606",
-                    "us-isob-east-1": "453391408702",
-                    "us-west-1": "520713654638",
-                    "us-west-2": "520713654638"
-                },
-                "repository": "sagemaker-tensorflow"
-            },
-            "1.8.0": {
-                "py_versions": [
-                    "py2"
-                ],
-                "registries": {
-                    "af-south-1": "313743910680",
-                    "ap-east-1": "057415533634",
-                    "ap-northeast-1": "520713654638",
-                    "ap-northeast-2": "520713654638",
-                    "ap-south-1": "520713654638",
-                    "ap-southeast-1": "520713654638",
-                    "ap-southeast-2": "520713654638",
-                    "ca-central-1": "520713654638",
-                    "cn-north-1": "422961961927",
-                    "cn-northwest-1": "423003514399",
-                    "eu-central-1": "520713654638",
-                    "eu-north-1": "520713654638",
-                    "eu-south-1": "048378556238",
-                    "eu-west-1": "520713654638",
-                    "eu-west-2": "520713654638",
-                    "eu-west-3": "520713654638",
-                    "me-south-1": "724002660598",
-                    "sa-east-1": "520713654638",
-                    "us-east-1": "520713654638",
-                    "us-east-2": "520713654638",
-                    "us-gov-west-1": "246785580436",
-                    "us-iso-east-1": "744548109606",
-                    "us-isob-east-1": "453391408702",
-                    "us-west-1": "520713654638",
-                    "us-west-2": "520713654638"
-                },
-                "repository": "sagemaker-tensorflow"
-            },
-            "1.9.0": {
+            "1.15.5": {
                 "py_versions": [
-                    "py2"
+                    "py3",
+                    "py36",
+                    "py37"
                 ],
                 "registries": {
-                    "af-south-1": "313743910680",
-                    "ap-east-1": "057415533634",
-                    "ap-northeast-1": "520713654638",
-                    "ap-northeast-2": "520713654638",
-                    "ap-south-1": "520713654638",
-                    "ap-southeast-1": "520713654638",
-                    "ap-southeast-2": "520713654638",
-                    "ca-central-1": "520713654638",
-                    "cn-north-1": "422961961927",
-                    "cn-northwest-1": "423003514399",
-                    "eu-central-1": "520713654638",
-                    "eu-north-1": "520713654638",
-                    "eu-south-1": "048378556238",
-                    "eu-west-1": "520713654638",
-                    "eu-west-2": "520713654638",
-                    "eu-west-3": "520713654638",
-                    "me-south-1": "724002660598",
-                    "sa-east-1": "520713654638",
-                    "us-east-1": "520713654638",
-                    "us-east-2": "520713654638",
-                    "us-gov-west-1": "246785580436",
-                    "us-iso-east-1": "744548109606",
-                    "us-isob-east-1": "453391408702",
-                    "us-west-1": "520713654638",
-                    "us-west-2": "520713654638"
+                    "af-south-1": "626614931356",
+                    "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",
+                    "ca-west-1": "204538143572",
+                    "cn-north-1": "727897471807",
+                    "cn-northwest-1": "727897471807",
+                    "eu-central-1": "763104351884",
+                    "eu-central-2": "380420809688",
+                    "eu-north-1": "763104351884",
+                    "eu-south-1": "692866216735",
+                    "eu-south-2": "503227376785",
+                    "eu-west-1": "763104351884",
+                    "eu-west-2": "763104351884",
+                    "eu-west-3": "763104351884",
+                    "il-central-1": "780543022126",
+                    "me-central-1": "914824155844",
+                    "me-south-1": "217643126080",
+                    "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"
                 },
-                "repository": "sagemaker-tensorflow"
+                "repository": "tensorflow-training"
             },
             "2.0.0": {
                 "py_versions": [
@@ -3053,7 +3065,6 @@
                 ],
                 "registries": {
                     "af-south-1": "626614931356",
-                    "il-central-1": "780543022126",
                     "ap-east-1": "871362719292",
                     "ap-northeast-1": "763104351884",
                     "ap-northeast-2": "763104351884",
@@ -3065,6 +3076,7 @@
                     "ap-southeast-3": "907027046896",
                     "ap-southeast-4": "457447274322",
                     "ca-central-1": "763104351884",
+                    "ca-west-1": "204538143572",
                     "cn-north-1": "727897471807",
                     "cn-northwest-1": "727897471807",
                     "eu-central-1": "763104351884",
@@ -3075,8 +3087,9 @@
                     "eu-west-1": "763104351884",
                     "eu-west-2": "763104351884",
                     "eu-west-3": "763104351884",
-                    "me-south-1": "217643126080",
+                    "il-central-1": "780543022126",
                     "me-central-1": "914824155844",
+                    "me-south-1": "217643126080",
                     "sa-east-1": "763104351884",
                     "us-east-1": "763104351884",
                     "us-east-2": "763104351884",
@@ -3085,8 +3098,7 @@
                     "us-iso-east-1": "886529160074",
                     "us-isob-east-1": "094389454867",
                     "us-west-1": "763104351884",
-                    "us-west-2": "763104351884",
-                    "ca-west-1": "204538143572"
+                    "us-west-2": "763104351884"
                 },
                 "repository": "tensorflow-training"
             },
@@ -3097,7 +3109,6 @@
                 ],
                 "registries": {
                     "af-south-1": "626614931356",
-                    "il-central-1": "780543022126",
                     "ap-east-1": "871362719292",
                     "ap-northeast-1": "763104351884",
                     "ap-northeast-2": "763104351884",
@@ -3109,6 +3120,7 @@
                     "ap-southeast-3": "907027046896",
                     "ap-southeast-4": "457447274322",
                     "ca-central-1": "763104351884",
+                    "ca-west-1": "204538143572",
                     "cn-north-1": "727897471807",
                     "cn-northwest-1": "727897471807",
                     "eu-central-1": "763104351884",
@@ -3119,8 +3131,9 @@
                     "eu-west-1": "763104351884",
                     "eu-west-2": "763104351884",
                     "eu-west-3": "763104351884",
-                    "me-south-1": "217643126080",
+                    "il-central-1": "780543022126",
                     "me-central-1": "914824155844",
+                    "me-south-1": "217643126080",
                     "sa-east-1": "763104351884",
                     "us-east-1": "763104351884",
                     "us-east-2": "763104351884",
@@ -3129,8 +3142,7 @@
                     "us-iso-east-1": "886529160074",
                     "us-isob-east-1": "094389454867",
                     "us-west-1": "763104351884",
-                    "us-west-2": "763104351884",
-                    "ca-west-1": "204538143572"
+                    "us-west-2": "763104351884"
                 },
                 "repository": "tensorflow-training"
             },
@@ -3141,7 +3153,6 @@
                 ],
                 "registries": {
                     "af-south-1": "626614931356",
-                    "il-central-1": "780543022126",
                     "ap-east-1": "871362719292",
                     "ap-northeast-1": "763104351884",
                     "ap-northeast-2": "763104351884",
@@ -3153,6 +3164,7 @@
                     "ap-southeast-3": "907027046896",
                     "ap-southeast-4": "457447274322",
                     "ca-central-1": "763104351884",
+                    "ca-west-1": "204538143572",
                     "cn-north-1": "727897471807",
                     "cn-northwest-1": "727897471807",
                     "eu-central-1": "763104351884",
@@ -3163,8 +3175,9 @@
                     "eu-west-1": "763104351884",
                     "eu-west-2": "763104351884",
                     "eu-west-3": "763104351884",
-                    "me-south-1": "217643126080",
+                    "il-central-1": "780543022126",
                     "me-central-1": "914824155844",
+                    "me-south-1": "217643126080",
                     "sa-east-1": "763104351884",
                     "us-east-1": "763104351884",
                     "us-east-2": "763104351884",
@@ -3173,18 +3186,17 @@
                     "us-iso-east-1": "886529160074",
                     "us-isob-east-1": "094389454867",
                     "us-west-1": "763104351884",
-                    "us-west-2": "763104351884",
-                    "ca-west-1": "204538143572"
+                    "us-west-2": "763104351884"
                 },
                 "repository": "tensorflow-training"
             },
             "2.0.3": {
                 "py_versions": [
-                    "py3"
+                    "py3",
+                    "py36"
                 ],
                 "registries": {
                     "af-south-1": "626614931356",
-                    "il-central-1": "780543022126",
                     "ap-east-1": "871362719292",
                     "ap-northeast-1": "763104351884",
                     "ap-northeast-2": "763104351884",
@@ -3196,6 +3208,7 @@
                     "ap-southeast-3": "907027046896",
                     "ap-southeast-4": "457447274322",
                     "ca-central-1": "763104351884",
+                    "ca-west-1": "204538143572",
                     "cn-north-1": "727897471807",
                     "cn-northwest-1": "727897471807",
                     "eu-central-1": "763104351884",
@@ -3206,8 +3219,9 @@
                     "eu-west-1": "763104351884",
                     "eu-west-2": "763104351884",
                     "eu-west-3": "763104351884",
-                    "me-south-1": "217643126080",
+                    "il-central-1": "780543022126",
                     "me-central-1": "914824155844",
+                    "me-south-1": "217643126080",
                     "sa-east-1": "763104351884",
                     "us-east-1": "763104351884",
                     "us-east-2": "763104351884",
@@ -3216,18 +3230,17 @@
                     "us-iso-east-1": "886529160074",
                     "us-isob-east-1": "094389454867",
                     "us-west-1": "763104351884",
-                    "us-west-2": "763104351884",
-                    "ca-west-1": "204538143572"
+                    "us-west-2": "763104351884"
                 },
                 "repository": "tensorflow-training"
             },
             "2.0.4": {
                 "py_versions": [
-                    "py3"
+                    "py3",
+                    "py36"
                 ],
                 "registries": {
                     "af-south-1": "626614931356",
-                    "il-central-1": "780543022126",
                     "ap-east-1": "871362719292",
                     "ap-northeast-1": "763104351884",
                     "ap-northeast-2": "763104351884",
@@ -3239,6 +3252,7 @@
                     "ap-southeast-3": "907027046896",
                     "ap-southeast-4": "457447274322",
                     "ca-central-1": "763104351884",
+                    "ca-west-1": "204538143572",
                     "cn-north-1": "727897471807",
                     "cn-northwest-1": "727897471807",
                     "eu-central-1": "763104351884",
@@ -3249,8 +3263,9 @@
                     "eu-west-1": "763104351884",
                     "eu-west-2": "763104351884",
                     "eu-west-3": "763104351884",
-                    "me-south-1": "217643126080",
+                    "il-central-1": "780543022126",
                     "me-central-1": "914824155844",
+                    "me-south-1": "217643126080",
                     "sa-east-1": "763104351884",
                     "us-east-1": "763104351884",
                     "us-east-2": "763104351884",
@@ -3259,8 +3274,7 @@
                     "us-iso-east-1": "886529160074",
                     "us-isob-east-1": "094389454867",
                     "us-west-1": "763104351884",
-                    "us-west-2": "763104351884",
-                    "ca-west-1": "204538143572"
+                    "us-west-2": "763104351884"
                 },
                 "repository": "tensorflow-training"
             },
@@ -3271,7 +3285,6 @@
                 ],
                 "registries": {
                     "af-south-1": "626614931356",
-                    "il-central-1": "780543022126",
                     "ap-east-1": "871362719292",
                     "ap-northeast-1": "763104351884",
                     "ap-northeast-2": "763104351884",
@@ -3283,6 +3296,7 @@
                     "ap-southeast-3": "907027046896",
                     "ap-southeast-4": "457447274322",
                     "ca-central-1": "763104351884",
+                    "ca-west-1": "204538143572",
                     "cn-north-1": "727897471807",
                     "cn-northwest-1": "727897471807",
                     "eu-central-1": "763104351884",
@@ -3293,8 +3307,9 @@
                     "eu-west-1": "763104351884",
                     "eu-west-2": "763104351884",
                     "eu-west-3": "763104351884",
-                    "me-south-1": "217643126080",
+                    "il-central-1": "780543022126",
                     "me-central-1": "914824155844",
+                    "me-south-1": "217643126080",
                     "sa-east-1": "763104351884",
                     "us-east-1": "763104351884",
                     "us-east-2": "763104351884",
@@ -3303,8 +3318,7 @@
                     "us-iso-east-1": "886529160074",
                     "us-isob-east-1": "094389454867",
                     "us-west-1": "763104351884",
-                    "us-west-2": "763104351884",
-                    "ca-west-1": "204538143572"
+                    "us-west-2": "763104351884"
                 },
                 "repository": "tensorflow-training"
             },
@@ -3315,7 +3329,6 @@
                 ],
                 "registries": {
                     "af-south-1": "626614931356",
-                    "il-central-1": "780543022126",
                     "ap-east-1": "871362719292",
                     "ap-northeast-1": "763104351884",
                     "ap-northeast-2": "763104351884",
@@ -3327,6 +3340,7 @@
                     "ap-southeast-3": "907027046896",
                     "ap-southeast-4": "457447274322",
                     "ca-central-1": "763104351884",
+                    "ca-west-1": "204538143572",
                     "cn-north-1": "727897471807",
                     "cn-northwest-1": "727897471807",
                     "eu-central-1": "763104351884",
@@ -3337,8 +3351,9 @@
                     "eu-west-1": "763104351884",
                     "eu-west-2": "763104351884",
                     "eu-west-3": "763104351884",
-                    "me-south-1": "217643126080",
+                    "il-central-1": "780543022126",
                     "me-central-1": "914824155844",
+                    "me-south-1": "217643126080",
                     "sa-east-1": "763104351884",
                     "us-east-1": "763104351884",
                     "us-east-2": "763104351884",
@@ -3347,18 +3362,17 @@
                     "us-iso-east-1": "886529160074",
                     "us-isob-east-1": "094389454867",
                     "us-west-1": "763104351884",
-                    "us-west-2": "763104351884",
-                    "ca-west-1": "204538143572"
+                    "us-west-2": "763104351884"
                 },
                 "repository": "tensorflow-training"
             },
             "2.1.2": {
                 "py_versions": [
-                    "py3"
+                    "py3",
+                    "py36"
                 ],
                 "registries": {
                     "af-south-1": "626614931356",
-                    "il-central-1": "780543022126",
                     "ap-east-1": "871362719292",
                     "ap-northeast-1": "763104351884",
                     "ap-northeast-2": "763104351884",
@@ -3370,6 +3384,7 @@
                     "ap-southeast-3": "907027046896",
                     "ap-southeast-4": "457447274322",
                     "ca-central-1": "763104351884",
+                    "ca-west-1": "204538143572",
                     "cn-north-1": "727897471807",
                     "cn-northwest-1": "727897471807",
                     "eu-central-1": "763104351884",
@@ -3380,8 +3395,9 @@
                     "eu-west-1": "763104351884",
                     "eu-west-2": "763104351884",
                     "eu-west-3": "763104351884",
-                    "me-south-1": "217643126080",
+                    "il-central-1": "780543022126",
                     "me-central-1": "914824155844",
+                    "me-south-1": "217643126080",
                     "sa-east-1": "763104351884",
                     "us-east-1": "763104351884",
                     "us-east-2": "763104351884",
@@ -3390,18 +3406,17 @@
                     "us-iso-east-1": "886529160074",
                     "us-isob-east-1": "094389454867",
                     "us-west-1": "763104351884",
-                    "us-west-2": "763104351884",
-                    "ca-west-1": "204538143572"
+                    "us-west-2": "763104351884"
                 },
                 "repository": "tensorflow-training"
             },
             "2.1.3": {
                 "py_versions": [
-                    "py3"
+                    "py3",
+                    "py36"
                 ],
                 "registries": {
                     "af-south-1": "626614931356",
-                    "il-central-1": "780543022126",
                     "ap-east-1": "871362719292",
                     "ap-northeast-1": "763104351884",
                     "ap-northeast-2": "763104351884",
@@ -3413,6 +3428,7 @@
                     "ap-southeast-3": "907027046896",
                     "ap-southeast-4": "457447274322",
                     "ca-central-1": "763104351884",
+                    "ca-west-1": "204538143572",
                     "cn-north-1": "727897471807",
                     "cn-northwest-1": "727897471807",
                     "eu-central-1": "763104351884",
@@ -3423,8 +3439,9 @@
                     "eu-west-1": "763104351884",
                     "eu-west-2": "763104351884",
                     "eu-west-3": "763104351884",
-                    "me-south-1": "217643126080",
+                    "il-central-1": "780543022126",
                     "me-central-1": "914824155844",
+                    "me-south-1": "217643126080",
                     "sa-east-1": "763104351884",
                     "us-east-1": "763104351884",
                     "us-east-2": "763104351884",
@@ -3433,8 +3450,7 @@
                     "us-iso-east-1": "886529160074",
                     "us-isob-east-1": "094389454867",
                     "us-west-1": "763104351884",
-                    "us-west-2": "763104351884",
-                    "ca-west-1": "204538143572"
+                    "us-west-2": "763104351884"
                 },
                 "repository": "tensorflow-training"
             },
@@ -3444,7 +3460,6 @@
                 ],
                 "registries": {
                     "af-south-1": "626614931356",
-                    "il-central-1": "780543022126",
                     "ap-east-1": "871362719292",
                     "ap-northeast-1": "763104351884",
                     "ap-northeast-2": "763104351884",
@@ -3456,6 +3471,7 @@
                     "ap-southeast-3": "907027046896",
                     "ap-southeast-4": "457447274322",
                     "ca-central-1": "763104351884",
+                    "ca-west-1": "204538143572",
                     "cn-north-1": "727897471807",
                     "cn-northwest-1": "727897471807",
                     "eu-central-1": "763104351884",
@@ -3466,8 +3482,9 @@
                     "eu-west-1": "763104351884",
                     "eu-west-2": "763104351884",
                     "eu-west-3": "763104351884",
-                    "me-south-1": "217643126080",
+                    "il-central-1": "780543022126",
                     "me-central-1": "914824155844",
+                    "me-south-1": "217643126080",
                     "sa-east-1": "763104351884",
                     "us-east-1": "763104351884",
                     "us-east-2": "763104351884",
@@ -3476,8 +3493,7 @@
                     "us-iso-east-1": "886529160074",
                     "us-isob-east-1": "094389454867",
                     "us-west-1": "763104351884",
-                    "us-west-2": "763104351884",
-                    "ca-west-1": "204538143572"
+                    "us-west-2": "763104351884"
                 },
                 "repository": "tensorflow-training"
             },
@@ -3487,7 +3503,6 @@
                 ],
                 "registries": {
                     "af-south-1": "626614931356",
-                    "il-central-1": "780543022126",
                     "ap-east-1": "871362719292",
                     "ap-northeast-1": "763104351884",
                     "ap-northeast-2": "763104351884",
@@ -3499,6 +3514,7 @@
                     "ap-southeast-3": "907027046896",
                     "ap-southeast-4": "457447274322",
                     "ca-central-1": "763104351884",
+                    "ca-west-1": "204538143572",
                     "cn-north-1": "727897471807",
                     "cn-northwest-1": "727897471807",
                     "eu-central-1": "763104351884",
@@ -3509,8 +3525,9 @@
                     "eu-west-1": "763104351884",
                     "eu-west-2": "763104351884",
                     "eu-west-3": "763104351884",
-                    "me-south-1": "217643126080",
+                    "il-central-1": "780543022126",
                     "me-central-1": "914824155844",
+                    "me-south-1": "217643126080",
                     "sa-east-1": "763104351884",
                     "us-east-1": "763104351884",
                     "us-east-2": "763104351884",
@@ -3519,8 +3536,7 @@
                     "us-iso-east-1": "886529160074",
                     "us-isob-east-1": "094389454867",
                     "us-west-1": "763104351884",
-                    "us-west-2": "763104351884",
-                    "ca-west-1": "204538143572"
+                    "us-west-2": "763104351884"
                 },
                 "repository": "tensorflow-training"
             },
@@ -3530,7 +3546,6 @@
                 ],
                 "registries": {
                     "af-south-1": "626614931356",
-                    "il-central-1": "780543022126",
                     "ap-east-1": "871362719292",
                     "ap-northeast-1": "763104351884",
                     "ap-northeast-2": "763104351884",
@@ -3542,6 +3557,7 @@
                     "ap-southeast-3": "907027046896",
                     "ap-southeast-4": "457447274322",
                     "ca-central-1": "763104351884",
+                    "ca-west-1": "204538143572",
                     "cn-north-1": "727897471807",
                     "cn-northwest-1": "727897471807",
                     "eu-central-1": "763104351884",
@@ -3552,8 +3568,9 @@
                     "eu-west-1": "763104351884",
                     "eu-west-2": "763104351884",
                     "eu-west-3": "763104351884",
-                    "me-south-1": "217643126080",
+                    "il-central-1": "780543022126",
                     "me-central-1": "914824155844",
+                    "me-south-1": "217643126080",
                     "sa-east-1": "763104351884",
                     "us-east-1": "763104351884",
                     "us-east-2": "763104351884",
@@ -3562,8 +3579,7 @@
                     "us-iso-east-1": "886529160074",
                     "us-isob-east-1": "094389454867",
                     "us-west-1": "763104351884",
-                    "us-west-2": "763104351884",
-                    "ca-west-1": "204538143572"
+                    "us-west-2": "763104351884"
                 },
                 "repository": "tensorflow-training"
             },
@@ -3573,7 +3589,6 @@
                 ],
                 "registries": {
                     "af-south-1": "626614931356",
-                    "il-central-1": "780543022126",
                     "ap-east-1": "871362719292",
                     "ap-northeast-1": "763104351884",
                     "ap-northeast-2": "763104351884",
@@ -3585,6 +3600,7 @@
                     "ap-southeast-3": "907027046896",
                     "ap-southeast-4": "457447274322",
                     "ca-central-1": "763104351884",
+                    "ca-west-1": "204538143572",
                     "cn-north-1": "727897471807",
                     "cn-northwest-1": "727897471807",
                     "eu-central-1": "763104351884",
@@ -3595,8 +3611,9 @@
                     "eu-west-1": "763104351884",
                     "eu-west-2": "763104351884",
                     "eu-west-3": "763104351884",
-                    "me-south-1": "217643126080",
+                    "il-central-1": "780543022126",
                     "me-central-1": "914824155844",
+                    "me-south-1": "217643126080",
                     "sa-east-1": "763104351884",
                     "us-east-1": "763104351884",
                     "us-east-2": "763104351884",
@@ -3605,8 +3622,7 @@
                     "us-iso-east-1": "886529160074",
                     "us-isob-east-1": "094389454867",
                     "us-west-1": "763104351884",
-                    "us-west-2": "763104351884",
-                    "ca-west-1": "204538143572"
+                    "us-west-2": "763104351884"
                 },
                 "repository": "tensorflow-training"
             },
@@ -3616,7 +3632,6 @@
                 ],
                 "registries": {
                     "af-south-1": "626614931356",
-                    "il-central-1": "780543022126",
                     "ap-east-1": "871362719292",
                     "ap-northeast-1": "763104351884",
                     "ap-northeast-2": "763104351884",
@@ -3628,6 +3643,7 @@
                     "ap-southeast-3": "907027046896",
                     "ap-southeast-4": "457447274322",
                     "ca-central-1": "763104351884",
+                    "ca-west-1": "204538143572",
                     "cn-north-1": "727897471807",
                     "cn-northwest-1": "727897471807",
                     "eu-central-1": "763104351884",
@@ -3638,8 +3654,9 @@
                     "eu-west-1": "763104351884",
                     "eu-west-2": "763104351884",
                     "eu-west-3": "763104351884",
-                    "me-south-1": "217643126080",
+                    "il-central-1": "780543022126",
                     "me-central-1": "914824155844",
+                    "me-south-1": "217643126080",
                     "sa-east-1": "763104351884",
                     "us-east-1": "763104351884",
                     "us-east-2": "763104351884",
@@ -3648,8 +3665,7 @@
                     "us-iso-east-1": "886529160074",
                     "us-isob-east-1": "094389454867",
                     "us-west-1": "763104351884",
-                    "us-west-2": "763104351884",
-                    "ca-west-1": "204538143572"
+                    "us-west-2": "763104351884"
                 },
                 "repository": "tensorflow-training"
             },
@@ -3659,7 +3675,6 @@
                 ],
                 "registries": {
                     "af-south-1": "626614931356",
-                    "il-central-1": "780543022126",
                     "ap-east-1": "871362719292",
                     "ap-northeast-1": "763104351884",
                     "ap-northeast-2": "763104351884",
@@ -3671,6 +3686,7 @@
                     "ap-southeast-3": "907027046896",
                     "ap-southeast-4": "457447274322",
                     "ca-central-1": "763104351884",
+                    "ca-west-1": "204538143572",
                     "cn-north-1": "727897471807",
                     "cn-northwest-1": "727897471807",
                     "eu-central-1": "763104351884",
@@ -3681,8 +3697,9 @@
                     "eu-west-1": "763104351884",
                     "eu-west-2": "763104351884",
                     "eu-west-3": "763104351884",
-                    "me-south-1": "217643126080",
+                    "il-central-1": "780543022126",
                     "me-central-1": "914824155844",
+                    "me-south-1": "217643126080",
                     "sa-east-1": "763104351884",
                     "us-east-1": "763104351884",
                     "us-east-2": "763104351884",
@@ -3691,8 +3708,7 @@
                     "us-iso-east-1": "886529160074",
                     "us-isob-east-1": "094389454867",
                     "us-west-1": "763104351884",
-                    "us-west-2": "763104351884",
-                    "ca-west-1": "204538143572"
+                    "us-west-2": "763104351884"
                 },
                 "repository": "tensorflow-training"
             },
@@ -3702,7 +3718,6 @@
                 ],
                 "registries": {
                     "af-south-1": "626614931356",
-                    "il-central-1": "780543022126",
                     "ap-east-1": "871362719292",
                     "ap-northeast-1": "763104351884",
                     "ap-northeast-2": "763104351884",
@@ -3714,6 +3729,7 @@
                     "ap-southeast-3": "907027046896",
                     "ap-southeast-4": "457447274322",
                     "ca-central-1": "763104351884",
+                    "ca-west-1": "204538143572",
                     "cn-north-1": "727897471807",
                     "cn-northwest-1": "727897471807",
                     "eu-central-1": "763104351884",
@@ -3724,8 +3740,9 @@
                     "eu-west-1": "763104351884",
                     "eu-west-2": "763104351884",
                     "eu-west-3": "763104351884",
-                    "me-south-1": "217643126080",
+                    "il-central-1": "780543022126",
                     "me-central-1": "914824155844",
+                    "me-south-1": "217643126080",
                     "sa-east-1": "763104351884",
                     "us-east-1": "763104351884",
                     "us-east-2": "763104351884",
@@ -3734,8 +3751,7 @@
                     "us-iso-east-1": "886529160074",
                     "us-isob-east-1": "094389454867",
                     "us-west-1": "763104351884",
-                    "us-west-2": "763104351884",
-                    "ca-west-1": "204538143572"
+                    "us-west-2": "763104351884"
                 },
                 "repository": "tensorflow-training"
             },
@@ -3745,7 +3761,6 @@
                 ],
                 "registries": {
                     "af-south-1": "626614931356",
-                    "il-central-1": "780543022126",
                     "ap-east-1": "871362719292",
                     "ap-northeast-1": "763104351884",
                     "ap-northeast-2": "763104351884",
@@ -3757,6 +3772,7 @@
                     "ap-southeast-3": "907027046896",
                     "ap-southeast-4": "457447274322",
                     "ca-central-1": "763104351884",
+                    "ca-west-1": "204538143572",
                     "cn-north-1": "727897471807",
                     "cn-northwest-1": "727897471807",
                     "eu-central-1": "763104351884",
@@ -3767,8 +3783,9 @@
                     "eu-west-1": "763104351884",
                     "eu-west-2": "763104351884",
                     "eu-west-3": "763104351884",
-                    "me-south-1": "217643126080",
+                    "il-central-1": "780543022126",
                     "me-central-1": "914824155844",
+                    "me-south-1": "217643126080",
                     "sa-east-1": "763104351884",
                     "us-east-1": "763104351884",
                     "us-east-2": "763104351884",
@@ -3777,8 +3794,7 @@
                     "us-iso-east-1": "886529160074",
                     "us-isob-east-1": "094389454867",
                     "us-west-1": "763104351884",
-                    "us-west-2": "763104351884",
-                    "ca-west-1": "204538143572"
+                    "us-west-2": "763104351884"
                 },
                 "repository": "tensorflow-training"
             },
@@ -3788,7 +3804,6 @@
                 ],
                 "registries": {
                     "af-south-1": "626614931356",
-                    "il-central-1": "780543022126",
                     "ap-east-1": "871362719292",
                     "ap-northeast-1": "763104351884",
                     "ap-northeast-2": "763104351884",
@@ -3800,6 +3815,7 @@
                     "ap-southeast-3": "907027046896",
                     "ap-southeast-4": "457447274322",
                     "ca-central-1": "763104351884",
+                    "ca-west-1": "204538143572",
                     "cn-north-1": "727897471807",
                     "cn-northwest-1": "727897471807",
                     "eu-central-1": "763104351884",
@@ -3810,8 +3826,9 @@
                     "eu-west-1": "763104351884",
                     "eu-west-2": "763104351884",
                     "eu-west-3": "763104351884",
-                    "me-south-1": "217643126080",
+                    "il-central-1": "780543022126",
                     "me-central-1": "914824155844",
+                    "me-south-1": "217643126080",
                     "sa-east-1": "763104351884",
                     "us-east-1": "763104351884",
                     "us-east-2": "763104351884",
@@ -3820,8 +3837,7 @@
                     "us-iso-east-1": "886529160074",
                     "us-isob-east-1": "094389454867",
                     "us-west-1": "763104351884",
-                    "us-west-2": "763104351884",
-                    "ca-west-1": "204538143572"
+                    "us-west-2": "763104351884"
                 },
                 "repository": "tensorflow-training"
             },
@@ -3831,7 +3847,6 @@
                 ],
                 "registries": {
                     "af-south-1": "626614931356",
-                    "il-central-1": "780543022126",
                     "ap-east-1": "871362719292",
                     "ap-northeast-1": "763104351884",
                     "ap-northeast-2": "763104351884",
@@ -3843,6 +3858,7 @@
                     "ap-southeast-3": "907027046896",
                     "ap-southeast-4": "457447274322",
                     "ca-central-1": "763104351884",
+                    "ca-west-1": "204538143572",
                     "cn-north-1": "727897471807",
                     "cn-northwest-1": "727897471807",
                     "eu-central-1": "763104351884",
@@ -3853,8 +3869,9 @@
                     "eu-west-1": "763104351884",
                     "eu-west-2": "763104351884",
                     "eu-west-3": "763104351884",
-                    "me-south-1": "217643126080",
+                    "il-central-1": "780543022126",
                     "me-central-1": "914824155844",
+                    "me-south-1": "217643126080",
                     "sa-east-1": "763104351884",
                     "us-east-1": "763104351884",
                     "us-east-2": "763104351884",
@@ -3863,8 +3880,7 @@
                     "us-iso-east-1": "886529160074",
                     "us-isob-east-1": "094389454867",
                     "us-west-1": "763104351884",
-                    "us-west-2": "763104351884",
-                    "ca-west-1": "204538143572"
+                    "us-west-2": "763104351884"
                 },
                 "repository": "tensorflow-training"
             },
@@ -3874,7 +3890,6 @@
                 ],
                 "registries": {
                     "af-south-1": "626614931356",
-                    "il-central-1": "780543022126",
                     "ap-east-1": "871362719292",
                     "ap-northeast-1": "763104351884",
                     "ap-northeast-2": "763104351884",
@@ -3886,6 +3901,7 @@
                     "ap-southeast-3": "907027046896",
                     "ap-southeast-4": "457447274322",
                     "ca-central-1": "763104351884",
+                    "ca-west-1": "204538143572",
                     "cn-north-1": "727897471807",
                     "cn-northwest-1": "727897471807",
                     "eu-central-1": "763104351884",
@@ -3896,8 +3912,9 @@
                     "eu-west-1": "763104351884",
                     "eu-west-2": "763104351884",
                     "eu-west-3": "763104351884",
-                    "me-south-1": "217643126080",
+                    "il-central-1": "780543022126",
                     "me-central-1": "914824155844",
+                    "me-south-1": "217643126080",
                     "sa-east-1": "763104351884",
                     "us-east-1": "763104351884",
                     "us-east-2": "763104351884",
@@ -3906,8 +3923,7 @@
                     "us-iso-east-1": "886529160074",
                     "us-isob-east-1": "094389454867",
                     "us-west-1": "763104351884",
-                    "us-west-2": "763104351884",
-                    "ca-west-1": "204538143572"
+                    "us-west-2": "763104351884"
                 },
                 "repository": "tensorflow-training"
             },
@@ -3917,7 +3933,6 @@
                 ],
                 "registries": {
                     "af-south-1": "626614931356",
-                    "il-central-1": "780543022126",
                     "ap-east-1": "871362719292",
                     "ap-northeast-1": "763104351884",
                     "ap-northeast-2": "763104351884",
@@ -3929,6 +3944,7 @@
                     "ap-southeast-3": "907027046896",
                     "ap-southeast-4": "457447274322",
                     "ca-central-1": "763104351884",
+                    "ca-west-1": "204538143572",
                     "cn-north-1": "727897471807",
                     "cn-northwest-1": "727897471807",
                     "eu-central-1": "763104351884",
@@ -3939,8 +3955,9 @@
                     "eu-west-1": "763104351884",
                     "eu-west-2": "763104351884",
                     "eu-west-3": "763104351884",
-                    "me-south-1": "217643126080",
+                    "il-central-1": "780543022126",
                     "me-central-1": "914824155844",
+                    "me-south-1": "217643126080",
                     "sa-east-1": "763104351884",
                     "us-east-1": "763104351884",
                     "us-east-2": "763104351884",
@@ -3949,8 +3966,7 @@
                     "us-iso-east-1": "886529160074",
                     "us-isob-east-1": "094389454867",
                     "us-west-1": "763104351884",
-                    "us-west-2": "763104351884",
-                    "ca-west-1": "204538143572"
+                    "us-west-2": "763104351884"
                 },
                 "repository": "tensorflow-training"
             },
@@ -3960,7 +3976,6 @@
                 ],
                 "registries": {
                     "af-south-1": "626614931356",
-                    "il-central-1": "780543022126",
                     "ap-east-1": "871362719292",
                     "ap-northeast-1": "763104351884",
                     "ap-northeast-2": "763104351884",
@@ -3972,6 +3987,7 @@
                     "ap-southeast-3": "907027046896",
                     "ap-southeast-4": "457447274322",
                     "ca-central-1": "763104351884",
+                    "ca-west-1": "204538143572",
                     "cn-north-1": "727897471807",
                     "cn-northwest-1": "727897471807",
                     "eu-central-1": "763104351884",
@@ -3982,8 +3998,9 @@
                     "eu-west-1": "763104351884",
                     "eu-west-2": "763104351884",
                     "eu-west-3": "763104351884",
-                    "me-south-1": "217643126080",
+                    "il-central-1": "780543022126",
                     "me-central-1": "914824155844",
+                    "me-south-1": "217643126080",
                     "sa-east-1": "763104351884",
                     "us-east-1": "763104351884",
                     "us-east-2": "763104351884",
@@ -3992,8 +4009,7 @@
                     "us-iso-east-1": "886529160074",
                     "us-isob-east-1": "094389454867",
                     "us-west-1": "763104351884",
-                    "us-west-2": "763104351884",
-                    "ca-west-1": "204538143572"
+                    "us-west-2": "763104351884"
                 },
                 "repository": "tensorflow-training"
             },
@@ -4003,7 +4019,6 @@
                 ],
                 "registries": {
                     "af-south-1": "626614931356",
-                    "il-central-1": "780543022126",
                     "ap-east-1": "871362719292",
                     "ap-northeast-1": "763104351884",
                     "ap-northeast-2": "763104351884",
@@ -4015,6 +4030,7 @@
                     "ap-southeast-3": "907027046896",
                     "ap-southeast-4": "457447274322",
                     "ca-central-1": "763104351884",
+                    "ca-west-1": "204538143572",
                     "cn-north-1": "727897471807",
                     "cn-northwest-1": "727897471807",
                     "eu-central-1": "763104351884",
@@ -4025,8 +4041,9 @@
                     "eu-west-1": "763104351884",
                     "eu-west-2": "763104351884",
                     "eu-west-3": "763104351884",
-                    "me-south-1": "217643126080",
+                    "il-central-1": "780543022126",
                     "me-central-1": "914824155844",
+                    "me-south-1": "217643126080",
                     "sa-east-1": "763104351884",
                     "us-east-1": "763104351884",
                     "us-east-2": "763104351884",
@@ -4035,8 +4052,7 @@
                     "us-iso-east-1": "886529160074",
                     "us-isob-east-1": "094389454867",
                     "us-west-1": "763104351884",
-                    "us-west-2": "763104351884",
-                    "ca-west-1": "204538143572"
+                    "us-west-2": "763104351884"
                 },
                 "repository": "tensorflow-training"
             },
@@ -4046,7 +4062,6 @@
                 ],
                 "registries": {
                     "af-south-1": "626614931356",
-                    "il-central-1": "780543022126",
                     "ap-east-1": "871362719292",
                     "ap-northeast-1": "763104351884",
                     "ap-northeast-2": "763104351884",
@@ -4058,6 +4073,7 @@
                     "ap-southeast-3": "907027046896",
                     "ap-southeast-4": "457447274322",
                     "ca-central-1": "763104351884",
+                    "ca-west-1": "204538143572",
                     "cn-north-1": "727897471807",
                     "cn-northwest-1": "727897471807",
                     "eu-central-1": "763104351884",
@@ -4068,8 +4084,9 @@
                     "eu-west-1": "763104351884",
                     "eu-west-2": "763104351884",
                     "eu-west-3": "763104351884",
-                    "me-south-1": "217643126080",
+                    "il-central-1": "780543022126",
                     "me-central-1": "914824155844",
+                    "me-south-1": "217643126080",
                     "sa-east-1": "763104351884",
                     "us-east-1": "763104351884",
                     "us-east-2": "763104351884",
@@ -4078,8 +4095,7 @@
                     "us-iso-east-1": "886529160074",
                     "us-isob-east-1": "094389454867",
                     "us-west-1": "763104351884",
-                    "us-west-2": "763104351884",
-                    "ca-west-1": "204538143572"
+                    "us-west-2": "763104351884"
                 },
                 "repository": "tensorflow-training"
             },
@@ -4089,7 +4105,6 @@
                 ],
                 "registries": {
                     "af-south-1": "626614931356",
-                    "il-central-1": "780543022126",
                     "ap-east-1": "871362719292",
                     "ap-northeast-1": "763104351884",
                     "ap-northeast-2": "763104351884",
@@ -4101,6 +4116,7 @@
                     "ap-southeast-3": "907027046896",
                     "ap-southeast-4": "457447274322",
                     "ca-central-1": "763104351884",
+                    "ca-west-1": "204538143572",
                     "cn-north-1": "727897471807",
                     "cn-northwest-1": "727897471807",
                     "eu-central-1": "763104351884",
@@ -4111,8 +4127,9 @@
                     "eu-west-1": "763104351884",
                     "eu-west-2": "763104351884",
                     "eu-west-3": "763104351884",
-                    "me-south-1": "217643126080",
+                    "il-central-1": "780543022126",
                     "me-central-1": "914824155844",
+                    "me-south-1": "217643126080",
                     "sa-east-1": "763104351884",
                     "us-east-1": "763104351884",
                     "us-east-2": "763104351884",
@@ -4121,8 +4138,7 @@
                     "us-iso-east-1": "886529160074",
                     "us-isob-east-1": "094389454867",
                     "us-west-1": "763104351884",
-                    "us-west-2": "763104351884",
-                    "ca-west-1": "204538143572"
+                    "us-west-2": "763104351884"
                 },
                 "repository": "tensorflow-training"
             },
@@ -4132,7 +4148,6 @@
                 ],
                 "registries": {
                     "af-south-1": "626614931356",
-                    "il-central-1": "780543022126",
                     "ap-east-1": "871362719292",
                     "ap-northeast-1": "763104351884",
                     "ap-northeast-2": "763104351884",
@@ -4144,6 +4159,7 @@
                     "ap-southeast-3": "907027046896",
                     "ap-southeast-4": "457447274322",
                     "ca-central-1": "763104351884",
+                    "ca-west-1": "204538143572",
                     "cn-north-1": "727897471807",
                     "cn-northwest-1": "727897471807",
                     "eu-central-1": "763104351884",
@@ -4154,6 +4170,8 @@
                     "eu-west-1": "763104351884",
                     "eu-west-2": "763104351884",
                     "eu-west-3": "763104351884",
+                    "il-central-1": "780543022126",
+                    "me-central-1": "914824155844",
                     "me-south-1": "217643126080",
                     "sa-east-1": "763104351884",
                     "us-east-1": "763104351884",
@@ -4163,8 +4181,7 @@
                     "us-iso-east-1": "886529160074",
                     "us-isob-east-1": "094389454867",
                     "us-west-1": "763104351884",
-                    "us-west-2": "763104351884",
-                    "ca-west-1": "204538143572"
+                    "us-west-2": "763104351884"
                 },
                 "repository": "tensorflow-training"
             },
@@ -4174,25 +4191,30 @@
                 ],
                 "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",
+                    "ca-west-1": "204538143572",
                     "cn-north-1": "727897471807",
                     "cn-northwest-1": "727897471807",
                     "eu-central-1": "763104351884",
+                    "eu-central-2": "380420809688",
                     "eu-north-1": "763104351884",
                     "eu-south-1": "692866216735",
+                    "eu-south-2": "503227376785",
                     "eu-west-1": "763104351884",
                     "eu-west-2": "763104351884",
                     "eu-west-3": "763104351884",
+                    "il-central-1": "780543022126",
+                    "me-central-1": "914824155844",
                     "me-south-1": "217643126080",
                     "sa-east-1": "763104351884",
                     "us-east-1": "763104351884",
@@ -4202,8 +4224,7 @@
                     "us-iso-east-1": "886529160074",
                     "us-isob-east-1": "094389454867",
                     "us-west-1": "763104351884",
-                    "us-west-2": "763104351884",
-                    "ca-west-1": "204538143572"
+                    "us-west-2": "763104351884"
                 },
                 "repository": "tensorflow-training"
             },
@@ -4213,30 +4234,38 @@
                 ],
                 "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",
+                    "ca-west-1": "204538143572",
+                    "cn-north-1": "727897471807",
+                    "cn-northwest-1": "727897471807",
                     "eu-central-1": "763104351884",
+                    "eu-central-2": "380420809688",
                     "eu-north-1": "763104351884",
                     "eu-south-1": "692866216735",
+                    "eu-south-2": "503227376785",
                     "eu-west-1": "763104351884",
                     "eu-west-2": "763104351884",
                     "eu-west-3": "763104351884",
+                    "il-central-1": "780543022126",
+                    "me-central-1": "914824155844",
                     "me-south-1": "217643126080",
                     "sa-east-1": "763104351884",
                     "us-east-1": "763104351884",
                     "us-east-2": "763104351884",
+                    "us-gov-east-1": "446045086412",
+                    "us-gov-west-1": "442386744353",
                     "us-west-1": "763104351884",
-                    "us-west-2": "763104351884",
-                    "ca-west-1": "204538143572"
+                    "us-west-2": "763104351884"
                 },
                 "repository": "tensorflow-training"
             },
@@ -4246,25 +4275,30 @@
                 ],
                 "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",
+                    "ca-west-1": "204538143572",
                     "cn-north-1": "727897471807",
                     "cn-northwest-1": "727897471807",
                     "eu-central-1": "763104351884",
+                    "eu-central-2": "380420809688",
                     "eu-north-1": "763104351884",
                     "eu-south-1": "692866216735",
+                    "eu-south-2": "503227376785",
                     "eu-west-1": "763104351884",
                     "eu-west-2": "763104351884",
                     "eu-west-3": "763104351884",
+                    "il-central-1": "780543022126",
+                    "me-central-1": "914824155844",
                     "me-south-1": "217643126080",
                     "sa-east-1": "763104351884",
                     "us-east-1": "763104351884",
@@ -4274,8 +4308,7 @@
                     "us-iso-east-1": "886529160074",
                     "us-isob-east-1": "094389454867",
                     "us-west-1": "763104351884",
-                    "us-west-2": "763104351884",
-                    "ca-west-1": "204538143572"
+                    "us-west-2": "763104351884"
                 },
                 "repository": "tensorflow-training"
             },
@@ -4324,4 +4357,4 @@
             }
         }
     }
-}
+}
\ No newline at end of file

From a5c6229b0f70410f44f9937e83f5c4e873caa179 Mon Sep 17 00:00:00 2001
From: jiapinw <95885824+jiapinw@users.noreply.github.com>
Date: Thu, 9 May 2024 14:54:41 -0700
Subject: [PATCH 25/58] Add tensorflow_serving support for mlflow models and
 enable lineage tracking for mlflow models (#4662)

* Initial commit for tensorflow_serving support of MLflow

* Add integ tests for mlflow tf_serving

* fix style issues

* remove unused attributes from tf builder

* Add deep ping for tf_serving local mode

* Initial commit for lineage impl

* Initial commit for tensorflow_serving support of MLflow

* Add integ tests for mlflow tf_serving

* fix style issues

* remove unused attributes from tf builder

* Add deep ping for tf_serving local mode

* Add integ tests and uts

* fix local mode for tf_serving

* Allow lineage tracking only in sagemaker endpoint mode

* fix regex pattern

* fix style issues

* fix regex pattern and hard coded py version in ut

* fix missing session

* Resolve pr comments and fix regex for mlflow registry and ids
---
 requirements/extras/test_requirements.txt     |   1 +
 src/sagemaker/serve/builder/model_builder.py  |  27 +-
 .../serve/builder/tf_serving_builder.py       | 129 ++++++
 .../serve/mode/local_container_mode.py        |  17 +-
 .../serve/mode/sagemaker_endpoint_mode.py     |  11 +
 .../serve/model_format/mlflow/constants.py    |  20 +-
 .../serve/model_format/mlflow/utils.py        |  66 +++-
 .../tensorflow_serving/__init__.py            |   0
 .../tensorflow_serving/inference.py           | 147 +++++++
 .../tensorflow_serving/prepare.py             |  67 ++++
 .../model_server/tensorflow_serving/server.py | 139 +++++++
 .../serve/utils/lineage_constants.py          |  28 ++
 src/sagemaker/serve/utils/lineage_utils.py    | 277 +++++++++++++
 src/sagemaker/serve/utils/predictors.py       |  46 +++
 .../serve_resources/mlflow/tensorflow/MLmodel |  17 +
 .../mlflow/tensorflow/conda.yaml              |  11 +
 .../mlflow/tensorflow/python_env.yaml         |   7 +
 .../mlflow/tensorflow/registered_model_meta   |   2 +
 .../mlflow/tensorflow/requirements.txt        |   4 +
 .../mlflow/tensorflow/tf2model/fingerprint.pb |   1 +
 .../mlflow/tensorflow/tf2model/saved_model.pb | Bin 0 -> 22966 bytes
 .../variables/variables.data-00000-of-00001   | Bin 0 -> 553 bytes
 .../tf2model/variables/variables.index        | Bin 0 -> 248 bytes
 tests/integ/sagemaker/serve/constants.py      |   1 +
 .../test_serve_mlflow_pytorch_flavor_happy.py |  26 +-
 ...st_serve_mlflow_tensorflow_flavor_happy.py | 175 ++++++++
 .../test_serve_mlflow_xgboost_flavor_happy.py |  26 +-
 .../serve/builder/test_model_builder.py       | 320 +++++++++++++++
 .../test_tensorflow_serving_builder.py        |  75 ++++
 .../model_format/mlflow/test_mlflow_utils.py  |  78 +++-
 .../tensorflow_serving/test_tf_prepare.py     | 116 ++++++
 .../tensorflow_serving/test_tf_server.py      | 100 +++++
 .../serve/utils/test_lineage_utils.py         | 374 ++++++++++++++++++
 33 files changed, 2275 insertions(+), 33 deletions(-)
 create mode 100644 src/sagemaker/serve/builder/tf_serving_builder.py
 create mode 100644 src/sagemaker/serve/model_server/tensorflow_serving/__init__.py
 create mode 100644 src/sagemaker/serve/model_server/tensorflow_serving/inference.py
 create mode 100644 src/sagemaker/serve/model_server/tensorflow_serving/prepare.py
 create mode 100644 src/sagemaker/serve/model_server/tensorflow_serving/server.py
 create mode 100644 src/sagemaker/serve/utils/lineage_constants.py
 create mode 100644 src/sagemaker/serve/utils/lineage_utils.py
 create mode 100644 tests/data/serve_resources/mlflow/tensorflow/MLmodel
 create mode 100644 tests/data/serve_resources/mlflow/tensorflow/conda.yaml
 create mode 100644 tests/data/serve_resources/mlflow/tensorflow/python_env.yaml
 create mode 100644 tests/data/serve_resources/mlflow/tensorflow/registered_model_meta
 create mode 100644 tests/data/serve_resources/mlflow/tensorflow/requirements.txt
 create mode 100644 tests/data/serve_resources/mlflow/tensorflow/tf2model/fingerprint.pb
 create mode 100644 tests/data/serve_resources/mlflow/tensorflow/tf2model/saved_model.pb
 create mode 100644 tests/data/serve_resources/mlflow/tensorflow/tf2model/variables/variables.data-00000-of-00001
 create mode 100644 tests/data/serve_resources/mlflow/tensorflow/tf2model/variables/variables.index
 create mode 100644 tests/integ/sagemaker/serve/test_serve_mlflow_tensorflow_flavor_happy.py
 create mode 100644 tests/unit/sagemaker/serve/builder/test_tensorflow_serving_builder.py
 create mode 100644 tests/unit/sagemaker/serve/model_server/tensorflow_serving/test_tf_prepare.py
 create mode 100644 tests/unit/sagemaker/serve/model_server/tensorflow_serving/test_tf_server.py
 create mode 100644 tests/unit/sagemaker/serve/utils/test_lineage_utils.py

diff --git a/requirements/extras/test_requirements.txt b/requirements/extras/test_requirements.txt
index 2713f88c04..ad6322365f 100644
--- a/requirements/extras/test_requirements.txt
+++ b/requirements/extras/test_requirements.txt
@@ -36,3 +36,4 @@ onnx>=1.15.0
 nbformat>=5.9,<6
 accelerate>=0.24.1,<=0.27.0
 schema==0.7.5
+tensorflow>=2.1,<=2.16
diff --git a/src/sagemaker/serve/builder/model_builder.py b/src/sagemaker/serve/builder/model_builder.py
index 06b3d70aeb..52268ea40c 100644
--- a/src/sagemaker/serve/builder/model_builder.py
+++ b/src/sagemaker/serve/builder/model_builder.py
@@ -29,6 +29,7 @@
 from sagemaker.serializers import NumpySerializer, TorchTensorSerializer
 from sagemaker.deserializers import JSONDeserializer, TorchTensorDeserializer
 from sagemaker.serve.builder.schema_builder import SchemaBuilder
+from sagemaker.serve.builder.tf_serving_builder import TensorflowServing
 from sagemaker.serve.mode.function_pointers import Mode
 from sagemaker.serve.mode.sagemaker_endpoint_mode import SageMakerEndpointMode
 from sagemaker.serve.mode.local_container_mode import LocalContainerMode
@@ -59,6 +60,7 @@
 from sagemaker.serve.spec.inference_spec import InferenceSpec
 from sagemaker.serve.utils import task
 from sagemaker.serve.utils.exceptions import TaskNotFoundException
+from sagemaker.serve.utils.lineage_utils import _maintain_lineage_tracking_for_mlflow_model
 from sagemaker.serve.utils.predictors import _get_local_mode_predictor
 from sagemaker.serve.utils.hardware_detector import (
     _get_gpu_info,
@@ -89,12 +91,13 @@
     ModelServer.TORCHSERVE,
     ModelServer.TRITON,
     ModelServer.DJL_SERVING,
+    ModelServer.TENSORFLOW_SERVING,
 }
 
 
-# pylint: disable=attribute-defined-outside-init, disable=E1101
+# pylint: disable=attribute-defined-outside-init, disable=E1101, disable=R0901
 @dataclass
-class ModelBuilder(Triton, DJL, JumpStart, TGI, Transformers):
+class ModelBuilder(Triton, DJL, JumpStart, TGI, Transformers, TensorflowServing):
     """Class that builds a deployable model.
 
     Args:
@@ -493,6 +496,12 @@ def _model_builder_register_wrapper(self, *args, **kwargs):
         self.pysdk_model.model_package_arn = new_model_package.model_package_arn
         new_model_package.deploy = self._model_builder_deploy_model_package_wrapper
         self.model_package = new_model_package
+        if getattr(self, "_is_mlflow_model", False) and self.mode == Mode.SAGEMAKER_ENDPOINT:
+            _maintain_lineage_tracking_for_mlflow_model(
+                mlflow_model_path=self.model_metadata[MLFLOW_MODEL_PATH],
+                s3_upload_path=self.s3_upload_path,
+                sagemaker_session=self.sagemaker_session,
+            )
         return new_model_package
 
     def _model_builder_deploy_model_package_wrapper(self, *args, **kwargs):
@@ -551,12 +560,19 @@ def _model_builder_deploy_wrapper(
 
         if "endpoint_logging" not in kwargs:
             kwargs["endpoint_logging"] = True
-        return self._original_deploy(
+        predictor = self._original_deploy(
             *args,
             instance_type=instance_type,
             initial_instance_count=initial_instance_count,
             **kwargs,
         )
+        if getattr(self, "_is_mlflow_model", False) and self.mode == Mode.SAGEMAKER_ENDPOINT:
+            _maintain_lineage_tracking_for_mlflow_model(
+                mlflow_model_path=self.model_metadata[MLFLOW_MODEL_PATH],
+                s3_upload_path=self.s3_upload_path,
+                sagemaker_session=self.sagemaker_session,
+            )
+        return predictor
 
     def _overwrite_mode_in_deploy(self, overwrite_mode: str):
         """Mode overwritten by customer during model.deploy()"""
@@ -728,7 +744,7 @@ def build(  # pylint: disable=R0911
                 " for production at this moment."
             )
             self._initialize_for_mlflow()
-            _validate_input_for_mlflow(self.model_server)
+            _validate_input_for_mlflow(self.model_server, self.env_vars.get("MLFLOW_MODEL_FLAVOR"))
 
         if isinstance(self.model, str):
             model_task = None
@@ -767,6 +783,9 @@ def build(  # pylint: disable=R0911
         if self.model_server == ModelServer.TRITON:
             return self._build_for_triton()
 
+        if self.model_server == ModelServer.TENSORFLOW_SERVING:
+            return self._build_for_tensorflow_serving()
+
         raise ValueError("%s model server is not supported" % self.model_server)
 
     def save(
diff --git a/src/sagemaker/serve/builder/tf_serving_builder.py b/src/sagemaker/serve/builder/tf_serving_builder.py
new file mode 100644
index 0000000000..42c548f4e4
--- /dev/null
+++ b/src/sagemaker/serve/builder/tf_serving_builder.py
@@ -0,0 +1,129 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You
+# may not use this file except in compliance with the License. A copy of
+# the License is located at
+#
+#     http://aws.amazon.com/apache2.0/
+#
+# or in the "license" file accompanying this file. This file is
+# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
+# ANY KIND, either express or implied. See the License for the specific
+# language governing permissions and limitations under the License.
+"""Holds mixin logic to support deployment of Model ID"""
+from __future__ import absolute_import
+import logging
+import os
+from pathlib import Path
+from abc import ABC, abstractmethod
+
+from sagemaker import Session
+from sagemaker.serve.detector.pickler import save_pkl
+from sagemaker.serve.model_server.tensorflow_serving.prepare import prepare_for_tf_serving
+from sagemaker.tensorflow import TensorFlowModel, TensorFlowPredictor
+
+logger = logging.getLogger(__name__)
+
+_TF_SERVING_MODEL_BUILDER_ENTRY_POINT = "inference.py"
+_CODE_FOLDER = "code"
+
+
+# pylint: disable=attribute-defined-outside-init, disable=E1101
+class TensorflowServing(ABC):
+    """TensorflowServing build logic for ModelBuilder()"""
+
+    def __init__(self):
+        self.model = None
+        self.serve_settings = None
+        self.sagemaker_session = None
+        self.model_path = None
+        self.dependencies = None
+        self.modes = None
+        self.mode = None
+        self.model_server = None
+        self.image_uri = None
+        self._is_custom_image_uri = False
+        self.image_config = None
+        self.vpc_config = None
+        self._original_deploy = None
+        self.secret_key = None
+        self.engine = None
+        self.pysdk_model = None
+        self.schema_builder = None
+        self.env_vars = None
+
+    @abstractmethod
+    def _prepare_for_mode(self):
+        """Prepare model artifacts based on mode."""
+
+    @abstractmethod
+    def _get_client_translators(self):
+        """Set up client marshaller based on schema builder."""
+
+    def _save_schema_builder(self):
+        """Save schema builder for tensorflow serving."""
+        if not os.path.exists(self.model_path):
+            os.makedirs(self.model_path)
+
+        code_path = Path(self.model_path).joinpath("code")
+        save_pkl(code_path, self.schema_builder)
+
+    def _get_tensorflow_predictor(
+        self, endpoint_name: str, sagemaker_session: Session
+    ) -> TensorFlowPredictor:
+        """Creates a TensorFlowPredictor object"""
+        serializer, deserializer = self._get_client_translators()
+
+        return TensorFlowPredictor(
+            endpoint_name=endpoint_name,
+            sagemaker_session=sagemaker_session,
+            serializer=serializer,
+            deserializer=deserializer,
+        )
+
+    def _validate_for_tensorflow_serving(self):
+        """Validate for tensorflow serving"""
+        if not getattr(self, "_is_mlflow_model", False):
+            raise ValueError("Tensorflow Serving is currently only supported for mlflow models.")
+
+    def _create_tensorflow_model(self):
+        """Creates a TensorFlow model object"""
+        self.pysdk_model = TensorFlowModel(
+            image_uri=self.image_uri,
+            image_config=self.image_config,
+            vpc_config=self.vpc_config,
+            model_data=self.s3_upload_path,
+            role=self.serve_settings.role_arn,
+            env=self.env_vars,
+            sagemaker_session=self.sagemaker_session,
+            predictor_cls=self._get_tensorflow_predictor,
+        )
+
+        self.pysdk_model.mode = self.mode
+        self.pysdk_model.modes = self.modes
+        self.pysdk_model.serve_settings = self.serve_settings
+
+        self._original_deploy = self.pysdk_model.deploy
+        self.pysdk_model.deploy = self._model_builder_deploy_wrapper
+        self._original_register = self.pysdk_model.register
+        self.pysdk_model.register = self._model_builder_register_wrapper
+        self.model_package = None
+        return self.pysdk_model
+
+    def _build_for_tensorflow_serving(self):
+        """Build the model for Tensorflow Serving"""
+        self._validate_for_tensorflow_serving()
+        self._save_schema_builder()
+
+        if not self.image_uri:
+            raise ValueError("image_uri is not set for tensorflow serving")
+
+        self.secret_key = prepare_for_tf_serving(
+            model_path=self.model_path,
+            shared_libs=self.shared_libs,
+            dependencies=self.dependencies,
+        )
+
+        self._prepare_for_mode()
+
+        return self._create_tensorflow_model()
diff --git a/src/sagemaker/serve/mode/local_container_mode.py b/src/sagemaker/serve/mode/local_container_mode.py
index 362a3804de..f940e2959c 100644
--- a/src/sagemaker/serve/mode/local_container_mode.py
+++ b/src/sagemaker/serve/mode/local_container_mode.py
@@ -11,6 +11,7 @@
 import docker
 
 from sagemaker.base_predictor import PredictorBase
+from sagemaker.serve.model_server.tensorflow_serving.server import LocalTensorflowServing
 from sagemaker.serve.spec.inference_spec import InferenceSpec
 from sagemaker.serve.builder.schema_builder import SchemaBuilder
 from sagemaker.serve.utils.logging_agent import pull_logs
@@ -34,7 +35,12 @@
 
 
 class LocalContainerMode(
-    LocalTorchServe, LocalDJLServing, LocalTritonServer, LocalTgiServing, LocalMultiModelServer
+    LocalTorchServe,
+    LocalDJLServing,
+    LocalTritonServer,
+    LocalTgiServing,
+    LocalMultiModelServer,
+    LocalTensorflowServing,
 ):
     """A class that holds methods to deploy model to a container in local environment"""
 
@@ -141,6 +147,15 @@ def create_server(
                 env_vars=env_vars if env_vars else self.env_vars,
             )
             self._ping_container = self._multi_model_server_deep_ping
+        elif self.model_server == ModelServer.TENSORFLOW_SERVING:
+            self._start_tensorflow_serving(
+                client=self.client,
+                image=image,
+                model_path=model_path if model_path else self.model_path,
+                secret_key=secret_key,
+                env_vars=env_vars if env_vars else self.env_vars,
+            )
+            self._ping_container = self._tensorflow_serving_deep_ping
 
         # allow some time for container to be ready
         time.sleep(10)
diff --git a/src/sagemaker/serve/mode/sagemaker_endpoint_mode.py b/src/sagemaker/serve/mode/sagemaker_endpoint_mode.py
index 0fdc425b31..24acfc6a2f 100644
--- a/src/sagemaker/serve/mode/sagemaker_endpoint_mode.py
+++ b/src/sagemaker/serve/mode/sagemaker_endpoint_mode.py
@@ -6,6 +6,7 @@
 import logging
 from typing import Type
 
+from sagemaker.serve.model_server.tensorflow_serving.server import SageMakerTensorflowServing
 from sagemaker.session import Session
 from sagemaker.serve.utils.types import ModelServer
 from sagemaker.serve.spec.inference_spec import InferenceSpec
@@ -24,6 +25,7 @@ class SageMakerEndpointMode(
     SageMakerDjlServing,
     SageMakerTgiServing,
     SageMakerMultiModelServer,
+    SageMakerTensorflowServing,
 ):
     """Holds the required method to deploy a model to a SageMaker Endpoint"""
 
@@ -107,4 +109,13 @@ def prepare(
                 image=image,
             )
 
+        if self.model_server == ModelServer.TENSORFLOW_SERVING:
+            return self._upload_tensorflow_serving_artifacts(
+                model_path=model_path,
+                sagemaker_session=sagemaker_session,
+                secret_key=secret_key,
+                s3_model_data_url=s3_model_data_url,
+                image=image,
+            )
+
         raise ValueError("%s model server is not supported" % self.model_server)
diff --git a/src/sagemaker/serve/model_format/mlflow/constants.py b/src/sagemaker/serve/model_format/mlflow/constants.py
index 00ef76170c..5d2ac484b5 100644
--- a/src/sagemaker/serve/model_format/mlflow/constants.py
+++ b/src/sagemaker/serve/model_format/mlflow/constants.py
@@ -19,6 +19,12 @@
     "py39": "1.13.1",
     "py310": "2.2.0",
 }
+MODEL_PACAKGE_ARN_REGEX = (
+    r"^arn:aws:sagemaker:[a-z0-9\-]+:[0-9]{12}:model-package\/[" r"a-zA-Z0-9\-_\/\.]+$"
+)
+MLFLOW_RUN_ID_REGEX = r"^runs:/[a-zA-Z0-9]+(/[a-zA-Z0-9]+)*$"
+MLFLOW_REGISTRY_PATH_REGEX = r"^models:/[a-zA-Z0-9\-_\.]+(/[0-9]+)*$"
+S3_PATH_REGEX = r"^s3:\/\/[a-zA-Z0-9\-_\.]+\/[a-zA-Z0-9\-_\/\.]*$"
 MLFLOW_MODEL_PATH = "MLFLOW_MODEL_PATH"
 MLFLOW_METADATA_FILE = "MLmodel"
 MLFLOW_PIP_DEPENDENCY_FILE = "requirements.txt"
@@ -34,8 +40,12 @@
     "spark": "pyspark",
     "onnx": "onnxruntime",
 }
-FLAVORS_WITH_FRAMEWORK_SPECIFIC_DLC_SUPPORT = [  # will extend to keras and tf
-    "sklearn",
-    "pytorch",
-    "xgboost",
-]
+TENSORFLOW_SAVED_MODEL_NAME = "saved_model.pb"
+FLAVORS_WITH_FRAMEWORK_SPECIFIC_DLC_SUPPORT = {
+    "sklearn": "sklearn",
+    "pytorch": "pytorch",
+    "xgboost": "xgboost",
+    "tensorflow": "tensorflow",
+    "keras": "tensorflow",
+}
+FLAVORS_DEFAULT_WITH_TF_SERVING = ["keras", "tensorflow"]
diff --git a/src/sagemaker/serve/model_format/mlflow/utils.py b/src/sagemaker/serve/model_format/mlflow/utils.py
index c9a8093a79..b67de08d78 100644
--- a/src/sagemaker/serve/model_format/mlflow/utils.py
+++ b/src/sagemaker/serve/model_format/mlflow/utils.py
@@ -13,7 +13,8 @@
 """Holds the util functions used for MLflow model format"""
 from __future__ import absolute_import
 
-from typing import Optional, Dict, Any
+from pathlib import Path
+from typing import Optional, Dict, Any, Union
 import yaml
 import logging
 import shutil
@@ -30,6 +31,8 @@
     DEFAULT_PYTORCH_VERSION,
     MLFLOW_METADATA_FILE,
     MLFLOW_PIP_DEPENDENCY_FILE,
+    FLAVORS_DEFAULT_WITH_TF_SERVING,
+    TENSORFLOW_SAVED_MODEL_NAME,
 )
 
 logger = logging.getLogger(__name__)
@@ -44,7 +47,8 @@ def _get_default_model_server_for_mlflow(deployment_flavor: str) -> ModelServer:
     Returns:
         str: The model server chosen for given model flavor.
     """
-    # TODO: implement real logic here based on mlflow flavor
+    if deployment_flavor in FLAVORS_DEFAULT_WITH_TF_SERVING:
+        return ModelServer.TENSORFLOW_SERVING
     return ModelServer.TORCHSERVE
 
 
@@ -344,15 +348,16 @@ def _select_container_for_mlflow_model(
             f"specific DLC support. Defaulting to generic image..."
         )
         return _get_default_image_for_mlflow(python_version, region, instance_type)
-    framework_version = _get_framework_version_from_requirements(
-        deployment_flavor, requirement_path
-    )
+
+    framework_to_use = FLAVORS_WITH_FRAMEWORK_SPECIFIC_DLC_SUPPORT.get(deployment_flavor)
+    framework_version = _get_framework_version_from_requirements(framework_to_use, requirement_path)
 
     logger.info("Auto-detected deployment flavor is %s", deployment_flavor)
+    logger.info("Auto-detected framework to use is %s", framework_to_use)
     logger.info("Auto-detected framework version is %s", framework_version)
 
     casted_versions = (
-        _cast_to_compatible_version(deployment_flavor, framework_version)
+        _cast_to_compatible_version(framework_to_use, framework_version)
         if framework_version
         else (None,)
     )
@@ -361,7 +366,7 @@ def _select_container_for_mlflow_model(
     for casted_version in casted_versions:
         try:
             image_uri = image_uris.retrieve(
-                framework=deployment_flavor,
+                framework=framework_to_use,
                 region=region,
                 version=casted_version,
                 image_scope="inference",
@@ -392,17 +397,60 @@ def _select_container_for_mlflow_model(
     )
 
 
-def _validate_input_for_mlflow(model_server: ModelServer) -> None:
+def _validate_input_for_mlflow(model_server: ModelServer, deployment_flavor: str) -> None:
     """Validates arguments provided with mlflow models.
 
     Args:
         - model_server (ModelServer): Model server used for orchestrating mlflow model.
+        - deployment_flavor (str): The flavor mlflow model will be deployed with.
 
     Raises:
     - ValueError: If model server is not torchserve.
     """
-    if model_server != ModelServer.TORCHSERVE:
+    if model_server != ModelServer.TORCHSERVE and model_server != ModelServer.TENSORFLOW_SERVING:
         raise ValueError(
             f"{model_server} is currently not supported for MLflow Model. "
             f"Please choose another model server."
         )
+    if (
+        model_server == ModelServer.TENSORFLOW_SERVING
+        and deployment_flavor not in FLAVORS_DEFAULT_WITH_TF_SERVING
+    ):
+        raise ValueError(
+            "Tensorflow Serving is currently only supported for the following "
+            "deployment flavors: {}".format(FLAVORS_DEFAULT_WITH_TF_SERVING)
+        )
+
+
+def _get_saved_model_path_for_tensorflow_and_keras_flavor(model_path: str) -> Optional[str]:
+    """Recursively searches for tensorflow saved model.
+
+    Args:
+        model_path (str): The root directory to start the search from.
+
+    Returns:
+        Optional[str]: The absolute path to the directory containing 'saved_model.pb'.
+    """
+    for dirpath, dirnames, filenames in os.walk(model_path):
+        if TENSORFLOW_SAVED_MODEL_NAME in filenames:
+            return os.path.abspath(dirpath)
+
+    return None
+
+
+def _move_contents(src_dir: Union[str, Path], dest_dir: Union[str, Path]) -> None:
+    """Moves all contents of a source directory to a specified destination directory.
+
+    Args:
+        src_dir (Union[str, Path]): The path to the source directory.
+        dest_dir (Union[str, Path]): The path to the destination directory.
+
+    """
+    _src_dir = Path(os.path.normpath(src_dir))
+    _dest_dir = Path(os.path.normpath(dest_dir))
+
+    _dest_dir.mkdir(parents=True, exist_ok=True)
+
+    for item in _src_dir.iterdir():
+        _dest_path = _dest_dir / item.name
+        shutil.move(str(item), str(_dest_path))
diff --git a/src/sagemaker/serve/model_server/tensorflow_serving/__init__.py b/src/sagemaker/serve/model_server/tensorflow_serving/__init__.py
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/src/sagemaker/serve/model_server/tensorflow_serving/inference.py b/src/sagemaker/serve/model_server/tensorflow_serving/inference.py
new file mode 100644
index 0000000000..928278e3c6
--- /dev/null
+++ b/src/sagemaker/serve/model_server/tensorflow_serving/inference.py
@@ -0,0 +1,147 @@
+"""This module is for SageMaker inference.py."""
+
+from __future__ import absolute_import
+import os
+import io
+import json
+import cloudpickle
+import shutil
+import platform
+from pathlib import Path
+from sagemaker.serve.validations.check_integrity import perform_integrity_check
+import logging
+
+logger = logging.getLogger(__name__)
+
+schema_builder = None
+SHARED_LIBS_DIR = Path(__file__).parent.parent.joinpath("shared_libs")
+SERVE_PATH = Path(__file__).parent.joinpath("serve.pkl")
+METADATA_PATH = Path(__file__).parent.joinpath("metadata.json")
+
+
+def input_handler(data, context):
+    """Pre-process request input before it is sent to TensorFlow Serving REST API
+
+    Args:
+        data (obj): the request data, in format of dict or string
+        context (Context): an object containing request and configuration details
+    Returns:
+        (dict): a JSON-serializable dict that contains request body and headers
+    """
+    read_data = data.read()
+    deserialized_data = None
+    try:
+        if hasattr(schema_builder, "custom_input_translator"):
+            deserialized_data = schema_builder.custom_input_translator.deserialize(
+                io.BytesIO(read_data), context.request_content_type
+            )
+        else:
+            deserialized_data = schema_builder.input_deserializer.deserialize(
+                io.BytesIO(read_data), context.request_content_type
+            )
+    except Exception as e:
+        logger.error("Encountered error: %s in deserialize_request." % e)
+        raise Exception("Encountered error in deserialize_request.") from e
+
+    try:
+        return json.dumps({"instances": _convert_for_serialization(deserialized_data)})
+    except Exception as e:
+        logger.error(
+            "Encountered error: %s in deserialize_request. "
+            "Deserialized data is not json serializable. " % e
+        )
+        raise Exception("Encountered error in deserialize_request.") from e
+
+
+def output_handler(data, context):
+    """Post-process TensorFlow Serving output before it is returned to the client.
+
+    Args:
+        data (obj): the TensorFlow serving response
+        context (Context): an object containing request and configuration details
+    Returns:
+        (bytes, string): data to return to client, response content type
+    """
+    if data.status_code != 200:
+        raise ValueError(data.content.decode("utf-8"))
+
+    response_content_type = context.accept_header
+    prediction = data.content
+    try:
+        prediction_dict = json.loads(prediction.decode("utf-8"))
+        if hasattr(schema_builder, "custom_output_translator"):
+            return (
+                schema_builder.custom_output_translator.serialize(
+                    prediction_dict["predictions"], response_content_type
+                ),
+                response_content_type,
+            )
+        else:
+            return schema_builder.output_serializer.serialize(prediction), response_content_type
+    except Exception as e:
+        logger.error("Encountered error: %s in serialize_response." % e)
+        raise Exception("Encountered error in serialize_response.") from e
+
+
+def _run_preflight_diagnostics():
+    _py_vs_parity_check()
+    _pickle_file_integrity_check()
+
+
+def _py_vs_parity_check():
+    container_py_vs = platform.python_version()
+    local_py_vs = os.getenv("LOCAL_PYTHON")
+
+    if not local_py_vs or container_py_vs.split(".")[1] != local_py_vs.split(".")[1]:
+        logger.warning(
+            f"The local python version {local_py_vs} differs from the python version "
+            f"{container_py_vs} on the container. Please align the two to avoid unexpected behavior"
+        )
+
+
+def _pickle_file_integrity_check():
+    with open(SERVE_PATH, "rb") as f:
+        buffer = f.read()
+
+    perform_integrity_check(buffer=buffer, metadata_path=METADATA_PATH)
+
+
+def _set_up_schema_builder():
+    """Sets up the schema_builder object."""
+    global schema_builder
+    with open(SERVE_PATH, "rb") as f:
+        schema_builder = cloudpickle.load(f)
+
+
+def _set_up_shared_libs():
+    """Sets up the shared libs path."""
+    if SHARED_LIBS_DIR.exists():
+        # before importing, place dynamic linked libraries in shared lib path
+        shutil.copytree(SHARED_LIBS_DIR, "/lib", dirs_exist_ok=True)
+
+
+def _convert_for_serialization(deserialized_data):
+    """Attempt to convert non-serializable objects to a serializable form.
+
+    Args:
+        deserialized_data: The object to convert.
+
+    Returns:
+        The converted object if it was not originally serializable, otherwise the original object.
+    """
+    import numpy as np
+    import pandas as pd
+
+    if isinstance(deserialized_data, np.ndarray):
+        return deserialized_data.tolist()
+    elif isinstance(deserialized_data, pd.DataFrame):
+        return deserialized_data.to_dict(orient="list")
+    elif isinstance(deserialized_data, pd.Series):
+        return deserialized_data.tolist()
+    return deserialized_data
+
+
+# on import, execute
+_run_preflight_diagnostics()
+_set_up_schema_builder()
+_set_up_shared_libs()
diff --git a/src/sagemaker/serve/model_server/tensorflow_serving/prepare.py b/src/sagemaker/serve/model_server/tensorflow_serving/prepare.py
new file mode 100644
index 0000000000..e9aa4aafff
--- /dev/null
+++ b/src/sagemaker/serve/model_server/tensorflow_serving/prepare.py
@@ -0,0 +1,67 @@
+"""Module for artifacts preparation for tensorflow_serving"""
+
+from __future__ import absolute_import
+from pathlib import Path
+import shutil
+from typing import List, Dict, Any
+
+from sagemaker.serve.model_format.mlflow.utils import (
+    _get_saved_model_path_for_tensorflow_and_keras_flavor,
+    _move_contents,
+)
+from sagemaker.serve.detector.dependency_manager import capture_dependencies
+from sagemaker.serve.validations.check_integrity import (
+    generate_secret_key,
+    compute_hash,
+)
+from sagemaker.remote_function.core.serialization import _MetaData
+
+
+def prepare_for_tf_serving(
+    model_path: str,
+    shared_libs: List[str],
+    dependencies: Dict[str, Any],
+) -> str:
+    """Prepares the model for serving.
+
+    Args:
+        model_path (str): Path to the model directory.
+        shared_libs (List[str]): List of shared libraries.
+        dependencies (Dict[str, Any]): Dictionary of dependencies.
+
+    Returns:
+        str: Secret key.
+    """
+
+    _model_path = Path(model_path)
+    if not _model_path.exists():
+        _model_path.mkdir()
+    elif not _model_path.is_dir():
+        raise Exception("model_dir is not a valid directory")
+
+    code_dir = _model_path.joinpath("code")
+    code_dir.mkdir(exist_ok=True)
+    shutil.copy2(Path(__file__).parent.joinpath("inference.py"), code_dir)
+
+    shared_libs_dir = _model_path.joinpath("shared_libs")
+    shared_libs_dir.mkdir(exist_ok=True)
+    for shared_lib in shared_libs:
+        shutil.copy2(Path(shared_lib), shared_libs_dir)
+
+    capture_dependencies(dependencies=dependencies, work_dir=code_dir)
+
+    saved_model_bundle_dir = _model_path.joinpath("1")
+    saved_model_bundle_dir.mkdir(exist_ok=True)
+    mlflow_saved_model_dir = _get_saved_model_path_for_tensorflow_and_keras_flavor(model_path)
+    if not mlflow_saved_model_dir:
+        raise ValueError("SavedModel is not found for Tensorflow or Keras flavor.")
+    _move_contents(src_dir=mlflow_saved_model_dir, dest_dir=saved_model_bundle_dir)
+
+    secret_key = generate_secret_key()
+    with open(str(code_dir.joinpath("serve.pkl")), "rb") as f:
+        buffer = f.read()
+    hash_value = compute_hash(buffer=buffer, secret_key=secret_key)
+    with open(str(code_dir.joinpath("metadata.json")), "wb") as metadata:
+        metadata.write(_MetaData(hash_value).to_json())
+
+    return secret_key
diff --git a/src/sagemaker/serve/model_server/tensorflow_serving/server.py b/src/sagemaker/serve/model_server/tensorflow_serving/server.py
new file mode 100644
index 0000000000..2392287c61
--- /dev/null
+++ b/src/sagemaker/serve/model_server/tensorflow_serving/server.py
@@ -0,0 +1,139 @@
+"""Module for Local Tensorflow Server"""
+
+from __future__ import absolute_import
+
+import requests
+import logging
+import platform
+from pathlib import Path
+from sagemaker.base_predictor import PredictorBase
+from sagemaker.session import Session
+from sagemaker.serve.utils.exceptions import LocalModelInvocationException
+from sagemaker.s3_utils import determine_bucket_and_prefix, parse_s3_url
+from sagemaker import fw_utils
+from sagemaker.serve.utils.uploader import upload
+from sagemaker.local.utils import get_docker_host
+
+logger = logging.getLogger(__name__)
+
+
+class LocalTensorflowServing:
+    """LocalTensorflowServing class."""
+
+    def _start_tensorflow_serving(
+        self, client: object, image: str, model_path: str, secret_key: str, env_vars: dict
+    ):
+        """Starts a local tensorflow serving container.
+
+        Args:
+            client: Docker client
+            image: Image to use
+            model_path: Path to the model
+            secret_key: Secret key to use for authentication
+            env_vars: Environment variables to set
+        """
+        self.container = client.containers.run(
+            image,
+            "serve",
+            detach=True,
+            auto_remove=True,
+            network_mode="host",
+            volumes={
+                Path(model_path): {
+                    "bind": "/opt/ml/model",
+                    "mode": "rw",
+                },
+            },
+            environment={
+                "SAGEMAKER_SUBMIT_DIRECTORY": "/opt/ml/model/code",
+                "SAGEMAKER_PROGRAM": "inference.py",
+                "SAGEMAKER_SERVE_SECRET_KEY": secret_key,
+                "LOCAL_PYTHON": platform.python_version(),
+                **env_vars,
+            },
+        )
+
+    def _invoke_tensorflow_serving(self, request: object, content_type: str, accept: str):
+        """Invokes a local tensorflow serving container.
+
+        Args:
+            request: Request to send
+            content_type: Content type to use
+            accept: Accept to use
+        """
+        try:
+            response = requests.post(
+                f"http://{get_docker_host()}:8080/invocations",
+                data=request,
+                headers={"Content-Type": content_type, "Accept": accept},
+                timeout=60,  # this is what SageMaker Hosting uses as timeout
+            )
+            response.raise_for_status()
+            return response.content
+        except Exception as e:
+            raise Exception("Unable to send request to the local container server") from e
+
+    def _tensorflow_serving_deep_ping(self, predictor: PredictorBase):
+        """Checks if the local tensorflow serving container is up and running.
+
+        If the container is not up and running, it will raise an exception.
+        """
+        response = None
+        try:
+            response = predictor.predict(self.schema_builder.sample_input)
+            return (True, response)
+            # pylint: disable=broad-except
+        except Exception as e:
+            if "422 Client Error: Unprocessable Entity for url" in str(e):
+                raise LocalModelInvocationException(str(e))
+            return (False, response)
+
+        return (True, response)
+
+
+class SageMakerTensorflowServing:
+    """SageMakerTensorflowServing class."""
+
+    def _upload_tensorflow_serving_artifacts(
+        self,
+        model_path: str,
+        sagemaker_session: Session,
+        secret_key: str,
+        s3_model_data_url: str = None,
+        image: str = None,
+    ):
+        """Uploads the model artifacts to S3.
+
+        Args:
+            model_path: Path to the model
+            sagemaker_session: SageMaker session
+            secret_key: Secret key to use for authentication
+            s3_model_data_url: S3 model data URL
+            image: Image to use
+        """
+        if s3_model_data_url:
+            bucket, key_prefix = parse_s3_url(url=s3_model_data_url)
+        else:
+            bucket, key_prefix = None, None
+
+        code_key_prefix = fw_utils.model_code_key_prefix(key_prefix, None, image)
+
+        bucket, code_key_prefix = determine_bucket_and_prefix(
+            bucket=bucket, key_prefix=code_key_prefix, sagemaker_session=sagemaker_session
+        )
+
+        logger.debug(
+            "Uploading the model resources to bucket=%s, key_prefix=%s.", bucket, code_key_prefix
+        )
+        s3_upload_path = upload(sagemaker_session, model_path, bucket, code_key_prefix)
+        logger.debug("Model resources uploaded to: %s", s3_upload_path)
+
+        env_vars = {
+            "SAGEMAKER_SUBMIT_DIRECTORY": "/opt/ml/model/code",
+            "SAGEMAKER_PROGRAM": "inference.py",
+            "SAGEMAKER_REGION": sagemaker_session.boto_region_name,
+            "SAGEMAKER_CONTAINER_LOG_LEVEL": "10",
+            "SAGEMAKER_SERVE_SECRET_KEY": secret_key,
+            "LOCAL_PYTHON": platform.python_version(),
+        }
+        return s3_upload_path, env_vars
diff --git a/src/sagemaker/serve/utils/lineage_constants.py b/src/sagemaker/serve/utils/lineage_constants.py
new file mode 100644
index 0000000000..51be20739f
--- /dev/null
+++ b/src/sagemaker/serve/utils/lineage_constants.py
@@ -0,0 +1,28 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You
+# may not use this file except in compliance with the License. A copy of
+# the License is located at
+#
+#     http://aws.amazon.com/apache2.0/
+#
+# or in the "license" file accompanying this file. This file is
+# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
+# ANY KIND, either express or implied. See the License for the specific
+# language governing permissions and limitations under the License.
+"""Holds constants used for lineage support"""
+from __future__ import absolute_import
+
+
+LINEAGE_POLLER_INTERVAL_SECS = 15
+LINEAGE_POLLER_MAX_TIMEOUT_SECS = 120
+MODEL_BUILDER_MLFLOW_MODEL_PATH_LINEAGE_ARTIFACT_TYPE = "ModelBuilderInputModelData"
+MLFLOW_S3_PATH = "S3"
+MLFLOW_MODEL_PACKAGE_PATH = "ModelPackage"
+MLFLOW_RUN_ID = "MLflowRunId"
+MLFLOW_LOCAL_PATH = "Local"
+MLFLOW_REGISTRY_PATH = "MLflowRegistry"
+ERROR = "Error"
+CODE = "Code"
+CONTRIBUTED_TO = "ContributedTo"
+VALIDATION_EXCEPTION = "ValidationException"
diff --git a/src/sagemaker/serve/utils/lineage_utils.py b/src/sagemaker/serve/utils/lineage_utils.py
new file mode 100644
index 0000000000..b2e28d26c3
--- /dev/null
+++ b/src/sagemaker/serve/utils/lineage_utils.py
@@ -0,0 +1,277 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You
+# may not use this file except in compliance with the License. A copy of
+# the License is located at
+#
+#     http://aws.amazon.com/apache2.0/
+#
+# or in the "license" file accompanying this file. This file is
+# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
+# ANY KIND, either express or implied. See the License for the specific
+# language governing permissions and limitations under the License.
+"""Holds the util functions used for lineage tracking"""
+from __future__ import absolute_import
+
+import os
+import time
+import re
+import logging
+from typing import Optional, Union
+
+from botocore.exceptions import ClientError
+
+from sagemaker import Session
+from sagemaker.lineage._api_types import ArtifactSummary
+from sagemaker.lineage.artifact import Artifact
+from sagemaker.lineage.association import Association
+from sagemaker.lineage.query import LineageSourceEnum
+from sagemaker.serve.model_format.mlflow.constants import (
+    MLFLOW_RUN_ID_REGEX,
+    MODEL_PACAKGE_ARN_REGEX,
+    S3_PATH_REGEX,
+    MLFLOW_REGISTRY_PATH_REGEX,
+)
+from sagemaker.serve.utils.lineage_constants import (
+    LINEAGE_POLLER_MAX_TIMEOUT_SECS,
+    LINEAGE_POLLER_INTERVAL_SECS,
+    MLFLOW_S3_PATH,
+    MODEL_BUILDER_MLFLOW_MODEL_PATH_LINEAGE_ARTIFACT_TYPE,
+    MLFLOW_LOCAL_PATH,
+    MLFLOW_MODEL_PACKAGE_PATH,
+    MLFLOW_RUN_ID,
+    MLFLOW_REGISTRY_PATH,
+    CONTRIBUTED_TO,
+    ERROR,
+    CODE,
+    VALIDATION_EXCEPTION,
+)
+
+logger = logging.getLogger(__name__)
+
+
+def _load_artifact_by_source_uri(
+    source_uri: str, artifact_type: str, sagemaker_session: Session
+) -> Optional[ArtifactSummary]:
+    """Load lineage artifact by source uri
+
+    Arguments:
+        source_uri (str): The s3 uri used for uploading transfomred model artifacts.
+        artifact_type (str): The type of the lineage artifact.
+        sagemaker_session (Session): Session object which manages interactions
+            with Amazon SageMaker APIs and any other AWS services needed. If not specified, the
+            function creates one using the default AWS configuration chain.
+
+    Returns:
+        ArtifactSummary: The Artifact Summary for the provided S3 URI.
+    """
+    artifacts = Artifact.list(source_uri=source_uri, sagemaker_session=sagemaker_session)
+    for artifact_summary in artifacts:
+        if artifact_summary.artifact_type == artifact_type:
+            return artifact_summary
+    return None
+
+
+def _poll_lineage_artifact(
+    s3_uri: str, artifact_type: str, sagemaker_session: Session
+) -> Optional[ArtifactSummary]:
+    """Polls lineage artifacts by s3 path.
+
+    Arguments:
+        s3_uri (str): The S3 URI to check for artifacts.
+        artifact_type (str): The type of the lineage artifact.
+        sagemaker_session (Session): Session object which manages interactions
+            with Amazon SageMaker APIs and any other AWS services needed. If not specified, the
+            function creates one using the default AWS configuration chain.
+
+    Returns:
+        Optional[ArtifactSummary]: The artifact summary if found, otherwise None.
+    """
+    logger.info("Polling lineage artifact for model data in %s", s3_uri)
+    start_time = time.time()
+    while time.time() - start_time < LINEAGE_POLLER_MAX_TIMEOUT_SECS:
+        result = _load_artifact_by_source_uri(s3_uri, artifact_type, sagemaker_session)
+        if result is not None:
+            return result
+        time.sleep(LINEAGE_POLLER_INTERVAL_SECS)
+
+
+def _get_mlflow_model_path_type(mlflow_model_path: str) -> str:
+    """Identify mlflow model path type.
+
+    Args:
+        mlflow_model_path (str): The string to be identified.
+
+    Returns:
+        str: Description of what the input string is identified as.
+    """
+    mlflow_rub_id_pattern = MLFLOW_RUN_ID_REGEX
+    mlflow_registry_id_pattern = MLFLOW_REGISTRY_PATH_REGEX
+    sagemaker_arn_pattern = MODEL_PACAKGE_ARN_REGEX
+    s3_pattern = S3_PATH_REGEX
+
+    if re.match(mlflow_rub_id_pattern, mlflow_model_path):
+        return MLFLOW_RUN_ID
+    if re.match(mlflow_registry_id_pattern, mlflow_model_path):
+        return MLFLOW_REGISTRY_PATH
+    if re.match(sagemaker_arn_pattern, mlflow_model_path):
+        return MLFLOW_MODEL_PACKAGE_PATH
+    if re.match(s3_pattern, mlflow_model_path):
+        return MLFLOW_S3_PATH
+    if os.path.exists(mlflow_model_path):
+        return MLFLOW_LOCAL_PATH
+
+    raise ValueError(f"Invalid MLflow model path: {mlflow_model_path}")
+
+
+def _create_mlflow_model_path_lineage_artifact(
+    mlflow_model_path: str,
+    sagemaker_session: Session,
+) -> Optional[Artifact]:
+    """Creates a lineage artifact for the given MLflow model path.
+
+    Args:
+        mlflow_model_path (str): The path to the MLflow model.
+        sagemaker_session (Session): The SageMaker session object.
+
+    Returns:
+        Optional[Artifact]: The created lineage artifact, or None if an error occurred.
+    """
+    _artifact_name = _get_mlflow_model_path_type(mlflow_model_path)
+    properties = dict(
+        model_builder_input_model_data_type=_artifact_name,
+    )
+    try:
+        return Artifact.create(
+            source_uri=mlflow_model_path,
+            artifact_type=MODEL_BUILDER_MLFLOW_MODEL_PATH_LINEAGE_ARTIFACT_TYPE,
+            artifact_name=_artifact_name,
+            properties=properties,
+            sagemaker_session=sagemaker_session,
+        )
+    except ClientError as e:
+        if e.response[ERROR][CODE] == VALIDATION_EXCEPTION:
+            logger.info("Artifact already exists")
+        else:
+            logger.warning("Failed to create mlflow model path lineage artifact: %s", e)
+            raise e
+
+
+def _retrieve_and_create_if_not_exist_mlflow_model_path_lineage_artifact(
+    mlflow_model_path: str,
+    sagemaker_session: Session,
+) -> Optional[Union[Artifact, ArtifactSummary]]:
+    """Retrieves an existing artifact for the given MLflow model path or
+
+    creates a new one if it doesn't exist.
+
+    Args:
+        mlflow_model_path (str): The path to the MLflow model.
+        sagemaker_session (Session): Session object which manages interactions
+            with Amazon SageMaker APIs and any other AWS services needed. If not specified, the
+            function creates one using the default AWS configuration chain.
+
+
+    Returns:
+        Optional[Union[Artifact, ArtifactSummary]]: The existing or newly created artifact,
+            or None if an error occurred.
+    """
+    _loaded_artifact = _load_artifact_by_source_uri(
+        mlflow_model_path, MODEL_BUILDER_MLFLOW_MODEL_PATH_LINEAGE_ARTIFACT_TYPE, sagemaker_session
+    )
+    if _loaded_artifact is not None:
+        return _loaded_artifact
+    return _create_mlflow_model_path_lineage_artifact(
+        mlflow_model_path,
+        sagemaker_session,
+    )
+
+
+def _add_association_between_artifacts(
+    mlflow_model_path_artifact_arn: str,
+    autogenerated_model_data_artifact_arn: str,
+    sagemaker_session: Session,
+) -> None:
+    """Add association between mlflow model path artifact and autogenerated model data artifact.
+
+    Arguments:
+        mlflow_model_path_artifact_arn (str): The mlflow model path artifact.
+        autogenerated_model_data_artifact_arn (str): The autogenerated model data artifact.
+        sagemaker_session (Session): Session object which manages interactions
+            with Amazon SageMaker APIs and any other AWS services needed. If not specified, the
+            function creates one using the default AWS configuration chain.
+    """
+    _association_type = CONTRIBUTED_TO
+    _source_arn = mlflow_model_path_artifact_arn
+    _destination_arn = autogenerated_model_data_artifact_arn
+    try:
+        logger.info(
+            "Adding association with source_arn: "
+            "%s, destination_arn: %s and association_type: %s.",
+            _source_arn,
+            _destination_arn,
+            _association_type,
+        )
+        Association.create(
+            source_arn=_source_arn,
+            destination_arn=_destination_arn,
+            association_type=_association_type,
+            sagemaker_session=sagemaker_session,
+        )
+    except ClientError as e:
+        if e.response[ERROR][CODE] == VALIDATION_EXCEPTION:
+            logger.info("Association already exists")
+        else:
+            raise e
+
+
+def _maintain_lineage_tracking_for_mlflow_model(
+    mlflow_model_path: str,
+    s3_upload_path: str,
+    sagemaker_session: Session,
+) -> None:
+    """Maintains lineage tracking for an MLflow model by creating or retrieving artifacts.
+
+    Args:
+        mlflow_model_path (str): The path to the MLflow model.
+        s3_upload_path (str): The S3 path where the transformed model data is uploaded.
+        sagemaker_session (Session): Session object which manages interactions
+            with Amazon SageMaker APIs and any other AWS services needed. If not specified, the
+            function creates one using the default AWS configuration chain.
+    """
+    artifact_for_transformed_model_data = _poll_lineage_artifact(
+        s3_uri=s3_upload_path,
+        artifact_type=LineageSourceEnum.MODEL_DATA.value,
+        sagemaker_session=sagemaker_session,
+    )
+    if artifact_for_transformed_model_data:
+        mlflow_model_artifact = (
+            _retrieve_and_create_if_not_exist_mlflow_model_path_lineage_artifact(
+                mlflow_model_path=mlflow_model_path,
+                sagemaker_session=sagemaker_session,
+            )
+        )
+        if mlflow_model_artifact:
+            _mlflow_model_artifact_arn = (
+                mlflow_model_artifact.artifact_arn
+            )  # pylint: disable=E1101, disable=C0301
+            _artifact_for_transformed_model_data_arn = (
+                artifact_for_transformed_model_data.artifact_arn
+            )  # pylint: disable=C0301
+            _add_association_between_artifacts(
+                mlflow_model_path_artifact_arn=_mlflow_model_artifact_arn,
+                autogenerated_model_data_artifact_arn=_artifact_for_transformed_model_data_arn,
+                sagemaker_session=sagemaker_session,
+            )
+        else:
+            logger.warning(
+                "Unable to add association between autogenerated lineage "
+                "artifact for transformed model data and mlflow model path"
+                " lineage artifacts."
+            )
+    else:
+        logger.warning(
+            "Lineage artifact for transformed model data is not auto-created within "
+            "%s seconds, skipping creation of lineage artifact for mlflow model path",
+            LINEAGE_POLLER_MAX_TIMEOUT_SECS,
+        )
diff --git a/src/sagemaker/serve/utils/predictors.py b/src/sagemaker/serve/utils/predictors.py
index e0ff8f8ee1..866167c2c6 100644
--- a/src/sagemaker/serve/utils/predictors.py
+++ b/src/sagemaker/serve/utils/predictors.py
@@ -209,6 +209,47 @@ def delete_predictor(self):
         self._mode_obj.destroy_server()
 
 
+class TensorflowServingLocalPredictor(PredictorBase):
+    """Lightweight predictor for local deployment in LOCAL_CONTAINER modes"""
+
+    # TODO: change mode_obj to union of IN_PROCESS and LOCAL_CONTAINER objs
+    def __init__(
+        self,
+        mode_obj: Type[LocalContainerMode],
+        serializer=IdentitySerializer(),
+        deserializer=BytesDeserializer(),
+    ):
+        self._mode_obj = mode_obj
+        self.serializer = serializer
+        self.deserializer = deserializer
+
+    def predict(self, data):
+        """Placeholder docstring"""
+        return self.deserializer.deserialize(
+            io.BytesIO(
+                self._mode_obj._invoke_tensorflow_serving(
+                    self.serializer.serialize(data),
+                    self.content_type,
+                    self.accept[0],
+                )
+            )
+        )
+
+    @property
+    def content_type(self):
+        """The MIME type of the data sent to the inference endpoint."""
+        return self.serializer.CONTENT_TYPE
+
+    @property
+    def accept(self):
+        """The content type(s) that are expected from the inference endpoint."""
+        return self.deserializer.ACCEPT
+
+    def delete_predictor(self):
+        """Shut down and remove the container that you created in LOCAL_CONTAINER mode"""
+        self._mode_obj.destroy_server()
+
+
 def _get_local_mode_predictor(
     model_server: ModelServer,
     mode_obj: Type[LocalContainerMode],
@@ -223,6 +264,11 @@ def _get_local_mode_predictor(
     if model_server == ModelServer.TRITON:
         return TritonLocalPredictor(mode_obj=mode_obj)
 
+    if model_server == ModelServer.TENSORFLOW_SERVING:
+        return TensorflowServingLocalPredictor(
+            mode_obj=mode_obj, serializer=serializer, deserializer=deserializer
+        )
+
     raise ValueError("%s model server is not supported yet!" % model_server)
 
 
diff --git a/tests/data/serve_resources/mlflow/tensorflow/MLmodel b/tests/data/serve_resources/mlflow/tensorflow/MLmodel
new file mode 100644
index 0000000000..f00412149d
--- /dev/null
+++ b/tests/data/serve_resources/mlflow/tensorflow/MLmodel
@@ -0,0 +1,17 @@
+artifact_path: model
+flavors:
+  python_function:
+    env:
+      conda: conda.yaml
+      virtualenv: python_env.yaml
+    loader_module: mlflow.tensorflow
+    python_version: 3.10.13
+  tensorflow:
+    code: null
+    model_type: tf2-module
+    saved_model_dir: tf2model
+mlflow_version: 2.11.1
+model_size_bytes: 23823
+model_uuid: 40d2323944294fce898d8693455f60e8
+run_id: 592132312fb84935b201de2c027c54c6
+utc_time_created: '2024-04-01 19:47:15.396517'
diff --git a/tests/data/serve_resources/mlflow/tensorflow/conda.yaml b/tests/data/serve_resources/mlflow/tensorflow/conda.yaml
new file mode 100644
index 0000000000..90d8c300a0
--- /dev/null
+++ b/tests/data/serve_resources/mlflow/tensorflow/conda.yaml
@@ -0,0 +1,11 @@
+channels:
+- conda-forge
+dependencies:
+- python=3.10.13
+- pip<=23.3.1
+- pip:
+  - mlflow==2.11.1
+  - cloudpickle==2.2.1
+  - numpy==1.26.4
+  - tensorflow==2.16.1
+name: mlflow-env
diff --git a/tests/data/serve_resources/mlflow/tensorflow/python_env.yaml b/tests/data/serve_resources/mlflow/tensorflow/python_env.yaml
new file mode 100644
index 0000000000..9e09178b6c
--- /dev/null
+++ b/tests/data/serve_resources/mlflow/tensorflow/python_env.yaml
@@ -0,0 +1,7 @@
+python: 3.10.13
+build_dependencies:
+- pip==23.3.1
+- setuptools==68.2.2
+- wheel==0.41.2
+dependencies:
+- -r requirements.txt
diff --git a/tests/data/serve_resources/mlflow/tensorflow/registered_model_meta b/tests/data/serve_resources/mlflow/tensorflow/registered_model_meta
new file mode 100644
index 0000000000..5423c0e6c7
--- /dev/null
+++ b/tests/data/serve_resources/mlflow/tensorflow/registered_model_meta
@@ -0,0 +1,2 @@
+model_name: model
+model_version: '2'
diff --git a/tests/data/serve_resources/mlflow/tensorflow/requirements.txt b/tests/data/serve_resources/mlflow/tensorflow/requirements.txt
new file mode 100644
index 0000000000..2ff55b8e87
--- /dev/null
+++ b/tests/data/serve_resources/mlflow/tensorflow/requirements.txt
@@ -0,0 +1,4 @@
+mlflow==2.11.1
+cloudpickle==2.2.1
+numpy==1.26.4
+tensorflow==2.16.1
diff --git a/tests/data/serve_resources/mlflow/tensorflow/tf2model/fingerprint.pb b/tests/data/serve_resources/mlflow/tensorflow/tf2model/fingerprint.pb
new file mode 100644
index 0000000000..ba1e240ba5
--- /dev/null
+++ b/tests/data/serve_resources/mlflow/tensorflow/tf2model/fingerprint.pb
@@ -0,0 +1 @@
+���������
ď����n�����/ �ʢ������
(��������32

\ No newline at end of file
diff --git a/tests/data/serve_resources/mlflow/tensorflow/tf2model/saved_model.pb b/tests/data/serve_resources/mlflow/tensorflow/tf2model/saved_model.pb
new file mode 100644
index 0000000000000000000000000000000000000000..e48f2b59cc1ac074b08eeba09a7114d70023e2e2
GIT binary patch
literal 22966
zcmeG^TWlOzR(-qr<>U9Y6L+o?$LVCUo}MJ`E<1_S?2cnQ$<BmK64_38>~3ny-DRgI
z?e3-@juVhz_hU3<7qk*WXcr-b#m8!-WmmI8!>}t@5ZYfPeh?BL2&CZy5JJEQ64G$)
zt#@@*bvw=s$tW{&x*oUAJ@?#m&ppqZBJ>A8B=9f8^v{Rk3dEO6rMogs0NKrub^iMZ
ze;NKh#^1~_z=yHcqq<h0Fd`2Pjge?<^zcL?F7>78uMzSsLf{^ZEH#_ut?FI1QC2r9
z+O0YrfK)?k*4mAtmL8`=5Z_TNZ7mZoaafr_&@Xh*83sf6x?ECQn$o<l)-@X2sMRVg
zIum&p;V(dZrB-dW*f1n&?N+_rGT{sxaA~5YRhzYj2QZw1(JSR9M`@*2e{`$5s;MP6
zY8?JKNWEFos;zSC5e0~stGov!Lci4=QlihI1pFR>#9M0Xt#*aeOEu`W!9N**WW7-<
zwTm1;W&-+K4Yk^=*P5E5+Jkbx*l@g<w0bFwz>qw^=`Hnfil=CXLSn0-mRyj`XAU6n
z4gva>*4WbS$}9J^;)8mvTx~V!1dJ7}TczI6Hp{zOGo7S|ftIvpt6WuE<yzHvlTKzP
zVYsAKG@O)LrKFU~ji!Yu$sC6<wNj})RJO}F_th<Bvs}^a24dnCK{ANcJ98yQ<I>>?
zWG>EG35eaS-KxV|5WB4wA0Qtk1Rph#x><75JwpRPZc+({Izw{AF<+Tsh^f2frg+Dq
z=?JVq|Luxe)b8WBv_>aedDwFt;yPFAKP*;=rU_hzA<j-gJkP>;p7Y&hNFs#FmGaL1
z)}KWps6pQ<&Rng5Q;n>f&?gnAbOZ+U6`@qsZ86E?Fd_z_sMV6vtdxuVDI0~P-mWPp
z#L{eL(qLRssFa(n8NOg&<s{`OS0J&b?mT60bichVVsFTR#TFfezGAIfDsu)+11A?>
zVEk9(%Pwf8z5&B)_ti#8D_ukG<fP!Me?x7GYot@3nA%1uEg^3~wY|;dPovaKOSqI%
zoTd3k=_K5Q<k~N_HSMw1i-^G0G0@bT<GR^y(mp{-<jq%E6e2SO{(`{7S_=i<X1j7*
zZJ_MIk!htBTvN0k`L(gtl%@4@71_v;d2(zWg<ob2_+$6nWX3_f^gKn8wMk=}?P?Jl
zB#?fa<t>_Mwi>uxS(HXl?hfcran09SHP^dI7}R#PVjG!D;Vj^MiJb6V!N22t2<76c
zwx#V}YgZ~L6^r*dQH2ae{z9JAC3YGzIQw<A)zTVO6BgV6SQF@6M-hdSoDkn3WAepq
zf1AvJy|z(!_BFM-+E5WS6A~W|d>*8=_Qw9#5(=Bc6Yv89{kWLl!1cj*Fh2kG4#qo0
zQnl1_RdaQkIRgEhWwer5X-<4_B|%BsXm54C@dy|c6$k@qtp8+A!myC72(+Z_@TE!z
z1Ow@FCwK_wAQNwDjU6qUkk92WoXg9Hcl3+1xwBj6H;R`^@`VfYmx>F7{{FiJs>Uu#
z_02Sko35hIjlSnLX$lgB7>%Y<*?x0>j^;@Ayt&DqS<a>i6)iH;cg|NX52WJDYbU>d
zIr|6SSp1dpx3B#1M`thp&QN;!Cm;XkU*g;*SNN{2<jo!rck~#8Xd&Y6?U32qz%pQO
zk!<W=|M}k*ZD{fVq5W_AtBX)Kaho%t9grFqFa6!(LSkd-A0w&D4?p^+*T1{oy7HlX
zb@A7~nO^)K>4VqdS0QEY6x7LP*&w1EaWPLn?K3~0b<+!{O6MI;Ne?puwK|2bz)=eV
z-?wHxo0n%ZW7d~}7&@5*euDs*@Nvd_8semDh=bxIp+72l&mhOlK|&$L85I9WcGA7_
z2B-r8D8rPEju>mV0Xx3VPIYkV064pL4IS03QN9ql5>(PQ)pn(&>{6$Ok<E3$n}VYS
zM*CAAoBt8ljeCw?F^3=e4g5zmcuxu1Y*;7i%+R)t<vYh5l(E@@8LwsYBA=!D*3=Ju
za*_}-K2At>(Ab&cAjx=p3?`Ift-|;8@{Xo#t5tPNYbfPXlj6RZJ=sC0Q5Y^DXE^AT
zJq1i@AO|b8YUL4L=~jlKZyhN8V=OAgvKKsi+i7|Vk|3?zzB5Aj#TW4+CiiEPxik4Q
z@|n5okp=Q~63L$Bo94{wX*i`I$W5)GRf}kqpmn3R+6_&4*ih?rL^+?olv}t21G<Jv
zbv^a(W2Ru4Hv~&&a8?)L1?;(y>t!(|v%eiPh2S@0%kRcgNg7%HS`4m2A4<gdb`7<1
zfK-$ws<p<pvWwS}as3DeQZbrRlp<~*iULEaI0Xy&W){~COhY`CpyVO`Bq`azFH-gw
zqD3?t0sc1({i!~hL>mT4ki!RPLZpBWd$=86T%)KC-(6aLb7}d;Rpsu|jXPJJ%aLUS
z(z6d|&n>O5ufDl_XZ`Bhth>GF`3il#1CAw$Oeh=Opj<$_`z^fjX2873<ISLXBje4G
zc{7KjcpU~)!?drc*2Prk@OH$$&ExH;eLIi0)V@9c48Aw%9`t()hMK5-H-xPvv{{#w
z3ED~=(kjUXi?v_CpGDzFfEJ2EJWytMum~3r^VwltE}(iL$CBgYNw-i)b`uIg!<r72
z16-UutqhpqGT`Lu2{`VQ0g6}$XguK*@a3Psn1>HS2q+(rfY~9PZb^vC64LFtC);bB
zY^PUGLB@}4D5$xiDk!&bfxgp(5BuOp1SVaL4Z$T^<#IHReyKE}Ihfj3Tifl55d1ra
zo2gc}q<?z6Fo3erQV!@VFsYVGA<bi4x#UWI*>L@&xq-V!C_2PEs=^*Qq`P}O&shun
zUr?_y)dl-~=^N}zw)aU#i;-P+%ooFg@Rs1VljMFSuFMB(Hq%_GVw>PYfR6Em(YBW@
zj<6y~NNB3NAK8eGeZlaZ6>^MqPQ(5O^K-V*Zd&gwiGT}_og@iX{im~sJa#;)F^mwl
z>hKs^zC$)*cnAu*S%ssB-{bBpoBAPLv0F|}f{oqlBmg^n&v%xy;~i=fN<Z}7TD(n1
zO$VoiTXqHxp(=(3S_8e6-0#QjTWz;Rhw0#ht6VKxZKF?gTg#q+L0!8t4MlWK3O$^&
zOsWtmMC2J?w((&l`CQPN@n1@X8Ps!7qQJP)Mzup}*0o}Jvuv8%XyKR&Ahskuwp>a@
z=ocV-DWqUBpihaCUHJDdvTF`4%a^i?n`sX1gRs&;WY_09*!&QT6-Hgw!KW(3D$hcQ
z^*pMt_9V@EGNA&*+LxYEto738b1dJoUio$eCOX;bMHF<`QP5q39yzRw#QVuG{FL;h
z+M(IVJ?B;+iJnDeOaurWnwxp4<b8EFZ}Yrjet^w&-3=Y~rtYT9rISK?>j*@0h0YV3
zRPyX^2pnED11Gc2NJ)U%9=vf&P63z9J~A!$`f?S`Ywj$SmV2;%%{F1S?ebA&<&CQ=
zKYRPun>W{mp?z!l`qh<n<%_FJx8Gnp!0a&z8_yBzIq0~KnBy8*9<s(2S;qg__lPqJ
zWe22~lMC2!8#gDtEq#QYv>SMlw;Fl9DbpQ0!VnDb=$!eN9krra!i*v$iTbmpvQg`b
zjal6njEGnv=k3nYqSGDyqT7l0$^MA=1j9SHl7+1Y)~T#{5jhRh%6+xT{by~pA~d%~
zt%Baix>iMVrdY;^L3D=rD~3*nCG0Bzrf!HPe<(fjUNm~3$NLzI?p%Ph7mC<Ro{iZk
z2J*r)q{%NhRdMg257ZC)I{L}2!N?AJyB5pnJul6;LC`73bo;`Kp)oDVAWMwhao{lL
zz`-W7Y>}BIg$<+RO)DlC><vh8a|hW-ctcr_qwwd_%n7G>r#5PH7Ha!2VUBpWcpmFQ
z!^JRNIw(thm{va3%hK*{u-rA=7X=<oLqBtXnc<3O6jpF2UxWZx7`zey13l-Cow%FP
zlDqKI;bO#=UK%aV%$82iOOGlgVK0rAcFS~|JX5F4OQYpu!rf2dZkiAK;rk&j=G;Y0
zPSD*nJQJwBp*%uwGbS6MQ%ozBcf{=xNVcmFs<nsJ4ohVW`ixgf&a_K-`V03|b4Ylp
zq8A!3eLT&9;gxG3Au5z7vvAk0j+eRpl5CG_FU5Wyle_eBP&3htB@r@re=>WBPXYN)
zpC;0bxsZaWqI;91%qE~K-ht6$cf8|Miyf2bQ^$=?@c_v2)nmai4cuax3KH@M<*AWu
zIw1OH={QX1C3}u~^0izC3Hi?St|zfOL3EuN_X*42*PTT8I!up?yC?iU?^Li^`d$DP
z0yM*C=Gb2GYK1c7!m={->?q@3p*^DhDDOY%qHq)g@aZ6J4&TN@-k!z)e+Gm2{todJ
zP4LWw3v-3jm=w^xUA{1l4)&uE?BGaY3(}><8fvGvVBc9T;j;V3eehSLw=@h>i1*u~
zT8MF5I1@T)%j!;-lS7cpmLSJ{zyvU-Zwuz_{Qd;`KRk`NM9J?XByD+_w=ZS;a?ZY-
z|7;}W^dM9(=Gg;Zr%n&z^`-a<(nA8%7e0CMwUF;kd|kZs{EDydO~L!ot{Eh>!&yED
zCq#BzsVoXtideW(tBNP42i(VH{+jVo9@EW(KC|Q+)zgS}WIg_?%wwXq)0uwWaHjhd
zk?axln`825M^()pK?iC_R_Pu|^QV~4>GQ!M;}!ogeIl55{Sf@K%6*@@e9Ue?ST@N0
zg+YBQ5gp^l_nLEs+1Ybh_KWTW>0Y5ZTm{@UrWZQV!R3CDwiiE#nd&Xn`8>(cPX&F_
zgie=71aweZ+ID?bbn!JTmocg(j9+%lFFoH@Er&#YwtYWh<ZD|`yhwI+EN;5;`5n3a
zu6bg^I-X(u+@#TTG-w7fj2ys*w6orebze^YdM5i#0)9;0kqy(m12%C_td5RV5R#9h
zCpkN~iy((u{&|Pjc2a<LucA(mO>nk@=m<H~_DehC0NsNTRurIH7n=|$4T9$_OrizM
z+Y<u*vL#^V!n$Bb#LW3gupsVq$5-w;=<4+MALA2-V7tv*fXf%7SoVB%^WnSKu_GHS
z2sqZ|*tJk*i9|qS=Wr!dISm2Wx12-+BvQa^sZBVv#;dCADy&(I@HNmr6#~bL-Bb1x
zVZTK>%e}nyUu?`SII`pqoASlaem`bY!+O{w+U&1Uis&_SpOrv!6D?j*1jwhcj+L!e
zMPa~<1{;^2=y^z#`JJd$dT5j&GUl1Jf?2PJ?yaLfB)6sx-7fz4+b!&PQ{L~P9fV>c
z8QRvRylioxaG_uL`db7}cUJe#`jqxkQ8XJ$rsO@H5?-uPrt=+`n7*Q9NVyncSWoVw
zD^j0lY_>PlZJIj+SY9j-L|b;>AwlOI`>C@>47+wB>_wxiTD0cq>Aa&So5|+w0lRZX
zgQKnACN3wMAEJS32OHP*5%=a3ve#kisqcQ}zXQ8sx91(P+w;~i_Yu3rN>q3HDb~B_
zv0FWQ+c`mgoG0f3b;bZ)QqKCTe||clJRhRl1Zq7z8P=Dq)c<i|5x3Nzfq)GXP6YS_
zv4lt0?B}27;B>cY02wPkUuM*HWYm5_)c`zk?K4vipbzGt)c~<iuNols9B0c0Rs)E_
za4V&zqlyMoYC4|15Y+%;T;ZXLZn7(Hg_&PgG8rfTU@_7sRmqh8EXgi?Fb9=g=}%8~
zrC&tZ#i{%vCir)&WD?^Fm0g*5>Z8Gt0d@n=2XfUG7jkqmATF=u(ZUWa9F`ZN{B$S|
zt)z&$Z_HCJSCv|wpD>~uE!o$x`b<>PnoY{pa*H>4g`Nz01^Ty)+ogYqz8+5*I-VGX
zM}?7HI)|Nm&toX$MhFWzaDgsB$~X-rN6!RSMGG+Q^9UHz^TjvuFsxal-&s6}sZ=hu
z3O`MQa^R}{%u!VAw=_I?rd-sk;A$Ih^HY(q4gwD|D1_DXclmbefbahc56y8L5g`9<
zmt?CC2z-;gFX3@C{rtS3rN-787WCn9RpP`Wn);4}6_#(~91cPf!(IFsBz`ai9*Kl*
z#m0k&;${d^#yK1M%Ta&_5BXsbqPg@PDMGO#Qh(%XNTFY?R%`ssA}oBQ(tA=QdwNeg
z2`8RJr)UUwq{(CFFyV~vIS83UFle2%fIMF}CWsZ%m_cFHCZ~yP9ccixFoN$Qu4n%p
zMWXKe`~vw&G&t`lpUZ{3%;!U1VuDL>i8huwqSy5~n^I&LoGzPycw-*`fj9O65O`xB
z00BRY9_0Lo<v;1d`@E7p^L$dveeA_!_pujv>pb=XZ=J_pz>g0`H9QEdr;}E&=r)*q
z#))ZEto{9nj>9g|VT+&M<p*cBgOKk>B$RP*;f9X<|50+{$npG!aO~0W37fa)9D8&*
zy};O`Y<`ZhKYkqx@`PLR?eW&;xph2a=z)GvD;`OpR$Fl5dY_pzG1t9`+`UQOzX|>2
z;Qfu~hkuzShDQ>~QOf4_q#V4uCt=ObljsSqX@;r{!1asFu_sYxWhGL#>WE2&Uk>h^
zFa0mjXd=l@>`7BFu_uiKeG(n!9dl(uoW6huk&14{pG5I|MI%2^l$}m@fsRWv47#u!
zf}y9oK*yyM23`16f}zI{dNqVH$lA9AWwLwH=U{qI!j$VL(KM$_$8i8Dvzsn6@w>fl
tL`{n!7WSm`Ani#JIQO3P8obK;?};@sv+qe~AnPzWrk+G69k*US`9G+vekT9`

literal 0
HcmV?d00001

diff --git a/tests/data/serve_resources/mlflow/tensorflow/tf2model/variables/variables.data-00000-of-00001 b/tests/data/serve_resources/mlflow/tensorflow/tf2model/variables/variables.data-00000-of-00001
new file mode 100644
index 0000000000000000000000000000000000000000..575da962826e4f881fc73e05a9254344f10756b1
GIT binary patch
literal 553
zcmccPY-iu6ec0A&OZ?vTu~oL8*rwPjZ(!PY*qd|T<EF=ZZJ4fG{byDxtzzSH<>KUE
z6k^TIPf5+;;^JTuV#~`f%8jo8vYDamN-iD_79o!K_~gW#ocMSyJ`PqPuHww}yu^~y
zqSRt7CJsg}1uj+&HX+7xAjK}km;}_pfzSa~!KB5&1;YGX+#H-j>?K8s$*J)MKml<s
zVGb@K{^Hc4vdp~n_>|PN#L}D+poEK%9T#7iW00q#laFhBn4?dqs}M(6Vo_#dQckLr
z64=T5<@$P#At6DYPN5;L!TPAG!S)hmN)jPc%z&=uhPxWb;el}sfgD~K#|X&bgK><3
rez4=><=_|MOinBSg;l%(kSl=T8Une32(Af`D}>;hQb9i~)N#-N7x;xN

literal 0
HcmV?d00001

diff --git a/tests/data/serve_resources/mlflow/tensorflow/tf2model/variables/variables.index b/tests/data/serve_resources/mlflow/tensorflow/tf2model/variables/variables.index
new file mode 100644
index 0000000000000000000000000000000000000000..57646ac3509ce55f5f5e1f7c8f4e8570badfe5d7
GIT binary patch
literal 248
zcmZQzVB=tvV&Y(Akl~AW_HcFf4)FK%3vqPvagFzP@^W<!iFXfj4DjG!7h+IQ(P(8c
zWm!G*8iNvVZhlH?j((E9o?}Qzkf&2<h-<KZm}8J9SZA1{PpB&gBT$Qi28(H}yV))_
z8R2sLYB_|MI6#<DL&5Z{m+n^vAYcR$Tq5>0aL$2G3pX$@GUqWY0?UC3iR<^5EobD|
T!pIL&h989gZs=Aib-xV&T!t}k

literal 0
HcmV?d00001

diff --git a/tests/integ/sagemaker/serve/constants.py b/tests/integ/sagemaker/serve/constants.py
index 794f7333a3..d5e7a56f83 100644
--- a/tests/integ/sagemaker/serve/constants.py
+++ b/tests/integ/sagemaker/serve/constants.py
@@ -32,6 +32,7 @@
     DATA_DIR, "serve_resources", "mlflow", "pytorch"
 )
 XGBOOST_MLFLOW_RESOURCE_DIR = os.path.join(DATA_DIR, "serve_resources", "mlflow", "xgboost")
+TENSORFLOW_MLFLOW_RESOURCE_DIR = os.path.join(DATA_DIR, "serve_resources", "mlflow", "tensorflow")
 TF_EFFICIENT_RESOURCE_DIR = os.path.join(DATA_DIR, "serve_resources", "tensorflow")
 HF_DIR = os.path.join(DATA_DIR, "serve_resources", "hf")
 
diff --git a/tests/integ/sagemaker/serve/test_serve_mlflow_pytorch_flavor_happy.py b/tests/integ/sagemaker/serve/test_serve_mlflow_pytorch_flavor_happy.py
index e7ebd9c5bf..e6beb76d6e 100644
--- a/tests/integ/sagemaker/serve/test_serve_mlflow_pytorch_flavor_happy.py
+++ b/tests/integ/sagemaker/serve/test_serve_mlflow_pytorch_flavor_happy.py
@@ -19,6 +19,8 @@
 import io
 import numpy as np
 
+from sagemaker.lineage.artifact import Artifact
+from sagemaker.lineage.association import Association
 from sagemaker.s3 import S3Uploader
 from sagemaker.serve.builder.model_builder import ModelBuilder, Mode
 from sagemaker.serve.builder.schema_builder import SchemaBuilder, CustomPayloadTranslator
@@ -35,6 +37,10 @@
 from tests.integ.utils import cleanup_model_resources
 import logging
 
+from sagemaker.serve.utils.lineage_constants import (
+    MODEL_BUILDER_MLFLOW_MODEL_PATH_LINEAGE_ARTIFACT_TYPE,
+)
+
 logger = logging.getLogger(__name__)
 
 ROLE_NAME = "SageMakerRole"
@@ -205,6 +211,19 @@ def test_happy_pytorch_sagemaker_endpoint_with_torch_serve(
             predictor = model.deploy(instance_type=cpu_instance_type, initial_instance_count=1)
             logger.info("Endpoint successfully deployed.")
             predictor.predict(test_image)
+            model_data_artifact = None
+            for artifact in Artifact.list(
+                source_uri=model_builder.s3_upload_path, sagemaker_session=sagemaker_session
+            ):
+                model_data_artifact = artifact
+            for association in Association.list(
+                destination_arn=model_data_artifact.artifact_arn,
+                sagemaker_session=sagemaker_session,
+            ):
+                assert (
+                    association.source_type == MODEL_BUILDER_MLFLOW_MODEL_PATH_LINEAGE_ARTIFACT_TYPE
+                )
+                break
         except Exception as e:
             caught_ex = e
         finally:
@@ -214,9 +233,4 @@ def test_happy_pytorch_sagemaker_endpoint_with_torch_serve(
                 endpoint_name=model.endpoint_name,
             )
             if caught_ex:
-                logger.exception(caught_ex)
-                ignore_if_worker_dies = "Worker died." in str(caught_ex)
-                # https://github.com/pytorch/serve/issues/3032
-                assert (
-                    ignore_if_worker_dies
-                ), f"{caught_ex} was thrown when running pytorch squeezenet sagemaker endpoint test"
+                raise caught_ex
diff --git a/tests/integ/sagemaker/serve/test_serve_mlflow_tensorflow_flavor_happy.py b/tests/integ/sagemaker/serve/test_serve_mlflow_tensorflow_flavor_happy.py
new file mode 100644
index 0000000000..c25cbd7e18
--- /dev/null
+++ b/tests/integ/sagemaker/serve/test_serve_mlflow_tensorflow_flavor_happy.py
@@ -0,0 +1,175 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You
+# may not use this file except in compliance with the License. A copy of
+# the License is located at
+#
+#     http://aws.amazon.com/apache2.0/
+#
+# or in the "license" file accompanying this file. This file is
+# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
+# ANY KIND, either express or implied. See the License for the specific
+# language governing permissions and limitations under the License.
+from __future__ import absolute_import
+
+import pytest
+import io
+import numpy as np
+
+from sagemaker.lineage.artifact import Artifact
+from sagemaker.lineage.association import Association
+from sagemaker.s3 import S3Uploader
+from sagemaker.serve.builder.model_builder import ModelBuilder, Mode
+from sagemaker.serve.builder.schema_builder import SchemaBuilder, CustomPayloadTranslator
+import tensorflow as tf
+from sklearn.datasets import fetch_california_housing
+
+
+from tests.integ.sagemaker.serve.constants import (
+    TENSORFLOW_MLFLOW_RESOURCE_DIR,
+    SERVE_SAGEMAKER_ENDPOINT_TIMEOUT,
+    PYTHON_VERSION_IS_NOT_310,
+)
+from tests.integ.timeout import timeout
+from tests.integ.utils import cleanup_model_resources
+import logging
+
+from sagemaker.serve.utils.lineage_constants import (
+    MODEL_BUILDER_MLFLOW_MODEL_PATH_LINEAGE_ARTIFACT_TYPE,
+)
+
+logger = logging.getLogger(__name__)
+
+ROLE_NAME = "SageMakerRole"
+
+
+@pytest.fixture
+def test_data():
+    dataset = fetch_california_housing(as_frame=True)["frame"]
+    dataset = dataset.dropna()
+    dataset_tf = tf.convert_to_tensor(dataset, dtype=tf.float32)
+    dataset_tf = dataset_tf[:50]
+    x_test, y_test = dataset_tf[:, :-1], dataset_tf[:, -1]
+    return x_test, y_test
+
+
+@pytest.fixture
+def custom_request_translator():
+    class MyRequestTranslator(CustomPayloadTranslator):
+        def serialize_payload_to_bytes(self, payload: object) -> bytes:
+            return self._convert_numpy_to_bytes(payload)
+
+        def deserialize_payload_from_stream(self, stream) -> object:
+            np_array = np.load(io.BytesIO(stream.read()))
+            return np_array
+
+        def _convert_numpy_to_bytes(self, np_array: np.ndarray) -> bytes:
+            buffer = io.BytesIO()
+            np.save(buffer, np_array)
+            return buffer.getvalue()
+
+    return MyRequestTranslator()
+
+
+@pytest.fixture
+def custom_response_translator():
+    class MyResponseTranslator(CustomPayloadTranslator):
+        def serialize_payload_to_bytes(self, payload: object) -> bytes:
+            import numpy as np
+
+            return self._convert_numpy_to_bytes(np.array(payload))
+
+        def deserialize_payload_from_stream(self, stream) -> object:
+            import tensorflow as tf
+
+            return tf.convert_to_tensor(np.load(io.BytesIO(stream.read())))
+
+        def _convert_numpy_to_bytes(self, np_array: np.ndarray) -> bytes:
+            buffer = io.BytesIO()
+            np.save(buffer, np_array)
+            return buffer.getvalue()
+
+    return MyResponseTranslator()
+
+
+@pytest.fixture
+def tensorflow_schema_builder(custom_request_translator, custom_response_translator, test_data):
+    input_data, output_data = test_data
+    return SchemaBuilder(
+        sample_input=input_data,
+        sample_output=output_data,
+        input_translator=custom_request_translator,
+        output_translator=custom_response_translator,
+    )
+
+
+@pytest.mark.skipif(
+    PYTHON_VERSION_IS_NOT_310,
+    reason="The goal of these test are to test the serving components of our feature",
+)
+def test_happy_tensorflow_sagemaker_endpoint_with_tensorflow_serving(
+    sagemaker_session,
+    tensorflow_schema_builder,
+    cpu_instance_type,
+    test_data,
+):
+    logger.info("Running in SAGEMAKER_ENDPOINT mode...")
+    caught_ex = None
+
+    iam_client = sagemaker_session.boto_session.client("iam")
+    role_arn = iam_client.get_role(RoleName=ROLE_NAME)["Role"]["Arn"]
+
+    model_artifacts_uri = "s3://{}/{}/{}/{}".format(
+        sagemaker_session.default_bucket(),
+        "model_builder_integ_test",
+        "mlflow",
+        "tensorflow",
+    )
+
+    model_path = S3Uploader.upload(
+        local_path=TENSORFLOW_MLFLOW_RESOURCE_DIR,
+        desired_s3_uri=model_artifacts_uri,
+        sagemaker_session=sagemaker_session,
+    )
+
+    model_builder = ModelBuilder(
+        mode=Mode.SAGEMAKER_ENDPOINT,
+        schema_builder=tensorflow_schema_builder,
+        role_arn=role_arn,
+        sagemaker_session=sagemaker_session,
+        model_metadata={"MLFLOW_MODEL_PATH": model_path},
+    )
+
+    model = model_builder.build(sagemaker_session=sagemaker_session)
+
+    with timeout(minutes=SERVE_SAGEMAKER_ENDPOINT_TIMEOUT):
+        try:
+            test_x, _ = test_data
+            logger.info("Deploying and predicting in SAGEMAKER_ENDPOINT mode...")
+            predictor = model.deploy(instance_type=cpu_instance_type, initial_instance_count=1)
+            logger.info("Endpoint successfully deployed.")
+            predictor.predict(test_x)
+            model_data_artifact = None
+            for artifact in Artifact.list(
+                source_uri=model_builder.s3_upload_path, sagemaker_session=sagemaker_session
+            ):
+                model_data_artifact = artifact
+            for association in Association.list(
+                destination_arn=model_data_artifact.artifact_arn,
+                sagemaker_session=sagemaker_session,
+            ):
+                assert (
+                    association.source_type == MODEL_BUILDER_MLFLOW_MODEL_PATH_LINEAGE_ARTIFACT_TYPE
+                )
+                break
+
+        except Exception as e:
+            caught_ex = e
+        finally:
+            cleanup_model_resources(
+                sagemaker_session=model_builder.sagemaker_session,
+                model_name=model.name,
+                endpoint_name=model.endpoint_name,
+            )
+            if caught_ex:
+                raise caught_ex
diff --git a/tests/integ/sagemaker/serve/test_serve_mlflow_xgboost_flavor_happy.py b/tests/integ/sagemaker/serve/test_serve_mlflow_xgboost_flavor_happy.py
index 5a73942afe..7b47440a97 100644
--- a/tests/integ/sagemaker/serve/test_serve_mlflow_xgboost_flavor_happy.py
+++ b/tests/integ/sagemaker/serve/test_serve_mlflow_xgboost_flavor_happy.py
@@ -16,6 +16,8 @@
 import io
 import numpy as np
 
+from sagemaker.lineage.artifact import Artifact
+from sagemaker.lineage.association import Association
 from sagemaker.s3 import S3Uploader
 from sagemaker.serve.builder.model_builder import ModelBuilder, Mode
 from sagemaker.serve.builder.schema_builder import SchemaBuilder, CustomPayloadTranslator
@@ -32,6 +34,10 @@
 from tests.integ.utils import cleanup_model_resources
 import logging
 
+from sagemaker.serve.utils.lineage_constants import (
+    MODEL_BUILDER_MLFLOW_MODEL_PATH_LINEAGE_ARTIFACT_TYPE,
+)
+
 logger = logging.getLogger(__name__)
 
 ROLE_NAME = "SageMakerRole"
@@ -187,6 +193,19 @@ def test_happy_xgboost_sagemaker_endpoint_with_torch_serve(
             predictor = model.deploy(instance_type=cpu_instance_type, initial_instance_count=1)
             logger.info("Endpoint successfully deployed.")
             predictor.predict(test_x)
+            model_data_artifact = None
+            for artifact in Artifact.list(
+                source_uri=model_builder.s3_upload_path, sagemaker_session=sagemaker_session
+            ):
+                model_data_artifact = artifact
+            for association in Association.list(
+                destination_arn=model_data_artifact.artifact_arn,
+                sagemaker_session=sagemaker_session,
+            ):
+                assert (
+                    association.source_type == MODEL_BUILDER_MLFLOW_MODEL_PATH_LINEAGE_ARTIFACT_TYPE
+                )
+                break
         except Exception as e:
             caught_ex = e
         finally:
@@ -196,9 +215,4 @@ def test_happy_xgboost_sagemaker_endpoint_with_torch_serve(
                 endpoint_name=model.endpoint_name,
             )
             if caught_ex:
-                logger.exception(caught_ex)
-                ignore_if_worker_dies = "Worker died." in str(caught_ex)
-                # https://github.com/pytorch/serve/issues/3032
-                assert (
-                    ignore_if_worker_dies
-                ), f"{caught_ex} was thrown when running pytorch squeezenet sagemaker endpoint test"
+                raise caught_ex
diff --git a/tests/unit/sagemaker/serve/builder/test_model_builder.py b/tests/unit/sagemaker/serve/builder/test_model_builder.py
index 1d199b7401..3ffbdd7c03 100644
--- a/tests/unit/sagemaker/serve/builder/test_model_builder.py
+++ b/tests/unit/sagemaker/serve/builder/test_model_builder.py
@@ -21,6 +21,7 @@
 from sagemaker.serve.mode.function_pointers import Mode
 from sagemaker.serve.utils import task
 from sagemaker.serve.utils.exceptions import TaskNotFoundException
+from sagemaker.serve.utils.predictors import TensorflowServingLocalPredictor
 from sagemaker.serve.utils.types import ModelServer
 from tests.unit.sagemaker.serve.constants import MOCK_IMAGE_CONFIG, MOCK_VPC_CONFIG
 
@@ -52,6 +53,7 @@
     ModelServer.TORCHSERVE,
     ModelServer.TRITON,
     ModelServer.DJL_SERVING,
+    ModelServer.TENSORFLOW_SERVING,
 }
 
 mock_session = MagicMock()
@@ -1677,6 +1679,7 @@ def test_build_task_override_with_invalid_model_provided(
             model_builder.build(sagemaker_session=mock_session)
 
     @patch("os.makedirs", Mock())
+    @patch("sagemaker.serve.builder.model_builder._maintain_lineage_tracking_for_mlflow_model")
     @patch("sagemaker.serve.builder.model_builder._detect_framework_and_version")
     @patch("sagemaker.serve.builder.model_builder.prepare_for_torchserve")
     @patch("sagemaker.serve.builder.model_builder.save_pkl")
@@ -1705,6 +1708,7 @@ def test_build_mlflow_model_local_input_happy(
         mock_save_pkl,
         mock_prepare_for_torchserve,
         mock_detect_fw_version,
+        mock_lineage_tracking,
     ):
         # setup mocks
 
@@ -1750,6 +1754,85 @@ def test_build_mlflow_model_local_input_happy(
         self.assertEqual(build_result.serve_settings, mock_setting_object)
         self.assertEqual(builder.env_vars["MLFLOW_MODEL_FLAVOR"], "sklearn")
 
+        build_result.deploy(
+            initial_instance_count=1, instance_type=mock_instance_type, mode=Mode.SAGEMAKER_ENDPOINT
+        )
+        mock_lineage_tracking.assert_called_once()
+
+    @patch("os.makedirs", Mock())
+    @patch("sagemaker.serve.builder.model_builder._detect_framework_and_version")
+    @patch("sagemaker.serve.builder.model_builder.prepare_for_torchserve")
+    @patch("sagemaker.serve.builder.model_builder.save_pkl")
+    @patch("sagemaker.serve.builder.model_builder._copy_directory_contents")
+    @patch("sagemaker.serve.builder.model_builder._generate_mlflow_artifact_path")
+    @patch("sagemaker.serve.builder.model_builder._get_all_flavor_metadata")
+    @patch("sagemaker.serve.builder.model_builder._select_container_for_mlflow_model")
+    @patch("sagemaker.serve.builder.model_builder._ServeSettings")
+    @patch("sagemaker.serve.builder.model_builder.SageMakerEndpointMode")
+    @patch("sagemaker.serve.builder.model_builder.Model")
+    @patch("builtins.open", new_callable=mock_open, read_data="data")
+    @patch("os.path.isfile", return_value=True)
+    @patch("os.path.exists")
+    def test_build_mlflow_model_local_input_happy_flavor_server_mismatch(
+        self,
+        mock_path_exists,
+        mock_is_file,
+        mock_open,
+        mock_sdk_model,
+        mock_sageMakerEndpointMode,
+        mock_serveSettings,
+        mock_detect_container,
+        mock_get_all_flavor_metadata,
+        mock_generate_mlflow_artifact_path,
+        mock_copy_directory_contents,
+        mock_save_pkl,
+        mock_prepare_for_torchserve,
+        mock_detect_fw_version,
+    ):
+        # setup mocks
+
+        mock_detect_container.return_value = mock_image_uri
+        mock_get_all_flavor_metadata.return_value = {"sklearn": "some_data"}
+        mock_generate_mlflow_artifact_path.return_value = "some_path"
+
+        mock_prepare_for_torchserve.return_value = mock_secret_key
+
+        # Mock _ServeSettings
+        mock_setting_object = mock_serveSettings.return_value
+        mock_setting_object.role_arn = mock_role_arn
+        mock_setting_object.s3_model_data_url = mock_s3_model_data_url
+
+        mock_path_exists.side_effect = lambda path: True if path == "test_path" else False
+
+        mock_mode = Mock()
+        mock_sageMakerEndpointMode.side_effect = lambda inference_spec, model_server: (
+            mock_mode if inference_spec is None and model_server == ModelServer.TORCHSERVE else None
+        )
+        mock_mode.prepare.return_value = (
+            model_data,
+            ENV_VAR_PAIR,
+        )
+
+        updated_env_var = deepcopy(ENV_VARS)
+        updated_env_var.update({"MLFLOW_MODEL_FLAVOR": "sklearn"})
+        mock_model_obj = Mock()
+        mock_sdk_model.return_value = mock_model_obj
+
+        mock_session.sagemaker_client._user_agent_creator.to_string = lambda: "sample agent"
+
+        # run
+        builder = ModelBuilder(
+            schema_builder=schema_builder,
+            model_metadata={"MLFLOW_MODEL_PATH": MODEL_PATH},
+            model_server=ModelServer.TENSORFLOW_SERVING,
+        )
+        with self.assertRaises(ValueError):
+            builder.build(
+                Mode.SAGEMAKER_ENDPOINT,
+                mock_role_arn,
+                mock_session,
+            )
+
     @patch("os.makedirs", Mock())
     @patch("sagemaker.serve.builder.model_builder.S3Downloader.list")
     @patch("sagemaker.serve.builder.model_builder._detect_framework_and_version")
@@ -1899,3 +1982,240 @@ def test_build_mlflow_model_s3_input_non_mlflow_case(
             mock_role_arn,
             mock_session,
         )
+
+    @patch("os.makedirs", Mock())
+    @patch("sagemaker.serve.builder.model_builder._maintain_lineage_tracking_for_mlflow_model")
+    @patch("sagemaker.serve.builder.tf_serving_builder.prepare_for_tf_serving")
+    @patch("sagemaker.serve.builder.model_builder.S3Downloader.list")
+    @patch("sagemaker.serve.builder.model_builder._detect_framework_and_version")
+    @patch("sagemaker.serve.builder.tf_serving_builder.save_pkl")
+    @patch("sagemaker.serve.builder.model_builder._download_s3_artifacts")
+    @patch("sagemaker.serve.builder.model_builder._generate_mlflow_artifact_path")
+    @patch("sagemaker.serve.builder.model_builder._get_all_flavor_metadata")
+    @patch("sagemaker.serve.builder.model_builder._select_container_for_mlflow_model")
+    @patch("sagemaker.serve.builder.model_builder._ServeSettings")
+    @patch("sagemaker.serve.builder.model_builder.SageMakerEndpointMode")
+    @patch("sagemaker.serve.builder.tf_serving_builder.TensorFlowModel")
+    @patch("builtins.open", new_callable=mock_open, read_data="data")
+    @patch("os.path.exists")
+    def test_build_mlflow_model_s3_input_tensorflow_serving_happy(
+        self,
+        mock_path_exists,
+        mock_open,
+        mock_sdk_model,
+        mock_sageMakerEndpointMode,
+        mock_serveSettings,
+        mock_detect_container,
+        mock_get_all_flavor_metadata,
+        mock_generate_mlflow_artifact_path,
+        mock_download_s3_artifacts,
+        mock_save_pkl,
+        mock_detect_fw_version,
+        mock_s3_downloader,
+        mock_prepare_for_tf_serving,
+        mock_lineage_tracking,
+    ):
+        # setup mocks
+        mock_s3_downloader.return_value = ["s3://some_path/MLmodel"]
+
+        mock_detect_container.return_value = mock_image_uri
+        mock_get_all_flavor_metadata.return_value = {"tensorflow": "some_data"}
+        mock_generate_mlflow_artifact_path.return_value = "some_path"
+
+        mock_prepare_for_tf_serving.return_value = mock_secret_key
+
+        # Mock _ServeSettings
+        mock_setting_object = mock_serveSettings.return_value
+        mock_setting_object.role_arn = mock_role_arn
+        mock_setting_object.s3_model_data_url = mock_s3_model_data_url
+
+        mock_path_exists.side_effect = lambda path: True if path == "test_path" else False
+
+        mock_mode = Mock()
+        mock_sageMakerEndpointMode.side_effect = lambda inference_spec, model_server: (
+            mock_mode
+            if inference_spec is None and model_server == ModelServer.TENSORFLOW_SERVING
+            else None
+        )
+        mock_mode.prepare.return_value = (
+            model_data,
+            ENV_VAR_PAIR,
+        )
+
+        updated_env_var = deepcopy(ENV_VARS)
+        updated_env_var.update({"MLFLOW_MODEL_FLAVOR": "tensorflow"})
+        mock_model_obj = Mock()
+        mock_sdk_model.return_value = mock_model_obj
+
+        mock_session.sagemaker_client._user_agent_creator.to_string = lambda: "sample agent"
+
+        # run
+        builder = ModelBuilder(
+            schema_builder=schema_builder, model_metadata={"MLFLOW_MODEL_PATH": "s3://test_path/"}
+        )
+        build_result = builder.build(sagemaker_session=mock_session)
+        self.assertEqual(mock_model_obj, build_result)
+        self.assertEqual(build_result.mode, Mode.SAGEMAKER_ENDPOINT)
+        self.assertEqual(build_result.modes, {str(Mode.SAGEMAKER_ENDPOINT): mock_mode})
+        self.assertEqual(build_result.serve_settings, mock_setting_object)
+        self.assertEqual(builder.env_vars["MLFLOW_MODEL_FLAVOR"], "tensorflow")
+
+        build_result.deploy(
+            initial_instance_count=1, instance_type=mock_instance_type, mode=Mode.SAGEMAKER_ENDPOINT
+        )
+        mock_lineage_tracking.assert_called_once()
+
+    @patch("os.makedirs", Mock())
+    @patch("sagemaker.serve.builder.tf_serving_builder.prepare_for_tf_serving")
+    @patch("sagemaker.serve.builder.model_builder.S3Downloader.list")
+    @patch("sagemaker.serve.builder.model_builder._detect_framework_and_version")
+    @patch("sagemaker.serve.builder.tf_serving_builder.save_pkl")
+    @patch("sagemaker.serve.builder.model_builder._download_s3_artifacts")
+    @patch("sagemaker.serve.builder.model_builder._generate_mlflow_artifact_path")
+    @patch("sagemaker.serve.builder.model_builder._get_all_flavor_metadata")
+    @patch("sagemaker.serve.builder.model_builder._select_container_for_mlflow_model")
+    @patch("sagemaker.serve.builder.model_builder._ServeSettings")
+    @patch("sagemaker.serve.builder.model_builder.LocalContainerMode")
+    @patch("sagemaker.serve.builder.tf_serving_builder.TensorFlowModel")
+    @patch("builtins.open", new_callable=mock_open, read_data="data")
+    @patch("os.path.exists")
+    def test_build_mlflow_model_s3_input_tensorflow_serving_local_mode_happy(
+        self,
+        mock_path_exists,
+        mock_open,
+        mock_sdk_model,
+        mock_local_container_mode,
+        mock_serveSettings,
+        mock_detect_container,
+        mock_get_all_flavor_metadata,
+        mock_generate_mlflow_artifact_path,
+        mock_download_s3_artifacts,
+        mock_save_pkl,
+        mock_detect_fw_version,
+        mock_s3_downloader,
+        mock_prepare_for_tf_serving,
+    ):
+        # setup mocks
+        mock_s3_downloader.return_value = ["s3://some_path/MLmodel"]
+
+        mock_detect_container.return_value = mock_image_uri
+        mock_get_all_flavor_metadata.return_value = {"tensorflow": "some_data"}
+        mock_generate_mlflow_artifact_path.return_value = "some_path"
+
+        mock_prepare_for_tf_serving.return_value = mock_secret_key
+
+        # Mock _ServeSettings
+        mock_setting_object = mock_serveSettings.return_value
+        mock_setting_object.role_arn = mock_role_arn
+        mock_setting_object.s3_model_data_url = mock_s3_model_data_url
+
+        mock_path_exists.side_effect = lambda path: True if path == "test_path" else False
+
+        mock_mode = Mock()
+        mock_mode.prepare.side_effect = lambda: None
+        mock_local_container_mode.return_value = mock_mode
+        mock_mode.prepare.return_value = (
+            model_data,
+            ENV_VAR_PAIR,
+        )
+
+        updated_env_var = deepcopy(ENV_VARS)
+        updated_env_var.update({"MLFLOW_MODEL_FLAVOR": "tensorflow"})
+        mock_model_obj = Mock()
+        mock_sdk_model.return_value = mock_model_obj
+
+        mock_session.sagemaker_client._user_agent_creator.to_string = lambda: "sample agent"
+
+        # run
+        builder = ModelBuilder(
+            mode=Mode.LOCAL_CONTAINER,
+            schema_builder=schema_builder,
+            model_metadata={"MLFLOW_MODEL_PATH": "s3://test_path/"},
+        )
+        build_result = builder.build(sagemaker_session=mock_session)
+        self.assertEqual(mock_model_obj, build_result)
+        self.assertEqual(build_result.mode, Mode.LOCAL_CONTAINER)
+        self.assertEqual(build_result.modes, {str(Mode.LOCAL_CONTAINER): mock_mode})
+        self.assertEqual(build_result.serve_settings, mock_setting_object)
+        self.assertEqual(builder.env_vars["MLFLOW_MODEL_FLAVOR"], "tensorflow")
+
+        predictor = build_result.deploy()
+        assert isinstance(predictor, TensorflowServingLocalPredictor)
+
+    @patch("os.makedirs", Mock())
+    @patch("sagemaker.serve.builder.tf_serving_builder.prepare_for_tf_serving")
+    @patch("sagemaker.serve.builder.model_builder.S3Downloader.list")
+    @patch("sagemaker.serve.builder.model_builder._detect_framework_and_version")
+    @patch("sagemaker.serve.builder.model_builder.save_pkl")
+    @patch("sagemaker.serve.builder.model_builder._download_s3_artifacts")
+    @patch("sagemaker.serve.builder.model_builder._generate_mlflow_artifact_path")
+    @patch("sagemaker.serve.builder.model_builder._get_all_flavor_metadata")
+    @patch("sagemaker.serve.builder.model_builder._select_container_for_mlflow_model")
+    @patch("sagemaker.serve.builder.model_builder._ServeSettings")
+    @patch("sagemaker.serve.builder.model_builder.SageMakerEndpointMode")
+    @patch("sagemaker.serve.builder.tf_serving_builder.TensorFlowModel")
+    @patch("builtins.open", new_callable=mock_open, read_data="data")
+    @patch("os.path.exists")
+    def test_build_tensorflow_serving_non_mlflow_case(
+        self,
+        mock_path_exists,
+        mock_open,
+        mock_sdk_model,
+        mock_sageMakerEndpointMode,
+        mock_serveSettings,
+        mock_detect_container,
+        mock_get_all_flavor_metadata,
+        mock_generate_mlflow_artifact_path,
+        mock_download_s3_artifacts,
+        mock_save_pkl,
+        mock_detect_fw_version,
+        mock_s3_downloader,
+        mock_prepare_for_tf_serving,
+    ):
+        mock_s3_downloader.return_value = []
+        mock_detect_container.return_value = mock_image_uri
+        mock_get_all_flavor_metadata.return_value = {"tensorflow": "some_data"}
+        mock_generate_mlflow_artifact_path.return_value = "some_path"
+
+        mock_prepare_for_tf_serving.return_value = mock_secret_key
+
+        # Mock _ServeSettings
+        mock_setting_object = mock_serveSettings.return_value
+        mock_setting_object.role_arn = mock_role_arn
+        mock_setting_object.s3_model_data_url = mock_s3_model_data_url
+
+        mock_path_exists.side_effect = lambda path: True if path == "test_path" else False
+
+        mock_mode = Mock()
+        mock_sageMakerEndpointMode.side_effect = lambda inference_spec, model_server: (
+            mock_mode
+            if inference_spec is None and model_server == ModelServer.TENSORFLOW_SERVING
+            else None
+        )
+        mock_mode.prepare.return_value = (
+            model_data,
+            ENV_VAR_PAIR,
+        )
+
+        updated_env_var = deepcopy(ENV_VARS)
+        updated_env_var.update({"MLFLOW_MODEL_FLAVOR": "tensorflow"})
+        mock_model_obj = Mock()
+        mock_sdk_model.return_value = mock_model_obj
+
+        mock_session.sagemaker_client._user_agent_creator.to_string = lambda: "sample agent"
+
+        # run
+        builder = ModelBuilder(
+            model=mock_fw_model,
+            schema_builder=schema_builder,
+            model_server=ModelServer.TENSORFLOW_SERVING,
+        )
+
+        self.assertRaisesRegex(
+            Exception,
+            "Tensorflow Serving is currently only supported for mlflow models.",
+            builder.build,
+            Mode.SAGEMAKER_ENDPOINT,
+            mock_role_arn,
+            mock_session,
+        )
diff --git a/tests/unit/sagemaker/serve/builder/test_tensorflow_serving_builder.py b/tests/unit/sagemaker/serve/builder/test_tensorflow_serving_builder.py
new file mode 100644
index 0000000000..9d51b04e08
--- /dev/null
+++ b/tests/unit/sagemaker/serve/builder/test_tensorflow_serving_builder.py
@@ -0,0 +1,75 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You
+# may not use this file except in compliance with the License. A copy of
+# the License is located at
+#
+#     http://aws.amazon.com/apache2.0/
+#
+# or in the "license" file accompanying this file. This file is
+# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
+# ANY KIND, either express or implied. See the License for the specific
+# language governing permissions and limitations under the License.
+from __future__ import absolute_import
+from unittest.mock import MagicMock, patch
+
+import unittest
+from pathlib import Path
+
+from sagemaker.serve import ModelBuilder, ModelServer
+
+
+class TestTransformersBuilder(unittest.TestCase):
+    def setUp(self):
+        self.instance = ModelBuilder()
+        self.instance.model_server = ModelServer.TENSORFLOW_SERVING
+        self.instance.model_path = "/fake/model/path"
+        self.instance.image_uri = "fake_image_uri"
+        self.instance.s3_upload_path = "s3://bucket/path"
+        self.instance.serve_settings = MagicMock(role_arn="fake_role_arn")
+        self.instance.schema_builder = MagicMock()
+        self.instance.env_vars = {}
+        self.instance.sagemaker_session = MagicMock()
+        self.instance.image_config = {}
+        self.instance.vpc_config = {}
+        self.instance.modes = {}
+
+    @patch("os.makedirs")
+    @patch("os.path.exists")
+    @patch("sagemaker.serve.builder.tf_serving_builder.save_pkl")
+    def test_save_schema_builder(self, mock_save_pkl, mock_exists, mock_makedirs):
+        mock_exists.return_value = False
+        self.instance._save_schema_builder()
+        mock_makedirs.assert_called_once_with(self.instance.model_path)
+        code_path = Path(self.instance.model_path).joinpath("code")
+        mock_save_pkl.assert_called_once_with(code_path, self.instance.schema_builder)
+
+    @patch("sagemaker.serve.builder.tf_serving_builder.TensorflowServing._get_client_translators")
+    @patch("sagemaker.serve.builder.tf_serving_builder.TensorFlowPredictor")
+    def test_get_tensorflow_predictor(self, mock_predictor, mock_get_marshaller):
+        endpoint_name = "test_endpoint"
+        predictor = self.instance._get_tensorflow_predictor(
+            endpoint_name, self.instance.sagemaker_session
+        )
+        mock_predictor.assert_called_once_with(
+            endpoint_name=endpoint_name,
+            sagemaker_session=self.instance.sagemaker_session,
+            serializer=self.instance.schema_builder.custom_input_translator,
+            deserializer=self.instance.schema_builder.custom_output_translator,
+        )
+        self.assertEqual(predictor, mock_predictor.return_value)
+
+    @patch("sagemaker.serve.builder.tf_serving_builder.TensorFlowModel")
+    def test_create_tensorflow_model(self, mock_model):
+        model = self.instance._create_tensorflow_model()
+        mock_model.assert_called_once_with(
+            image_uri=self.instance.image_uri,
+            image_config=self.instance.image_config,
+            vpc_config=self.instance.vpc_config,
+            model_data=self.instance.s3_upload_path,
+            role=self.instance.serve_settings.role_arn,
+            env=self.instance.env_vars,
+            sagemaker_session=self.instance.sagemaker_session,
+            predictor_cls=self.instance._get_tensorflow_predictor,
+        )
+        self.assertEqual(model, mock_model.return_value)
diff --git a/tests/unit/sagemaker/serve/model_format/mlflow/test_mlflow_utils.py b/tests/unit/sagemaker/serve/model_format/mlflow/test_mlflow_utils.py
index 154b6b7d95..27c3a5280f 100644
--- a/tests/unit/sagemaker/serve/model_format/mlflow/test_mlflow_utils.py
+++ b/tests/unit/sagemaker/serve/model_format/mlflow/test_mlflow_utils.py
@@ -13,6 +13,7 @@
 from __future__ import absolute_import
 
 import os
+from pathlib import Path
 from unittest.mock import patch, MagicMock, mock_open
 
 import pytest
@@ -21,6 +22,7 @@
 from sagemaker.serve import ModelServer
 from sagemaker.serve.model_format.mlflow.constants import (
     MLFLOW_PYFUNC,
+    TENSORFLOW_SAVED_MODEL_NAME,
 )
 from sagemaker.serve.model_format.mlflow.utils import (
     _get_default_model_server_for_mlflow,
@@ -35,6 +37,8 @@
     _select_container_for_mlflow_model,
     _validate_input_for_mlflow,
     _copy_directory_contents,
+    _move_contents,
+    _get_saved_model_path_for_tensorflow_and_keras_flavor,
 )
 
 
@@ -415,10 +419,15 @@ def test_select_container_for_mlflow_model_no_dlc_detected(
 
 
 def test_validate_input_for_mlflow():
-    _validate_input_for_mlflow(ModelServer.TORCHSERVE)
+    _validate_input_for_mlflow(ModelServer.TORCHSERVE, "pytorch")
 
     with pytest.raises(ValueError):
-        _validate_input_for_mlflow(ModelServer.DJL_SERVING)
+        _validate_input_for_mlflow(ModelServer.DJL_SERVING, "pytorch")
+
+
+def test_validate_input_for_mlflow_non_supported_flavor_with_tf_serving():
+    with pytest.raises(ValueError):
+        _validate_input_for_mlflow(ModelServer.TENSORFLOW_SERVING, "pytorch")
 
 
 @patch("sagemaker.serve.model_format.mlflow.utils.shutil.copy2")
@@ -472,3 +481,68 @@ def test_copy_directory_contents_handles_same_src_dst(
     mock_os_walk.assert_not_called()
     mock_os_makedirs.assert_not_called()
     mock_shutil_copy2.assert_not_called()
+
+
+@patch("os.path.abspath")
+@patch("os.walk")
+def test_get_saved_model_path_found(mock_os_walk, mock_os_abspath):
+    mock_os_walk.return_value = [
+        ("/root/folder1", ("subfolder",), ()),
+        ("/root/folder1/subfolder", (), (TENSORFLOW_SAVED_MODEL_NAME,)),
+    ]
+    expected_path = "/root/folder1/subfolder"
+    mock_os_abspath.return_value = expected_path
+
+    # Call the function
+    result = _get_saved_model_path_for_tensorflow_and_keras_flavor("/root/folder1")
+
+    # Assertions
+    mock_os_walk.assert_called_once_with("/root/folder1")
+    mock_os_abspath.assert_called_once_with("/root/folder1/subfolder")
+    assert result == expected_path
+
+
+@patch("os.path.abspath")
+@patch("os.walk")
+def test_get_saved_model_path_not_found(mock_os_walk, mock_os_abspath):
+    mock_os_walk.return_value = [
+        ("/root/folder2", ("subfolder",), ()),
+        ("/root/folder2/subfolder", (), ("not_saved_model.pb",)),
+    ]
+
+    result = _get_saved_model_path_for_tensorflow_and_keras_flavor("/root/folder2")
+
+    mock_os_walk.assert_called_once_with("/root/folder2")
+    mock_os_abspath.assert_not_called()
+    assert result is None
+
+
+@patch("sagemaker.serve.model_format.mlflow.utils.shutil.move")
+@patch("sagemaker.serve.model_format.mlflow.utils.Path.iterdir")
+@patch("sagemaker.serve.model_format.mlflow.utils.Path.mkdir")
+def test_move_contents_handles_same_src_dst(mock_mkdir, mock_iterdir, mock_shutil_move):
+    src_dir = "/fake/source/dir"
+    dest_dir = "/fake/source/./dir"
+
+    mock_iterdir.return_value = []
+
+    _move_contents(src_dir, dest_dir)
+
+    mock_mkdir.assert_called_once_with(parents=True, exist_ok=True)
+    mock_shutil_move.assert_not_called()
+
+
+@patch("sagemaker.serve.model_format.mlflow.utils.shutil.move")
+@patch("sagemaker.serve.model_format.mlflow.utils.Path.iterdir")
+@patch("sagemaker.serve.model_format.mlflow.utils.Path.mkdir")
+def test_move_contents_with_actual_files(mock_mkdir, mock_iterdir, mock_shutil_move):
+    src_dir = Path("/fake/source/dir")
+    dest_dir = Path("/fake/destination/dir")
+
+    file_path = src_dir / "testfile.txt"
+    mock_iterdir.return_value = [file_path]
+
+    _move_contents(src_dir, dest_dir)
+
+    mock_mkdir.assert_called_once_with(parents=True, exist_ok=True)
+    mock_shutil_move.assert_called_once_with(str(file_path), str(dest_dir / "testfile.txt"))
diff --git a/tests/unit/sagemaker/serve/model_server/tensorflow_serving/test_tf_prepare.py b/tests/unit/sagemaker/serve/model_server/tensorflow_serving/test_tf_prepare.py
new file mode 100644
index 0000000000..9915b19649
--- /dev/null
+++ b/tests/unit/sagemaker/serve/model_server/tensorflow_serving/test_tf_prepare.py
@@ -0,0 +1,116 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You
+# may not use this file except in compliance with the License. A copy of
+# the License is located at
+#
+#     http://aws.amazon.com/apache2.0/
+#
+# or in the "license" file accompanying this file. This file is
+# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
+# ANY KIND, either express or implied. See the License for the specific
+# language governing permissions and limitations under the License.
+from __future__ import absolute_import
+
+from unittest import TestCase
+from unittest.mock import Mock, patch, mock_open
+import pytest
+
+from sagemaker.serve.model_server.tensorflow_serving.prepare import prepare_for_tf_serving
+
+MODEL_PATH = "/path/to/your/model/dir"
+SHARED_LIBS = ["/path/to/shared/libs"]
+DEPENDENCIES = {"dependencies": "requirements.txt"}
+INFERENCE_SPEC = Mock()
+IMAGE_URI = "mock_image_uri"
+XGB_1P_IMAGE_URI = "246618743249.dkr.ecr.us-west-2.amazonaws.com/sagemaker-tensorflow:1.7-1"
+INFERENCE_SPEC.prepare = Mock(return_value=None)
+
+SECRET_KEY = "secret-key"
+
+mock_session = Mock()
+
+
+class PrepareForTensorflowServingTests(TestCase):
+    def setUp(self):
+        INFERENCE_SPEC.reset_mock()
+
+    @patch("builtins.open", new_callable=mock_open, read_data=b"{}")
+    @patch("sagemaker.serve.model_server.tensorflow_serving.prepare._move_contents")
+    @patch(
+        "sagemaker.serve.model_server.tensorflow_serving.prepare."
+        "_get_saved_model_path_for_tensorflow_and_keras_flavor"
+    )
+    @patch("sagemaker.serve.model_server.tensorflow_serving.prepare._MetaData")
+    @patch("sagemaker.serve.model_server.tensorflow_serving.prepare.compute_hash")
+    @patch("sagemaker.serve.model_server.tensorflow_serving.prepare.generate_secret_key")
+    @patch("sagemaker.serve.model_server.tensorflow_serving.prepare.capture_dependencies")
+    @patch("sagemaker.serve.model_server.tensorflow_serving.prepare.shutil")
+    @patch("sagemaker.serve.model_server.tensorflow_serving.prepare.Path")
+    def test_prepare_happy(
+        self,
+        mock_path,
+        mock_shutil,
+        mock_capture_dependencies,
+        mock_generate_secret_key,
+        mock_compute_hash,
+        mock_metadata,
+        mock_get_saved_model_path,
+        mock_move_contents,
+        mock_open,
+    ):
+
+        mock_path_instance = mock_path.return_value
+        mock_path_instance.exists.return_value = True
+        mock_path_instance.joinpath.return_value = Mock()
+        mock_get_saved_model_path.return_value = MODEL_PATH + "/1/"
+
+        mock_generate_secret_key.return_value = SECRET_KEY
+
+        secret_key = prepare_for_tf_serving(
+            model_path=MODEL_PATH,
+            shared_libs=SHARED_LIBS,
+            dependencies=DEPENDENCIES,
+        )
+
+        mock_path_instance.mkdir.assert_not_called()
+        self.assertEqual(secret_key, SECRET_KEY)
+
+    @patch("builtins.open", new_callable=mock_open, read_data=b"{}")
+    @patch("sagemaker.serve.model_server.tensorflow_serving.prepare._move_contents")
+    @patch(
+        "sagemaker.serve.model_server.tensorflow_serving.prepare."
+        "_get_saved_model_path_for_tensorflow_and_keras_flavor"
+    )
+    @patch("sagemaker.serve.model_server.tensorflow_serving.prepare._MetaData")
+    @patch("sagemaker.serve.model_server.tensorflow_serving.prepare.compute_hash")
+    @patch("sagemaker.serve.model_server.tensorflow_serving.prepare.generate_secret_key")
+    @patch("sagemaker.serve.model_server.tensorflow_serving.prepare.capture_dependencies")
+    @patch("sagemaker.serve.model_server.tensorflow_serving.prepare.shutil")
+    @patch("sagemaker.serve.model_server.tensorflow_serving.prepare.Path")
+    def test_prepare_saved_model_not_found(
+        self,
+        mock_path,
+        mock_shutil,
+        mock_capture_dependencies,
+        mock_generate_secret_key,
+        mock_compute_hash,
+        mock_metadata,
+        mock_get_saved_model_path,
+        mock_move_contents,
+        mock_open,
+    ):
+
+        mock_path_instance = mock_path.return_value
+        mock_path_instance.exists.return_value = True
+        mock_path_instance.joinpath.return_value = Mock()
+        mock_get_saved_model_path.return_value = None
+
+        with pytest.raises(
+            ValueError, match="SavedModel is not found for Tensorflow or Keras flavor."
+        ):
+            prepare_for_tf_serving(
+                model_path=MODEL_PATH,
+                shared_libs=SHARED_LIBS,
+                dependencies=DEPENDENCIES,
+            )
diff --git a/tests/unit/sagemaker/serve/model_server/tensorflow_serving/test_tf_server.py b/tests/unit/sagemaker/serve/model_server/tensorflow_serving/test_tf_server.py
new file mode 100644
index 0000000000..3d3bac0935
--- /dev/null
+++ b/tests/unit/sagemaker/serve/model_server/tensorflow_serving/test_tf_server.py
@@ -0,0 +1,100 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You
+# may not use this file except in compliance with the License. A copy of
+# the License is located at
+#
+#     http://aws.amazon.com/apache2.0/
+#
+# or in the "license" file accompanying this file. This file is
+# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
+# ANY KIND, either express or implied. See the License for the specific
+# language governing permissions and limitations under the License.
+from __future__ import absolute_import
+
+from pathlib import PosixPath
+import platform
+from unittest import TestCase
+from unittest.mock import Mock, patch, ANY
+
+import numpy as np
+
+from sagemaker.serve.model_server.tensorflow_serving.server import (
+    LocalTensorflowServing,
+    SageMakerTensorflowServing,
+)
+
+CPU_TF_IMAGE = "763104351884.dkr.ecr.us-west-2.amazonaws.com/tensorflow-inference:2.14.1-cpu"
+MODEL_PATH = "model_path"
+MODEL_REPO = f"{MODEL_PATH}/1"
+ENV_VAR = {"KEY": "VALUE"}
+_SHM_SIZE = "2G"
+PAYLOAD = np.random.rand(3, 4).astype(dtype=np.float32)
+S3_URI = "s3://mock_model_data_uri"
+DTYPE = "TYPE_FP32"
+SECRET_KEY = "secret_key"
+
+INFER_RESPONSE = {"outputs": [{"name": "output_name"}]}
+
+
+class TensorflowservingServerTests(TestCase):
+    def test_start_invoke_destroy_local_tensorflow_serving_server(self):
+        mock_container = Mock()
+        mock_docker_client = Mock()
+        mock_docker_client.containers.run.return_value = mock_container
+
+        local_tensorflow_server = LocalTensorflowServing()
+        mock_schema_builder = Mock()
+        mock_schema_builder.input_serializer.serialize.return_value = PAYLOAD
+        local_tensorflow_server.schema_builder = mock_schema_builder
+
+        local_tensorflow_server._start_tensorflow_serving(
+            client=mock_docker_client,
+            model_path=MODEL_PATH,
+            secret_key=SECRET_KEY,
+            env_vars=ENV_VAR,
+            image=CPU_TF_IMAGE,
+        )
+
+        mock_docker_client.containers.run.assert_called_once_with(
+            CPU_TF_IMAGE,
+            "serve",
+            detach=True,
+            auto_remove=True,
+            network_mode="host",
+            volumes={PosixPath("model_path"): {"bind": "/opt/ml/model", "mode": "rw"}},
+            environment={
+                "SAGEMAKER_SUBMIT_DIRECTORY": "/opt/ml/model/code",
+                "SAGEMAKER_PROGRAM": "inference.py",
+                "SAGEMAKER_SERVE_SECRET_KEY": "secret_key",
+                "LOCAL_PYTHON": platform.python_version(),
+                "KEY": "VALUE",
+            },
+        )
+
+    @patch("sagemaker.serve.model_server.tensorflow_serving.server.platform")
+    @patch("sagemaker.serve.model_server.tensorflow_serving.server.upload")
+    def test_upload_artifacts_sagemaker_triton_server(self, mock_upload, mock_platform):
+        mock_session = Mock()
+        mock_platform.python_version.return_value = "3.8"
+        mock_upload.side_effect = lambda session, repo, bucket, prefix: (
+            S3_URI
+            if session == mock_session and repo == MODEL_PATH and bucket == "mock_model_data_uri"
+            else None
+        )
+
+        (
+            s3_upload_path,
+            env_vars,
+        ) = SageMakerTensorflowServing()._upload_tensorflow_serving_artifacts(
+            model_path=MODEL_PATH,
+            sagemaker_session=mock_session,
+            secret_key=SECRET_KEY,
+            s3_model_data_url=S3_URI,
+            image=CPU_TF_IMAGE,
+        )
+
+        mock_upload.assert_called_once_with(mock_session, MODEL_PATH, "mock_model_data_uri", ANY)
+        self.assertEqual(s3_upload_path, S3_URI)
+        self.assertEqual(env_vars.get("SAGEMAKER_SERVE_SECRET_KEY"), SECRET_KEY)
+        self.assertEqual(env_vars.get("LOCAL_PYTHON"), "3.8")
diff --git a/tests/unit/sagemaker/serve/utils/test_lineage_utils.py b/tests/unit/sagemaker/serve/utils/test_lineage_utils.py
new file mode 100644
index 0000000000..25e4fe246e
--- /dev/null
+++ b/tests/unit/sagemaker/serve/utils/test_lineage_utils.py
@@ -0,0 +1,374 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You
+# may not use this file except in compliance with the License. A copy of
+# the License is located at
+#
+#     http://aws.amazon.com/apache2.0/
+#
+# or in the "license" file accompanying this file. This file is
+# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
+# ANY KIND, either express or implied. See the License for the specific
+# language governing permissions and limitations under the License.
+from __future__ import absolute_import
+
+from unittest.mock import call
+
+import pytest
+from botocore.exceptions import ClientError
+from mock import Mock, patch
+from sagemaker import Session
+from sagemaker.lineage.artifact import ArtifactSummary, Artifact
+from sagemaker.lineage.query import LineageSourceEnum
+
+from sagemaker.serve.utils.lineage_constants import (
+    MLFLOW_RUN_ID,
+    MLFLOW_MODEL_PACKAGE_PATH,
+    MLFLOW_S3_PATH,
+    MLFLOW_LOCAL_PATH,
+    LINEAGE_POLLER_MAX_TIMEOUT_SECS,
+    MODEL_BUILDER_MLFLOW_MODEL_PATH_LINEAGE_ARTIFACT_TYPE,
+    CONTRIBUTED_TO,
+    MLFLOW_REGISTRY_PATH,
+)
+from sagemaker.serve.utils.lineage_utils import (
+    _load_artifact_by_source_uri,
+    _poll_lineage_artifact,
+    _get_mlflow_model_path_type,
+    _create_mlflow_model_path_lineage_artifact,
+    _add_association_between_artifacts,
+    _retrieve_and_create_if_not_exist_mlflow_model_path_lineage_artifact,
+    _maintain_lineage_tracking_for_mlflow_model,
+)
+
+
+@patch("sagemaker.lineage.artifact.Artifact.list")
+def test_load_artifact_by_source_uri(mock_artifact_list):
+    source_uri = "s3://mybucket/mymodel"
+    sagemaker_session = Mock(spec=Session)
+
+    mock_artifact_1 = Mock(spec=ArtifactSummary)
+    mock_artifact_1.artifact_type = LineageSourceEnum.MODEL_DATA.value
+    mock_artifact_2 = Mock(spec=ArtifactSummary)
+    mock_artifact_2.artifact_type = LineageSourceEnum.IMAGE.value
+    mock_artifacts = [mock_artifact_1, mock_artifact_2]
+    mock_artifact_list.return_value = mock_artifacts
+
+    result = _load_artifact_by_source_uri(
+        source_uri, LineageSourceEnum.MODEL_DATA.value, sagemaker_session
+    )
+
+    mock_artifact_list.assert_called_once_with(
+        source_uri=source_uri, sagemaker_session=sagemaker_session
+    )
+    assert result == mock_artifact_1
+
+
+@patch("sagemaker.lineage.artifact.Artifact.list")
+def test_load_artifact_by_source_uri_no_match(mock_artifact_list):
+    source_uri = "s3://mybucket/mymodel"
+    sagemaker_session = Mock(spec=Session)
+
+    mock_artifact_1 = Mock(spec=ArtifactSummary)
+    mock_artifact_1.artifact_type = LineageSourceEnum.IMAGE.value
+    mock_artifact_2 = Mock(spec=ArtifactSummary)
+    mock_artifact_2.artifact_type = LineageSourceEnum.IMAGE.value
+    mock_artifacts = [mock_artifact_1, mock_artifact_2]
+    mock_artifact_list.return_value = mock_artifacts
+
+    result = _load_artifact_by_source_uri(
+        source_uri, LineageSourceEnum.MODEL_DATA.value, sagemaker_session
+    )
+
+    mock_artifact_list.assert_called_once_with(
+        source_uri=source_uri, sagemaker_session=sagemaker_session
+    )
+    assert result is None
+
+
+@patch("sagemaker.serve.utils.lineage_utils._load_artifact_by_source_uri")
+def test_poll_lineage_artifact_found(mock_load_artifact):
+    s3_uri = "s3://mybucket/mymodel"
+    sagemaker_session = Mock(spec=Session)
+    mock_artifact = Mock(spec=ArtifactSummary)
+
+    with patch("time.time") as mock_time:
+        mock_time.return_value = 0
+
+        mock_load_artifact.return_value = mock_artifact
+
+        result = _poll_lineage_artifact(
+            s3_uri, LineageSourceEnum.MODEL_DATA.value, sagemaker_session
+        )
+
+    assert result == mock_artifact
+    mock_load_artifact.assert_has_calls(
+        [
+            call(s3_uri, LineageSourceEnum.MODEL_DATA.value, sagemaker_session),
+        ]
+    )
+
+
+@patch("sagemaker.serve.utils.lineage_utils._load_artifact_by_source_uri")
+def test_poll_lineage_artifact_not_found(mock_load_artifact):
+    s3_uri = "s3://mybucket/mymodel"
+    artifact_type = LineageSourceEnum.MODEL_DATA.value
+    sagemaker_session = Mock(spec=Session)
+
+    with patch("time.time") as mock_time:
+        mock_time_values = [0.0, 1.0, LINEAGE_POLLER_MAX_TIMEOUT_SECS + 1.0]
+        mock_time.side_effect = mock_time_values
+
+        with patch("time.sleep"):
+            mock_load_artifact.side_effect = [None, None, None]
+
+            result = _poll_lineage_artifact(s3_uri, artifact_type, sagemaker_session)
+
+    assert result is None
+
+
+@pytest.mark.parametrize(
+    "mlflow_model_path, expected_output",
+    [
+        ("runs:/abc123", MLFLOW_RUN_ID),
+        ("models:/my-model/1", MLFLOW_REGISTRY_PATH),
+        (
+            "arn:aws:sagemaker:us-west-2:123456789012:model-package/my-model-package",
+            MLFLOW_MODEL_PACKAGE_PATH,
+        ),
+        ("s3://my-bucket/path/to/model", MLFLOW_S3_PATH),
+    ],
+)
+def test_get_mlflow_model_path_type_valid(mlflow_model_path, expected_output):
+    result = _get_mlflow_model_path_type(mlflow_model_path)
+    assert result == expected_output
+
+
+@patch("os.path.exists")
+def test_get_mlflow_model_path_type_valid_local_path(mock_path_exists):
+    valid_path = "/path/to/mlflow_model"
+    mock_path_exists.side_effect = lambda path: path == valid_path
+    result = _get_mlflow_model_path_type(valid_path)
+    assert result == MLFLOW_LOCAL_PATH
+
+
+def test_get_mlflow_model_path_type_invalid():
+    invalid_path = "invalid_path"
+    with pytest.raises(ValueError, match=f"Invalid MLflow model path: {invalid_path}"):
+        _get_mlflow_model_path_type(invalid_path)
+
+
+@patch("sagemaker.serve.utils.lineage_utils._get_mlflow_model_path_type")
+@patch("sagemaker.lineage.artifact.Artifact.create")
+def test_create_mlflow_model_path_lineage_artifact_success(
+    mock_artifact_create, mock_get_mlflow_path_type
+):
+    mlflow_model_path = "runs:/Ab12Cd34"
+    sagemaker_session = Mock(spec=Session)
+    mock_artifact = Mock(spec=Artifact)
+    mock_get_mlflow_path_type.return_value = "mlflow_run_id"
+    mock_artifact_create.return_value = mock_artifact
+
+    result = _create_mlflow_model_path_lineage_artifact(mlflow_model_path, sagemaker_session)
+
+    assert result == mock_artifact
+    mock_get_mlflow_path_type.assert_called_once_with(mlflow_model_path)
+    mock_artifact_create.assert_called_once_with(
+        source_uri=mlflow_model_path,
+        artifact_type=MODEL_BUILDER_MLFLOW_MODEL_PATH_LINEAGE_ARTIFACT_TYPE,
+        artifact_name="mlflow_run_id",
+        properties={"model_builder_input_model_data_type": "mlflow_run_id"},
+        sagemaker_session=sagemaker_session,
+    )
+
+
+@patch("sagemaker.serve.utils.lineage_utils._get_mlflow_model_path_type")
+@patch("sagemaker.lineage.artifact.Artifact.create")
+def test_create_mlflow_model_path_lineage_artifact_validation_exception(
+    mock_artifact_create, mock_get_mlflow_path_type
+):
+    mlflow_model_path = "runs:/Ab12Cd34"
+    sagemaker_session = Mock(spec=Session)
+    mock_get_mlflow_path_type.return_value = "mlflow_run_id"
+    mock_artifact_create.side_effect = ClientError(
+        error_response={"Error": {"Code": "ValidationException"}}, operation_name="CreateArtifact"
+    )
+
+    result = _create_mlflow_model_path_lineage_artifact(mlflow_model_path, sagemaker_session)
+
+    assert result is None
+
+
+@patch("sagemaker.serve.utils.lineage_utils._get_mlflow_model_path_type")
+@patch("sagemaker.lineage.artifact.Artifact.create")
+def test_create_mlflow_model_path_lineage_artifact_other_exception(
+    mock_artifact_create, mock_get_mlflow_path_type
+):
+    mlflow_model_path = "runs:/Ab12Cd34"
+    sagemaker_session = Mock(spec=Session)
+    mock_get_mlflow_path_type.return_value = "mlflow_run_id"
+    mock_artifact_create.side_effect = ClientError(
+        error_response={"Error": {"Code": "SomeOtherException"}}, operation_name="CreateArtifact"
+    )
+
+    with pytest.raises(ClientError):
+        _create_mlflow_model_path_lineage_artifact(mlflow_model_path, sagemaker_session)
+
+
+@patch("sagemaker.serve.utils.lineage_utils._create_mlflow_model_path_lineage_artifact")
+@patch("sagemaker.serve.utils.lineage_utils._load_artifact_by_source_uri")
+def test_retrieve_and_create_if_not_exist_mlflow_model_path_lineage_artifact_existing(
+    mock_load_artifact, mock_create_artifact
+):
+    mlflow_model_path = "runs:/Ab12Cd34"
+    sagemaker_session = Mock(spec=Session)
+    mock_artifact_summary = Mock(spec=ArtifactSummary)
+    mock_load_artifact.return_value = mock_artifact_summary
+
+    result = _retrieve_and_create_if_not_exist_mlflow_model_path_lineage_artifact(
+        mlflow_model_path, sagemaker_session
+    )
+
+    assert result == mock_artifact_summary
+    mock_load_artifact.assert_called_once_with(
+        mlflow_model_path, MODEL_BUILDER_MLFLOW_MODEL_PATH_LINEAGE_ARTIFACT_TYPE, sagemaker_session
+    )
+    mock_create_artifact.assert_not_called()
+
+
+@patch("sagemaker.serve.utils.lineage_utils._create_mlflow_model_path_lineage_artifact")
+@patch("sagemaker.serve.utils.lineage_utils._load_artifact_by_source_uri")
+def test_retrieve_and_create_if_not_exist_mlflow_model_path_lineage_artifact_create(
+    mock_load_artifact, mock_create_artifact
+):
+    mlflow_model_path = "runs:/Ab12Cd34"
+    sagemaker_session = Mock(spec=Session)
+    mock_artifact = Mock(spec=Artifact)
+    mock_load_artifact.return_value = None
+    mock_create_artifact.return_value = mock_artifact
+
+    result = _retrieve_and_create_if_not_exist_mlflow_model_path_lineage_artifact(
+        mlflow_model_path, sagemaker_session
+    )
+
+    assert result == mock_artifact
+    mock_load_artifact.assert_called_once_with(
+        mlflow_model_path, MODEL_BUILDER_MLFLOW_MODEL_PATH_LINEAGE_ARTIFACT_TYPE, sagemaker_session
+    )
+    mock_create_artifact.assert_called_once_with(mlflow_model_path, sagemaker_session)
+
+
+@patch("sagemaker.lineage.association.Association.create")
+def test_add_association_between_artifacts_success(mock_association_create):
+    mlflow_model_path_artifact_arn = "arn:aws:sagemaker:us-west-2:123456789012:artifact/123"
+    autogenerated_model_data_artifact_arn = "arn:aws:sagemaker:us-west-2:123456789012:artifact/456"
+    sagemaker_session = Mock(spec=Session)
+
+    _add_association_between_artifacts(
+        mlflow_model_path_artifact_arn,
+        autogenerated_model_data_artifact_arn,
+        sagemaker_session,
+    )
+
+    mock_association_create.assert_called_once_with(
+        source_arn=mlflow_model_path_artifact_arn,
+        destination_arn=autogenerated_model_data_artifact_arn,
+        association_type=CONTRIBUTED_TO,
+        sagemaker_session=sagemaker_session,
+    )
+
+
+@patch("sagemaker.lineage.association.Association.create")
+def test_add_association_between_artifacts_validation_exception(mock_association_create):
+    mlflow_model_path_artifact_arn = "arn:aws:sagemaker:us-west-2:123456789012:artifact/123"
+    autogenerated_model_data_artifact_arn = "arn:aws:sagemaker:us-west-2:123456789012:artifact/456"
+    sagemaker_session = Mock(spec=Session)
+    mock_association_create.side_effect = ClientError(
+        error_response={"Error": {"Code": "ValidationException"}},
+        operation_name="CreateAssociation",
+    )
+
+    _add_association_between_artifacts(
+        mlflow_model_path_artifact_arn,
+        autogenerated_model_data_artifact_arn,
+        sagemaker_session,
+    )
+
+
+@patch("sagemaker.lineage.association.Association.create")
+def test_add_association_between_artifacts_other_exception(mock_association_create):
+    mlflow_model_path_artifact_arn = "arn:aws:sagemaker:us-west-2:123456789012:artifact/123"
+    autogenerated_model_data_artifact_arn = "arn:aws:sagemaker:us-west-2:123456789012:artifact/456"
+    sagemaker_session = Mock(spec=Session)
+    mock_association_create.side_effect = ClientError(
+        error_response={"Error": {"Code": "SomeOtherException"}}, operation_name="CreateAssociation"
+    )
+
+    with pytest.raises(ClientError):
+        _add_association_between_artifacts(
+            mlflow_model_path_artifact_arn,
+            autogenerated_model_data_artifact_arn,
+            sagemaker_session,
+        )
+
+
+@patch("sagemaker.serve.utils.lineage_utils._poll_lineage_artifact")
+@patch(
+    "sagemaker.serve.utils.lineage_utils._retrieve_and_create_if_not_exist_mlflow_model_path_lineage_artifact"
+)
+@patch("sagemaker.serve.utils.lineage_utils._add_association_between_artifacts")
+def test_maintain_lineage_tracking_for_mlflow_model_success(
+    mock_add_association, mock_retrieve_create_artifact, mock_poll_artifact
+):
+    mlflow_model_path = "runs:/Ab12Cd34"
+    s3_upload_path = "s3://mybucket/path/to/model"
+    sagemaker_session = Mock(spec=Session)
+    mock_model_data_artifact = Mock(spec=ArtifactSummary)
+    mock_mlflow_model_artifact = Mock(spec=Artifact)
+    mock_poll_artifact.return_value = mock_model_data_artifact
+    mock_retrieve_create_artifact.return_value = mock_mlflow_model_artifact
+
+    _maintain_lineage_tracking_for_mlflow_model(
+        mlflow_model_path, s3_upload_path, sagemaker_session
+    )
+
+    mock_poll_artifact.assert_called_once_with(
+        s3_uri=s3_upload_path,
+        artifact_type=LineageSourceEnum.MODEL_DATA.value,
+        sagemaker_session=sagemaker_session,
+    )
+    mock_retrieve_create_artifact.assert_called_once_with(
+        mlflow_model_path=mlflow_model_path, sagemaker_session=sagemaker_session
+    )
+    mock_add_association.assert_called_once_with(
+        mlflow_model_path_artifact_arn=mock_mlflow_model_artifact.artifact_arn,
+        autogenerated_model_data_artifact_arn=mock_model_data_artifact.artifact_arn,
+        sagemaker_session=sagemaker_session,
+    )
+
+
+@patch("sagemaker.serve.utils.lineage_utils._poll_lineage_artifact")
+@patch(
+    "sagemaker.serve.utils.lineage_utils._retrieve_and_create_if_not_exist_mlflow_model_path_lineage_artifact"
+)
+@patch("sagemaker.serve.utils.lineage_utils._add_association_between_artifacts")
+def test_maintain_lineage_tracking_for_mlflow_model_no_model_data_artifact(
+    mock_add_association, mock_retrieve_create_artifact, mock_poll_artifact
+):
+    mlflow_model_path = "runs:/Ab12Cd34"
+    s3_upload_path = "s3://mybucket/path/to/model"
+    sagemaker_session = Mock(spec=Session)
+    mock_poll_artifact.return_value = None
+    mock_retrieve_create_artifact.return_value = None
+
+    _maintain_lineage_tracking_for_mlflow_model(
+        mlflow_model_path, s3_upload_path, sagemaker_session
+    )
+
+    mock_poll_artifact.assert_called_once_with(
+        s3_uri=s3_upload_path,
+        artifact_type=LineageSourceEnum.MODEL_DATA.value,
+        sagemaker_session=sagemaker_session,
+    )
+    mock_retrieve_create_artifact.assert_not_called()
+    mock_add_association.assert_not_called()

From 1279620f7ab7f7e695eb000c39ea07965b890b16 Mon Sep 17 00:00:00 2001
From: Jonathan Makunga <54963715+makungaj1@users.noreply.github.com>
Date: Fri, 10 May 2024 14:36:25 -0700
Subject: [PATCH 26/58] fix: model builder race condition on sagemaker session
 (#4673)

Co-authored-by: Jonathan Makunga <makung@amazon.com>
---
 src/sagemaker/serve/builder/model_builder.py | 2 --
 1 file changed, 2 deletions(-)

diff --git a/src/sagemaker/serve/builder/model_builder.py b/src/sagemaker/serve/builder/model_builder.py
index 52268ea40c..c43ab53e85 100644
--- a/src/sagemaker/serve/builder/model_builder.py
+++ b/src/sagemaker/serve/builder/model_builder.py
@@ -724,8 +724,6 @@ def build(  # pylint: disable=R0911
             self.role_arn = role_arn
         self.sagemaker_session = sagemaker_session or Session()
 
-        self.sagemaker_session.settings._local_download_dir = self.model_path
-
         # https://github.com/boto/botocore/blob/develop/botocore/useragent.py#L258
         # decorate to_string() due to
         # https://github.com/boto/botocore/blob/develop/botocore/client.py#L1014-L1015

From 2f2be057acb26a15f2f85c6a0b927a0102c7fb09 Mon Sep 17 00:00:00 2001
From: jiapinw <95885824+jiapinw@users.noreply.github.com>
Date: Mon, 13 May 2024 14:20:22 -0700
Subject: [PATCH 27/58] feat: Add telemetry support for mlflow models (#4674)

* Initial commit for telemetry support

* Fix style issues and add more logger messages

* fix value error messages in ut
---
 src/sagemaker/serve/builder/model_builder.py  |  3 +
 .../serve/model_format/mlflow/constants.py    |  6 +-
 .../serve/model_format/mlflow/utils.py        | 11 +++-
 src/sagemaker/serve/utils/lineage_utils.py    |  4 +-
 src/sagemaker/serve/utils/telemetry_logger.py | 22 ++++++++
 .../model_format/mlflow/test_mlflow_utils.py  | 55 +++++++++++++++++++
 .../serve/utils/test_telemetry_logger.py      | 36 ++++++++++++
 7 files changed, 131 insertions(+), 6 deletions(-)

diff --git a/src/sagemaker/serve/builder/model_builder.py b/src/sagemaker/serve/builder/model_builder.py
index c43ab53e85..42a2b994a8 100644
--- a/src/sagemaker/serve/builder/model_builder.py
+++ b/src/sagemaker/serve/builder/model_builder.py
@@ -669,6 +669,9 @@ def _initialize_for_mlflow(self) -> None:
         mlflow_path = self.model_metadata.get(MLFLOW_MODEL_PATH)
         if not _mlflow_input_is_local_path(mlflow_path):
             # TODO: extend to package arn, run id and etc.
+            logger.info(
+                "Start downloading model artifacts from %s to %s", mlflow_path, self.model_path
+            )
             _download_s3_artifacts(mlflow_path, self.model_path, self.sagemaker_session)
         else:
             _copy_directory_contents(mlflow_path, self.model_path)
diff --git a/src/sagemaker/serve/model_format/mlflow/constants.py b/src/sagemaker/serve/model_format/mlflow/constants.py
index 5d2ac484b5..28a3cbdc8d 100644
--- a/src/sagemaker/serve/model_format/mlflow/constants.py
+++ b/src/sagemaker/serve/model_format/mlflow/constants.py
@@ -19,12 +19,12 @@
     "py39": "1.13.1",
     "py310": "2.2.0",
 }
-MODEL_PACAKGE_ARN_REGEX = (
-    r"^arn:aws:sagemaker:[a-z0-9\-]+:[0-9]{12}:model-package\/[" r"a-zA-Z0-9\-_\/\.]+$"
+MODEL_PACKAGE_ARN_REGEX = (
+    r"^arn:aws:sagemaker:[a-z0-9\-]+:[0-9]{12}:model-package\/(.*?)(?:/(\d+))?$"
 )
 MLFLOW_RUN_ID_REGEX = r"^runs:/[a-zA-Z0-9]+(/[a-zA-Z0-9]+)*$"
 MLFLOW_REGISTRY_PATH_REGEX = r"^models:/[a-zA-Z0-9\-_\.]+(/[0-9]+)*$"
-S3_PATH_REGEX = r"^s3:\/\/[a-zA-Z0-9\-_\.]+\/[a-zA-Z0-9\-_\/\.]*$"
+S3_PATH_REGEX = r"^s3:\/\/[a-zA-Z0-9\-_\.]+(?:\/[a-zA-Z0-9\-_\/\.]*)?$"
 MLFLOW_MODEL_PATH = "MLFLOW_MODEL_PATH"
 MLFLOW_METADATA_FILE = "MLmodel"
 MLFLOW_PIP_DEPENDENCY_FILE = "requirements.txt"
diff --git a/src/sagemaker/serve/model_format/mlflow/utils.py b/src/sagemaker/serve/model_format/mlflow/utils.py
index b67de08d78..c92a6a8a27 100644
--- a/src/sagemaker/serve/model_format/mlflow/utils.py
+++ b/src/sagemaker/serve/model_format/mlflow/utils.py
@@ -278,7 +278,7 @@ def _download_s3_artifacts(s3_path: str, dst_path: str, session: Session) -> Non
                 os.makedirs(local_file_dir, exist_ok=True)
 
                 # Download the file
-                print(f"Downloading {key} to {local_file_path}")
+                logger.info(f"Downloading {key} to {local_file_path}")
                 s3.download_file(s3_bucket, key, local_file_path)
 
 
@@ -356,6 +356,15 @@ def _select_container_for_mlflow_model(
     logger.info("Auto-detected framework to use is %s", framework_to_use)
     logger.info("Auto-detected framework version is %s", framework_version)
 
+    if framework_version is None:
+        raise ValueError(
+            (
+                "Unable to auto detect framework version. Please provide framework %s as part of the "
+                "requirements.txt file for deployment flavor %s"
+            )
+            % (framework_to_use, deployment_flavor)
+        )
+
     casted_versions = (
         _cast_to_compatible_version(framework_to_use, framework_version)
         if framework_version
diff --git a/src/sagemaker/serve/utils/lineage_utils.py b/src/sagemaker/serve/utils/lineage_utils.py
index b2e28d26c3..3435e138c9 100644
--- a/src/sagemaker/serve/utils/lineage_utils.py
+++ b/src/sagemaker/serve/utils/lineage_utils.py
@@ -28,7 +28,7 @@
 from sagemaker.lineage.query import LineageSourceEnum
 from sagemaker.serve.model_format.mlflow.constants import (
     MLFLOW_RUN_ID_REGEX,
-    MODEL_PACAKGE_ARN_REGEX,
+    MODEL_PACKAGE_ARN_REGEX,
     S3_PATH_REGEX,
     MLFLOW_REGISTRY_PATH_REGEX,
 )
@@ -107,7 +107,7 @@ def _get_mlflow_model_path_type(mlflow_model_path: str) -> str:
     """
     mlflow_rub_id_pattern = MLFLOW_RUN_ID_REGEX
     mlflow_registry_id_pattern = MLFLOW_REGISTRY_PATH_REGEX
-    sagemaker_arn_pattern = MODEL_PACAKGE_ARN_REGEX
+    sagemaker_arn_pattern = MODEL_PACKAGE_ARN_REGEX
     s3_pattern = S3_PATH_REGEX
 
     if re.match(mlflow_rub_id_pattern, mlflow_model_path):
diff --git a/src/sagemaker/serve/utils/telemetry_logger.py b/src/sagemaker/serve/utils/telemetry_logger.py
index 64cbce03e8..8983a4b5c9 100644
--- a/src/sagemaker/serve/utils/telemetry_logger.py
+++ b/src/sagemaker/serve/utils/telemetry_logger.py
@@ -19,7 +19,16 @@
 
 from sagemaker import Session, exceptions
 from sagemaker.serve.mode.function_pointers import Mode
+from sagemaker.serve.model_format.mlflow.constants import MLFLOW_MODEL_PATH
 from sagemaker.serve.utils.exceptions import ModelBuilderException
+from sagemaker.serve.utils.lineage_constants import (
+    MLFLOW_LOCAL_PATH,
+    MLFLOW_S3_PATH,
+    MLFLOW_MODEL_PACKAGE_PATH,
+    MLFLOW_RUN_ID,
+    MLFLOW_REGISTRY_PATH,
+)
+from sagemaker.serve.utils.lineage_utils import _get_mlflow_model_path_type
 from sagemaker.serve.utils.types import ModelServer, ImageUriOption
 from sagemaker.serve.validations.check_image_uri import is_1p_image_uri
 from sagemaker.user_agent import SDK_VERSION
@@ -51,6 +60,14 @@
     str(ModelServer.TGI): 6,
 }
 
+MLFLOW_MODEL_PATH_CODE = {
+    MLFLOW_LOCAL_PATH: 1,
+    MLFLOW_S3_PATH: 2,
+    MLFLOW_MODEL_PACKAGE_PATH: 3,
+    MLFLOW_RUN_ID: 4,
+    MLFLOW_REGISTRY_PATH: 5,
+}
+
 
 def _capture_telemetry(func_name: str):
     """Placeholder docstring"""
@@ -78,6 +95,11 @@ def wrapper(self, *args, **kwargs):
             if self.sagemaker_session and self.sagemaker_session.endpoint_arn:
                 extra += f"&x-endpointArn={self.sagemaker_session.endpoint_arn}"
 
+            if getattr(self, "_is_mlflow_model", False):
+                mlflow_model_path = self.model_metadata[MLFLOW_MODEL_PATH]
+                mlflow_model_path_type = _get_mlflow_model_path_type(mlflow_model_path)
+                extra += f"&x-mlflowModelPathType={MLFLOW_MODEL_PATH_CODE[mlflow_model_path_type]}"
+
             start_timer = perf_counter()
             try:
                 response = func(self, *args, **kwargs)
diff --git a/tests/unit/sagemaker/serve/model_format/mlflow/test_mlflow_utils.py b/tests/unit/sagemaker/serve/model_format/mlflow/test_mlflow_utils.py
index 27c3a5280f..23d1315647 100644
--- a/tests/unit/sagemaker/serve/model_format/mlflow/test_mlflow_utils.py
+++ b/tests/unit/sagemaker/serve/model_format/mlflow/test_mlflow_utils.py
@@ -418,6 +418,61 @@ def test_select_container_for_mlflow_model_no_dlc_detected(
         )
 
 
+@patch("sagemaker.image_uris.retrieve")
+@patch("sagemaker.serve.model_format.mlflow.utils._cast_to_compatible_version")
+@patch("sagemaker.serve.model_format.mlflow.utils._get_framework_version_from_requirements")
+@patch(
+    "sagemaker.serve.model_format.mlflow.utils._get_python_version_from_parsed_mlflow_model_file"
+)
+@patch("sagemaker.serve.model_format.mlflow.utils._get_all_flavor_metadata")
+@patch("sagemaker.serve.model_format.mlflow.utils._generate_mlflow_artifact_path")
+def test_select_container_for_mlflow_model_no_framework_version_detected(
+    mock_generate_mlflow_artifact_path,
+    mock_get_all_flavor_metadata,
+    mock_get_python_version_from_parsed_mlflow_model_file,
+    mock_get_framework_version_from_requirements,
+    mock_cast_to_compatible_version,
+    mock_image_uris_retrieve,
+):
+    mlflow_model_src_path = "/path/to/mlflow_model"
+    deployment_flavor = "pytorch"
+    region = "us-west-2"
+    instance_type = "ml.m5.xlarge"
+
+    mock_requirements_path = "/path/to/requirements.txt"
+    mock_metadata_path = "/path/to/mlmodel"
+    mock_flavor_metadata = {"pytorch": {"some_key": "some_value"}}
+    mock_python_version = "3.8.6"
+
+    mock_generate_mlflow_artifact_path.side_effect = lambda path, artifact: (
+        mock_requirements_path if artifact == "requirements.txt" else mock_metadata_path
+    )
+    mock_get_all_flavor_metadata.return_value = mock_flavor_metadata
+    mock_get_python_version_from_parsed_mlflow_model_file.return_value = mock_python_version
+    mock_get_framework_version_from_requirements.return_value = None
+
+    with pytest.raises(
+        ValueError,
+        match="Unable to auto detect framework version. Please provide framework "
+        "pytorch as part of the requirements.txt file for deployment flavor "
+        "pytorch",
+    ):
+        _select_container_for_mlflow_model(
+            mlflow_model_src_path, deployment_flavor, region, instance_type
+        )
+
+        mock_generate_mlflow_artifact_path.assert_any_call(
+            mlflow_model_src_path, "requirements.txt"
+        )
+        mock_generate_mlflow_artifact_path.assert_any_call(mlflow_model_src_path, "MLmodel")
+        mock_get_all_flavor_metadata.assert_called_once_with(mock_metadata_path)
+        mock_get_framework_version_from_requirements.assert_called_once_with(
+            deployment_flavor, mock_requirements_path
+        )
+        mock_cast_to_compatible_version.assert_not_called()
+        mock_image_uris_retrieve.assert_not_called()
+
+
 def test_validate_input_for_mlflow():
     _validate_input_for_mlflow(ModelServer.TORCHSERVE, "pytorch")
 
diff --git a/tests/unit/sagemaker/serve/utils/test_telemetry_logger.py b/tests/unit/sagemaker/serve/utils/test_telemetry_logger.py
index e8273dd9a1..33af575e8f 100644
--- a/tests/unit/sagemaker/serve/utils/test_telemetry_logger.py
+++ b/tests/unit/sagemaker/serve/utils/test_telemetry_logger.py
@@ -14,6 +14,7 @@
 import unittest
 from unittest.mock import Mock, patch
 from sagemaker.serve import Mode, ModelServer
+from sagemaker.serve.model_format.mlflow.constants import MLFLOW_MODEL_PATH
 from sagemaker.serve.utils.telemetry_logger import (
     _send_telemetry,
     _capture_telemetry,
@@ -32,9 +33,13 @@
     "763104351884.dkr.ecr.us-east-1.amazonaws.com/"
     "huggingface-pytorch-inference:2.0.0-transformers4.28.1-cpu-py310-ubuntu20.04"
 )
+MOCK_PYTORCH_CONTAINER = (
+    "763104351884.dkr.ecr.us-west-2.amazonaws.com/pytorch-inference:2.0.1-cpu-py310"
+)
 MOCK_HUGGINGFACE_ID = "meta-llama/Llama-2-7b-hf"
 MOCK_EXCEPTION = LocalModelOutOfMemoryException("mock raise ex")
 MOCK_ENDPOINT_ARN = "arn:aws:sagemaker:us-west-2:123456789012:endpoint/test"
+MOCK_MODEL_METADATA_FOR_MLFLOW = {MLFLOW_MODEL_PATH: "s3://some_path"}
 
 
 class ModelBuilderMock:
@@ -239,3 +244,34 @@ def test_construct_url_with_failure_reason_and_extra_info(self):
             f"&x-extra={mock_extra_info}"
         )
         self.assertEquals(ret_url, expected_base_url)
+
+    @patch("sagemaker.serve.utils.telemetry_logger._send_telemetry")
+    def test_capture_telemetry_decorator_mlflow_success(self, mock_send_telemetry):
+        mock_model_builder = ModelBuilderMock()
+        mock_model_builder.serve_settings.telemetry_opt_out = False
+        mock_model_builder.image_uri = MOCK_PYTORCH_CONTAINER
+        mock_model_builder._is_mlflow_model = True
+        mock_model_builder.model_metadata = MOCK_MODEL_METADATA_FOR_MLFLOW
+        mock_model_builder._is_custom_image_uri = False
+        mock_model_builder.mode = Mode.SAGEMAKER_ENDPOINT
+        mock_model_builder.model_server = ModelServer.TORCHSERVE
+        mock_model_builder.sagemaker_session.endpoint_arn = MOCK_ENDPOINT_ARN
+
+        mock_model_builder.mock_deploy()
+
+        args = mock_send_telemetry.call_args.args
+        latency = str(args[5]).split("latency=")[1]
+        expected_extra_str = (
+            f"{MOCK_FUNC_NAME}"
+            "&x-modelServer=1"
+            "&x-imageTag=pytorch-inference:2.0.1-cpu-py310"
+            f"&x-sdkVersion={SDK_VERSION}"
+            f"&x-defaultImageUsage={ImageUriOption.DEFAULT_IMAGE.value}"
+            f"&x-endpointArn={MOCK_ENDPOINT_ARN}"
+            f"&x-mlflowModelPathType=2"
+            f"&x-latency={latency}"
+        )
+
+        mock_send_telemetry.assert_called_once_with(
+            "1", 3, MOCK_SESSION, None, None, expected_extra_str
+        )

From 2bd30c41d27cfd6333341aaedd6ebef93066c1fb Mon Sep 17 00:00:00 2001
From: Haotian An <33510317+Captainia@users.noreply.github.com>
Date: Tue, 14 May 2024 16:13:09 -0400
Subject: [PATCH 28/58] feat: add new images for HF TGI release (#4677)

* chore: add new images for HF TGI release

* test
---
 .../huggingface-llm-neuronx.json              | 31 ++++++-
 .../image_uri_config/huggingface-llm.json     | 49 +++++++++-
 .../image_uri_config/huggingface-neuronx.json | 89 ++++++++++++++++++-
 src/sagemaker/jumpstart/utils.py              |  2 +-
 .../image_uris/test_huggingface_llm.py        |  2 +
 5 files changed, 168 insertions(+), 5 deletions(-)

diff --git a/src/sagemaker/image_uri_config/huggingface-llm-neuronx.json b/src/sagemaker/image_uri_config/huggingface-llm-neuronx.json
index 9da18c1b56..9efbbea305 100644
--- a/src/sagemaker/image_uri_config/huggingface-llm-neuronx.json
+++ b/src/sagemaker/image_uri_config/huggingface-llm-neuronx.json
@@ -4,7 +4,7 @@
             "inf2"
         ],
         "version_aliases": {
-            "0.0": "0.0.16"
+            "0.0": "0.0.22"
         },
         "versions": {
             "0.0.16": {
@@ -180,6 +180,35 @@
                 "container_version": {
                     "inf2": "ubuntu22.04"
                 }
+            },
+            "0.0.22": {
+                "py_versions": [
+                    "py310"
+                ],
+                "registries": {
+                    "ap-northeast-1": "763104351884",
+                    "ap-south-1": "763104351884",
+                    "ap-south-2": "772153158452",
+                    "ap-southeast-1": "763104351884",
+                    "ap-southeast-2": "763104351884",
+                    "ap-southeast-4": "457447274322",
+                    "eu-central-1": "763104351884",
+                    "eu-central-2": "380420809688",
+                    "eu-south-2": "503227376785",
+                    "eu-west-1": "763104351884",
+                    "eu-west-3": "763104351884",
+                    "il-central-1": "780543022126",
+                    "sa-east-1": "763104351884",
+                    "us-east-1": "763104351884",
+                    "us-east-2": "763104351884",
+                    "us-west-2": "763104351884",
+                    "ca-west-1": "204538143572"
+                },
+                "tag_prefix": "2.1.2-optimum0.0.22",
+                "repository": "huggingface-pytorch-tgi-inference",
+                "container_version": {
+                    "inf2": "ubuntu22.04"
+                }
             }
         }
     }
diff --git a/src/sagemaker/image_uri_config/huggingface-llm.json b/src/sagemaker/image_uri_config/huggingface-llm.json
index d357367e6e..3e3f450d23 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.1"
+            "2.0": "2.0.2"
         },
         "versions": {
             "0.6.0": {
@@ -625,6 +625,53 @@
                 "container_version": {
                     "gpu": "cu121-ubuntu22.04"
                 }
+            },
+            "2.0.2": {
+                "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.3.0-tgi2.0.2",
+                "repository": "huggingface-pytorch-tgi-inference",
+                "container_version": {
+                    "gpu": "cu121-ubuntu22.04"
+                }
             }
         }
     }
diff --git a/src/sagemaker/image_uri_config/huggingface-neuronx.json b/src/sagemaker/image_uri_config/huggingface-neuronx.json
index 0d8b7268b1..3721d75c5f 100644
--- a/src/sagemaker/image_uri_config/huggingface-neuronx.json
+++ b/src/sagemaker/image_uri_config/huggingface-neuronx.json
@@ -5,7 +5,8 @@
         ],
         "version_aliases": {
             "4.28": "4.28.1",
-            "4.34": "4.34.1"
+            "4.34": "4.34.1",
+            "4.36": "4.36.2"
         },
         "versions": {
             "4.28.1": {
@@ -79,6 +80,42 @@
                         "sdk2.15.0"
                     ]
                 }
+            },
+            "4.36.2": {
+                "version_aliases": {
+                    "pytorch1.13": "pytorch1.13.1"
+                },
+                "pytorch1.13.1": {
+                    "py_versions": [
+                        "py310"
+                    ],
+                    "repository": "huggingface-pytorch-inference-neuronx",
+                    "registries": {
+                        "ap-northeast-1": "763104351884",
+                        "ap-south-1": "763104351884",
+                        "ap-south-2": "772153158452",
+                        "ap-southeast-1": "763104351884",
+                        "ap-southeast-2": "763104351884",
+                        "ap-southeast-4": "457447274322",
+                        "eu-central-1": "763104351884",
+                        "eu-central-2": "380420809688",
+                        "eu-south-2": "503227376785",
+                        "eu-west-1": "763104351884",
+                        "eu-west-3": "763104351884",
+                        "il-central-1": "780543022126",
+                        "sa-east-1": "763104351884",
+                        "us-east-1": "763104351884",
+                        "us-east-2": "763104351884",
+                        "us-west-2": "763104351884",
+                        "ca-west-1": "204538143572"
+                    },
+                    "container_version": {
+                        "inf": "ubuntu20.04"
+                    },
+                    "sdk_versions": [
+                        "sdk2.18.0"
+                    ]
+                }
             }
         }
     },
@@ -198,7 +235,8 @@
             },
             "4.36.2": {
                 "version_aliases": {
-                    "pytorch1.13": "pytorch1.13.1"
+                    "pytorch1.13": "pytorch1.13.1",
+                    "pytorch2.1": "pytorch2.1.2"
                 },
                 "pytorch1.13.1": {
                     "py_versions": [
@@ -246,6 +284,53 @@
                     "sdk_versions": [
                         "sdk2.16.1"
                     ]
+                },
+                "pytorch2.1.2": {
+                    "py_versions": [
+                        "py310"
+                    ],
+                    "repository": "huggingface-pytorch-inference-neuronx",
+                    "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-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",
+                        "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"
+                    },
+                    "container_version": {
+                        "inf": "ubuntu20.04"
+                    },
+                    "sdk_versions": [
+                        "sdk2.18.0"
+                    ]
                 }
             }
         }
diff --git a/src/sagemaker/jumpstart/utils.py b/src/sagemaker/jumpstart/utils.py
index 63cfac0939..3cdb7f297e 100644
--- a/src/sagemaker/jumpstart/utils.py
+++ b/src/sagemaker/jumpstart/utils.py
@@ -123,7 +123,7 @@ def get_jumpstart_gated_content_bucket(
 def get_jumpstart_content_bucket(
     region: str = constants.JUMPSTART_DEFAULT_REGION_NAME,
 ) -> str:
-    """Returns regionalized content bucket name for JumpStart.
+    """Returns the regionalized content bucket name for JumpStart.
 
     Raises:
         ValueError: If JumpStart is not launched in ``region``.
diff --git a/tests/unit/sagemaker/image_uris/test_huggingface_llm.py b/tests/unit/sagemaker/image_uris/test_huggingface_llm.py
index 2ef981a109..5a9662c164 100644
--- a/tests/unit/sagemaker/image_uris/test_huggingface_llm.py
+++ b/tests/unit/sagemaker/image_uris/test_huggingface_llm.py
@@ -33,6 +33,7 @@
         "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",
+        "2.0.2": "2.3.0-tgi2.0.2-gpu-py310-cu121-ubuntu22.04",
     },
     "inf2": {
         "0.0.16": "1.13.1-optimum0.0.16-neuronx-py310-ubuntu22.04",
@@ -41,6 +42,7 @@
         "0.0.19": "1.13.1-optimum0.0.19-neuronx-py310-ubuntu22.04",
         "0.0.20": "1.13.1-optimum0.0.20-neuronx-py310-ubuntu22.04",
         "0.0.21": "1.13.1-optimum0.0.21-neuronx-py310-ubuntu22.04",
+        "0.0.22": "2.1.2-optimum0.0.22-neuronx-py310-ubuntu22.04",
     },
 }
 

From 6647c23c6d36235d4db793814578def0acb6f600 Mon Sep 17 00:00:00 2001
From: Prateek M Desai <prateekmdesai04@gmail.com>
Date: Tue, 14 May 2024 13:47:22 -0700
Subject: [PATCH 29/58] feature: AutoGluon 1.1.0 image_uris update (#4679)

Co-authored-by: Ubuntu <ubuntu@ip-172-31-9-154.us-west-2.compute.internal>
---
 src/sagemaker/image_uri_config/autogluon.json | 90 ++++++++++++++++++-
 1 file changed, 88 insertions(+), 2 deletions(-)

diff --git a/src/sagemaker/image_uri_config/autogluon.json b/src/sagemaker/image_uri_config/autogluon.json
index 57ce47f94c..1ea6441774 100644
--- a/src/sagemaker/image_uri_config/autogluon.json
+++ b/src/sagemaker/image_uri_config/autogluon.json
@@ -11,7 +11,8 @@
             "0.6": "0.6.2",
             "0.7": "0.7.0",
             "0.8": "0.8.2",
-            "1.0": "1.0.0"
+            "1.0": "1.0.0",
+            "1.1": "1.1.0"
         },
         "versions": {
             "0.3.1": {
@@ -480,6 +481,47 @@
                 "py_versions": [
                     "py310"
                 ]
+            },
+            "1.1.0": {
+                "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",
+                    "ap-southeast-4": "457447274322",
+                    "ca-central-1": "763104351884",
+                    "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-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"
+                },
+                "repository": "autogluon-training",
+                "processors": [
+                    "cpu",
+                    "gpu"
+                ],
+                "py_versions": [
+                    "py310"
+                ]
             }
         }
     },
@@ -491,7 +533,8 @@
             "0.6": "0.6.2",
             "0.7": "0.7.0",
             "0.8": "0.8.2",
-            "1.0": "1.0.0"
+            "1.0": "1.0.0",
+            "1.1": "1.1.0"
         },
         "versions": {
             "0.3.1": {
@@ -987,6 +1030,49 @@
                 "py_versions": [
                     "py310"
                 ]
+            },
+            "1.1.0": {
+                "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",
+                    "ap-southeast-4": "457447274322",
+                    "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-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"
+                },
+                "repository": "autogluon-inference",
+                "processors": [
+                    "cpu",
+                    "gpu"
+                ],
+                "py_versions": [
+                    "py310"
+                ]
             }
         }
     }

From e488c90bd60d69026e63c69bc0e678f37d07d40a Mon Sep 17 00:00:00 2001
From: Mufaddal Rohawala <89424143+mufaddal-rohawala@users.noreply.github.com>
Date: Tue, 14 May 2024 19:31:06 -0700
Subject: [PATCH 30/58] change: add debug logs to workflow container dist
 creation (#4682)

---
 tests/integ/sagemaker/conftest.py | 30 +++++++++++++++++-------------
 1 file changed, 17 insertions(+), 13 deletions(-)

diff --git a/tests/integ/sagemaker/conftest.py b/tests/integ/sagemaker/conftest.py
index 2dc9f7df4d..2640f51515 100644
--- a/tests/integ/sagemaker/conftest.py
+++ b/tests/integ/sagemaker/conftest.py
@@ -176,7 +176,7 @@ def conda_env_yml():
         os.remove(conda_yml_file_name)
 
 
-def _build_container(sagemaker_session, py_version, docker_templete):
+def _build_container(sagemaker_session, py_version, docker_template):
     """Build a dummy test container locally and push a container to an ecr repo"""
 
     region = sagemaker_session.boto_region_name
@@ -189,9 +189,9 @@ def _build_container(sagemaker_session, py_version, docker_templete):
         print("building source archive...")
         source_archive = _generate_sagemaker_sdk_tar(tmpdir)
         with open(os.path.join(tmpdir, "Dockerfile"), "w") as file:
-            file.writelines(
-                docker_templete.format(py_version=py_version, source_archive=source_archive)
-            )
+            content = docker_template.format(py_version=py_version, source_archive=source_archive)
+            print(f"Dockerfile contents: \n{content}\n")
+            file.writelines(content)
 
         docker_client = docker.from_env()
 
@@ -209,6 +209,7 @@ def _build_container(sagemaker_session, py_version, docker_templete):
             raise
 
     if _is_repository_exists(ecr_client, REPO_NAME):
+        print("pushing to session configured account id!")
         sts_client = sagemaker_session.boto_session.client(
             "sts", region_name=region, endpoint_url=sts_regional_endpoint(region)
         )
@@ -218,6 +219,7 @@ def _build_container(sagemaker_session, py_version, docker_templete):
             account_id, sagemaker_session.boto_region_name, REPO_NAME, image_tag
         )
     else:
+        print(f"pushing to account id: {REPO_ACCOUNT_ID}")
         ecr_image = _ecr_image_uri(
             REPO_ACCOUNT_ID,
             sagemaker_session.boto_region_name,
@@ -232,7 +234,7 @@ def _build_container(sagemaker_session, py_version, docker_templete):
     return ecr_image
 
 
-def _build_auto_capture_client_container(py_version, docker_templete):
+def _build_auto_capture_client_container(py_version, docker_template):
     """Build a test docker container that will act as a client for auto_capture tests"""
     with _tmpdir() as tmpdir:
         print("building docker image locally in ", tmpdir)
@@ -240,9 +242,9 @@ def _build_auto_capture_client_container(py_version, docker_templete):
         source_archive = _generate_sdk_tar_with_public_version(tmpdir)
         _move_auto_capture_test_file(tmpdir)
         with open(os.path.join(tmpdir, "Dockerfile"), "w") as file:
-            file.writelines(
-                docker_templete.format(py_version=py_version, source_archive=source_archive)
-            )
+            content = docker_template.format(py_version=py_version, source_archive=source_archive)
+            print(f"Dockerfile contents: \n{content}\n")
+            file.writelines(content)
 
         docker_client = docker.from_env()
 
@@ -276,11 +278,13 @@ def _generate_sagemaker_sdk_tar(destination_folder):
     """
     Run setup.py sdist to generate the PySDK tar file
     """
-    subprocess.run(
-        f"python3 setup.py egg_info --egg-base {destination_folder} sdist -d {destination_folder} -k",
-        shell=True,
-        check=True,
-    )
+    command = f"python3 setup.py egg_info --egg-base {destination_folder} sdist -d {destination_folder} -k"
+    print(f"Running command: {command}")
+    result = subprocess.run(command, shell=True, check=True, capture_output=True)
+    if result.returncode != 0:
+        print(f"Command failed with return code: {result.returncode}")
+        print(f"Standard output: {result.stdout.decode()}")
+        print(f"Standard error: {result.stderr.decode()}")
     destination_folder_contents = os.listdir(destination_folder)
     source_archive = [file for file in destination_folder_contents if file.endswith("tar.gz")][0]
 

From e50ce3206543e9c52781cd1b676a33d89b63670d Mon Sep 17 00:00:00 2001
From: ci <ci>
Date: Wed, 15 May 2024 19:45:10 +0000
Subject: [PATCH 31/58] prepare release v2.220.0

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

diff --git a/CHANGELOG.md b/CHANGELOG.md
index b41a169cdc..5d85261dd7 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,5 +1,21 @@
 # Changelog
 
+## v2.220.0 (2024-05-15)
+
+### Features
+
+ * AutoGluon 1.1.0 image_uris update
+ * add new images for HF TGI release
+ * Add telemetry support for mlflow models
+
+### Bug Fixes and Other Changes
+
+ * add debug logs to workflow container dist creation
+ * model builder race condition on sagemaker session
+ * Add tensorflow_serving support for mlflow models and enable lineage tracking for mlflow models
+ * update image_uri_configs  05-09-2024 07:17:41 PST
+ * skip flakey tests pending investigation
+
 ## v2.219.0 (2024-05-08)
 
 ### Features
diff --git a/VERSION b/VERSION
index 6f59b6ff9a..96d37ee3f0 100644
--- a/VERSION
+++ b/VERSION
@@ -1 +1 @@
-2.219.1.dev0
+2.220.0

From d4f3c91f4b79e2ccfdeaa73ea57cdd4c8dc6fe22 Mon Sep 17 00:00:00 2001
From: ci <ci>
Date: Wed, 15 May 2024 19:45:12 +0000
Subject: [PATCH 32/58] update development version to v2.220.1.dev0

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

diff --git a/VERSION b/VERSION
index 96d37ee3f0..bd29d7bb5b 100644
--- a/VERSION
+++ b/VERSION
@@ -1 +1 @@
-2.220.0
+2.220.1.dev0

From 65cc586969cf402ced7d76e52caa71a0cfeaebb1 Mon Sep 17 00:00:00 2001
From: Samrudhi Sharma <154457034+samruds@users.noreply.github.com>
Date: Wed, 15 May 2024 13:47:30 -0700
Subject: [PATCH 33/58] fix: Image URI should take precedence for HF models
 (#4684)

* Fix: Image URI should take precedence for HF models

* Fix formatting

* Fix formatting

* Fix formatting

* Increase coverage -  UT pass
---
 .../serve/builder/transformers_builder.py     | 10 +++--
 .../builder/test_transformers_builder.py      | 44 +++++++++++++++++++
 2 files changed, 51 insertions(+), 3 deletions(-)

diff --git a/src/sagemaker/serve/builder/transformers_builder.py b/src/sagemaker/serve/builder/transformers_builder.py
index 3d84e314df..ead9b7425f 100644
--- a/src/sagemaker/serve/builder/transformers_builder.py
+++ b/src/sagemaker/serve/builder/transformers_builder.py
@@ -132,17 +132,20 @@ def _create_transformers_model(self) -> Type[Model]:
                 vpc_config=self.vpc_config,
             )
 
-        if self.mode == Mode.LOCAL_CONTAINER:
+        if not self.image_uri and self.mode == Mode.LOCAL_CONTAINER:
             self.image_uri = pysdk_model.serving_image_uri(
                 self.sagemaker_session.boto_region_name, "local"
             )
-        else:
+        elif not self.image_uri:
             self.image_uri = pysdk_model.serving_image_uri(
                 self.sagemaker_session.boto_region_name, self.instance_type
             )
 
         logger.info("Detected %s. Proceeding with the the deployment.", self.image_uri)
 
+        if not pysdk_model.image_uri:
+            pysdk_model.image_uri = self.image_uri
+
         self._original_deploy = pysdk_model.deploy
         pysdk_model.deploy = self._transformers_model_builder_deploy_wrapper
         return pysdk_model
@@ -251,13 +254,14 @@ def _set_instance(self, **kwargs):
         if self.mode == Mode.SAGEMAKER_ENDPOINT:
             if self.nb_instance_type and "instance_type" not in kwargs:
                 kwargs.update({"instance_type": self.nb_instance_type})
+                logger.info("Setting instance type to %s", self.nb_instance_type)
             elif self.instance_type and "instance_type" not in kwargs:
                 kwargs.update({"instance_type": self.instance_type})
+                logger.info("Setting instance type to %s", self.instance_type)
             else:
                 raise ValueError(
                     "Instance type must be provided when deploying to SageMaker Endpoint mode."
                 )
-            logger.info("Setting instance type to %s", self.instance_type)
 
     def _get_supported_version(self, hf_config, hugging_face_version, base_fw):
         """Uses the hugging face json config to pick supported versions"""
diff --git a/tests/unit/sagemaker/serve/builder/test_transformers_builder.py b/tests/unit/sagemaker/serve/builder/test_transformers_builder.py
index e17364f22d..b7e3db79d6 100644
--- a/tests/unit/sagemaker/serve/builder/test_transformers_builder.py
+++ b/tests/unit/sagemaker/serve/builder/test_transformers_builder.py
@@ -58,6 +58,10 @@
 mock_schema_builder = MagicMock()
 mock_schema_builder.sample_input = mock_sample_input
 mock_schema_builder.sample_output = mock_sample_output
+MOCK_IMAGE_CONFIG = (
+    "763104351884.dkr.ecr.us-west-2.amazonaws.com/"
+    "huggingface-pytorch-inference:2.0.0-transformers4.28.1-gpu-py310-cu118-ubuntu20.04-v1.0"
+)
 
 
 class TestTransformersBuilder(unittest.TestCase):
@@ -100,3 +104,43 @@ def test_build_deploy_for_transformers_local_container_and_remote_container(
 
         with self.assertRaises(ValueError) as _:
             model.deploy(mode=Mode.IN_PROCESS)
+
+    @patch(
+        "sagemaker.serve.builder.transformers_builder._get_nb_instance",
+        return_value="ml.g5.24xlarge",
+    )
+    @patch("sagemaker.serve.builder.transformers_builder._capture_telemetry", side_effect=None)
+    def test_image_uri(
+        self,
+        mock_get_nb_instance,
+        mock_telemetry,
+    ):
+        builder = ModelBuilder(
+            model=mock_model_id,
+            schema_builder=mock_schema_builder,
+            mode=Mode.LOCAL_CONTAINER,
+            image_uri=MOCK_IMAGE_CONFIG,
+        )
+
+        builder._prepare_for_mode = MagicMock()
+        builder._prepare_for_mode.side_effect = None
+
+        model = builder.build()
+        builder.serve_settings.telemetry_opt_out = True
+
+        builder.modes[str(Mode.LOCAL_CONTAINER)] = MagicMock()
+        predictor = model.deploy(model_data_download_timeout=1800)
+
+        assert builder.image_uri == MOCK_IMAGE_CONFIG
+        assert builder.env_vars["MODEL_LOADING_TIMEOUT"] == "1800"
+        assert isinstance(predictor, TransformersLocalModePredictor)
+
+        assert builder.nb_instance_type == "ml.g5.24xlarge"
+
+        builder._original_deploy = MagicMock()
+        builder._prepare_for_mode.return_value = (None, {})
+        predictor = model.deploy(mode=Mode.SAGEMAKER_ENDPOINT, role="mock_role_arn")
+        assert "HF_MODEL_ID" in model.env
+
+        with self.assertRaises(ValueError) as _:
+            model.deploy(mode=Mode.IN_PROCESS)

From 63a9ac307d9628b5bae8f9e2c88f4657e7970786 Mon Sep 17 00:00:00 2001
From: Haixin Wang <98612668+haixiw@users.noreply.github.com>
Date: Wed, 15 May 2024 15:22:08 -0700
Subject: [PATCH 34/58] feat: onboard tei image config to pysdk (#4681)

* feat: onboard tei image config to pysdk

* fix formatting issue

* minor fix func name

* fix unit tests

---------

Co-authored-by: Mufaddal Rohawala <89424143+mufaddal-rohawala@users.noreply.github.com>
---
 src/sagemaker/huggingface/llm_utils.py        |  7 +++
 .../image_uri_config/huggingface-tei.json     | 59 +++++++++++++++++++
 .../image_uris/test_huggingface_llm.py        | 24 ++++++++
 3 files changed, 90 insertions(+)
 create mode 100644 src/sagemaker/image_uri_config/huggingface-tei.json

diff --git a/src/sagemaker/huggingface/llm_utils.py b/src/sagemaker/huggingface/llm_utils.py
index de5e624dbc..974cffcddf 100644
--- a/src/sagemaker/huggingface/llm_utils.py
+++ b/src/sagemaker/huggingface/llm_utils.py
@@ -65,6 +65,13 @@ def get_huggingface_llm_image_uri(
             image_scope="inference",
             inference_tool="neuronx",
         )
+    if backend == "huggingface-tei":
+        return image_uris.retrieve(
+            "huggingface-tei",
+            region=region,
+            version=version,
+            image_scope="inference",
+        )
     if backend == "lmi":
         version = version or "0.24.0"
         return image_uris.retrieve(framework="djl-deepspeed", region=region, version=version)
diff --git a/src/sagemaker/image_uri_config/huggingface-tei.json b/src/sagemaker/image_uri_config/huggingface-tei.json
new file mode 100644
index 0000000000..b7c597df18
--- /dev/null
+++ b/src/sagemaker/image_uri_config/huggingface-tei.json
@@ -0,0 +1,59 @@
+{
+    "inference": {
+        "processors": [
+            "gpu"
+        ],
+        "version_aliases": {
+            "1.2": "1.2.3"
+        },
+        "versions": {
+            "1.2.3": {
+                "py_versions": [
+                    "py310"
+                ],
+                "registries": {
+                    "af-south-1": "510948584623",
+                    "ap-east-1": "651117190479",
+                    "ap-northeast-1": "354813040037",
+                    "ap-northeast-2": "366743142698",
+                    "ap-northeast-3": "867004704886",
+                    "ap-south-1": "720646828776",
+                    "ap-south-2": "628508329040",
+                    "ap-southeast-1": "121021644041",
+                    "ap-southeast-2": "783357654285",
+                    "ap-southeast-3": "951798379941",
+                    "ap-southeast-4": "106583098589",
+                    "ca-central-1": "341280168497",
+                    "ca-west-1": "190319476487",
+                    "cn-north-1": "450853457545",
+                    "cn-northwest-1": "451049120500",
+                    "eu-central-1": "492215442770",
+                    "eu-central-2": "680994064768",
+                    "eu-north-1": "662702820516",
+                    "eu-south-1": "978288397137",
+                    "eu-south-2": "104374241257",
+                    "eu-west-1": "141502667606",
+                    "eu-west-2": "764974769150",
+                    "eu-west-3": "659782779980",
+                    "il-central-1": "898809789911",
+                    "me-central-1": "272398656194",
+                    "me-south-1": "801668240914",
+                    "sa-east-1": "737474898029",
+                    "us-east-1": "683313688378",
+                    "us-east-2": "257758044811",
+                    "us-gov-east-1": "237065988967",
+                    "us-gov-west-1": "414596584902",
+                    "us-iso-east-1": "833128469047",
+                    "us-isob-east-1": "281123927165",
+                    "us-west-1": "746614075791",
+                    "us-west-2": "246618743249"
+                },
+                "tag_prefix": "2.0.1-tei1.2.3",
+                "repository": "tei",
+                "container_version": {
+                    "gpu": "cu122-ubuntu22.04"
+                }
+            }
+        }
+    }
+}
\ No newline at end of file
diff --git a/tests/unit/sagemaker/image_uris/test_huggingface_llm.py b/tests/unit/sagemaker/image_uris/test_huggingface_llm.py
index 5a9662c164..b1e8e8253e 100644
--- a/tests/unit/sagemaker/image_uris/test_huggingface_llm.py
+++ b/tests/unit/sagemaker/image_uris/test_huggingface_llm.py
@@ -18,6 +18,11 @@
 from tests.unit.sagemaker.image_uris import expected_uris, conftest
 
 LMI_VERSIONS = ["0.24.0"]
+TEI_VERSIONS_MAPPING = {
+    "gpu": {
+        "1.2.3": "2.0.1-tei1.2.3-gpu-py310-cu122-ubuntu22.04",
+    },
+}
 HF_VERSIONS_MAPPING = {
     "gpu": {
         "0.6.0": "2.0.0-tgi0.6.0-gpu-py39-cu118-ubuntu20.04",
@@ -68,6 +73,25 @@ def test_huggingface_uris(load_config):
             assert expected == uri
 
 
+@pytest.mark.parametrize("load_config", ["huggingface-tei.json"], indirect=True)
+def test_huggingface_tei_uris(load_config):
+    VERSIONS = load_config["inference"]["versions"]
+    device = load_config["inference"]["processors"][0]
+    backend = "huggingface-tei"
+    for version in VERSIONS:
+        ACCOUNTS = load_config["inference"]["versions"][version]["registries"]
+        for region in ACCOUNTS.keys():
+            uri = get_huggingface_llm_image_uri(backend, region=region, version=version)
+            expected = expected_uris.huggingface_llm_framework_uri(
+                "tei",
+                ACCOUNTS[region],
+                version,
+                TEI_VERSIONS_MAPPING[device][version],
+                region=region,
+            )
+            assert expected == uri
+
+
 @pytest.mark.parametrize("load_config", ["huggingface-llm.json"], indirect=True)
 def test_lmi_uris(load_config):
     VERSIONS = load_config["inference"]["versions"]

From 8002d7f92c42d4118631f9c82d51dcaf1b5dc96c Mon Sep 17 00:00:00 2001
From: Jonathan Makunga <54963715+makungaj1@users.noreply.github.com>
Date: Wed, 15 May 2024 19:18:16 -0700
Subject: [PATCH 35/58] fix: model builder limited container support for
 endpoint mode. (#4683)

* Allow ModelBuilder's endpoint mode for Jumpstart models packaged with containers other than TGI and DJL

* increase coverage

* Add JS Support for MMS Serving

* Add JS Support for MMS Serving

* Unit tests

* Refactoring

* Refactoring

* Refactoring

---------

Co-authored-by: Jonathan Makunga <makung@amazon.com>
---
 .../serve/builder/jumpstart_builder.py        | 109 +++++++++++++-----
 .../multi_model_server/prepare.py             |  27 +++++
 .../serve/builder/test_js_builder.py          |  86 ++++++++++++++
 3 files changed, 191 insertions(+), 31 deletions(-)

diff --git a/src/sagemaker/serve/builder/jumpstart_builder.py b/src/sagemaker/serve/builder/jumpstart_builder.py
index e3368869fe..8cb42689fe 100644
--- a/src/sagemaker/serve/builder/jumpstart_builder.py
+++ b/src/sagemaker/serve/builder/jumpstart_builder.py
@@ -23,6 +23,7 @@
 from sagemaker import model_uris
 from sagemaker.serve.model_server.djl_serving.prepare import prepare_djl_js_resources
 from sagemaker.serve.model_server.djl_serving.utils import _get_admissible_tensor_parallel_degrees
+from sagemaker.serve.model_server.multi_model_server.prepare import prepare_mms_js_resources
 from sagemaker.serve.model_server.tgi.prepare import prepare_tgi_js_resources, _create_dir_structure
 from sagemaker.serve.mode.function_pointers import Mode
 from sagemaker.serve.utils.exceptions import (
@@ -35,6 +36,7 @@
 from sagemaker.serve.utils.predictors import (
     DjlLocalModePredictor,
     TgiLocalModePredictor,
+    TransformersLocalModePredictor,
 )
 from sagemaker.serve.utils.local_hardware import (
     _get_nb_instance,
@@ -90,6 +92,7 @@ def __init__(self):
         self.existing_properties = None
         self.prepared_for_tgi = None
         self.prepared_for_djl = None
+        self.prepared_for_mms = None
         self.schema_builder = None
         self.nb_instance_type = None
         self.ram_usage_model_load = None
@@ -137,7 +140,11 @@ def _js_builder_deploy_wrapper(self, *args, **kwargs) -> Type[PredictorBase]:
 
             if overwrite_mode == Mode.SAGEMAKER_ENDPOINT:
                 self.mode = self.pysdk_model.mode = Mode.SAGEMAKER_ENDPOINT
-                if not hasattr(self, "prepared_for_djl") or not hasattr(self, "prepared_for_tgi"):
+                if (
+                    not hasattr(self, "prepared_for_djl")
+                    or not hasattr(self, "prepared_for_tgi")
+                    or not hasattr(self, "prepared_for_mms")
+                ):
                     self.pysdk_model.model_data, env = self._prepare_for_mode()
             elif overwrite_mode == Mode.LOCAL_CONTAINER:
                 self.mode = self.pysdk_model.mode = Mode.LOCAL_CONTAINER
@@ -160,6 +167,13 @@ def _js_builder_deploy_wrapper(self, *args, **kwargs) -> Type[PredictorBase]:
                         dependencies=self.dependencies,
                         model_data=self.pysdk_model.model_data,
                     )
+                elif not hasattr(self, "prepared_for_mms"):
+                    self.js_model_config, self.prepared_for_mms = prepare_mms_js_resources(
+                        model_path=self.model_path,
+                        js_id=self.model,
+                        dependencies=self.dependencies,
+                        model_data=self.pysdk_model.model_data,
+                    )
 
                 self._prepare_for_mode()
                 env = {}
@@ -179,6 +193,10 @@ def _js_builder_deploy_wrapper(self, *args, **kwargs) -> Type[PredictorBase]:
                 predictor = TgiLocalModePredictor(
                     self.modes[str(Mode.LOCAL_CONTAINER)], serializer, deserializer
                 )
+            elif self.model_server == ModelServer.MMS:
+                predictor = TransformersLocalModePredictor(
+                    self.modes[str(Mode.LOCAL_CONTAINER)], serializer, deserializer
+                )
 
             ram_usage_before = _get_ram_usage_mb()
             self.modes[str(Mode.LOCAL_CONTAINER)].create_server(
@@ -254,6 +272,24 @@ def _build_for_tgi_jumpstart(self):
 
         self.pysdk_model.env.update(env)
 
+    def _build_for_mms_jumpstart(self):
+        """Placeholder docstring"""
+
+        env = {}
+        if self.mode == Mode.LOCAL_CONTAINER:
+            if not hasattr(self, "prepared_for_mms"):
+                self.js_model_config, self.prepared_for_mms = prepare_mms_js_resources(
+                    model_path=self.model_path,
+                    js_id=self.model,
+                    dependencies=self.dependencies,
+                    model_data=self.pysdk_model.model_data,
+                )
+            self._prepare_for_mode()
+        elif self.mode == Mode.SAGEMAKER_ENDPOINT and hasattr(self, "prepared_for_mms"):
+            self.pysdk_model.model_data, env = self._prepare_for_mode()
+
+        self.pysdk_model.env.update(env)
+
     def _tune_for_js(self, sharded_supported: bool, max_tuning_duration: int = 1800):
         """Tune for Jumpstart Models in Local Mode.
 
@@ -264,11 +300,6 @@ def _tune_for_js(self, sharded_supported: bool, max_tuning_duration: int = 1800)
         returns:
             Tuned Model.
         """
-        if self.mode != Mode.LOCAL_CONTAINER:
-            logger.warning(
-                "Tuning is only a %s capability. Returning original model.", Mode.LOCAL_CONTAINER
-            )
-            return self.pysdk_model
 
         num_shard_env_var_name = "SM_NUM_GPUS"
         if "OPTION_TENSOR_PARALLEL_DEGREE" in self.pysdk_model.env.keys():
@@ -437,42 +468,58 @@ def _build_for_jumpstart(self):
         self.secret_key = None
         self.jumpstart = True
 
-        pysdk_model = self._create_pre_trained_js_model()
+        self.pysdk_model = self._create_pre_trained_js_model()
+        self.pysdk_model.tune = lambda *args, **kwargs: self._default_tune()
 
-        image_uri = pysdk_model.image_uri
+        logger.info(
+            "JumpStart ID %s is packaged with Image URI: %s", self.model, self.pysdk_model.image_uri
+        )
 
-        logger.info("JumpStart ID %s is packaged with Image URI: %s", self.model, image_uri)
+        if self.mode != Mode.SAGEMAKER_ENDPOINT:
+            if self._is_gated_model(self.pysdk_model):
+                raise ValueError(
+                    "JumpStart Gated Models are only supported in SAGEMAKER_ENDPOINT mode."
+                )
 
-        if self._is_gated_model(pysdk_model) and self.mode != Mode.SAGEMAKER_ENDPOINT:
-            raise ValueError(
-                "JumpStart Gated Models are only supported in SAGEMAKER_ENDPOINT mode."
-            )
+            if "djl-inference" in self.pysdk_model.image_uri:
+                logger.info("Building for DJL JumpStart Model ID...")
+                self.model_server = ModelServer.DJL_SERVING
+                self.image_uri = self.pysdk_model.image_uri
 
-        if "djl-inference" in image_uri:
-            logger.info("Building for DJL JumpStart Model ID...")
-            self.model_server = ModelServer.DJL_SERVING
+                self._build_for_djl_jumpstart()
 
-            self.pysdk_model = pysdk_model
-            self.image_uri = self.pysdk_model.image_uri
+                self.pysdk_model.tune = self.tune_for_djl_jumpstart
+            elif "tgi-inference" in self.pysdk_model.image_uri:
+                logger.info("Building for TGI JumpStart Model ID...")
+                self.model_server = ModelServer.TGI
+                self.image_uri = self.pysdk_model.image_uri
 
-            self._build_for_djl_jumpstart()
+                self._build_for_tgi_jumpstart()
 
-            self.pysdk_model.tune = self.tune_for_djl_jumpstart
-        elif "tgi-inference" in image_uri:
-            logger.info("Building for TGI JumpStart Model ID...")
-            self.model_server = ModelServer.TGI
+                self.pysdk_model.tune = self.tune_for_tgi_jumpstart
+            elif "huggingface-pytorch-inference:" in self.pysdk_model.image_uri:
+                logger.info("Building for MMS JumpStart Model ID...")
+                self.model_server = ModelServer.MMS
+                self.image_uri = self.pysdk_model.image_uri
 
-            self.pysdk_model = pysdk_model
-            self.image_uri = self.pysdk_model.image_uri
+                self._build_for_mms_jumpstart()
+            else:
+                raise ValueError(
+                    "JumpStart Model ID was not packaged "
+                    "with djl-inference, tgi-inference, or mms-inference container."
+                )
 
-            self._build_for_tgi_jumpstart()
+        return self.pysdk_model
 
-            self.pysdk_model.tune = self.tune_for_tgi_jumpstart
-        else:
-            raise ValueError(
-                "JumpStart Model ID was not packaged with djl-inference or tgi-inference container."
-            )
+    def _default_tune(self):
+        """Logs a warning message if tune is invoked on endpoint mode.
 
+        Returns:
+            Jumpstart Model: ``This`` model
+        """
+        logger.warning(
+            "Tuning is only a %s capability. Returning original model.", Mode.LOCAL_CONTAINER
+        )
         return self.pysdk_model
 
     def _is_gated_model(self, model) -> bool:
diff --git a/src/sagemaker/serve/model_server/multi_model_server/prepare.py b/src/sagemaker/serve/model_server/multi_model_server/prepare.py
index 7a16cc0a43..7059d9026d 100644
--- a/src/sagemaker/serve/model_server/multi_model_server/prepare.py
+++ b/src/sagemaker/serve/model_server/multi_model_server/prepare.py
@@ -15,7 +15,9 @@
 from __future__ import absolute_import
 import logging
 from pathlib import Path
+from typing import List
 
+from sagemaker.serve.model_server.tgi.prepare import _copy_jumpstart_artifacts
 from sagemaker.serve.utils.local_hardware import _check_disk_space, _check_docker_disk_usage
 
 logger = logging.getLogger(__name__)
@@ -36,3 +38,28 @@ def _create_dir_structure(model_path: str) -> tuple:
     _check_docker_disk_usage()
 
     return model_path, code_dir
+
+
+def prepare_mms_js_resources(
+    model_path: str,
+    js_id: str,
+    shared_libs: List[str] = None,
+    dependencies: str = None,
+    model_data: str = None,
+) -> tuple:
+    """Prepare serving when a JumpStart model id is given
+
+    Args:
+        model_path (str) : Argument
+        js_id (str): Argument
+        shared_libs (List[]) : Argument
+        dependencies (str) : Argument
+        model_data (str) : Argument
+
+    Returns:
+        ( str ) :
+
+    """
+    model_path, code_dir = _create_dir_structure(model_path)
+
+    return _copy_jumpstart_artifacts(model_data, js_id, code_dir)
diff --git a/tests/unit/sagemaker/serve/builder/test_js_builder.py b/tests/unit/sagemaker/serve/builder/test_js_builder.py
index 2a0c791215..2065e86818 100644
--- a/tests/unit/sagemaker/serve/builder/test_js_builder.py
+++ b/tests/unit/sagemaker/serve/builder/test_js_builder.py
@@ -63,6 +63,10 @@
     "123456789712.dkr.ecr.us-west-2.amazonaws.com/huggingface-pytorch-tgi"
     "-inference:2.1.1-tgi1.4.0-gpu-py310-cu121-ubuntu20.04"
 )
+mock_invalid_image_uri = (
+    "123456789712.dkr.ecr.us-west-2.amazonaws.com/invalid"
+    "-inference:2.1.1-tgi1.4.0-gpu-py310-cu121-ubuntu20.04"
+)
 mock_djl_image_uri = (
     "123456789712.dkr.ecr.us-west-2.amazonaws.com/djl-inference:0.24.0-neuronx-sdk2.14.1"
 )
@@ -82,6 +86,88 @@
 
 
 class TestJumpStartBuilder(unittest.TestCase):
+    @patch("sagemaker.serve.builder.jumpstart_builder._capture_telemetry", side_effect=None)
+    @patch(
+        "sagemaker.serve.builder.jumpstart_builder.JumpStart._is_jumpstart_model_id",
+        return_value=True,
+    )
+    @patch(
+        "sagemaker.serve.builder.jumpstart_builder.JumpStart._create_pre_trained_js_model",
+        return_value=MagicMock(),
+    )
+    @patch(
+        "sagemaker.serve.builder.jumpstart_builder.prepare_tgi_js_resources",
+        return_value=({"model_type": "t5", "n_head": 71}, True),
+    )
+    @patch("sagemaker.serve.builder.jumpstart_builder._get_ram_usage_mb", return_value=1024)
+    @patch(
+        "sagemaker.serve.builder.jumpstart_builder._get_nb_instance", return_value="ml.g5.24xlarge"
+    )
+    def test__build_for_jumpstart_value_error(
+        self,
+        mock_get_nb_instance,
+        mock_get_ram_usage_mb,
+        mock_prepare_for_tgi,
+        mock_pre_trained_model,
+        mock_is_jumpstart_model,
+        mock_telemetry,
+    ):
+        builder = ModelBuilder(
+            model="facebook/invalid",
+            schema_builder=mock_schema_builder,
+            mode=Mode.LOCAL_CONTAINER,
+        )
+
+        mock_pre_trained_model.return_value.image_uri = mock_invalid_image_uri
+
+        self.assertRaises(
+            ValueError,
+            lambda: builder.build(),
+        )
+
+    @patch("sagemaker.serve.builder.jumpstart_builder._capture_telemetry", side_effect=None)
+    @patch(
+        "sagemaker.serve.builder.jumpstart_builder.JumpStart._is_jumpstart_model_id",
+        return_value=True,
+    )
+    @patch(
+        "sagemaker.serve.builder.jumpstart_builder.JumpStart._create_pre_trained_js_model",
+        return_value=MagicMock(),
+    )
+    @patch(
+        "sagemaker.serve.builder.jumpstart_builder.prepare_mms_js_resources",
+        return_value=({"model_type": "t5", "n_head": 71}, True),
+    )
+    @patch("sagemaker.serve.builder.jumpstart_builder._get_ram_usage_mb", return_value=1024)
+    @patch(
+        "sagemaker.serve.builder.jumpstart_builder._get_nb_instance", return_value="ml.g5.24xlarge"
+    )
+    def test__build_for_mms_jumpstart(
+        self,
+        mock_get_nb_instance,
+        mock_get_ram_usage_mb,
+        mock_prepare_for_mms,
+        mock_pre_trained_model,
+        mock_is_jumpstart_model,
+        mock_telemetry,
+    ):
+        builder = ModelBuilder(
+            model="facebook/galactica-mock-model-id",
+            schema_builder=mock_schema_builder,
+            mode=Mode.LOCAL_CONTAINER,
+        )
+
+        mock_pre_trained_model.return_value.image_uri = (
+            "763104351884.dkr.ecr.us-west-2.amazonaws.com/huggingface"
+            "-pytorch-inference:2.1.0-transformers4.37.0-gpu-py310-cu118"
+            "-ubuntu20.04"
+        )
+
+        builder.build()
+        builder.serve_settings.telemetry_opt_out = True
+
+        mock_prepare_for_mms.assert_called()
+
     @patch("sagemaker.serve.builder.jumpstart_builder._capture_telemetry", side_effect=None)
     @patch(
         "sagemaker.serve.builder.jumpstart_builder.JumpStart._is_jumpstart_model_id",

From 9faa8bee2c62840137256d7855736c4537c50c5f Mon Sep 17 00:00:00 2001
From: Mufaddal Rohawala <89424143+mufaddal-rohawala@users.noreply.github.com>
Date: Thu, 16 May 2024 09:29:16 -0700
Subject: [PATCH 36/58] change: Add more debuging (#4687)

---
 tests/integ/sagemaker/conftest.py | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

diff --git a/tests/integ/sagemaker/conftest.py b/tests/integ/sagemaker/conftest.py
index 2640f51515..043b0c703e 100644
--- a/tests/integ/sagemaker/conftest.py
+++ b/tests/integ/sagemaker/conftest.py
@@ -278,13 +278,14 @@ def _generate_sagemaker_sdk_tar(destination_folder):
     """
     Run setup.py sdist to generate the PySDK tar file
     """
-    command = f"python3 setup.py egg_info --egg-base {destination_folder} sdist -d {destination_folder} -k"
+    command = f"python3 setup.py egg_info --egg-base {destination_folder} sdist -d {destination_folder} -k --verbose"
     print(f"Running command: {command}")
     result = subprocess.run(command, shell=True, check=True, capture_output=True)
     if result.returncode != 0:
         print(f"Command failed with return code: {result.returncode}")
-        print(f"Standard output: {result.stdout.decode()}")
-        print(f"Standard error: {result.stderr.decode()}")
+
+    print(f"Standard output: {result.stdout.decode()}")
+    print(f"Standard error: {result.stderr.decode()}")
     destination_folder_contents = os.listdir(destination_folder)
     source_archive = [file for file in destination_folder_contents if file.endswith("tar.gz")][0]
 

From 6280d81dc56734ccfbf5b5defe0eb5533796bfe1 Mon Sep 17 00:00:00 2001
From: Haixin Wang <98612668+haixiw@users.noreply.github.com>
Date: Thu, 16 May 2024 12:46:07 -0700
Subject: [PATCH 37/58] change: cover tei with image_uris.retrieve API (#4689)

---
 src/sagemaker/image_uris.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/src/sagemaker/image_uris.py b/src/sagemaker/image_uris.py
index 143ecc9bdb..1ca73d0af9 100644
--- a/src/sagemaker/image_uris.py
+++ b/src/sagemaker/image_uris.py
@@ -37,6 +37,7 @@
 ECR_URI_TEMPLATE = "{registry}.dkr.{hostname}/{repository}"
 HUGGING_FACE_FRAMEWORK = "huggingface"
 HUGGING_FACE_LLM_FRAMEWORK = "huggingface-llm"
+HUGGING_FACE_TEI_FRAMEWORK = "huggingface-tei"
 HUGGING_FACE_LLM_NEURONX_FRAMEWORK = "huggingface-llm-neuronx"
 XGBOOST_FRAMEWORK = "xgboost"
 SKLEARN_FRAMEWORK = "sklearn"
@@ -477,6 +478,7 @@ def _validate_version_and_set_if_needed(version, config, framework):
     if version is None and framework in [
         DATA_WRANGLER_FRAMEWORK,
         HUGGING_FACE_LLM_FRAMEWORK,
+        HUGGING_FACE_TEI_FRAMEWORK,
         HUGGING_FACE_LLM_NEURONX_FRAMEWORK,
         STABILITYAI_FRAMEWORK,
     ]:

From 06e6f9d69abb5bbb6192bb0be3fdaf9475b3c681 Mon Sep 17 00:00:00 2001
From: Jonathan Makunga <54963715+makungaj1@users.noreply.github.com>
Date: Thu, 16 May 2024 15:26:50 -0700
Subject: [PATCH 38/58] fix: JS Model with non-TGI/non-DJL deployment failure
 (#4688)

* Debug

* Debug

* Debug

* Debug

* Debug

* Debug

* fix docstyle

* Refactoring

* Add Integ tests

---------

Co-authored-by: Jonathan Makunga <makung@amazon.com>
---
 .../serve/builder/jumpstart_builder.py        | 90 +++++++++----------
 .../sagemaker/serve/test_serve_js_happy.py    | 50 +++++++++++
 2 files changed, 92 insertions(+), 48 deletions(-)

diff --git a/src/sagemaker/serve/builder/jumpstart_builder.py b/src/sagemaker/serve/builder/jumpstart_builder.py
index 8cb42689fe..bc31e8d323 100644
--- a/src/sagemaker/serve/builder/jumpstart_builder.py
+++ b/src/sagemaker/serve/builder/jumpstart_builder.py
@@ -300,6 +300,11 @@ def _tune_for_js(self, sharded_supported: bool, max_tuning_duration: int = 1800)
         returns:
             Tuned Model.
         """
+        if self.mode == Mode.SAGEMAKER_ENDPOINT:
+            logger.warning(
+                "Tuning is only a %s capability. Returning original model.", Mode.LOCAL_CONTAINER
+            )
+            return self.pysdk_model
 
         num_shard_env_var_name = "SM_NUM_GPUS"
         if "OPTION_TENSOR_PARALLEL_DEGREE" in self.pysdk_model.env.keys():
@@ -468,58 +473,47 @@ def _build_for_jumpstart(self):
         self.secret_key = None
         self.jumpstart = True
 
-        self.pysdk_model = self._create_pre_trained_js_model()
-        self.pysdk_model.tune = lambda *args, **kwargs: self._default_tune()
-
-        logger.info(
-            "JumpStart ID %s is packaged with Image URI: %s", self.model, self.pysdk_model.image_uri
-        )
-
-        if self.mode != Mode.SAGEMAKER_ENDPOINT:
-            if self._is_gated_model(self.pysdk_model):
-                raise ValueError(
-                    "JumpStart Gated Models are only supported in SAGEMAKER_ENDPOINT mode."
-                )
-
-            if "djl-inference" in self.pysdk_model.image_uri:
-                logger.info("Building for DJL JumpStart Model ID...")
-                self.model_server = ModelServer.DJL_SERVING
-                self.image_uri = self.pysdk_model.image_uri
-
-                self._build_for_djl_jumpstart()
-
-                self.pysdk_model.tune = self.tune_for_djl_jumpstart
-            elif "tgi-inference" in self.pysdk_model.image_uri:
-                logger.info("Building for TGI JumpStart Model ID...")
-                self.model_server = ModelServer.TGI
-                self.image_uri = self.pysdk_model.image_uri
-
-                self._build_for_tgi_jumpstart()
+        pysdk_model = self._create_pre_trained_js_model()
+        image_uri = pysdk_model.image_uri
 
-                self.pysdk_model.tune = self.tune_for_tgi_jumpstart
-            elif "huggingface-pytorch-inference:" in self.pysdk_model.image_uri:
-                logger.info("Building for MMS JumpStart Model ID...")
-                self.model_server = ModelServer.MMS
-                self.image_uri = self.pysdk_model.image_uri
+        logger.info("JumpStart ID %s is packaged with Image URI: %s", self.model, image_uri)
 
-                self._build_for_mms_jumpstart()
-            else:
-                raise ValueError(
-                    "JumpStart Model ID was not packaged "
-                    "with djl-inference, tgi-inference, or mms-inference container."
-                )
-
-        return self.pysdk_model
+        if self._is_gated_model(pysdk_model) and self.mode != Mode.SAGEMAKER_ENDPOINT:
+            raise ValueError(
+                "JumpStart Gated Models are only supported in SAGEMAKER_ENDPOINT mode."
+            )
 
-    def _default_tune(self):
-        """Logs a warning message if tune is invoked on endpoint mode.
+        if "djl-inference" in image_uri:
+            logger.info("Building for DJL JumpStart Model ID...")
+            self.model_server = ModelServer.DJL_SERVING
+            self.pysdk_model = pysdk_model
+            self.image_uri = self.pysdk_model.image_uri
+
+            self._build_for_djl_jumpstart()
+
+            self.pysdk_model.tune = self.tune_for_djl_jumpstart
+        elif "tgi-inference" in image_uri:
+            logger.info("Building for TGI JumpStart Model ID...")
+            self.model_server = ModelServer.TGI
+            self.pysdk_model = pysdk_model
+            self.image_uri = self.pysdk_model.image_uri
+
+            self._build_for_tgi_jumpstart()
+
+            self.pysdk_model.tune = self.tune_for_tgi_jumpstart
+        elif "huggingface-pytorch-inference:" in image_uri:
+            logger.info("Building for MMS JumpStart Model ID...")
+            self.model_server = ModelServer.MMS
+            self.pysdk_model = pysdk_model
+            self.image_uri = self.pysdk_model.image_uri
+
+            self._build_for_mms_jumpstart()
+        elif self.mode != Mode.SAGEMAKER_ENDPOINT:
+            raise ValueError(
+                "JumpStart Model ID was not packaged "
+                "with djl-inference, tgi-inference, or mms-inference container."
+            )
 
-        Returns:
-            Jumpstart Model: ``This`` model
-        """
-        logger.warning(
-            "Tuning is only a %s capability. Returning original model.", Mode.LOCAL_CONTAINER
-        )
         return self.pysdk_model
 
     def _is_gated_model(self, model) -> bool:
diff --git a/tests/integ/sagemaker/serve/test_serve_js_happy.py b/tests/integ/sagemaker/serve/test_serve_js_happy.py
index 7835c8ae3c..ad0527fcc0 100644
--- a/tests/integ/sagemaker/serve/test_serve_js_happy.py
+++ b/tests/integ/sagemaker/serve/test_serve_js_happy.py
@@ -34,6 +34,14 @@
 JS_MODEL_ID = "huggingface-textgeneration1-gpt-neo-125m-fp16"
 ROLE_NAME = "SageMakerRole"
 
+SAMPLE_MMS_PROMPT = [
+    "How cute your dog is!",
+    "Your dog is so cute.",
+    "The mitochondria is the powerhouse of the cell.",
+]
+SAMPLE_MMS_RESPONSE = {"embedding": []}
+JS_MMS_MODEL_ID = "huggingface-sentencesimilarity-bge-m3"
+
 
 @pytest.fixture
 def happy_model_builder(sagemaker_session):
@@ -46,6 +54,17 @@ def happy_model_builder(sagemaker_session):
     )
 
 
+@pytest.fixture
+def happy_mms_model_builder(sagemaker_session):
+    iam_client = sagemaker_session.boto_session.client("iam")
+    return ModelBuilder(
+        model=JS_MMS_MODEL_ID,
+        schema_builder=SchemaBuilder(SAMPLE_MMS_PROMPT, SAMPLE_MMS_RESPONSE),
+        role_arn=iam_client.get_role(RoleName=ROLE_NAME)["Role"]["Arn"],
+        sagemaker_session=sagemaker_session,
+    )
+
+
 @pytest.mark.skipif(
     PYTHON_VERSION_IS_NOT_310,
     reason="The goal of these test are to test the serving components of our feature",
@@ -75,3 +94,34 @@ def test_happy_tgi_sagemaker_endpoint(happy_model_builder, gpu_instance_type):
             )
             if caught_ex:
                 raise caught_ex
+
+
+@pytest.mark.skipif(
+    PYTHON_VERSION_IS_NOT_310,
+    reason="The goal of these test are to test the serving components of our feature",
+)
+@pytest.mark.slow_test
+def test_happy_mms_sagemaker_endpoint(happy_mms_model_builder, gpu_instance_type):
+    logger.info("Running in SAGEMAKER_ENDPOINT mode...")
+    caught_ex = None
+    model = happy_mms_model_builder.build()
+
+    with timeout(minutes=SERVE_SAGEMAKER_ENDPOINT_TIMEOUT):
+        try:
+            logger.info("Deploying and predicting in SAGEMAKER_ENDPOINT mode...")
+            predictor = model.deploy(instance_type=gpu_instance_type, endpoint_logging=False)
+            logger.info("Endpoint successfully deployed.")
+
+            updated_sample_input = happy_mms_model_builder.schema_builder.sample_input
+
+            predictor.predict(updated_sample_input)
+        except Exception as e:
+            caught_ex = e
+        finally:
+            cleanup_model_resources(
+                sagemaker_session=happy_mms_model_builder.sagemaker_session,
+                model_name=model.name,
+                endpoint_name=model.endpoint_name,
+            )
+            if caught_ex:
+                raise caught_ex

From c9b55a4b1d42313f8545c1515fb4b41f1d5f7a70 Mon Sep 17 00:00:00 2001
From: Samrudhi Sharma <154457034+samruds@users.noreply.github.com>
Date: Fri, 17 May 2024 12:48:08 -0700
Subject: [PATCH 39/58] Feat: Pull latest tei container for sentence
 similiarity models on HuggingFace hub (#4686)

* Update: Pull latest tei container for sentence similiarity models

* Fix formatting

* Address PR comments

* Fix formatting

* Fix check

* Switch sentence similarity to be deployed on tgi

* Fix formatting

* Fix formatting

* Fix formatting

* Fix formatting

* Introduce TEI builder with TGI server

* Fix formmatting

* Add integ test

* Fix formatting

* Add integ test

* Add integ test

* Add integ test

* Add integ test

* Add integ test

* Fix formatting

* Move to G5 for integ test

* Fix formatting

* Integ test updates

* Integ test updates

* Integ test updates

* Fix formatting

* Integ test updates

* Move back to generate for ping

* Integ test updates

* Integ test updates
---
 src/sagemaker/serve/builder/model_builder.py  |  11 +-
 src/sagemaker/serve/builder/tei_builder.py    | 222 ++++++++++++++++++
 tests/integ/sagemaker/serve/test_serve_tei.py | 123 ++++++++++
 .../serve/builder/test_model_builder.py       |  38 +++
 .../serve/builder/test_tei_builder.py         | 152 ++++++++++++
 .../builder/test_transformers_builder.py      |   2 +-
 6 files changed, 543 insertions(+), 5 deletions(-)
 create mode 100644 src/sagemaker/serve/builder/tei_builder.py
 create mode 100644 tests/integ/sagemaker/serve/test_serve_tei.py
 create mode 100644 tests/unit/sagemaker/serve/builder/test_tei_builder.py

diff --git a/src/sagemaker/serve/builder/model_builder.py b/src/sagemaker/serve/builder/model_builder.py
index 42a2b994a8..b24f30fb3e 100644
--- a/src/sagemaker/serve/builder/model_builder.py
+++ b/src/sagemaker/serve/builder/model_builder.py
@@ -36,6 +36,7 @@
 from sagemaker.serve.detector.pickler import save_pkl, save_xgboost
 from sagemaker.serve.builder.serve_settings import _ServeSettings
 from sagemaker.serve.builder.djl_builder import DJL
+from sagemaker.serve.builder.tei_builder import TEI
 from sagemaker.serve.builder.tgi_builder import TGI
 from sagemaker.serve.builder.jumpstart_builder import JumpStart
 from sagemaker.serve.builder.transformers_builder import Transformers
@@ -95,9 +96,9 @@
 }
 
 
-# pylint: disable=attribute-defined-outside-init, disable=E1101, disable=R0901
+# pylint: disable=attribute-defined-outside-init, disable=E1101, disable=R0901, disable=R1705
 @dataclass
-class ModelBuilder(Triton, DJL, JumpStart, TGI, Transformers, TensorflowServing):
+class ModelBuilder(Triton, DJL, JumpStart, TGI, Transformers, TensorflowServing, TEI):
     """Class that builds a deployable model.
 
     Args:
@@ -753,7 +754,7 @@ def build(  # pylint: disable=R0911
                 model_task = self.model_metadata.get("HF_TASK")
             if self._is_jumpstart_model_id():
                 return self._build_for_jumpstart()
-            if self._is_djl():  # pylint: disable=R1705
+            if self._is_djl():
                 return self._build_for_djl()
             else:
                 hf_model_md = get_huggingface_model_metadata(
@@ -764,8 +765,10 @@ def build(  # pylint: disable=R0911
                     model_task = hf_model_md.get("pipeline_tag")
                 if self.schema_builder is None and model_task is not None:
                     self._hf_schema_builder_init(model_task)
-                if model_task == "text-generation":  # pylint: disable=R1705
+                if model_task == "text-generation":
                     return self._build_for_tgi()
+                if model_task == "sentence-similarity":
+                    return self._build_for_tei()
                 elif self._can_fit_on_single_gpu():
                     return self._build_for_transformers()
                 elif (
diff --git a/src/sagemaker/serve/builder/tei_builder.py b/src/sagemaker/serve/builder/tei_builder.py
new file mode 100644
index 0000000000..50d3866468
--- /dev/null
+++ b/src/sagemaker/serve/builder/tei_builder.py
@@ -0,0 +1,222 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You
+# may not use this file except in compliance with the License. A copy of
+# the License is located at
+#
+#     http://aws.amazon.com/apache2.0/
+#
+# or in the "license" file accompanying this file. This file is
+# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
+# ANY KIND, either express or implied. See the License for the specific
+# language governing permissions and limitations under the License.
+"""Holds mixin logic to support deployment of Model ID"""
+from __future__ import absolute_import
+import logging
+from typing import Type
+from abc import ABC, abstractmethod
+
+from sagemaker import image_uris
+from sagemaker.model import Model
+from sagemaker.djl_inference.model import _get_model_config_properties_from_hf
+
+from sagemaker.huggingface import HuggingFaceModel
+from sagemaker.serve.utils.local_hardware import (
+    _get_nb_instance,
+)
+from sagemaker.serve.model_server.tgi.prepare import _create_dir_structure
+from sagemaker.serve.utils.predictors import TgiLocalModePredictor
+from sagemaker.serve.utils.types import ModelServer
+from sagemaker.serve.mode.function_pointers import Mode
+from sagemaker.serve.utils.telemetry_logger import _capture_telemetry
+from sagemaker.base_predictor import PredictorBase
+
+logger = logging.getLogger(__name__)
+
+_CODE_FOLDER = "code"
+
+
+class TEI(ABC):
+    """TEI build logic for ModelBuilder()"""
+
+    def __init__(self):
+        self.model = None
+        self.serve_settings = None
+        self.sagemaker_session = None
+        self.model_path = None
+        self.dependencies = None
+        self.modes = None
+        self.mode = None
+        self.model_server = None
+        self.image_uri = None
+        self._is_custom_image_uri = False
+        self.image_config = None
+        self.vpc_config = None
+        self._original_deploy = None
+        self.hf_model_config = None
+        self._default_tensor_parallel_degree = None
+        self._default_data_type = None
+        self._default_max_tokens = None
+        self.pysdk_model = None
+        self.schema_builder = None
+        self.env_vars = None
+        self.nb_instance_type = None
+        self.ram_usage_model_load = None
+        self.secret_key = None
+        self.jumpstart = None
+        self.role_arn = None
+
+    @abstractmethod
+    def _prepare_for_mode(self):
+        """Placeholder docstring"""
+
+    @abstractmethod
+    def _get_client_translators(self):
+        """Placeholder docstring"""
+
+    def _set_to_tgi(self):
+        """Placeholder docstring"""
+        if self.model_server != ModelServer.TGI:
+            messaging = (
+                "HuggingFace Model ID support on model server: "
+                f"{self.model_server} is not currently supported. "
+                f"Defaulting to {ModelServer.TGI}"
+            )
+            logger.warning(messaging)
+            self.model_server = ModelServer.TGI
+
+    def _create_tei_model(self, **kwargs) -> Type[Model]:
+        """Placeholder docstring"""
+        if self.nb_instance_type and "instance_type" not in kwargs:
+            kwargs.update({"instance_type": self.nb_instance_type})
+
+        if not self.image_uri:
+            self.image_uri = image_uris.retrieve(
+                "huggingface-tei",
+                image_scope="inference",
+                instance_type=kwargs.get("instance_type"),
+                region=self.sagemaker_session.boto_region_name,
+            )
+
+        pysdk_model = HuggingFaceModel(
+            image_uri=self.image_uri,
+            image_config=self.image_config,
+            vpc_config=self.vpc_config,
+            env=self.env_vars,
+            role=self.role_arn,
+            sagemaker_session=self.sagemaker_session,
+        )
+
+        logger.info("Detected %s. Proceeding with the the deployment.", self.image_uri)
+
+        self._original_deploy = pysdk_model.deploy
+        pysdk_model.deploy = self._tei_model_builder_deploy_wrapper
+        return pysdk_model
+
+    @_capture_telemetry("tei.deploy")
+    def _tei_model_builder_deploy_wrapper(self, *args, **kwargs) -> Type[PredictorBase]:
+        """Placeholder docstring"""
+        timeout = kwargs.get("model_data_download_timeout")
+        if timeout:
+            self.pysdk_model.env.update({"MODEL_LOADING_TIMEOUT": str(timeout)})
+
+        if "mode" in kwargs and kwargs.get("mode") != self.mode:
+            overwrite_mode = kwargs.get("mode")
+            # mode overwritten by customer during model.deploy()
+            logger.warning(
+                "Deploying in %s Mode, overriding existing configurations set for %s mode",
+                overwrite_mode,
+                self.mode,
+            )
+
+            if overwrite_mode == Mode.SAGEMAKER_ENDPOINT:
+                self.mode = self.pysdk_model.mode = Mode.SAGEMAKER_ENDPOINT
+            elif overwrite_mode == Mode.LOCAL_CONTAINER:
+                self._prepare_for_mode()
+                self.mode = self.pysdk_model.mode = Mode.LOCAL_CONTAINER
+            else:
+                raise ValueError("Mode %s is not supported!" % overwrite_mode)
+
+        serializer = self.schema_builder.input_serializer
+        deserializer = self.schema_builder._output_deserializer
+        if self.mode == Mode.LOCAL_CONTAINER:
+            timeout = kwargs.get("model_data_download_timeout")
+
+            predictor = TgiLocalModePredictor(
+                self.modes[str(Mode.LOCAL_CONTAINER)], serializer, deserializer
+            )
+
+            self.modes[str(Mode.LOCAL_CONTAINER)].create_server(
+                self.image_uri,
+                timeout if timeout else 1800,
+                None,
+                predictor,
+                self.pysdk_model.env,
+                jumpstart=False,
+            )
+
+            return predictor
+
+        if "mode" in kwargs:
+            del kwargs["mode"]
+        if "role" in kwargs:
+            self.pysdk_model.role = kwargs.get("role")
+            del kwargs["role"]
+
+        # set model_data to uncompressed s3 dict
+        self.pysdk_model.model_data, env_vars = self._prepare_for_mode()
+        self.env_vars.update(env_vars)
+        self.pysdk_model.env.update(self.env_vars)
+
+        # if the weights have been cached via local container mode -> set to offline
+        if str(Mode.LOCAL_CONTAINER) in self.modes:
+            self.pysdk_model.env.update({"TRANSFORMERS_OFFLINE": "1"})
+        else:
+            # if has not been built for local container we must use cache
+            # that hosting has write access to.
+            self.pysdk_model.env["TRANSFORMERS_CACHE"] = "/tmp"
+            self.pysdk_model.env["HUGGINGFACE_HUB_CACHE"] = "/tmp"
+
+        if "endpoint_logging" not in kwargs:
+            kwargs["endpoint_logging"] = True
+
+        if not self.nb_instance_type and "instance_type" not in kwargs:
+            raise ValueError(
+                "Instance type must be provided when deploying " "to SageMaker Endpoint mode."
+            )
+
+        if "initial_instance_count" not in kwargs:
+            kwargs.update({"initial_instance_count": 1})
+
+        predictor = self._original_deploy(*args, **kwargs)
+
+        predictor.serializer = serializer
+        predictor.deserializer = deserializer
+        return predictor
+
+    def _build_for_hf_tei(self):
+        """Placeholder docstring"""
+        self.nb_instance_type = _get_nb_instance()
+
+        _create_dir_structure(self.model_path)
+        if not hasattr(self, "pysdk_model"):
+            self.env_vars.update({"HF_MODEL_ID": self.model})
+            self.hf_model_config = _get_model_config_properties_from_hf(
+                self.model, self.env_vars.get("HUGGING_FACE_HUB_TOKEN")
+            )
+
+        self.pysdk_model = self._create_tei_model()
+
+        if self.mode == Mode.LOCAL_CONTAINER:
+            self._prepare_for_mode()
+
+        return self.pysdk_model
+
+    def _build_for_tei(self):
+        """Placeholder docstring"""
+        self.secret_key = None
+
+        self._set_to_tgi()
+
+        self.pysdk_model = self._build_for_hf_tei()
+        return self.pysdk_model
diff --git a/tests/integ/sagemaker/serve/test_serve_tei.py b/tests/integ/sagemaker/serve/test_serve_tei.py
new file mode 100644
index 0000000000..19ee0b57de
--- /dev/null
+++ b/tests/integ/sagemaker/serve/test_serve_tei.py
@@ -0,0 +1,123 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You
+# may not use this file except in compliance with the License. A copy of
+# the License is located at
+#
+#     http://aws.amazon.com/apache2.0/
+#
+# or in the "license" file accompanying this file. This file is
+# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
+# ANY KIND, either express or implied. See the License for the specific
+# language governing permissions and limitations under the License.
+from __future__ import absolute_import
+
+import pytest
+from sagemaker.serve.builder.schema_builder import SchemaBuilder
+from sagemaker.serve.builder.model_builder import ModelBuilder, Mode
+
+from tests.integ.sagemaker.serve.constants import (
+    HF_DIR,
+    PYTHON_VERSION_IS_NOT_310,
+    SERVE_SAGEMAKER_ENDPOINT_TIMEOUT,
+)
+
+from tests.integ.timeout import timeout
+from tests.integ.utils import cleanup_model_resources
+import logging
+
+logger = logging.getLogger(__name__)
+
+sample_input = {
+    "inputs": "The man worked as a [MASK].",
+}
+
+loaded_response = [
+    {
+        "score": 0.0974755585193634,
+        "token": 10533,
+        "token_str": "carpenter",
+        "sequence": "the man worked as a carpenter.",
+    },
+    {
+        "score": 0.052383411675691605,
+        "token": 15610,
+        "token_str": "waiter",
+        "sequence": "the man worked as a waiter.",
+    },
+    {
+        "score": 0.04962712526321411,
+        "token": 13362,
+        "token_str": "barber",
+        "sequence": "the man worked as a barber.",
+    },
+    {
+        "score": 0.0378861166536808,
+        "token": 15893,
+        "token_str": "mechanic",
+        "sequence": "the man worked as a mechanic.",
+    },
+    {
+        "score": 0.037680838257074356,
+        "token": 18968,
+        "token_str": "salesman",
+        "sequence": "the man worked as a salesman.",
+    },
+]
+
+
+@pytest.fixture
+def model_input():
+    return {"inputs": "The man worked as a [MASK]."}
+
+
+@pytest.fixture
+def model_builder_model_schema_builder():
+    return ModelBuilder(
+        model_path=HF_DIR,
+        model="BAAI/bge-m3",
+        schema_builder=SchemaBuilder(sample_input, loaded_response),
+        model_metadata={
+            "HF_TASK": "sentence-similarity",
+        },
+    )
+
+
+@pytest.fixture
+def model_builder(request):
+    return request.getfixturevalue(request.param)
+
+
+@pytest.mark.skipif(
+    PYTHON_VERSION_IS_NOT_310,
+    reason="Testing feature needs latest metadata",
+)
+@pytest.mark.parametrize("model_builder", ["model_builder_model_schema_builder"], indirect=True)
+def test_tei_sagemaker_endpoint(sagemaker_session, model_builder, model_input):
+    logger.info("Running in SAGEMAKER_ENDPOINT mode...")
+    caught_ex = None
+
+    iam_client = sagemaker_session.boto_session.client("iam")
+    role_arn = iam_client.get_role(RoleName="SageMakerRole")["Role"]["Arn"]
+
+    model = model_builder.build(
+        mode=Mode.SAGEMAKER_ENDPOINT, role_arn=role_arn, sagemaker_session=sagemaker_session
+    )
+
+    with timeout(minutes=SERVE_SAGEMAKER_ENDPOINT_TIMEOUT):
+        try:
+            logger.info("Deploying and predicting in SAGEMAKER_ENDPOINT mode...")
+            predictor = model.deploy(instance_type="ml.g5.2xlarge", initial_instance_count=1)
+            predictor.predict(model_input)
+            assert predictor is not None
+        except Exception as e:
+            caught_ex = e
+        finally:
+            cleanup_model_resources(
+                sagemaker_session=model_builder.sagemaker_session,
+                model_name=model.name,
+                endpoint_name=model.endpoint_name,
+            )
+            if caught_ex:
+                logger.exception(caught_ex)
+                assert False, f"{caught_ex} was thrown when running tei sagemaker endpoint test"
diff --git a/tests/unit/sagemaker/serve/builder/test_model_builder.py b/tests/unit/sagemaker/serve/builder/test_model_builder.py
index 3ffbdd7c03..0c06b5ae8e 100644
--- a/tests/unit/sagemaker/serve/builder/test_model_builder.py
+++ b/tests/unit/sagemaker/serve/builder/test_model_builder.py
@@ -1476,6 +1476,44 @@ def test_text_generation(
 
         mock_build_for_tgi.assert_called_once()
 
+    @patch("sagemaker.serve.builder.model_builder.ModelBuilder._build_for_tei")
+    @patch("sagemaker.image_uris.retrieve")
+    @patch("sagemaker.djl_inference.model.urllib")
+    @patch("sagemaker.djl_inference.model.json")
+    @patch("sagemaker.huggingface.llm_utils.urllib")
+    @patch("sagemaker.huggingface.llm_utils.json")
+    @patch("sagemaker.model_uris.retrieve")
+    @patch("sagemaker.serve.builder.model_builder._ServeSettings")
+    def test_sentence_similarity(
+        self,
+        mock_serveSettings,
+        mock_model_uris_retrieve,
+        mock_llm_utils_json,
+        mock_llm_utils_urllib,
+        mock_model_json,
+        mock_model_urllib,
+        mock_image_uris_retrieve,
+        mock_build_for_tei,
+    ):
+        mock_setting_object = mock_serveSettings.return_value
+        mock_setting_object.role_arn = mock_role_arn
+        mock_setting_object.s3_model_data_url = mock_s3_model_data_url
+
+        mock_model_uris_retrieve.side_effect = KeyError
+        mock_llm_utils_json.load.return_value = {"pipeline_tag": "sentence-similarity"}
+        mock_llm_utils_urllib.request.Request.side_effect = Mock()
+
+        mock_model_json.load.return_value = {"some": "config"}
+        mock_model_urllib.request.Request.side_effect = Mock()
+        mock_build_for_tei.side_effect = Mock()
+
+        mock_image_uris_retrieve.return_value = "https://some-image-uri"
+
+        model_builder = ModelBuilder(model="bloom-560m", schema_builder=schema_builder)
+        model_builder.build(sagemaker_session=mock_session)
+
+        mock_build_for_tei.assert_called_once()
+
     @patch("sagemaker.serve.builder.model_builder.ModelBuilder._build_for_transformers", Mock())
     @patch("sagemaker.serve.builder.model_builder.ModelBuilder._try_fetch_gpu_info")
     @patch("sagemaker.image_uris.retrieve")
diff --git a/tests/unit/sagemaker/serve/builder/test_tei_builder.py b/tests/unit/sagemaker/serve/builder/test_tei_builder.py
new file mode 100644
index 0000000000..79a8f23324
--- /dev/null
+++ b/tests/unit/sagemaker/serve/builder/test_tei_builder.py
@@ -0,0 +1,152 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You
+# may not use this file except in compliance with the License. A copy of
+# the License is located at
+#
+#     http://aws.amazon.com/apache2.0/
+#
+# or in the "license" file accompanying this file. This file is
+# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
+# ANY KIND, either express or implied. See the License for the specific
+# language governing permissions and limitations under the License.
+from __future__ import absolute_import
+from unittest.mock import MagicMock, patch
+
+import unittest
+from sagemaker.serve.builder.model_builder import ModelBuilder
+from sagemaker.serve.mode.function_pointers import Mode
+from tests.unit.sagemaker.serve.constants import MOCK_VPC_CONFIG
+
+from sagemaker.serve.utils.predictors import TgiLocalModePredictor
+
+mock_model_id = "bert-base-uncased"
+mock_prompt = "The man worked as a [MASK]."
+mock_sample_input = {"inputs": mock_prompt}
+mock_sample_output = [
+    {
+        "score": 0.0974755585193634,
+        "token": 10533,
+        "token_str": "carpenter",
+        "sequence": "the man worked as a carpenter.",
+    },
+    {
+        "score": 0.052383411675691605,
+        "token": 15610,
+        "token_str": "waiter",
+        "sequence": "the man worked as a waiter.",
+    },
+    {
+        "score": 0.04962712526321411,
+        "token": 13362,
+        "token_str": "barber",
+        "sequence": "the man worked as a barber.",
+    },
+    {
+        "score": 0.0378861166536808,
+        "token": 15893,
+        "token_str": "mechanic",
+        "sequence": "the man worked as a mechanic.",
+    },
+    {
+        "score": 0.037680838257074356,
+        "token": 18968,
+        "token_str": "salesman",
+        "sequence": "the man worked as a salesman.",
+    },
+]
+mock_schema_builder = MagicMock()
+mock_schema_builder.sample_input = mock_sample_input
+mock_schema_builder.sample_output = mock_sample_output
+MOCK_IMAGE_CONFIG = (
+    "763104351884.dkr.ecr.us-west-2.amazonaws.com/"
+    "huggingface-pytorch-inference:2.0.0-transformers4.28.1-gpu-py310-cu118-ubuntu20.04-v1.0"
+)
+
+
+class TestTEIBuilder(unittest.TestCase):
+    @patch(
+        "sagemaker.serve.builder.tei_builder._get_nb_instance",
+        return_value="ml.g5.24xlarge",
+    )
+    @patch("sagemaker.serve.builder.tei_builder._capture_telemetry", side_effect=None)
+    def test_build_deploy_for_tei_local_container_and_remote_container(
+        self,
+        mock_get_nb_instance,
+        mock_telemetry,
+    ):
+        builder = ModelBuilder(
+            model=mock_model_id,
+            schema_builder=mock_schema_builder,
+            mode=Mode.LOCAL_CONTAINER,
+            vpc_config=MOCK_VPC_CONFIG,
+            model_metadata={
+                "HF_TASK": "sentence-similarity",
+            },
+        )
+
+        builder._prepare_for_mode = MagicMock()
+        builder._prepare_for_mode.side_effect = None
+
+        model = builder.build()
+        builder.serve_settings.telemetry_opt_out = True
+
+        builder.modes[str(Mode.LOCAL_CONTAINER)] = MagicMock()
+        predictor = model.deploy(model_data_download_timeout=1800)
+
+        assert model.vpc_config == MOCK_VPC_CONFIG
+        assert builder.env_vars["MODEL_LOADING_TIMEOUT"] == "1800"
+        assert isinstance(predictor, TgiLocalModePredictor)
+
+        assert builder.nb_instance_type == "ml.g5.24xlarge"
+
+        builder._original_deploy = MagicMock()
+        builder._prepare_for_mode.return_value = (None, {})
+        predictor = model.deploy(mode=Mode.SAGEMAKER_ENDPOINT, role="mock_role_arn")
+        assert "HF_MODEL_ID" in model.env
+
+        with self.assertRaises(ValueError) as _:
+            model.deploy(mode=Mode.IN_PROCESS)
+
+    @patch(
+        "sagemaker.serve.builder.tei_builder._get_nb_instance",
+        return_value="ml.g5.24xlarge",
+    )
+    @patch("sagemaker.serve.builder.tei_builder._capture_telemetry", side_effect=None)
+    def test_image_uri_override(
+        self,
+        mock_get_nb_instance,
+        mock_telemetry,
+    ):
+        builder = ModelBuilder(
+            model=mock_model_id,
+            schema_builder=mock_schema_builder,
+            mode=Mode.LOCAL_CONTAINER,
+            image_uri=MOCK_IMAGE_CONFIG,
+            model_metadata={
+                "HF_TASK": "sentence-similarity",
+            },
+        )
+
+        builder._prepare_for_mode = MagicMock()
+        builder._prepare_for_mode.side_effect = None
+
+        model = builder.build()
+        builder.serve_settings.telemetry_opt_out = True
+
+        builder.modes[str(Mode.LOCAL_CONTAINER)] = MagicMock()
+        predictor = model.deploy(model_data_download_timeout=1800)
+
+        assert builder.image_uri == MOCK_IMAGE_CONFIG
+        assert builder.env_vars["MODEL_LOADING_TIMEOUT"] == "1800"
+        assert isinstance(predictor, TgiLocalModePredictor)
+
+        assert builder.nb_instance_type == "ml.g5.24xlarge"
+
+        builder._original_deploy = MagicMock()
+        builder._prepare_for_mode.return_value = (None, {})
+        predictor = model.deploy(mode=Mode.SAGEMAKER_ENDPOINT, role="mock_role_arn")
+        assert "HF_MODEL_ID" in model.env
+
+        with self.assertRaises(ValueError) as _:
+            model.deploy(mode=Mode.IN_PROCESS)
diff --git a/tests/unit/sagemaker/serve/builder/test_transformers_builder.py b/tests/unit/sagemaker/serve/builder/test_transformers_builder.py
index b7e3db79d6..d63eabf2a3 100644
--- a/tests/unit/sagemaker/serve/builder/test_transformers_builder.py
+++ b/tests/unit/sagemaker/serve/builder/test_transformers_builder.py
@@ -110,7 +110,7 @@ def test_build_deploy_for_transformers_local_container_and_remote_container(
         return_value="ml.g5.24xlarge",
     )
     @patch("sagemaker.serve.builder.transformers_builder._capture_telemetry", side_effect=None)
-    def test_image_uri(
+    def test_image_uri_override(
         self,
         mock_get_nb_instance,
         mock_telemetry,

From 9b7874b3246800bc620e09f8f78370230e7fd3d3 Mon Sep 17 00:00:00 2001
From: Samrudhi Sharma <154457034+samruds@users.noreply.github.com>
Date: Mon, 20 May 2024 09:25:31 -0700
Subject: [PATCH 40/58] Fix: Add Image URI overrides for transformers models
 (#4693)

* Fix: Add Image URI overrides for transformers models

* Increase coverage

* Fix formatting
---
 .../serve/builder/transformers_builder.py     | 28 +++++++++++++++----
 .../serve/test_serve_transformers.py          |  2 +-
 .../builder/test_transformers_builder.py      | 26 +++++++++++++++++
 3 files changed, 50 insertions(+), 6 deletions(-)

diff --git a/src/sagemaker/serve/builder/transformers_builder.py b/src/sagemaker/serve/builder/transformers_builder.py
index ead9b7425f..614290b132 100644
--- a/src/sagemaker/serve/builder/transformers_builder.py
+++ b/src/sagemaker/serve/builder/transformers_builder.py
@@ -78,6 +78,25 @@ def _prepare_for_mode(self):
         """Abstract method"""
 
     def _create_transformers_model(self) -> Type[Model]:
+        """Initializes HF model with or without image_uri"""
+        if self.image_uri is None:
+            pysdk_model = self._get_hf_metadata_create_model()
+        else:
+            pysdk_model = HuggingFaceModel(
+                image_uri=self.image_uri,
+                vpc_config=self.vpc_config,
+                env=self.env_vars,
+                role=self.role_arn,
+                sagemaker_session=self.sagemaker_session,
+            )
+
+        logger.info("Detected %s. Proceeding with the the deployment.", self.image_uri)
+
+        self._original_deploy = pysdk_model.deploy
+        pysdk_model.deploy = self._transformers_model_builder_deploy_wrapper
+        return pysdk_model
+
+    def _get_hf_metadata_create_model(self) -> Type[Model]:
         """Initializes the model after fetching image
 
         1. Get the metadata for deciding framework
@@ -132,22 +151,21 @@ def _create_transformers_model(self) -> Type[Model]:
                 vpc_config=self.vpc_config,
             )
 
-        if not self.image_uri and self.mode == Mode.LOCAL_CONTAINER:
+        if self.mode == Mode.LOCAL_CONTAINER:
             self.image_uri = pysdk_model.serving_image_uri(
                 self.sagemaker_session.boto_region_name, "local"
             )
-        elif not self.image_uri:
+        else:
             self.image_uri = pysdk_model.serving_image_uri(
                 self.sagemaker_session.boto_region_name, self.instance_type
             )
 
-        logger.info("Detected %s. Proceeding with the the deployment.", self.image_uri)
+        if pysdk_model is None or self.image_uri is None:
+            raise ValueError("PySDK model unable to be created, try overriding image_uri")
 
         if not pysdk_model.image_uri:
             pysdk_model.image_uri = self.image_uri
 
-        self._original_deploy = pysdk_model.deploy
-        pysdk_model.deploy = self._transformers_model_builder_deploy_wrapper
         return pysdk_model
 
     @_capture_telemetry("transformers.deploy")
diff --git a/tests/integ/sagemaker/serve/test_serve_transformers.py b/tests/integ/sagemaker/serve/test_serve_transformers.py
index 64029f7290..33a1ae6708 100644
--- a/tests/integ/sagemaker/serve/test_serve_transformers.py
+++ b/tests/integ/sagemaker/serve/test_serve_transformers.py
@@ -127,4 +127,4 @@ def test_pytorch_transformers_sagemaker_endpoint(
                 logger.exception(caught_ex)
                 assert (
                     False
-                ), f"{caught_ex} was thrown when running pytorch transformers sagemaker endpoint test"
+                ), f"{caught_ex} thrown when running pytorch transformers sagemaker endpoint test"
diff --git a/tests/unit/sagemaker/serve/builder/test_transformers_builder.py b/tests/unit/sagemaker/serve/builder/test_transformers_builder.py
index d63eabf2a3..9ea797adc2 100644
--- a/tests/unit/sagemaker/serve/builder/test_transformers_builder.py
+++ b/tests/unit/sagemaker/serve/builder/test_transformers_builder.py
@@ -144,3 +144,29 @@ def test_image_uri_override(
 
         with self.assertRaises(ValueError) as _:
             model.deploy(mode=Mode.IN_PROCESS)
+
+    @patch("sagemaker.serve.builder.model_builder.ModelBuilder._build_for_transformers")
+    @patch(
+        "sagemaker.serve.builder.transformers_builder._get_nb_instance",
+        return_value="ml.g5.24xlarge",
+    )
+    @patch("sagemaker.serve.builder.transformers_builder._capture_telemetry", side_effect=None)
+    @patch(
+        "sagemaker.huggingface.llm_utils.get_huggingface_model_metadata",
+        return_value=None,
+    )
+    def test_failure_hf_md(
+        self, mock_model_md, mock_get_nb_instance, mock_telemetry, mock_build_for_transformers
+    ):
+        builder = ModelBuilder(
+            model=mock_model_id,
+            schema_builder=mock_schema_builder,
+            mode=Mode.LOCAL_CONTAINER,
+        )
+
+        builder._prepare_for_mode = MagicMock()
+        builder._prepare_for_mode.side_effect = None
+
+        builder.build()
+
+        mock_build_for_transformers.assert_called_once()

From a26224a09fe75d1c6babae965af643767002c061 Mon Sep 17 00:00:00 2001
From: ci <ci>
Date: Mon, 20 May 2024 20:51:00 +0000
Subject: [PATCH 41/58] prepare release v2.221.0

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

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 5d85261dd7..9261d394e8 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,5 +1,19 @@
 # Changelog
 
+## v2.221.0 (2024-05-20)
+
+### Features
+
+ * onboard tei image config to pysdk
+
+### Bug Fixes and Other Changes
+
+ * JS Model with non-TGI/non-DJL deployment failure
+ * cover tei with image_uris.retrieve API
+ * Add more debuging
+ * model builder limited container support for endpoint mode.
+ * Image URI should take precedence for HF models
+
 ## v2.220.0 (2024-05-15)
 
 ### Features
diff --git a/VERSION b/VERSION
index bd29d7bb5b..1155bf691b 100644
--- a/VERSION
+++ b/VERSION
@@ -1 +1 @@
-2.220.1.dev0
+2.221.0

From 4e83ccee59e4cc050ca8a9db54d6336705f75b43 Mon Sep 17 00:00:00 2001
From: ci <ci>
Date: Mon, 20 May 2024 20:51:02 +0000
Subject: [PATCH 42/58] update development version to v2.221.1.dev0

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

diff --git a/VERSION b/VERSION
index 1155bf691b..1ab82bbae0 100644
--- a/VERSION
+++ b/VERSION
@@ -1 +1 @@
-2.221.0
+2.221.1.dev0

From 828cdc35ebc02050ec95c63edfc510440ec1ac26 Mon Sep 17 00:00:00 2001
From: Haixin Wang <98612668+haixiw@users.noreply.github.com>
Date: Tue, 21 May 2024 11:15:28 -0700
Subject: [PATCH 43/58] Add tei cpu image (#4695)

* Add tei cpu image

* fix format issue

* fix unit tests

* fix typo

* fix typo
---
 src/sagemaker/huggingface/llm_utils.py        |  7 +++
 .../image_uri_config/huggingface-tei-cpu.json | 59 +++++++++++++++++++
 src/sagemaker/image_uris.py                   |  6 +-
 .../image_uris/test_huggingface_llm.py        | 12 +++-
 4 files changed, 79 insertions(+), 5 deletions(-)
 create mode 100644 src/sagemaker/image_uri_config/huggingface-tei-cpu.json

diff --git a/src/sagemaker/huggingface/llm_utils.py b/src/sagemaker/huggingface/llm_utils.py
index 974cffcddf..9927d1d293 100644
--- a/src/sagemaker/huggingface/llm_utils.py
+++ b/src/sagemaker/huggingface/llm_utils.py
@@ -72,6 +72,13 @@ def get_huggingface_llm_image_uri(
             version=version,
             image_scope="inference",
         )
+    if backend == "huggingface-tei-cpu":
+        return image_uris.retrieve(
+            "huggingface-tei-cpu",
+            region=region,
+            version=version,
+            image_scope="inference",
+        )
     if backend == "lmi":
         version = version or "0.24.0"
         return image_uris.retrieve(framework="djl-deepspeed", region=region, version=version)
diff --git a/src/sagemaker/image_uri_config/huggingface-tei-cpu.json b/src/sagemaker/image_uri_config/huggingface-tei-cpu.json
new file mode 100644
index 0000000000..d68b0d6307
--- /dev/null
+++ b/src/sagemaker/image_uri_config/huggingface-tei-cpu.json
@@ -0,0 +1,59 @@
+{
+    "inference": {
+        "processors": [
+            "cpu"
+        ],
+        "version_aliases": {
+            "1.2": "1.2.3"
+        },
+        "versions": {
+            "1.2.3": {
+                "py_versions": [
+                    "py310"
+                ],
+                "registries": {
+                    "af-south-1": "510948584623",
+                    "ap-east-1": "651117190479",
+                    "ap-northeast-1": "354813040037",
+                    "ap-northeast-2": "366743142698",
+                    "ap-northeast-3": "867004704886",
+                    "ap-south-1": "720646828776",
+                    "ap-south-2": "628508329040",
+                    "ap-southeast-1": "121021644041",
+                    "ap-southeast-2": "783357654285",
+                    "ap-southeast-3": "951798379941",
+                    "ap-southeast-4": "106583098589",
+                    "ca-central-1": "341280168497",
+                    "ca-west-1": "190319476487",
+                    "cn-north-1": "450853457545",
+                    "cn-northwest-1": "451049120500",
+                    "eu-central-1": "492215442770",
+                    "eu-central-2": "680994064768",
+                    "eu-north-1": "662702820516",
+                    "eu-south-1": "978288397137",
+                    "eu-south-2": "104374241257",
+                    "eu-west-1": "141502667606",
+                    "eu-west-2": "764974769150",
+                    "eu-west-3": "659782779980",
+                    "il-central-1": "898809789911",
+                    "me-central-1": "272398656194",
+                    "me-south-1": "801668240914",
+                    "sa-east-1": "737474898029",
+                    "us-east-1": "683313688378",
+                    "us-east-2": "257758044811",
+                    "us-gov-east-1": "237065988967",
+                    "us-gov-west-1": "414596584902",
+                    "us-iso-east-1": "833128469047",
+                    "us-isob-east-1": "281123927165",
+                    "us-west-1": "746614075791",
+                    "us-west-2": "246618743249"
+                },
+                "tag_prefix": "2.0.1-tei1.2.3",
+                "repository": "tei-cpu",
+                "container_version": {
+                    "cpu": "ubuntu22.04"
+                }
+            }
+        }
+    }
+}
\ No newline at end of file
diff --git a/src/sagemaker/image_uris.py b/src/sagemaker/image_uris.py
index 1ca73d0af9..be5167dcc7 100644
--- a/src/sagemaker/image_uris.py
+++ b/src/sagemaker/image_uris.py
@@ -37,7 +37,8 @@
 ECR_URI_TEMPLATE = "{registry}.dkr.{hostname}/{repository}"
 HUGGING_FACE_FRAMEWORK = "huggingface"
 HUGGING_FACE_LLM_FRAMEWORK = "huggingface-llm"
-HUGGING_FACE_TEI_FRAMEWORK = "huggingface-tei"
+HUGGING_FACE_TEI_GPU_FRAMEWORK = "huggingface-tei"
+HUGGING_FACE_TEI_CPU_FRAMEWORK = "huggingface-tei-cpu"
 HUGGING_FACE_LLM_NEURONX_FRAMEWORK = "huggingface-llm-neuronx"
 XGBOOST_FRAMEWORK = "xgboost"
 SKLEARN_FRAMEWORK = "sklearn"
@@ -478,7 +479,8 @@ def _validate_version_and_set_if_needed(version, config, framework):
     if version is None and framework in [
         DATA_WRANGLER_FRAMEWORK,
         HUGGING_FACE_LLM_FRAMEWORK,
-        HUGGING_FACE_TEI_FRAMEWORK,
+        HUGGING_FACE_TEI_GPU_FRAMEWORK,
+        HUGGING_FACE_TEI_CPU_FRAMEWORK,
         HUGGING_FACE_LLM_NEURONX_FRAMEWORK,
         STABILITYAI_FRAMEWORK,
     ]:
diff --git a/tests/unit/sagemaker/image_uris/test_huggingface_llm.py b/tests/unit/sagemaker/image_uris/test_huggingface_llm.py
index b1e8e8253e..fa10fd24fe 100644
--- a/tests/unit/sagemaker/image_uris/test_huggingface_llm.py
+++ b/tests/unit/sagemaker/image_uris/test_huggingface_llm.py
@@ -22,6 +22,9 @@
     "gpu": {
         "1.2.3": "2.0.1-tei1.2.3-gpu-py310-cu122-ubuntu22.04",
     },
+    "cpu": {
+        "1.2.3": "2.0.1-tei1.2.3-cpu-py310-ubuntu22.04",
+    },
 }
 HF_VERSIONS_MAPPING = {
     "gpu": {
@@ -73,17 +76,20 @@ def test_huggingface_uris(load_config):
             assert expected == uri
 
 
-@pytest.mark.parametrize("load_config", ["huggingface-tei.json"], indirect=True)
+@pytest.mark.parametrize(
+    "load_config", ["huggingface-tei.json", "huggingface-tei-cpu.json"], indirect=True
+)
 def test_huggingface_tei_uris(load_config):
     VERSIONS = load_config["inference"]["versions"]
     device = load_config["inference"]["processors"][0]
-    backend = "huggingface-tei"
+    backend = "huggingface-tei" if device == "gpu" else "huggingface-tei-cpu"
+    repo = "tei" if device == "gpu" else "tei-cpu"
     for version in VERSIONS:
         ACCOUNTS = load_config["inference"]["versions"][version]["registries"]
         for region in ACCOUNTS.keys():
             uri = get_huggingface_llm_image_uri(backend, region=region, version=version)
             expected = expected_uris.huggingface_llm_framework_uri(
-                "tei",
+                repo,
                 ACCOUNTS[region],
                 version,
                 TEI_VERSIONS_MAPPING[device][version],

From 18e76c7751ba377bf22e8b81116ebe3f1070aaf8 Mon Sep 17 00:00:00 2001
From: Jonathan Makunga <54963715+makungaj1@users.noreply.github.com>
Date: Tue, 21 May 2024 14:25:38 -0700
Subject: [PATCH 44/58] Feat: Add TEI support for ModelBuilder (#4694)

* Add TEI Serving

* Add TEI Serving

* Add TEI Serving

* Add TEI Serving

* Add TEI Serving

* Add TEI Serving

* Notebook testing

* Notebook testing

* Notebook testing

* Refactoring

* Refactoring

* UT

* UT

* Refactoring

* Test coverage

* Refactoring

* Refactoring

---------

Co-authored-by: Jonathan Makunga <makung@amazon.com>
---
 src/sagemaker/serve/builder/model_builder.py  |   2 +-
 src/sagemaker/serve/builder/tei_builder.py    |  18 +-
 .../serve/mode/local_container_mode.py        |  15 ++
 .../serve/mode/sagemaker_endpoint_mode.py     |  27 ++-
 .../serve/model_server/tei/__init__.py        |   0
 .../serve/model_server/tei/server.py          | 160 ++++++++++++++++++
 src/sagemaker/serve/utils/predictors.py       |  43 +++++
 src/sagemaker/serve/utils/telemetry_logger.py |   1 +
 src/sagemaker/serve/utils/types.py            |   1 +
 tests/integ/sagemaker/serve/test_serve_tei.py |  44 +----
 .../serve/builder/test_tei_builder.py         |   6 +-
 .../serve/model_server/tei/__init__.py        |   0
 .../serve/model_server/tei/test_server.py     | 150 ++++++++++++++++
 13 files changed, 409 insertions(+), 58 deletions(-)
 create mode 100644 src/sagemaker/serve/model_server/tei/__init__.py
 create mode 100644 src/sagemaker/serve/model_server/tei/server.py
 create mode 100644 tests/unit/sagemaker/serve/model_server/tei/__init__.py
 create mode 100644 tests/unit/sagemaker/serve/model_server/tei/test_server.py

diff --git a/src/sagemaker/serve/builder/model_builder.py b/src/sagemaker/serve/builder/model_builder.py
index b24f30fb3e..1fe75065d5 100644
--- a/src/sagemaker/serve/builder/model_builder.py
+++ b/src/sagemaker/serve/builder/model_builder.py
@@ -169,7 +169,7 @@ class ModelBuilder(Triton, DJL, JumpStart, TGI, Transformers, TensorflowServing,
             in order for model builder to build the artifacts correctly (according
             to the model server). Possible values for this argument are
             ``TORCHSERVE``, ``MMS``, ``TENSORFLOW_SERVING``, ``DJL_SERVING``,
-            ``TRITON``, and``TGI``.
+            ``TRITON``,``TGI``, and ``TEI``.
         model_metadata (Optional[Dict[str, Any]): Dictionary used to override model metadata.
             Currently, ``HF_TASK`` is overridable for HuggingFace model. HF_TASK should be set for
             new models without task metadata in the Hub, adding unsupported task types will throw
diff --git a/src/sagemaker/serve/builder/tei_builder.py b/src/sagemaker/serve/builder/tei_builder.py
index 50d3866468..6aba3c9da2 100644
--- a/src/sagemaker/serve/builder/tei_builder.py
+++ b/src/sagemaker/serve/builder/tei_builder.py
@@ -25,7 +25,7 @@
     _get_nb_instance,
 )
 from sagemaker.serve.model_server.tgi.prepare import _create_dir_structure
-from sagemaker.serve.utils.predictors import TgiLocalModePredictor
+from sagemaker.serve.utils.predictors import TeiLocalModePredictor
 from sagemaker.serve.utils.types import ModelServer
 from sagemaker.serve.mode.function_pointers import Mode
 from sagemaker.serve.utils.telemetry_logger import _capture_telemetry
@@ -74,16 +74,16 @@ def _prepare_for_mode(self):
     def _get_client_translators(self):
         """Placeholder docstring"""
 
-    def _set_to_tgi(self):
+    def _set_to_tei(self):
         """Placeholder docstring"""
-        if self.model_server != ModelServer.TGI:
+        if self.model_server != ModelServer.TEI:
             messaging = (
                 "HuggingFace Model ID support on model server: "
                 f"{self.model_server} is not currently supported. "
-                f"Defaulting to {ModelServer.TGI}"
+                f"Defaulting to {ModelServer.TEI}"
             )
             logger.warning(messaging)
-            self.model_server = ModelServer.TGI
+            self.model_server = ModelServer.TEI
 
     def _create_tei_model(self, **kwargs) -> Type[Model]:
         """Placeholder docstring"""
@@ -142,7 +142,7 @@ def _tei_model_builder_deploy_wrapper(self, *args, **kwargs) -> Type[PredictorBa
         if self.mode == Mode.LOCAL_CONTAINER:
             timeout = kwargs.get("model_data_download_timeout")
 
-            predictor = TgiLocalModePredictor(
+            predictor = TeiLocalModePredictor(
                 self.modes[str(Mode.LOCAL_CONTAINER)], serializer, deserializer
             )
 
@@ -180,7 +180,9 @@ def _tei_model_builder_deploy_wrapper(self, *args, **kwargs) -> Type[PredictorBa
         if "endpoint_logging" not in kwargs:
             kwargs["endpoint_logging"] = True
 
-        if not self.nb_instance_type and "instance_type" not in kwargs:
+        if self.nb_instance_type and "instance_type" not in kwargs:
+            kwargs.update({"instance_type": self.nb_instance_type})
+        elif not self.nb_instance_type and "instance_type" not in kwargs:
             raise ValueError(
                 "Instance type must be provided when deploying " "to SageMaker Endpoint mode."
             )
@@ -216,7 +218,7 @@ def _build_for_tei(self):
         """Placeholder docstring"""
         self.secret_key = None
 
-        self._set_to_tgi()
+        self._set_to_tei()
 
         self.pysdk_model = self._build_for_hf_tei()
         return self.pysdk_model
diff --git a/src/sagemaker/serve/mode/local_container_mode.py b/src/sagemaker/serve/mode/local_container_mode.py
index f940e2959c..f040c61c1d 100644
--- a/src/sagemaker/serve/mode/local_container_mode.py
+++ b/src/sagemaker/serve/mode/local_container_mode.py
@@ -21,6 +21,7 @@
 from sagemaker.serve.model_server.djl_serving.server import LocalDJLServing
 from sagemaker.serve.model_server.triton.server import LocalTritonServer
 from sagemaker.serve.model_server.tgi.server import LocalTgiServing
+from sagemaker.serve.model_server.tei.server import LocalTeiServing
 from sagemaker.serve.model_server.multi_model_server.server import LocalMultiModelServer
 from sagemaker.session import Session
 
@@ -69,6 +70,7 @@ def __init__(
         self.container = None
         self.secret_key = None
         self._ping_container = None
+        self._invoke_serving = None
 
     def load(self, model_path: str = None):
         """Placeholder docstring"""
@@ -156,6 +158,19 @@ def create_server(
                 env_vars=env_vars if env_vars else self.env_vars,
             )
             self._ping_container = self._tensorflow_serving_deep_ping
+        elif self.model_server == ModelServer.TEI:
+            tei_serving = LocalTeiServing()
+            tei_serving._start_tei_serving(
+                client=self.client,
+                image=image,
+                model_path=model_path if model_path else self.model_path,
+                secret_key=secret_key,
+                env_vars=env_vars if env_vars else self.env_vars,
+            )
+            tei_serving.schema_builder = self.schema_builder
+            self.container = tei_serving.container
+            self._ping_container = tei_serving._tei_deep_ping
+            self._invoke_serving = tei_serving._invoke_tei_serving
 
         # allow some time for container to be ready
         time.sleep(10)
diff --git a/src/sagemaker/serve/mode/sagemaker_endpoint_mode.py b/src/sagemaker/serve/mode/sagemaker_endpoint_mode.py
index 24acfc6a2f..b8f1d0529b 100644
--- a/src/sagemaker/serve/mode/sagemaker_endpoint_mode.py
+++ b/src/sagemaker/serve/mode/sagemaker_endpoint_mode.py
@@ -6,6 +6,7 @@
 import logging
 from typing import Type
 
+from sagemaker.serve.model_server.tei.server import SageMakerTeiServing
 from sagemaker.serve.model_server.tensorflow_serving.server import SageMakerTensorflowServing
 from sagemaker.session import Session
 from sagemaker.serve.utils.types import ModelServer
@@ -37,6 +38,8 @@ def __init__(self, inference_spec: Type[InferenceSpec], model_server: ModelServe
         self.inference_spec = inference_spec
         self.model_server = model_server
 
+        self._tei_serving = SageMakerTeiServing()
+
     def load(self, model_path: str):
         """Placeholder docstring"""
         path = Path(model_path)
@@ -66,8 +69,9 @@ def prepare(
                 + "session to be created or supply `sagemaker_session` into @serve.invoke."
             ) from e
 
+        upload_artifacts = None
         if self.model_server == ModelServer.TORCHSERVE:
-            return self._upload_torchserve_artifacts(
+            upload_artifacts = self._upload_torchserve_artifacts(
                 model_path=model_path,
                 sagemaker_session=sagemaker_session,
                 secret_key=secret_key,
@@ -76,7 +80,7 @@ def prepare(
             )
 
         if self.model_server == ModelServer.TRITON:
-            return self._upload_triton_artifacts(
+            upload_artifacts = self._upload_triton_artifacts(
                 model_path=model_path,
                 sagemaker_session=sagemaker_session,
                 secret_key=secret_key,
@@ -85,7 +89,7 @@ def prepare(
             )
 
         if self.model_server == ModelServer.DJL_SERVING:
-            return self._upload_djl_artifacts(
+            upload_artifacts = self._upload_djl_artifacts(
                 model_path=model_path,
                 sagemaker_session=sagemaker_session,
                 s3_model_data_url=s3_model_data_url,
@@ -93,7 +97,7 @@ def prepare(
             )
 
         if self.model_server == ModelServer.TGI:
-            return self._upload_tgi_artifacts(
+            upload_artifacts = self._upload_tgi_artifacts(
                 model_path=model_path,
                 sagemaker_session=sagemaker_session,
                 s3_model_data_url=s3_model_data_url,
@@ -102,7 +106,7 @@ def prepare(
             )
 
         if self.model_server == ModelServer.MMS:
-            return self._upload_server_artifacts(
+            upload_artifacts = self._upload_server_artifacts(
                 model_path=model_path,
                 sagemaker_session=sagemaker_session,
                 s3_model_data_url=s3_model_data_url,
@@ -110,7 +114,7 @@ def prepare(
             )
 
         if self.model_server == ModelServer.TENSORFLOW_SERVING:
-            return self._upload_tensorflow_serving_artifacts(
+            upload_artifacts = self._upload_tensorflow_serving_artifacts(
                 model_path=model_path,
                 sagemaker_session=sagemaker_session,
                 secret_key=secret_key,
@@ -118,4 +122,15 @@ def prepare(
                 image=image,
             )
 
+        if self.model_server == ModelServer.TEI:
+            upload_artifacts = self._tei_serving._upload_tei_artifacts(
+                model_path=model_path,
+                sagemaker_session=sagemaker_session,
+                s3_model_data_url=s3_model_data_url,
+                image=image,
+            )
+
+        if upload_artifacts:
+            return upload_artifacts
+
         raise ValueError("%s model server is not supported" % self.model_server)
diff --git a/src/sagemaker/serve/model_server/tei/__init__.py b/src/sagemaker/serve/model_server/tei/__init__.py
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/src/sagemaker/serve/model_server/tei/server.py b/src/sagemaker/serve/model_server/tei/server.py
new file mode 100644
index 0000000000..67fca0e847
--- /dev/null
+++ b/src/sagemaker/serve/model_server/tei/server.py
@@ -0,0 +1,160 @@
+"""Module for Local TEI Serving"""
+
+from __future__ import absolute_import
+
+import requests
+import logging
+from pathlib import Path
+from docker.types import DeviceRequest
+from sagemaker import Session, fw_utils
+from sagemaker.serve.utils.exceptions import LocalModelInvocationException
+from sagemaker.base_predictor import PredictorBase
+from sagemaker.s3_utils import determine_bucket_and_prefix, parse_s3_url, s3_path_join
+from sagemaker.s3 import S3Uploader
+from sagemaker.local.utils import get_docker_host
+
+
+MODE_DIR_BINDING = "/opt/ml/model/"
+_SHM_SIZE = "2G"
+_DEFAULT_ENV_VARS = {
+    "TRANSFORMERS_CACHE": "/opt/ml/model/",
+    "HUGGINGFACE_HUB_CACHE": "/opt/ml/model/",
+}
+
+logger = logging.getLogger(__name__)
+
+
+class LocalTeiServing:
+    """LocalTeiServing class"""
+
+    def _start_tei_serving(
+        self, client: object, image: str, model_path: str, secret_key: str, env_vars: dict
+    ):
+        """Starts a local tei serving container.
+
+        Args:
+            client: Docker client
+            image: Image to use
+            model_path: Path to the model
+            secret_key: Secret key to use for authentication
+            env_vars: Environment variables to set
+        """
+        if env_vars and secret_key:
+            env_vars["SAGEMAKER_SERVE_SECRET_KEY"] = secret_key
+
+        self.container = client.containers.run(
+            image,
+            shm_size=_SHM_SIZE,
+            device_requests=[DeviceRequest(count=-1, capabilities=[["gpu"]])],
+            network_mode="host",
+            detach=True,
+            auto_remove=True,
+            volumes={
+                Path(model_path).joinpath("code"): {
+                    "bind": MODE_DIR_BINDING,
+                    "mode": "rw",
+                },
+            },
+            environment=_update_env_vars(env_vars),
+        )
+
+    def _invoke_tei_serving(self, request: object, content_type: str, accept: str):
+        """Invokes a local tei serving container.
+
+        Args:
+            request: Request to send
+            content_type: Content type to use
+            accept: Accept to use
+        """
+        try:
+            response = requests.post(
+                f"http://{get_docker_host()}:8080/invocations",
+                data=request,
+                headers={"Content-Type": content_type, "Accept": accept},
+                timeout=600,
+            )
+            response.raise_for_status()
+            return response.content
+        except Exception as e:
+            raise Exception("Unable to send request to the local container server") from e
+
+    def _tei_deep_ping(self, predictor: PredictorBase):
+        """Checks if the local tei serving container is up and running.
+
+        If the container is not up and running, it will raise an exception.
+        """
+        response = None
+        try:
+            response = predictor.predict(self.schema_builder.sample_input)
+            return (True, response)
+            # pylint: disable=broad-except
+        except Exception as e:
+            if "422 Client Error: Unprocessable Entity for url" in str(e):
+                raise LocalModelInvocationException(str(e))
+            return (False, response)
+
+        return (True, response)
+
+
+class SageMakerTeiServing:
+    """SageMakerTeiServing class"""
+
+    def _upload_tei_artifacts(
+        self,
+        model_path: str,
+        sagemaker_session: Session,
+        s3_model_data_url: str = None,
+        image: str = None,
+        env_vars: dict = None,
+    ):
+        """Uploads the model artifacts to S3.
+
+        Args:
+            model_path: Path to the model
+            sagemaker_session: SageMaker session
+            s3_model_data_url: S3 model data URL
+            image: Image to use
+            env_vars: Environment variables to set
+        """
+        if s3_model_data_url:
+            bucket, key_prefix = parse_s3_url(url=s3_model_data_url)
+        else:
+            bucket, key_prefix = None, None
+
+        code_key_prefix = fw_utils.model_code_key_prefix(key_prefix, None, image)
+
+        bucket, code_key_prefix = determine_bucket_and_prefix(
+            bucket=bucket, key_prefix=code_key_prefix, sagemaker_session=sagemaker_session
+        )
+
+        code_dir = Path(model_path).joinpath("code")
+
+        s3_location = s3_path_join("s3://", bucket, code_key_prefix, "code")
+
+        logger.debug("Uploading TEI Model Resources uncompressed to: %s", s3_location)
+
+        model_data_url = S3Uploader.upload(
+            str(code_dir),
+            s3_location,
+            None,
+            sagemaker_session,
+        )
+
+        model_data = {
+            "S3DataSource": {
+                "CompressionType": "None",
+                "S3DataType": "S3Prefix",
+                "S3Uri": model_data_url + "/",
+            }
+        }
+
+        return (model_data, _update_env_vars(env_vars))
+
+
+def _update_env_vars(env_vars: dict) -> dict:
+    """Placeholder docstring"""
+    updated_env_vars = {}
+    updated_env_vars.update(_DEFAULT_ENV_VARS)
+    if env_vars:
+        updated_env_vars.update(env_vars)
+    return updated_env_vars
diff --git a/src/sagemaker/serve/utils/predictors.py b/src/sagemaker/serve/utils/predictors.py
index 866167c2c6..25a995eb48 100644
--- a/src/sagemaker/serve/utils/predictors.py
+++ b/src/sagemaker/serve/utils/predictors.py
@@ -209,6 +209,49 @@ def delete_predictor(self):
         self._mode_obj.destroy_server()
 
 
+class TeiLocalModePredictor(PredictorBase):
+    """Lightweight Tei predictor for local deployment in IN_PROCESS and LOCAL_CONTAINER modes"""
+
+    def __init__(
+        self,
+        mode_obj: Type[LocalContainerMode],
+        serializer=JSONSerializer(),
+        deserializer=JSONDeserializer(),
+    ):
+        self._mode_obj = mode_obj
+        self.serializer = serializer
+        self.deserializer = deserializer
+
+    def predict(self, data):
+        """Placeholder docstring"""
+        return [
+            self.deserializer.deserialize(
+                io.BytesIO(
+                    self._mode_obj._invoke_serving(
+                        self.serializer.serialize(data),
+                        self.content_type,
+                        self.deserializer.ACCEPT[0],
+                    )
+                ),
+                self.content_type,
+            )
+        ]
+
+    @property
+    def content_type(self):
+        """The MIME type of the data sent to the inference endpoint."""
+        return self.serializer.CONTENT_TYPE
+
+    @property
+    def accept(self):
+        """The content type(s) that are expected from the inference endpoint."""
+        return self.deserializer.ACCEPT
+
+    def delete_predictor(self):
+        """Shut down and remove the container that you created in LOCAL_CONTAINER mode"""
+        self._mode_obj.destroy_server()
+
+
 class TensorflowServingLocalPredictor(PredictorBase):
     """Lightweight predictor for local deployment in LOCAL_CONTAINER modes"""
 
diff --git a/src/sagemaker/serve/utils/telemetry_logger.py b/src/sagemaker/serve/utils/telemetry_logger.py
index 8983a4b5c9..99aeb4ff26 100644
--- a/src/sagemaker/serve/utils/telemetry_logger.py
+++ b/src/sagemaker/serve/utils/telemetry_logger.py
@@ -58,6 +58,7 @@
     str(ModelServer.DJL_SERVING): 4,
     str(ModelServer.TRITON): 5,
     str(ModelServer.TGI): 6,
+    str(ModelServer.TEI): 7,
 }
 
 MLFLOW_MODEL_PATH_CODE = {
diff --git a/src/sagemaker/serve/utils/types.py b/src/sagemaker/serve/utils/types.py
index 661093f249..3ac80aa7ea 100644
--- a/src/sagemaker/serve/utils/types.py
+++ b/src/sagemaker/serve/utils/types.py
@@ -18,6 +18,7 @@ def __str__(self):
     DJL_SERVING = 4
     TRITON = 5
     TGI = 6
+    TEI = 7
 
 
 class _DjlEngine(Enum):
diff --git a/tests/integ/sagemaker/serve/test_serve_tei.py b/tests/integ/sagemaker/serve/test_serve_tei.py
index 19ee0b57de..5cf1a3635c 100644
--- a/tests/integ/sagemaker/serve/test_serve_tei.py
+++ b/tests/integ/sagemaker/serve/test_serve_tei.py
@@ -28,47 +28,14 @@
 
 logger = logging.getLogger(__name__)
 
-sample_input = {
-    "inputs": "The man worked as a [MASK].",
-}
-
-loaded_response = [
-    {
-        "score": 0.0974755585193634,
-        "token": 10533,
-        "token_str": "carpenter",
-        "sequence": "the man worked as a carpenter.",
-    },
-    {
-        "score": 0.052383411675691605,
-        "token": 15610,
-        "token_str": "waiter",
-        "sequence": "the man worked as a waiter.",
-    },
-    {
-        "score": 0.04962712526321411,
-        "token": 13362,
-        "token_str": "barber",
-        "sequence": "the man worked as a barber.",
-    },
-    {
-        "score": 0.0378861166536808,
-        "token": 15893,
-        "token_str": "mechanic",
-        "sequence": "the man worked as a mechanic.",
-    },
-    {
-        "score": 0.037680838257074356,
-        "token": 18968,
-        "token_str": "salesman",
-        "sequence": "the man worked as a salesman.",
-    },
-]
+sample_input = {"inputs": "What is Deep Learning?"}
+
+loaded_response = []
 
 
 @pytest.fixture
 def model_input():
-    return {"inputs": "The man worked as a [MASK]."}
+    return {"inputs": "What is Deep Learning?"}
 
 
 @pytest.fixture
@@ -77,9 +44,6 @@ def model_builder_model_schema_builder():
         model_path=HF_DIR,
         model="BAAI/bge-m3",
         schema_builder=SchemaBuilder(sample_input, loaded_response),
-        model_metadata={
-            "HF_TASK": "sentence-similarity",
-        },
     )
 
 
diff --git a/tests/unit/sagemaker/serve/builder/test_tei_builder.py b/tests/unit/sagemaker/serve/builder/test_tei_builder.py
index 79a8f23324..4a75174bfc 100644
--- a/tests/unit/sagemaker/serve/builder/test_tei_builder.py
+++ b/tests/unit/sagemaker/serve/builder/test_tei_builder.py
@@ -18,7 +18,7 @@
 from sagemaker.serve.mode.function_pointers import Mode
 from tests.unit.sagemaker.serve.constants import MOCK_VPC_CONFIG
 
-from sagemaker.serve.utils.predictors import TgiLocalModePredictor
+from sagemaker.serve.utils.predictors import TeiLocalModePredictor
 
 mock_model_id = "bert-base-uncased"
 mock_prompt = "The man worked as a [MASK]."
@@ -96,7 +96,7 @@ def test_build_deploy_for_tei_local_container_and_remote_container(
 
         assert model.vpc_config == MOCK_VPC_CONFIG
         assert builder.env_vars["MODEL_LOADING_TIMEOUT"] == "1800"
-        assert isinstance(predictor, TgiLocalModePredictor)
+        assert isinstance(predictor, TeiLocalModePredictor)
 
         assert builder.nb_instance_type == "ml.g5.24xlarge"
 
@@ -139,7 +139,7 @@ def test_image_uri_override(
 
         assert builder.image_uri == MOCK_IMAGE_CONFIG
         assert builder.env_vars["MODEL_LOADING_TIMEOUT"] == "1800"
-        assert isinstance(predictor, TgiLocalModePredictor)
+        assert isinstance(predictor, TeiLocalModePredictor)
 
         assert builder.nb_instance_type == "ml.g5.24xlarge"
 
diff --git a/tests/unit/sagemaker/serve/model_server/tei/__init__.py b/tests/unit/sagemaker/serve/model_server/tei/__init__.py
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/tests/unit/sagemaker/serve/model_server/tei/test_server.py b/tests/unit/sagemaker/serve/model_server/tei/test_server.py
new file mode 100644
index 0000000000..16dcf12b5a
--- /dev/null
+++ b/tests/unit/sagemaker/serve/model_server/tei/test_server.py
@@ -0,0 +1,150 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You
+# may not use this file except in compliance with the License. A copy of
+# the License is located at
+#
+#     http://aws.amazon.com/apache2.0/
+#
+# or in the "license" file accompanying this file. This file is
+# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
+# ANY KIND, either express or implied. See the License for the specific
+# language governing permissions and limitations under the License.
+from __future__ import absolute_import
+
+from pathlib import PosixPath
+from unittest import TestCase
+from unittest.mock import Mock, patch
+
+from docker.types import DeviceRequest
+from sagemaker.serve.model_server.tei.server import LocalTeiServing, SageMakerTeiServing
+from sagemaker.serve.utils.exceptions import LocalModelInvocationException
+
+TEI_IMAGE = (
+    "246618743249.dkr.ecr.us-west-2.amazonaws.com/tei:2.0.1-tei1.2.3-gpu-py310-cu122-ubuntu22.04"
+)
+MODEL_PATH = "model_path"
+ENV_VAR = {"KEY": "VALUE"}
+PAYLOAD = {
+    "inputs": {
+        "sourceSentence": "How cute your dog is!",
+        "sentences": ["The mitochondria is the powerhouse of the cell.", "Your dog is so cute."],
+    }
+}
+S3_URI = "s3://mock_model_data_uri"
+SECRET_KEY = "secret_key"
+INFER_RESPONSE = []
+
+
+class TeiServerTests(TestCase):
+    @patch("sagemaker.serve.model_server.tei.server.requests")
+    def test_start_invoke_destroy_local_tei_server(self, mock_requests):
+        mock_container = Mock()
+        mock_docker_client = Mock()
+        mock_docker_client.containers.run.return_value = mock_container
+
+        local_tei_server = LocalTeiServing()
+        mock_schema_builder = Mock()
+        mock_schema_builder.input_serializer.serialize.return_value = PAYLOAD
+        local_tei_server.schema_builder = mock_schema_builder
+
+        local_tei_server._start_tei_serving(
+            client=mock_docker_client,
+            model_path=MODEL_PATH,
+            secret_key=SECRET_KEY,
+            image=TEI_IMAGE,
+            env_vars=ENV_VAR,
+        )
+
+        mock_docker_client.containers.run.assert_called_once_with(
+            TEI_IMAGE,
+            shm_size="2G",
+            device_requests=[DeviceRequest(count=-1, capabilities=[["gpu"]])],
+            network_mode="host",
+            detach=True,
+            auto_remove=True,
+            volumes={PosixPath("model_path/code"): {"bind": "/opt/ml/model/", "mode": "rw"}},
+            environment={
+                "TRANSFORMERS_CACHE": "/opt/ml/model/",
+                "HUGGINGFACE_HUB_CACHE": "/opt/ml/model/",
+                "KEY": "VALUE",
+                "SAGEMAKER_SERVE_SECRET_KEY": "secret_key",
+            },
+        )
+
+        mock_response = Mock()
+        mock_requests.post.side_effect = lambda *args, **kwargs: mock_response
+        mock_response.content = INFER_RESPONSE
+
+        res = local_tei_server._invoke_tei_serving(
+            request=PAYLOAD, content_type="application/json", accept="application/json"
+        )
+
+        self.assertEqual(res, INFER_RESPONSE)
+
+    def test_tei_deep_ping(self):
+        mock_predictor = Mock()
+        mock_response = Mock()
+        mock_schema_builder = Mock()
+
+        mock_predictor.predict.side_effect = lambda *args, **kwargs: mock_response
+        mock_schema_builder.sample_input = PAYLOAD
+
+        local_tei_server = LocalTeiServing()
+        local_tei_server.schema_builder = mock_schema_builder
+        res = local_tei_server._tei_deep_ping(mock_predictor)
+
+        self.assertEqual(res, (True, mock_response))
+
+    def test_tei_deep_ping_invoke_ex(self):
+        mock_predictor = Mock()
+        mock_schema_builder = Mock()
+
+        mock_predictor.predict.side_effect = lambda *args, **kwargs: exec(
+            'raise(ValueError("422 Client Error: Unprocessable Entity for url:"))'
+        )
+        mock_schema_builder.sample_input = PAYLOAD
+
+        local_tei_server = LocalTeiServing()
+        local_tei_server.schema_builder = mock_schema_builder
+
+        self.assertRaises(
+            LocalModelInvocationException, lambda: local_tei_server._tei_deep_ping(mock_predictor)
+        )
+
+    def test_tei_deep_ping_ex(self):
+        mock_predictor = Mock()
+
+        mock_predictor.predict.side_effect = lambda *args, **kwargs: Exception()
+
+        local_tei_server = LocalTeiServing()
+        res = local_tei_server._tei_deep_ping(mock_predictor)
+
+        self.assertEqual(res, (False, None))
+
+    @patch("sagemaker.serve.model_server.tei.server.S3Uploader")
+    def test_upload_artifacts_sagemaker_tei_server(self, mock_uploader):
+        mock_session = Mock()
+        mock_uploader.upload.side_effect = (
+            lambda *args, **kwargs: "s3://sagemaker-us-west-2-123456789123/tei-2024-05-20-16-05-36-027/code"
+        )
+
+        s3_upload_path, env_vars = SageMakerTeiServing()._upload_tei_artifacts(
+            model_path=MODEL_PATH,
+            sagemaker_session=mock_session,
+            s3_model_data_url=S3_URI,
+            image=TEI_IMAGE,
+        )
+
+        mock_uploader.upload.assert_called_once()
+        self.assertEqual(
+            s3_upload_path,
+            {
+                "S3DataSource": {
+                    "CompressionType": "None",
+                    "S3DataType": "S3Prefix",
+                    "S3Uri": "s3://sagemaker-us-west-2-123456789123/tei-2024-05-20-16-05-36-027/code/",
+                }
+            },
+        )
+        self.assertIsNotNone(env_vars)

From 6196b75eb8592f86e658708bfc9133c945949412 Mon Sep 17 00:00:00 2001
From: Tom Bousso <tombousso@gmail.com>
Date: Wed, 22 May 2024 13:25:51 -0700
Subject: [PATCH 45/58] Convert pytorchddp distribution to smdistributed
 distribution (#4698)

* rewrite pytorchddp to smdistributed

* remove instance type check

* Update estimator.py

* remove validate_pytorch_distribution

* fix

* fix unit tests

* fix formatting

* check instance type not None
---
 src/sagemaker/fw_utils.py          | 99 +-----------------------------
 src/sagemaker/pytorch/estimator.py | 14 +++++
 tests/unit/test_fw_utils.py        | 75 +---------------------
 tests/unit/test_pytorch.py         |  5 +-
 4 files changed, 21 insertions(+), 172 deletions(-)

diff --git a/src/sagemaker/fw_utils.py b/src/sagemaker/fw_utils.py
index 33018becdd..be3658365a 100644
--- a/src/sagemaker/fw_utils.py
+++ b/src/sagemaker/fw_utils.py
@@ -145,22 +145,6 @@
     ],
 }
 
-PYTORCHDDP_SUPPORTED_FRAMEWORK_VERSIONS = [
-    "1.10",
-    "1.10.0",
-    "1.10.2",
-    "1.11",
-    "1.11.0",
-    "1.12",
-    "1.12.0",
-    "1.12.1",
-    "1.13.1",
-    "2.0.0",
-    "2.0.1",
-    "2.1.0",
-    "2.2.0",
-]
-
 TORCH_DISTRIBUTED_GPU_SUPPORTED_FRAMEWORK_VERSIONS = [
     "1.13.1",
     "2.0.0",
@@ -795,7 +779,6 @@ def _validate_smdataparallel_args(
 
     Raises:
         ValueError: if
-            (`instance_type` is not in SM_DATAPARALLEL_SUPPORTED_INSTANCE_TYPES or
             `py_version` is not python3 or
             `framework_version` is not in SM_DATAPARALLEL_SUPPORTED_FRAMEWORK_VERSION
     """
@@ -806,17 +789,10 @@ def _validate_smdataparallel_args(
     if not smdataparallel_enabled:
         return
 
-    is_instance_type_supported = instance_type in SM_DATAPARALLEL_SUPPORTED_INSTANCE_TYPES
-
     err_msg = ""
 
-    if not is_instance_type_supported:
-        # instance_type is required
-        err_msg += (
-            f"Provided instance_type {instance_type} is not supported by smdataparallel.\n"
-            "Please specify one of the supported instance types:"
-            f"{SM_DATAPARALLEL_SUPPORTED_INSTANCE_TYPES}\n"
-        )
+    if not instance_type:
+        err_msg += "Please specify an instance_type for smdataparallel.\n"
 
     if not image_uri:
         # ignore framework_version & py_version if image_uri is set
@@ -928,13 +904,6 @@ def validate_distribution(
             )
             if framework_name and framework_name == "pytorch":
                 # We need to validate only for PyTorch framework
-                validate_pytorch_distribution(
-                    distribution=validated_distribution,
-                    framework_name=framework_name,
-                    framework_version=framework_version,
-                    py_version=py_version,
-                    image_uri=image_uri,
-                )
                 validate_torch_distributed_distribution(
                     instance_type=instance_type,
                     distribution=validated_distribution,
@@ -968,13 +937,6 @@ def validate_distribution(
         )
         if framework_name and framework_name == "pytorch":
             # We need to validate only for PyTorch framework
-            validate_pytorch_distribution(
-                distribution=validated_distribution,
-                framework_name=framework_name,
-                framework_version=framework_version,
-                py_version=py_version,
-                image_uri=image_uri,
-            )
             validate_torch_distributed_distribution(
                 instance_type=instance_type,
                 distribution=validated_distribution,
@@ -1023,63 +985,6 @@ def validate_distribution_for_instance_type(instance_type, distribution):
         raise ValueError(err_msg)
 
 
-def validate_pytorch_distribution(
-    distribution, framework_name, framework_version, py_version, image_uri
-):
-    """Check if pytorch distribution strategy is correctly invoked by the user.
-
-    Args:
-        distribution (dict): A dictionary with information to enable distributed training.
-            (Defaults to None if distributed training is not enabled.) For example:
-
-            .. code:: python
-
-                {
-                    "pytorchddp": {
-                        "enabled": True
-                    }
-                }
-        framework_name (str): A string representing the name of framework selected.
-        framework_version (str): A string representing the framework version selected.
-        py_version (str): A string representing the python version selected.
-        image_uri (str): A string representing a Docker image URI.
-
-    Raises:
-        ValueError: if
-            `py_version` is not python3 or
-            `framework_version` is not in PYTORCHDDP_SUPPORTED_FRAMEWORK_VERSIONS
-    """
-    if framework_name and framework_name != "pytorch":
-        # We need to validate only for PyTorch framework
-        return
-
-    pytorch_ddp_enabled = False
-    if "pytorchddp" in distribution:
-        pytorch_ddp_enabled = distribution.get("pytorchddp").get("enabled", False)
-    if not pytorch_ddp_enabled:
-        # Distribution strategy other than pytorchddp is selected
-        return
-
-    err_msg = ""
-    if not image_uri:
-        # ignore framework_version and py_version if image_uri is set
-        # in case image_uri is not set, then both are mandatory
-        if framework_version not in PYTORCHDDP_SUPPORTED_FRAMEWORK_VERSIONS:
-            err_msg += (
-                f"Provided framework_version {framework_version} is not supported by"
-                " pytorchddp.\n"
-                "Please specify one of the supported framework versions:"
-                f" {PYTORCHDDP_SUPPORTED_FRAMEWORK_VERSIONS} \n"
-            )
-        if "py3" not in py_version:
-            err_msg += (
-                f"Provided py_version {py_version} is not supported by pytorchddp.\n"
-                "Please specify py_version>=py3"
-            )
-    if err_msg:
-        raise ValueError(err_msg)
-
-
 def validate_torch_distributed_distribution(
     instance_type,
     distribution,
diff --git a/src/sagemaker/pytorch/estimator.py b/src/sagemaker/pytorch/estimator.py
index a4e24d1ff0..412926279c 100644
--- a/src/sagemaker/pytorch/estimator.py
+++ b/src/sagemaker/pytorch/estimator.py
@@ -276,6 +276,20 @@ def __init__(
             kwargs["entry_point"] = entry_point
 
         if distribution is not None:
+            # rewrite pytorchddp to smdistributed
+            if "pytorchddp" in distribution:
+                if "smdistributed" in distribution:
+                    raise ValueError(
+                        "Cannot use both pytorchddp and smdistributed "
+                        "distribution options together.",
+                        distribution,
+                    )
+
+                # convert pytorchddp distribution into smdistributed distribution
+                distribution = distribution.copy()
+                distribution["smdistributed"] = {"dataparallel": distribution["pytorchddp"]}
+                del distribution["pytorchddp"]
+
             distribution = validate_distribution(
                 distribution,
                 self.instance_groups,
diff --git a/tests/unit/test_fw_utils.py b/tests/unit/test_fw_utils.py
index e955d68227..97d4e6ec2a 100644
--- a/tests/unit/test_fw_utils.py
+++ b/tests/unit/test_fw_utils.py
@@ -854,17 +854,14 @@ def test_validate_smdataparallel_args_raises():
 
     # Cases {PT|TF2}
     # 1. None instance type
-    # 2. incorrect instance type
-    # 3. incorrect python version
-    # 4. incorrect framework version
+    # 2. incorrect python version
+    # 3. incorrect framework version
 
     bad_args = [
         (None, "tensorflow", "2.3.1", "py3", smdataparallel_enabled),
-        ("ml.p3.2xlarge", "tensorflow", "2.3.1", "py3", smdataparallel_enabled),
         ("ml.p3dn.24xlarge", "tensorflow", "2.3.1", "py2", smdataparallel_enabled),
         ("ml.p3.16xlarge", "tensorflow", "1.3.1", "py3", smdataparallel_enabled),
         (None, "pytorch", "1.6.0", "py3", smdataparallel_enabled),
-        ("ml.p3.2xlarge", "pytorch", "1.6.0", "py3", smdataparallel_enabled),
         ("ml.p3dn.24xlarge", "pytorch", "1.6.0", "py2", smdataparallel_enabled),
         ("ml.p3.16xlarge", "pytorch", "1.5.0", "py3", smdataparallel_enabled),
     ]
@@ -966,74 +963,6 @@ def test_validate_smdataparallel_args_not_raises():
         )
 
 
-def test_validate_pytorchddp_not_raises():
-    # Case 1: Framework is not PyTorch
-    fw_utils.validate_pytorch_distribution(
-        distribution=None,
-        framework_name="tensorflow",
-        framework_version="2.9.1",
-        py_version="py3",
-        image_uri="custom-container",
-    )
-    # Case 2: Framework is PyTorch, but distribution is not PyTorchDDP
-    pytorchddp_disabled = {"pytorchddp": {"enabled": False}}
-    fw_utils.validate_pytorch_distribution(
-        distribution=pytorchddp_disabled,
-        framework_name="pytorch",
-        framework_version="1.10",
-        py_version="py3",
-        image_uri="custom-container",
-    )
-    # Case 3: Framework is PyTorch, Distribution is PyTorchDDP enabled, supported framework and py versions
-    pytorchddp_enabled = {"pytorchddp": {"enabled": True}}
-    pytorchddp_supported_fw_versions = [
-        "1.10",
-        "1.10.0",
-        "1.10.2",
-        "1.11",
-        "1.11.0",
-        "1.12",
-        "1.12.0",
-        "1.12.1",
-        "1.13.1",
-        "2.0.0",
-        "2.0.1",
-        "2.1.0",
-        "2.2.0",
-    ]
-    for framework_version in pytorchddp_supported_fw_versions:
-        fw_utils.validate_pytorch_distribution(
-            distribution=pytorchddp_enabled,
-            framework_name="pytorch",
-            framework_version=framework_version,
-            py_version="py3",
-            image_uri="custom-container",
-        )
-
-
-def test_validate_pytorchddp_raises():
-    pytorchddp_enabled = {"pytorchddp": {"enabled": True}}
-    # Case 1: Unsupported framework version
-    with pytest.raises(ValueError):
-        fw_utils.validate_pytorch_distribution(
-            distribution=pytorchddp_enabled,
-            framework_name="pytorch",
-            framework_version="1.8",
-            py_version="py3",
-            image_uri=None,
-        )
-
-    # Case 2: Unsupported Py version
-    with pytest.raises(ValueError):
-        fw_utils.validate_pytorch_distribution(
-            distribution=pytorchddp_enabled,
-            framework_name="pytorch",
-            framework_version="1.10",
-            py_version="py2",
-            image_uri=None,
-        )
-
-
 def test_validate_torch_distributed_not_raises():
     # Case 1: Framework is PyTorch, but torch_distributed is not enabled
     torch_distributed_disabled = {"torch_distributed": {"enabled": False}}
diff --git a/tests/unit/test_pytorch.py b/tests/unit/test_pytorch.py
index 5ada026ef8..618d0d7ea8 100644
--- a/tests/unit/test_pytorch.py
+++ b/tests/unit/test_pytorch.py
@@ -801,14 +801,15 @@ def test_pytorch_ddp_distribution_configuration(
         distribution=pytorch.distribution
     )
     expected_torch_ddp = {
-        "sagemaker_pytorch_ddp_enabled": True,
+        "sagemaker_distributed_dataparallel_enabled": True,
+        "sagemaker_distributed_dataparallel_custom_mpi_options": "",
         "sagemaker_instance_type": test_instance_type,
     }
     assert actual_pytorch_ddp == expected_torch_ddp
 
 
 def test_pytorch_ddp_distribution_configuration_unsupported(sagemaker_session):
-    unsupported_framework_version = "1.9.1"
+    unsupported_framework_version = "1.5.0"
     unsupported_py_version = "py2"
     with pytest.raises(ValueError) as error:
         _pytorch_estimator(

From d4e42dbf1729edcd9ed8295ce583fae4c37c32a2 Mon Sep 17 00:00:00 2001
From: ci <ci>
Date: Wed, 22 May 2024 22:51:18 +0000
Subject: [PATCH 46/58] prepare release v2.221.1

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

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 9261d394e8..63e5114f10 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,5 +1,12 @@
 # Changelog
 
+## v2.221.1 (2024-05-22)
+
+### Bug Fixes and Other Changes
+
+ * Convert pytorchddp distribution to smdistributed distribution
+ * Add tei cpu image
+
 ## v2.221.0 (2024-05-20)
 
 ### Features
diff --git a/VERSION b/VERSION
index 1ab82bbae0..73572ec54a 100644
--- a/VERSION
+++ b/VERSION
@@ -1 +1 @@
-2.221.1.dev0
+2.221.1

From 3676d3e90c34fc338dca2414fa50c3ad898a1671 Mon Sep 17 00:00:00 2001
From: ci <ci>
Date: Wed, 22 May 2024 22:51:20 +0000
Subject: [PATCH 47/58] update development version to v2.221.2.dev0

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

diff --git a/VERSION b/VERSION
index 73572ec54a..e55266069e 100644
--- a/VERSION
+++ b/VERSION
@@ -1 +1 @@
-2.221.1
+2.221.2.dev0

From e60c488ff05ba4e0f47b9cd00b51ba9e44baa013 Mon Sep 17 00:00:00 2001
From: Jonathan Makunga <54963715+makungaj1@users.noreply.github.com>
Date: Fri, 24 May 2024 11:56:12 -0700
Subject: [PATCH 48/58] Update: SM Endpoint Routing Strategy Support. (#4702)

* RoutingConfig

* Refactoring

* Docstring

* UT

* Refactoring

* Refactoring

---------

Co-authored-by: Jonathan Makunga <makung@amazon.com>
---
 src/sagemaker/enums.py                       | 12 ++++++++
 src/sagemaker/huggingface/model.py           |  1 +
 src/sagemaker/jumpstart/factory/model.py     |  2 ++
 src/sagemaker/jumpstart/model.py             |  6 +++-
 src/sagemaker/jumpstart/types.py             |  3 ++
 src/sagemaker/model.py                       | 17 ++++++++++-
 src/sagemaker/serve/builder/model_builder.py |  2 +-
 src/sagemaker/utils.py                       | 31 ++++++++++++++++++++
 tests/unit/sagemaker/model/test_deploy.py    |  5 ++++
 tests/unit/test_utils.py                     | 29 ++++++++++++++++++
 10 files changed, 105 insertions(+), 3 deletions(-)

diff --git a/src/sagemaker/enums.py b/src/sagemaker/enums.py
index 5b4d0d6790..f02b275cbe 100644
--- a/src/sagemaker/enums.py
+++ b/src/sagemaker/enums.py
@@ -28,3 +28,15 @@ class EndpointType(Enum):
     INFERENCE_COMPONENT_BASED = (
         "InferenceComponentBased"  # Amazon SageMaker Inference Component Based Endpoint
     )
+
+
+class RoutingStrategy(Enum):
+    """Strategy for routing https traffics."""
+
+    RANDOM = "RANDOM"
+    """The endpoint routes each request to a randomly chosen instance.
+    """
+    LEAST_OUTSTANDING_REQUESTS = "LEAST_OUTSTANDING_REQUESTS"
+    """The endpoint routes requests to the specific instances that have
+    more capacity to process them.
+    """
diff --git a/src/sagemaker/huggingface/model.py b/src/sagemaker/huggingface/model.py
index f71dca0ac8..662baecae6 100644
--- a/src/sagemaker/huggingface/model.py
+++ b/src/sagemaker/huggingface/model.py
@@ -334,6 +334,7 @@ def deploy(
             endpoint_type=kwargs.get("endpoint_type", None),
             resources=kwargs.get("resources", None),
             managed_instance_scaling=kwargs.get("managed_instance_scaling", None),
+            routing_config=kwargs.get("routing_config", None),
         )
 
     def register(
diff --git a/src/sagemaker/jumpstart/factory/model.py b/src/sagemaker/jumpstart/factory/model.py
index 28746990e3..41c1ca4437 100644
--- a/src/sagemaker/jumpstart/factory/model.py
+++ b/src/sagemaker/jumpstart/factory/model.py
@@ -555,6 +555,7 @@ def get_deploy_kwargs(
     resources: Optional[ResourceRequirements] = None,
     managed_instance_scaling: Optional[str] = None,
     endpoint_type: Optional[EndpointType] = None,
+    routing_config: Optional[Dict[str, Any]] = None,
 ) -> JumpStartModelDeployKwargs:
     """Returns kwargs required to call `deploy` on `sagemaker.estimator.Model` object."""
 
@@ -586,6 +587,7 @@ def get_deploy_kwargs(
         accept_eula=accept_eula,
         endpoint_logging=endpoint_logging,
         resources=resources,
+        routing_config=routing_config,
     )
 
     deploy_kwargs = _add_sagemaker_session_to_kwargs(kwargs=deploy_kwargs)
diff --git a/src/sagemaker/jumpstart/model.py b/src/sagemaker/jumpstart/model.py
index 4529bc11b9..994193de3e 100644
--- a/src/sagemaker/jumpstart/model.py
+++ b/src/sagemaker/jumpstart/model.py
@@ -14,7 +14,7 @@
 
 from __future__ import absolute_import
 
-from typing import Dict, List, Optional, Union
+from typing import Dict, List, Optional, Union, Any
 from botocore.exceptions import ClientError
 
 from sagemaker import payloads
@@ -496,6 +496,7 @@ def deploy(
         resources: Optional[ResourceRequirements] = None,
         managed_instance_scaling: Optional[str] = None,
         endpoint_type: EndpointType = EndpointType.MODEL_BASED,
+        routing_config: Optional[Dict[str, Any]] = None,
     ) -> PredictorBase:
         """Creates endpoint by calling base ``Model`` class `deploy` method.
 
@@ -590,6 +591,8 @@ def deploy(
                 endpoint.
             endpoint_type (EndpointType): The type of endpoint used to deploy models.
                 (Default: EndpointType.MODEL_BASED).
+            routing_config (Optional[Dict]): Settings the control how the endpoint routes
+                incoming traffic to the instances that the endpoint hosts.
 
         Raises:
             MarketplaceModelSubscriptionError: If the caller is not subscribed to the model.
@@ -625,6 +628,7 @@ def deploy(
             managed_instance_scaling=managed_instance_scaling,
             endpoint_type=endpoint_type,
             model_type=self.model_type,
+            routing_config=routing_config,
         )
         if (
             self.model_type == JumpStartModelType.PROPRIETARY
diff --git a/src/sagemaker/jumpstart/types.py b/src/sagemaker/jumpstart/types.py
index 05c38da266..bea125d423 100644
--- a/src/sagemaker/jumpstart/types.py
+++ b/src/sagemaker/jumpstart/types.py
@@ -1614,6 +1614,7 @@ class JumpStartModelDeployKwargs(JumpStartKwargs):
         "endpoint_logging",
         "resources",
         "endpoint_type",
+        "routing_config",
     ]
 
     SERIALIZATION_EXCLUSION_SET = {
@@ -1658,6 +1659,7 @@ def __init__(
         endpoint_logging: Optional[bool] = None,
         resources: Optional[ResourceRequirements] = None,
         endpoint_type: Optional[EndpointType] = None,
+        routing_config: Optional[Dict[str, Any]] = None,
     ) -> None:
         """Instantiates JumpStartModelDeployKwargs object."""
 
@@ -1690,6 +1692,7 @@ def __init__(
         self.endpoint_logging = endpoint_logging
         self.resources = resources
         self.endpoint_type = endpoint_type
+        self.routing_config = routing_config
 
 
 class JumpStartEstimatorInitKwargs(JumpStartKwargs):
diff --git a/src/sagemaker/model.py b/src/sagemaker/model.py
index fd21b6342e..1bb6cb2e5c 100644
--- a/src/sagemaker/model.py
+++ b/src/sagemaker/model.py
@@ -20,7 +20,7 @@
 import os
 import re
 import copy
-from typing import List, Dict, Optional, Union
+from typing import List, Dict, Optional, Union, Any
 
 import sagemaker
 from sagemaker import (
@@ -66,6 +66,7 @@
     resolve_nested_dict_value_from_config,
     format_tags,
     Tags,
+    _resolve_routing_config,
 )
 from sagemaker.async_inference import AsyncInferenceConfig
 from sagemaker.predictor_async import AsyncPredictor
@@ -1309,6 +1310,7 @@ def deploy(
         resources: Optional[ResourceRequirements] = None,
         endpoint_type: EndpointType = EndpointType.MODEL_BASED,
         managed_instance_scaling: Optional[str] = None,
+        routing_config: Optional[Dict[str, Any]] = None,
         **kwargs,
     ):
         """Deploy this ``Model`` to an ``Endpoint`` and optionally return a ``Predictor``.
@@ -1406,6 +1408,15 @@ def deploy(
                 Endpoint. (Default: None).
             endpoint_type (Optional[EndpointType]): The type of an endpoint used to deploy models.
                 (Default: EndpointType.MODEL_BASED).
+            routing_config (Optional[Dict[str, Any]): Settings the control how the endpoint routes incoming
+                traffic to the instances that the endpoint hosts.
+                Currently, support dictionary key ``RoutingStrategy``.
+
+                .. code:: python
+
+                    {
+                        "RoutingStrategy":  sagemaker.enums.RoutingStrategy.RANDOM
+                    }
         Raises:
              ValueError: If arguments combination check failed in these circumstances:
                 - If no role is specified or
@@ -1458,6 +1469,8 @@ def deploy(
         if self.role is None:
             raise ValueError("Role can not be null for deploying a model")
 
+        routing_config = _resolve_routing_config(routing_config)
+
         if (
             inference_recommendation_id is not None
             or self.inference_recommender_job_results is not None
@@ -1543,6 +1556,7 @@ def deploy(
                     model_data_download_timeout=model_data_download_timeout,
                     container_startup_health_check_timeout=container_startup_health_check_timeout,
                     managed_instance_scaling=managed_instance_scaling_config,
+                    routing_config=routing_config,
                 )
 
                 self.sagemaker_session.endpoint_from_production_variants(
@@ -1625,6 +1639,7 @@ def deploy(
                 volume_size=volume_size,
                 model_data_download_timeout=model_data_download_timeout,
                 container_startup_health_check_timeout=container_startup_health_check_timeout,
+                routing_config=routing_config,
             )
             if endpoint_name:
                 self.endpoint_name = endpoint_name
diff --git a/src/sagemaker/serve/builder/model_builder.py b/src/sagemaker/serve/builder/model_builder.py
index 1fe75065d5..44bc46b00b 100644
--- a/src/sagemaker/serve/builder/model_builder.py
+++ b/src/sagemaker/serve/builder/model_builder.py
@@ -169,7 +169,7 @@ class ModelBuilder(Triton, DJL, JumpStart, TGI, Transformers, TensorflowServing,
             in order for model builder to build the artifacts correctly (according
             to the model server). Possible values for this argument are
             ``TORCHSERVE``, ``MMS``, ``TENSORFLOW_SERVING``, ``DJL_SERVING``,
-            ``TRITON``,``TGI``, and ``TEI``.
+            ``TRITON``, ``TGI``, and ``TEI``.
         model_metadata (Optional[Dict[str, Any]): Dictionary used to override model metadata.
             Currently, ``HF_TASK`` is overridable for HuggingFace model. HF_TASK should be set for
             new models without task metadata in the Hub, adding unsupported task types will throw
diff --git a/src/sagemaker/utils.py b/src/sagemaker/utils.py
index 0436c0afea..430effefa3 100644
--- a/src/sagemaker/utils.py
+++ b/src/sagemaker/utils.py
@@ -44,6 +44,7 @@
     _log_sagemaker_config_single_substitution,
     _log_sagemaker_config_merge,
 )
+from sagemaker.enums import RoutingStrategy
 from sagemaker.session_settings import SessionSettings
 from sagemaker.workflow import is_pipeline_variable, is_pipeline_parameter_string
 from sagemaker.workflow.entities import PipelineVariable
@@ -1655,3 +1656,33 @@ def deep_override_dict(
     )
     flattened_dict1.update(flattened_dict2)
     return unflatten_dict(flattened_dict1) if flattened_dict1 else {}
+
+
+def _resolve_routing_config(routing_config: Optional[Dict[str, Any]]) -> Optional[Dict[str, Any]]:
+    """Resolve Routing Config
+
+    Args:
+        routing_config (Optional[Dict[str, Any]]): The routing config.
+
+    Returns:
+        Optional[Dict[str, Any]]: The resolved routing config.
+
+    Raises:
+        ValueError: If the RoutingStrategy is invalid.
+    """
+
+    if routing_config:
+        routing_strategy = routing_config.get("RoutingStrategy", None)
+        if routing_strategy:
+            if isinstance(routing_strategy, RoutingStrategy):
+                return {"RoutingStrategy": routing_strategy.name}
+            if isinstance(routing_strategy, str) and (
+                routing_strategy.upper() == RoutingStrategy.RANDOM.name
+                or routing_strategy.upper() == RoutingStrategy.LEAST_OUTSTANDING_REQUESTS.name
+            ):
+                return {"RoutingStrategy": routing_strategy.upper()}
+            raise ValueError(
+                "RoutingStrategy must be either RoutingStrategy.RANDOM "
+                "or RoutingStrategy.LEAST_OUTSTANDING_REQUESTS"
+            )
+    return None
diff --git a/tests/unit/sagemaker/model/test_deploy.py b/tests/unit/sagemaker/model/test_deploy.py
index 953cbe775c..69ea2c1f56 100644
--- a/tests/unit/sagemaker/model/test_deploy.py
+++ b/tests/unit/sagemaker/model/test_deploy.py
@@ -125,6 +125,7 @@ def test_deploy(name_from_base, prepare_container_def, production_variant, sagem
         volume_size=None,
         model_data_download_timeout=None,
         container_startup_health_check_timeout=None,
+        routing_config=None,
     )
 
     sagemaker_session.create_model.assert_called_with(
@@ -184,6 +185,7 @@ def test_deploy_accelerator_type(
         volume_size=None,
         model_data_download_timeout=None,
         container_startup_health_check_timeout=None,
+        routing_config=None,
     )
 
     sagemaker_session.endpoint_from_production_variants.assert_called_with(
@@ -506,6 +508,7 @@ def test_deploy_serverless_inference(production_variant, create_sagemaker_model,
         volume_size=None,
         model_data_download_timeout=None,
         container_startup_health_check_timeout=None,
+        routing_config=None,
     )
 
     sagemaker_session.endpoint_from_production_variants.assert_called_with(
@@ -938,6 +941,7 @@ def test_deploy_customized_volume_size_and_timeout(
         volume_size=volume_size_gb,
         model_data_download_timeout=model_data_download_timeout_sec,
         container_startup_health_check_timeout=startup_health_check_timeout_sec,
+        routing_config=None,
     )
 
     sagemaker_session.create_model.assert_called_with(
@@ -987,6 +991,7 @@ def test_deploy_with_resources(sagemaker_session, name_from_base, production_var
         volume_size=None,
         model_data_download_timeout=None,
         container_startup_health_check_timeout=None,
+        routing_config=None,
     )
     sagemaker_session.endpoint_from_production_variants.assert_called_with(
         name=name_from_base(MODEL_NAME),
diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py
index 81d8279e6d..d81984e81f 100644
--- a/tests/unit/test_utils.py
+++ b/tests/unit/test_utils.py
@@ -30,6 +30,7 @@
 from mock import call, patch, Mock, MagicMock, PropertyMock
 
 import sagemaker
+from sagemaker.enums import RoutingStrategy
 from sagemaker.experiments._run_context import _RunContext
 from sagemaker.session_settings import SessionSettings
 from sagemaker.utils import (
@@ -50,6 +51,7 @@
     _is_bad_link,
     custom_extractall_tarfile,
     can_model_package_source_uri_autopopulate,
+    _resolve_routing_config,
 )
 from tests.unit.sagemaker.workflow.helpers import CustomStep
 from sagemaker.workflow.parameters import ParameterString, ParameterInteger
@@ -1866,3 +1868,30 @@ def test_deep_override_skip_keys(self):
         expected_result = {"a": 1, "b": {"x": 20, "y": 3, "z": 30}, "c": [4, 5]}
 
         self.assertEqual(deep_override_dict(dict1, dict2, skip_keys=["c", "d"]), expected_result)
+
+
+@pytest.mark.parametrize(
+    "routing_config, expected",
+    [
+        ({"RoutingStrategy": RoutingStrategy.RANDOM}, {"RoutingStrategy": "RANDOM"}),
+        ({"RoutingStrategy": "RANDOM"}, {"RoutingStrategy": "RANDOM"}),
+        (
+            {"RoutingStrategy": RoutingStrategy.LEAST_OUTSTANDING_REQUESTS},
+            {"RoutingStrategy": "LEAST_OUTSTANDING_REQUESTS"},
+        ),
+        (
+            {"RoutingStrategy": "LEAST_OUTSTANDING_REQUESTS"},
+            {"RoutingStrategy": "LEAST_OUTSTANDING_REQUESTS"},
+        ),
+        ({"RoutingStrategy": None}, None),
+        (None, None),
+    ],
+)
+def test_resolve_routing_config(routing_config, expected):
+    res = _resolve_routing_config(routing_config)
+
+    assert res == expected
+
+
+def test_resolve_routing_config_ex():
+    pytest.raises(ValueError, lambda: _resolve_routing_config({"RoutingStrategy": "Invalid"}))

From cbbbb32c439dc7b3654e5b73568033aaca665b4c Mon Sep 17 00:00:00 2001
From: sagemaker-bot <sagemaker-bot@amazon.com>
Date: Wed, 29 May 2024 14:17:35 +0000
Subject: [PATCH 49/58] change: update image_uri_configs  05-29-2024 07:17:35
 PST

---
 src/sagemaker/image_uri_config/pytorch.json | 44 ++++++++++++++++++++-
 1 file changed, 43 insertions(+), 1 deletion(-)

diff --git a/src/sagemaker/image_uri_config/pytorch.json b/src/sagemaker/image_uri_config/pytorch.json
index f04f9c6b69..b846d51246 100644
--- a/src/sagemaker/image_uri_config/pytorch.json
+++ b/src/sagemaker/image_uri_config/pytorch.json
@@ -1315,7 +1315,8 @@
             "1.13": "1.13.1",
             "2.0": "2.0.1",
             "2.1": "2.1.0",
-            "2.2": "2.2.0"
+            "2.2": "2.2.0",
+            "2.3": "2.3.0"
         },
         "versions": {
             "0.4.0": {
@@ -2288,6 +2289,47 @@
                     "us-west-2": "763104351884"
                 },
                 "repository": "pytorch-training"
+            },
+            "2.3.0": {
+                "py_versions": [
+                    "py311"
+                ],
+                "registries": {
+                    "af-south-1": "626614931356",
+                    "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",
+                    "ca-west-1": "204538143572",
+                    "cn-north-1": "727897471807",
+                    "cn-northwest-1": "727897471807",
+                    "eu-central-1": "763104351884",
+                    "eu-central-2": "380420809688",
+                    "eu-north-1": "763104351884",
+                    "eu-south-1": "692866216735",
+                    "eu-south-2": "503227376785",
+                    "eu-west-1": "763104351884",
+                    "eu-west-2": "763104351884",
+                    "eu-west-3": "763104351884",
+                    "il-central-1": "780543022126",
+                    "me-central-1": "914824155844",
+                    "me-south-1": "217643126080",
+                    "sa-east-1": "763104351884",
+                    "us-east-1": "763104351884",
+                    "us-east-2": "763104351884",
+                    "us-gov-east-1": "446045086412",
+                    "us-gov-west-1": "442386744353",
+                    "us-west-1": "763104351884",
+                    "us-west-2": "763104351884"
+                },
+                "repository": "pytorch-training"
             }
         }
     }

From c48d7c8c4d871dc4bc07ef773f0e58df6f1d4cc8 Mon Sep 17 00:00:00 2001
From: Zhaoqi <jzhaoqwa@amazon.com>
Date: Wed, 29 May 2024 14:24:22 -0400
Subject: [PATCH 50/58] Making project name in workflow files dynamic (#4708)

---
 .github/workflows/codebuild-ci.yml | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/.github/workflows/codebuild-ci.yml b/.github/workflows/codebuild-ci.yml
index 85919f0afe..8c6bd6b337 100644
--- a/.github/workflows/codebuild-ci.yml
+++ b/.github/workflows/codebuild-ci.yml
@@ -55,7 +55,7 @@ jobs:
       - name: Run Codestyle & Doc Tests
         uses: aws-actions/aws-codebuild-run-build@v1
         with:
-          project-name: sagemaker-python-sdk-ci-codestyle-doc-tests
+          project-name: ${{ github.event.repository.name }}-ci-codestyle-doc-tests
           source-version-override: 'refs/pull/${{ github.event.pull_request.number }}/head^{${{ github.event.pull_request.head.sha }}}'
   unit-tests:
     runs-on: ubuntu-latest
@@ -74,7 +74,7 @@ jobs:
       - name: Run Unit Tests
         uses: aws-actions/aws-codebuild-run-build@v1
         with:
-          project-name: sagemaker-python-sdk-ci-unit-tests
+          project-name: ${{ github.event.repository.name }}-ci-unit-tests
           source-version-override: 'refs/pull/${{ github.event.pull_request.number }}/head^{${{ github.event.pull_request.head.sha }}}'
           env-vars-for-codebuild: |
             PY_VERSION
@@ -93,5 +93,5 @@ jobs:
       - name: Run Integ Tests
         uses: aws-actions/aws-codebuild-run-build@v1
         with:
-          project-name: sagemaker-python-sdk-ci-integ-tests
+          project-name: ${{ github.event.repository.name }}-ci-integ-tests
           source-version-override: 'refs/pull/${{ github.event.pull_request.number }}/head^{${{ github.event.pull_request.head.sha }}}'

From b68a81078c06a590ee3027dcf180e36ac9719de2 Mon Sep 17 00:00:00 2001
From: Kalyani Nikure <110067132+knikure@users.noreply.github.com>
Date: Mon, 3 Jun 2024 11:54:28 -0700
Subject: [PATCH 51/58] fix: Fix ci unit-tests (#4713)

---
 tests/conftest.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/tests/conftest.py b/tests/conftest.py
index 0309781e7b..7bab05dfb3 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -253,7 +253,9 @@ def mxnet_eia_latest_py_version():
 
 @pytest.fixture(scope="module", params=["py2", "py3"])
 def pytorch_training_py_version(pytorch_training_version, request):
-    if Version(pytorch_training_version) >= Version("2.0"):
+    if Version(pytorch_training_version) >= Version("2.3"):
+        return "py311"
+    elif Version(pytorch_training_version) >= Version("2.0"):
         return "py310"
     elif Version(pytorch_training_version) >= Version("1.13"):
         return "py39"

From dc8e9cd7ce7e6676f9a30634f0dfe1c1bfbdd3c5 Mon Sep 17 00:00:00 2001
From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com>
Date: Mon, 3 Jun 2024 13:17:02 -0700
Subject: [PATCH 52/58] chore(deps): bump requests from 2.31.0 to 2.32.2 in
 /tests/data/serve_resources/mlflow/pytorch (#4709)

Bumps [requests](https://github.com/psf/requests) from 2.31.0 to 2.32.2.
- [Release notes](https://github.com/psf/requests/releases)
- [Changelog](https://github.com/psf/requests/blob/main/HISTORY.md)
- [Commits](https://github.com/psf/requests/compare/v2.31.0...v2.32.2)

---
updated-dependencies:
- dependency-name: requests
  dependency-type: direct:production
...

Signed-off-by: dependabot[bot] <support@github.com>
Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>
---
 tests/data/serve_resources/mlflow/pytorch/requirements.txt | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/data/serve_resources/mlflow/pytorch/requirements.txt b/tests/data/serve_resources/mlflow/pytorch/requirements.txt
index a10494de69..fef7d61013 100644
--- a/tests/data/serve_resources/mlflow/pytorch/requirements.txt
+++ b/tests/data/serve_resources/mlflow/pytorch/requirements.txt
@@ -10,7 +10,7 @@ opt-einsum==3.3.0
 packaging==21.3
 pandas==2.2.1
 pyyaml==6.0.1
-requests==2.31.0
+requests==2.32.2
 torch==2.0.1
 torchvision==0.15.2
 tqdm==4.66.3

From c29ca55823be00432c0f09fe9bc54d9bffd4eecf Mon Sep 17 00:00:00 2001
From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com>
Date: Mon, 3 Jun 2024 14:47:43 -0700
Subject: [PATCH 53/58] chore(deps): bump apache-airflow from 2.9.0 to 2.9.1 in
 /requirements/extras (#4703)

* chore(deps): bump apache-airflow in /requirements/extras

Bumps [apache-airflow](https://github.com/apache/airflow) from 2.9.0 to 2.9.1.
- [Release notes](https://github.com/apache/airflow/releases)
- [Changelog](https://github.com/apache/airflow/blob/main/RELEASE_NOTES.rst)
- [Commits](https://github.com/apache/airflow/compare/2.9.0...2.9.1)

---
updated-dependencies:
- dependency-name: apache-airflow
  dependency-type: direct:production
...

Signed-off-by: dependabot[bot] <support@github.com>

* Update tox.ini to bump apache-airflow

---------

Signed-off-by: dependabot[bot] <support@github.com>
Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>
Co-authored-by: Kalyani Nikure <110067132+knikure@users.noreply.github.com>
---
 requirements/extras/test_requirements.txt | 2 +-
 tox.ini                                   | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/requirements/extras/test_requirements.txt b/requirements/extras/test_requirements.txt
index ad6322365f..889ff72779 100644
--- a/requirements/extras/test_requirements.txt
+++ b/requirements/extras/test_requirements.txt
@@ -12,7 +12,7 @@ awslogs==0.14.0
 black==24.3.0
 stopit==1.1.2
 # Update tox.ini to have correct version of airflow constraints file
-apache-airflow==2.9.0
+apache-airflow==2.9.1
 apache-airflow-providers-amazon==7.2.1
 attrs>=23.1.0,<24
 fabric==2.6.0
diff --git a/tox.ini b/tox.ini
index 718e968013..6e1f9ce956 100644
--- a/tox.ini
+++ b/tox.ini
@@ -81,7 +81,7 @@ passenv =
 # Can be used to specify which tests to run, e.g.: tox -- -s
 commands =
     python -c "import os; os.system('install-custom-pkgs --install-boto-wheels')"
-    pip install 'apache-airflow==2.9.0' --constraint "https://raw.githubusercontent.com/apache/airflow/constraints-2.9.0/constraints-3.8.txt"
+    pip install 'apache-airflow==2.9.1' --constraint "https://raw.githubusercontent.com/apache/airflow/constraints-2.9.1/constraints-3.8.txt"
     pip install 'torch==2.0.1+cpu' -f 'https://download.pytorch.org/whl/torch_stable.html'
     pip install 'torchvision==0.15.2+cpu' -f 'https://download.pytorch.org/whl/torch_stable.html'
     pip install 'dill>=0.3.8'

From 864eb7190e2d0a7a97a77e8856ccb5060b167cb4 Mon Sep 17 00:00:00 2001
From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com>
Date: Mon, 3 Jun 2024 14:47:58 -0700
Subject: [PATCH 54/58] chore(deps): bump mlflow from 2.10.2 to 2.12.1 in
 /tests/data/serve_resources/mlflow/pytorch (#4690)

Bumps [mlflow](https://github.com/mlflow/mlflow) from 2.10.2 to 2.12.1.
- [Release notes](https://github.com/mlflow/mlflow/releases)
- [Changelog](https://github.com/mlflow/mlflow/blob/master/CHANGELOG.md)
- [Commits](https://github.com/mlflow/mlflow/compare/v2.10.2...v2.12.1)

---
updated-dependencies:
- dependency-name: mlflow
  dependency-type: direct:production
...

Signed-off-by: dependabot[bot] <support@github.com>
Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>
---
 tests/data/serve_resources/mlflow/pytorch/requirements.txt | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/data/serve_resources/mlflow/pytorch/requirements.txt b/tests/data/serve_resources/mlflow/pytorch/requirements.txt
index fef7d61013..895e2173bf 100644
--- a/tests/data/serve_resources/mlflow/pytorch/requirements.txt
+++ b/tests/data/serve_resources/mlflow/pytorch/requirements.txt
@@ -1,4 +1,4 @@
-mlflow==2.10.2
+mlflow==2.12.1
 astunparse==1.6.3
 cffi==1.16.0
 cloudpickle==2.2.1

From a6a6cfdb901870ba9bc4f82f78f6ad5d2bbbc2f1 Mon Sep 17 00:00:00 2001
From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com>
Date: Mon, 3 Jun 2024 15:49:19 -0700
Subject: [PATCH 55/58] chore(deps): bump mlflow from 2.11.1 to 2.12.1 in
 /tests/data/serve_resources/mlflow/xgboost (#4692)

Bumps [mlflow](https://github.com/mlflow/mlflow) from 2.11.1 to 2.12.1.
- [Release notes](https://github.com/mlflow/mlflow/releases)
- [Changelog](https://github.com/mlflow/mlflow/blob/master/CHANGELOG.md)
- [Commits](https://github.com/mlflow/mlflow/compare/v2.11.1...v2.12.1)

---
updated-dependencies:
- dependency-name: mlflow
  dependency-type: direct:production
...

Signed-off-by: dependabot[bot] <support@github.com>
Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>
---
 tests/data/serve_resources/mlflow/xgboost/requirements.txt | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/data/serve_resources/mlflow/xgboost/requirements.txt b/tests/data/serve_resources/mlflow/xgboost/requirements.txt
index 8150c9fedf..18d687aec6 100644
--- a/tests/data/serve_resources/mlflow/xgboost/requirements.txt
+++ b/tests/data/serve_resources/mlflow/xgboost/requirements.txt
@@ -1,4 +1,4 @@
-mlflow==2.11.1
+mlflow==2.12.1
 lz4==4.3.2
 numpy==1.24.4
 pandas==2.0.3

From 2b3571729cf7c191d284fa0a9f5d8ad4e7b4801e Mon Sep 17 00:00:00 2001
From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com>
Date: Mon, 3 Jun 2024 16:39:43 -0700
Subject: [PATCH 56/58] chore(deps): bump mlflow from 2.11.1 to 2.12.1 in
 /tests/data/serve_resources/mlflow/tensorflow (#4691)

Bumps [mlflow](https://github.com/mlflow/mlflow) from 2.11.1 to 2.12.1.
- [Release notes](https://github.com/mlflow/mlflow/releases)
- [Changelog](https://github.com/mlflow/mlflow/blob/master/CHANGELOG.md)
- [Commits](https://github.com/mlflow/mlflow/compare/v2.11.1...v2.12.1)

---
updated-dependencies:
- dependency-name: mlflow
  dependency-type: direct:production
...

Signed-off-by: dependabot[bot] <support@github.com>
Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>
---
 tests/data/serve_resources/mlflow/tensorflow/requirements.txt | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/data/serve_resources/mlflow/tensorflow/requirements.txt b/tests/data/serve_resources/mlflow/tensorflow/requirements.txt
index 2ff55b8e87..d4ff5b4782 100644
--- a/tests/data/serve_resources/mlflow/tensorflow/requirements.txt
+++ b/tests/data/serve_resources/mlflow/tensorflow/requirements.txt
@@ -1,4 +1,4 @@
-mlflow==2.11.1
+mlflow==2.12.1
 cloudpickle==2.2.1
 numpy==1.26.4
 tensorflow==2.16.1

From cfe98a96a535b771685ce0001632f7904031c5a8 Mon Sep 17 00:00:00 2001
From: Tyler Osterberg <tyoster@amazon.com>
Date: Tue, 4 Jun 2024 21:08:02 -0700
Subject: [PATCH 57/58] change: Updates for DJL 0.28.0 release (#4701)

---
 src/sagemaker/image_uri_config/djl-lmi.json   | 39 +++++++++++++++++++
 .../image_uri_config/djl-neuronx.json         | 18 +++++++++
 .../image_uri_config/djl-tensorrtllm.json     | 32 +++++++++++++++
 3 files changed, 89 insertions(+)
 create mode 100644 src/sagemaker/image_uri_config/djl-lmi.json

diff --git a/src/sagemaker/image_uri_config/djl-lmi.json b/src/sagemaker/image_uri_config/djl-lmi.json
new file mode 100644
index 0000000000..9d2cdd699a
--- /dev/null
+++ b/src/sagemaker/image_uri_config/djl-lmi.json
@@ -0,0 +1,39 @@
+{
+    "scope": [
+        "inference"
+    ],
+    "versions": {
+        "0.28.0": {
+            "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": "djl-inference",
+            "tag_prefix": "0.28.0-lmi10.0.0-cu124"
+        }
+    }
+}
\ No newline at end of file
diff --git a/src/sagemaker/image_uri_config/djl-neuronx.json b/src/sagemaker/image_uri_config/djl-neuronx.json
index a63acc87e4..6038946e28 100644
--- a/src/sagemaker/image_uri_config/djl-neuronx.json
+++ b/src/sagemaker/image_uri_config/djl-neuronx.json
@@ -3,6 +3,24 @@
         "inference"
     ],
     "versions": {
+        "0.28.0": {
+            "registries": {
+                "ap-northeast-1": "763104351884",
+                "ap-south-1": "763104351884",
+                "ap-southeast-1": "763104351884",
+                "ap-southeast-2": "763104351884",
+                "eu-central-1": "763104351884",
+                "eu-west-1": "763104351884",
+                "eu-west-3": "763104351884",
+                "sa-east-1": "763104351884",
+                "us-east-1": "763104351884",
+                "us-east-2": "763104351884",
+                "us-west-2": "763104351884",
+                "ca-west-1": "204538143572"
+            },
+            "repository": "djl-inference",
+            "tag_prefix": "0.28.0-neuronx-sdk2.18.2"
+        },
          "0.27.0": {
             "registries": {
                 "ap-northeast-1": "763104351884",
diff --git a/src/sagemaker/image_uri_config/djl-tensorrtllm.json b/src/sagemaker/image_uri_config/djl-tensorrtllm.json
index e125cbd419..6cde6109bb 100644
--- a/src/sagemaker/image_uri_config/djl-tensorrtllm.json
+++ b/src/sagemaker/image_uri_config/djl-tensorrtllm.json
@@ -3,6 +3,38 @@
         "inference"
     ],
     "versions": {
+        "0.28.0": {
+            "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": "djl-inference",
+            "tag_prefix": "0.28.0-tensorrtllm0.9.0-cu122"
+        },
         "0.27.0": {
             "registries": {
                 "af-south-1": "626614931356",

From 73ae14afd3b13384204f70eee8d143d1efb2c2d7 Mon Sep 17 00:00:00 2001
From: Jonathan Makunga <makung@amazon.com>
Date: Wed, 5 Jun 2024 08:35:15 -0700
Subject: [PATCH 58/58] Sync Branch

---
 tests/unit/sagemaker/jumpstart/test_notebook_utils.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/tests/unit/sagemaker/jumpstart/test_notebook_utils.py b/tests/unit/sagemaker/jumpstart/test_notebook_utils.py
index 8544de709a..301afe4d53 100644
--- a/tests/unit/sagemaker/jumpstart/test_notebook_utils.py
+++ b/tests/unit/sagemaker/jumpstart/test_notebook_utils.py
@@ -1,4 +1,6 @@
 from __future__ import absolute_import
+
+import datetime
 import json
 
 from unittest import TestCase