Keep the server generation loop alive when a batched request fails

[rajanshxrma]

Fixes #1505 (and one of the concrete crashes reported there).

The problem

As reported in #1505, when Thread-1 (_generate) dies from any uncaught exception, mlx_lm.server turns into a zombie: ThreadingHTTPServer keeps accepting connections, but every request blocks forever in ResponseGenerator.generate() on a queue no worker will ever service, in-flight requests never resolve, and /health still returns {"status": "ok"}. I reproduced this end-to-end on current main with nothing but two plain HTTP requests (M1, 16GB, Llama-3.2-1B-Instruct-4bit, --prompt-concurrency 2):

1. request A: long prompt (~5k tokens), no sampling penalties
2. ~1s later, request B: long prompt (~3k tokens), with repetition_penalty

Result: TypeError: 'NoneType' object is not iterable at generate.py:1428, generation thread dead, both requests hang, all subsequent requests hang, /health still ok — one of the exact exceptions listed in the issue.

Root cause of that TypeError

PromptProcessingBatch.extend uses None placeholders when no sequence so far has logits processors:

• A enters the prompt batch alone → self.logits_processors = [None]
• B (with processors) joins a step later → [None, [processor]]
• both finish prefill in the same step → split().generate() → GenerationBatch._step iterates entry None → TypeError

Everything else in the file treats "no processors" as an empty list (filter() uses [[]] * len(keep)); extend is the one place that violates the invariant. The sampler None placeholders are fine — the use site falls back (self.samplers[e] or self.fallback_sampler) — so those are left untouched.

Changes

• generate.py: empty-list placeholders in PromptProcessingBatch.extend, consistent with filter(). This also keeps the any(...)-gated fast path (all-empty still skips the processor loop).
• server.py: the design fix for server: any uncaught exception in _generate leaves HTTP threads serving while every completion hangs forever #1505 —
  • batch admission path: the whole per-request block (state machine, cache fetch, insert_segments) is now covered by the same per-request exception handling that _tokenize and _serve_single already had, so one bad request fails with an error response instead of taking down the thread. The context is only handed to the client after the request is actually registered in the batch.
  • batch serving path: an exception now fails the in-flight requests with the exception (clients get an error instead of a hang), logs the traceback, closes and resets the batch generator, and keeps serving. Queued requests are retried against a fresh generator; a poison request can't loop because it fails from inside batch_results on the next crash.
  • ResponseGenerator.generate() raises immediately if the generation thread is not alive, so any residual way to kill the thread turns into error responses rather than infinite client hangs.

This also covers the other exceptions listed in #1505 ([metal::malloc] resource-limit errors under sustained load, etc.): the batch is dropped, the affected requests get errors, and the server keeps serving instead of bricking.

Tests

• tests/test_generate.py::test_batch_extend_mixed_logits_processors — deterministic regression test for the mixed-placeholder crash (staggered inserts sized to co-finish prefill); fails with TypeError on current main, passes with the fix.
• tests/test_server.py::TestResponseGeneratorResilience — injects a RuntimeError into BatchGenerator.next, asserts the caller receives the exception, the generation thread survives, and a subsequent request completes normally.
• Full tests/test_server.py (25) and tests/test_generate.py (27) pass locally on an M1.
• End-to-end: the two-request scenario above completes cleanly on the fixed build (both finish_reason: length, no traceback, follow-up requests fine).

[TadeuszWolfGang]

I independently validated the deterministic mixed-logits_processors fix at
exact commit 58d38e4f57ba4f3d2495a3ed9b8a0da7ba129960, tested through stacked
head 6e3f367c4b4b6693ef1006d2b9fc23dc4695842b.

Environment: Apple M5 Max, macOS 26.5.2, Python 3.12.13, mlx 0.31.2. I used the
locally cached mlx-community/Qwen1.5-0.5B-Chat-4bit fixture and ran offline.

I ran the same API-level BatchGenerator sequence on both sides:

1. insert a 20-token request without logits processors;
2. advance one prefill step;
3. insert a 15-token request with repetition_penalty=1.5;
4. let both enter generation on the same step.

Result:

• pre-fix commit 15b522f fails deterministically with
  TypeError: 'NoneType' object is not iterable in
  GenerationBatch._step();
• the Keep the server generation loop alive when a batched request fails #1513 code completes both request UIDs.

I also ran test_generation_thread_survives_batch_exception, which confirms
that an injected batch exception is returned to the in-flight request, the
generation thread stays alive, and a subsequent request completes. The full
tests/test_server.py file passed 27/27 locally.

Scope: this independently confirms the None placeholder root cause and the
in-process exception-recovery test. It is not a two-request HTTP end-to-end
replay, and I did not inject a live [metal::malloc] Resource limit failure.

[TadeuszWolfGang]

Independent validation for the deterministic mixed-logits_processors failure and the in-process recovery path is recorded above at exact head 58d38e4. The full tests/test_server.py suite also passes locally with the stacked follow-up changes.

#1598 is stacked on this branch, but #1513 still has no review. Could a maintainer take a look when available?

[TadeuszWolfGang]

Independent production reproduction of #1505 on stock mlx-lm 0.31.3, captured
live on 2026-07-30 before restarting the server.

