Skip to content

New slurm customization parameters (account, containers) #1209

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
Kipok merged 17 commits into from
Feb 27, 2026
Merged

New slurm customization parameters (account, containers) #1209

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
17 commits
Select commit Hold shift + click to select a range
b1c2066
Add an option to specify non-default slurm account
Kipok Feb 3, 2026
5a89b68
Add overrides for containers
Kipok Feb 3, 2026
be1b700
Debugging
Kipok Feb 4, 2026
3d7aebd
Client side cuonting
Kipok Feb 4, 2026
ef451dd
Fixed toeknizer
Kipok Feb 4, 2026
9a9ba45
Debugging
Kipok Feb 4, 2026
3b64aa7
Merge remote-tracking branch 'origin/main' into igitman/account-arg
gwarmstrong Feb 25, 2026
bbf2ea7
Merge remote-tracking branch 'origin/main' into igitman/account-arg
gwarmstrong Feb 25, 2026
dd0d94e
fix: add missing account= arg to test_server_metadata_from_num_tasks
gwarmstrong Feb 25, 2026
51f059e
Merge branch 'igitman/account-arg' of https://github.com/NVIDIA/NeMo-…
Kipok Feb 26, 2026
316352b
Roll-back tmp changes
Kipok Feb 26, 2026
dbea2ca
.get -> []
Kipok Feb 26, 2026
a6fe5ab
Merge branch 'main' into igitman/account-arg
Kipok Feb 26, 2026
a7cac9d
fix: add timeouts to API-calling tests to prevent CI hangs
gwarmstrong Feb 26, 2026
5e7965a
fix: also pass inference timeout to judge subprocess in test
gwarmstrong Feb 26, 2026
4374543
fix: disable litellm retries in API tests to prevent timeout compounding
gwarmstrong Feb 26, 2026
92600fb
tst: update test parameters to limit retries and use a different mode…
gwarmstrong Feb 27, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions nemo_skills/inference/model/tool_call.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -145,6 +145,7 @@ async def generate_async(
endpoint_type=endpoint_type,
**generation_kwargs,
)

if isinstance(tokens_to_generate, int):
tokens_to_generate -= generation["num_generated_tokens"]

