From ea18a625813d36069956ba742e8f519eabee05b2 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Mustafa=20Alp=20Y=C4=B1lmaz?=
 <96022931+mustafalpyilmaz@users.noreply.github.com>
Date: Thu, 12 Mar 2026 17:48:36 +0300
Subject: [PATCH] fix: guard against ZeroDivisionError in mlx_lm stats (#1707)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

## 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),