Distributed JACCL server crashes when two requests arrive concurrently

[nachogala]

Distributed JACCL server crashes when two requests arrive concurrently

Summary

With a distributed mlx_lm.server deployment using the JACCL backend, a
single request completes successfully, but two HTTP requests submitted at the
same time cause both distributed ranks to terminate. The clients receive an
empty reply (curl 52/HTTP 000).

Environment

• macOS on two Apple Silicon hosts with unified memory
• mlx==0.32.0
• mlx-metal==0.32.0
• mlx-lm==0.31.3
• Python 3.11
• Distributed launch via mlx.launch with a two-host JACCL hostfile
• OpenAI-compatible HTTP server endpoint

Reproduction

1. Start the distributed server with the JACCL backend and wait for
   /v1/models to return HTTP 200.
2. Submit one normal completion request; it returns HTTP 200.
3. Submit two otherwise valid completion requests concurrently (for example,
   two simultaneous POST /v1/chat/completions requests).
4. Repeat with both requests using the same generation mode. The failure is
   not specific to mixing reasoning/thinking modes.

Expected behavior

The server should serialize or safely schedule concurrent requests, or return
an explicit overload response without terminating either distributed rank.

Actual behavior

Both requests fail with an empty HTTP response. The distributed server logs
include messages equivalent to:

_pickle.UnpicklingError: invalid load key
[jaccl] Recv failed with error code -12
rank 0 exited with code 255

The peer rank exits as part of the same failure. After the crash, the server
endpoint is unavailable until restarted.

Additional observations

• A single request succeeds immediately before the concurrent pair.
• The failure reproduces with thinking=false + thinking=false,
  thinking=true + thinking=true, and mixed modes.
• The failure occurs in both supervised and manually launched server modes.
• Serializing admission at the HTTP boundary avoids the concurrent crash, but
  that is only a workaround and removes parallel request handling.

Could the JACCL transport or distributed server path support concurrent
request admission? If concurrency is currently unsupported, an explicit
server-side queue/lock or a documented limitation would prevent a process-wide
crash.

[nachogala]

Additional reproduction result: serialization of HTTP admission does not eliminate the distributed failure.

Environment: two-rank GLM 5.2 MXFP4 distributed run with the DSA attention patch applied identically on both ranks. A FIFO proxy serialized POST /v1/completions and /v1/chat/completions, so requests were admitted strictly one at a time (no concurrent server admission). The proxy itself returned two upstream failures (HTTP 502) because rank 0 still terminated during the workload.

Observed rank-0 error:

[METAL] Command buffer execution failed: Ignored (for causing prior/excessive GPU errors) (00000004:kIOGPUCommandBufferCallbackErrorSubmissionsIgnored)

Rank 1 was then terminated by the launcher. Health before the workload and a clean restart health check both passed (HTTP 200), and final shutdown was clean.

This suggests the problem is deeper than only shared request-state corruption: even serialized admission can trigger a Metal/JACCL command-buffer failure. The earlier UnpicklingError/JACCL -12 remains reproducible with true concurrent admission, but this serialized case is a separate failure mode worth investigating in the transport/Metal path.

[nachogala]

Minimal two-rank reproducer (JACCL concurrent admission)

I uploaded a self-contained lab-only harness here:

https://gist.github.com/nachogala/245ad84ecd98833997ba7cd0bd203a31

SHA-256: f9e0ebdf17cf2109b01a215d7e981eceb71b40575e9966d2da993c1a5de5a620

The script launches exactly two ranks with mlx.launch --hostfile ... --backend jaccl, waits for /v1/models, then submits two identical POST /v1/chat/completions requests concurrently. It exits non-zero when either request fails and preserves the launcher log path. It does not modify routing, model files, or any production service.

Reproduction environment from our run: macOS Apple Silicon, mlx==0.32.0, mlx-metal==0.32.0, mlx-lm==0.31.3, Python 3.11, two-host JACCL hostfile.

Expected failure from the captured distributed logs: both clients receive an empty reply/HTTP 000; rank logs contain:

```
_pickle.UnpicklingError: invalid load key
[jaccl] Recv failed with error code -12
rank 0 exited with code 255
```

We reproduced this across 80/80 concurrent pairs (OFF/OFF, ON/ON, mixed, and delayed variants). Serializing admission is a workaround, but a separate serialized run also exposed a Metal command-buffer failure (kIOGPUCommandBufferCallbackErrorSubmissionsIgnored), so the harness intentionally tests the original concurrent-admission failure without a proxy.

[nachogala]

Follow-up: after removing an orphaned lab rank that left stale JACCL state, the same two-rank GLM-5.2 mxfp4 lab setup achieved:

• JACCL all_sum smoke: 10/10 (fast synch 0 and 1), barrier 10/10, 16 MiB tensor 20/20.
• Fresh distributed start and /v1/models health: PASS.
• FIFO admission proxy (serializing requests before mlx-lm): two concurrent completions HTTP 200/200.
• Clean shutdown and re-start: health HTTP 200 and a post-restart completion HTTP 200.

This configuration did not reproduce the Metal command-buffer/JACCL crash during the short serialized run. A ~35k-token prompt progressed to 26,624/35,018 tokens but exceeded a 300s client budget; a later 40k-target run hit a fresh-start rank-0 exit 255 before the request, so long-context performance is not yet claimed. No production endpoint was involved.