fix(core): keep environment variable names intact when parsing private hub model documents

Private hub model documents are converted between UpperCamelCase and snake_case
several times while a JumpStartModelSpecs object is built. The hub document
parser snake cases the instance type variants, the hub parsers turn the spec
back into UpperCamelCase, and JumpStartModelSpecs snake cases it again, twice,
when is_hub_content is set. walk_and_apply_json rewrote every key on every
pass, including the environment variable names that are stored as keys under
Variants.Properties.EnvironmentVariables, so SM_VLLM_MAX_MODEL_LEN came out as
s_m__v_l_l_m__m_a_x__m_o_d_e_l__l_e_n and instance specific overrides never
merged over the base defaults. The public hub path performs no key conversion,
which is why the same model deploys correctly from the public catalog.

walk_and_apply_json already has a stop list so that metric definitions are left
alone, but it compared only the converted key, which meant a single list could
not stop both conversion directions. Compare the original key as well and add
environment_variables to the default stop list. The helper exists twice, in
common_utils and in jumpstart.hub.parser_utils, and JumpStartInstanceTypeVariants
uses the common_utils copy, so both copies receive the same change.

Fixes #6191

Author: Om-singhaI

sagemaker-core/src/sagemaker/core/common_utils.py

@@ -2205,19 +2205,24 @@ def camel_to_snake(camel_case_string: str) -> str:
 
 
 def walk_and_apply_json(
-    json_obj: Dict[Any, Any], apply, stop_keys: Optional[List[str]] = ["metrics"]
+    json_obj: Dict[Any, Any],
+    apply,
+    stop_keys: Optional[List[str]] = ["metrics", "environment_variables"],
 ) -> Dict[Any, Any]:
     """Recursively walks a json object and applies a given function to the keys.
 
     stop_keys (Optional[list[str]]): List of field keys that should stop the application function.
         Any children of these keys will not have the application function applied to them.
+        A key stops the walk if either its original or its converted form is in stop_keys, so
+        the same list works for camel_to_snake and snake_to_upper_camel passes. Environment
+        variable names are user facing values stored as keys and must never be converted.
     """
 
     def _walk_and_apply_json(json_obj, new):
         if isinstance(json_obj, dict) and isinstance(new, dict):
             for key, value in json_obj.items():
                 new_key = apply(key)
-                if (stop_keys and new_key not in stop_keys) or stop_keys is None:
+                if stop_keys is None or (key not in stop_keys and new_key not in stop_keys):
                     if isinstance(value, dict):
                         new[new_key] = {}
                         _walk_and_apply_json(value, new=new[new_key])

sagemaker-core/src/sagemaker/core/jumpstart/hub/parser_utils.py

@@ -33,19 +33,24 @@ def snake_to_upper_camel(snake_case_string: str) -> str:
 
 
 def walk_and_apply_json(
-    json_obj: Dict[Any, Any], apply, stop_keys: Optional[List[str]] = ["metrics"]
+    json_obj: Dict[Any, Any],
+    apply,
+    stop_keys: Optional[List[str]] = ["metrics", "environment_variables"],
 ) -> Dict[Any, Any]:
     """Recursively walks a json object and applies a given function to the keys.
 
     stop_keys (Optional[list[str]]): List of field keys that should stop the application function.
         Any children of these keys will not have the application function applied to them.
+        A key stops the walk if either its original or its converted form is in stop_keys, so
+        the same list works for camel_to_snake and snake_to_upper_camel passes. Environment
+        variable names are user facing values stored as keys and must never be converted.
     """
 
     def _walk_and_apply_json(json_obj, new):
         if isinstance(json_obj, dict) and isinstance(new, dict):
             for key, value in json_obj.items():
                 new_key = apply(key)
-                if (stop_keys and new_key not in stop_keys) or stop_keys is None:
+                if stop_keys is None or (key not in stop_keys and new_key not in stop_keys):
                     if isinstance(value, dict):
                         new[new_key] = {}
                         _walk_and_apply_json(value, new=new[new_key])

sagemaker-core/tests/unit/jumpstart/artifacts/test_environment_variables.py

