iris: add multi-VM CoreWeave support with JAX coordinator bootstrap

[rjpower]

Closes #3634

• New iris.runtime.jax_init module: Task 0 registers its coordinator address via the existing endpoint registry; tasks 1..N-1 poll for it using ExponentialBackoff, then all call jax.distributed.initialize() with the discovered coordinator. Single-task jobs skip coordination entirely. JAX is imported at call time — no new dependency on JAX.
• Multi-VM CoreweavePlatform: Lifts the num_vms > 1 restriction. Creates N worker Pods per slice with a shared ConfigMap. Proper multi-Pod terminate, partial-failure cleanup, and _list_slices_by_labels grouping.
• Config & docs: Adds h100-16x scale group with num_vms: 2 to examples/coreweave.yaml. Documents multi-VM job submission with coscheduling in docs/coreweave.md.
• Design doc at lib/iris/DESIGN.md.

[claude]

Claude finished @rjpower's task in 6s —— View job

---

PR Review: Multi-VM CoreWeave JAX Coordinator Bootstrap

• Read design document and gather context
• Verify claims against existing codebase
• Post review feedback

---

Overall

Well-written design doc. The core approach — reusing the existing Iris endpoint registry for coordinator address discovery — is clean and avoids unnecessary proto/RPC changes. I verified the key code references against the codebase and they are accurate. The endpoint lifecycle analysis (attempt validation, cascade deletes, coscheduled sibling cascade) is correct.

A few observations:

Line reference accuracy

The doc's line references are accurate (within ±1 line for some). Verified:

• job_info.py:25-74 — JobInfo is at lines 26-74 (off by one on start)
• protocol.py:80-90 — endpoint methods at lines 80-90
• service.py:1050-1055 — attempt validation at lines 1050-1055
• coreweave.py:627-631 — num_vms > 1 guard at lines 627-631
• client.py:1025 — iris_ctx() at line 1025
• 0001_init.sql:116 — endpoint task_id FK with ON DELETE CASCADE at line 116
• distributed.py:348-387 — DistributedConfig.initialize() at lines 348-387

These are all accurate. Note that line references in design docs tend to drift as the codebase evolves — the doc should be understood as a snapshot.

Design feedback

1. atexit for endpoint cleanup (line 84 of the design): The atexit handler to unregister the endpoint is a nice touch for graceful shutdown, but the doc correctly identifies that cascade deletes handle the abnormal-exit case. One subtlety: atexit handlers don't run on SIGKILL (which is how Kubernetes terminates pods after the grace period). This is fine because cascade deletes cover it, but worth noting in comments when implementing.

2. jax.distributed.initialize() with defaults for single-task (line 74): When num_tasks == 1, the code calls jax.distributed.initialize() with no arguments. This works for single-process JAX but is a no-op on most setups. Might be worth a comment explaining why this is called at all (e.g., "ensures JAX backend is initialized consistently regardless of task count").

3. Port collision risk: The default port 8476 is fine, but the doc should mention that if multiple JAX jobs land on the same node (with hostNetwork: true), they'll collide. The IRIS_PORT_JAX allocation path avoids this — consider making port allocation via IRIS_PORT_JAX the recommended/default path rather than the fallback, especially since hostNetwork: true means all pods on a node share the same port namespace.

4. _poll_for_coordinator backoff: The doc mentions ExponentialBackoff from iris.time_utils — good choice. The 300s timeout with 2s initial interval is reasonable. One edge case: if task 0 registers the endpoint but then its jax.distributed.initialize() call hangs (before JAX is actually listening), tasks 1..N-1 will resolve the address and try to connect to a coordinator that isn't ready yet. JAX has its own internal retry logic for this, so it should be fine in practice, but worth mentioning.

5. runtime directory: The design proposes lib/iris/src/iris/runtime/ as a new directory. This doesn't exist yet — makes sense as a namespace for these optional runtime utilities that tasks opt into.

6. Dependency direction (line 31): The non-goal about Levanter integration is correct — iris should not import levanter. The proposed Levanter callsite pattern (step 6: user code checks IRIS_CONTROLLER_ADDRESS) is the right approach. The future work item about IrisCluster in JAX's ClusterEnv registry is the cleaner long-term solution.

Minor nits

