From 7df3c6dd2219d50774b5aca48e0246e8d08b63d3 Mon Sep 17 00:00:00 2001 From: Yevhenii Date: Tue, 6 Oct 2026 13:27:13 +0300 Subject: [PATCH] fix(server): leak no semaphores when a signal stops the server --- python/freetoken/server/launch.py | 8 ++++++ python/freetoken/server/supervisor.py | 7 +++++ tests/server/test_launch.py | 37 +++++++++++++++++++++++++++ tests/server/test_supervisor.py | 31 ++++++++++++++++++++++ 4 files changed, 83 insertions(+) create mode 100644 tests/server/test_launch.py diff --git a/python/freetoken/server/launch.py b/python/freetoken/server/launch.py index bb0c256f46..5179e8e238 100644 --- a/python/freetoken/server/launch.py +++ b/python/freetoken/server/launch.py @@ -59,6 +59,14 @@ def _run_scheduler(args: ServerArgs, ack_queue: mp.Queue[str]) -> None: if args.shell_mode: _detach_process_group() + import threading + + from tqdm import tqdm + + # Only rank 0 draws a bar, so no bar lock spans processes; tqdm's default one is a multiprocessing + # semaphore, which this worker leaves to resource_tracker as leaked when a signal ends it. + tqdm.set_lock(threading.RLock()) + # published (not bound) here: the engine binds it after the allocator setup from freetoken.gpu_select import set_assigned_gpu diff --git a/python/freetoken/server/supervisor.py b/python/freetoken/server/supervisor.py index 322419a65f..20deb7d132 100644 --- a/python/freetoken/server/supervisor.py +++ b/python/freetoken/server/supervisor.py @@ -55,6 +55,12 @@ class BackendHandle: processes: List[Any] = field(default_factory=list) expected_acks: int = 0 + def release_acks(self) -> None: + """Drop the ack queue once readiness has read it for the last time. A multiprocessing queue's + semaphores are unlinked only when it is collected, and a main process that a signal ends never + collects it, so resource_tracker reports them as leaked.""" + self.ack_queue = None + class WorkerDied(Exception): def __init__(self, message: str) -> None: @@ -172,6 +178,7 @@ def _shutting_down() -> bool: on_failure(exc.message) return + handle.release_acks() on_ready() while True: diff --git a/tests/server/test_launch.py b/tests/server/test_launch.py new file mode 100644 index 0000000000..defba4f307 --- /dev/null +++ b/tests/server/test_launch.py @@ -0,0 +1,37 @@ +"""The scheduler worker's own process setup, before it builds the engine.""" + +from __future__ import annotations + +import queue +import threading + +import pytest +import torch +from tqdm import tqdm + +import freetoken.gpu_select +import freetoken.scheduler +from freetoken.distributed import DistributedInfo +from freetoken.server import launch +from freetoken.server.args import ServerArgs +from freetoken.utils import progress + + +def test_the_scheduler_worker_draws_its_bars_under_a_lock_of_its_own(monkeypatch): + """Only rank 0 draws a bar, so no bar lock spans processes. tqdm's default one is a multiprocessing + semaphore, which a worker ended by a signal leaves to resource_tracker as leaked.""" + + class EngineBuilt(Exception): + pass + + def scheduler(args): + raise EngineBuilt(tqdm.get_lock()) + + monkeypatch.setattr(freetoken.scheduler, "Scheduler", scheduler) + monkeypatch.setattr(freetoken.gpu_select, "set_assigned_gpu", lambda target: None) + monkeypatch.setattr(progress, "_PROGRESS_SINK", None) + monkeypatch.setattr(tqdm, "_lock", tqdm.get_lock()) + args = ServerArgs(model_path="unused", tp_info=DistributedInfo(0, 1), dtype=torch.bfloat16) + with pytest.raises(EngineBuilt) as built: + launch._run_scheduler(args, queue.Queue()) + assert type(built.value.args[0]) is type(threading.RLock()) diff --git a/tests/server/test_supervisor.py b/tests/server/test_supervisor.py index 3568d90505..1ee1ca17bf 100644 --- a/tests/server/test_supervisor.py +++ b/tests/server/test_supervisor.py @@ -211,6 +211,37 @@ def on_ready() -> None: assert "failure" not in seen # graceful stop: the death was not reported +def test_supervisor_releases_the_ack_queue_once_the_backend_is_ready(): + """Readiness is the ack queue's last read. A multiprocessing queue's semaphores are unlinked only when + it is collected, and a main process ended by a signal never collects it.""" + import multiprocessing + import weakref + + from freetoken.server.supervisor import run_backend_supervisor + + class Proc: + name = "freetoken-TP0-scheduler" + alive = True + + def is_alive(self) -> bool: + return self.alive + + proc = Proc() + ack_queue = multiprocessing.get_context("spawn").Queue() + ack_queue.put("scheduler ready") + queue_at_ready = weakref.ref(ack_queue) + handle = BackendHandle(ack_queue=ack_queue, processes=[proc], expected_acks=1) + del ack_queue + seen: dict = {} + + def on_ready() -> None: + seen["queue at ready"] = queue_at_ready() + proc.alive = False + + run_backend_supervisor(handle, LoadProgress(), on_ready=on_ready, poll=0.01, is_shutting_down=lambda: True) + assert seen == {"queue at ready": None} + + def test_supervisor_silent_on_startup_death_during_shutdown(): """A worker dying mid-load while an orderly stop is already in progress must not be reported as a load failure either."""