feat(gpu-arbiter): event-driven admission wakeup replacing 2s poll tick (#1864 A3)

tests/test_gpu_arbiter_wakeup.py

@@ -0,0 +1,85 @@
+"""Tests for event-driven admission wakeup — Slice A3 of taOS #1864.
+
+Verifies that releasing a reservation or completing a task immediately
+wakes the drain loop (asyncio.Event path) instead of waiting for the
+full poll tick.  The 2 s sleep is demoted to fallback timeout only.
+"""
+
+import asyncio
+import time
+
+import pytest
+
+from tinyagentos.scheduler.gpu_arbiter import GpuArbiter
+from tinyagentos.scheduler.types import Capability, Priority, Task
+from tinyagentos.vram_reservation import VramReservationManager
+
+
+@pytest.mark.asyncio
+async def test_release_triggers_immediate_drain():
+    """A reservation release immediately wakes the drain loop.
+
+    The fallback tick is set absurdly long (60 s) so the test can only
+    pass through the wake-event path.  When the first task releases its
+    VRAM reservation, the second task (queued) should admit in < 0.5 s.
+    """
+    mgr = VramReservationManager(probe=lambda: (1024, 16384))
+    arbiter = GpuArbiter(vram_reservation=mgr, drain_tick_seconds=60.0)
+    await arbiter.start()
+    try:
+        release_first = asyncio.Event()
+
+        async def hold(_res):
+            await release_first.wait()
+            return "first"
+
+        async def fast(_res):
+            return "second"
+
+        t1 = Task(
+            capability=Capability.LLM_CHAT, payload=hold,
+            preferred_resources=[], priority=Priority.BACKGROUND, submitter="a",
+        )
+        t2 = Task(
+            capability=Capability.LLM_CHAT, payload=fast,
+            preferred_resources=[], priority=Priority.BACKGROUND, submitter="b",
+        )
+        f1 = asyncio.ensure_future(arbiter.submit_gpu(t1, required_vram_mb=1024))
+        await asyncio.sleep(0.05)
+        f2 = asyncio.ensure_future(arbiter.submit_gpu(t2, required_vram_mb=1024))
+        await asyncio.sleep(0.05)               # t2 is queued (VRAM exhausted)
+        start = time.monotonic()
+        release_first.set()                     # t1 finishes, releases 1024 MiB
+        assert await asyncio.wait_for(f2, timeout=1.0) == "second"
+        assert time.monotonic() - start < 0.5   # event path, not the 60 s tick
+        assert await f1 == "first"
+    finally:
+        await arbiter.stop()
+
+
+@pytest.mark.asyncio
+async def test_poll_tick_still_drains_without_signal():
+    """The fallback poll tick admits work even when no event fires.
+
+    Capacity appears out of band (external process freed VRAM, probe
+    changes) — only the periodic tick can discover it.  With a short
+    drain_tick_seconds, admission happens within 1 s.
+    """
+    free = {"mb": 0}
+    mgr = VramReservationManager(probe=lambda: (free["mb"], 16384))
+    arbiter = GpuArbiter(vram_reservation=mgr, drain_tick_seconds=0.1)
+    await arbiter.start()
+    try:
+        async def fast(_res):
+            return "ok"
+
+        t = Task(
+            capability=Capability.LLM_CHAT, payload=fast,
+            preferred_resources=[], priority=Priority.BACKGROUND, submitter="a",
+        )
+        f = asyncio.ensure_future(arbiter.submit_gpu(t, required_vram_mb=1024))
+        await asyncio.sleep(0.05)
+        free["mb"] = 8192                       # external process freed VRAM
+        assert await asyncio.wait_for(f, timeout=1.0) == "ok"
+    finally:
+        await arbiter.stop()

tinyagentos/scheduler/gpu_arbiter.py

@@ -86,6 +86,7 @@ def __init__(
         vram_reservation: VramReservationManager | None = None,
         max_queue_size: int = 100,
         eviction_enabled: bool = True,
+        drain_tick_seconds: float = 2.0,
     ):
         self._scheduler = scheduler
         self._cluster_manager = cluster_manager
@@ -132,6 +133,10 @@ def probe_adapter() -> tuple[int, int] | None:
         self._paused: bool = False
         self._paused_at: float | None = None
 
+        # ── taOS #1864 A3: event-driven wakeup ───────────────────────────
+        self._drain_tick_seconds = drain_tick_seconds
+        self._wake: asyncio.Event = asyncio.Event()
+
     async def start(self) -> None:
         if self._queue_processor_task is not None:
             return
@@ -242,6 +247,7 @@ def _release_reservation(self, task_id: str) -> None:
         if reservation_id is not None:
             self._vram.release(reservation_id)
             logger.debug("gpu-arbiter: released reservation %s for task %s", reservation_id, task_id)
+            self._signal_capacity()  # taOS #1864 A3: wake drain loop immediately
 
     async def _reserve_and_check(self, task_id: str, required_vram_mb: int) -> GpuAdmission:
         """Atomically check admission and reserve VRAM if admitted.
@@ -439,6 +445,7 @@ async def _run_gpu_task(
             async with self._running_lock:
                 entry = self._running.pop(task.id, None)
                 self._running_tasks.pop(task.id, None)
+            self._signal_capacity()  # taOS #1864 A3: wake drain on completion
             # Release the lease only for normal completion (not eviction).
             # When _evict_task evicts us it pops self._running first and
             # handles the lease — our pop above returns None, so we skip.
@@ -533,11 +540,25 @@ async def _evict_task(self, task_id: str) -> int:
                      task_id, _pri, _vram, asyncio_task is not None)
         return 1
 
+    def _signal_capacity(self) -> None:
+        """Wake the drain loop on every reservation release and op completion.
+
+        Decision 7 (taOS #1864 A3): a release immediately wakes the drain so
+        a queued op admits in << 2 s instead of ≤ 2 s.  ``Event.set()`` is
+        idempotent — multiple callers may fire it without harm.
+        """
+        self._wake.set()
+
     async def _process_queue(self) -> None:
         try:
             while True:
-                await asyncio.sleep(2)
-                if not self._paused:  # taOS #796: skip drain while paused
+                try:
+                    await asyncio.wait_for(self._wake.wait(),
+                                           timeout=self._drain_tick_seconds)
+                except asyncio.TimeoutError:
+                    pass
+                self._wake.clear()
+                if not self._paused:
                     await self._drain_queue()
         except asyncio.CancelledError:
             raise