diff --git a/sagemaker-serve/src/sagemaker/serve/model_builder.py b/sagemaker-serve/src/sagemaker/serve/model_builder.py index 24e747f206..0f5400b13c 100644 --- a/sagemaker-serve/src/sagemaker/serve/model_builder.py +++ b/sagemaker-serve/src/sagemaker/serve/model_builder.py @@ -599,6 +599,10 @@ def _initialize_defaults(self) -> None: self.framework = None if not hasattr(self, "framework_version"): self.framework_version = None + if not hasattr(self, "content_types"): + self.content_types = None + if not hasattr(self, "response_types"): + self.response_types = None def _fetch_default_instance_type_for_custom_model(self) -> str: hosting_configs = self._fetch_hosting_configs_for_custom_model() 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..14195e445d --- /dev/null +++ b/sagemaker-serve/tests/unit/test_model_builder_register.py @@ -0,0 +1,78 @@ +# 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. +"""Unit tests for ModelBuilder.register().""" + +from unittest.mock import MagicMock, patch + +import pytest + +from sagemaker.serve.model_builder import ModelBuilder + +ROLE_ARN = "arn:aws:iam::123456789012:role/SageMakerRole" +IMAGE_URI = "123456789012.dkr.ecr.us-west-2.amazonaws.com/inference:latest" +MODEL_PACKAGE_ARN = "arn:aws:sagemaker:us-west-2:123456789012:model-package/group/1" + + +@pytest.fixture +def model_builder(): + session = MagicMock(boto_region_name="us-west-2") + # register() is telemetry-wrapped and reads sagemaker_config. + session.sagemaker_config = {} + return ModelBuilder( + image_uri=IMAGE_URI, + s3_model_data_url="s3://bucket/model.tar.gz", + role_arn=ROLE_ARN, + sagemaker_session=session, + instance_type="ml.m5.large", + ) + + +def _register(model_builder, **kwargs): + """Call register() and return the arguments it passes to create the model package.""" + with patch( + "sagemaker.serve.model_builder.create_model_package_from_containers", + return_value={"ModelPackageArn": MODEL_PACKAGE_ARN}, + ) as create_model_package, patch.object( + model_builder, "_prepare_container_def", return_value={"Image": IMAGE_URI} + ): + arn = model_builder.register(model_package_group_name="group", **kwargs) + assert arn == MODEL_PACKAGE_ARN + return create_model_package.call_args.kwargs + + +@pytest.mark.parametrize( + "types", + [ + {}, + {"content_types": ["text/csv"]}, + {"response_types": ["application/json"]}, + ], + ids=["neither", "content_types_only", "response_types_only"], +) +def test_register_without_content_or_response_types(model_builder, types): + request_args = _register(model_builder, **types) + + for name in ("content_types", "response_types"): + if name in types: + assert request_args[name] == types[name] + else: + assert name not in request_args + + +def test_register_reuses_content_and_response_types_from_earlier_call(model_builder): + _register(model_builder, content_types=["text/csv"], response_types=["application/json"]) + + request_args = _register(model_builder) + + assert request_args["content_types"] == ["text/csv"] + assert request_args["response_types"] == ["application/json"]