fix: preserve all core container fields in ModelBuilder CreateModel requests

sagemaker-serve/src/sagemaker/serve/model_builder.py

@@ -3082,38 +3082,42 @@ def _prepare_pipeline_container_defs(self):
         return containers
 
     def _core_container_to_dict(self, container):
-        """Convert core ContainerDefinition to dict using container_def helper."""
-        from sagemaker.core.utils.utils import Unassigned
+        """Convert a core ContainerDefinition into a CreateModel container dict.
 
-        # Helper to check if value is Unassigned
-        def get_value(obj, attr, default=None):
-            if not hasattr(obj, attr):
-                return default
-            val = getattr(obj, attr)
-            return default if isinstance(val, Unassigned) else val
+        The core serializer already emits the exact CreateModel request shape for
+        every field a ``ContainerDefinition`` supports, so the container is routed
+        through it instead of copying a fixed subset of fields. Copying a subset
+        dropped fields such as ``ModelPackageName`` and ``ModelDataSource``.
 
-        return container_def(
-            container.image,
-            get_value(container, "model_data_url"),
-            get_value(container, "environment", {}),
-            image_config=get_value(container, "image_config"),
-        )
+        Args:
+            container (sagemaker.core.shapes.ContainerDefinition): The container to convert.
+
+        Returns:
+            dict: A container definition usable with the CreateModel API.
+        """
+        from sagemaker.core.utils.utils import serialize
+
+        return serialize(container) or {}
 
     def _create_sagemaker_model(self):
         """Create a SageMaker Model Entity."""
         container_def = self._prepare_container_def()
 
+        # A container does not always carry an image: a model-package reference
+        # identifies its image through ModelPackageName instead, so Image is
+        # looked up rather than indexed.
+        resolved_image_uri = (
+            container_def.get("Image")
+            if isinstance(container_def, dict)
+            else container_def[0].get("Image")
+        )
+
         if not isinstance(self.sagemaker_session, PipelineSession):
             # _base_name, model_name are not needed under PipelineSession.
             # the model_data may be Pipeline variable
             # which may break the _base_name generation
-            image_uri = (
-                container_def["Image"]
-                if isinstance(container_def, dict)
-                else container_def[0]["Image"]
-            )
             self._ensure_base_name_if_needed(
-                image_uri=image_uri,
+                image_uri=resolved_image_uri,
                 script_uri=self.source_dir,
                 model_uri=self._get_model_uri(),
             )
@@ -3146,9 +3150,6 @@ def _create_sagemaker_model(self):
         # A pipeline-variable image is only resolved at execution time, so it
         # cannot be inspected here; set enable_network_isolation explicitly then.
         enable_network_isolation = self._enable_network_isolation