Expand Down
5 changes: 4 additions & 1 deletion nemo_skills/pipeline/convert.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -180,6 +180,8 @@ def convert(
partition: str = typer.Option(
None, help="Can specify if need interactive jobs or a specific non-default partition"
),
account: str = typer.Option(None, help="Can specify a non-default Slurm account"),
container: str = typer.Option(None, help="Override container image for the conversion job"),
qos: str = typer.Option(None, help="Specify Slurm QoS, e.g. to request interactive nodes"),
time_min: str = typer.Option(None, help="If specified, will use as a time-min slurm parameter"),
mount_paths: str = typer.Option(None, help="Comma separated list of paths to mount on the remote machine"),
Expand Down Expand Up @@ -320,12 +322,13 @@ def convert(
cmd=conversion_cmd,
task_name=expname,
log_dir=log_dir,
container=container_map[(convert_from, convert_to)],
container=container or container_map[(convert_from, convert_to)],
num_gpus=num_gpus,
num_nodes=1, # always running on a single node, might need to change that in the future
num_tasks=1,
cluster_config=cluster_config,
partition=partition,
account=account,
run_after=run_after,
reuse_code=reuse_code,
reuse_code_exp=reuse_code_exp,
Expand Down
23 changes: 22 additions & 1 deletion nemo_skills/pipeline/eval.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -55,6 +55,9 @@ def _create_llm_judge_tasks(
cluster,
config_dir,
partition,
account,
main_container,
sandbox_container,
with_sandbox,
keep_mounts_for_sandbox,
run_after,
Expand Down Expand Up @@ -95,6 +98,9 @@ def _create_llm_judge_tasks(
cluster=cluster,
config_dir=config_dir,
partition=partition,
account=account,
main_container=main_container,
sandbox_container=sandbox_container,
with_sandbox=with_sandbox,
keep_mounts_for_sandbox=keep_mounts_for_sandbox,
run_after=run_after,
Expand Down Expand Up @@ -184,6 +190,12 @@ def eval(
server_container: str = typer.Option(
None, help="Override container image for the hosted server (if server_gpus is set)"
),
main_container: str = typer.Option(None, help="Override container image for the main evaluation client"),
sandbox_container: str = typer.Option(None, help="Override container image for the sandbox"),
judge_container: str = typer.Option(None, help="Override container image for GPU-based judges (comet, nvembed)"),
Comment on lines +193 to +195
Copy link
Collaborator

@gwarmstrong gwarmstrong Feb 3, 2026 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think it's a little bulky to have separate override arguments for each container everywhere. Not sure that there is a better solution though. If we wanted to have overrides like we do for tools, e.g.,

++container_overrides.sandbox = "..."
++container_overrides.judge = "..."

But then the choice of key is unclear--since our "job components", e.g., Judge, main, sandbox, ... don't map cleanly to a container name (e.g., "judge" -> containers[judge_server_type], main -> containers["nemo-skills"], sandbox -> containers["sandbox"]).

So I think with the current structure, what you've done the best choice, but maybe we can eventually work toward something a little more general here.

judge_server_container: str = typer.Option(
None, help="Override container image for the hosted judge server (if judge_server_gpus is set)"
),
extra_judge_args: str = typer.Option(
"", help="Additional arguments for judge (passed to generate script, so should start with ++)"
),
Expand All @@ -210,6 +222,7 @@ def eval(
"Can provide a list directly when using through Python",
),
partition: str = typer.Option(None, help="Cluster partition to use"),
account: str = typer.Option(None, help="Can specify a non-default Slurm account"),
qos: str = typer.Option(None, help="Specify Slurm QoS, e.g. to request interactive nodes"),
time_min: str = typer.Option(None, help="If specified, will use as a time-min slurm parameter"),
mount_paths: str = typer.Option(None, help="Comma separated list of paths to mount on the remote machine"),
Expand Down Expand Up @@ -349,6 +362,7 @@ def eval(
"server_nodes": judge_server_nodes,
"server_args": judge_server_args,
"server_entrypoint": judge_server_entrypoint,
"server_container": judge_server_container,
"generation_type": judge_generation_type,
"generation_module": judge_generation_module,
}
Expand Down Expand Up @@ -430,14 +444,16 @@ def eval(
cmd=pipeline_utils.wrap_python_path(cmd=combine_cmds(cmds, single_node_mode)),
task_name=f"{expname}-{'-'.join(job_benchmarks)}",
log_dir=log_dir,
container=cluster_config["containers"]["nemo-skills"],
container=main_container or cluster_config["containers"]["nemo-skills"],
cluster_config=cluster_config,
partition=partition,
account=account,
server_config=job_server_config,
with_sandbox=job_needs_sandbox or with_sandbox,
keep_mounts_for_sandbox=job_needs_sandbox_to_keep_mounts or keep_mounts_for_sandbox,
sandbox_port=None if get_random_port else 6000,
sandbox_env_overrides=job_sandbox_env_overrides,
sandbox_container=sandbox_container,
run_after=run_after,
reuse_code_exp=reuse_code_exp,
reuse_code=reuse_code,
Expand Down Expand Up @@ -508,6 +524,8 @@ def eval(
judge_server_gpus=judge_server_gpus,
judge_server_nodes=judge_server_nodes,
partition=partition,
account=account,
judge_container=judge_container,
run_after=run_after,
reuse_code_exp=reuse_code_exp,
reuse_code=reuse_code,
Expand Down Expand Up @@ -535,6 +553,9 @@ def eval(
cluster=cluster,
config_dir=config_dir,
partition=partition,
account=account,
main_container=main_container,
sandbox_container=sandbox_container,
with_sandbox=with_sandbox,
keep_mounts_for_sandbox=keep_mounts_for_sandbox,
run_after=run_after,
Expand Down
14 changes: 12 additions & 2 deletions nemo_skills/pipeline/generate.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -55,11 +55,14 @@ def _create_job_unified(
installation_command: Optional[str],
with_sandbox: bool,
partition: Optional[str],
account: Optional[str],
keep_mounts_for_sandbox: bool,
task_name: str,
log_dir: str,
sbatch_kwargs: Optional[Dict] = None,
sandbox_env_overrides: Optional[List[str]] = None,
main_container: Optional[str] = None,
sandbox_container: Optional[str] = None,
) -> List[CommandGroup]:
"""
Create CommandGroups for n models (unified for n=1 and n>1).
Expand Down Expand Up @@ -147,7 +150,7 @@ def _create_job_unified(

sandbox_cmd = Command(
script=sandbox_script,
container=cluster_config["containers"]["sandbox"],
container=sandbox_container or cluster_config["containers"]["sandbox"],
name=f"{task_name}_sandbox",
)
components.append(sandbox_cmd)
Expand Down Expand Up @@ -178,7 +181,7 @@ def _create_job_unified(

client_cmd = Command(
script=client_script,
container=cluster_config["containers"]["nemo-skills"],
container=main_container or cluster_config["containers"]["nemo-skills"],
name=f"{task_name}",
)
components.append(client_cmd)
Expand All @@ -191,6 +194,7 @@ def _create_job_unified(
commands=components,
hardware=HardwareConfig(
partition=partition,
account=account,
num_gpus=group_gpus,
num_nodes=group_nodes,
num_tasks=group_tasks,
Expand Down Expand Up @@ -272,6 +276,8 @@ def generate(
help="Container image(s). CLI: space-separated. Python API: string or list. "
"Single value broadcasts to all models.",
),
main_container: str = typer.Option(None, help="Override container image for the main generation client"),
sandbox_container: str = typer.Option(None, help="Override container image for the sandbox"),
dependent_jobs: int = typer.Option(0, help="Specify this to launch that number of dependent jobs"),
mount_paths: str = typer.Option(None, help="Comma separated list of paths to mount on the remote machine"),
num_random_seeds: int = typer.Option(
Expand All @@ -297,6 +303,7 @@ def generate(
partition: str = typer.Option(
None, help="Can specify if need interactive jobs or a specific non-default partition"
),
account: str = typer.Option(None, help="Can specify a non-default Slurm account"),
qos: str = typer.Option(None, help="Specify Slurm QoS, e.g. to request interactive nodes"),
time_min: str = typer.Option(None, help="If specified, will use as a time-min slurm parameter"),
run_after: List[str] = typer.Option(
Expand Down Expand Up @@ -589,11 +596,14 @@ def convert_server_type_to_string(server_type):
installation_command=installation_command,
with_sandbox=with_sandbox,
partition=partition,
account=account,
keep_mounts_for_sandbox=keep_mounts_for_sandbox,
task_name=task_name,
log_dir=log_dir,
sbatch_kwargs=sbatch_kwargs,
sandbox_env_overrides=sandbox_env_overrides,
main_container=main_container,
sandbox_container=sandbox_container,
)

# Use unique internal job name for dependency tracking, but same task_name
Expand Down
21 changes: 19 additions & 2 deletions nemo_skills/pipeline/nemo_evaluator.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -125,6 +125,7 @@ def nemo_evaluator(
job_gpus: int = typer.Option(0, help="GPUs to allocate for the evaluator client when no servers are hosted"),
job_nodes: int = typer.Option(1, help="Nodes to allocate for the evaluator job"),
partition: str = typer.Option(None, help="Cluster partition to use"),
account: str = typer.Option(None, help="Can specify a non-default Slurm account"),
qos: str = typer.Option(None, help="Slurm QoS"),
mount_paths: str = typer.Option(None, help="Comma separated list of paths to mount on the remote machine"),
log_dir: str = typer.Option(None, help="Custom location for logs"),
Expand Down Expand Up @@ -325,6 +326,7 @@ def nemo_evaluator(
job_nodes=job_nodes,
cluster_config=cluster_config,
partition=partition,
account=account,
qos=qos,
exclusive=exclusive,
)
Expand All @@ -346,6 +348,7 @@ def nemo_evaluator(
commands=[main_server_cmd, client_cmd],
hardware=_hardware_for_group(
task_ctx.partition,
task_ctx.account,
task_ctx.server_gpus or None,
task_ctx.server_nodes or 1,
task_ctx.qos,
Expand All @@ -358,6 +361,7 @@ def nemo_evaluator(
commands=[judge_server_cmd],
hardware=_hardware_for_group(
task_ctx.partition,
task_ctx.account,
task_ctx.judge_server_gpus or None,
task_ctx.judge_server_nodes or 1,
task_ctx.qos,
Expand Down Expand Up @@ -393,7 +397,12 @@ def nemo_evaluator(
CommandGroup(
commands=sg_cmds,
hardware=_hardware_for_group(
task_ctx.partition, group_num_gpus, group_num_nodes, task_ctx.qos, task_ctx.exclusive
task_ctx.partition,
task_ctx.account,
group_num_gpus,
group_num_nodes,
task_ctx.qos,
task_ctx.exclusive,
),
name=f"{task_ctx.expname}-{task_ctx.idx}",
log_dir=log_dir,
Expand Down Expand Up @@ -543,17 +552,24 @@ class _TaskCreationContext:
job_nodes: int
cluster_config: Dict
partition: Optional[str]
account: Optional[str]
qos: Optional[str]
exclusive: bool


def _hardware_for_group(
partition: Optional[str], num_gpus: Optional[int], num_nodes: int, qos: Optional[str], exclusive: bool
partition: Optional[str],
account: Optional[str],
num_gpus: Optional[int],
num_nodes: int,
qos: Optional[str],
exclusive: bool,
) -> HardwareConfig:
"""Create HardwareConfig for a CommandGroup.

Args:
partition: SLURM partition name
account: SLURM account name
num_gpus: Number of GPUs (None means no GPU allocation)
num_nodes: Number of nodes
qos: SLURM QoS setting
Expand All @@ -564,6 +580,7 @@ def _hardware_for_group(
"""
return HardwareConfig(
partition=partition,
account=account,
num_gpus=num_gpus,
num_nodes=num_nodes,
sbatch_kwargs={
Expand Down
4 changes: 4 additions & 0 deletions nemo_skills/pipeline/run_cmd.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -58,6 +58,7 @@ def run_cmd(
partition: str = typer.Option(
None, help="Can specify if need interactive jobs or a specific non-default partition"
),
account: str = typer.Option(None, help="Can specify a non-default Slurm account"),
qos: str = typer.Option(None, help="Specify Slurm QoS, e.g. to request interactive nodes"),
time_min: str = typer.Option(None, help="If specified, will use as a time-min slurm parameter"),
num_gpus: int | None = typer.Option(None, help="Number of GPUs per node to use"),
Expand All @@ -77,6 +78,7 @@ def run_cmd(
server_container: str = typer.Option(
None, help="Override container image for the hosted server (if server_gpus is set)"
),
sandbox_container: str = typer.Option(None, help="Override container image for the sandbox"),
dependent_jobs: int = typer.Option(0, help="Specify this to launch that number of dependent jobs"),
mount_paths: str = typer.Option(None, help="Comma separated list of paths to mount on the remote machine"),
run_after: List[str] = typer.Option(
Expand Down Expand Up @@ -197,10 +199,12 @@ def run_cmd(
container=containers,
cluster_config=cluster_config,
partition=partition,
account=account,
server_config=server_config,
with_sandbox=with_sandbox,
keep_mounts_for_sandbox=keep_mounts_for_sandbox,
sandbox_port=None if get_random_port else 6000,
sandbox_container=sandbox_container,
run_after=run_after,
reuse_code=reuse_code,
reuse_code_exp=reuse_code_exp,
Expand Down
13 changes: 12 additions & 1 deletion nemo_skills/pipeline/start_server.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -125,10 +125,13 @@ def launch_server(
tail_logs=False,
cmd="",
partition=None,
account=None,
with_sandbox=False,
keep_mounts_for_sandbox=False,
server_port=None,
sandbox_port=None,
main_container=None,
sandbox_container=None,
sbatch_kwargs=None,
):
"""Launch a model server in the background.
Expand Down Expand Up @@ -174,13 +177,15 @@ def launch_server(
cmd=cmd,
task_name="server",
log_dir=log_dir,
container=cluster_config["containers"]["nemo-skills"],
container=main_container or cluster_config["containers"]["nemo-skills"],
cluster_config=cluster_config,
partition=partition,
account=account,
server_config=server_config,
with_sandbox=with_sandbox,
keep_mounts_for_sandbox=keep_mounts_for_sandbox,
sandbox_port=sandbox_port,
sandbox_container=sandbox_container,
sbatch_kwargs=sbatch_kwargs,
)
exp.run(detach=True, tail_logs=tail_logs)
Expand Down Expand Up @@ -213,7 +218,10 @@ def start_server(
"If not specified, will use the default entrypoint for the server type.",
),
server_container: str = typer.Option(None, help="Override container image for the hosted server"),
main_container: str = typer.Option(None, help="Override container image for the main task (e.g., chat interface)"),
sandbox_container: str = typer.Option(None, help="Override container image for the sandbox"),
partition: str = typer.Option(None, help="Cluster partition to use"),
account: str = typer.Option(None, help="Can specify a non-default Slurm account"),
qos: str = typer.Option(None, help="Specify Slurm QoS, e.g. to request interactive nodes"),
time_min: str = typer.Option(None, help="If specified, will use as a time-min slurm parameter"),
mount_paths: str = typer.Option(None, help="Comma separated list of paths to mount on the remote machine"),
Expand Down Expand Up @@ -274,10 +282,13 @@ def start_server(
tail_logs=True,
cmd=cmd,
partition=partition,
account=account,
with_sandbox=with_sandbox,
keep_mounts_for_sandbox=keep_mounts_for_sandbox,
server_port=server_port,
sandbox_port=sandbox_port,
main_container=main_container,
sandbox_container=sandbox_container,
sbatch_kwargs=parse_kwargs(sbatch_kwargs, exclusive=exclusive, qos=qos, time_min=time_min),
)

Expand Down
2 changes: 2 additions & 0 deletions nemo_skills/pipeline/utils/declarative.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -264,6 +264,7 @@ class HardwareConfig:
"""Hardware configuration for a group of tasks."""

partition: Optional[str] = None
account: Optional[str] = None
num_gpus: Optional[int] = None
num_nodes: Optional[int] = None
num_tasks: Optional[int] = 1
Expand Down Expand Up @@ -585,6 +586,7 @@ def _create_executor(
log_dir=log_dir,
log_prefix=exec_config["log_prefix"],
partition=hardware.partition if hardware else None,
account=hardware.account if hardware else None,
heterogeneous=heterogeneous,
het_group=het_group,
total_het_groups=total_het_groups,
Expand Down
Loading