fixes for model builder (#5631)

* fixes for model builder

* add nova model support

* fix env_vars merge, update integ test for LORA two-step deployment, fix unit tests for nova model support

- env_vars: append recipe/nova config to existing env_vars instead of skipping
- integ test: verify both base IC and adapter IC creation for LORA models
- unit tests: add _is_nova_model mock to accommodate nova model support changes

* update codegen to mark MinMemoryRequiredInMb as optional

DescribeInferenceComponent returns empty ComputeResourceRequirements
for adapter ICs (created with BaseInferenceComponentName), but the
service model still marks MinMemoryRequiredInMb as required. Add a
REQUIRED_TO_OPTIONAL_OVERRIDES config in the codegen so re-running
shapes generation produces the correct Optional field.

* add retry for adapter IC creation on transient endpoint-not-found

* model builder fixes

* Skip test_deploy_from_training_job: parallel cleanup race condition under investigation

---------

Co-authored-by: Joshua Towner <jjtowner@amazon.com>
Co-authored-by: Roja Reddy Sareddy <rsareddy@amazon.com>

Author: 3 people

sagemaker-core/src/sagemaker/core/shapes/shapes.py

@@ -8577,7 +8577,7 @@ class InferenceComponentComputeResourceRequirements(Base):
     max_memory_required_in_mb: The maximum MB of memory to allocate to run a model that you assign to an inference component.
     """
 
-    min_memory_required_in_mb: int
+    min_memory_required_in_mb: Optional[int] = Unassigned()
     number_of_cpu_cores_required: Optional[float] = Unassigned()
     number_of_accelerator_devices_required: Optional[float] = Unassigned()
     max_memory_required_in_mb: Optional[int] = Unassigned()

sagemaker-core/src/sagemaker/core/tools/constants.py

@@ -107,3 +107,9 @@
 CONFIG_SCHEMA_FILE_NAME = "config_schema.py"
 
 API_COVERAGE_JSON_FILE_PATH = os.getcwd() + "/src/sagemaker/core/tools/api_coverage.json"
+
+# Members that the service model marks as required but the API returns as optional.
+# E.g. DescribeInferenceComponent returns empty ComputeResourceRequirements for adapter ICs.
+REQUIRED_TO_OPTIONAL_OVERRIDES = {
+    "InferenceComponentComputeResourceRequirements": ["MinMemoryRequiredInMb"],
+}

sagemaker-core/src/sagemaker/core/tools/shapes_extractor.py

@@ -16,7 +16,11 @@
 from functools import lru_cache
 from typing import Optional, Any
 
-from sagemaker.core.tools.constants import BASIC_JSON_TYPES_TO_PYTHON_TYPES, SHAPE_DAG_FILE_PATH
+from sagemaker.core.tools.constants import (
+    BASIC_JSON_TYPES_TO_PYTHON_TYPES,
+    REQUIRED_TO_OPTIONAL_OVERRIDES,
+    SHAPE_DAG_FILE_PATH,
+)
 from sagemaker.core.utils.utils import (
     reformat_file_with_black,
     convert_to_snake_case,
@@ -216,6 +220,11 @@ def generate_shape_members(self, shape, required_override=()):
         shape_dict = self.combined_shapes[shape]
         members = shape_dict["members"]
         required_args = list(required_override) or shape_dict.get("required", [])
+        # Remove members that are known to be optional despite the service model
+        required_args = [
+            r for r in required_args
+            if r not in REQUIRED_TO_OPTIONAL_OVERRIDES.get(shape, [])
+        ]
         init_data_body = {}
         # bring the required members in front
         ordered_members = {key: members[key] for key in required_args if key in members}

sagemaker-serve/src/sagemaker/serve/model_builder.py

@@ -112,34 +112,45 @@ def test_build_from_training_job(self, training_job_name):
         assert model_builder.image_uri is not None
         assert model_builder.instance_type is not None
 
+    @pytest.mark.skip(reason="Skipped: parallel cleanup race condition under investigation")
     def test_deploy_from_training_job(self, training_job_name, endpoint_name, cleanup_endpoints):
-        """Test deploying model from training job and adapter."""
-        from sagemaker.core.resources import TrainingJob
+        """Test deploying model from training job.
+
+        For LORA models, this verifies the two-step deployment:
+        base IC + adapter IC are both created on the same endpoint.
+        """
+        from sagemaker.core.resources import TrainingJob, InferenceComponent
         from sagemaker.serve import ModelBuilder
         import time
 
         training_job = TrainingJob.get(training_job_name=training_job_name)
-        model_builder = ModelBuilder(model=training_job)
-        model = model_builder.build(model_name=f"test-model-{int(time.time())}-{random.randint(100, 10000)}")
-        endpoint = model_builder.deploy(endpoint_name=endpoint_name)
+        model_builder = ModelBuilder(model=training_job, instance_type="ml.g5.4xlarge")
+        model_builder.build(model_name=f"test-model-{int(time.time())}-{random.randint(100, 10000)}")
+
+        peft_type = model_builder._fetch_peft()
+        adapter_name = f"{endpoint_name}-adapter"
+
+        endpoint = model_builder.deploy(
+            endpoint_name=endpoint_name,
+            inference_component_name=adapter_name if peft_type == "LORA" else None,
+        )
 
         cleanup_endpoints.append(endpoint_name)
 
         assert endpoint is not None
         assert endpoint.endpoint_arn is not None
         assert endpoint.endpoint_status == "InService"
 
-        # Deploy adapter to the same endpoint
-        adapter_name = f"{endpoint_name}-adapter-{int(time.time())}-{random.randint(100, 100000)}"
-        model_builder2 = ModelBuilder(model=training_job)
-        model_builder2.build()
-        endpoint2 = model_builder2.deploy(
-            endpoint_name=endpoint_name,
-            inference_component_name=adapter_name
-        )
+        if peft_type == "LORA":
+            # Verify base IC was created
+            base_ic_name = f"{endpoint_name}-inference-component"
+            base_ic = InferenceComponent.get(inference_component_name=base_ic_name)
+            assert base_ic is not None
+            assert base_ic.inference_component_status == "InService"
 
-        assert endpoint2 is not None
-        assert endpoint2.endpoint_name == endpoint_name
+            # Verify adapter IC was created
+            adapter_ic = InferenceComponent.get(inference_component_name=adapter_name)
+            assert adapter_ic is not None
 
     def test_fetch_endpoint_names_for_base_model(self, training_job_name):
         """Test fetching endpoint names for base model."""

sagemaker-serve/tests/integ/test_model_customization_deployment.py

@@ -45,8 +45,10 @@ def setUp(self):
     @patch("sagemaker.serve.model_builder.ModelBuilder._resolve_model_artifact_uri")
     @patch("sagemaker.serve.model_builder.ModelBuilder._fetch_peft")
     @patch("sagemaker.serve.model_builder.ModelBuilder._is_model_customization")
+    @patch("sagemaker.serve.model_builder.ModelBuilder._is_nova_model", return_value=False)
     def test_base_model_artifact_uri_propagated_to_inference_component(
         self,
+        mock_is_nova_model,
         mock_is_model_customization,
         mock_fetch_peft,
         mock_resolve_artifact,
@@ -133,8 +135,10 @@ def test_base_model_artifact_uri_propagated_to_inference_component(
     @patch("sagemaker.serve.model_builder.ModelBuilder._resolve_model_artifact_uri")
     @patch("sagemaker.serve.model_builder.ModelBuilder._fetch_peft")
     @patch("sagemaker.serve.model_builder.ModelBuilder._is_model_customization")
+    @patch("sagemaker.serve.model_builder.ModelBuilder._is_nova_model", return_value=False)
     def test_fine_tuned_model_artifact_uri_propagated_to_inference_component(
         self,
+        mock_is_nova_model,
         mock_is_model_customization,
         mock_fetch_peft,
         mock_resolve_artifact,
@@ -220,8 +224,10 @@ def test_fine_tuned_model_artifact_uri_propagated_to_inference_component(
     @patch("sagemaker.serve.model_builder.ModelBuilder._resolve_model_artifact_uri")
     @patch("sagemaker.serve.model_builder.ModelBuilder._fetch_peft")
     @patch("sagemaker.serve.model_builder.ModelBuilder._is_model_customization")
+    @patch("sagemaker.serve.model_builder.ModelBuilder._is_nova_model", return_value=False)
     def test_lora_adapter_no_artifact_uri_propagated(
         self,
+        mock_is_nova_model,
         mock_is_model_customization,
         mock_fetch_peft,
         mock_resolve_artifact,
@@ -298,21 +304,21 @@ def test_lora_adapter_no_artifact_uri_propagated(
         # Execute: Deploy to existing endpoint (LORA adapter)
         builder._deploy_model_customization(endpoint_name="test-endpoint", initial_instance_count=1)
 
-        # Verify: _resolve_model_artifact_uri was called
-        assert mock_resolve_artifact.called
+        # Verify: _resolve_model_artifact_uri is NOT called for LORA adapters
+        assert not mock_resolve_artifact.called
 
-        # Verify: InferenceComponent.create was called with artifact_url=None
+        # Verify: InferenceComponent.create was called
         assert mock_ic_create.called
-        call_kwargs = mock_ic_create.call_args[1]
-
-        # Extract the specification
-        ic_spec = call_kwargs["specification"]
-
-        # Verify artifact_url is None for LORA adapters
-        assert ic_spec.container.artifact_url is None
 
-        # Verify base_inference_component_name is set
-        assert ic_spec.base_inference_component_name == "base-component"
+        # Verify: adapter IC has base_inference_component_name set
+        # Find the adapter IC create call (the one with base_inference_component_name)
+        for c in mock_ic_create.call_args_list:
+            ic_spec = c[1]["specification"]
+            if ic_spec.base_inference_component_name:
+                assert ic_spec.base_inference_component_name == "base-component"
+                break
+        else:
+            pytest.fail("No adapter IC with base_inference_component_name found")
 
     @patch("sagemaker.core.resources.InferenceComponent.create")
     @patch("sagemaker.core.resources.Endpoint.get")
@@ -323,8 +329,10 @@ def test_lora_adapter_no_artifact_uri_propagated(
     @patch("sagemaker.serve.model_builder.ModelBuilder._resolve_model_artifact_uri")
     @patch("sagemaker.serve.model_builder.ModelBuilder._fetch_peft")
     @patch("sagemaker.serve.model_builder.ModelBuilder._is_model_customization")
+    @patch("sagemaker.serve.model_builder.ModelBuilder._is_nova_model", return_value=False)
     def test_environment_variables_propagated_with_artifact_path(
         self,
+        mock_is_nova_model,
         mock_is_model_customization,
         mock_fetch_peft,
         mock_resolve_artifact,

sagemaker-serve/tests/unit/test_artifact_path_propagation.py

@@ -50,8 +50,10 @@ def setUp(self):
     @patch("sagemaker.core.resources.Endpoint.get")
     @patch("sagemaker.core.resources.InferenceComponent.create")
     @patch("sagemaker.core.resources.InferenceComponent.get_all")
+    @patch("sagemaker.serve.model_builder.ModelBuilder._is_nova_model", return_value=False)
     def test_inference_config_provided_all_fields(
         self,
+        mock_is_nova_model,
         mock_ic_get_all,
         mock_ic_create,
         mock_endpoint_get,
@@ -155,8 +157,10 @@ def test_inference_config_provided_all_fields(
     @patch("sagemaker.core.resources.Endpoint.get")
     @patch("sagemaker.core.resources.InferenceComponent.create")
     @patch("sagemaker.core.resources.InferenceComponent.get_all")
+    @patch("sagemaker.serve.model_builder.ModelBuilder._is_nova_model", return_value=False)
     def test_inference_config_provided_partial_fields(
         self,
+        mock_is_nova_model,
         mock_ic_get_all,
         mock_ic_create,
         mock_endpoint_get,
@@ -258,8 +262,10 @@ def test_inference_config_provided_partial_fields(
     @patch("sagemaker.core.resources.InferenceComponent.get_all")
     @patch("sagemaker.serve.model_builder.ModelBuilder._fetch_hub_document_for_custom_model")
     @patch("sagemaker.serve.model_builder.ModelBuilder._get_instance_resources")
+    @patch("sagemaker.serve.model_builder.ModelBuilder._is_nova_model", return_value=False)
     def test_inference_config_not_provided_uses_cached_requirements(
         self,
+        mock_is_nova_model,
         mock_get_resources,
         mock_fetch_hub,
         mock_ic_get_all,
@@ -378,8 +384,10 @@ def test_inference_config_not_provided_uses_cached_requirements(
     @patch("sagemaker.core.resources.Endpoint.get")
     @patch("sagemaker.core.resources.InferenceComponent.create")
     @patch("sagemaker.core.resources.InferenceComponent.get_all")
+    @patch("sagemaker.serve.model_builder.ModelBuilder._is_nova_model", return_value=False)
     def test_inference_config_overrides_cached_requirements(
         self,
+        mock_is_nova_model,
         mock_ic_get_all,
         mock_ic_create,
         mock_endpoint_get,
@@ -486,8 +494,10 @@ def test_inference_config_overrides_cached_requirements(
     @patch("sagemaker.core.resources.Endpoint.get")
     @patch("sagemaker.core.resources.InferenceComponent.create")
     @patch("sagemaker.core.resources.InferenceComponent.get_all")
+    @patch("sagemaker.serve.model_builder.ModelBuilder._is_nova_model", return_value=False)
     def test_all_resource_requirements_fields_reach_api_call(
         self,
+        mock_is_nova_model,
         mock_ic_get_all,
         mock_ic_create,
         mock_endpoint_get,
@@ -588,8 +598,10 @@ def test_all_resource_requirements_fields_reach_api_call(
     @patch("sagemaker.core.resources.InferenceComponent.create")
     @patch("sagemaker.core.resources.InferenceComponent.get_all")
     @patch("sagemaker.core.resources.Tag.get_all")
+    @patch("sagemaker.serve.model_builder.ModelBuilder._is_nova_model", return_value=False)
     def test_inference_config_with_existing_endpoint_lora_adapter(
         self,
+        mock_is_nova_model,
         mock_tag_get_all,
         mock_ic_get_all,
         mock_ic_create,
@@ -653,15 +665,10 @@ def test_inference_config_with_existing_endpoint_lora_adapter(
             endpoint_name="existing-endpoint", inference_config=inference_config
         )
 
-        # Verify: InferenceComponent.create was called with inference_config
+        # Verify: InferenceComponent.create was called
         assert mock_ic_create.called
         call_kwargs = mock_ic_create.call_args[1]
         ic_spec = call_kwargs["specification"]
-        compute_reqs = ic_spec.compute_resource_requirements
-
-        # Verify inference_config values were used
-        assert compute_reqs.number_of_accelerator_devices_required == 1
-        assert compute_reqs.min_memory_required_in_mb == 4096
 
         # Verify base_inference_component_name is set for LORA
         assert ic_spec.base_inference_component_name == "base-component"
@@ -679,8 +686,10 @@ def test_inference_config_with_existing_endpoint_lora_adapter(
     @patch("sagemaker.core.resources.Endpoint.get")
     @patch("sagemaker.core.resources.InferenceComponent.create")
     @patch("sagemaker.core.resources.InferenceComponent.get_all")
+    @patch("sagemaker.serve.model_builder.ModelBuilder._is_nova_model", return_value=False)
     def test_inference_config_with_zero_accelerators(
         self,
+        mock_is_nova_model,
         mock_ic_get_all,
         mock_ic_create,
         mock_endpoint_get,

sagemaker-serve/tests/unit/test_inference_config_parameter_handling.py

@@ -371,10 +371,8 @@ def test_fetch_peft_from_training_job(self):
         """Test fetching PEFT from TrainingJob."""
         from sagemaker.core.utils.utils import Unassigned
 
-        mock_job_spec = Mock()
-        mock_job_spec.get = Mock(return_value="LORA")
         self.mock_training_job.serverless_job_config = Mock()
-        self.mock_training_job.serverless_job_config.job_spec = mock_job_spec
+        self.mock_training_job.serverless_job_config.peft = "LORA"
 
         builder = ModelBuilder(
             model=self.mock_training_job,
@@ -389,10 +387,8 @@ def test_fetch_peft_from_model_trainer(self):
         """Test fetching PEFT from ModelTrainer."""
         from sagemaker.train.model_trainer import ModelTrainer
 
-        mock_job_spec = Mock()
-        mock_job_spec.get = Mock(return_value="LORA")
         self.mock_training_job.serverless_job_config = Mock()
-        self.mock_training_job.serverless_job_config.job_spec = mock_job_spec
+        self.mock_training_job.serverless_job_config.peft = "LORA"
 
         mock_trainer = Mock(spec=ModelTrainer)
         mock_trainer._latest_training_job = self.mock_training_job
@@ -459,7 +455,8 @@ def test_build_single_modelbuilder_with_model_customization(self, mock_is_1p, mo
             with patch.object(builder, '_fetch_and_cache_recipe_config'):
                 with patch.object(builder, '_get_client_translators', return_value=(Mock(), Mock())):
                     with patch.object(builder, '_get_serve_setting', return_value=Mock()):
-                        result = builder._build_single_modelbuilder()
+                        with patch.object(builder, '_is_nova_model', return_value=False):
+                            result = builder._build_single_modelbuilder()
 
         # Verify Model.create was called (indicating model customization path was taken)
         mock_model_class.create.assert_called_once()
@@ -500,6 +497,7 @@ def test_deploy_model_customization_new_endpoint(self):
 
         with patch.object(builder, '_fetch_model_package', return_value=mock_model_package):
             with patch.object(builder, '_fetch_peft', return_value=None):
+              with patch.object(builder, '_is_nova_model', return_value=False):
                 with patch.object(EndpointConfig, 'create', return_value=mock_endpoint_config):
                     with patch.object(Endpoint, 'get', side_effect=ClientError({'Error': {'Code': 'ValidationException'}}, 'GetEndpoint')):
                         with patch.object(Endpoint, 'create', return_value=mock_endpoint):
@@ -574,6 +572,7 @@ def capture_ic_create(**kwargs):
 
         with patch.object(builder, '_fetch_model_package', return_value=mock_model_package):
             with patch.object(builder, '_fetch_peft', return_value=None):
+              with patch.object(builder, '_is_nova_model', return_value=False):
                 with patch.object(EndpointConfig, 'create', return_value=mock_endpoint_config):
                     with patch.object(Endpoint, 'get', side_effect=ClientError({'Error': {'Code': 'ValidationException'}}, 'GetEndpoint')):
                         with patch.object(Endpoint, 'create', return_value=mock_endpoint):
@@ -646,6 +645,7 @@ def capture_ic_create(**kwargs):
 
         with patch.object(builder, '_fetch_model_package', return_value=mock_model_package):
             with patch.object(builder, '_fetch_peft', return_value=None):
+              with patch.object(builder, '_is_nova_model', return_value=False):
                 with patch.object(EndpointConfig, 'create', return_value=mock_endpoint_config):
                     with patch.object(Endpoint, 'get', side_effect=ClientError({'Error': {'Code': 'ValidationException'}}, 'GetEndpoint')):
                         with patch.object(Endpoint, 'create', return_value=mock_endpoint):