fix: accept CUSTOMIZATION capability in JumpStart hub documents

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