From eaa31d97f2164e749cd6c3c78203e4f19beba1da Mon Sep 17 00:00:00 2001
From: gauravkriplani <127468958+gauravkriplani@users.noreply.github.com>
Date: Sun, 11 Oct 2026 22:09:20 -0700
Subject: [PATCH] fix: expose inference component fields on ModelBuilder

---
 .../src/sagemaker/serve/model_builder.py          | 15 +++++++++++++--
 .../tests/unit/test_model_builder_workflows.py    | 14 ++++++++++----
 2 files changed, 23 insertions(+), 6 deletions(-)

diff --git a/sagemaker-serve/src/sagemaker/serve/model_builder.py b/sagemaker-serve/src/sagemaker/serve/model_builder.py
index 10a29fa5f8..f640052153 100644
--- a/sagemaker-serve/src/sagemaker/serve/model_builder.py
+++ b/sagemaker-serve/src/sagemaker/serve/model_builder.py
@@ -245,6 +245,8 @@ class ModelBuilder(_InferenceRecommenderMixin, _ModelBuilderServers, _ModelBuild
         inference_spec: Custom inference specification with load() and invoke() functions.
         schema_builder: Defines input/output schema for serialization and deserialization.
         modelbuilder_list: List of ModelBuilder objects for multi-model deployments.
+        resource_requirements: Compute resources allocated to an inference component.
+        inference_component_name: Name to use when building an inference component.
         pipeline_models: List of Model objects for creating inference pipelines.
         role_arn: IAM role ARN for SageMaker to assume.
         sagemaker_session: Session object for managing SageMaker API interactions.
@@ -315,6 +317,17 @@ class ModelBuilder(_InferenceRecommenderMixin, _ModelBuilderServers, _ModelBuild
             "Used when deploying multiple models to a single endpoint."
         },
     )
+    resource_requirements: Optional[ResourceRequirements] = field(
+        default=None,
+        metadata={
+            "help": "Compute resources to allocate to an inference component, including CPU, "
+            "memory, accelerator devices, and copies."
+        },
+    )
+    inference_component_name: Optional[str] = field(
+        default=None,
+        metadata={"help": "Name to use when building an inference component."},
+    )
     role_arn: Optional[str] = field(
         default=None,
         metadata={
@@ -1837,8 +1850,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"):
diff --git a/sagemaker-serve/tests/unit/test_model_builder_workflows.py b/sagemaker-serve/tests/unit/test_model_builder_workflows.py
index 395d251fee..c70f251fb8 100644
--- a/sagemaker-serve/tests/unit/test_model_builder_workflows.py
+++ b/sagemaker-serve/tests/unit/test_model_builder_workflows.py
@@ -128,16 +128,17 @@ def test_build_with_modelbuilder_list_builds_inference_components(
         mock_get_serve.return_value = Mock()
 
         # Create ModelBuilder with inference component
+        resource_requirements = ResourceRequirements(
+            requests={"memory": 1024, "copies": 1}, limits={}
+        )
         mb1 = ModelBuilder(
             inference_spec=Mock(spec=InferenceSpec),
             role_arn=MOCK_ROLE_ARN,
             sagemaker_session=self.mock_session,
             image_uri=MOCK_IMAGE_URI,
             mode=Mode.SAGEMAKER_ENDPOINT,
-        )
-        mb1.inference_component_name = "ic-1"
-        mb1.resource_requirements = ResourceRequirements(
-            requests={"memory": 1024, "copies": 1}, limits={}
+            resource_requirements=resource_requirements,
+            inference_component_name="ic-1",
         )
         mb1.model_name = "test-model-1"
         mb1.model_server = ModelServer.TORCHSERVE
@@ -160,6 +161,11 @@ def test_build_with_modelbuilder_list_builds_inference_components(
         self.assertIsNotNone(result)
         self.assertIn("InferenceComponents", builder._deployables)
         self.assertEqual(len(builder._deployables["InferenceComponents"]), 1)
+        self.assertEqual(builder._deployables["InferenceComponents"][0]["Name"], "ic-1")
+        self.assertIs(
+            builder._deployables["InferenceComponents"][0]["ResourceRequirements"],
+            resource_requirements,
+        )
 
 
 class TestModelBuilderDeployMethod(unittest.TestCase):