From e4f161c9a6e4ad765b7a7402e33798a9d816108a Mon Sep 17 00:00:00 2001 From: Joseph Zhang Date: Tue, 20 Oct 2026 13:31:00 -0400 Subject: [PATCH 1/2] 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 --- sagemaker-core/src/sagemaker/core/common_utils.py | 4 ++-- sagemaker-core/tests/unit/test_common_utils.py | 10 ++++++++++ 2 files changed, 12 insertions(+), 2 deletions(-) diff --git a/sagemaker-core/src/sagemaker/core/common_utils.py b/sagemaker-core/src/sagemaker/core/common_utils.py index e8704ef600..f76faba505 100644 --- a/sagemaker-core/src/sagemaker/core/common_utils.py +++ b/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 diff --git a/sagemaker-core/tests/unit/test_common_utils.py b/sagemaker-core/tests/unit/test_common_utils.py index 58863335bb..83145ab242 100644 --- a/sagemaker-core/tests/unit/test_common_utils.py +++ b/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 From af314ab417c09eb0776997b75ebe28782bb73143 Mon Sep 17 00:00:00 2001 From: Joseph Zhang Date: Tue, 20 Oct 2026 15:54:54 -0400 Subject: [PATCH 2/2] style: apply black formatting to sagemaker-core and sagemaker-serve unit tests --- X-AI-Prompt: Run formatters and fix issues X-AI-Tool: Claude --- .../tests/unit/local/test_entities.py | 4 ++- sagemaker-serve/tests/unit/conftest.py | 1 + .../test_jumpstart_default_instance_type.py | 25 ++++++++----------- 3 files changed, 15 insertions(+), 15 deletions(-) diff --git a/sagemaker-core/tests/unit/local/test_entities.py b/sagemaker-core/tests/unit/local/test_entities.py index 2c8373c8d1..2d33e35f2c 100644 --- a/sagemaker-core/tests/unit/local/test_entities.py +++ b/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 diff --git a/sagemaker-serve/tests/unit/conftest.py b/sagemaker-serve/tests/unit/conftest.py index ab164faa4b..0a46f7cdfb 100644 --- a/sagemaker-serve/tests/unit/conftest.py +++ b/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 diff --git a/sagemaker-serve/tests/unit/test_jumpstart_default_instance_type.py b/sagemaker-serve/tests/unit/test_jumpstart_default_instance_type.py index e73fe84795..273aca2619 100644 --- a/sagemaker-serve/tests/unit/test_jumpstart_default_instance_type.py +++ b/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()