fix: add BLACKWELL_96 pool, RTX PRO 6000 Max-Q type, and deprecated NVIDIA_ aliases

Author: KAJdev

src/runpod_flash/core/resources/gpu.py

@@ -64,6 +64,9 @@ class GpuGroup(Enum):
     HOPPER_141 = "HOPPER_141"
     """NVIDIA H200"""
 
+    BLACKWELL_96 = "BLACKWELL_96"
+    """NVIDIA RTX PRO 6000 Blackwell Server Edition, NVIDIA RTX PRO 6000 Blackwell Workstation Edition, NVIDIA RTX PRO 6000 Blackwell Max-Q Workstation Edition"""
+
     BLACKWELL_180 = "BLACKWELL_180"
     """NVIDIA B200"""
 
@@ -172,6 +175,9 @@ class GpuType(Enum):
     RTX_PRO_6000_BLACKWELL_WORKSTATION_EDITION = (
         "NVIDIA RTX PRO 6000 Blackwell Workstation Edition"
     )
+    RTX_PRO_6000_BLACKWELL_MAX_Q_WORKSTATION_EDITION = (
+        "NVIDIA RTX PRO 6000 Blackwell Max-Q Workstation Edition"
+    )
     H100_80GB_HBM3 = "NVIDIA H100 80GB HBM3"
     RTX_A4000 = "NVIDIA RTX A4000"
     RTX_A4500 = "NVIDIA RTX A4500"
@@ -200,6 +206,26 @@ def is_gpu_type(cls, gpu_type: str) -> bool:
         return gpu_type in {m.value for m in cls}
 
 
+# deprecated aliases (old NVIDIA_ prefix names) set outside the class
+# so they don't interfere with cloudpickle serialization
+GpuType.NVIDIA_GEFORCE_RTX_4090 = GpuType.GEFORCE_RTX_4090
+GpuType.NVIDIA_GEFORCE_RTX_5090 = GpuType.GEFORCE_RTX_5090
+GpuType.NVIDIA_RTX_6000_ADA_GENERATION = GpuType.RTX_6000_ADA_GENERATION
+GpuType.NVIDIA_H100_80GB_HBM3 = GpuType.H100_80GB_HBM3
+GpuType.NVIDIA_RTX_A4000 = GpuType.RTX_A4000
+GpuType.NVIDIA_RTX_A4500 = GpuType.RTX_A4500
+GpuType.NVIDIA_RTX_4000_ADA_GENERATION = GpuType.RTX_4000_ADA_GENERATION
+GpuType.NVIDIA_RTX_2000_ADA_GENERATION = GpuType.RTX_2000_ADA_GENERATION
+GpuType.NVIDIA_RTX_A5000 = GpuType.RTX_A5000
+GpuType.NVIDIA_L4 = GpuType.L4
+GpuType.NVIDIA_GEFORCE_RTX_3090 = GpuType.GEFORCE_RTX_3090
+GpuType.NVIDIA_A40 = GpuType.A40
+GpuType.NVIDIA_RTX_A6000 = GpuType.RTX_A6000
+GpuType.NVIDIA_A100_80GB_PCIe = GpuType.A100_80GB_PCIe
+GpuType.NVIDIA_A100_SXM4_80GB = GpuType.A100_SXM4_80GB
+GpuType.NVIDIA_H200 = GpuType.H200
+
+
 POOLS_TO_TYPES = {
     GpuGroup.ADA_24: [GpuType.GEFORCE_RTX_4090],
     GpuGroup.ADA_32_PRO: [GpuType.GEFORCE_RTX_5090],
@@ -219,6 +245,11 @@ def is_gpu_type(cls, gpu_type: str) -> bool:
     GpuGroup.AMPERE_48: [GpuType.A40, GpuType.RTX_A6000],
     GpuGroup.AMPERE_80: [GpuType.A100_80GB_PCIe, GpuType.A100_SXM4_80GB],
     GpuGroup.HOPPER_141: [GpuType.H200],
+    GpuGroup.BLACKWELL_96: [
+        GpuType.RTX_PRO_6000_BLACKWELL_SERVER_EDITION,
+        GpuType.RTX_PRO_6000_BLACKWELL_WORKSTATION_EDITION,
+        GpuType.RTX_PRO_6000_BLACKWELL_MAX_Q_WORKSTATION_EDITION,
+    ],
     GpuGroup.BLACKWELL_180: [GpuType.B200],
 }
 

tests/unit/resources/test_gpu_ids.py

@@ -1,4 +1,4 @@
-from runpod_flash.core.resources.gpu import GpuGroup, GpuType
+from runpod_flash.core.resources.gpu import GpuGroup, GpuType, POOLS_TO_TYPES
 
 
 class TestGpuIdsImports:
@@ -23,3 +23,44 @@ def test_from_gpu_ids_str_pool_only_returns_group(self):
     def test_gpu_type_is_gpu_type_checks_enum_values(self):
         assert GpuType.is_gpu_type("L4") is False
         assert GpuType.is_gpu_type("NVIDIA L4") is True
+
+    def test_blackwell_groups_round_trip(self):
+        gpu_ids = GpuGroup.to_gpu_ids_str([GpuGroup.BLACKWELL_96])
+        assert "BLACKWELL_96" in gpu_ids
+        parsed = GpuGroup.from_gpu_ids_str(gpu_ids)
+        assert parsed == [GpuGroup.BLACKWELL_96]
+
+        gpu_ids = GpuGroup.to_gpu_ids_str([GpuGroup.BLACKWELL_180])
+        assert "BLACKWELL_180" in gpu_ids
+        parsed = GpuGroup.from_gpu_ids_str(gpu_ids)
+        assert parsed == [GpuGroup.BLACKWELL_180]
+
+    def test_b200_type_maps_to_blackwell_180(self):
+        gpu_ids = GpuGroup.to_gpu_ids_str([GpuType.B200])
+        assert "BLACKWELL_180" in gpu_ids
+        # b200 is the only type in BLACKWELL_180, so no negations needed
+        # and from_gpu_ids_str returns the group
+        parsed = GpuGroup.from_gpu_ids_str(gpu_ids)
+        assert parsed == [GpuGroup.BLACKWELL_180]
+
+    def test_rtx_pro_6000_type_maps_to_blackwell_96(self):
+        gpu_ids = GpuGroup.to_gpu_ids_str(
+            [GpuType.RTX_PRO_6000_BLACKWELL_SERVER_EDITION]
+        )
+        assert "BLACKWELL_96" in gpu_ids
+        # other RTX PRO 6000 variants are negated
+        parsed = GpuGroup.from_gpu_ids_str(gpu_ids)
+        assert GpuType.RTX_PRO_6000_BLACKWELL_SERVER_EDITION in parsed
+
+    def test_every_gpu_type_has_pool_mapping(self):
+        all_mapped = set()
+        for types in POOLS_TO_TYPES.values():
+            all_mapped.update(types)
+        for gpu_type in GpuType.all():
+            assert gpu_type in all_mapped, f"{gpu_type.name} has no pool mapping"
+
+    def test_deprecated_aliases_resolve_to_canonical(self):
+        assert GpuType.NVIDIA_L4 is GpuType.L4
+        assert GpuType.NVIDIA_A40 is GpuType.A40
+        assert GpuType.NVIDIA_H200 is GpuType.H200
+        assert GpuType.NVIDIA_GEFORCE_RTX_4090 is GpuType.GEFORCE_RTX_4090