Repository navigation
feat(cluster): make TP and pipeline serving request-safe - #6
Closed
jonathan308 wants to merge 24 commits into
Closed
jonathan308 wants to merge 24 commits into
jonathan308 wants to merge 24 commits into
Conversation
This was referenced Aug 26, 2026
jundot
force-pushed
the
feat/cluster-v2-control
branch
from
September 8, 2026 04:27
9f0c9c3 to
e3eac11
Compare
* distributed: verified teardown with SIGKILL retry, rank-marker sweep, and launch-manifest reaper DistributedJobSupervisor._terminate ignored the result of the final _wait_for_process_group_exit after SIGKILL, so a rank surviving the kill (unkillable Metal/JACCL wait, D-state, PermissionError-on-killpg) kept its unified memory resident while engine_pool had already released the model's budget (G1). killpg also only reaches the local process group, leaving remote ranks to a fragile SSH-teardown chain (G2), and all teardown lived in the coordinator process, so a crashed omlx-server orphaned everything with no reaper anywhere (G8). - _terminate now verifies the post-SIGKILL group exit, retries the kill once, then sweeps leftover rank processes by their runtime-marker pids (local via read_runtime_markers + marker_owner_is_live, remote via read_remote_marker + an SSH kill of the recorded pid). Unverified teardown raises DistributedTeardownError and keeps the process reference so a later stop() retries instead of dropping state. - start() writes an on-disk launch manifest (job id, pgid, coordinator pid, rank hosts, marker dir) removed only after verified group exit; reap_orphaned_launches() consumes manifests whose coordinator is dead and reaps their ranks. It runs at server startup (lifespan) and at the top of every supervisor start(). - engine_pool._unload_engine's distributed retry path now actually sees teardown failures because the supervisor raises them. - start()'s failure path no longer lets a teardown error mask the launch error. * distributed: unit tests for verified teardown, rank sweep, and launch-manifest reaper Covers: SIGKILL retry round, DistributedTeardownError with state kept for retry, manifest write/removal around verified stop, marker-pid sweep of local and remote leftover ranks, foreign-deployment marker immunity, and reap_orphaned_launches reaping a dead coordinator's job while leaving active jobs alone. The two pre-existing stop tests now stub the rank sweep (it would otherwise SSH to the fake 'studio' host). * distributed: rank-side force-cancel hook driven by a coordinator cancel file RuntimeTelemetry gains a force-cancel surface: the coordinator drops <deployment_id>-cancel.json next to the runtime markers; the rank's telemetry heartbeat (and each batch step, throttled to 1s) consumes new epochs and cancels every active uid through BatchGenerator.remove — the same call MLX-LM handlers make on disconnect, applied by the batch loop at a step boundary and shared with peer ranks, so cancellation never severs a collective mid-request. An ack file records each consumed epoch. This is the rank-side backstop for a handler wedged in a collective that never reaches its disconnect finally (G3), wired through the existing telemetry.cancel_uids hook point. * distributed: require rank-side quiescence evidence in pool unload gate _entry_has_active_requests now consults rank_side_active_requests() when the engine provides it: a distributed entry whose coordinator counter drained but whose rank-0 telemetry still reports active requests is not quiescent, closing the gap where abort_all_requests reported success while ranks kept generating (G5, audit hook #7). Missing marker evidence falls through to the supervisor failure paths. * distributed: surface abort (not transport error) when a closed stream raises abort_request() closes the request's response, so the SSE iterator can raise an httpx stream error before the loop-top flag check; check the abort flag first in the streaming HTTPError handlers. * distributed: graduated PeerWatchdog response — abort before exit A confirmed peer loss used to fire on_lost (marker + os._exit(1)) immediately. The watchdog now runs a graduated ladder once the failure tolerance is reached: log, fire the new on_abort hook (request cancellation at a step boundary), then re-probe for up to abort_grace seconds — a peer flap that recovers inside the window no longer suicides the healthy rank. on_lost remains the last stage. Defaults preserve the existing detection timing (same 2-strike tolerance); with on_abort unset or abort_grace=0 the behavior is byte-for-byte the old one. * distributed: widen launcher watchdog, wire abort stages on both watchdogs _watch_launcher_parent now also watches an optional launcher/coordinator lease file (env OMLX_CLUSTER_LAUNCHER_LEASE; off by default, and a lease that never appeared yet is not 'stale'), catching a launcher that is alive but wedged — invisible to the PPID check alone. Both the launcher watchdog and the rank-0 peer watchdog fire an abort stage before os._exit: writing <deployment_id>-cancel.json into the runtime state dir, which the rank's telemetry heartbeat consumes to force-cancel in-flight requests through BatchGenerator.remove at a step boundary (the cancel-file contract is consumed by the telemetry hook in pr/distributed-request-abort; merge that branch first). The peer watchdog's abort grace is configurable via OMLX_CLUSTER_PEER_ABORT_GRACE (default 5.0s). * runtime: jaccl-if-all-rdma_ctl-else-ring selection; N-node manifest + ring verified-teardown tests omlx/cluster/backends.py: select_cluster_backend (binary fail-safe rule with named blockers) and members_from_host_records. Routes: POST /admin/api/cluster/backend-selection for the plan view, and backend='auto' on /replan resolved from posted host records. Tests prove the ring backend runs through the same DistributedJobSupervisor verified teardown (TERM→KILL×2 escalation raising DistributedTeardownError, marker-pid rank sweep over all N hosts) and that the launch manifest records every host of a 3-node ring. * fix(cluster): close thundermlx stability gaps — unwire before kill, MLX/JACCL env guards - Rank SIGTERM/SIGINT handlers now release wired Metal memory (set_wired_limit(0) + clear_cache) under a 10s alarm escape hatch before the graceful KeyboardInterrupt path; launcher-watchdog and peer-watchdog force-exit paths release before os._exit. A SIGKILL cannot unwire — this closes the 'wedged rank strands wired memory until reboot' class (ported from thundermlx run_with_watchdog.py) - Rank env defaults in the launch hostfile: MLX_MAX_OPS_PER_BUFFER=16 / MLX_MAX_MB_PER_BUFFER=512 (bounds Metal command buffers under the GPU driver's ~10s execution timeout — the SIGABRT/orphan crash class) and JACCL_PROGRESS_TIMEOUT_MS=30000 / JACCL_TIMEOUT_ACTION=teardown-exit (explicit ProgressGuard pairing). All yield to operator-set env - 6 new worker tests + env-default tests; 1447 cluster tests green * Let rank-ready markers own the load barrier * Gate serving on every rank-ready marker * Separate rank control from tensor collectives * Use sustained cluster readiness canary * Clear stale cluster engines after reused PGIDs * Ignore cluster control JSON in runtime marker reader * Bootstrap JACCL through resilient side channel * Route JACCL bootstrap through macOS sidecar * Hold a process lease for JACCL RDMA * fix(cluster): verify rank exit on unload * fix(cluster): keep control transport on the runtime layer * fix(cluster): fail closed when rank exit cannot be proven * test(cluster): mark runtime-first UI seams explicitly * fix(cluster): authenticate rank-control admission without listener DoS * fix(cluster): replace public plan hash with launch-scoped control secret * fix(cluster): mutually authenticate rank-control sessions and payloads --------- Co-authored-by: jonathan308 <[email protected]>
…FT (jundot#3327) Remove the serialized prefix preload pass and report model and visible-stream TTFT separately. Stream validated Qwen tool envelopes in order while preserving terminal structured calls and existing behavior for unsupported engines. Signed-off-by: jonathan308 <[email protected]> Co-Authored-By: jonathan308 <[email protected]>
…erator reaper - abort_request(request_id): flags the request (its generator stops at the next yield boundary) and closes only that request's backend connection, replacing the all-or-nothing client close for single-request aborts (G3). Signature mirrors AsyncEngineCore.abort_request. - abort_all_requests: flags all tracked requests, writes the rank-side cancel file, closes the shared client, then waits (bounded, default 15s) for BOTH the coordinator counter and rank-0 marker telemetry to report zero active requests — unload no longer trusts client close as drain evidence (G5). - reap_orphaned_generators: pop-only sweep parallel to AsyncEngineCore._reap_orphaned_collectors; requests whose backend finished but whose consumer abandoned the generator are reaped after a grace period, so _active_requests can't leak and wedge (or spuriously unblock) quiescence-gated unload (G4). Driven from has_active_requests and abort_all_requests. - ReadTimeout now also writes the rank-side cancel file: 300s of inactivity means a stalled rank, and the old path let the request run to completion with nobody reading (G5). - Non-distributed engines are untouched; all changes are confined to DistributedBatchedEngine.
…han reaper telemetry: force_cancel_all removes active uids through the registered BatchGenerator, degrades safely with no generator/uids or a failing batch loop, and the cancel file is consumed exactly once per epoch with an ack while foreign-deployment files are ignored. engine: abort_request flags one request and closes only its connection; an aborted stream surfaces DistributedRequestAborted at the next yield boundary and drains the counter; abort_all_requests writes the rank-side cancel file and swaps the client; _wait_for_backend_drain requires rank-side marker evidence and expires unconfirmed when ranks survive; the orphan reaper drops only finished-and-abandoned requests and runs from has_active_requests; a read timeout drops the cancel file. Note: tests/test_distributed_engine.py imports omlx.engine.batched -> engine_core -> mlx, so it needs the full env (WS3 build); the same holds on the base branch.
jonathan308
force-pushed
the
feat/cluster-v2-serving
branch
from
September 8, 2026 05:11
f155c20 to
dcb7189
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.
Stack position
4/4 of the backend-first Cluster v2 split requested in jundot#3118.
Depends on fork PR #5 (
feat/cluster-v2-control). This draft exists to run the final serving/cache CI incrementally before promotion upstream.Scope
Deliberately excluded
Validation before draft
git diff --check: passed