fix: expose inference component fields on ModelBuilder

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={
@@ -1848,8 +1861,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"):

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):