From 02779553ee8a7d691e2b1027c569c4de196162af Mon Sep 17 00:00:00 2001
From: Hogne <227774406+hognek@users.noreply.github.com>
Date: Sat, 18 Jul 2026 12:22:58 +0200
Subject: [PATCH 1/2] feat(gpu-arbiter): event-driven admission wakeup
 replacing 2s poll tick (#1864 A3)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Replace the hardcoded asyncio.sleep(2) in _process_queue with an
asyncio.Event-based wake path.  The drain loop now blocks on
asyncio.wait_for(self._wake.wait(), timeout=drain_tick_seconds) so a
reservation release or task completion immediately wakes the drain
loop.  The 2 s constant becomes the configurable fallback timeout
(drain_tick_seconds, default 2.0).

- __init__: new drain_tick_seconds param, self._wake Event
- _signal_capacity(): new method — calls self._wake.set()
- _process_queue: event-driven wait + clear, fallback timeout
- _release_reservation: calls _signal_capacity on actual release
- _run_gpu_task finally: calls _signal_capacity on completion
- tests/test_gpu_arbiter_wakeup.py: 2 new tests
  * test_release_triggers_immediate_drain (60 s tick, < 0.5 s admit)
  * test_poll_tick_still_drains_without_signal (0.1 s tick)
---
 tests/test_gpu_arbiter_wakeup.py     | 85 ++++++++++++++++++++++++++++
 tinyagentos/scheduler/gpu_arbiter.py | 25 +++++++-
 2 files changed, 108 insertions(+), 2 deletions(-)
 create mode 100644 tests/test_gpu_arbiter_wakeup.py

diff --git a/tests/test_gpu_arbiter_wakeup.py b/tests/test_gpu_arbiter_wakeup.py
new file mode 100644
index 000000000..a9d5a90ce
--- /dev/null
+++ b/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()
diff --git a/tinyagentos/scheduler/gpu_arbiter.py b/tinyagentos/scheduler/gpu_arbiter.py
index a01a1bb10..009d7feef 100644
--- a/tinyagentos/scheduler/gpu_arbiter.py
+++ b/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

From 6971b36523be5e6969b89901072570a97b03aebe Mon Sep 17 00:00:00 2001
From: Hogne <227774406+hognek@users.noreply.github.com>
Date: Sat, 18 Jul 2026 12:39:19 +0200
Subject: [PATCH 2/2] fix(gpu-arbiter): clamp drain_tick_seconds floor + remove
 double _signal_capacity wake
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Kilo review on #1986:
1. Clamp drain_tick_seconds to >= 0.01 in __init__ to prevent
   asyncio.wait_for(timeout=0) ValueError crash in _process_queue.
2. Remove redundant _signal_capacity() in _run_gpu_task finally —
   _release_reservation already calls it at line 250, making the
   line-448 call a double-wake on every task completion.

32/32 GPU arbiter tests pass.
---
 tinyagentos/scheduler/gpu_arbiter.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tinyagentos/scheduler/gpu_arbiter.py b/tinyagentos/scheduler/gpu_arbiter.py
index 009d7feef..89114564d 100644
--- a/tinyagentos/scheduler/gpu_arbiter.py
+++ b/tinyagentos/scheduler/gpu_arbiter.py
@@ -134,7 +134,7 @@ def probe_adapter() -> tuple[int, int] | None:
         self._paused_at: float | None = None
 
         # ── taOS #1864 A3: event-driven wakeup ───────────────────────────
-        self._drain_tick_seconds = drain_tick_seconds
+        self._drain_tick_seconds = max(drain_tick_seconds, 0.01)
         self._wake: asyncio.Event = asyncio.Event()
 
     async def start(self) -> None:
@@ -445,7 +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_reservation above calls _signal_capacity (taOS #1864 A3).
             # 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.