From 3dfd6e89f299a6209bc9d33e98e55fe6849346cb Mon Sep 17 00:00:00 2001 From: Lucas Jia Date: Tue, 13 Oct 2026 13:02:55 -0700 Subject: [PATCH] fix(serve): expose orchestrator deployment fields Restore ModelBuilder constructor support for resource requirements and inference component names, preserving them in custom orchestrator deployables. Fixes #6198 --- .../src/sagemaker/serve/model_builder.py | 27 +++++++-- .../tests/unit/test_model_builder_build.py | 60 +++++++++++++++++++ 2 files changed, 83 insertions(+), 4 deletions(-) diff --git a/sagemaker-serve/src/sagemaker/serve/model_builder.py b/sagemaker-serve/src/sagemaker/serve/model_builder.py index 10a29fa5f8..e3bf4a0fff 100644 --- a/sagemaker-serve/src/sagemaker/serve/model_builder.py +++ b/sagemaker-serve/src/sagemaker/serve/model_builder.py @@ -261,6 +261,11 @@ class ModelBuilder(_InferenceRecommenderMixin, _ModelBuilderServers, _ModelBuild network: Network configuration including VPC settings and network isolation. instance_type: EC2 instance type for deployment (e.g., 'ml.m5.large'). mode: Deployment mode (SAGEMAKER_ENDPOINT, LOCAL_CONTAINER, or IN_PROCESS). + resource_requirements: Compute resources for an inference component. When provided with + a custom orchestrator, the orchestrator is built as an inference component. + (Default: None). + inference_component_name: Name for the inference component created from this + ModelBuilder, including a custom orchestrator inference component. (Default: None). Note: ModelBuilder returns sagemaker.core.resources.Model and sagemaker.core.resources.Endpoint @@ -451,6 +456,22 @@ class ModelBuilder(_InferenceRecommenderMixin, _ModelBuilderServers, _ModelBuild metadata={"help": "DEPRECATED: Use configure_for_torchserve() instead"}, ) + # Keep these fields trailing to preserve the positional order of existing constructor arguments. + resource_requirements: Optional[ResourceRequirements] = field( + default=None, + metadata={ + "help": "Compute resources for an inference component. When provided with a custom " + "orchestrator, the orchestrator is built as an inference component." + }, + ) + inference_component_name: Optional[str] = field( + default=None, + metadata={ + "help": "Name for the inference component created from this ModelBuilder, including " + "a custom orchestrator inference component." + }, + ) + def _create_session_with_region(self): """Create a SageMaker session with the correct region.""" if hasattr(self, "region") and self.region: @@ -1837,8 +1858,6 @@ def _initialize_jumpstart_config(self) -> None: self.inference_ami_version = None if not hasattr(self, "model_version"): self.model_version = None - if not hasattr(self, "resource_requirements"): - self.resource_requirements = None if not hasattr(self, "model_kms_key"): self.model_kms_key = None if not hasattr(self, "hub_name"): @@ -3928,12 +3947,12 @@ def _deploy_core_endpoint(self, **kwargs): if "model_data_download_timeout" in kwargs: self.model_data_download_timeout = model_data_download_timeout - resources = kwargs.get("resources", getattr(self, "resource_requirements", None)) + resources = kwargs.get("resources", self.resource_requirements) if "resources" in kwargs: self.resource_requirements = resources inference_component_name = kwargs.get( - "inference_component_name", getattr(self, "inference_component_name", None) + "inference_component_name", self.inference_component_name ) if "inference_component_name" in kwargs: self.inference_component_name = inference_component_name diff --git a/sagemaker-serve/tests/unit/test_model_builder_build.py b/sagemaker-serve/tests/unit/test_model_builder_build.py index d0d5eebc28..7ab89c8197 100644 --- a/sagemaker-serve/tests/unit/test_model_builder_build.py +++ b/sagemaker-serve/tests/unit/test_model_builder_build.py @@ -9,12 +9,22 @@ import os from sagemaker.core.exceptions import UnexpectedStatusException +from sagemaker.core.inference_config import ResourceRequirements from sagemaker.serve.model_builder import ModelBuilder from sagemaker.serve.utils.types import ModelServer from sagemaker.serve.mode.function_pointers import Mode +from sagemaker.serve.spec.inference_base import CustomOrchestrator from sagemaker.core.resources import Model +class ConcreteOrchestrator(CustomOrchestrator): + """Custom orchestrator implementation for ModelBuilder tests.""" + + def handle(self, data, context=None): + """Return the supplied data unchanged.""" + return data + + class TestModelBuilderSaveModel(unittest.TestCase): """Test ModelBuilder _save_model_inference_spec method.""" @@ -182,6 +192,56 @@ def setUp(self): self.mock_session.boto_session = Mock() self.mock_session.boto_session.region_name = "us-west-2" + def test_constructor_accepts_orchestrator_deployment_fields(self): + """Verify ModelBuilder accepts inference component settings at construction time.""" + resource_requirements = ResourceRequirements( + requests={"num_cpus": 2, "memory": 4096}, limits={"memory": 8192} + ) + + builder = ModelBuilder( + inference_spec=ConcreteOrchestrator(), + resource_requirements=resource_requirements, + inference_component_name="custom-orchestrator", + role_arn="arn:aws:iam::123456789012:role/TestRole", + sagemaker_session=self.mock_session, + ) + + self.assertIs(builder.resource_requirements, resource_requirements) + self.assertEqual(builder.inference_component_name, "custom-orchestrator") + + def test_custom_orchestrator_build_retains_deployment_fields(self): + """Custom orchestrator deployables retain constructor deployment settings.""" + resource_requirements = ResourceRequirements( + requests={"num_cpus": 2, "memory": 4096}, limits={"memory": 8192} + ) + builder = ModelBuilder( + inference_spec=ConcreteOrchestrator(), + resource_requirements=resource_requirements, + inference_component_name="custom-orchestrator", + role_arn="arn:aws:iam::123456789012:role/TestRole", + sagemaker_session=self.mock_session, + ) + built_model = Mock(spec=Model) + + with ( + patch.object(builder, "_get_serve_setting", return_value=Mock()), + patch.object(builder, "_get_processing_unit", return_value="cpu"), + patch.object(builder, "_get_smd_image_uri", return_value="test-smd-image"), + patch.object(builder, "_build_single_modelbuilder", return_value=built_model), + ): + result = builder.build() + + self.assertIs(result, builder) + self.assertEqual( + builder._deployables["CustomOrchestrator"], + { + "Name": "custom-orchestrator", + "Mode": "InferenceComponent", + "ResourceRequirements": resource_requirements, + "Model": built_model, + }, + ) + def test_build_warns_on_multiple_calls(self): """Test build() warns when called multiple times.""" builder = ModelBuilder(