Add slurm qos setting (#757)

* add slurm qos setting * set defaults to none
FAIR-Chem · Jul 11, 2024 · 7c3f0d2 · 7c3f0d2
1 parent 51a439e
commit 7c3f0d2
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 1 deletion.
diff --git a/src/fairchem/core/_cli.py b/src/fairchem/core/_cli.py
@@ -81,6 +81,8 @@ def main():
             tasks_per_node=(args.num_gpus if args.distributed else 1),
             nodes=args.num_nodes,
             slurm_additional_parameters=slurm_add_params,
+            slurm_qos=args.slurm_qos,
+            slurm_account=args.slurm_account,
         )
         for config in configs:
             config["slurm"] = copy.deepcopy(executor.parameters)

diff --git a/src/fairchem/core/common/flags.py b/src/fairchem/core/common/flags.py
@@ -92,10 +92,22 @@ def add_core_args(self) -> None:
         )
         self.parser.add_argument(
             "--slurm-partition",
-            default="ocp",
+            default=None,
             type=str,
             help="Name of partition",
         )
+        self.parser.add_argument(
+            "--slurm-account",
+            default=None,
+            type=str,
+            help="Name of account",
+        )
+        self.parser.add_argument(
+            "--slurm-qos",
+            default=None,
+            type=str,
+            help="Name of qos",
+        )
         self.parser.add_argument(
             "--slurm-mem", default=80, type=int, help="Memory (in gigabytes)"
         )