fix: accept CUSTOMIZATION capability in JumpStart hub documents

[akappagan]

Issue

SageMaker Python SDK v3 (3.24.0) cannot parse the HubContentDocument of any JumpStart model whose Capabilities include CUSTOMIZATION. CapabilityEnum does not list that value, so pydantic validation rejects the whole document.

Because HubContentDocument.Capabilities is typed List[CapabilityEnum], a single unrecognized value fails validation of the entire document at HubContentDocument(**document_json), breaking every v3 training path that loads the document before any training logic runs — ModelTrainer.from_jumpstart_config and the JumpStartTrainDefaults getters.

In us-west-2 SageMakerPublicHub, 42 of 754 public models carry CUSTOMIZATION, including three with TrainingSupported=true: meta-textgeneration-llama-3-1-8b-instruct, meta-textgeneration-llama-3-2-1b-instruct, meta-textgeneration-llama-3-2-3b-instruct.

CUSTOMIZATION is a valid, documented capability (listed in the HubContentDocument model-properties reference alongside BEDROCK_CONSOLE, TRAINING, FINE_TUNING, HYPERPOD_DEPLOYMENT).

Reproduction

# pip install "sagemaker>=3,<4"   (3.24.0)
from sagemaker.core.jumpstart.configs import JumpStartConfig
from sagemaker.core.jumpstart.document import get_hub_content_and_document
from sagemaker.core.helper.session_helper import Session
get_hub_content_and_document(
    JumpStartConfig(model_id="meta-textgeneration-llama-3-1-8b-instruct"), Session()
)  # us-west-2

pydantic_core.ValidationError: 1 validation error for HubContentDocument
Capabilities.1
  Input should be 'BEDROCK_CONSOLE', 'HYPERPOD_DEPLOYMENT', 'SAGEMAKER_EVALUATE',
  'SAGEMAKER_MODEL_OPTIMIZE', 'TRAINING', 'FINE_TUNING', 'VALIDATION' or
  'INCREMENTAL_TRAINING' [type=enum, input_value='CUSTOMIZATION', input_type=str]

Description of changes

1. Add CUSTOMIZATION to CapabilityEnum in sagemaker-core/src/sagemaker/core/jumpstart/models.py.
2. Make Capabilities parsing forward-compatible: the field type is widened to List[Union[CapabilityEnum, str]] and a mode="before" validator coerces known values to CapabilityEnum members while keeping unknown values as plain strings with a warning. This prevents a future service-side capability value from again failing validation of the entire document for a pinned SDK version. Applied to both HubContentDocument.Capabilities and the HostingComponentsModel.Capabilities base-class field it inherits from.

Known values still deserialize to CapabilityEnum; behavior is unchanged for documents that only use recognized capabilities. This is backwards compatible.

Testing

Added regression tests in sagemaker-core/tests/unit/jumpstart/test_models.py:

• CUSTOMIZATION is a member of CapabilityEnum.
• A document with CUSTOMIZATION parses.
• Every documented CapabilityEnum value parses.
• An unknown capability is kept as a string with a warning instead of failing (on both HubContentDocument and HostingComponentsModel).

All unit tests in tests/unit/jumpstart/test_models.py and tests/unit/jumpstart/test_document.py pass; black (line-length 100) and flake8 are clean.

Merge Checklist

• I have read the CONTRIBUTING doc
• I used the commit message format described in CONTRIBUTING
• I have passed the region in to all S3 and STS clients that I've initialized as part of this change (N/A)
• I have updated any necessary documentation, including READMEs and API docs (N/A)

Tests

• I have added tests that prove my fix is effective or that my feature works
• I have added unit and/or integration tests as appropriate to ensure backward compatibility of the changes

By submitting this pull request, I confirm that you can use, modify, copy, and redistribute this contribution, under the terms of your choice.

[github-actions]

🤖 AI code review

This PR adds CUSTOMIZATION to CapabilityEnum and makes Capabilities parsing forward-compatible by widening the field type to List[Union[CapabilityEnum, str]] and adding mode="before" validators that coerce known values to enum members while retaining unknown values as strings with a warning. The change is well-scoped, backward compatible, and covered by regression tests. The duplicated validator across HostingComponentsModel and HubContentDocument is required by pydantic v2 (a redefined field needs its validator re-declared), so that duplication is appropriate. I found no blocking issues; one minor note on a weak test assertion.

_{Reviewed commit f7b7d84. Automated review; verify before acting.}

[github-actions]

🟡 Low: isinstance(document.Capabilities[-1], str) is a tautology for distinguishing a raw string from an enum member: CapabilityEnum is a StrEnum (subclasses str), so coerced enum members are also str instances and would pass this assertion too. If the intent is to verify the unknown value is not coerced to a CapabilityEnum, assert not isinstance(document.Capabilities[-1], CapabilityEnum) (or type(document.Capabilities[-1]) is str). The in check above still gives real coverage, so this is minor.

[akappagan]

CI status: the 4 red checks are pre-existing or infra, not caused by this change

This PR only adds CUSTOMIZATION to CapabilityEnum and a forward-compatible validator in sagemaker-core/src/sagemaker/core/jumpstart/models.py, plus unit tests. Everything that exercises the change is green; the red checks are unrelated. Evidence per check:

Red check	Actual failure	Caused by this PR?
codestyle-doc-tests (sagemaker-core)	The repo-wide black --check ./ (black 26.3.1, py314 target) reformats tests/unit/local/test_entities.py, a file this PR does not touch. In the same run flake8, pylint, and twine all passed.	No
integ-tests (sagemaker-mlops)	89 passed, 2 failed: tests/integ/workflow/test_deployment_steps.py::test_deployment_steps_execute_end_to_end (expected exactly one EndpointConfigName for this run, got []) and tests/integ/test_transform_job.py::test_transform_job (Transform job Failed) — live-AWS deploy/transform tests in the mlops submodule.	No
integ-tests-us-east-1	Live-AWS integ across submodules; the sampled block shows 11 passed, 1 skipped, the non-zero exit is an AWS-flaky submodule run.	No
fast-integ-tests	sagemaker-train shallow integ against live AWS.	No

Green checks that validate this change:

• unit-tests (sagemaker-core) — runs the new CapabilityEnum / HubContentDocument tests
• integ-tests (sagemaker-core) — the changed submodule, integ level
• codecov/patch 96% of diff (target 70%), codecov/project 72% (target 65%)
• review (public), review (internal), CodeQL, Analyze (python), fortress-scan, docs/readthedocs

Reproduced locally with the CI pinned black==26.3.1: tests/unit/local/test_entities.py → "would reformat"; the two files in this diff → clean. The remaining failures need either a repo-wide formatting sweep of unrelated files or the maintainers' CI environment for the live-AWS integ jobs, so they are out of scope for this change.