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:
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),
|
||||
|
||||
Reference in New Issue
Block a user