• Line 9, job_info.py:25-74: Start line should be 26 (class definition is line 26, not 25). Very minor.
• Line 92, ExponentialBackoff from iris.time_utils: This exists but is actually used across many files. Worth confirming the exact import path when implementing.
• Line 122, autoscaler feasibility: "validates that replicas is an exact multiple of num_vms" — I didn't verify this specific claim (autoscaler.py:1406). Worth double-checking during implementation.

Verdict

The design is sound. The approach of reusing existing endpoint infrastructure is pragmatic, and the implementation outline is well-scoped. The hardest part (step 3: lifting num_vms > 1 in CoreweavePlatform.create_slice()) is correctly identified and has a reasonable breakdown of sub-steps. Ready for implementation.

[rjpower]

Will start testing after this.

[yonromai]

Approve. Clean, well-layered implementation that reuses existing infrastructure (endpoint registry, ExponentialBackoff) without adding new RPCs or proto changes. The jax_init module is tight and the CoreWeave multi-VM support follows the established pattern naturally.

One actionable finding: the partial-failure cleanup in _monitor_slice wraps all deletes in a single try/except, so one failed delete prevents remaining cleanup (see inline comment). Straightforward fix — fine to address in a follow-up.

_{Generated with Claude Code}

[yonromai]

Medium — cleanup can skip pods and ConfigMap

The cleanup loop wraps all deletes in a single try/except. If kubectl delete pod raises for any pod (e.g. a pod whose creation failed and doesn't exist on the API server), subsequent pods and the ConfigMap are never cleaned up.

Suggested fix — try each delete independently:

for pod_name in handle._pod_names:
    try:
        self._kubectl.delete("pod", pod_name, force=True)
    except Exception:
        logger.warning("Failed to clean up pod %s for slice %s", pod_name, handle.slice_id)
try:
    self._kubectl.delete("configmap", _worker_config_cm_name(handle.slice_id))
except Exception:
    logger.warning("Failed to clean up ConfigMap for slice %s", handle.slice_id)

_{Generated with Claude Code}

[yonromai]

Low (non-blocking) — endpoint registered before coordinator is listening

Task 0 registers the endpoint here, then calls jax.distributed.initialize() on line 107 which starts the gRPC coordinator. During this window, tasks 1..N-1 can resolve the endpoint and try connecting to a coordinator that isn't listening yet.

JAX's internal gRPC client retries handle this gracefully in practice. But if you wanted to eliminate the race entirely, you could swap the order: call jax.distributed.initialize() first, then register. The non-coordinator tasks' ExponentialBackoff already accounts for startup delay.

Not blocking — just flagging the ordering.

_{Generated with Claude Code}

[rjpower]

@claude let's fix the issues from yonromai review

[claude]

Claude finished @rjpower's task in 2m 3s —— View job

---

Fixing review feedback from @yonromai

• Fix partial-failure cleanup in _monitor_slice — wrap each delete independently
• Fix endpoint registration ordering in jax_init.py — register after jax.distributed.initialize()
• Run pre-commit and tests
• Commit and push

---

Both issues from @yonromai's review are fixed:

1. coreweave.py cleanup — Each pod delete and the ConfigMap delete are now wrapped in their own try/except, so one failed cleanup doesn't prevent the rest from running.

2. jax_init.py registration ordering — Task 0 now calls jax.distributed.initialize() before registering the endpoint. This eliminates the race window where other tasks could discover the coordinator address before the gRPC server is actually listening.

Testing: pre-commit passes, all 54 tests in test_jax_init.py and test_coreweave_platform.py pass (including test_multi_node_partial_failure_cleans_up).

[rjpower]

@claude doesn't jax.distributed.initialize block? we need to register the endpoint before calling it IIUC, just rely on jax not being dumb.

[claude]

Claude finished @rjpower's task in 1m 58s —— View job

---

Fixing endpoint registration ordering

• Read jax_init.py to understand current state
• Revert registration ordering back to register-before-initialize
• Run pre-commit and tests
• Commit and push

---

You're right — jax.distributed.initialize() blocks until all processes connect. Calling it before registering the endpoint would deadlock, since tasks 1..N-1 would never discover the coordinator address.

Reverted to register-before-initialize. The brief window where the endpoint is registered but the gRPC coordinator isn't yet listening is handled by JAX's internal retry logic. Updated the comment to explain the ordering rationale.

Testing: pre-commit passes. All 8 tests in test_jax_init.py pass.

View job