@@ -0,0 +1,82 @@
+# 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.
+"""Unit tests for JumpStart environment variable retrieval."""
+from __future__ import absolute_import
+
+import datetime
+import json
+import os
+from unittest.mock import patch
+
+import pytest
+
+from sagemaker.core.jumpstart.artifacts.environment_variables import (
+    _retrieve_default_environment_variables,
+)
+from sagemaker.core.jumpstart.hub.interfaces import DescribeHubContentResponse
+from sagemaker.core.jumpstart.hub.parsers import (
+    make_model_specs_from_describe_hub_content_response,
+)
+
+HUB_ARN = "arn:aws:sagemaker:us-east-1:123456789012:hub/my-private-hub"
+
+
+@pytest.fixture
+def private_hub_model_specs():
+    """Model specs parsed from a private hub document, the way the JumpStart cache builds them."""
+    path = os.path.join(os.path.dirname(__file__), "..", "hub_content_document.json")
+    with open(path, "r") as f:
+        hub_content_document = json.load(f)
+    response = DescribeHubContentResponse(
+        {
+            "CreationTime": datetime.datetime(2024, 1, 1),
+            "DocumentSchemaVersion": "2.0.0",
+            "HubArn": HUB_ARN,
+            "HubContentArn": (
+                "arn:aws:sagemaker:us-east-1:123456789012:hub-content/"
+                "my-private-hub/Model/meta-textgeneration-llama-2-13b-f/1.0.0"
+            ),
+            "HubContentName": "meta-textgeneration-llama-2-13b-f",
+            "HubContentType": "Model",
+            "HubContentVersion": "1.0.0",
+            "HubContentStatus": "Available",
+            "HubName": "my-private-hub",
+            "HubContentDocument": json.dumps(hub_content_document),
+        }
+    )
+    return make_model_specs_from_describe_hub_content_response(response)
+
+
+@patch(
+    "sagemaker.core.jumpstart.artifacts.environment_variables.verify_model_region_and_return_specs"
+)
+def test_private_hub_instance_specific_environment_variable_overrides_default(
+    mock_verify_model_region_and_return_specs, private_hub_model_specs
+):
+    """An instance specific value must replace the model default rather than being added
+    under a mangled second key, which is what ended up in the container environment before."""
+    mock_verify_model_region_and_return_specs.return_value = private_hub_model_specs
+
+    environment_variables = _retrieve_default_environment_variables(
+        model_id="meta-textgeneration-llama-2-13b-f",
+        model_version="1.0.0",
+        hub_arn=HUB_ARN,
+        region="us-east-1",
+        instance_type="ml.g5.48xlarge",
+        sagemaker_session=None,
+    )
+
+    gpu_keys = [key for key in environment_variables if key.replace("_", "").lower() == "smnumgpus"]
+    assert gpu_keys == ["SM_NUM_GPUS"]
+    assert environment_variables["SM_NUM_GPUS"] == "8"
+    assert mock_verify_model_region_and_return_specs.call_args.kwargs["hub_arn"] == HUB_ARN

sagemaker-core/tests/unit/jumpstart/hub/test_parser_utils.py

@@ -0,0 +1,89 @@
+# 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 sagemaker.core.jumpstart.hub.parser_utils import (
+    camel_to_snake,
+    snake_to_upper_camel,
+    walk_and_apply_json,
+)
+
+
+VARIANTS = {
+    "Variants": {
+        "g5": {
+            "Properties": {
+                "ImageUri": "image",
+                "EnvironmentVariables": {
+                    "SM_VLLM_MAX_MODEL_LEN": "131072",
+                    "HF_HUB_OFFLINE": "1",
+                },
+                "Metrics": [{"Name": "loss", "Regex": "loss=(.*)"}],
+            }
+        }
+    }
+}
+
+
+def test_walk_and_apply_json_camel_to_snake_preserves_environment_variable_names():
+    result = walk_and_apply_json(VARIANTS, camel_to_snake)
+
+    properties = result["variants"]["g5"]["properties"]
+    assert properties["image_uri"] == "image"
+    assert properties["environment_variables"] == {
+        "SM_VLLM_MAX_MODEL_LEN": "131072",
+        "HF_HUB_OFFLINE": "1",
+    }
+    assert properties["metrics"] == [{"Name": "loss", "Regex": "loss=(.*)"}]
+
+
+def test_walk_and_apply_json_snake_to_upper_camel_preserves_environment_variable_names():
+    snake = walk_and_apply_json(VARIANTS, camel_to_snake)
+
+    result = walk_and_apply_json(snake, snake_to_upper_camel)
+
+    properties = result["Variants"]["G5"]["Properties"]
+    assert properties["ImageUri"] == "image"
+    assert properties["EnvironmentVariables"] == {
+        "SM_VLLM_MAX_MODEL_LEN": "131072",
+        "HF_HUB_OFFLINE": "1",
+    }
+
+
+def test_walk_and_apply_json_round_trip_preserves_environment_variable_names():
+    """The private hub pipeline converts the same document several times in both directions."""
+    result = walk_and_apply_json(VARIANTS, camel_to_snake)
+    result = walk_and_apply_json(result, snake_to_upper_camel)
+    result = walk_and_apply_json(result, camel_to_snake)
+    result = walk_and_apply_json(result, camel_to_snake)
+
+    assert result["variants"]["g5"]["properties"]["environment_variables"] == {
+        "SM_VLLM_MAX_MODEL_LEN": "131072",
+        "HF_HUB_OFFLINE": "1",
+    }
+
+
+def test_walk_and_apply_json_explicit_stop_keys_still_honored():
+    result = walk_and_apply_json(
+        {"Outer": {"Inner": {"KeepMe": 1}}}, camel_to_snake, stop_keys=["inner"]
+    )
+
+    assert result == {"outer": {"inner": {"KeepMe": 1}}}
+
+
+def test_walk_and_apply_json_no_stop_keys_converts_everything():
+    result = walk_and_apply_json(VARIANTS, camel_to_snake, stop_keys=None)
+
+    properties = result["variants"]["g5"]["properties"]
+    assert "SM_VLLM_MAX_MODEL_LEN" not in properties["environment_variables"]
+    assert properties["metrics"] == [{"name": "loss", "regex": "loss=(.*)"}]

