Add GPU directive support check to SlurmSystem and use it in command gen (#541)

Author: TaekyungHeo

src/cloudai/systems/slurm/slurm_system.py

@@ -139,6 +139,7 @@ class SlurmSystem(BaseModel, System):
     cmd_shell: CommandShell = Field(default=CommandShell(), exclude=True)
     extra_srun_args: Optional[str] = None
     extra_sbatch_args: list[str] = []
+    supports_gpu_directives_cache: Optional[bool] = Field(default=None, exclude=True)
 
     data_repository: Optional[DataRepositoryConfig] = None
 
@@ -165,6 +166,29 @@ def groups(self) -> Dict[str, Dict[str, List[SlurmNode]]]:
 
         return groups
 
+    @property
+    def supports_gpu_directives(self) -> bool:
+        if self.supports_gpu_directives_cache is not None:
+            return self.supports_gpu_directives_cache
+
+        stdout, stderr = self.fetch_command_output("scontrol show config")
+        if stderr:
+            logging.warning(f"Error checking GPU support: {stderr}")
+            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
+
+        self.supports_gpu_directives_cache = has_gres_gpu and has_gpu_gres_type
+        return self.supports_gpu_directives_cache
+
     @field_serializer("install_path", "output_path")
     def _path_serializer(self, v: Path) -> str:
         return str(v)

src/cloudai/systems/slurm/strategy/slurm_command_gen_strategy.py

@@ -378,9 +378,10 @@ def _append_sbatch_directives(self, batch_script_content: List[str], tr: TestRun
 
         hostfile = self._append_nodes_related_directives(batch_script_content, tr)
 
-        if self.system.gpus_per_node:
+        if self.system.gpus_per_node and self.system.supports_gpu_directives:
             batch_script_content.append(f"#SBATCH --gpus-per-node={self.system.gpus_per_node}")
             batch_script_content.append(f"#SBATCH --gres=gpu:{self.system.gpus_per_node}")
+
         if self.system.ntasks_per_node:
             batch_script_content.append(f"#SBATCH --ntasks-per-node={self.system.ntasks_per_node}")
         if tr.time_limit:

src/cloudai/workloads/nemo_run/slurm_command_gen_strategy.py

@@ -134,6 +134,19 @@ def generate_test_command(
 
         cmd_args_dict["trainer"]["num_nodes"] = num_nodes
 
+        if self.system.gpus_per_node:
+            trainer_config = cmd_args_dict.get("trainer", {})
+            if "devices" in trainer_config:
+                user_devices = trainer_config["devices"]
+                if user_devices != self.system.gpus_per_node:
+                    logging.warning(
+                        f"User-specified trainer.devices ({user_devices}) differs from "
+                        f"system gpus_per_node ({self.system.gpus_per_node})"
+                    )
+            cmd_args_dict["trainer"]["devices"] = self.system.gpus_per_node
+        else:
+            logging.debug("SlurmSystem.gpus_per_node is not set. Skipping trainer.devices injection.")
+
         self.append_flattened_dict("", cmd_args_dict, command)
 
         if tr.test.extra_cmd_args:

tests/test_slurm_system.py

@@ -530,3 +530,44 @@ def test_per_step_isolation(self, mock_get_nodes: Mock, slurm_system: SlurmSyste
         res = strategy.get_cached_nodes_spec(test_run)
         assert mock_get_nodes.call_count == 2
         assert res == (2, ["node03", "node04"])
+
+
+@pytest.mark.parametrize(
+    "scontrol_output,expected_support",
+    [
+        # Case 1: gres/gpu in AccountingStorageTRES and gpu in GresTypes - 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
+        (
+            """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
+        (
+            """Configuration data as of 2023-06-14T16:28:09
+AccountingStorageTRES   = cpu,mem,energy,node,billing,fs/disk,vmem,pages
+GresTypes               = gpu""",
+            False,
+        ),
+        # Case 4: No gres/gpu in AccountingStorageTRES and 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
+GresTypes               = (null)""",
+            False,
+        ),
+    ],
+)
+@patch("cloudai.systems.slurm.slurm_system.SlurmSystem.fetch_command_output")
+def test_supports_gpu_directives(
+    mock_fetch_command_output, scontrol_output: str, expected_support: bool, slurm_system: SlurmSystem
+):
+    mock_fetch_command_output.return_value = (scontrol_output, "")
+    assert slurm_system.supports_gpu_directives == expected_support