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

[Om-singhaI]

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

Issue

Fixes #6191

Description of changes

When a JumpStart model is resolved through a private hub (hub_name=...), the environment variable names stored under HostingInstanceTypeVariants.Variants.<instance>.Properties.EnvironmentVariables arrive mangled. SM_VLLM_MAX_MODEL_LEN becomes s_m__v_l_l_m__m_a_x__m_o_d_e_l__l_e_n, so the instance specific override no longer collides with the base default and both end up in the container environment. reported on 2.245.0 and confirmed by AWS Premium Support on 2.257.6 and 3.20.0 as well and noted that the public hub from_json path performs no key conversion, which is why the same model deploys correctly from the public catalog.

Root cause

Private hub documents are converted between UpperCamelCase and snake_case several times while the JumpStartModelSpecs object is built. For a document with a top level HostingInstanceTypeVariants block the variants go through four passes:

1. HubModelDocument.from_json constructs JumpStartInstanceTypeVariants(..., is_hub_content=True), whose from_describe_hub_content_response runs walk_and_apply_json(response, camel_to_snake) (types.py, the copy of the helper imported from common_utils).
2. hub/parsers.py _to_json serialises every data holder in the spec and walks it back with snake_to_upper_camel (the copy of the helper in hub/parser_utils.py).
3. JumpStartModelSpecs(..., is_hub_content=True) walks the whole spec with camel_to_snake again in JumpStartMetadataBaseFields.from_json (the base class that JumpStartModelSpecs inherits).
4. The variants are rebuilt through JumpStartInstanceTypeVariants(..., is_hub_content=True), which walks them with camel_to_snake once more.

Variants that live inside InferenceConfigComponents only see camel_to_snake passes, which is why the issue reports the double underscore form (s_m__v_l_l_m__...); the top level block additionally round trips through UpperCamelCase and ends up as s_m_v_l_l_m_m_a_x_m_o_d_e_l_l_e_n. Either way the names are destroyed.

walk_and_apply_json already has a stop_keys list (default ["metrics"]) so that metric definitions are left untouched, but it compared only the converted key against the list. That means one list cannot stop both conversion directions: the snake_to_upper_camel pass would need EnvironmentVariables while the camel_to_snake passes need environment_variables. Adding a stop key at a single call site is therefore not enough; I verified that patching only the call at JumpStartInstanceTypeVariants.from_describe_hub_content_response still produces mangled names after the remaining passes.

Fix

In walk_and_apply_json, treat a key as a stop key if either its original or its converted form is in stop_keys, and add environment_variables to the default stop list. With that, every pass in both directions leaves the children of EnvironmentVariables / environment_variables verbatim, while sibling keys such as ImageUri and ResourceRequirements are still converted as before.

Metric definitions behave as they did for real documents: their children were already kept verbatim by the first camel_to_snake pass, and the only keys that appear under a metrics block are Name and Regex, for which the UpperCamelCase conversion is the identity. Stopping on the original metrics key during the snake_to_upper_camel pass additionally protects any multiword child key there (the old code would have turned a hypothetical MetricName into Metricname), which is the intended meaning of the stop list anyway.

The helper exists twice, verbatim, in sagemaker/core/common_utils.py and in sagemaker/core/jumpstart/hub/parser_utils.py. types.py imports the common_utils copy while hub/interfaces.py and hub/parsers.py import the parser_utils copy, so both copies receive the identical change. I did not make one delegate to the other: common_utils importing from sagemaker.core.jumpstart.hub would pull in the jumpstart package init, which imports types.py, which imports common_utils. Consolidating the two helpers would be a reasonable follow up but is out of scope here.

The stop_keys parameter keeps its signature, and stop_keys=None still converts everything. Two edge cases of an explicitly passed list do change: because the original key is now compared as well, a list that happens to name an unconverted key form now stops there too, and an empty list now converts every key (the old condition was falsy for an empty list, so it converted the top level keys but left nested children untouched). No code in sagemaker-core, sagemaker-serve, sagemaker-train or sagemaker-mlops passes stop_keys; every caller relies on the default.

Reproduction on master

Building a DescribeHubContentResponse from sagemaker-core/tests/unit/jumpstart/hub_content_document.json (with SM_VLLM_MAX_MODEL_LEN added to the ml.g5.12xlarge variant) and calling make_model_specs_from_describe_hub_content_response on master gives:

