From de649e1375f6378bd6c9c296680c07787b4b128e Mon Sep 17 00:00:00 2001 From: Lucas Jia Date: Tue, 13 Oct 2026 13:03:32 -0700 Subject: [PATCH] fix(serve): handle partial base model metadata Normalize optional BaseModel fields, preserve latest Hub content lookup, and safely resolve hosting configurations through an explicit recipe override or unambiguous inference. Prevent unset model metadata from leaking into generated resource tags. Fixes #5667 --- .../deploy_sagemaker_endpoint.rst | 25 +++ migration.md | 6 + .../src/sagemaker/serve/model_builder.py | 168 +++++++++++------ .../test_model_builder_partial_base_model.py | 169 ++++++++++++++++++ 4 files changed, 314 insertions(+), 54 deletions(-) create mode 100644 sagemaker-serve/tests/unit/test_model_builder_partial_base_model.py diff --git a/docs/model_customization/deploy_sagemaker_endpoint.rst b/docs/model_customization/deploy_sagemaker_endpoint.rst index b420830797..0923dbd950 100644 --- a/docs/model_customization/deploy_sagemaker_endpoint.rst +++ b/docs/model_customization/deploy_sagemaker_endpoint.rst @@ -77,6 +77,31 @@ production deployments. endpoint = model_builder.deploy(endpoint_name="my-endpoint") +Select a Deployment Recipe +---------------------------- + +``ModelBuilder`` normally reads the recipe name from the model package's +base-model metadata. If that metadata omits the recipe, it first uses top-level +``HostingConfigs`` from the Hub document. When no top-level configuration +exists, it automatically uses the only recipe that publishes +``HostingConfigs``. If multiple recipes are deployable, pass ``recipe_name`` +explicitly: + +.. code-block:: python + + from sagemaker.serve import ModelBuilder + + model_builder = ModelBuilder( + model=model_package, + recipe_name="my-deployment-recipe", + ) + endpoint = model_builder.build().deploy(endpoint_name="my-endpoint") + +The override is only needed to disambiguate recipe-level hosting +configurations; existing model packages that publish a recipe name or top-level +hosting configuration require no change. + + Reuse Deployed Resources -------------------------- diff --git a/migration.md b/migration.md index 0ea3e9bef1..aeca5a7df4 100644 --- a/migration.md +++ b/migration.md @@ -334,6 +334,12 @@ predictor = model_builder.deploy( ) ``` +For customized model packages, V3 reads the deployment recipe from the package's base-model +metadata. If the recipe is omitted, ``ModelBuilder`` prefers top-level Hub ``HostingConfigs`` and +otherwise infers the recipe only when exactly one recipe is deployable. Pass +``recipe_name=""`` to ``ModelBuilder`` when multiple recipes publish hosting +configurations. + ### 3. Processing Jobs **V2 Processor:** diff --git a/sagemaker-serve/src/sagemaker/serve/model_builder.py b/sagemaker-serve/src/sagemaker/serve/model_builder.py index 10a29fa5f8..5ad3df2dbf 100644 --- a/sagemaker-serve/src/sagemaker/serve/model_builder.py +++ b/sagemaker-serve/src/sagemaker/serve/model_builder.py @@ -57,7 +57,7 @@ ModelCard, ModelPackageModelCard, ) -from sagemaker.core.utils.utils import logger +from sagemaker.core.utils.utils import Unassigned, logger from sagemaker.core.exceptions import CapacityError, UnexpectedStatusException from sagemaker.core.helper import session_helper from sagemaker.core.helper.iam_role_resolver import resolve_and_validate_role @@ -261,6 +261,8 @@ class ModelBuilder(_InferenceRecommenderMixin, _ModelBuilderServers, _ModelBuild network: Network configuration including VPC settings and network isolation. instance_type: EC2 instance type for deployment (e.g., 'ml.m5.large'). mode: Deployment mode (SAGEMAKER_ENDPOINT, LOCAL_CONTAINER, or IN_PROCESS). + recipe_name: Optional recipe name used to select recipe-level hosting configuration when + base model metadata omits it or the Hub document publishes multiple deployable recipes. Note: ModelBuilder returns sagemaker.core.resources.Model and sagemaker.core.resources.Endpoint @@ -450,6 +452,14 @@ class ModelBuilder(_InferenceRecommenderMixin, _ModelBuilderServers, _ModelBuild default=None, metadata={"help": "DEPRECATED: Use configure_for_torchserve() instead"}, ) + recipe_name: Optional[str] = field( + default=None, + metadata={ + "help": "Recipe name override for resolving recipe-level hosting configurations. " + "Use this when base model metadata omits a recipe and the Hub document publishes " + "multiple recipes with hosting configurations." + }, + ) def _create_session_with_region(self): """Create a SageMaker session with the correct region.""" @@ -750,22 +760,47 @@ def _infer_instance_type_from_jumpstart(self) -> str: raise ValueError(error_msg) + @staticmethod + def _optional_base_model_string(value: Any) -> Optional[str]: + """Normalize an optional string from generated ``BaseModel`` metadata.""" + if value is None or isinstance(value, Unassigned) or value == "": + return None + return value + + def _effective_recipe_name(self, base_model: Any) -> Optional[str]: + """Return the public recipe override or the normalized BaseModel recipe.""" + return self._optional_base_model_string( + self.recipe_name + ) or self._optional_base_model_string(getattr(base_model, "recipe_name", None)) + def _fetch_hub_document_for_custom_model(self) -> dict: from sagemaker.core.shapes import BaseModel as CoreBaseModel base_model: CoreBaseModel = ( self._fetch_model_package().inference_specification.containers[0].base_model ) + hub_content_name = self._optional_base_model_string(base_model.hub_content_name) + if hub_content_name is None: + raise ValueError( + "Cannot resolve hosting configuration because the model package's BaseModel " + "does not specify hub_content_name." + ) + hub_content_version = self._optional_base_model_string(base_model.hub_content_version) + # Prefer an explicitly configured hub (e.g. a private hub used for # testing or for sourcing pre-release hosting configs); fall back to the # public JumpStart hub when none is set. hub_name = getattr(self, "hub_name", None) or "SageMakerPublicHub" - hub_content = HubContent.get( - hub_content_type="Model", - hub_name=hub_name, - hub_content_name=base_model.hub_content_name, - hub_content_version=base_model.hub_content_version, - ) + get_kwargs = { + "hub_content_type": "Model", + "hub_name": hub_name, + "hub_content_name": hub_content_name, + } + # Omitting HubContentVersion asks SageMaker for the latest version. Passing + # None, an empty string, or Unassigned would serialize an invalid request. + if hub_content_version is not None: + get_kwargs["hub_content_version"] = hub_content_version + hub_content = HubContent.get(**get_kwargs) return json.loads(hub_content.hub_content_document) def _fetch_hosting_configs_for_custom_model(self) -> dict: @@ -1011,28 +1046,41 @@ def _is_gpu_instance(self, instance_type: str) -> bool: @staticmethod def _extract_hosting_configs_from_hub( - hub_document: Dict[str, Any], recipe_name: str + hub_document: Dict[str, Any], recipe_name: Optional[str] ) -> List[Dict[str, Any]]: - """Discover the raw ``HostingConfigs`` for a recipe from an already-fetched hub document. - - The authoritative hosting configurations live under - ``RecipeCollection[].HostingConfigs``. Falls back to the top-level - ``HostingConfigs`` when the matching recipe entry does not carry its own (some models - publish there instead). Returns ``[]`` when neither is present — callers decide whether an - empty result is an error (fine-tuned) or a signal to try another path (e.g. Nova). - - This is the single source of truth for config discovery so that - :meth:`_resolve_recipe_hosting_configs` (used by list/get/set_deployment_config) and - :meth:`_fetch_and_cache_recipe_config` (used by build) never diverge — in particular the - top-level fallback is honored on BOTH the selection and the build paths. + """Discover raw ``HostingConfigs`` using explicit, top-level, or unique recipe metadata. + + A named recipe is selected first and retains the historical top-level fallback. When no + recipe is available, top-level ``HostingConfigs`` takes precedence. Only if top-level + configs are absent does ModelBuilder infer a recipe, and only when exactly one recipe + publishes non-empty hosting configs. """ - for recipe in hub_document.get("RecipeCollection", []): - if recipe.get("Name") == recipe_name: - configs = recipe.get("HostingConfigs") - if configs: - return configs - break - return hub_document.get("HostingConfigs") or [] + recipe_name = ModelBuilder._optional_base_model_string(recipe_name) + recipes = hub_document.get("RecipeCollection", []) or [] + if recipe_name is not None: + for recipe in recipes: + if recipe.get("Name") == recipe_name: + configs = recipe.get("HostingConfigs") + if configs: + return configs + break + + top_level_configs = hub_document.get("HostingConfigs") + if top_level_configs: + return top_level_configs + if recipe_name is not None: + return [] + + candidates = [recipe for recipe in recipes if recipe.get("HostingConfigs")] + if len(candidates) == 1: + return candidates[0]["HostingConfigs"] + if len(candidates) > 1: + candidate_names = [str(recipe.get("Name") or "") for recipe in candidates] + raise ValueError( + "Multiple recipes publish HostingConfigs: " + f"{', '.join(candidate_names)}. Specify recipe_name when creating ModelBuilder." + ) + return [] def _resolve_recipe_hosting_configs(self) -> List[Dict[str, Any]]: """Return the raw ``HostingConfigs`` list published for this model's recipe. @@ -1052,7 +1100,7 @@ def _resolve_recipe_hosting_configs(self) -> List[Dict[str, Any]]: "Model is not supported for deployment: the model package has no inference " "container to resolve a recipe from." ) - recipe_name = getattr(containers[0].base_model, "recipe_name", None) or "" + recipe_name = self._effective_recipe_name(containers[0].base_model) configs = self._extract_hosting_configs_from_hub(hub_document, recipe_name) if configs: @@ -1286,7 +1334,7 @@ def _fetch_and_cache_recipe_config(self): hub_document = self._fetch_hub_document_for_custom_model() model_package = self._fetch_model_package() container = model_package.inference_specification.containers[0] - recipe_name = getattr(container.base_model, "recipe_name", None) or "" + recipe_name = self._effective_recipe_name(container.base_model) if not self.s3_upload_path: from sagemaker.serve.utils.model_package_utils import get_s3_uri_from_inference_spec @@ -1546,7 +1594,11 @@ def _base_model_name(self) -> Optional[str]: model_package = self._fetch_model_package() if model_package is not None: base_model = model_package.inference_specification.containers[0].base_model - return getattr(base_model, "hub_content_name", None) if base_model else None + return ( + self._optional_base_model_string(getattr(base_model, "hub_content_name", None)) + if base_model + else None + ) if isinstance(self.model, BaseTrainer): return self.model.base_model_name # Raw S3 checkpoint: base model identity is supplied via model_metadata. @@ -1572,8 +1624,13 @@ def _is_nova_model(self) -> bool: if containers: base_model = getattr(containers[0], "base_model", None) if base_model: - recipe_name = getattr(base_model, "recipe_name", None) or "" - hub_content_name = getattr(base_model, "hub_content_name", None) or "" + recipe_name = self._effective_recipe_name(base_model) or "" + hub_content_name = ( + self._optional_base_model_string( + getattr(base_model, "hub_content_name", None) + ) + or "" + ) if (isinstance(recipe_name, str) and "nova" in recipe_name.lower()) or ( isinstance(hub_content_name, str) and "nova" in hub_content_name.lower() ): @@ -1645,16 +1702,9 @@ def _get_nova_hosting_config_from_hub_document(self, instance_type=None): return None container = self._fetch_model_package().inference_specification.containers[0] - recipe_name = getattr(container.base_model, "recipe_name", None) or "" - - hosting_configs = None - for recipe in hub_document.get("RecipeCollection", []): - if recipe.get("Name") == recipe_name: - hosting_configs = recipe.get("HostingConfigs") - break - if not hosting_configs: - hosting_configs = hub_document.get("HostingConfigs") + recipe_name = self._effective_recipe_name(container.base_model) + hosting_configs = self._extract_hosting_configs_from_hub(hub_document, recipe_name) if not hosting_configs: return None @@ -1690,9 +1740,9 @@ def _get_nova_hosting_config(self, instance_type=None): model_package = self._fetch_model_package() if model_package: - hub_content_name = model_package.inference_specification.containers[ - 0 - ].base_model.hub_content_name + hub_content_name = self._optional_base_model_string( + model_package.inference_specification.containers[0].base_model.hub_content_name + ) else: # No model package (e.g. SMTJ trainer): resolve from base_model_name base_model_name = self._base_model_name() @@ -1747,7 +1797,9 @@ def _validate_nova_smi_config(self) -> None: base_model = getattr(containers[0], "base_model", None) if not base_model: return - hub_content_name = getattr(base_model, "hub_content_name", None) + hub_content_name = self._optional_base_model_string( + getattr(base_model, "hub_content_name", None) + ) if not hub_content_name: return @@ -3286,7 +3338,13 @@ def fetch_endpoint_names_for_base_model(self) -> Set[str]: model_package = self._fetch_model_package() if is_restricted_model_package(model_package): return set() - recipe_name = model_package.inference_specification.containers[0].base_model.recipe_name + base_model = model_package.inference_specification.containers[0].base_model + recipe_name = self._effective_recipe_name(base_model) + if recipe_name is None: + raise ValueError( + "Cannot fetch endpoints for the base model because no recipe is available. " + "Specify recipe_name when creating ModelBuilder." + ) endpoint_names = set() logger.error(f"recipe_name: {recipe_name}") for inference_component in InferenceComponent.get_all(): @@ -6465,8 +6523,8 @@ def _deploy_model_customization( base_model_recipe_name = None if peft_type == "LORA": container = model_package.inference_specification.containers[0] - base_model_recipe_name = getattr( - getattr(container, "base_model", None), "recipe_name", None + base_model_recipe_name = self._effective_recipe_name( + getattr(container, "base_model", None) ) if not base_model_recipe_name: raise ValueError( @@ -6681,7 +6739,7 @@ def _fetch_peft(self) -> Optional[str]: if getattr(container, "is_checkpoint", None) is False: return None recipe_name = ( - getattr(getattr(container, "base_model", None), "recipe_name", "") or "" + self._effective_recipe_name(getattr(container, "base_model", None)) or "" ) if "lora" in recipe_name.lower(): return "LORA" @@ -6766,14 +6824,16 @@ def _deploy_nova_model( # The jumpstart-model-id tag always applies (resolved from the model # package or the trainer's base_model_name). The recipe-name tag is only # available when a model package is present. - tags = [ - {"key": "sagemaker-sdk:jumpstart-model-id", "value": self._base_model_name()}, - ] + tags = [] + base_model_name = self._optional_base_model_string(self._base_model_name()) + if base_model_name is not None: + tags.append({"key": "sagemaker-sdk:jumpstart-model-id", "value": base_model_name}) model_package = self._fetch_model_package() if model_package is not None: base_model = model_package.inference_specification.containers[0].base_model - if base_model is not None and base_model.recipe_name: - tags.append({"key": "sagemaker-sdk:recipe-name", "value": base_model.recipe_name}) + recipe_name = self._effective_recipe_name(base_model) + if recipe_name is not None: + tags.append({"key": "sagemaker-sdk:recipe-name", "value": recipe_name}) # Merge tags accumulated via add_tags (e.g. the model-source reuse tag). # Those are stored in {"Key": ..., "Value": ...} form, so normalize to the diff --git a/sagemaker-serve/tests/unit/test_model_builder_partial_base_model.py b/sagemaker-serve/tests/unit/test_model_builder_partial_base_model.py new file mode 100644 index 0000000000..834f5fd6dc --- /dev/null +++ b/sagemaker-serve/tests/unit/test_model_builder_partial_base_model.py @@ -0,0 +1,169 @@ +"""Regression tests for partially populated BaseModel metadata in ModelBuilder.""" + +import json +import unittest +from dataclasses import fields +from unittest.mock import Mock, patch + +from sagemaker.core.utils.utils import Unassigned +from sagemaker.serve.model_builder import ModelBuilder + + +def _model_package( + *, hub_content_name="test-model", hub_content_version=Unassigned(), recipe_name=Unassigned() +): + base_model = Mock() + base_model.hub_content_name = hub_content_name + base_model.hub_content_version = hub_content_version + base_model.recipe_name = recipe_name + container = Mock() + container.base_model = base_model + model_package = Mock() + model_package.inference_specification.containers = [container] + return model_package + + +def _builder(*, recipe_name=None): + builder = ModelBuilder.__new__(ModelBuilder) + builder.recipe_name = recipe_name + return builder + + +def test_fetch_hub_document_omits_absent_hub_content_version(): + """An absent BaseModel version asks HubContent.get for the latest version.""" + for version in (None, "", Unassigned()): + with unittest.TestCase().subTest(version=version): + builder = _builder() + builder._fetch_model_package = Mock( + return_value=_model_package(hub_content_version=version) + ) + + with patch("sagemaker.serve.model_builder.HubContent.get") as get_hub_content: + get_hub_content.return_value.hub_content_document = json.dumps( + {"HostingConfigs": []} + ) + builder._fetch_hub_document_for_custom_model() + + get_hub_content.assert_called_once_with( + hub_content_type="Model", + hub_name="SageMakerPublicHub", + hub_content_name="test-model", + ) + + +def test_fetch_hub_document_rejects_absent_hub_content_name(): + for hub_content_name in (None, "", Unassigned()): + with unittest.TestCase().subTest(hub_content_name=hub_content_name): + builder = _builder() + builder._fetch_model_package = Mock( + return_value=_model_package(hub_content_name=hub_content_name) + ) + + with patch("sagemaker.serve.model_builder.HubContent.get") as get_hub_content: + with unittest.TestCase().assertRaisesRegex( + ValueError, "does not specify hub_content_name" + ): + builder._fetch_hub_document_for_custom_model() + + get_hub_content.assert_not_called() + + +def test_fetch_hub_document_preserves_explicit_hub_content_version(): + builder = _builder() + builder._fetch_model_package = Mock(return_value=_model_package(hub_content_version="2.0.1")) + + with patch("sagemaker.serve.model_builder.HubContent.get") as get_hub_content: + get_hub_content.return_value.hub_content_document = json.dumps({"HostingConfigs": []}) + builder._fetch_hub_document_for_custom_model() + + get_hub_content.assert_called_once_with( + hub_content_type="Model", + hub_name="SageMakerPublicHub", + hub_content_name="test-model", + hub_content_version="2.0.1", + ) + + +def test_recipe_name_preserves_existing_constructor_parameters(): + """The additive override does not shift any existing positional parameter.""" + init_fields = [field.name for field in fields(ModelBuilder) if field.init] + assert init_fields.index("recipe_name") > init_fields.index("image_config") + + +def test_recipe_absent_prefers_top_level_hosting_configs(): + builder = _builder() + top_level = [{"InstanceType": "ml.top.xlarge"}] + recipe_level = [{"InstanceType": "ml.recipe.xlarge"}] + hub_document = { + "HostingConfigs": top_level, + "RecipeCollection": [{"Name": "recipe-a", "HostingConfigs": recipe_level}], + } + + assert builder._extract_hosting_configs_from_hub(hub_document, None) == top_level + + +def test_recipe_absent_infers_unique_recipe_with_hosting_configs(): + builder = _builder() + expected = [{"InstanceType": "ml.g5.xlarge"}] + hub_document = { + "RecipeCollection": [ + {"Name": "without-hosting"}, + {"Name": "deployable-recipe", "HostingConfigs": expected}, + ] + } + builder._fetch_hub_document_for_custom_model = Mock(return_value=hub_document) + builder._fetch_model_package = Mock(return_value=_model_package()) + + assert builder._resolve_recipe_hosting_configs() == expected + + +def test_recipe_name_override_selects_explicit_recipe(): + builder = _builder(recipe_name="recipe-b") + expected = [{"InstanceType": "ml.p5.48xlarge"}] + hub_document = { + "RecipeCollection": [ + {"Name": "recipe-a", "HostingConfigs": [{"InstanceType": "ml.g5.xlarge"}]}, + {"Name": "recipe-b", "HostingConfigs": expected}, + ] + } + builder._fetch_hub_document_for_custom_model = Mock(return_value=hub_document) + builder._fetch_model_package = Mock(return_value=_model_package(recipe_name=None)) + + assert builder._resolve_recipe_hosting_configs() == expected + + +def test_recipe_absent_rejects_multiple_hosting_candidates(): + builder = _builder() + hub_document = { + "RecipeCollection": [ + {"Name": "recipe-a", "HostingConfigs": [{"InstanceType": "ml.g5.xlarge"}]}, + {"Name": "recipe-b", "HostingConfigs": [{"InstanceType": "ml.p5.48xlarge"}]}, + ] + } + builder._fetch_hub_document_for_custom_model = Mock(return_value=hub_document) + builder._fetch_model_package = Mock(return_value=_model_package(recipe_name="")) + + with unittest.TestCase().assertRaisesRegex(ValueError, "recipe-a.*recipe-b.*recipe_name"): + builder._resolve_recipe_hosting_configs() + + +def test_nova_endpoint_tags_omit_unassigned_recipe_name(): + builder = _builder() + builder.instance_type = "ml.g5.12xlarge" + builder.built_model = Mock(model_name="model-name") + builder._base_model_name = Mock(return_value="nova-textgeneration-lite") + builder._fetch_model_package = Mock(return_value=_model_package()) + builder._tags = None + endpoint = Mock() + + with patch("sagemaker.serve.model_builder.EndpointConfig.create"), patch( + "sagemaker.serve.model_builder.Endpoint.create", return_value=endpoint + ) as create_endpoint: + assert builder._deploy_nova_model("endpoint", wait=False) is endpoint + + assert create_endpoint.call_args.kwargs["tags"] == [ + { + "key": "sagemaker-sdk:jumpstart-model-id", + "value": "nova-textgeneration-lite", + } + ]