fix: add g7e, p6-b200 and trn2 to volume size unsupported instance families (#6505)

* fix: add g7e, p6-b200 and trn2 to volume size unsupported instance families

---
X-AI-Prompt: Add g7e, p6-b200, and trn2 to volume size unsupported instances
X-AI-Tool: Claude

* style: apply black formatting to sagemaker-core and sagemaker-serve unit tests

---
X-AI-Prompt: Run formatters and fix issues
X-AI-Tool: Claude

---------

Co-authored-by: Joseph Zhang <cjz@amazon.com>

Author: cj-zhang

sagemaker-core/src/sagemaker/core/common_utils.py

@@ -1645,10 +1645,10 @@ def volume_size_supported(instance_type: str) -> bool:
             raise ValueError(f"Failed to parse instance type '{instance_type}'")
 
         # Any instance type with a "d" in the instance family (i.e. c5d, p4d, etc)
-        # + g5 or g6 or p5 does not support attaching an EBS volume.
+        # + g5, g6, g7e, p5, p6-b200, trn1 or trn2 does not support attaching an EBS volume.
         family = parts[0]
 
-        unsupported_families = ["g5", "g6", "p5", "trn1"]
+        unsupported_families = ["g5", "g6", "g7e", "p5", "p6-b200", "trn1", "trn2"]
 
         return "d" not in family and not any(
             family.startswith(prefix) for prefix in unsupported_families

sagemaker-core/tests/unit/local/test_entities.py

@@ -762,7 +762,9 @@ def test_wait_custom_timeout(self, mock_sleep, mock_get_host, mock_perform_reque
     @patch("sagemaker.core.local.entities._perform_request")
     @patch("sagemaker.core.local.entities.get_docker_host")
     @patch("time.sleep")
-    def test_wait_uses_default_timeout_constant(self, mock_sleep, mock_get_host, mock_perform_request):
+    def test_wait_uses_default_timeout_constant(
+        self, mock_sleep, mock_get_host, mock_perform_request
+    ):
         """Test that _wait_for_serving_container defaults to HEALTH_CHECK_TIMEOUT_LIMIT"""
         mock_get_host.return_value = "localhost"
         # Succeed on first attempt so we can verify it was called

sagemaker-core/tests/unit/test_common_utils.py

@@ -1100,6 +1100,16 @@ def test_volume_size_supported_g5(self):
 
         assert volume_size_supported("ml.g5.xlarge") is False
 
+    @pytest.mark.parametrize(
+        "instance_type",
+        ["ml.g7e.2xlarge", "ml.p6-b200.48xlarge", "ml.trn2.48xlarge"],
+    )
+    def test_volume_size_supported_unsupported_families(self, instance_type):
+        """Test g7e, p6-b200 and trn2 instances."""
+        from sagemaker.core.common_utils import volume_size_supported
+
+        assert volume_size_supported(instance_type) is False
+
     def test_volume_size_supported_local(self):
         """Test local mode."""
         from sagemaker.core.common_utils import volume_size_supported

sagemaker-serve/tests/unit/conftest.py

@@ -26,6 +26,7 @@
 the suite three times (py310/py311/py312) and these failures almost always hit
 only one of the three.
 """
+
 from __future__ import annotations
 
 from contextlib import ExitStack

sagemaker-serve/tests/unit/test_jumpstart_default_instance_type.py

@@ -35,9 +35,7 @@ def _load_fixture_spec_dict():
     # Exclude ml.m5.large (the buggy code's hardcoded fallback) from the
     # supported list so membership assertions discriminate red from green.
     spec["supported_inference_instance_types"] = [_FIXTURE_DEFAULT_INSTANCE_TYPE] + [
-        t
-        for t in spec.get("supported_inference_instance_types", [])
-        if t != "ml.m5.large"
+        t for t in spec.get("supported_inference_instance_types", []) if t != "ml.m5.large"
     ]
     return spec
 
@@ -50,7 +48,6 @@ def _get_fixture_spec(*args, **kwargs):
     return JumpStartModelSpecs(_load_fixture_spec_dict())
 
 
-
 def _mock_session():
     """Mock sagemaker session following the package's shared-setup convention."""
     mock_session = Mock()
@@ -122,14 +119,14 @@ def spy(self):
         with patch.object(ModelBuilder, "_get_default_instance_type", spy):
             _build_jumpstart_builder()
 
-        self.assertTrue(seen.get("has_region"), "self.region missing during instance-type detection")
+        self.assertTrue(
+            seen.get("has_region"), "self.region missing during instance-type detection"
+        )
         self.assertEqual(seen.get("region"), "us-west-2")
 
     def test_user_provided_instance_type_wins_and_skips_lookup(self):
         """Explicit instance_type must be honored without consulting the spec."""
-        builder, get_deploy_kwargs_mock = _build_jumpstart_builder(
-            instance_type="ml.p5.48xlarge"
-        )
+        builder, get_deploy_kwargs_mock = _build_jumpstart_builder(instance_type="ml.p5.48xlarge")
 
         self.assertEqual(builder.instance_type, "ml.p5.48xlarge")
         self.assertTrue(builder._user_provided_instance_type)
@@ -157,9 +154,9 @@ def test_deploy_passes_spec_default_to_internal_deploy(self):
         builder, _ = _build_jumpstart_builder()
         builder.built_model = Mock()
 
-        with patch.object(ModelBuilder, "_is_model_customization", return_value=False), patch.object(
-            ModelBuilder, "_deploy", return_value=Mock()
-        ) as mock_deploy:
+        with patch.object(
+            ModelBuilder, "_is_model_customization", return_value=False
+        ), patch.object(ModelBuilder, "_deploy", return_value=Mock()) as mock_deploy:
             builder.deploy(endpoint_name="test-endpoint")
 
         mock_deploy.assert_called_once()
@@ -265,9 +262,9 @@ def test_deploy_passes_user_override_to_internal_deploy(self):
         builder, _ = _build_jumpstart_builder(instance_type=self.USER_INSTANCE_TYPE)
         builder.built_model = Mock()
 
-        with patch.object(ModelBuilder, "_is_model_customization", return_value=False), patch.object(
-            ModelBuilder, "_deploy", return_value=Mock()
-        ) as mock_deploy:
+        with patch.object(
+            ModelBuilder, "_is_model_customization", return_value=False
+        ), patch.object(ModelBuilder, "_deploy", return_value=Mock()) as mock_deploy:
             builder.deploy(endpoint_name="test-endpoint")
 
         mock_deploy.assert_called_once()