EBS Volue

Author: Jonathan Makunga

src/sagemaker/utils.py

@@ -1454,10 +1454,15 @@ def volume_size_supported(instance_type: str) -> bool:
         if len(parts) != 2:
             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
-        # does not support attaching an EBS volume.
+        # 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.
         family = parts[0]
-        return "d" not in family and not family.startswith("g5")
+        return (
+            "d" not in family
+            and not family.startswith("g5")
+            and not family.startswith("g6")
+            and not family.startswith("p5")
+        )
     except Exception as e:
         raise ValueError(f"Failed to parse instance type '{instance_type}': {str(e)}")