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

ea18a625813d36069956ba742e8f519eabee05b2 · 2026-03-12 17:48:36 +0300 · Mustafa Alp Yılmaz

## 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
```

Files touched

M src/exo/worker/engines/mlx/generator/batch_generate.py

Diff

commit ea18a625813d36069956ba742e8f519eabee05b2
Author: Mustafa Alp Yılmaz <96022931+mustafalpyilmaz@users.noreply.github.com>
Date:   Thu Mar 12 17:48:36 2026 +0300

    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
    ```
---
 src/exo/worker/engines/mlx/generator/batch_generate.py | 7 +++++--
 1 file changed, 5 insertions(+), 2 deletions(-)

diff --git a/src/exo/worker/engines/mlx/generator/batch_generate.py b/src/exo/worker/engines/mlx/generator/batch_generate.py
index 1911a97c..3462a873 100644
--- a/src/exo/worker/engines/mlx/generator/batch_generate.py
+++ b/src/exo/worker/engines/mlx/generator/batch_generate.py
@@ -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),

← 0782d90e fix: show partial download progress on initial dashboard loa · back to Exo · Extract mlx revision from uv lock (#1715) f28b2fd0 →