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

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)

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)

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)

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)

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)

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)

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)

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)