feat(slurm): support external vLLM services in one allocation

[yfw]

What does this PR do ?

Summary

Adds an opt-in framework for running external Gym vLLM service pools alongside NeMo RL in one Slurm heterogeneous job.

This provides gang scheduling while keeping service nodes isolated from the NeMo RL Ray cluster:

• Hetgroup 0 runs NeMo RL and the pool load balancers.
• Hetgroup 1 runs independent Ray/vLLM replicas.
• Training starts only after all required backends and load balancers are healthy.

Motivation

External judges previously required separate Slurm jobs or had to share the training Ray cluster. Separate jobs are not gang scheduled and may become available at different times, while shared clusters risk resource contention and accidental node reuse.

This change allocates training and serving resources together while preserving separate Ray clusters.

Implementation

• Adds a generic registration interface for arbitrary named vLLM pools.
• Validates node counts, TP topology, ports, paths, URL placeholders, containers, and required Python packages before startup.
• Starts each replica in an isolated private Ray cluster.
• Supports per-pool environment variables and vLLM arguments.
• Injects resolved load-balancer URLs into the training command.
• Provides health-aware, affinity-aware, least-loaded routing and per-request retry across live backends.
• Monitors required processes and terminates the heterogeneous job if a required service fails.
• Forwards shutdown signals and drains in-flight load-balancer requests.
• Documents the interface and integration requirements.

Model-specific configuration remains in the consuming recipe or launcher; the framework only manages allocation, lifecycle, routing, and endpoint injection.

Compatibility

The feature is opt-in and does not modify ray.sub. Existing launch workflows are unchanged unless they explicitly use tools/external_gym_vllm/run_in_allocation.sh.

External replicas currently consume whole nodes, so each pool’s tensor-parallel size must be divisible by the configured GPUs per node.

Issues

List issues that this PR closes (syntax):

Usage

• You can potentially add a usage example below

# Add a code snippet demonstrating how to use this

Before your PR is "Ready for review"

Pre checks:

• Make sure you read and followed Contributor guidelines
• Did you write any new necessary tests?
• Did you run the unit tests and functional tests locally? Visit our Testing Guide for how to run tests
• Did you add or update any necessary documentation? Visit our Document Development Guide for how to write, build and test the docs.

Additional Information

• ...

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[yfw]

/ok to test e95d138

[yfw]

/ok to test 91c7b7d

[yfw]

/ok to test c96ef34

[terrykong]

offline discussion was to generalize w/ hetgroups to re-use gang scheduling (will also help with head node on a separate node cc @macandro96 )

[yfw]

/ok to test 481603e

[yfw]

/ok to test 7d965a1

[yfw]

/ok to test b08fa5b

[yfw]

/ok to test 690181a

[yfw]

/ok to test 830ca30

[terrykong]

