fix: accept core shapes and omitted content types in ModelBuilder.register

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):

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(

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,