From 64060e63157ff64d3a32a3b358b86dbf4f4d5702 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?oliver=20k=C3=B6nig?= Date: Fri, 6 Feb 2026 16:28:21 +0100 Subject: [PATCH] fix: Use nargs for `custom_bash_cmds` (#2261) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: oliver könig Signed-off-by: NeMo Bot --- scripts/performance/argument_parser.py | 7 ++++--- scripts/performance/setup_experiment.py | 2 +- scripts/performance/utils/executors.py | 4 ++-- 3 files changed, 7 insertions(+), 6 deletions(-) diff --git a/scripts/performance/argument_parser.py b/scripts/performance/argument_parser.py index 0b6fdca692..97b550bdfa 100644 --- a/scripts/performance/argument_parser.py +++ b/scripts/performance/argument_parser.py @@ -422,9 +422,10 @@ def parse_cli_args(): slurm_args.add_argument( "-cb", "--custom_bash_cmds", - type=list_of_strings, - help="Comma separated string of bash commands", - default=[], + nargs="*", + action="append", + help="List of bash commands to execute before the main command", + default=None, ) slurm_args.add_argument( "--gres", diff --git a/scripts/performance/setup_experiment.py b/scripts/performance/setup_experiment.py index f40f60717c..cc3204e825 100755 --- a/scripts/performance/setup_experiment.py +++ b/scripts/performance/setup_experiment.py @@ -206,7 +206,7 @@ def main( custom_mounts: List[str], custom_env_vars: Dict[str, str], custom_srun_args: List[str], - custom_bash_cmds: List[str], + custom_bash_cmds: List[List[str]], nccl_ub: bool, pretrained_checkpoint: Optional[str], num_gpus: int, diff --git a/scripts/performance/utils/executors.py b/scripts/performance/utils/executors.py index bd8eaffc20..4592b355f0 100644 --- a/scripts/performance/utils/executors.py +++ b/scripts/performance/utils/executors.py @@ -64,7 +64,7 @@ def slurm_executor( nemo_home: str = DEFAULT_NEMO_HOME, wandb_key: str = None, network: str = None, - custom_bash_cmds: List[str] = None, + custom_bash_cmds: List[List[str]] = None, additional_slurm_params: Dict[str, Any] = None, gres: Optional[str] = None, ) -> run.SlurmExecutor: @@ -79,7 +79,7 @@ def slurm_executor( #SBATCH --nodelist=node001,node002 #SBATCH --constraint=gpu """ - custom_bash_cmds = [] if custom_bash_cmds is None else custom_bash_cmds + custom_bash_cmds = [] if custom_bash_cmds is None else [" ".join(cmd) for cmd in custom_bash_cmds] mounts = [] # Explicitly request GPU resources to ensure proper allocation # Without --gres=gpu:N, some clusters only allocate 1 GPU regardless of ntasks_per_node