hosting_instance_type_variants.variants["ml.g5.12xlarge"]["properties"]["environment_variables"]
  == {'s_m__n_u_m__g_p_u_s': '4', 's_m__v_l_l_m__m_a_x__m_o_d_e_l__l_e_n': '4096'}
get_instance_specific_environment_variables("ml.g5.12xlarge")
  == {'s_m__n_u_m__g_p_u_s': '4', 's_m__v_l_l_m__m_a_x__m_o_d_e_l__l_e_n': '4096'}

With this change the same call returns {'SM_NUM_GPUS': '4', 'SM_VLLM_MAX_MODEL_LEN': '4096'}.

The same fixture also reproduces the user facing symptom without any modification: it declares a default SM_NUM_GPUS=4 in InferenceEnvironmentVariables and an override SM_NUM_GPUS=8 on the ml.g5.48xlarge variant. Feeding the parsed spec to _retrieve_default_environment_variables(instance_type="ml.g5.48xlarge") on master yields both SM_NUM_GPUS: '4' and s_m__n_u_m__g_p_u_s: '8'; with this change it yields a single SM_NUM_GPUS: '8'.

Testing done

New unit tests:

• sagemaker-core/tests/unit/jumpstart/hub/test_parser_utils.py (new file, 5 tests): camel_to_snake pass, snake_to_upper_camel pass and a four pass round trip all keep environment variable names verbatim while still converting sibling keys; explicit stop_keys and stop_keys=None keep their existing behaviour.
• sagemaker-core/tests/unit/jumpstart/hub/test_parsers.py (2 tests added to TestParsers): parse the real fixture through DescribeHubContentResponse and make_model_specs_from_describe_hub_content_response, once via the InferenceConfigComponents path that the fixture already exercises and once via a top level HostingInstanceTypeVariants block, and assert the names under the variants and from get_instance_specific_environment_variables survive verbatim.
• sagemaker-core/tests/unit/jumpstart/artifacts/test_environment_variables.py (new file, 1 test): mirrors the issue end to end. It parses the fixture the way the JumpStart cache does for hub content, returns that spec from a patched verify_model_region_and_return_specs, and asserts that _retrieve_default_environment_variables for ml.g5.48xlarge contains exactly one SM_NUM_GPUS key holding the instance override 8. On master this assertion fails with ['SM_NUM_GPUS', 's_m__n_u_m__g_p_u_s'].

With the two source files checked out from the parent commit, the three test modules give 6 failed, 20 passed (the six preservation tests fail with the mangled keys; the two stop key behaviour tests pass on both versions). With the fix they give 26 passed.

Surrounding suites (sagemaker-core/tests/unit/jumpstart, test_jumpstart_types.py, test_jumpstart_types_coverage.py, test_jumpstart_types_extended.py, test_jumpstart_utils.py, test_common_utils.py): 952 passed, 4 skipped, 1 failed. The single failure, tests/unit/jumpstart/test_search_unit.py::test_search_public_hub_models, fails identically with the pristine source files (ValueError: Must setup local AWS configuration with a region supported by SageMaker.); it needs a configured AWS region and is unrelated to this change.

black (24.10.0, line length 100 from sagemaker-core/pyproject.toml) leaves parser_utils.py and the three test files unchanged. common_utils.py was already not black clean on master; the changed hunk in it is not part of what black would reformat. flake8 with sagemaker-core/tox.ini reports only pre existing findings in untouched lines (common_utils.py:472 F841, common_utils.py:714 W293, and the unused patch import that test_parsers.py already had on master).

No CHANGELOG entry was added because sagemaker-core/CHANGELOG.md is updated by the release commits.

Notes for reviewers

While reproducing I noticed a separate defect in the same top level variants path: the snake_to_upper_camel round trip is not lossless for instance keys containing digits. g4dn becomes G4Dn and then g4_dn, and ml.g5.12xlarge becomes ml._g5.12_xlarge, so family and instance level lookups on top level HostingInstanceTypeVariants / TrainingInstanceTypeVariants blocks from a private hub miss for those keys (the fixture's training variants show ['g4_dn', 'g5', 'g6', 'g6_e', ..., 'p3_dn', 'p4_d', 'p4_de', ..., 'ml._g4_dn.12_xlarge'] after parsing). That is independent of environment variable names and needs a different fix (the variant keys must stay verbatim while their Properties children are still converted), so it is left for a follow up rather than widening this change.