sagemaker-core/tests/unit/jumpstart/hub/test_parsers.py

@@ -11,6 +11,10 @@
 # ANY KIND, either express or implied. See the License for the specific
 # language governing permissions and limitations under the License.
 
+import copy
+import datetime
+import json
+import os
 import pytest
 from unittest.mock import Mock, patch
 from sagemaker.core.jumpstart.hub.parsers import (
@@ -367,3 +371,87 @@ def test_make_model_specs_from_describe_hub_content_response_with_payloads(self)
         result = make_model_specs_from_describe_hub_content_response(response)
 
         assert result is not None
+
+    def test_make_model_specs_preserves_instance_variant_environment_variable_names(
+        self, hub_content_document
+    ):
+        """Environment variable names under HostingInstanceTypeVariants must survive verbatim."""
+        variants = hub_content_document["InferenceConfigComponents"]["tgi"][
+            "HostingInstanceTypeVariants"
+        ]["Variants"]
+        variants["ml.g5.12xlarge"]["Properties"]["EnvironmentVariables"][
+            "SM_VLLM_MAX_MODEL_LEN"
+        ] = "4096"
+        variants["g5"]["Properties"]["EnvironmentVariables"] = {"HF_HUB_OFFLINE": "1"}
+
+        specs = make_model_specs_from_describe_hub_content_response(
+            _describe_hub_content_response(hub_content_document)
+        )
+
+        instance_type_variants = specs.hosting_instance_type_variants
+        assert instance_type_variants.variants["ml.g5.12xlarge"]["properties"][
+            "environment_variables"
+        ] == {"SM_NUM_GPUS": "4", "SM_VLLM_MAX_MODEL_LEN": "4096"}
+        assert instance_type_variants.get_instance_specific_environment_variables(
+            "ml.g5.12xlarge"
+        ) == {
+            "HF_HUB_OFFLINE": "1",
+            "SM_NUM_GPUS": "4",
+            "SM_VLLM_MAX_MODEL_LEN": "4096",
+        }
+
+    def test_make_model_specs_preserves_top_level_variant_environment_variable_names(
+        self, hub_content_document
+    ):
+        """Top level HostingInstanceTypeVariants go through one more conversion pass than
+        inference config components and must also keep environment variable names verbatim."""
+        document = copy.deepcopy(hub_content_document)
+        for key in ("InferenceConfigs", "InferenceConfigComponents", "InferenceConfigRankings"):
+            document.pop(key, None)
+        document["HostingInstanceTypeVariants"] = {
+            "Variants": {
+                "g5": {
+                    "Properties": {
+                        "ImageUri": "image",
+                        "EnvironmentVariables": {"SM_VLLM_MAX_MODEL_LEN": "131072"},
+                    }
+                }
+            }
+        }
+
+        specs = make_model_specs_from_describe_hub_content_response(
+            _describe_hub_content_response(document)
+        )
+
+        instance_type_variants = specs.hosting_instance_type_variants
+        assert instance_type_variants.variants["g5"]["properties"]["image_uri"] == "image"
+        assert instance_type_variants.get_instance_specific_environment_variables(
+            "ml.g5.12xlarge"
+        ) == {"SM_VLLM_MAX_MODEL_LEN": "131072"}
+
+
+def _describe_hub_content_response(hub_content_document):
+    return DescribeHubContentResponse(
+        {
+            "CreationTime": datetime.datetime(2024, 1, 1),
+            "DocumentSchemaVersion": "2.0.0",
+            "HubArn": "arn:aws:sagemaker:us-east-1:123456789012:hub/my-private-hub",
+            "HubContentArn": (
+                "arn:aws:sagemaker:us-east-1:123456789012:hub-content/"
+                "my-private-hub/Model/meta-textgeneration-llama-2-13b-f/1.0.0"
+            ),
+            "HubContentName": "meta-textgeneration-llama-2-13b-f",
+            "HubContentType": "Model",
+            "HubContentVersion": "1.0.0",
+            "HubContentStatus": "Available",
+            "HubName": "my-private-hub",
+            "HubContentDocument": json.dumps(hub_content_document),
+        }
+    )
+
+
+@pytest.fixture
+def hub_content_document():
+    path = os.path.join(os.path.dirname(__file__), "..", "hub_content_document.json")
+    with open(path, "r") as f:
+        return json.load(f)