infra: K8s GPU cluster setup with KAI scheduler, KubeRay, and JobSet

[terrykong]

Summary

Add K8s infrastructure for deploying nemo-rl GPU workloads:

• nvkind cluster setup: scripts for local GPU-enabled K8s cluster with NVIDIA GPU support
• Helmfile: KAI scheduler (gang scheduling + fairshare), KubeRay operator, JobSet controller, nvidia-device-plugin (kind) / GPU operator (prod)
• Workload examples: monolithic RayJob, disaggregated RL+Gym via KubeRay RayClusters, disaggregated RL+Gym via JobSet
• Disaggregated architecture: ConfigMap endpoint registry for vLLM/Gym URL exchange, peer-watcher sidecar for failure cascading, standalone Gym server
• Code changes: K8sEndpointRegistry for service discovery, remote_gym_url in NemoGymConfig, standalone Gym server moved to nemo_rl/distributed/

Architecture

Two deployment patterns for disaggregated RL+Gym:

1. KubeRay RayClusters — two independent RayClusters with peer-watcher sidecars for failure cascading
2. JobSet — single JobSet with native failurePolicy and init containers (no peer-watcher needed)

Both use a ConfigMap endpoint registry for dynamic vLLM URL exchange.

Test plan

• nvkind cluster creation with 2 GPUs
• Helmfile sync (KAI, KubeRay, JobSet, device plugin)
• KAI queue creation (high-prio/low-prio hierarchy)
• Disagg RayClusters: both clusters reach Ready, peer-watcher failure cascading works both directions
• Disagg JobSet: all 6 pods gang-scheduled, init containers work, failurePolicy tears down on crash, successPolicy triggers on driver exit 0
• GRPO initialization: config loaded, vLLM workers started (CUDA graphs captured), Megatron workers started, endpoint registry created, vLLM URLs published
• dependsOn + KAI deadlock identified and documented (use init containers instead)
• Full E2E training (blocked by Gym standalone server handshake — needs matching container image)

