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