feat: disaggregated generation and NemoGym - #2321
Draft
terrykong wants to merge 84 commits into
Draft
Conversation
This was referenced Apr 23, 2026
terrykong
force-pushed
the
hemil/k8s-infra-cp
branch
from
April 23, 2026 09:07
5b8fbbd to
65f4fcd
Compare
terrykong
commented
Apr 23, 2026
terrykong
left a comment
terrykong
left a comment
Collaborator
Author
There was a problem hiding this comment.
Review Summary
This is a substantial PR (~18.8k insertions, 88 files) adding GB300 infra support, the nrl-k8s CLI tool, and several new nemo_rl/ runtime modules for disaggregated deployment. The core changes look solid overall. A few items to address:
General findings (not tied to specific diff lines)
Copyright headers missing in tools/nrl_k8s/ — Most files under tools/nrl_k8s/src/nrl_k8s/ (cli.py, k8s.py, config.py, schema.py, orchestrate.py, submit.py, manifest.py, rayjob.py, dev.py, inspect.py, workdir.py, __init__.py, _logging.py, _retry.py) have no NVIDIA copyright header. The submitters/exec_.py and submitters/portforward.py have abbreviated headers (missing the "Unless required by applicable law..." clause). The copyright skill requires the full Apache 2.0 header on all new Python files (except tests). All new nemo_rl/*.py files in this PR have the correct full header, so this is just an oversight in tools/nrl_k8s/.
Missing unit tests for new modules — Three new nemo_rl/ modules have no unit tests:
k8s_endpoint_registry.py(260 lines)standalone_gym_server.py(197 lines)generation_control_server.py(253 lines)
Only remote_generation.py has coverage via test_remote_generation_http.py. The registry and control server have HTTP/K8s-client surfaces that could be unit-tested with mocks (same pattern as the existing test). standalone_gym_server.py is mostly a thin entry point, so lower priority.
Generated by Claude Code
Local K8s GPU dev environment using nvkind (NVIDIA's kind wrapper): - nvkind cluster setup scripts (install-nvkind.sh, create-cluster.sh) - Custom config template with extraMounts for dev code mounting - Helmfile with kind/prod environments (device plugin vs GPU operator) - KAI scheduler for gang scheduling, KubeRay for RayCluster management - Example manifests: gang-scheduled pods, RayClusters, SFT RayJobs - SETUP.md with prerequisites, quick start, and architecture docs Tested: SFT RayJob (train/loss 4.06 < 5.9), KAI all-or-nothing gang scheduling, two simultaneous 1-GPU SFT jobs. Signed-off-by: Terry Kong <terryk@nvidia.com>
Signed-off-by: Terry Kong <terryk@nvidia.com>
…overy (infra part) Signed-off-by: Terry Kong <terryk@nvidia.com>
When RL and Gym run on separate RayClusters, either cluster failing or being deleted triggers teardown of both clusters to release resources. - peer-watcher.py: pure Python sidecar (no deps beyond stdlib), deployed as a ConfigMap volume mount on each head pod - Monitors peer RayCluster status via K8s API (polls every 10s) - Tears down after MAX_PEER_FAILURES (default 3) consecutive failures - Also monitors ConfigMap "error" key for application-level error signaling - Handles transient K8s API errors as failures (not false-healthy) - Added signal_error() to K8sEndpointRegistry - Updated disagg manifests with peer-watcher sidecar containers - Updated RBAC with "delete" verb for rayclusters Tested: deleting either cluster triggers teardown of both within ~10s. Signed-off-by: Terry Kong <terryk@nvidia.com>
…share configs - Kyverno policy: RayCluster/RayJob must have kai.scheduler/queue label. Validates at CRD level (not pod) since KubeRay operator creates pods. Optional Policy 2 for user→queue access control via ConfigMap. - kube-prometheus-stack: Prometheus + Grafana for fairshare monitoring. Pre-built Grafana dashboard showing GPU allocation vs fair share, preemption events, and scheduling latency per queue. - ServiceMonitors for KAI scheduler, binder, and queue-controller. - Example queue configs: - kai-queue.yaml: 2-GPU kind cluster (2 teams, equal quotas) - kai-queue-prod.yaml: 256-GPU prod (3 departments, 6 teams) - preemptMinRuntime: 4h (protect long training runs from priority preemption) - reclaimMinRuntime: 15m (fast fairness reclaim of over-quota resources) - SETUP.md: fairshare docs, preempt vs reclaim explanation, Grafana access. Tested: Kyverno rejects RayCluster without queue label, accepts with. Team A 2-GPU job reclaimed when Team B submitted to its guaranteed quota. Signed-off-by: Terry Kong <terryk@nvidia.com>
- Upgrade KAI scheduler v0.13.4 → v0.14.0 (adds Ray topology-aware
scheduling, segment-size annotation support for PyTorchJob)
- Update chart URL from NVIDIA/KAI-Scheduler to kai-scheduler/KAI-Scheduler
- Fix Grafana dashboard metric names (add kai_ prefix to match actual
Prometheus metric names). Verified: Grafana queries return live data.
- New: extensions/k8s_cli/ — standalone Python CLI (pip installable):
- nrl-k8s fairshare — show queue config (quota, limit, weight, priority)
- nrl-k8s occupancy — show GPU allocation per node and per queue
- nrl-k8s submit — submit gang-scheduled RayJob with optional
--segment-size for topology-aware scheduling
- 6 unit tests (mocked K8s API), all passing
- Add TODO for NVL72 topology testing with links to relevant PRs/issues
Tested: KAI v0.14.0 gang scheduling works, CLI commands verified against
live cluster, Grafana dashboard loads and queries return data.
Signed-off-by: Terry Kong <terryk@nvidia.com>
- Merge disagg_rl_raycluster.yaml + disagg_gym_raycluster.yaml into single disagg-rayclusters.yaml (always deployed together) - Inline peer-watcher Python script directly in sidecar container args (eliminates ConfigMap setup step, each deployment is self-contained) - Remove 7 redundant workload YAMLs (sft_rayjob, kai_scheduled_*, raycluster-blocker, standalone peer-watcher.py) - Update SETUP.md: simplified quick start, updated architecture tree, removed ConfigMap peer-watcher setup step 15 files → 8 files in examples/. Infrastructure configs unchanged. Tested: inlined peer-watcher works — deleting either cluster triggers teardown of both within 10s. Signed-off-by: Terry Kong <terryk@nvidia.com>
- Remove extensions/k8s_cli/ (not needed for now) - Rename queues: org → root-org, priority-team → high-prio, community → low-prio - Simplify Kyverno policy comments - Update SETUP.md to remove CLI section Signed-off-by: Terry Kong <terryk@nvidia.com>
Not needed right now — queue enforcement can be added back later if required. Signed-off-by: Terry Kong <terryk@nvidia.com>
Signed-off-by: Terry Kong <terryk@nvidia.com>
Signed-off-by: Terry Kong <terryk@nvidia.com>
…a part) Signed-off-by: Terry Kong <terryk@nvidia.com>
Three workload deployment patterns: 1. rayjob-monolithic.yaml — single-cluster RayJob (1 GPU, KubeRay) 2. disagg-rayclusters.yaml — two KubeRay RayClusters + peer-watcher 3. disagg-jobset.yaml — single JobSet with native failure/startup policies Also adds JobSet controller (v0.11.1) to the helmfile. Signed-off-by: Terry Kong <terryk@nvidia.com>
Shows the full Ray cluster pattern for Gym: separate head and worker pods within the JobSet, with dependsOn ordering and DNS discovery. Signed-off-by: Terry Kong <terryk@nvidia.com>
KAI gang-schedules all pods in a JobSet together (one PodGroup with minMember=total pods). This deadlocks with dependsOn: KAI waits for all pods to exist, but JobSet won't create dependent pods until the head is Ready. Fix: drop dependsOn, use init containers that poll ray health-check (same pattern KubeRay uses). Tested: all 6 pods schedule, init containers wait for heads, driver submits a Ray job successfully, successPolicy triggers on driver exit 0, failurePolicy tears down everything on gym-head crash. Signed-off-by: Terry Kong <terryk@nvidia.com>
Signed-off-by: Terry Kong <terryk@nvidia.com>
- Add git safe.directory for Gym submodule (uv build fails otherwise) - Add uv pip install kubernetes (needed by endpoint registry) - Increase readiness probe failureThreshold (uv install takes time) Signed-off-by: Terry Kong <terryk@nvidia.com>
Replace --working-dir and --runtime-env-json with a simple cd wrapper. Both --working-dir and runtime_env.working_dir cause Ray to zip and upload the entire directory to GCS, which is extremely slow for large repos (1GB+). Since the code is already on all nodes via hostPath, wrapping the entrypoint with cd avoids the upload entirely. Before: ray job submit --working-dir /workspace/nemo-rl -- python ... → scans entire tree, uploads 1GB+ to GCS, takes minutes After: ray job submit -- bash -c "cd /workspace/nemo-rl && python ..." → instant submission, no upload Signed-off-by: Terry Kong <terryk@nvidia.com>
…pefail - Use timestamp-based submission_id to avoid GCS collision across redeploys - Disable wandb (not configured in kind dev cluster) - Add set -eo pipefail for proper exit code propagation through tee - Persist driver logs to hostPath for post-mortem debugging Tested: GRPO training loads config, connects to Ray cluster, loads datasets, initializes compute cluster. Fails with "Not enough GPUs" (expected — kind cluster has 2 GPUs, config expects 8). Signed-off-by: Terry Kong <terryk@nvidia.com>
terrykong
force-pushed
the
hemil/k8s-infra-cp
branch
from
April 26, 2026 20:15
25d9bc2 to
46694b3
Compare
…ine env vars Signed-off-by: Terry Kong <terryk@nvidia.com>
terrykong
force-pushed
the
hemil/k8s-infra-cp
branch
from
April 26, 2026 20:48
46694b3 to
f0add82
Compare
…ble heredoc Signed-off-by: Terry Kong <terryk@nvidia.com>
terrykong
force-pushed
the
hemil/k8s-infra-cp
branch
from
April 26, 2026 21:32
f0add82 to
cb393cb
Compare
Signed-off-by: Terry Kong <terryk@nvidia.com>
terrykong
force-pushed
the
hemil/k8s-infra-cp
branch
from
April 26, 2026 22:22
cb393cb to
92a97e7
Compare
Signed-off-by: Terry Kong <terryk@nvidia.com>
The ray-edit ClusterRole only covered ray.io CRDs. Running nrl-k8s from the dev pod also needs create/delete on computedomains (resource.nvidia.com) and resourceclaimtemplates (resource.k8s.io). Signed-off-by: Terry Kong <terryk@nvidia.com>
--replace on the exec path (runMode=batch) was a no-op — the old process kept running and the new one queued behind it. Add ExecSubmitter.stop_all_running() which scans pidfiles on the head pod and SIGTERMs any live processes, then wire it into submit_training(). Signed-off-by: Terry Kong <terryk@nvidia.com>
The entrypoint spawns a process tree (bash → python, tee) that shares a PGID. Killing just the top-level PID leaves children running under init. Use kill -TERM -$pgid to take down the entire group. Signed-off-by: Terry Kong <terryk@nvidia.com>
After SIGTERM-ing the process group, wait up to 10s for all processes to actually die. This gives Ray time to notice the driver disconnected and GC actors (e.g. vLLM engines) on workers, freeing GPU memory before the new run starts. Signed-off-by: Terry Kong <terryk@nvidia.com>
Killing the driver process group alone doesn't reliably clean up Ray actors on workers (vLLM engines etc.) — Ray's GC depends on heartbeat timeouts. Now both exec and port-forward paths use the dashboard API to stop jobs, which properly tears down actors. The exec path still kills the driver process group as well for head-pod cleanup. Signed-off-by: Terry Kong <terryk@nvidia.com>
- Remove unused 'import inspect as ins' in cli.py _tail_daemon. - Add PATH export before krew plugin installs in get-kubectl.sh so krew is reachable on a fresh machine. Signed-off-by: Terry Kong <terryk@nvidia.com>
terrykong
force-pushed
the
hemil/k8s-infra-cp
branch
from
April 27, 2026 17:06
5166294 to
31aabed
Compare
Signed-off-by: Terry Kong <terryk@nvidia.com>
Signed-off-by: Terry Kong <terryk@nvidia.com>
…art) Signed-off-by: Terry Kong <terryk@nvidia.com>
…overy (disagg part) Signed-off-by: Terry Kong <terryk@nvidia.com>
…gg part) Signed-off-by: Terry Kong <terryk@nvidia.com>
Signed-off-by: Terry Kong <terryk@nvidia.com>
The standalone_gym_server has been moved to nemo_rl/distributed/ so the Gym submodule change from tk/standalone-server is no longer needed. The standalone server could potentially be upstreamed to Gym in the future, but for now it lives in nemo-rl since it's highly in flux. Signed-off-by: Terry Kong <terryk@nvidia.com>
Simplifies the cross-cluster disagg path to a single transport (OpenAI /v1/completions per DP shard) and adds an idempotent teardown endpoint so the generation cluster can accept a new training cluster's NCCL weight-sync group without a pod restart. Transport simplification (drops ~225 LoC of unused branching): * Remove the "head" tensor /generate endpoint on the control server. * Remove the per-shard /generate endpoint on each vLLM worker. * Remove _send_generate_to_head, _generate_direct_to_shards, _fast_serialize, _fast_deserialize, _shard_generate_urls, _get_session, _aio_session from RemoteGeneration; collapse generate()/generate_async() to the json path. * Drop disagg_mode and disagg_generation_routing config keys. Disaggregated HTTP mode is now implicit when remote_generation_url is set and colocated.enabled=false. Add /reset_collective so the generation server can tear down an orphaned NCCL group (e.g. after a training restart) without losing the running vLLM engines. Plumbed through VllmGeneration -> VllmGenerationWorker( Sync/Async) -> VllmInternalWorkerExtension -> StatelessProcessGroup.destroy, which aborts the NCCL communicator and drops the TCP store. Also adds unit coverage for the surviving json path: response parsing for token_id:NNN logprobs with None substitution, round-robin shard selection across calls, and greedy sampling overrides. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> Signed-off-by: Hemil Desai <hemild@nvidia.com> Signed-off-by: Terry Kong <terryk@nvidia.com>
Training clients already round-robin across DP shards themselves in RemoteGeneration._generate_json_completions, fetched once at startup via the control server's /dp_openai_server_base_urls endpoint. This commit publishes the same shard URL list to NemoGym (via dp_openai_server_base_urls), so the gym's vllm_model plugin — which already supports base_url: Union[str, List[str]] and sticky-maps session -> client round-robin — talks directly to DP shards too. Nothing uses the control server's OpenAI reverse proxy anymore, so delete: * /v1/completions, /v1/chat/completions, /tokenize route handlers * _select_shard, _route_request, _get_session, _pending_counts, _rr_index * RoutingStrategy enum (ROUND_ROBIN / LEAST_PENDING / PREFIX_HASH) * --routing-strategy CLI arg on the standalone launcher * aiohttp/json/hashlib/asyncio/StreamingResponse imports that only the router needed The GenerationControlServer is now a pure control-plane surface: weight sync (init/reset/update_weights_from_collective), lifecycle (prepare_for_generation, finish_generation), refit metadata, logger metrics, and /dp_openai_server_base_urls for shard discovery. Server file shrinks from ~346 → 226 LoC; one fewer network hop on every rollout; one fewer component for NCCL/aiohttp failures to land on. Also restores `import io` in remote_generation.py (dropped too eagerly in the previous commit — still used by prepare_refit_info's torch.save). Tested end-to-end on the 3-cluster disagg IF setup: training registered all 8 DP shard URLs in the endpoint-registry ConfigMap, gym picked them up, first two rollout+training steps completed (44.9s / 53.8s — within the previous router-mediated run's 44.6s mean over 200 steps). Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> Signed-off-by: Hemil Desai <hemild@nvidia.com> Signed-off-by: Terry Kong <terryk@nvidia.com>
… part) Signed-off-by: Terry Kong <terryk@nvidia.com>
…d exceptions - Revert async_utils.py and utils.py (algorithm changes unrelated to infra) - torch.load weights_only=True in generation_control_server.py - Narrow except Exception to requests.RequestException in remote_generation.py - Add pragma: no cover on ray.remote function in remote_generation.py - Fix return type annotation on k8s_endpoint_registry._build_owner_reference Signed-off-by: Terry Kong <terryk@nvidia.com>
Signed-off-by: Terry Kong <terryk@nvidia.com>
Signed-off-by: Terry Kong <terryk@nvidia.com>
terrykong
force-pushed
the
hemil/k8s-infra-cp
branch
from
April 27, 2026 18:09
31aabed to
6199d67
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Builds on #2330. Adds disaggregated (HTTP-based) generation and NemoGym support, allowing RL training and inference/gym to run on separate Ray clusters.
/v1/completions, control-plane calls (weight sync, lifecycle) via GenerationControlServerTest plan
pytest tests/unit/models/generation/test_remote_generation_http.py -x -q