diff --git a/src/fairchem/core/_cli.py b/src/fairchem/core/_cli.py index 168b63df4..047cade3c 100644 --- 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 index 536ffe357..ac4bd2f84 100644 --- 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)" )