fix(server): leak no semaphores when a signal stops the server

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
 

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:

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())

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."""