-        resolved_image_uri = (
-            container_def["Image"] if isinstance(container_def, dict) else container_def[0]["Image"]
-        )
         if (
             not enable_network_isolation
             and isinstance(resolved_image_uri, str)

sagemaker-serve/tests/unit/test_model_builder_deploy.py

@@ -14,6 +14,13 @@
 from sagemaker.core.inference_config import (
     ResourceRequirements,
 )
+from sagemaker.core.shapes import (
+    ContainerDefinition,
+    ImageConfig,
+    ModelDataSource,
+    S3ModelDataSource,
+)
+from sagemaker.core.workflow.pipeline_context import PipelineSession
 
 
 class TestModelBuilderContainerDef(unittest.TestCase):
@@ -104,27 +111,242 @@ def test_prepare_container_def_base_invalid_pipeline_models(self):
 
         self.assertIn("must be sagemaker.core.resources.Model instances", str(context.exception))
 
-    @patch("sagemaker.serve.model_builder.container_def")
-    def test_core_container_to_dict(self, mock_def):
+    def test_core_container_to_dict(self):
         """Test _core_container_to_dict converts container properly."""
-        from sagemaker.core.utils.utils import Unassigned
-
-        mock_container = Mock()
-        mock_container.image = "test-image"
-        mock_container.model_data_url = "s3://bucket/model.tar.gz"
-        mock_container.environment = {"KEY": "value"}
-        mock_container.image_config = Unassigned()
+        container = ContainerDefinition(
+            image="test-image",
+            model_data_url="s3://bucket/model.tar.gz",
+            environment={"KEY": "value"},
+        )
 
         builder = ModelBuilder(
             model=Mock(),
             role_arn="arn:aws:iam::123456789012:role/TestRole",
             sagemaker_session=self.mock_session,
         )
 
-        mock_def.return_value = {"Image": "test-image"}
-        builder._core_container_to_dict(mock_container)
+        result = builder._core_container_to_dict(container)
+
+        self.assertEqual(
+            result,
+            {
+                "Image": "test-image",
+                "ModelDataUrl": "s3://bucket/model.tar.gz",
+                "Environment": {"KEY": "value"},
+            },
+        )
+
+
+@patch("sagemaker.serve.model_builder.resolve_nested_dict_value_from_config")
+@patch("sagemaker.serve.model_builder.resolve_value_from_config")
+@patch.object(ModelBuilder, "_init_sagemaker_session_if_does_not_exist")
+class TestCoreContainerDefinitionFields(unittest.TestCase):
+    """Core container fields must survive the CreateModel conversion (issue #6476).
+
+    ``_core_container_to_dict`` used to copy a fixed subset of fields, which
+    dropped ``ModelPackageName`` and ``ModelDataSource`` from the generated
+    ``CreateModel`` request.
+    """
+
+    ROLE_ARN = "arn:aws:iam::123456789012:role/TestRole"
+    IMAGE = "123456789012.dkr.ecr.us-west-2.amazonaws.com/sagemaker-scikit-learn:1.4-2"
+    PACKAGE_ARN = "arn:aws:sagemaker:us-west-2:123456789012:model-package/example/1"
+
+    def _container_defs(self, container, mock_resolve, mock_resolve_nested):
+        """Run _create_sagemaker_model and return the container_defs sent to create_model."""
+        mock_resolve.side_effect = lambda value, *args, **kwargs: value
+        mock_resolve_nested.side_effect = lambda value, *args, **kwargs: value
+
+        session = Mock()
+        session.boto_region_name = "us-west-2"
+        session.default_bucket.return_value = "test-bucket"
+        session.default_bucket_prefix = None
+        session.config = {}
+        session.sagemaker_config = {}
+
+        builder = ModelBuilder(
+            model=[
+                Model(
+                    model_name="source-model",
+                    primary_container=container,
+                    execution_role_arn=self.ROLE_ARN,
+                )
+            ],
+            role_arn=self.ROLE_ARN,
+            sagemaker_session=session,
+            instance_type="ml.m5.large",
+        )
+        builder.model_name = "test-model"
+        builder.sagemaker_session = Mock(spec=PipelineSession)
+
+        builder._create_sagemaker_model()
+
+        builder.sagemaker_session.create_model.assert_called_once()
+        return builder.sagemaker_session.create_model.call_args.kwargs["container_defs"]
+
+    def test_model_package_container_keeps_model_package_name(
+        self, mock_init, mock_resolve, mock_resolve_nested
+    ):
+        """A model-package reference must build a request instead of failing on Image."""
+        container_defs = self._container_defs(
+            ContainerDefinition(model_package_name=self.PACKAGE_ARN),
+            mock_resolve,
+            mock_resolve_nested,
+        )
+
+        self.assertEqual(container_defs, [{"ModelPackageName": self.PACKAGE_ARN}])
+
+    def test_model_data_source_container_keeps_model_data_source(
+        self, mock_init, mock_resolve, mock_resolve_nested
+    ):
+        """A ModelDataSource must not be silently dropped from the request."""
+        container_defs = self._container_defs(
+            ContainerDefinition(
+                image=self.IMAGE,
+                model_data_source=ModelDataSource(
+                    s3_data_source=S3ModelDataSource(
+                        s3_data_type="S3Prefix",
+                        compression_type="None",
+                        s3_uri="s3://example-bucket/model/",
+                    )
+                ),
+                environment={"EXAMPLE": "true"},
+            ),
+            mock_resolve,
+            mock_resolve_nested,
+        )
+
+        self.assertEqual(
+            container_defs,
+            [
+                {
+                    "Image": self.IMAGE,
+                    "ModelDataSource": {
+                        "S3DataSource": {
+                            "S3DataType": "S3Prefix",
+                            "CompressionType": "None",
+                            "S3Uri": "s3://example-bucket/model/",
+                        }
+                    },
+                    "Environment": {"EXAMPLE": "true"},
+                }
+            ],
+        )
+
+    def test_plain_image_container_is_unchanged(self, mock_init, mock_resolve, mock_resolve_nested):
+        """A plain image container must keep producing the same request fields."""
+        container_defs = self._container_defs(
+            ContainerDefinition(
+                image=self.IMAGE,
+                model_data_url="s3://example-bucket/model.tar.gz",
+                environment={"A": "1"},
+            ),
+            mock_resolve,
+            mock_resolve_nested,
+        )
+
+        self.assertEqual(
+            container_defs,
+            [
+                {
+                    "Image": self.IMAGE,
+                    "ModelDataUrl": "s3://example-bucket/model.tar.gz",
+                    "Environment": {"A": "1"},
+                }
+            ],
+        )
+
+    def test_image_config_is_serialized_to_a_plain_dict(
+        self, mock_init, mock_resolve, mock_resolve_nested
+    ):
+        """ImageConfig must reach the request as a dict, not as a shape object."""
+        container_defs = self._container_defs(
+            ContainerDefinition(
+                image=self.IMAGE,
+                image_config=ImageConfig(repository_access_mode="Vpc"),
+            ),
+            mock_resolve,
+            mock_resolve_nested,
+        )
+
+        self.assertEqual(
+            container_defs,
+            [{"Image": self.IMAGE, "ImageConfig": {"RepositoryAccessMode": "Vpc"}}],
+        )
+
+    @patch("sagemaker.serve.model_builder.Model.get")
+    def test_model_package_container_under_interactive_session(
+        self, mock_model_get, mock_init, mock_resolve, mock_resolve_nested
+    ):
+        """Outside a PipelineSession the base-name lookup must tolerate a missing Image."""
+        mock_resolve.side_effect = lambda value, *args, **kwargs: value
+        mock_resolve_nested.side_effect = lambda value, *args, **kwargs: value
+
+        session = Mock()
+        session.boto_region_name = "us-west-2"
+        session.default_bucket.return_value = "test-bucket"
+        session.default_bucket_prefix = None
+        session.config = {}
+        session.sagemaker_config = {}
+
+        builder = ModelBuilder(
+            model=[
+                Model(
+                    model_name="source-model",
+                    primary_container=ContainerDefinition(model_package_name=self.PACKAGE_ARN),
+                    execution_role_arn=self.ROLE_ARN,
+                )
+            ],
+            role_arn=self.ROLE_ARN,
+            sagemaker_session=session,
+            instance_type="ml.m5.large",
+        )
+        builder.model_name = "test-model"
+
+        builder._create_sagemaker_model()
+
+        kwargs = session.create_model.call_args.kwargs
+        self.assertEqual(kwargs["container_defs"], [{"ModelPackageName": self.PACKAGE_ARN}])
+
+    def test_model_package_container_passes_real_boto_param_validation(
+        self, mock_init, mock_resolve, mock_resolve_nested
+    ):
+        """End-to-end proof against the real SageMaker service model (offline).
+
+        boto3's client-side ParamValidator needs no network and no credentials. The
+        request the old converter produced carried ``Image: Unassigned()`` and was
+        rejected by that validator; the request produced now is accepted.
+        """
+        import boto3
+        from botocore import validate
+
+        from sagemaker.core.utils.utils import Unassigned
+
+        shape = (
+            boto3.client("sagemaker", region_name="us-west-2")
+            .meta.service_model.operation_model("CreateModel")
+            .input_shape
+        )
+        validator = validate.ParamValidator()
+        base = {
+            "ModelName": "test-model",
+            "ExecutionRoleArn": self.ROLE_ARN,
+        }
+
+        # What the old field copy produced for a model-package reference -- boto rejects it.
+        old_report = validator.validate(
+            {**base, "Containers": [{"Image": Unassigned(), "Environment": {}}]}, shape
+        )
+        self.assertTrue(old_report.has_errors())
+
+        container_defs = self._container_defs(
+            ContainerDefinition(model_package_name=self.PACKAGE_ARN),
+            mock_resolve,
+            mock_resolve_nested,
+        )
 
-        mock_def.assert_called_once()
+        new_report = validator.validate({**base, "Containers": container_defs}, shape)
+        self.assertFalse(new_report.has_errors(), new_report.generate_report())
 
 
 class TestModelBuilderDeploy(unittest.TestCase):