Add support for cluster account and gpus_per_node in command generation

NVIDIA · Sep 27, 2024 · f3b4891 · f3b4891
1 parent 648c6a4
commit f3b4891
Showing 1 changed file with 7 additions and 0 deletions.
diff --git a/src/cloudai/schema/test_template/nemo_launcher/slurm_command_gen_strategy.py b/src/cloudai/schema/test_template/nemo_launcher/slurm_command_gen_strategy.py
@@ -80,6 +80,13 @@ def gen_exec_command(
         del self.final_cmd_args["repository_commit_hash"]
         del self.final_cmd_args["docker_image_url"]
 
+        self.final_cmd_args["cluster.account"] = (
+            self.slurm_system.account if self.slurm_system.account is not None else "null"
+        )
+        self.final_cmd_args["cluster.gpus_per_node"] = (
+            self.slurm_system.gpus_per_node if self.slurm_system.gpus_per_node is not None else "null"
+        )
+
         cmd_args_str = self._generate_cmd_args_str(self.final_cmd_args, nodes)
 
         full_cmd = f"python {launcher_path}/launcher_scripts/main.py {cmd_args_str}"