From bf307e980d66a80a68bb8cb98b8f1740637ff413 Mon Sep 17 00:00:00 2001 From: Om Singhal Date: Sat, 22 Aug 2026 00:03:24 -0400 Subject: [PATCH] fix(serve): pass instance_type to _deploy once when creating an inference component Deploying a CustomOrchestrator in InferenceComponent mode through ModelBuilder.deploy() always raised TypeError: _deploy() got multiple values for keyword argument 'instance_type' In _deploy_for_ic the branch that creates a new inference component read instance_type and initial_instance_count out of kwargs with kwargs.get, passed them to self._deploy as explicit keyword arguments, and then spread the unchanged **kwargs into the same call. kwargs.get leaves the key in place, so Python saw each of those arguments twice. The only caller of this branch, ModelBuilder.deploy, always supplies both keys, so the failure was deterministic. Consume the two keys with kwargs.pop instead so each is forwarded exactly once. The defaults (self.instance_type and 1) are unchanged and kwargs is a local dict, so the caller is unaffected. _deploy_core_endpoint reads the values back through kwargs.get and receives the same data as before. Add unit tests covering the create branch of _deploy_for_ic with explicit instance arguments (fails with the TypeError before this change) and with the default fallbacks. Fixes #6199 --- .../src/sagemaker/serve/model_builder.py | 4 +- .../tests/unit/test_model_builder_deploy.py | 68 +++++++++++++++++++ 2 files changed, 70 insertions(+), 2 deletions(-) diff --git a/sagemaker-serve/src/sagemaker/serve/model_builder.py b/sagemaker-serve/src/sagemaker/serve/model_builder.py index 0d1e6e4746..c3d96c6399 100644 --- a/sagemaker-serve/src/sagemaker/serve/model_builder.py +++ b/sagemaker-serve/src/sagemaker/serve/model_builder.py @@ -4230,8 +4230,8 @@ def _deploy_for_ic(self, ic_data: Dict[str, Any], endpoint_name: str, **kwargs) endpoint_type=EndpointType.INFERENCE_COMPONENT_BASED, resources=resource_requirements, inference_component_name=ic_name, - instance_type=kwargs.get("instance_type", self.instance_type), - initial_instance_count=kwargs.get("initial_instance_count", 1), + instance_type=kwargs.pop("instance_type", self.instance_type), + initial_instance_count=kwargs.pop("initial_instance_count", 1), **kwargs, ) diff --git a/sagemaker-serve/tests/unit/test_model_builder_deploy.py b/sagemaker-serve/tests/unit/test_model_builder_deploy.py index 3a0fca3d8e..1a6d0817ba 100644 --- a/sagemaker-serve/tests/unit/test_model_builder_deploy.py +++ b/sagemaker-serve/tests/unit/test_model_builder_deploy.py @@ -521,6 +521,74 @@ def test_does_ic_exist_false(self): self.assertFalse(result) + def test_deploy_for_ic_creates_new_ic_with_explicit_instance_kwargs(self): + """Test _deploy_for_ic forwards instance_type and initial_instance_count to _deploy once.""" + builder = ModelBuilder( + model=Mock(), + role_arn="arn:aws:iam::123456789012:role/TestRole", + sagemaker_session=self.mock_session, + ) + builder._does_ic_exist = Mock(return_value=False) + builder._deploy = Mock(return_value="endpoint") + built_model = Mock() + resource_requirements = ResourceRequirements(requests={"memory": 1024, "copies": 1}) + ic_data = { + "Name": "test-ic", + "ResourceRequirements": resource_requirements, + "Model": built_model, + } + + result = builder._deploy_for_ic( + ic_data=ic_data, + endpoint_name="test-endpoint", + container_timeout_in_seconds=600, + instance_type="ml.g5.xlarge", + initial_instance_count=2, + ) + + self.assertEqual(result, "endpoint") + builder._deploy.assert_called_once_with( + built_model=built_model, + endpoint_name="test-endpoint", + endpoint_type=EndpointType.INFERENCE_COMPONENT_BASED, + resources=resource_requirements, + inference_component_name="test-ic", + instance_type="ml.g5.xlarge", + initial_instance_count=2, + container_timeout_in_seconds=600, + ) + + def test_deploy_for_ic_creates_new_ic_with_default_instance_kwargs(self): + """Test _deploy_for_ic falls back to builder instance_type and one instance.""" + builder = ModelBuilder( + model=Mock(), + role_arn="arn:aws:iam::123456789012:role/TestRole", + sagemaker_session=self.mock_session, + ) + builder.instance_type = "ml.c5.xlarge" + builder._does_ic_exist = Mock(return_value=False) + builder._deploy = Mock(return_value="endpoint") + built_model = Mock() + resource_requirements = ResourceRequirements(requests={"memory": 1024, "copies": 1}) + ic_data = { + "Name": "test-ic", + "ResourceRequirements": resource_requirements, + "Model": built_model, + } + + result = builder._deploy_for_ic(ic_data=ic_data, endpoint_name="test-endpoint") + + self.assertEqual(result, "endpoint") + builder._deploy.assert_called_once_with( + built_model=built_model, + endpoint_name="test-endpoint", + endpoint_type=EndpointType.INFERENCE_COMPONENT_BASED, + resources=resource_requirements, + inference_component_name="test-ic", + instance_type="ml.c5.xlarge", + initial_instance_count=1, + ) + class TestModelBuilderResetState(unittest.TestCase): """Test ModelBuilder _reset_build_state method."""