[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.

[terrykong]

/ok to test ab15151

[terrykong]

Review Summary

This PR adds K8s infrastructure for disaggregated RL-Gym deployments, including:

• K8sEndpointRegistry — ConfigMap-backed service discovery between RL and Gym clusters
• standalone_gym_server — standalone NeMo Gym server for independent deployment
• Remote mode in NemoGym — connect to an external Gym HTTP service instead of spawning local subprocesses
• Kind cluster setup scripts, Helm charts, and K8s manifest examples (JobSet + KubeRay approaches)
• Gym submodule bump (23cdeb38 → 01a9765f) for disaggregated mode support

Action Items

⚠️ Merge conflicts: This PR is currently in a CONFLICTING state with main. Please rebase and resolve conflicts before merging.

📝 PR description: The PR body is currently the template placeholder. It would be helpful to fill in a description of the motivation, architecture decisions, and test plan — this helps reviewers and future readers understand the context.

📖 Documentation: Per CONTRIBUTING.md, new key features should include a documentation update. This PR introduces a significant new deployment architecture — consider adding a doc page covering the disaggregated mode, K8s prerequisites, and configuration options (env.remote_gym_url, env.disagg_job_id).

Additional Notes (not tied to specific diff lines)

• Internal image references: All example manifests use nvcr.io/nvidian/nemo-rl:e5a729c-47084432 which is only accessible within NVIDIA networks. If these examples are intended for external contributors too, consider adding a comment noting users should substitute their own image.
• Missing nvcr-secret setup: SETUP.md Quick Start doesn't mention creating the nvcr-secret imagePullSecret that all manifests require. First-time users will hit ImagePullBackOff.
• Shell script hardening: create-cluster.sh uses || true after nvkind cluster create which silently swallows all failures (not just the expected /proc/driver/nvidia patching issue). install-nvkind.sh installs nvkind at @latest without version pinning.

See inline comments for specific code-level findings.

Generated by Claude Code

[terrykong]

k8s_endpoint_registry.py:37

kubernetes is imported at module level but isn't declared in pyproject.toml (not in [project.dependencies] or any optional extra). Users importing K8sEndpointRegistry will get an ImportError unless they manually pip install kubernetes.

The disagg-jobset.yaml works around this with a runtime uv pip install kubernetes (line 212), but it would be cleaner to declare it as an optional dependency — e.g., a k8s extra in pyproject.toml.

[terrykong]

Fixed in 0881933. Added k8s extra to pyproject.toml.

[terrykong]

k8s_endpoint_registry.py:118-126

The read-modify-patch pattern here is susceptible to lost updates under concurrency. If both RL and Gym call set() on the same ConfigMap near-simultaneously, one side may read a stale resourceVersion. The patch_namespaced_config_map call will then return a 409 Conflict, which isn't caught (only 404 is handled in the outer except), so the write silently fails with an unhandled exception.

Consider catching 409 in the outer handler with a retry loop, or using a strategic merge patch that targets only the specific key rather than replacing the full data dict.

[terrykong]

Fixed in 0881933. Added 409 handling on the patch path with retry.

[terrykong]

k8s_endpoint_registry.py:141-144

The recursive self.set(key, value) on a 409 during the fallback create has no depth limit. In a tight race between RL and Gym both hitting this path, this could recurse until stack overflow. A simple iterative retry with a cap (e.g., for attempt in range(3): ...) would be safer.

[terrykong]

Fixed in 0881933. Converted to iterative retry with _max_retries=5. Raises RuntimeError if exhausted.

[terrykong]

nemo_gym.py:51-56

The URL parsing using rsplit(":", 1) works for simple host:port inputs but breaks on URLs with path components. For example, http://gym-service:8080/v1 → host="gym-service", port_str="8080/v1" → int("8080/v1".rstrip("/")) raises ValueError.

Consider using urllib.parse.urlparse for robustness:

from urllib.parse import urlparse
parsed = urlparse(remote_url if "://" in remote_url else f"http://{remote_url}")
host = parsed.hostname or remote_url.rstrip("/")
port = parsed.port or 8080

[terrykong]

Fixed in 0881933. Using urllib.parse.urlparse now.

[terrykong]

nemo_gym.py:63-66

Two things here:

1. Mutation leak: .pop("rollout_max_attempts_to_avoid_lp_nan", 1) mutates the caller's initial_global_config_dict (potentially the live config from run_grpo_nemo_gym.py). Consider using .get() instead of .pop(), or copying the dict first.

2. Missing validation: Unlike the colocated path (line 120), there's no assert >= 1 check here. If a user sets this to 0, run_rollouts() silently produces no results (the while trial < max_attempts loop at line 154 is never entered).

[terrykong]

Fixed in 0881933. Changed .pop() to .get() and added assert >= 1 matching the colocated path.

[terrykong]

run_grpo_nemo_gym.py:231-234

registry.get("gym_head_server") blocks for up to 600s and raises TimeoutError on failure. This exception isn't caught — it will crash the training job after all setup (Ray init, model loading, etc.) is complete, wasting GPU time.

Consider wrapping this in a try/except with a user-friendly error message, e.g.:

try:
    remote_gym_url = registry.get("gym_head_server")
except TimeoutError:
    raise RuntimeError(
        f"Timed out waiting for Gym cluster to register. "
        f"Check that the Gym cluster is running with --job-id={disagg_job_id}"
    )

[terrykong]

Fixed in 0881933. Wrapped in try/except TimeoutError with a message pointing to the Gym cluster.

[terrykong]

standalone_gym_server.py:125-136

_strip_interpolations silently removes any config key whose value contains ${. This is meant to handle unresolvable OmegaConf interpolations from the full GRPO config, but it could silently drop legitimate values — the server would then start with missing settings and fail later in an opaque way.

Consider logging each dropped key so users can diagnose config issues:

if isinstance(v, str) and "${" in v:
    print(f"Dropping unresolvable interpolation: {k}={v}")
    continue

[terrykong]

Fixed in 0881933. Now logs each dropped key with its path and value.

[terrykong]

endpoint-registry-rbac.yaml:13-19

This Role grants delete on all ConfigMaps and all RayClusters in the namespace. A few observations:

• delete on ConfigMaps is only needed by the peer-watcher teardown path. The endpoint registry itself only needs get, create, update, patch.
• delete on RayClusters is only needed by the peer-watcher sidecar on head pods, but this SA is also used by GPU worker pods (disagg-rayclusters.yaml lines 167, 200) which don't need it.

Consider splitting into two ServiceAccounts — one for head pods (full permissions) and one for workers (ConfigMap read/write only). This follows the principle of least privilege.

[terrykong]

Added a NOTE comment acknowledging this is an example and suggesting production hardening. Kept as-is for simplicity.

[terrykong]

disagg-rayclusters.yaml:131-140

The continue on line 136 skips the ConfigMap error-flag check (lines 137-140). If the peer cluster is unhealthy and signals an error via ConfigMap simultaneously, the error flag won't trigger teardown until either the peer recovers or MAX_FAIL is reached.

Moving the ConfigMap check before the continue would catch both conditions:

if state in ("failed", "suspended") or (code != 0 and state == ""):
    fails += 1
    ...
    if fails >= MAX_FAIL:
        teardown(...)
    # Check ConfigMap error even when peer is failing
    if JOB_ID:
        cm = kube(f"/api/v1/...")
        if cm.get("data", {}).get("error", ""):
            teardown(f"Error via ConfigMap: {cm['data']['error']}")
    continue

[terrykong]

Fixed in 0881933. Moved ConfigMap check before the continue in both peer-watcher sidecars.

[terrykong]

nvkind-config-values-dev.yaml:6-7

Nit: /home/terryk/nemo-rl is a developer-specific path checked into the repo. Other contributors using this file will silently get an empty DirectoryOrCreate mount. Consider adding a comment noting this path must be edited per-user, or adding this file to .gitignore and providing a .example template instead.

[terrykong]

Fixed in 0881933. Added TODO comment noting the path must be edited per-user.

[terrykong]

Superseded by #2321 which includes all the infra work from this branch plus nrl-k8s CLI improvements (dev pod, DRA auto-management, per-user cluster names, stale RayJob detection).