ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

fix: drain stale events from resp_queue after generation cancel

Browse source
This commit is contained in:
Roland Tannous 2026-03-07 08:12:16 +00:00
commit 717d5d621d
1 changed files with 22 additions and 1 deletions
Download patch file Download diff file Expand all files Collapse all files

23
studio/backend/core/inference/orchestrator.py
View file

@ -229,6 +229,24 @@ class InferenceOrchestrator:
except (EOFError, OSError, ValueError):
return events
def _drain_until_gen_done(self, timeout: float = 5.0) -> None:
"""Consume resp_queue events until gen_done/gen_error, discarding them.
Called after cancel to ensure stale tokens from the cancelled
generation don't leak into the next request.
"""
deadline = time.monotonic() + timeout
while time.monotonic() < deadline:
resp = self._read_resp(timeout=min(0.5, deadline - time.monotonic()))
if resp is None:
if not self._ensure_subprocess_alive():
return
continue
rtype = resp.get("type", "")
if rtype in ("gen_done", "gen_error"):
return
logger.warning("Timed out waiting for gen_done after cancel")
# ------------------------------------------------------------------
# Public API — same interface as InferenceBackend
# ------------------------------------------------------------------
@ -498,8 +516,11 @@ class InferenceOrchestrator:
if rtype == "token":
# Check cancel from route (e.g. SSE connection closed)
if cancel_event is not None and cancel_event.is_set():
# Set the subprocess mp.Event — instant, no queue needed
self._cancel_generation()
# Wait for the subprocess to acknowledge cancellation
# (gen_done/gen_error) so stale events don't leak into
# the next generation request.
self._drain_until_gen_done(timeout=5.0)
return
yield resp.get("text", "")