Update instance type regex to also include hyphens

For commit: aws/sagemaker-python-sdk-staging@824675b

Author: zhaoqizqwang

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

@@ -1555,7 +1555,7 @@ def get_instance_type_family(instance_type: str) -> str:
     """
     instance_type_family = ""
     if isinstance(instance_type, str):
-        match = re.match(r"^ml[\._]([a-z\d]+)\.?\w*$", instance_type)
+        match = re.match(r"^ml[\._]([a-z\d\-]+)\.?\w*$", instance_type)
         if match is not None:
             instance_type_family = match[1]
     return instance_type_family

sagemaker-core/src/sagemaker/core/fw_utils.py

@@ -25,7 +25,6 @@
 
 from packaging import version
 
-from sagemaker.core import image_uris
 import sagemaker.core.common_utils as utils
 from sagemaker.core.deprecations import deprecation_warn_base, renamed_kwargs, renamed_warning
 from sagemaker.core.instance_group import InstanceGroup
@@ -1208,6 +1207,8 @@ def create_image_uri(
     Returns:
         the image uri
     """
+    from sagemaker.core import image_uris
+    
     renamed_warning("The method create_image_uri")
     return image_uris.retrieve(
         framework=framework,

sagemaker-core/tests/unit/test_fw_utils.py

@@ -225,7 +225,7 @@ def test_validate_mp_config_ddp_and_horovod(self):
 class TestTarAndUploadDir:
     """Test tar_and_upload_dir function."""
 
-    @patch("sagemaker.core.fw_utils.sagemaker_utils.create_tar_file")
+    @patch("sagemaker.core.common_utils.create_tar_file")
     def test_tar_and_upload_dir_s3_source(self, mock_create_tar):
         """Test with S3 source directory."""
         mock_session = Mock()
@@ -242,7 +242,7 @@ def test_tar_and_upload_dir_s3_source(self, mock_create_tar):
         assert result.script_name == "train.py"
         mock_create_tar.assert_not_called()
 
-    @patch("sagemaker.core.fw_utils.sagemaker_utils.create_tar_file")
+    @patch("sagemaker.core.common_utils.create_tar_file")
     @patch("sagemaker.core.fw_utils.tempfile.mkdtemp")
     @patch("sagemaker.core.fw_utils.shutil.rmtree")
     def test_tar_and_upload_dir_local_file(
@@ -495,6 +495,8 @@ def test_is_gpu_instance_true(self):
         assert _is_gpu_instance("ml.p3.2xlarge") is True
         assert _is_gpu_instance("ml.g4dn.xlarge") is True
         assert _is_gpu_instance("local_gpu") is True
+        assert _is_gpu_instance("ml.p6-b200.48xlarge") is True
+        assert _is_gpu_instance("ml.g6e-12xlarge.xlarge") is True
 
     def test_is_gpu_instance_false(self):
         """Test _is_gpu_instance with non-GPU instance."""
@@ -505,6 +507,7 @@ def test_is_trainium_instance_true(self):
         """Test _is_trainium_instance with Trainium instance."""
         assert _is_trainium_instance("ml.trn1.2xlarge") is True
         assert _is_trainium_instance("ml.trn1.32xlarge") is True
+        assert _is_trainium_instance("ml.trn1-n.2xlarge") is True
 
     def test_is_trainium_instance_false(self):
         """Test _is_trainium_instance with non-Trainium instance."""
@@ -523,6 +526,7 @@ def test_region_supports_profiler(self):
 
     def test_instance_type_supports_profiler(self):
         """Test _instance_type_supports_profiler."""
+        assert _instance_type_supports_profiler("ml.trn1-n.xlarge") is True
         assert _instance_type_supports_profiler("ml.trn1.2xlarge") is True
         assert _instance_type_supports_profiler("ml.p3.2xlarge") is False
 

sagemaker-serve/src/sagemaker/serve/model_builder_utils.py

@@ -1504,7 +1504,7 @@ def _is_inferentia_or_trainium(self, instance_type: Optional[str]) -> bool:
             bool: Whether the given instance type is Inferentia or Trainium.
         """
         if isinstance(instance_type, str):
-            match = re.match(r"^ml[\._]([a-z\d]+)\.?\w*$", instance_type)
+            match = re.match(r"^ml[\._]([a-z\d\-]+)\.?\w*$", instance_type)
             if match:
                 if match[1].startswith("inf") or match[1].startswith("trn"):
                     return True