Allow het servers for nemo-rl jobs

[smahdavi4]

Allow heterogenous servers for nemo-rl jobs

Summary by CodeRabbit

• New Features
  • Added LLM-as-a-judge support to GRPO workflow with flexible server configuration.
  • New CLI options for judge server (model, address, type, GPUs, nodes, args) and multi-server control.
  • Support for local or remote judge hosting with automatic port handling.
  • Scheduler now launches and orders judge server tasks to reserve GPU resources before training.
  • Multi-server scheduling and improved sandbox/overlap behavior for GPU-aware runs.

[greptile-apps]

_{1 file reviewed, 1 comment}

_{Edit Code Review Agent Settings | Greptile}

[greptile-apps]

all servers launched in the loop use the same server_port from server_config, causing port conflicts when n_servers > 1

each server instance needs a unique port

Suggested change

	for server_idx in range(n_servers):
	server_cmd, num_server_tasks = get_server_command(**server_config, cluster_config=cluster_config)
	server_executor = get_executor(
	cluster_config=cluster_config,
	container=server_container,
	num_nodes=server_config["num_nodes"],
	tasks_per_node=num_server_tasks,
	gpus_per_node=server_config["num_gpus"],
	partition=partition,
	dependencies=dependencies,
	job_name=task_name,
	log_dir=log_dir,
	log_prefix=f"server_{server_idx}" if n_servers > 1 else "server",
	extra_package_dirs=extra_package_dirs,
	sbatch_kwargs=sbatch_kwargs,
	heterogeneous=heterogeneous,
	het_group=het_group,
	total_het_groups=total_het_groups,
	overlap=(not client_num_gpus), # Only overlap when the main task does not have gpus
	with_ray=False,
	ray_template=ray_template,
	)
	cmd_to_add = server_cmd
	if cluster_config["executor"] != "slurm" and num_server_tasks > 1:
	cmd_to_add = f"mpirun --allow-run-as-root -np {num_server_tasks} bash -c {shlex.quote(server_cmd)}"
	commands.append(cmd_to_add)
	executors.append(server_executor)
	het_group_indices.append(het_group)
	het_group += 1
	LOG.info("Server %d command: %s", server_idx, server_cmd)
	for server_idx in range(n_servers):
	# Get a unique port for each server if launching multiple
	current_server_config = server_config.copy()
	if n_servers > 1:
	current_server_config["server_port"] = get_free_port(strategy="random")
	server_cmd, num_server_tasks = get_server_command(**current_server_config, cluster_config=cluster_config)
	server_executor = get_executor(
	cluster_config=cluster_config,
	container=server_container,
	num_nodes=current_server_config["num_nodes"],
	tasks_per_node=num_server_tasks,
	gpus_per_node=current_server_config["num_gpus"],
	partition=partition,
	dependencies=dependencies,
	job_name=task_name,
	log_dir=log_dir,
	log_prefix=f"server_{server_idx}" if n_servers > 1 else "server",
	extra_package_dirs=extra_package_dirs,
	sbatch_kwargs=sbatch_kwargs,
	heterogeneous=heterogeneous,
	het_group=het_group,
	total_het_groups=total_het_groups,
	overlap=(not client_num_gpus), # Only overlap when the main task does not have gpus
	with_ray=False,
	ray_template=ray_template,
	)

[coderabbitai]

📝 Walkthrough

Walkthrough

Adds LLM-as-a-judge server support to the GRPO workflow: new CLI server options, server_config/client_server_args construction for local or remote hosting, port allocation, environment injection, and multi-server scheduling integrated into task creation and ordering.

Changes

Cohort / File(s)	Summary
GRPO Server Configuration nemo_skills/pipeline/nemo_rl/grpo.py	Adds server CLI options (server_model, server_address, server_type, server_gpus, server_nodes, n_servers, server_args), imports (json, SupportedServers, get_free_port, should_get_random_port), and logic to build server_config or client_server_args, allocate ports for local hosting, inject JUDGE_SERVER_ARGS, and pass server info into task submission.
Task Scheduling & Multi-Server Support nemo_skills/pipeline/utils/exp.py	Extends add_task(...) signature with n_servers. Adds add_server_tasks helper, multi-server command/executor construction, GPU-aware server-first scheduling when server needs GPUs and client lacks them, overlap/het_group adjustments, sandbox tweaks, and accounting for n_servers in het group totals.

