Refactor supports_gpu_directives to focus on GresTypes (#556)

Author: TaekyungHeo

src/cloudai/systems/slurm/slurm_system.py

@@ -178,17 +178,13 @@ def supports_gpu_directives(self) -> bool:
             self.supports_gpu_directives_cache = True
             return True
 
-        has_gres_gpu = False
-        has_gpu_gres_type = False
-
         for line in stdout.splitlines():
-            if "AccountingStorageTRES" in line and "gres/gpu" in line:
-                has_gres_gpu = True
-            if "GresTypes" in line and "gpu" in line and "(null)" not in line:
-                has_gpu_gres_type = True
+            if "GresTypes" in line and "gpu" in line:
+                self.supports_gpu_directives_cache = True
+                return True
 
-        self.supports_gpu_directives_cache = has_gres_gpu and has_gpu_gres_type
-        return self.supports_gpu_directives_cache
+        self.supports_gpu_directives_cache = False
+        return False
 
     @field_serializer("install_path", "output_path")
     def _path_serializer(self, v: Path) -> str:

tests/test_slurm_system.py

@@ -535,33 +535,29 @@ def test_per_step_isolation(self, mock_get_nodes: Mock, slurm_system: SlurmSyste
 @pytest.mark.parametrize(
     "scontrol_output,expected_support",
     [
-        # Case 1: gres/gpu in AccountingStorageTRES and gpu in GresTypes - should be supported
+        # Case 1: GresTypes includes gpu - should be supported
         (
             """Configuration data as of 2023-06-14T16:28:09
-AccountingStorageTRES   = cpu,mem,energy,node,billing,fs/disk,vmem,pages,gres/gpu,gres/gpumem,gres/gpuutil
 GresTypes               = gpu""",
             True,
         ),
-        # Case 2: gres/gpu in AccountingStorageTRES but GresTypes is (null) - should NOT be supported
+        # Case 2: GresTypes is (null) - should NOT be supported
         (
             """Configuration data as of 2023-06-14T16:28:09
-AccountingStorageTRES   = cpu,mem,energy,node,billing,fs/disk,vmem,pages,gres/gpu,gres/gpumem,gres/gpuutil
 GresTypes               = (null)""",
             False,
         ),
-        # Case 3: No gres/gpu in AccountingStorageTRES - should NOT be supported
+        # Case 3: GresTypes does not include gpu - should NOT be supported
         (
             """Configuration data as of 2023-06-14T16:28:09
-AccountingStorageTRES   = cpu,mem,energy,node,billing,fs/disk,vmem,pages
-GresTypes               = gpu""",
+GresTypes               = cpu""",
             False,
         ),
-        # Case 4: No gres/gpu in AccountingStorageTRES and GresTypes is (null) - should NOT be supported
+        # Case 4: GresTypes includes multiple types including gpu - should be supported
         (
             """Configuration data as of 2023-06-14T16:28:09
-AccountingStorageTRES   = cpu,mem,energy,node,billing,fs/disk,vmem,pages
-GresTypes               = (null)""",
-            False,
+GresTypes               = cpu,gpu,fpga""",
+            True,
         ),
     ],
 )
@@ -571,3 +567,14 @@ def test_supports_gpu_directives(
 ):
     mock_fetch_command_output.return_value = (scontrol_output, "")
     assert slurm_system.supports_gpu_directives == expected_support
+
+
+@pytest.mark.parametrize(
+    "cache_value",
+    [True, False],
+)
+@patch("cloudai.systems.slurm.slurm_system.SlurmSystem.fetch_command_output")
+def test_supports_gpu_directives_cache(mock_fetch_command_output, cache_value: bool, slurm_system: SlurmSystem):
+    slurm_system.supports_gpu_directives_cache = cache_value
+    assert slurm_system.supports_gpu_directives is cache_value
+    mock_fetch_command_output.assert_not_called()