From f7b7d843e650d06e16aa8ef38cfb1fd30806a2a0 Mon Sep 17 00:00:00 2001 From: Aniruddha Kappagantu Date: Tue, 13 Oct 2026 15:46:31 +0000 Subject: [PATCH] fix: accept CUSTOMIZATION capability in JumpStart hub documents CapabilityEnum was missing the CUSTOMIZATION member, so pydantic rejected the entire HubContentDocument for any JumpStart model whose Capabilities include CUSTOMIZATION. This broke every v3 training path that loads the document (ModelTrainer.from_jumpstart_config and the JumpStartTrainDefaults getters) for 42 of 754 public models in us-west-2, including three with TrainingSupported=true. Add CUSTOMIZATION to CapabilityEnum, and make Capabilities parsing forward-compatible: unknown values are kept as plain strings with a warning instead of failing validation, so a future service-side capability cannot break a pinned SDK version again. Applied to both the HubContentDocument field and the HostingComponentsModel base-class field. --- .../src/sagemaker/core/jumpstart/models.py | 47 ++++++++++++++- .../tests/unit/jumpstart/test_models.py | 57 ++++++++++++++++++- 2 files changed, 100 insertions(+), 4 deletions(-) diff --git a/sagemaker-core/src/sagemaker/core/jumpstart/models.py b/sagemaker-core/src/sagemaker/core/jumpstart/models.py index 372498e98c..10dff537cd 100644 --- a/sagemaker-core/src/sagemaker/core/jumpstart/models.py +++ b/sagemaker-core/src/sagemaker/core/jumpstart/models.py @@ -14,12 +14,15 @@ from __future__ import absolute_import, annotations +import logging from enum import Enum from typing import Any, Dict, List, Optional, Union -from pydantic import Field, ConfigDict +from pydantic import Field, ConfigDict, field_validator from sagemaker.core.jumpstart.configs import BaseConfig +logger = logging.getLogger(__name__) + class StrEnum(str, Enum): """A string-valued enumeration.""" @@ -406,6 +409,31 @@ class CapabilityEnum(StrEnum): FINE_TUNING = "FINE_TUNING" VALIDATION = "VALIDATION" INCREMENTAL_TRAINING = "INCREMENTAL_TRAINING" + CUSTOMIZATION = "CUSTOMIZATION" + + +def _coerce_capabilities(value: Any) -> Any: + """Coerce ``Capabilities`` values, tolerating unknown entries. + + Known values are converted to ``CapabilityEnum`` members. Unknown values + are kept as plain strings with a warning, so a new service-side capability + cannot fail validation of the entire ``HubContentDocument`` for a pinned + SDK version. + """ + if not isinstance(value, list): + return value + coerced = [] + for item in value: + try: + coerced.append(CapabilityEnum(item)) + except ValueError: + logger.warning( + "Unknown capability %r is not a member of CapabilityEnum; " + "keeping it as a string.", + item, + ) + coerced.append(item) + return coerced class DemoNotebookModel(BaseConfig): @@ -444,7 +472,14 @@ class HostingComponentsModel(BaseConfig): DynamicContainerDeploymentSupported: Optional[bool] = None HostingAdditionalDataSources: Optional[Dict[str, List[HostingAdditionalDataSourceModel]]] = None BedrockIOMappingId: Optional[str] = None - Capabilities: Optional[List[CapabilityEnum]] = None + Capabilities: Optional[List[Union[CapabilityEnum, str]]] = None + + @field_validator("Capabilities", mode="before") + @classmethod + def _validate_capabilities(cls, value: Any) -> Any: + """Keep unknown capability values as strings instead of failing.""" + return _coerce_capabilities(value) + HostingEcrUri: Optional[str] = None HostingArtifactS3DataType: Optional[HostingArtifactS3DataTypeEnum] = None HostingArtifactCompressionType: Optional[HostingArtifactCompressionTypeEnum] = None @@ -506,4 +541,10 @@ class HubContentDocument(HostingComponentsModel, TrainingComponentsModel): TrainingConfigRankings: Optional[Dict[str, RankingSchemaModel]] = None Resources: Optional[List[ResourceModel]] = None Highlights: Optional[List[str]] = None - Capabilities: List[CapabilityEnum] = None + Capabilities: List[Union[CapabilityEnum, str]] = None + + @field_validator("Capabilities", mode="before") + @classmethod + def _validate_capabilities(cls, value: Any) -> Any: + """Keep unknown capability values as strings instead of failing.""" + return _coerce_capabilities(value) diff --git a/sagemaker-core/tests/unit/jumpstart/test_models.py b/sagemaker-core/tests/unit/jumpstart/test_models.py index 3cdf80afb0..573c5b505f 100644 --- a/sagemaker-core/tests/unit/jumpstart/test_models.py +++ b/sagemaker-core/tests/unit/jumpstart/test_models.py @@ -16,10 +16,21 @@ import json import os -from sagemaker.core.jumpstart.models import HubContentDocument +from sagemaker.core.jumpstart.models import ( + CapabilityEnum, + HostingComponentsModel, + HubContentDocument, +) TEST_HUB_CONTENT_DOCUMENT = "hub_content_document.json" +MINIMAL_DOCUMENT = { + "ModelTypes": ["OPEN_WEIGHTS"], + "Url": "https://example.com", + "TrainingSupported": True, + "IncrementalTrainingSupported": False, +} + def test_init_hub_content_document(): """Test HubContentDocument initialization.""" @@ -28,3 +39,47 @@ def test_init_hub_content_document(): hub_content_document = json.load(f) hub_content_document_instance = HubContentDocument(**hub_content_document) assert isinstance(hub_content_document_instance, HubContentDocument) + + +def test_capability_enum_has_customization(): + """CUSTOMIZATION is a valid, documented capability value.""" + + assert CapabilityEnum.CUSTOMIZATION.value == "CUSTOMIZATION" + + +def test_hub_content_document_parses_customization_capability(): + """A document whose Capabilities include CUSTOMIZATION parses successfully.""" + + document = HubContentDocument(**MINIMAL_DOCUMENT, Capabilities=["TRAINING", "CUSTOMIZATION"]) + assert CapabilityEnum.CUSTOMIZATION in document.Capabilities + + +def test_hub_content_document_parses_all_known_capabilities(): + """Every documented CapabilityEnum value parses without error.""" + + known = [capability.value for capability in CapabilityEnum] + document = HubContentDocument(**MINIMAL_DOCUMENT, Capabilities=known) + assert all(isinstance(capability, CapabilityEnum) for capability in document.Capabilities) + assert len(document.Capabilities) == len(known) + + +def test_hub_content_document_tolerates_unknown_capability(caplog): + """An unknown capability is kept as a string with a warning, not rejected.""" + + with caplog.at_level("WARNING", logger="sagemaker.core.jumpstart.models"): + document = HubContentDocument( + **MINIMAL_DOCUMENT, Capabilities=["TRAINING", "SOME_FUTURE_CAPABILITY"] + ) + + assert "SOME_FUTURE_CAPABILITY" in document.Capabilities + assert isinstance(document.Capabilities[-1], str) + assert CapabilityEnum.TRAINING in document.Capabilities + assert "SOME_FUTURE_CAPABILITY" in caplog.text + + +def test_hosting_components_model_tolerates_unknown_capability(): + """The base-class Capabilities field also tolerates unknown values.""" + + component = HostingComponentsModel(Capabilities=["CUSTOMIZATION", "SOME_FUTURE_CAPABILITY"]) + assert CapabilityEnum.CUSTOMIZATION in component.Capabilities + assert "SOME_FUTURE_CAPABILITY" in component.Capabilities