From 6c5a33469f4c517cfcaba9df00cfcd5dcde21526 Mon Sep 17 00:00:00 2001 From: mac Date: Wed, 26 Aug 2026 23:11:33 +0300 Subject: [PATCH 1/2] daemon: never orphan the /bench/run child on stream cancellation MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The SSE generator had no try/finally: a client disconnect mid-stream closed the generator and left the bench child running — holding the GPU this route exists to hand over, and writing the shared profile — and never reaped. Spawn the child in its own process group and guarantee teardown on any exit from the stream: SIGTERM the group, 5 s grace, SIGKILL escalation, unconditional reap (shielded so a second cancellation mid-teardown cannot skip it). Wire format is unchanged. Found via the freetoken-mlx downstream audit (docs/AUDIT.md, defect 1). --- python/freetoken/daemon/app.py | 83 +++++++++++-- tests/daemon/test_daemon_app.py | 210 ++++++++++++++++++++++++++++++++ 2 files changed, 282 insertions(+), 11 deletions(-) create mode 100644 tests/daemon/test_daemon_app.py diff --git a/python/freetoken/daemon/app.py b/python/freetoken/daemon/app.py index d7a0a53c6..9fba87faf 100644 --- a/python/freetoken/daemon/app.py +++ b/python/freetoken/daemon/app.py @@ -13,6 +13,7 @@ import functools import json import os +import signal import sys from concurrent.futures import ThreadPoolExecutor from typing import Any, Callable @@ -26,6 +27,9 @@ from .version import DAEMON_VERSION + + + class StartBody(BaseModel): model: str port: int | None = None @@ -53,6 +57,30 @@ class CancelBody(BaseModel): id: str +class BenchBody(BaseModel): + # Raw `ft bench bw` args (e.g. ["--dtype", "nvfp4", "--threshold", "2.5"]); empty = all dtypes. + args: list[str] = [] +class StopBody(BaseModel): + force: bool = False + + +class SwitchBody(StartBody): + force: bool = False + + +class AccountingAckBody(BaseModel): + receiptId: str + + +class CheckpointBody(BaseModel): + id: str + args: list[str] = [] + + +class CancelBody(BaseModel): + id: str + + class BenchBody(BaseModel): # Raw `ft bench bw` args (e.g. ["--dtype", "nvfp4", "--threshold", "2.5"]); empty = all dtypes. args: list[str] = [] @@ -113,6 +141,25 @@ def _parse_ftbench(line: str) -> dict | None: return None +async def _terminate_and_reap_bench_child(proc: asyncio.subprocess.Process, grace_s: float = 5.0) -> None: + if proc.returncode is not None: + return + try: + os.killpg(proc.pid, signal.SIGTERM) + except (ProcessLookupError, PermissionError): + pass + try: + await asyncio.wait_for(proc.wait(), timeout=grace_s) + return + except asyncio.TimeoutError: + pass + try: + os.killpg(proc.pid, signal.SIGKILL) + except (ProcessLookupError, PermissionError): + pass + await proc.wait() + + def build_app( *, manager, @@ -350,7 +397,13 @@ async def gen(): argv = [sys.executable, "-m", "freetoken.cli", "bench", "bw", *body.args] try: proc = await asyncio.create_subprocess_exec( - *argv, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.STDOUT, env=env + *argv, + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.STDOUT, + env=env, + # Own process group: lets us tear down the whole bench tree (it may + # spawn its own children) when the client disconnects mid-stream. + start_new_session=True, ) except Exception as exc: # noqa: BLE001 yield _bench_sse("error", {"message": f"failed to spawn bench: {exc}"}) @@ -358,16 +411,24 @@ async def gen(): tail: collections.deque = collections.deque(maxlen=8) # last non-progress lines (errors) out_path: str | None = None assert proc.stdout is not None - async for raw in proc.stdout: - line = raw.decode(errors="replace").rstrip() - prog = _parse_ftbench(line) - if prog is not None: - yield _bench_sse("progress", prog) - elif line.startswith("FTBENCH_OUT "): - out_path = line[len("FTBENCH_OUT "):] - elif line: - tail.append(line) - rc = await proc.wait() + rc: int | None = None + try: + async for raw in proc.stdout: + line = raw.decode(errors="replace").rstrip() + prog = _parse_ftbench(line) + if prog is not None: + yield _bench_sse("progress", prog) + elif line.startswith("FTBENCH_OUT "): + out_path = line[len("FTBENCH_OUT "):] + elif line: + tail.append(line) + rc = await proc.wait() + finally: + # Client disconnect / task cancellation unwinds here; never orphan the + # bench child: terminate, 5s grace, kill escalation, always reap. + await asyncio.shield(_terminate_and_reap_bench_child(proc)) + if rc is None: + return if rc != 0: yield _bench_sse("error", {"message": "\n".join(tail) or f"bench exited {rc}"}) return diff --git a/tests/daemon/test_daemon_app.py b/tests/daemon/test_daemon_app.py new file mode 100644 index 000000000..f64220da6 --- /dev/null +++ b/tests/daemon/test_daemon_app.py @@ -0,0 +1,210 @@ +"""Tests for /bench/run child-process ownership.""" + +from __future__ import annotations + +import asyncio +import json +import os +import sys +import time +from concurrent.futures import ThreadPoolExecutor + +import pytest + +from freetoken.daemon import app as app_mod +from freetoken.daemon.app import BenchBody, build_app +from freetoken.daemon.logring import LogRing + + +# --------------------------------------------------------------------------- fakes + + +class _Mgr: + def status(self): + return {"running": False} + + def start(self, model, port, args): + return {"running": True} + + def stop(self, *a): + return {} + + def current_pid(self): + return None + + +def _build_app(monkeypatch, spawn_fn): + """Build app with `asyncio.create_subprocess_exec` swapped for ``spawn_fn``; return (app, endpoint).""" + monkeypatch.setattr(app_mod.asyncio, "create_subprocess_exec", spawn_fn) + application = build_app( + manager=_Mgr(), + ring=LogRing(), + probe=None, + footprint_fn=lambda pid: {}, + lifecycle_pool=ThreadPoolExecutor(1), + proxy_pool=ThreadPoolExecutor(1), + ) + route = next(r for r in application.routes if getattr(r, "path", None) == "/bench/run") + return application, route.endpoint + + +def _spawn_recorder(script: str): + """Return (spawn_fn, pids). The app module's create_subprocess_exec is swapped for spawn_fn, + which asserts group ownership, records child pid, and launches ``script``.""" + pids: list[int] = [] + real_create = asyncio.create_subprocess_exec + + def spawn(*argv, **kwargs): + assert kwargs.get("start_new_session") is True, "child must own its process group" + assert kwargs.get("stdout") == asyncio.subprocess.PIPE + + class _Spawn: + def __await__(self): + proc = yield from real_create( + sys.executable, + "-c", + script, + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.STDOUT, + start_new_session=True, + ).__await__() + pids.append(proc.pid) + return proc + + return _Spawn() + + return spawn, pids + + +def _alive(pid: int) -> bool: + try: + os.kill(pid, 0) + return True + except ProcessLookupError: + return False + + +async def _wait_gone(pid: int, timeout: float = 5.0) -> bool: + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + if not _alive(pid): + return True + await asyncio.sleep(0.05) + return not _alive(pid) + + +async def _collect(body, bench_run) -> list[tuple[str, dict]]: + resp = await bench_run(body=body) + out = [] + async for chunk in resp.body_iterator: + event, data = None, None + for raw_line in chunk.splitlines(): + line = raw_line.decode(errors="replace") if isinstance(raw_line, (bytes, bytearray)) else raw_line + if line.startswith("event: "): + event = line[len("event: ") :] + elif line.startswith("data: "): + data = json.loads(line[len("data: ") :]) + out.append((event, data)) + return out + + +# --------------------------------------------------------------------------- bench SSE frames + + +def test_bench_spawn_error_frame(monkeypatch): + async def boom(*argv, **kwargs): + raise RuntimeError("no such binary") + + _, bench_run = _build_app(monkeypatch, boom) + events = asyncio.run(_collect(BenchBody(), bench_run)) + assert len(events) == 1 + event, data = events[0] + assert event == "error" + assert "failed to spawn bench" in data["message"] + + +def test_bench_progress_frames(monkeypatch): + script = 'print("FTBENCH 1 2 dtype", flush=True)\nprint("FTBENCH 2 2 dtype", flush=True)\n' + spawn, pids = _spawn_recorder(script) + monkeypatch.setattr(app_mod, "_read_bench_profile", lambda *_a, **_k: None) + _, bench_run = _build_app(monkeypatch, spawn) + events = asyncio.run(_collect(BenchBody(), bench_run)) + kinds = [e for e, _ in events] + assert kinds == ["progress", "progress", "error"] + assert events[-1][1]["message"] == "bench finished but no profile was written" + for pid in pids: + assert not _alive(pid) + + +def test_bench_result_frame(monkeypatch): + script = 'print("FTBENCH 1 1 dtype", flush=True)\n' + spawn, _ = _spawn_recorder(script) + monkeypatch.setattr(app_mod, "_read_bench_profile", lambda *_a, **_k: {"result": {"ok": True}}) + _, bench_run = _build_app(monkeypatch, spawn) + events = asyncio.run(_collect(BenchBody(), bench_run)) + kinds = [e for e, _ in events] + assert kinds == ["progress", "result"] + assert events[-1][1] == {"result": {"ok": True}} + + +# --------------------------------------------------------------------------- orphaning on close + + +TERM_DIES = ( + 'import sys, time\n' + 'print("FTBENCH 1 99 dtype", flush=True)\n' + "time.sleep(60)\n" +) +TERM_TRAPS = ( + "import signal, time\n" + "signal.signal(signal.SIGTERM, lambda *a: None)\n" + 'print("FTBENCH 1 99 dtype", flush=True)\n' + "time.sleep(60)\n" +) + + +async def _run_close_case(monkeypatch, script): + """Park the bench generator at its first SSE yield while the child still runs, + then aclose() it and assert the child was terminated and reaped.""" + spawn, pids = _spawn_recorder(script) + _, bench_run = _build_app(monkeypatch, spawn) + resp = await bench_run(body=BenchBody()) + agen = resp.body_iterator + frame = await agen.__anext__() + if isinstance(frame, (bytes, bytearray)): + frame = frame.decode(errors="replace") + assert frame.startswith("event: progress"), frame + while not pids or not _alive(pids[0]): + await asyncio.sleep(0.02) + pid = pids[0] + return agen, pid + + +def test_generator_close_terminates_and_reaps_child(monkeypatch): + """Reproduces the pre-fix bug: closing the response generator mid-stream used to + abandon the sleeping bench child. After the fix, aclose() terminates and reaps.""" + + async def run(): + agen, pid = await _run_close_case(monkeypatch, TERM_DIES) + t0 = time.monotonic() + await agen.aclose() + elapsed = time.monotonic() - t0 + assert elapsed < 4.5, "well-behaved child should die on SIGTERM, not wait out grace" + assert await _wait_gone(pid), "bench child orphaned after generator close" + + asyncio.run(run()) + + +@pytest.mark.slow +def test_kill_only_on_ignored_terminate(monkeypatch): + """A child that traps SIGTERM survives the grace period and dies by SIGKILL.""" + + async def run(): + agen, pid = await _run_close_case(monkeypatch, TERM_TRAPS) + t0 = time.monotonic() + await agen.aclose() + elapsed = time.monotonic() - t0 + assert elapsed >= 4.5, "SIGKILL must wait out the terminate grace" + assert await _wait_gone(pid), "term-trapping child killed after grace" + + asyncio.run(run()) From 9a11c6e0dbc07ccbbef30ed0802e4b6e498aa21b Mon Sep 17 00:00:00 2001 From: mac Date: Wed, 26 Aug 2026 23:11:57 +0300 Subject: [PATCH 2/2] daemon: bound request-model inputs StartBody.port (and SwitchBody's, via inheritance) was an unbounded int flowing into serve lifecycle mutation; constrain it to a real socket port (1..65535, None = default). List args defaults move to Field(default_factory=list). Found via the freetoken-mlx downstream audit (docs/AUDIT.md, defect 5). --- python/freetoken/daemon/app.py | 11 +++++----- tests/daemon/test_daemon_app.py | 36 +++++++++++++++++++++++++++++++-- 2 files changed, 40 insertions(+), 7 deletions(-) diff --git a/python/freetoken/daemon/app.py b/python/freetoken/daemon/app.py index 9fba87faf..974cde745 100644 --- a/python/freetoken/daemon/app.py +++ b/python/freetoken/daemon/app.py @@ -20,7 +20,7 @@ from fastapi import Depends, FastAPI, Header, HTTPException, Request from fastapi.responses import JSONResponse, StreamingResponse -from pydantic import BaseModel +from pydantic import BaseModel, Field from .accounting import AccountingOutboxError, AccountingPrepareError from .serve_manager import Conflict @@ -32,8 +32,8 @@ class StartBody(BaseModel): model: str - port: int | None = None - args: list[str] = [] + port: int | None = Field(default=None, ge=1, le=65535) + args: list[str] = Field(default_factory=list) class StopBody(BaseModel): @@ -50,7 +50,7 @@ class AccountingAckBody(BaseModel): class CheckpointBody(BaseModel): id: str - args: list[str] = [] + args: list[str] = Field(default_factory=list) class CancelBody(BaseModel): @@ -59,7 +59,8 @@ class CancelBody(BaseModel): class BenchBody(BaseModel): # Raw `ft bench bw` args (e.g. ["--dtype", "nvfp4", "--threshold", "2.5"]); empty = all dtypes. - args: list[str] = [] + args: list[str] = Field(default_factory=list) + class StopBody(BaseModel): force: bool = False diff --git a/tests/daemon/test_daemon_app.py b/tests/daemon/test_daemon_app.py index f64220da6..496cea38f 100644 --- a/tests/daemon/test_daemon_app.py +++ b/tests/daemon/test_daemon_app.py @@ -1,4 +1,4 @@ -"""Tests for /bench/run child-process ownership.""" +"""Tests for daemon app request models and /bench/run child-process ownership.""" from __future__ import annotations @@ -9,10 +9,11 @@ import time from concurrent.futures import ThreadPoolExecutor +import pydantic import pytest from freetoken.daemon import app as app_mod -from freetoken.daemon.app import BenchBody, build_app +from freetoken.daemon.app import BenchBody, CheckpointBody, StartBody, build_app from freetoken.daemon.logring import LogRing @@ -108,6 +109,37 @@ async def _collect(body, bench_run) -> list[tuple[str, dict]]: return out +# --------------------------------------------------------------------------- request models + + +def test_start_body_port_bounds(): + with pytest.raises(pydantic.ValidationError): + StartBody(model="m", port=0) + with pytest.raises(pydantic.ValidationError): + StartBody(model="m", port=-1) + with pytest.raises(pydantic.ValidationError): + StartBody(model="m", port=65536) + assert StartBody.model_validate({"model": "m"}).port is None + assert StartBody(model="m", port=1).port == 1 + assert StartBody(model="m", port=65535).port == 65535 + + +def test_invalid_port_fails_before_manager_start(): + with pytest.raises(pydantic.ValidationError): + StartBody.model_validate({"model": "m", "port": 70000}) + assert _Mgr().status().get("running") is False + + +@pytest.mark.parametrize("cls", [StartBody, CheckpointBody, BenchBody]) +def test_list_defaults_are_independent(cls): + kwargs = {"model": "m"} if cls is StartBody else {"id": "x"} if cls is CheckpointBody else {} + a = cls.model_validate(kwargs) + b = cls.model_validate(kwargs) + a.args.append("--dtype") + assert b.args == [] + assert cls.model_validate(kwargs).args == [] + + # --------------------------------------------------------------------------- bench SSE frames