Fix

Author: nargokul

sagemaker-mlops/tests/unit/sagemaker/mlops/feature_store/test_telemetry_integration.py

@@ -831,34 +831,39 @@ def _resolve_compute_requirements_from_config(
                 final_accelerators = user_resource_requirements.num_accelerators
 
         # Determine accelerator count for GPU instances if not provided
-        if final_accelerators is None:
-            # Check if this is a GPU instance type
-            # GPU families: g5, g4dn, g6, p3, p4d, p4de, p5
-            gpu_patterns = ['.g5.', '.g4dn.', '.g6.', '.p3.', '.p4d.', '.p4de.', '.p5.']
-            is_gpu_instance = any(pattern in instance_type for pattern in gpu_patterns)
-
-            if is_gpu_instance:
-                # Try to infer accelerator count from instance type
-                accelerator_count = self._infer_accelerator_count_from_instance_type(instance_type)
-                if accelerator_count is not None:
-                    final_accelerators = accelerator_count
-                    logger.info(
-                        f"Inferred {final_accelerators} accelerator device(s) for instance type {instance_type}"
-                    )
-                else:
-                    # Cannot determine accelerator count - raise descriptive error
-                    raise ValueError(
-                        f"Instance type '{instance_type}' requires accelerator device count specification.\n"
-                        f"Please provide ResourceRequirements with number of accelerators:\n\n"
-                        f"    from sagemaker.core.inference_config import ResourceRequirements\n\n"
-                        f"    resource_requirements = ResourceRequirements(\n"
-                        f"        requests={{\n"
-                        f"            'num_accelerators': <number_of_gpus>,\n"
-                        f"            'memory': {final_min_memory}\n"
-                        f"        }}\n"
-                        f"    )\n\n"
-                        f"For {instance_type}, check AWS documentation for the number of GPUs available."
-                    )
+        # Also strip accelerator count for CPU instances (AWS rejects it)
+        gpu_patterns = ['.g5.', '.g4dn.', '.g6.', '.p3.', '.p4d.', '.p4de.', '.p5.', '.trn', '.inf']
+        is_gpu_instance = any(pattern in instance_type for pattern in gpu_patterns)
+
+        if not is_gpu_instance:
+            # CPU instance - must NOT include accelerator count
+            if final_accelerators is not None:
+                logger.info(
+                    f"Removing accelerator count ({final_accelerators}) for CPU instance type {instance_type}"
+                )
+            final_accelerators = None
+        elif final_accelerators is None:
+            # GPU instance without accelerator count - try to infer
+            accelerator_count = self._infer_accelerator_count_from_instance_type(instance_type)
+            if accelerator_count is not None:
+                final_accelerators = accelerator_count
+                logger.info(
+                    f"Inferred {final_accelerators} accelerator device(s) for instance type {instance_type}"
+                )
+            else:
+                # Cannot determine accelerator count - raise descriptive error
+                raise ValueError(
+                    f"Instance type '{instance_type}' requires accelerator device count specification.\n"
+                    f"Please provide ResourceRequirements with number of accelerators:\n\n"
+                    f"    from sagemaker.core.inference_config import ResourceRequirements\n\n"
+                    f"    resource_requirements = ResourceRequirements(\n"
+                    f"        requests={{\n"
+                    f"            'num_accelerators': <number_of_gpus>,\n"
+                    f"            'memory': {final_min_memory}\n"
+                    f"        }}\n"
+                    f"    )\n\n"
+                    f"For {instance_type}, check AWS documentation for the number of GPUs available."
+                )
 
         # Validate requirements are compatible with instance type
         # Only validate user-provided requirements (defaults are already adjusted above)

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

@@ -900,7 +900,7 @@ def test_both_cpu_and_memory_incompatible(self, mock_get_resources, mock_fetch_h
     @patch('sagemaker.serve.model_builder.ModelBuilder._fetch_hub_document_for_custom_model')
     @patch('sagemaker.serve.model_builder.ModelBuilder._get_instance_resources')
     def test_zero_accelerator_count_explicit(self, mock_get_resources, mock_fetch_hub):
-        """Test that explicitly setting 0 accelerators works for CPU instances."""
+        """Test that explicitly setting 0 accelerators on CPU instance is stripped."""
         # Setup
         mock_fetch_hub.return_value = {
             "HostingConfigs": [
@@ -927,7 +927,7 @@ def test_zero_accelerator_count_explicit(self, mock_get_resources, mock_fetch_hu
             instance_type="ml.m5.xlarge"
         )
 
-        # User explicitly sets 0 accelerators
+        # User explicitly sets 0 accelerators on a CPU instance
         user_requirements = ResourceRequirements(
             requests={
                 "num_accelerators": 0,
@@ -942,8 +942,9 @@ def test_zero_accelerator_count_explicit(self, mock_get_resources, mock_fetch_hu
             user_resource_requirements=user_requirements
         )
 
-        # Verify: Should accept 0 accelerators
-        assert requirements.number_of_accelerator_devices_required == 0
+        # Verify: Accelerator count is stripped for CPU instances
+        from sagemaker.core.utils.utils import Unassigned
+        assert isinstance(requirements.number_of_accelerator_devices_required, Unassigned)
 
 
 if __name__ == "__main__":