Sequence Diagram

sequenceDiagram
    participant CLI as GRPO CLI
    participant Scheduler as Task Scheduler
    participant GPU as GPU Allocator
    participant Server as Judge Server
    participant Trainer as Training Task

    CLI->>CLI: Parse server CLI params
    CLI->>Scheduler: add_task(server_config, n_servers=N)

    Scheduler->>GPU: Query client_num_gpus
    GPU-->>Scheduler: GPU availability

    alt Server needs GPUs & Client has none
        Scheduler->>Scheduler: schedule server tasks FIRST
        loop N servers
            Scheduler->>Server: create server executor (port/model/gpus)
            Server->>Server: allocate GPUs
        end
        Scheduler->>Trainer: schedule training after servers
    else Client has GPUs
        Scheduler->>Trainer: schedule training task
        loop N servers
            Scheduler->>Server: add server tasks (post-main)
        end
    end

    Scheduler->>Scheduler: set log_prefix, het_groups, overlap flags
    Scheduler-->>CLI: return configured tasks

Loading

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~60 minutes

🚥 Pre-merge checks | ✅ 2 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 66.67% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (2 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title accurately reflects the main change: adding support for heterogeneous servers in nemo-rl jobs, which aligns with the new server configuration parameters and logic introduced in GRPO workflow and task scheduling.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

• 📝 Generate docstrings (stacked PR)
• 📝 Generate docstrings (commit on current branch)

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Post copyable unit tests in a comment
• Commit unit tests in branch smahdavi/het-job-judg

Tip

Try Coding Plans. Let us write the prompt for your AI agent so you can ship faster (with fewer bugs).
Share your feedback on Discord.

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[coderabbitai]

Actionable comments posted: 3

🤖 Fix all issues with AI agents

In `@nemo_skills/pipeline/nemo_rl/grpo.py`:
- Around line 405-420: When server_type is provided but you intend to host the
model (server_address is None), ensure server_model is required and non-empty to
avoid passing model_path=None into get_server_command; add a validation check
(e.g., assert or raise ValueError) before building server_config that
server_model is not None/empty, referencing the variables server_type,
server_address, and server_model and the block that constructs server_config so
the code fails fast with a clear message instead of producing "None" in the
command.

In `@nemo_skills/pipeline/utils/exp.py`:
- Around line 530-534: The code should fail-fast on a missing num_gpus key:
replace int(server_config.get("num_gpus", 0)) with
int(server_config["num_gpus"]) in the server_needs_gpus calculation (keep the
existing server_config is not None check), so server_needs_gpus and
server_goes_first reflect the required presence of server_config["num_gpus"]
used later by get_server_command; this ensures a clear KeyError rather than
silently using 0.
- Around line 536-572: The add_server_tasks function currently mutates the
shared server_config by calling server_config.pop("container", ...); instead,
make a shallow copy of server_config at the start of add_server_tasks (e.g.,
local_server_cfg = server_config.copy()), read the container with
local_server_cfg.get("container") and if you need to remove it from what you
pass onward, delete it from the copy (del local_server_cfg["container"]) before
calling get_server_command(**local_server_cfg, cluster_config=cluster_config)
and passing local_server_cfg to other helpers (or pass container separately) so
the original server_config remains unchanged and unexpected keys are not
forwarded to get_server_command.

🧹 Nitpick comments (3)

nemo_skills/pipeline/utils/exp.py (2)

594-612: client_num_gpus is reassigned with different semantics — consider a distinct variable name.

Line 531 sets client_num_gpus = num_gpus or 0 to decide server ordering/overlap. Line 594 reassigns it to 0 when server_config is not None and num_nodes == 1, changing the meaning from "does the client need GPUs at all" to "how many GPUs to allocate for the main SLURM srun." This shadowing makes the control flow hard to follow — a reader (or future editor) must track which assignment is live at each usage site.

Consider using a distinct name like main_task_gpus for the line 594 assignment.

---

436-436: Consider validating n_servers >= 1 when server_config is provided.

If n_servers=0 is passed with a non-None server_config, no server tasks are created, yet server_config is still consumed (e.g., popping "container", computing server_goes_first). This would silently produce a misconfigured job. A simple guard early in the function would prevent this.

nemo_skills/pipeline/nemo_rl/grpo.py (1)

409-411: Consider using raise ValueError instead of assert for argument validation.

assert statements can be disabled with python -O. While unlikely in a CLI context, using explicit raise ValueError(...) or typer.BadParameter(...) is more robust and consistent with the existing validation patterns in this file (e.g., line 399).

Also applies to: 428-429

[coderabbitai]

⚠️ Potential issue | 🟠 Major

Missing validation: server_model should be required when server_type is specified.

When server_type is provided but server_model is omitted (defaults to None), model_path=None is passed into get_server_command, producing a command string containing the literal string "None". This would fail at runtime with a confusing error.

Proposed fix

     server_config = None
     if server_type is not None:
+        if server_model is None:
+            raise ValueError("server_model is required when server_type is specified")
         get_random_port = should_get_random_port(server_gpus, exclusive)

🤖 Prompt for AI Agents

In `@nemo_skills/pipeline/nemo_rl/grpo.py` around lines 405 - 420, When
server_type is provided but you intend to host the model (server_address is
None), ensure server_model is required and non-empty to avoid passing
model_path=None into get_server_command; add a validation check (e.g., assert or
raise ValueError) before building server_config that server_model is not
None/empty, referencing the variables server_type, server_address, and
server_model and the block that constructs server_config so the code fails fast
with a clear message instead of producing "None" in the command.

[Kipok]

@gwarmstrong do we need to update the declarative code path to reflect these changes?

[gwarmstrong]

It will need to be updated when we want to use the feature on the declarative path, but at the moment I'm not sure there is value to adding it to the declarative path purely for parity sake.

Is there any way to ensure it is covered by some test case (gpu or slurm probably?) that way when we convert to declarative, we can make sure the functionality isn't dropped?

[greptile-apps]

_{2 files reviewed, 1 comment}

_{Edit Code Review Agent Settings | Greptile}

[greptile-apps]

client_num_gpus is calculated here inside the loop, but was already defined at line 531. This shadows the outer variable and is calculated inside the wrong scope (should be outside the for loop at lines 588-617).

Suggested change

	client_num_gpus = num_gpus if (server_config is None or num_nodes > 1) else 0
	client_num_gpus = num_gpus if (server_config is None or num_nodes > 1) else 0

Move this line before line 588 (before the for cur_idx, (cur_cmd... loop starts).

[coderabbitai]

Actionable comments posted: 3

♻️ Duplicate comments (1)

nemo_skills/pipeline/utils/exp.py (1)

541-541: ⚠️ Potential issue | 🟠 Major

Use direct key access for required server_config["num_gpus"].

Line 541 currently masks missing num_gpus and can route scheduling logic incorrectly before failing later.

Proposed fix

-    server_needs_gpus = server_config is not None and int(server_config.get("num_gpus", 0)) > 0
+    server_needs_gpus = server_config is not None and int(server_config["num_gpus"]) > 0

As per coding guidelines, "Don't use .get() for accessing dictionary keys if the code expects them to be present; use direct access data[key_name] to fail with a clear error instead of silently corrupting data".

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@nemo_skills/pipeline/utils/exp.py` at line 541, The code currently uses
server_config.get("num_gpus", 0) which silently treats missing keys as zero;
change the expression that computes server_needs_gpus to use direct key access
so missing data fails fast: keep the None check on server_config, then use
int(server_config["num_gpus"]) > 0 when computing server_needs_gpus (variable
name: server_needs_gpus, object: server_config) so a missing "num_gpus" raises a
KeyError instead of masking the problem.

🧹 Nitpick comments (1)

nemo_skills/pipeline/utils/exp.py (1)

539-546: Add a SLURM regression test for server ordering with n_servers > 1.

Given the new ordering and heterogeneous-group logic, add coverage for:

• server-first path (client_num_gpus=0, server GPUs > 0),
• server-last path (client GPUs > 0),
• optional sandbox enabled.

This should prevent future regressions in het-group ordering and resource assignment.

Based on learnings, "When enabling new modality or adding complicated evaluation/metrics logic in benchmarks, consider adding the dataset into slurm tests for comprehensive evaluation".

Also applies to: 589-698

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@nemo_skills/pipeline/utils/exp.py` around lines 539 - 546, Add a SLURM
regression test that validates het-group ordering when n_servers > 1 by
exercising both server-first and server-last paths and toggling sandbox: create
tests that set server_config with num_gpus>0 and client num_gpus=0 (to assert
server_goes_first true), set client num_gpus>0 with server_config num_gpus>0 (to
assert server_goes_first false), and run each with sandbox enabled/disabled;
verify the resulting job heterogenous-group ordering and resource allocation
match expectations (use the same code paths that compute server_needs_gpus,
client_num_gpus, and server_goes_first) and ensure the test fails on
regressions.

🤖 Prompt for all review comments with AI agents

Verify each finding against the current code and only fix it if needed.

Inline comments:
In `@nemo_skills/pipeline/utils/exp.py`:
- Around line 589-592: The server-first insertion (server_goes_first ->
add_server_tasks) mutates executors so later code that computes sandbox
num_nodes from executors[0] can pick the server's node count; compute the
sandbox node count from the intended main-task executor before potentially
calling add_server_tasks (or locate the first non-server/main executor instead
of using executors[0]) so num_nodes is derived from the main task; update the
logic around server_goes_first, add_server_tasks, and the sandbox num_nodes
calculation to use that precomputed/main-task executor reference (referencing
symbols: server_goes_first, add_server_tasks, executors, and num_nodes).
- Line 447: The parameter n_servers is accepted but can be zero/negative or
ignored when server_config is None (e.g., code using range(n_servers)), so add
input validation at the start of the function that uses n_servers: if n_servers
is not a positive int raise a ValueError; additionally, if n_servers > 0 ensure
server_config is not None and raise a ValueError if it is, so user-specified
servers are never silently ignored; update any code paths that iterate with
range(n_servers) to rely on this validation.
- Around line 542-546: Compute the effective client GPU count once (e.g.,
rename/assign client_num_gpus -> effective_client_gpus using the same
num_gpus/server_needs_gpus inputs) and reuse that variable everywhere instead of
recomputing; update the server ordering boolean (server_goes_first) to use
effective_client_gpus and replace the later GPU-allocation logic that currently
recomputes client GPUs (the block around the "main executor GPU allocation"
code) to reference effective_client_gpus so ordering, overlap, and allocation
decisions are consistent across client_num_gpus, server_goes_first,
server_needs_gpus and the main executor allocation logic.

---

Duplicate comments:
In `@nemo_skills/pipeline/utils/exp.py`:
- Line 541: The code currently uses server_config.get("num_gpus", 0) which
silently treats missing keys as zero; change the expression that computes
server_needs_gpus to use direct key access so missing data fails fast: keep the
None check on server_config, then use int(server_config["num_gpus"]) > 0 when
computing server_needs_gpus (variable name: server_needs_gpus, object:
server_config) so a missing "num_gpus" raises a KeyError instead of masking the
problem.

---

Nitpick comments:
In `@nemo_skills/pipeline/utils/exp.py`:
- Around line 539-546: Add a SLURM regression test that validates het-group
ordering when n_servers > 1 by exercising both server-first and server-last
paths and toggling sandbox: create tests that set server_config with num_gpus>0
and client num_gpus=0 (to assert server_goes_first true), set client num_gpus>0
with server_config num_gpus>0 (to assert server_goes_first false), and run each
with sandbox enabled/disabled; verify the resulting job heterogenous-group
ordering and resource allocation match expectations (use the same code paths
that compute server_needs_gpus, client_num_gpus, and server_goes_first) and
ensure the test fails on regressions.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Run ID: 7f322e1c-a31f-4ad6-aea2-6254af182f0e

📥 Commits

Reviewing files that changed from the base of the PR and between 4486321 and 3ba53cd.

📒 Files selected for processing (1)

• nemo_skills/pipeline/utils/exp.py

[coderabbitai]

⚠️ Potential issue | 🟠 Major

Validate n_servers inputs to avoid silent no-op behavior.

range(n_servers) silently skips server creation for 0/negative values, and n_servers is currently accepted even when server_config is None.

Proposed fix

 def add_task(
@@
     ray_template: str | None = None,
 ):
@@
+    if n_servers < 1:
+        raise ValueError("n_servers must be >= 1")
+    if server_config is None and n_servers != 1:
+        raise ValueError("n_servers is only supported when server_config is provided")

As per coding guidelines, "Avoid cases where user-passed parameters are unused; code should fail if user specifies an unsupported argument or if a required argument is missing."

Also applies to: 557-557

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@nemo_skills/pipeline/utils/exp.py` at line 447, The parameter n_servers is
accepted but can be zero/negative or ignored when server_config is None (e.g.,
code using range(n_servers)), so add input validation at the start of the
function that uses n_servers: if n_servers is not a positive int raise a
ValueError; additionally, if n_servers > 0 ensure server_config is not None and
raise a ValueError if it is, so user-specified servers are never silently
ignored; update any code paths that iterate with range(n_servers) to rely on
this validation.

[coderabbitai]

⚠️ Potential issue | 🟠 Major

Compute effective client GPU count once and reuse it.

server_goes_first/server overlap are decided from Line 542, but main executor GPU allocation is recomputed at Line 609 with different logic. This can make ordering and overlap inconsistent with actual resource requests.

Proposed fix

-    client_num_gpus = num_gpus or 0
+    client_num_gpus = (num_gpus or 0) if (server_config is None or num_nodes > 1) else 0
     # For ray heterogenous jobs, nemo-run assumes the first het group is the main task
     # So we send the server last if the job needs gpus
     server_goes_first = server_needs_gpus and not client_num_gpus
...
-                client_num_gpus = num_gpus if (server_config is None or num_nodes > 1) else 0
                 executors.append(
                     get_executor(
...
                         gpus_per_node=client_num_gpus,
...
                         overlap=(not client_num_gpus),  # Only when the main task does not have gpus

As per coding guidelines, "Keep code simple and elegant; reuse/extend existing functionality when possible, minimize conditional checks..."

Also applies to: 576-576, 609-629

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@nemo_skills/pipeline/utils/exp.py` around lines 542 - 546, Compute the
effective client GPU count once (e.g., rename/assign client_num_gpus ->
effective_client_gpus using the same num_gpus/server_needs_gpus inputs) and
reuse that variable everywhere instead of recomputing; update the server
ordering boolean (server_goes_first) to use effective_client_gpus and replace
the later GPU-allocation logic that currently recomputes client GPUs (the block
around the "main executor GPU allocation" code) to reference
effective_client_gpus so ordering, overlap, and allocation decisions are
consistent across client_num_gpus, server_goes_first, server_needs_gpus and the
main executor allocation logic.

[coderabbitai]

⚠️ Potential issue | 🟠 Major

Server-first ordering breaks sandbox node selection assumptions.

When Line 590 runs, executors[0] becomes a server executor, but Line 659 still derives sandbox num_nodes from executors[0]. That can size sandbox by server nodes instead of main-task nodes.

Proposed fix

-                num_nodes=executors[0].nodes if cluster_config["executor"] == "slurm" else 1,
+                num_nodes=num_nodes if cluster_config["executor"] == "slurm" else 1,

Also applies to: 659-659

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@nemo_skills/pipeline/utils/exp.py` around lines 589 - 592, The server-first
insertion (server_goes_first -> add_server_tasks) mutates executors so later
code that computes sandbox num_nodes from executors[0] can pick the server's
node count; compute the sandbox node count from the intended main-task executor
before potentially calling add_server_tasks (or locate the first non-server/main
executor instead of using executors[0]) so num_nodes is derived from the main
task; update the logic around server_goes_first, add_server_tasks, and the
sandbox num_nodes calculation to use that precomputed/main-task executor
reference (referencing symbols: server_goes_first, add_server_tasks, executors,
and num_nodes).