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.

---
X-AI-Prompt: Address t.corp.amazon.com/P535211426 by fixing the SageMaker Python SDK v3 CapabilityEnum so JumpStart hub documents with the CUSTOMIZATION capability parse, and make capability parsing forward-compatible.
X-AI-Tool: MeshClaw/Kiro

Author: Aniruddha Kappagantu

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)

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