The exception is the one #1505 lists verbatim, down to the same number —
RuntimeError: [metal::malloc] Resource limit (499000) exceeded — reproduced
independently on different hardware (M5 Max / 128 GB here vs M3 Max / 36 GB in
the issue), unprovoked, under ordinary chat traffic. This PR's own end-to-end
reproduction is the TypeError from mixed logits_processors; this covers the
second exception in that list. It is the fourth occurrence on this host since
2026-07-05, all four with byte-identical tracebacks.

What happened

mlx_lm.server serving mlx-community/Qwen3.6-27B-8bit
(--max-tokens 32768 --decode-concurrency 64 --prompt-concurrency 16),
started 12:00:06, two streaming completions admitted at 16:34:37 and 16:35:52,
prefills of 8,338 and 14,815 tokens completed normally. Then:

Exception in thread Thread-1 (_generate):
Traceback (most recent call last):
  File ".../mlx_lm/server.py", line 853, in _generate
    prompt_responses, gen_responses = batch_generator.next()
  File ".../mlx_lm/generate.py", line 1855, in next
    return self._next()
  File ".../mlx_lm/generate.py", line 1775, in _next
    generation_responses = self._generation_batch.next()
  File ".../mlx_lm/generate.py", line 1415, in next
    tokens, logprobs = self._step()
  File ".../mlx_lm/generate.py", line 1369, in _step
    mx.async_eval(self._next_tokens, self._next_logprobs, token_context)
RuntimeError: [metal::malloc] Resource limit (499000) exceeded.

Nothing was generated after that. Measured on the still-zombied process, before
restarting it:

Probe	Result
GET /v1/models	200 in 0.021 s
4-token POST /v1/chat/completions	no response in 60 s (http_code=000)

py-spy on the zombied PID

• No Thread-1 (_generate) thread exists — zero matches for _generate in
  the dump. It is dead, not blocked.
• MainThread is in serve_forever (server.py:1729), still accepting.
• Request threads are parked forever in Queue.get(), at two distinct sites:
  • server.py:1037, inside _inner(), waiting for the next chunk — a request
    admitted before the death, hung mid-stream;
  • server.py:1048, ctx = response_queue.get(), waiting for the admission
    ack — requests submitted after the death.

One detail worth noting for anyone reading server logs for this failure: a
completion submitted after the death blocks at server.py:1048, which is
reached before _set_stream_headers, so log_request never runs for it. No
POST /v1/chat/completions line is written at all. The absence of POST lines in
a zombie window understates the number of hung requests rather than showing
there were none.

Recurrence

Same exception, same stack, four times:

#	Generation thread died after	Server restarted	Gap	Served while zombied
1	2026-07-05 07:07:40	2026-07-07 11:42:03	~52 h 34 min	5× GET /v1/models 200, 1× GET /health 200
2	2026-07-21 11:12:44	2026-07-21 17:14:42	~6 h 01 min	7× GET /v1/models 200
3	2026-07-27 08:00:05	2026-07-28 10:41:33	~26 h 41 min	1× GET /v1/models 200
4	2026-07-30, 16:39:53–16:57:54	2026-07-30 22:15:30	~5 h 35 min	2× GET /v1/models 200

The GET /health 200 during window 1 matches the health-endpoint behaviour
described in #1505 and in this PR.

These instances are supervised by launchd with KeepAlive enabled, which cannot
help: the process never exits, only the thread does. Every recovery above was a
manual or scheduled process restart — occurrence 4 ended with an explicit
launchctl kickstart at 22:15:30, after which a 4-token completion returned 200
in 1.554 s with no configuration change.

How this maps to the PR

1. Pre-death request hung mid-stream (server.py:1037) → the batch-serving
   handler fails in-flight requests with the exception.
2. Post-death requests hung before headers (server.py:1048) →
   ResponseGenerator.generate() raising when the generation thread is not
   alive turns these into errors.
3. No recovery without a restart → generator closed, batch reset, loop keeps
   serving.

Scope of this evidence

• The [metal::malloc] exhaustion is a host condition, not an mlx-lm defect:
  this 128 GB host had a second large model resident at the time. The claim here
  is only that the current failure handling turns a recoverable per-request
  error into an indefinitely bricked server, and that this PR changes that.
• I did not run the fixed build under the same memory pressure, so this is
  not a validation of recovery for this trigger. A second instance on the same
  host running the stacked Keep the server generation loop alive when a batched request fails #1513+Detect and recover from livelocked batch generation (#1493) #1598 build answered a 4-token completion in
  0.478 s at the same moment, but it was not subjected to the same load.
• Gaps for occurrences 1–3 are wall-clock intervals that include host sleep.
  Only occurrence 4 is confirmed continuous uptime.
• Only the captured event is version-verified as 0.31.3.

Environment: Apple M5 Max, 128 GB, macOS 26.5.2, Python 3.12.13, mlx-lm 0.31.3,
mlx 0.31.2, stock uv tool install with no patches.

Full capture (py-spy dump, log excerpts, probe transcripts, code proof) is
archived and I can share a redacted subset on request. Happy to run an
instrumented or patched build against the same workload if that is useful.