Reviewed by a team of specialised agents (RL/guidelines, bug-finder, tests, design, Gym integration, devil's advocate), with every finding adversarially challenged before posting. Thanks for this — the hetgroup rework paid off. A few things worth calling out as done well before the list:

• The PR touches zero pre-existing runtime files. ray.sub and docs/cluster.md are byte-identical to main; the only non-new files are a docs paragraph and two pyrefly.toml lines. That makes it easy to reason about and trivial to revert.
• vllm_pool_lb.py is well-factored. Backend / BackendPool / LoadBalancer are three clean responsibilities with constructor-injected dependencies, and pick() / _extract_affinity_key() are pure — which is why the tests can assert real routing behaviour without a Slurm allocation. We checked whether anything wanted extracting and concluded no: the state doesn't partition.
• The up-front validation block fails before a single service starts, with specific messages. On an 80-node allocation that's the difference between failing in 2 seconds and burning an hour.
• The shell is properly tested — four run_in_allocation.sh tests and five pool_config.sh tests run the real scripts as subprocesses and assert their [FATAL] output, plus the registry and watchdog tests.
• The generalization is data-driven, not conditional-driven. Pool differences ride on POOL_ENV_VARS / POOL_VLLM_ARGS consumed by one uniform loop — no per-pool-type branching anywhere, and no eval.
• The previous round's tautological hetgroup test is properly fixed. Verified by mutation: swapping every --het-group=0↔=1 in the source now makes test_launcher_routes_generic_pools_to_explicit_hetgroups fail. That's the right way to write it.

Five comments below, scoped to the launcher and config surface — the places where a mistake costs a whole allocation rather than a retry. None of them block; the two I'd most want addressed are the served-model-name default and the exported GPUS_PER_NODE, since both fail late and neither is discoverable from the README as written.

@terrykong — on your hetgroup comment: the gang-scheduling half landed cleanly; every srun is pinned to its component and Slurm does the co-scheduling. The head-node half got answered from a different direction while this PR was in flight: #3388 landed DEDICATED_RAY_HEAD on main (ray.sub:381, :703-705), and it passes straight through this launcher — ray.sub is started via env … bash "${RAY_SUB}", so setting it in the sbatch environment needs no launcher change. It composes better than by accident: the launcher already pins every load balancer to ray_nodes[0], which under DEDICATED_RAY_HEAD=1 is precisely the GPU-free node. Nothing in the new README mentions this, and the sizing rule (hetgroup 0 = N+1, cluster.num_nodes = N, which ray.sub warns it "cannot validate for you") is exactly what users get wrong — worth a row. What still isn't expressible is a third component for a genuinely GPU-free head: :36 hard-rejects SLURM_HET_SIZE != 2, and docs/cluster.md notes the head's GPUs today are "idled, not freed". Relaxing that to >= 2 with the group indices as env vars would make it a config change later rather than a script edit. cc @macandro96.

On the description: "backend failover" is doing slightly more work than the code currently supports. The load balancer really does retry across healthy backends, but check_service_steps tears the whole job down within ~5s of any replica step exiting, so the "a backend dies and the pool keeps serving" case can't be reached today — per-request failover off a live-but-erroring backend still works. The README's step 7 ("stops training if any required external service exits") states the opposite policy, so the two just need to agree on which one you meant.

One structural note: the framework currently ships with no in-tree caller — nothing outside tools/external_gym_vllm/ invokes register_external_vllm_pool, so the contract (including the served-model-name coupling in the comment below) isn't exercised by anything that runs in CI. Landing the ultra_launch.sh wiring or an examples/ snippet in a follow-up would give it one real consumer.

One non-blocking ask: if you have a smoke run handy, the wall-clock from launch to [INFO] External vLLM pools are healthy would tell users whether POOL_STARTUP_TIMEOUT=3600 has headroom or is uncomfortably tight. It's the one shipped constant nothing else validates.

Generated by Claude Code

[yfw]

Re POOL_STARTUP_TIMEOUT=3600, looking at some recent runs with 8 TP=8 GenRM replicas and 4 TP=4 NL2Bash replicas, it took ~16 mins on average for external vllms to become healthy, so startup time of 3600 leaves substantial headroom.

[yfw]

/ok to test 540db66

[terrykong]

All five findings verified as addressed in 540db66 — I checked the code rather than the reply:

• LOG_DIR — both action items: the SLURM_RESTART_COUNT suffix now mirrors ray.sub, and head_ip_* is added to the startup cleanup.
• Replica exit — wait "${VLLM_PID}" now always logs and exits non-zero, including the clean-exit case that previously shrank the pool silently.
• POOL_SERVED_MODEL_NAME — the coupling to Gym's model field is documented, and the README example now passes --served-model-name explicitly, which is the part that stops people hitting it.
• GPUS_PER_NODE — documented as set-before-registration and exported to ray.sub; the example sets it explicitly, which also closes the registration-time ordering hazard.
• Docs — /lustre literals replaced with EXTERNAL_VLLM_SHARED_ROOT, NUM_EXTERNAL_SERVICE_NODES documented plus a runtime warning when it's skipped, and the reasoning-parser ↔ uses_reasoning_parser contract spelled out.

Bonus: documenting DEDICATED_RAY_HEAD and the hetgroup-0 N+1 sizing rule answers @terrykong's original ask directly — that was the thing most likely to bite the first user.

Nice touch adding assertions for each fix rather than just making the change.

LGTM.

Generated by Claude Code