From 8eddfc4a81eea09bfca052b2c3e465f0fa1b00ea Mon Sep 17 00:00:00 2001 From: Mohamed Zeidan Date: Thu, 15 Oct 2026 15:32:26 -0700 Subject: [PATCH] fix: accept core shapes and omitted content types in ModelBuilder.register ModelBuilder.register() had two independent AttributeError defects on the same registry/deploy-from-package workflow. 1. register() rejected the ModelMetrics type its own signature declares. The signature types model_metrics (and metadata_properties, drift_check_baselines, model_life_cycle) as sagemaker.core.shapes classes, but get_model_package_args() converted them by calling value._to_request_dict(). The core shapes are pydantic models that do not implement that method, so the documented type raised "AttributeError: 'ModelMetrics' object has no attribute '_to_request_dict'". Only the older sagemaker.core.model_metrics compatibility class, which is not the declared type, has the method. get_model_package_args() now routes these four arguments through a small _to_request_dict() helper that keeps using _to_request_dict() when the value implements it and otherwise falls back to the core serializer, sagemaker.core.utils.utils.serialize(). The compatibility classes are unaffected, so they stay supported for the rest of their deprecation period, and dicts already in request form pass through. Replacing the call with serialize() unconditionally, as suggested on the issue, was rejected: the compatibility classes carry flat fields (model_statistics, model_constraints, ...) that serialize() renders as ModelStatistics/ModelConstraints instead of the nested ModelQuality.Statistics block the API expects, which would have broken the path that works today. 2. register() failed when content_types and response_types were omitted. Both are declared optional with None defaults, but the method only assigned self.content_types / self.response_types when the arguments were not None and then read those attributes unconditionally, raising "AttributeError: 'ModelBuilder' object has no attribute 'content_types'" on a first registration that omits them. The read now uses getattr(self, ..., None). This keeps the existing behaviour of reusing values supplied to an earlier register() call, and avoids adding new ModelBuilder dataclass fields next to the unrelated, already-public content_type field. Fixes #6479 Fixes #6480 --- .../src/sagemaker/core/model_registry.py | 29 ++- .../tests/unit/test_model_registry.py | 103 +++++++++ .../src/sagemaker/serve/model_builder.py | 12 +- .../tests/unit/test_model_builder_register.py | 217 ++++++++++++++++++ 4 files changed, 353 insertions(+), 8 deletions(-) create mode 100644 sagemaker-serve/tests/unit/test_model_builder_register.py diff --git a/sagemaker-core/src/sagemaker/core/model_registry.py b/sagemaker-core/src/sagemaker/core/model_registry.py index cee6c2a28f..d72141ac33 100644 --- a/sagemaker-core/src/sagemaker/core/model_registry.py +++ b/sagemaker-core/src/sagemaker/core/model_registry.py @@ -15,12 +15,33 @@ MODEL_PACKAGE_VALIDATION_PROFILES_PATH, ) from sagemaker.core.resources import ModelPackageModelCard +from sagemaker.core.utils.utils import serialize from botocore.exceptions import ClientError import logging logger = LOGGER = logging.getLogger("sagemaker") +def _to_request_dict(value): + """Convert a model registry argument into its CreateModelPackage request shape. + + Supports both the modern ``sagemaker.core.shapes`` pydantic shapes, which are + converted with the core serializer, and the older compatibility classes in + ``sagemaker.core.model_metrics`` and friends, which implement their own + ``_to_request_dict()``. Plain dicts are assumed to already be in request form + and are passed through the serializer unchanged. + + Args: + value: A core shape, a compatibility class instance, or a request dict. + + Returns: + dict: The serialized request shape. + """ + if hasattr(value, "_to_request_dict"): + return value._to_request_dict() + return serialize(value) + + def get_model_package_args( content_types=None, response_types=None, @@ -76,11 +97,11 @@ def get_model_package_args( if model_package_group_name is not None: model_package_args["model_package_group_name"] = model_package_group_name if model_metrics is not None: - model_package_args["model_metrics"] = model_metrics._to_request_dict() + model_package_args["model_metrics"] = _to_request_dict(model_metrics) if drift_check_baselines is not None: - model_package_args["drift_check_baselines"] = drift_check_baselines._to_request_dict() + model_package_args["drift_check_baselines"] = _to_request_dict(drift_check_baselines) if metadata_properties is not None: - model_package_args["metadata_properties"] = metadata_properties._to_request_dict() + model_package_args["metadata_properties"] = _to_request_dict(metadata_properties) if approval_status is not None: model_package_args["approval_status"] = approval_status if description is not None: @@ -102,7 +123,7 @@ def get_model_package_args( if source_uri is not None: model_package_args["source_uri"] = source_uri if model_life_cycle is not None: - model_package_args["model_life_cycle"] = model_life_cycle._to_request_dict() + model_package_args["model_life_cycle"] = _to_request_dict(model_life_cycle) if model_card is not None: original_req = {} if isinstance(model_card, ModelPackageModelCard): diff --git a/sagemaker-core/tests/unit/test_model_registry.py b/sagemaker-core/tests/unit/test_model_registry.py index 56ec2d1339..5f932762db 100644 --- a/sagemaker-core/tests/unit/test_model_registry.py +++ b/sagemaker-core/tests/unit/test_model_registry.py @@ -158,6 +158,109 @@ def test_get_model_package_args_model_package_model_card(self): "ModelCardStatus": "Approved", } + def test_get_model_package_args_with_core_shape_model_metrics(self): + """Core shapes.ModelMetrics is serialized with the core serializer. + + The public ModelBuilder.register() signature declares model_metrics as + sagemaker.core.shapes.ModelMetrics, which has no _to_request_dict(). + """ + from sagemaker.core.shapes import MetricsSource, ModelMetrics, ModelQuality + + model_metrics = ModelMetrics( + model_quality=ModelQuality( + statistics=MetricsSource( + content_type="application/json", + s3_uri="s3://bucket/evaluation.json", + ) + ) + ) + + args = get_model_package_args( + image_uri="test-image:latest", + model_metrics=model_metrics, + ) + + assert args["model_metrics"] == { + "ModelQuality": { + "Statistics": { + "ContentType": "application/json", + "S3Uri": "s3://bucket/evaluation.json", + } + } + } + + def test_get_model_package_args_with_compatibility_model_metrics(self): + """The older compatibility class keeps using its own _to_request_dict().""" + from sagemaker.core.model_metrics import MetricsSource, ModelMetrics + + model_metrics = ModelMetrics( + model_statistics=MetricsSource( + content_type="application/json", + s3_uri="s3://bucket/evaluation.json", + ) + ) + + args = get_model_package_args( + image_uri="test-image:latest", + model_metrics=model_metrics, + ) + + assert args["model_metrics"]["ModelQuality"] == { + "Statistics": { + "ContentType": "application/json", + "S3Uri": "s3://bucket/evaluation.json", + } + } + + def test_get_model_package_args_with_core_shape_drift_check_baselines(self): + """Core shapes.DriftCheckBaselines is serialized with the core serializer.""" + from sagemaker.core.shapes import ( + DriftCheckBaselines, + DriftCheckModelQuality, + MetricsSource, + ) + + drift_check_baselines = DriftCheckBaselines( + model_quality=DriftCheckModelQuality( + statistics=MetricsSource( + content_type="application/json", + s3_uri="s3://bucket/statistics.json", + ) + ) + ) + + args = get_model_package_args( + image_uri="test-image:latest", + drift_check_baselines=drift_check_baselines, + ) + + assert args["drift_check_baselines"] == { + "ModelQuality": { + "Statistics": { + "ContentType": "application/json", + "S3Uri": "s3://bucket/statistics.json", + } + } + } + + def test_get_model_package_args_with_dict_model_metrics(self): + """A dict already in request form is passed through unchanged.""" + model_metrics = { + "ModelQuality": { + "Statistics": { + "ContentType": "application/json", + "S3Uri": "s3://bucket/evaluation.json", + } + } + } + + args = get_model_package_args( + image_uri="test-image:latest", + model_metrics=model_metrics, + ) + + assert args["model_metrics"] == model_metrics + def test_get_create_model_package_request_minimal(self): """Test get_create_model_package_request with minimal parameters""" request = get_create_model_package_request( diff --git a/sagemaker-serve/src/sagemaker/serve/model_builder.py b/sagemaker-serve/src/sagemaker/serve/model_builder.py index 10a29fa5f8..20f0784081 100644 --- a/sagemaker-serve/src/sagemaker/serve/model_builder.py +++ b/sagemaker-serve/src/sagemaker/serve/model_builder.py @@ -6875,9 +6875,13 @@ def register( Args: content_types (list[str] or list[PipelineVariable]): The supported MIME types - for the input data. + for the input data. When omitted, any value supplied to an earlier + ``register()`` call is reused, otherwise it is left out of the request + (default: None). response_types (list[str] or list[PipelineVariable]): The supported MIME types - for the output data. + for the output data. When omitted, any value supplied to an earlier + ``register()`` call is reused, otherwise it is left out of the request + (default: None). inference_instances (list[str] or list[PipelineVariable]): A list of the instance types that are used to generate inferences in real-time (default: None). transform_instances (list[str] or list[PipelineVariable]): A list of the instance @@ -6986,8 +6990,8 @@ def register( container_def_list = container_def if isinstance(container_def, list) else [container_def] model_pkg_args = get_model_package_args( - self.content_types, - self.response_types, + getattr(self, "content_types", None), + getattr(self, "response_types", None), inference_instances=inference_instances, transform_instances=transform_instances, model_package_name=model_package_name, diff --git a/sagemaker-serve/tests/unit/test_model_builder_register.py b/sagemaker-serve/tests/unit/test_model_builder_register.py new file mode 100644 index 0000000000..f2f4d3a277 --- /dev/null +++ b/sagemaker-serve/tests/unit/test_model_builder_register.py @@ -0,0 +1,217 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"). You +# may not use this file except in compliance with the License. A copy of +# the License is located at +# +# http://aws.amazon.com/apache2.0/ +# +# or in the "license" file accompanying this file. This file is +# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF +# ANY KIND, either express or implied. See the License for the specific +# language governing permissions and limitations under the License. +"""Tests for ModelBuilder.register() argument handling.""" + +from __future__ import absolute_import + +from unittest.mock import MagicMock, patch + +import pytest + +from sagemaker.core.model_metrics import MetricsSource as CompatMetricsSource +from sagemaker.core.model_metrics import ModelMetrics as CompatModelMetrics +from sagemaker.core.shapes import MetricsSource, ModelMetrics, ModelQuality +from sagemaker.serve.model_builder import ModelBuilder + +IMAGE_URI = "123456789012.dkr.ecr.us-west-2.amazonaws.com/img:latest" +MODEL_DATA_URL = "s3://bucket/model.tar.gz" +METRICS_S3_URI = "s3://bucket/evaluation.json" + +EXPECTED_MODEL_METRICS = { + "ModelQuality": { + "Statistics": { + "ContentType": "application/json", + "S3Uri": METRICS_S3_URI, + } + } +} + + +@pytest.fixture(autouse=True) +def _ambient_execution_role(): + """Stand in a resolved serving role. + + ModelBuilder resolves a default serving role at construction when none is + supplied; these tests build with a mock session and no role. + """ + with patch( + "sagemaker.serve.model_builder.resolve_and_validate_role", + return_value="arn:aws:iam::123456789012:role/ambient-execution-role", + ): + yield + + +@pytest.fixture +def builder(): + """A ModelBuilder wired to a mock session, ready to register().""" + session = MagicMock() + # register() is telemetry-wrapped and reads sagemaker_config. + session.sagemaker_config = {} + session.boto_region_name = "us-west-2" + mb = ModelBuilder( + image_uri=IMAGE_URI, + s3_model_data_url=MODEL_DATA_URL, + sagemaker_session=session, + instance_type="ml.m5.large", + ) + mb.framework = None + mb.framework_version = None + return mb + + +def _register(mb, **kwargs): + """Run register() offline and return the args passed to the registry helper.""" + with patch( + "sagemaker.serve.model_builder.create_model_package_from_containers", + return_value={"ModelPackageArn": "arn:aws:sagemaker:us-west-2:1:model-package/g/1"}, + ), patch( + "sagemaker.serve.model_builder.update_container_with_inference_params", + side_effect=lambda **kw: kw.get("container_def"), + ), patch.object( + mb, "_prepare_container_def", return_value={"Image": IMAGE_URI} + ), patch( + "sagemaker.serve.model_builder.get_model_package_args", return_value={} + ) as mock_args: + mb.register(model_package_group_name="test-group", **kwargs) + return mock_args.call_args + + +class TestRegisterContentAndResponseTypes: + """content_types / response_types are optional and must not raise. + + They are declared with None defaults, so omitting them must leave them out + of the request rather than raising AttributeError. + """ + + def test_register_without_content_and_response_types(self, builder): + """Omitting both must succeed and pass None through.""" + call_args = _register(builder) + + assert call_args.args[0] is None + assert call_args.args[1] is None + + def test_register_without_content_and_response_types_leaves_builder_clean(self, builder): + """Omitting both must not invent attributes on the builder.""" + _register(builder) + + assert not hasattr(builder, "content_types") + assert not hasattr(builder, "response_types") + + def test_register_with_content_and_response_types(self, builder): + """Supplying both must forward them to the registry helper.""" + call_args = _register( + builder, + content_types=["application/json"], + response_types=["text/csv"], + ) + + assert call_args.args[0] == ["application/json"] + assert call_args.args[1] == ["text/csv"] + + def test_register_reuses_types_from_earlier_registration(self, builder): + """A later call that omits the types must reuse the earlier values.""" + _register( + builder, + content_types=["application/json"], + response_types=["text/csv"], + ) + + call_args = _register(builder) + + assert call_args.args[0] == ["application/json"] + assert call_args.args[1] == ["text/csv"] + + +class TestRegisterModelMetrics: + """register() must accept the ModelMetrics type its signature declares.""" + + def test_register_with_core_shape_model_metrics(self, builder): + """sagemaker.core.shapes.ModelMetrics is the declared type and must work.""" + model_metrics = ModelMetrics( + model_quality=ModelQuality( + statistics=MetricsSource( + content_type="application/json", + s3_uri=METRICS_S3_URI, + ) + ) + ) + + call_args = _register(builder, model_metrics=model_metrics) + + assert call_args.kwargs["model_metrics"] is model_metrics + + def test_register_with_compatibility_model_metrics(self, builder): + """The older compatibility class must keep working.""" + model_metrics = CompatModelMetrics( + model_statistics=CompatMetricsSource( + content_type="application/json", + s3_uri=METRICS_S3_URI, + ) + ) + + call_args = _register(builder, model_metrics=model_metrics) + + assert call_args.kwargs["model_metrics"] is model_metrics + + +class TestRegisterModelMetricsRequestShape: + """The serialized request must carry the CreateModelPackage ModelMetrics shape.""" + + def _capture_request(self, mb, **kwargs): + """Run register() without stubbing the registry helper and return the request.""" + with patch( + "sagemaker.serve.model_builder.create_model_package_from_containers", + return_value={"ModelPackageArn": "arn:aws:sagemaker:us-west-2:1:model-package/g/1"}, + ) as mock_create, patch( + "sagemaker.serve.model_builder.update_container_with_inference_params", + side_effect=lambda **kw: kw.get("container_def"), + ), patch.object( + mb, "_prepare_container_def", return_value={"Image": IMAGE_URI} + ): + mb.register(model_package_group_name="test-group", **kwargs) + return mock_create.call_args.kwargs + + def test_core_shape_model_metrics_serializes_to_request_shape(self, builder): + """The declared core shape must produce the documented ModelMetrics structure.""" + model_metrics = ModelMetrics( + model_quality=ModelQuality( + statistics=MetricsSource( + content_type="application/json", + s3_uri=METRICS_S3_URI, + ) + ) + ) + + request = self._capture_request(builder, model_metrics=model_metrics) + + assert request["model_metrics"] == EXPECTED_MODEL_METRICS + + def test_compatibility_model_metrics_serializes_to_request_shape(self, builder): + """The compatibility class must keep producing the same ModelQuality block.""" + model_metrics = CompatModelMetrics( + model_statistics=CompatMetricsSource( + content_type="application/json", + s3_uri=METRICS_S3_URI, + ) + ) + + request = self._capture_request(builder, model_metrics=model_metrics) + + assert request["model_metrics"]["ModelQuality"] == EXPECTED_MODEL_METRICS["ModelQuality"] + + def test_omitted_types_are_absent_from_request(self, builder): + """Omitted content/response types must not appear in the request at all.""" + request = self._capture_request(builder) + + assert "content_types" not in request + assert "response_types" not in request