Repository navigation
feat(cluster): harden rank control, teardown, and recovery - #5
Closed
jonathan308 wants to merge 29 commits into
Closed
jonathan308 wants to merge 29 commits into
jonathan308 wants to merge 29 commits into
Conversation
This was referenced Aug 26, 2026
jonathan308
force-pushed
the
feat/cluster-v2-planning
branch
from
September 8, 2026 03:33
059ec4a to
3218bf1
Compare
…oughput) (jundot#3287) * fix: use concurrent pread for Qwen4-Exp PLE row lookups _SafeTensorMMap.rows() gathers SSD-backed PLE (n-gram bank) rows with numpy fancy-indexing over an mmap'd buffer opened with MADV_RANDOM. Each row touch that isn't resident blocks synchronously on a page fault, and the gather walks the index array in a single thread, so N row lookups serialize into N blocking page faults with zero overlap. The hardware has plenty of parallelism to hide that latency; the code just never asks for it. Measured on an M4 Max (internal NVMe): 500 random 4KB reads, serial: 9,063 IOPS (110us/read) 500 random 4KB reads, 32 threads: 98,771 IOPS (x10.9) 500 random 4KB reads, 64 threads: 99,626 IOPS (x11.0) This patch replaces the mmap fancy-index gather with a thread pool of os.pread() calls on the same fd (threads, not asyncio/io_uring, because os.pread releases the GIL for the syscall duration). Verified against real PLE shard files that the patched and original code paths return byte-identical output. End-to-end on a 40k-token prompt (cold, no prefix-cache reuse), tested against two independent qwen4_exp checkpoints: jedisct1/Qwen3.8-Flash-Next-oQ4e-MTP-128k (2-3bit PLE, batched engine via a bundled omlx_support shim using the same rows() logic): cold TTFT 216.85s -> 6.54s (x33) warm TTFT 12.63s -> 6.02s (x2.1) generation 43.70 -> 59.92 tok/s Vontra/Qwen3.8-Flash-Next-MLX-oQ4-MTP (4bit PLE, vlm engine, this vendored implementation directly): isolated row-read throughput: 9,166 -> 109,636 rows/s (x12.0) steady-state TTFT: ~6-7.5s (first-ever cold-disk-cache run was an outlier at 108s, not reproducible on subsequent runs) Generation speed improves too because PLE is looked up on every decode step, not just during prefill. Scope: only affects OMLX_QWEN4_PLE_MODE=mmap (or equivalent SSD-backed PLE offload). resident mode doesn't go through this code path. Ran the existing qwen4_exp test suite (tests/test_mlx_vlm_qwen4_exp_compat.py, tests/test_vlm_qwen4_exp_loader.py, tests/test_qwen4_exp_residency.py, tests/test_qwen4_exp_mtp_admin_gate.py) — 41/41 pass, no regressions. * perf(qwen4_exp): gate concurrent PLE prefetch to cold pages, keep warm mmap gather Builds on jundot#3287 by scott-kim-7: the serialized-page-fault root cause and the concurrent os.pread thread pool are theirs. This variant keeps both, but gates the pool with a per-reader seen-page bitmap so only never-touched pages are prefetched concurrently and fully-warm lookups keep the original near-free mmap memcpy gather (measured 0.14us/row vs 22us/row for unconditional per-row pread on warm pages). Includes a byte-identity regression test over a multi-page gather with straddling rows, and handles the short-read nit raised on jundot#3287. * perf(qwen4_exp): re-arm the PLE seen-page bitmap when a warm gather lags memcpy The seen-page bitmap is write-once, so a page the OS evicted after we marked it seen would fall back to a serial minor fault forever. With a 37.5GB table under serving pressure eviction is the steady state (measured ~24% resident after 2.5h), not an edge case. mincore() detects eviction precisely but costs ~4.5ms per shard-span call, so instead every fully-seen gather self-instruments: wall time over a memcpy budget (floor + per-row allowance) is the signature of minor faults, so clear the reader's bitmap and let the next gather re-warm through the concurrent pread path. Re-arm is rate-limited to once per reader per minute and logs one INFO line (elapsed vs budget, per-shard counter) so the event is visible in field logs. Verified: 50 warm gathers worst-case 0.84ms vs an 8.69ms budget (no false re-arm), re-arm itself free, re-warm costs one concurrent pread round (9us/row), steady warm speed resumes at 0.18us/row. --------- Co-authored-by: 谢永梁 <[email protected]> Co-authored-by: Claude Opus 5 <[email protected]>
…undot#3256) * cluster v2 Module D checkpoint: path_map schema v2 + modelsync + manifest endpoint - deployment.py: path_map field, model_path_for(), schema v2 (v1 decodes to empty map), path_map rides the worker contract, decode_worker_path_map() - registry.py: deployments.json schema v1->v2 in-memory migration, persisted atomically on load, migrated_from reported - planner.py: ShardPlan.path_map passthrough (excluded from plan_hash) - routes.py: path_map on plan/deployment requests, placement signature folds path_map when present, staging honors per-node destinations - launch.py: preflight probes each node at its own path_map path - inference_worker.py: rank resolves its model path from the signed contract - modelsync.py: ModelSyncManager (status/sync, auto|rsync|download), manifest builder + GET /api/cluster/models/{model_id}/manifest, progress events - tests: 32 new offline tests (manifest fixture, path_map, v1 migration, sync decisions with mocked transports) * fix: route rsync through shared cluster SSH policy; fake path_map decoder in worker test harness - modelsync.build_rsync_argv now builds its -e command from ssh_policy.cluster_ssh_options() (seam test compliance: no literal ssh argv outside the shared policy) - test_cluster_inference_worker harness fakes decode_worker_path_map next to decode_worker_contract; the fake plan string is not a real contract * planner: RAM-proportional largest-remainder allocator for N-node plans allocate_layers_proportional (exact Fraction arithmetic, min one layer per node, deterministic tie-breaks) plus plan_proportional_pipeline, sharing the assignment/fit/hash tail with plan_unequal_pipeline via _finish_pipeline_plan. The balanced planner's hash payload is byte-identical to before. Routes gain an additive allocation field ('balanced' default) on plan/deployment/replan requests; proportional is pipeline-only. * fix(cluster): preserve proportional allocation routing * test: proportional allocator — ratios, min-1, tie-breaks, ceilings, KV fit, route wiring * runtime: one-action POST /admin/api/cluster/replan (deactivate→re-plan→reload) Extract the activation pipeline into _activate_and_report so the replan endpoint and the explicit /deployments approval share one code path. The replan endpoint derives nodes/hosts/backend from the current deployment when not posted, renders a signed preview with a placement diff, and on approval runs the quiescence-gated unload (engine-pool prepare_cluster_ reload) whose distributed supervisor teardown is verified before the registry swap and eager reload. * test: replan endpoint — preview/apply dance, busy refusal, derivation, N-node membership change * test: fix busy-pool replan test — an unchanged plan short-circuits the reload * cluster: thread path_map through the replan endpoint Agent E flagged the seam: ClusterReplanRequest had no path_map and the replan derivation dropped the current deployment's per-node paths, silently reverting a heterogeneous-path cluster to same-path resolution. Add path_map to ClusterReplanRequest, derive it from the current deployment when not posted (tracked in the 'derived' map), and pass it into the effective ClusterDeploymentRequest. Explicit override (including an empty map) wins. Tests updated + carry-forward coverage added. * Preserve measured rank speeds across replans * Preserve TP topology across cluster replans * feat(cluster): recommend tensor on known-fast links in catalogue The catalogue tie-break preferred pipeline at equal width because it survives slow links. When the paired registry shows every requested peer reachable over JACCL/Thunderbolt RDMA, tensor parallelism is the better recommendation — it splits per-token compute across the group. - _splits/assess_model/catalogue_for_cluster gain prefer_tensor (default False: API behavior unchanged for unknown links) - catalogue endpoint derives it from the paired device registry: every non-self requested node must be paired with jaccl/thunderbolt caps; unknown or unpaired peers fail closed to pipeline - 4 new tests; 33 catalogue + 1702 cluster/model tests green * fix(cluster): strip IPv6 zone ids, order routable IPs first in hostfiles macOS announces Thunderbolt interfaces as fe80::…%en10; mlx's address parser cannot read the zone id and ranks died at startup with 'Can't parse address fe80::…%en10:56758'. _validate_ip now strips the zone (machine-local by definition, never valid on the wire) and the hostfile orders routable addresses ahead of link-local fallbacks, so a Mac with configured 10.x Thunderbolt IPs never advertises the ambiguous one first. 3 tests pin the exact failure. * fix(cluster): reconcile portable paths and planning lint * test(cluster): mark backend-first replan seam explicitly * fix(cluster): protect remote rsync paths from shell interpretation --------- Co-authored-by: jonathan308 <[email protected]>
… 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.
…-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).
…el 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.
_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.
… 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.
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.
…dogs _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).
… 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.
…LX/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
jundot
force-pushed
the
feat/cluster-v2-control
branch
from
September 8, 2026 04:27
9f0c9c3 to
e3eac11
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
3 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
3/4 of the backend-first Cluster v2 split requested in jundot#3118.
Depends on fork PR #4 (
feat/cluster-v2-planning). This draft keeps CI and review incremental before the earlier upstream dependencies merge.Scope
Deliberately excluded
Validation before draft
git diff --check: passed