fix(serve): expose orchestrator deployment fields

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

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(