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