fix: overriding instance specific fields in config components (#1478)

* fix: instance specific variables override

* format

Author: Captainia

src/sagemaker/jumpstart/factory/model.py

@@ -677,6 +677,7 @@ def _add_config_name_to_deploy_kwargs(
         kwargs.config_name = _select_inference_config_from_training_config(
             specs=specs, training_config_name=training_config_name
         )
+        return kwargs
 
     if specs.inference_configs:
         default_config_name = specs.inference_configs.get_top_config_from_ranking().config_name

src/sagemaker/utils.py

@@ -1660,8 +1660,6 @@ def nested_set_dict(d: Dict[str, Any], keys: List[str], value: Any) -> None:
     if len(keys) == 1:
         d[key] = value
         return
-    if not d:
-        return
 
     d = d.setdefault(key, {})
     nested_set_dict(d, keys[1:], value)

tests/unit/sagemaker/jumpstart/constants.py

@@ -7758,7 +7758,7 @@
                 "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"
+                        "pytorch-hosting-neuronx:1.13.1-neuronx-py310-sdk2.14.1-ubuntu20.04"
                     }
                 },
                 "variants": {"inf2": {"regional_properties": {"image_uri": "$neuron-ecr-uri"}}},
@@ -7772,7 +7772,7 @@
                 "regional_aliases": {
                     "us-west-2": {
                         "gpu-ecr-uri": "763104351884.dkr.ecr.us-west-2.amazonaws.com/"
-                        "pytorch-hosting-neuronx:1.13.1-neuronx-py310-sdk2.14.1-ubuntu20.04"
+                        "huggingface-pytorch-hosting:2.0.0-transformers4.28.1-gpu-py310-cu118-ubuntu20.04"
                     }
                 },
                 "variants": {
@@ -7796,7 +7796,7 @@
                 "regional_aliases": {
                     "us-west-2": {
                         "gpu-ecr-uri": "763104351884.dkr.ecr.us-west-2.amazonaws.com/"
-                        "pytorch-hosting-neuronx:1.13.1-neuronx-py310-sdk2.14.1-ubuntu20.04"
+                        "pytorch-hosting:1.13.1-py310-sdk2.14.1-ubuntu20.04"
                     }
                 },
                 "variants": {
@@ -7922,7 +7922,7 @@
                 "regional_aliases": {
                     "us-west-2": {
                         "gpu-ecr-uri": "763104351884.dkr.ecr.us-west-2.amazonaws.com/"
-                        "pytorch-hosting-neuronx:1.13.1-neuronx-py310-sdk2.14.1-ubuntu20.04"
+                        "huggingface-pytorch-training:1.13.1-neuronx-py310-sdk2.14.1-ubuntu20.04"
                     }
                 },
                 "variants": {
@@ -7953,7 +7953,7 @@
                 "regional_aliases": {
                     "us-west-2": {
                         "gpu-ecr-uri": "763104351884.dkr.ecr.us-west-2.amazonaws.com/"
-                        "pytorch-hosting-neuronx:1.13.1-neuronx-py310-sdk2.14.1-ubuntu20.04"
+                        "pytorch-training:1.13.1-py310-sdk2.14.1-ubuntu20.04"
                     }
                 },
                 "variants": {

tests/unit/sagemaker/jumpstart/estimator/test_estimator.py

@@ -1953,7 +1953,8 @@ def test_estimator_initialization_with_config_name(
         mock_estimator_init.assert_called_once_with(
             instance_type="ml.p2.xlarge",
             instance_count=1,
-            image_uri="763104351884.dkr.ecr.us-west-2.amazonaws.com/pytorch-training:1.5.0-gpu-py3",
+            image_uri="763104351884.dkr.ecr.us-west-2.amazonaws.com/"
+            "huggingface-pytorch-training:1.13.1-neuronx-py310-sdk2.14.1-ubuntu20.04",
             model_uri="s3://jumpstart-cache-prod-us-west-2/artifacts/meta-textgeneration-llama-2-7b/"
             "gpu-training/model/",
             source_dir="s3://jumpstart-cache-prod-us-west-2/source-directory-tarballs/pytorch/transfer_learning/"
@@ -2005,7 +2006,8 @@ def test_estimator_set_config_name(
         mock_estimator_init.assert_called_with(
             instance_type="ml.p2.xlarge",
             instance_count=1,
-            image_uri="763104351884.dkr.ecr.us-west-2.amazonaws.com/pytorch-training:1.5.0-gpu-py3",
+            image_uri="763104351884.dkr.ecr.us-west-2.amazonaws.com/"
+            "pytorch-training:1.13.1-py310-sdk2.14.1-ubuntu20.04",
             model_uri="s3://jumpstart-cache-prod-us-west-2/artifacts/meta-textgeneration-llama-2-7b/"
             "gpu-training-budget/model/",
             source_dir="s3://jumpstart-cache-prod-us-west-2/source-directory-tarballs/pytorch/"
@@ -2060,7 +2062,8 @@ def test_estimator_default_inference_config(
         mock_estimator_deploy.assert_called_once_with(
             instance_type="ml.p2.xlarge",
             initial_instance_count=1,
-            image_uri="763104351884.dkr.ecr.us-west-2.amazonaws.com/pytorch-inference:1.5.0-gpu-py3",
+            image_uri="763104351884.dkr.ecr.us-west-2.amazonaws.com/huggingface-pytorch-hosting"
+            ":2.0.0-transformers4.28.1-gpu-py310-cu118-ubuntu20.04",
             source_dir="s3://jumpstart-cache-prod-us-west-2/source-directory-tarballs/"
             "pytorch/inference/eqa/v1.0.0/sourcedir.tar.gz",
             entry_point="inference.py",
@@ -2167,7 +2170,8 @@ def test_estimator_deploy_with_config(
         mock_estimator_deploy.assert_called_once_with(
             instance_type="ml.p2.xlarge",
             initial_instance_count=1,
-            image_uri="763104351884.dkr.ecr.us-west-2.amazonaws.com/pytorch-inference:1.5.0-gpu-py3",
+            image_uri="763104351884.dkr.ecr.us-west-2.amazonaws.com/"
+            "pytorch-hosting:1.13.1-py310-sdk2.14.1-ubuntu20.04",
             source_dir="s3://jumpstart-cache-prod-us-west-2/source-directory-tarballs/"
             "pytorch/inference/eqa/v1.0.0/sourcedir.tar.gz",
             entry_point="inference.py",

tests/unit/sagemaker/jumpstart/test_types.py

@@ -1143,7 +1143,7 @@ def test_inference_configs_parsing():
                 "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"
+                        "pytorch-hosting-neuronx:1.13.1-neuronx-py310-sdk2.14.1-ubuntu20.04"
                     }
                 },
                 "variants": {"inf2": {"regional_properties": {"image_uri": "$neuron-ecr-uri"}}},