fix: guard against ZeroDivisionError in mlx_lm stats (#1707)

## Problem

Running short completions (like `max_tokens=1` health check probes) can
finish so fast that `mlx_lm`'s internal `generation_time` rounds to
zero. When that happens, `BatchGenerator.stats()` in
`mlx_lm/generate.py` divides `generation_tokens / generation_time` and
throws a `ZeroDivisionError`, which kills the runner process.

EXO already handles this on its side — lines 289-293 in
`batch_generate.py` guard the TPS calculation with `if
generation_elapsed > 0`. But the call to `self._exo_gen.stats()` on line
294 goes into mlx_lm's *separate* timing code, which doesn't have the
same guard. Two different timers, only one is protected.

In my case, this was triggered by health check probes (content: `"a"`,
`max_tokens=1`). The generation completed in sub-microsecond time,
`generation_time` was exactly `0`, and the runner crashed. Since the
health check command stays in the queue and retries after recovery, it
created an infinite crash loop — every ~15 seconds the runner would load
the model, get the same health check, and die again.

## Fix

Wrap `self._exo_gen.stats()` in a `try/except ZeroDivisionError`. If it
throws, set `mlx_stats` to `None` and fall back to `0.0` for
`prompt_tps`. The only fields EXO reads from `mlx_stats` are
`prompt_tps` and `prompt_time` — losing them on a sub-microsecond
generation has no practical impact.

## Traceback

```
File "batch_generate.py", line 294, in step
    mlx_stats = self._exo_gen.stats()
File "mlx_lm/generate.py", line 1224, in stats
    self._stats.generation_tokens / self._stats.generation_time
ZeroDivisionError: division by zero
```
This commit is contained in:
Mustafa Alp Yılmaz
2026-03-12 17:48:36 +03:00
committed by GitHub
parent 0782d90ec5
commit ea18a62581
@@ -291,10 +291,13 @@ class ExoBatchGenerator:
if generation_elapsed > 0
else 0.0
)
mlx_stats = self._exo_gen.stats()
try:
mlx_stats = self._exo_gen.stats()
except ZeroDivisionError:
mlx_stats = None
stats = GenerationStats(
prompt_tps=float(mlx_stats.prompt_tps)
if mlx_stats.prompt_time > 0
if mlx_stats is not None and mlx_stats.prompt_time > 0
else 0.0,
generation_tps=float(generation_tps),
prompt_tokens=len(state.all_prompt_tokens),