diff --git a/tests/test_gpu_arbiter_queue_ops.py b/tests/test_gpu_arbiter_queue_ops.py new file mode 100644 index 000000000..28426ddc7 --- /dev/null +++ b/tests/test_gpu_arbiter_queue_ops.py @@ -0,0 +1,112 @@ +"""Tests for GPU arbiter queue ops — op shape, position, cancel, snapshot (taOS #1864 A2).""" + +import asyncio +import pytest +from tinyagentos.scheduler.gpu_arbiter import GpuArbiter +from tinyagentos.scheduler.types import Capability, Priority, Task +from tinyagentos.vram_reservation import VramReservationManager + + +def _mgr(free_mb: int, total_mb: int = 16384) -> VramReservationManager: + return VramReservationManager(probe=lambda: (free_mb, total_mb)) + + +def _task(priority=Priority.BACKGROUND, submitter="t"): + async def payload(_res): + await asyncio.sleep(0.05) + return "ok" + return Task(capability=Capability.LLM_CHAT, payload=payload, + preferred_resources=[], priority=priority, submitter=submitter) + + +@pytest.mark.asyncio +async def test_submit_gpu_defaults_backward_compatible(): + arbiter = GpuArbiter(vram_reservation=_mgr(8192)) + result = await arbiter.submit_gpu(_task(), required_vram_mb=1024) + assert result == "ok" # old call shape, no new kwargs + + +@pytest.mark.asyncio +async def test_queue_position_global_for_loads(): + arbiter = GpuArbiter(vram_reservation=_mgr(0)) # everything queues + t1, t2, t3 = _task(), _task(), _task() + f1 = asyncio.ensure_future(arbiter.submit_gpu( + t1, required_vram_mb=1024, op="load", model="a", backend_name="b1")) + f2 = asyncio.ensure_future(arbiter.submit_gpu( + t2, required_vram_mb=1024, op="load", model="b", backend_name="b1")) + f3 = asyncio.ensure_future(arbiter.submit_gpu( + t3, required_vram_mb=1024, op="load", model="c", backend_name="b1")) + await asyncio.sleep(0.05) # let them enqueue + assert arbiter.queue_position(t1.id) == 1 + assert arbiter.queue_position(t2.id) == 2 + assert arbiter.queue_position(t3.id) == 3 + for f in (f1, f2, f3): + f.cancel() + + +@pytest.mark.asyncio +async def test_queue_position_per_model_for_inference(): + arbiter = GpuArbiter(vram_reservation=_mgr(0)) + ta, tb, ta2 = _task(), _task(), _task() + fs = [asyncio.ensure_future(arbiter.submit_gpu( + t, required_vram_mb=1024, op="inference", model=m, backend_name="b1")) + for t, m in ((ta, "m-a"), (tb, "m-b"), (ta2, "m-a"))] + await asyncio.sleep(0.05) + assert arbiter.queue_position(ta.id) == 1 + assert arbiter.queue_position(tb.id) == 1 # only m-b entries count + assert arbiter.queue_position(ta2.id) == 2 # behind ta on m-a + for f in fs: + f.cancel() + + +@pytest.mark.asyncio +async def test_queue_snapshot_non_destructive_and_shaped(): + arbiter = GpuArbiter(vram_reservation=_mgr(0)) + t1 = _task(submitter="pull:x") + f = asyncio.ensure_future(arbiter.submit_gpu( + t1, required_vram_mb=1024, op="load", model="qwen", backend_name="b1")) + await asyncio.sleep(0.05) + snap1 = arbiter.queue_snapshot() + snap2 = arbiter.queue_snapshot() + entry = snap1[0] + assert entry["op"] == "load" and entry["model"] == "qwen" + assert entry["backend_name"] == "b1" and entry["submitter"] == "pull:x" + assert entry["position"] == 1 + assert [e["task_id"] for e in snap1] == [e["task_id"] for e in snap2] + stats = await arbiter.stats() + assert stats["queue_depth"] == 1 # snapshot did not drain + f.cancel() + + +@pytest.mark.asyncio +async def test_cancel_queued_op_removes_and_cancels_future(): + arbiter = GpuArbiter(vram_reservation=_mgr(0)) + t1 = _task() + f = asyncio.ensure_future(arbiter.submit_gpu( + t1, required_vram_mb=1024, op="load", model="m", backend_name="b1")) + await asyncio.sleep(0.05) + assert await arbiter.cancel_op(t1.id) is True + with pytest.raises(asyncio.CancelledError): + await f + assert arbiter.queue_position(t1.id) is None + assert await arbiter.cancel_op(t1.id) is False # idempotent-ish: gone + + +@pytest.mark.asyncio +async def test_cancel_running_op_delegates_to_evict(): + mgr = _mgr(8192) + arbiter = GpuArbiter(vram_reservation=mgr) + started = asyncio.Event() + + async def payload(_res): + started.set() + await asyncio.sleep(30) + + t1 = Task(capability=Capability.LLM_CHAT, payload=payload, + preferred_resources=[], priority=Priority.BACKGROUND, submitter="t") + f = asyncio.ensure_future(arbiter.submit_gpu(t1, required_vram_mb=1024)) + await started.wait() + assert await arbiter.cancel_op(t1.id) is True + await asyncio.sleep(0.05) + assert mgr.reserved_vram_mb == 0 # reservation released + f.cancel() diff --git a/tinyagentos/scheduler/gpu_arbiter.py b/tinyagentos/scheduler/gpu_arbiter.py index a01a1bb10..dd4240210 100644 --- a/tinyagentos/scheduler/gpu_arbiter.py +++ b/tinyagentos/scheduler/gpu_arbiter.py @@ -55,6 +55,10 @@ class _QueuedGpuTask: evictable: bool = field(compare=False) required_gpu_arch: str | None = field(default=None, compare=False) queued_at: float = field(default_factory=time.time, compare=False) + op: str = field(default="inference", compare=False) + model: str | None = field(default=None, compare=False) + backend_name: str | None = field(default=None, compare=False) + resource_id: str | None = field(default=None, compare=False) @dataclass @@ -96,6 +100,8 @@ def __init__( self._running: dict[str, tuple[Task, str | None, int, int]] = {} self._running_tasks: dict[str, asyncio.Task] = {} self._running_lock = asyncio.Lock() + self._queued_entries: dict[str, _QueuedGpuTask] = {} + self._cancelled_ids: set[str] = set() # --- Single VRAM authority (taOS #185) --- # The arbiter does NOT keep its own VRAM ledger. It reserves against a @@ -305,6 +311,9 @@ async def submit_gpu( self, task: Task, required_vram_mb: int = 0, evictable: bool = False, resource_id: str | None = None, required_gpu_arch: str | None = None, + op: str = "inference", + model: str | None = None, + backend_name: str | None = None, ) -> object: """Submit a GPU task with optional hardware-architecture requirements. @@ -314,9 +323,18 @@ async def submit_gpu( evictable: Whether lower-priority tasks can be evicted for this. resource_id: Specific cluster resource to target. required_gpu_arch: CUDA compute capability required (e.g. ``"sm_86"``). + op: Operation kind — ``"inference"`` or ``"load"``. + model: Backend model name (e.g. ``"qwen2.5:7b"``). + backend_name: Config backend name (e.g. ``"local-ollama"``). """ self._submitted += 1 + # Validate op kind — only 'inference' and 'load' are known. + if op not in ("inference", "load"): + raise ValueError( + f"unknown gpu_op: {op!r} (expected 'inference' or 'load')" + ) + # Hardware architecture check (taOS #796) if required_gpu_arch: arch_ok, arch_reason = self._check_gpu_arch_compatibility( @@ -345,8 +363,11 @@ async def submit_gpu( priority=int(task.priority), seq=self._seq, task=task, required_vram_mb=required_vram_mb, evictable=evictable, required_gpu_arch=required_gpu_arch, + op=op, model=model, backend_name=backend_name, + resource_id=resource_id, ) await self._queue.put(entry) + self._queued_entries[task.id] = entry self._queued += 1 loop = asyncio.get_running_loop() done: asyncio.Future = loop.create_future() @@ -354,6 +375,8 @@ async def submit_gpu( try: return await done except asyncio.CancelledError: + self._queued_entries.pop(task.id, None) + self._cancelled_ids.add(task.id) self._evicted += 1 raise # Admitted — reservation already held by _reserve_and_check, so @@ -375,6 +398,8 @@ def _check_cluster_admission(self, required_vram_mb: int) -> GpuAdmission: for worker in self._cluster_manager.get_workers(): if worker.status != "online": continue + if worker.free_vram_mb is None: + continue # non-NVIDIA worker — no VRAM probe worker_leases = sum( l.required_vram_mb for l in leases if self._resource_on_worker(l.resource_id, worker.name) and l.required_vram_mb > 0 @@ -391,6 +416,48 @@ def _check_cluster_admission(self, required_vram_mb: int) -> GpuAdmission: reason=f"no cluster worker with {required_vram_mb} MiB free VRAM", ) + async def _renew_lease_loop( + self, lease_id: str, stop_event: asyncio.Event, + task_to_cancel: asyncio.Task | None = None, + ) -> None: + """Periodically renew a lease until *stop_event* is set. + + Renews every 200 s so the lease never expires mid-run (the claim + TTL is 300 s). When *task_to_cancel* is provided and a renewal + failure or expiry is detected the referenced task is cancelled so + it doesn't keep executing GPU work without a valid lease (taOS + #1984 — lease-loss detection). + """ + RENEW_INTERVAL = 200 + while not stop_event.is_set(): + try: + await asyncio.wait_for(stop_event.wait(), timeout=RENEW_INTERVAL) + return # stop_event set — exit cleanly + except asyncio.TimeoutError: + pass # it's time to renew + try: + if self._cluster_manager is not None: + renewed = await self._cluster_manager.renew_lease( + lease_id, ttl_seconds=300, + ) + if renewed is None: + logger.warning( + "gpu-arbiter: lease %s expired mid-renewal — " + "cancelling running task", lease_id, + ) + if task_to_cancel is not None and not task_to_cancel.done(): + task_to_cancel.cancel() + return + logger.debug("gpu-arbiter: renewed lease %s", lease_id) + except Exception: + logger.exception( + "gpu-arbiter: failed to renew lease %s — " + "cancelling running task", lease_id, + ) + if task_to_cancel is not None and not task_to_cancel.done(): + task_to_cancel.cancel() + return + async def _run_gpu_task( self, task: Task, required_vram_mb: int, evictable: bool, resource_id: str | None, ) -> object: @@ -405,6 +472,8 @@ async def _run_gpu_task( claim_lease fails (taOS #1705 — reservation leak fix). """ lease_id: str | None = None + renew_task: asyncio.Task | None = None + renew_stop: asyncio.Event | None = None try: if self._cluster_manager is not None and resource_id is not None: lease = await self._cluster_manager.claim_lease( @@ -416,6 +485,14 @@ async def _run_gpu_task( f"GPU lease claim failed for {resource_id} (task {task.id})" ) lease_id = lease.lease_id + # Start periodic lease renewal so the lease doesn't expire + # mid-run (taOS #1864 defect 4 — fixed 300 s TTL). + renew_stop = asyncio.Event() + current = asyncio.current_task() + renew_task = asyncio.create_task( + self._renew_lease_loop(lease_id, renew_stop, current), + name=f"gpu-arbiter-renew-{task.id}", + ) current = asyncio.current_task() async with self._running_lock: self._running[task.id] = (task, lease_id, int(task.priority), required_vram_mb) @@ -432,6 +509,15 @@ async def _run_gpu_task( task.id, task.priority, required_vram_mb) raise finally: + # Stop lease renewal (taOS #1864 defect 4). + if renew_stop is not None: + renew_stop.set() + if renew_task is not None: + renew_task.cancel() + try: + await renew_task + except asyncio.CancelledError: + pass # Release the VRAM reservation whether we completed, errored, # or were cancelled. _evict_task handles its own reservation # release so idempotency matters. @@ -535,10 +621,37 @@ async def _evict_task(self, task_id: str) -> int: async def _process_queue(self) -> None: try: + MAX_CONSECUTIVE_FAILURES = 10 + COOLDOWN_BACKOFF = 300 # seconds to wait before restarting the processor while True: - await asyncio.sleep(2) - if not self._paused: # taOS #796: skip drain while paused - await self._drain_queue() + consecutive_failures = 0 + while True: + await asyncio.sleep(2) + if not self._paused: # taOS #796: skip drain while paused + try: + await self._drain_queue() + consecutive_failures = 0 # reset on success + except (NoResourceAvailableError, asyncio.TimeoutError, OSError): + consecutive_failures += 1 + backoff = min(2 * (2 ** consecutive_failures), 60) + logger.exception( + "gpu-arbiter: _drain_queue raised (consecutive=%d/%d) — " + "backing off %ds", + consecutive_failures, MAX_CONSECUTIVE_FAILURES, backoff, + ) + if consecutive_failures >= MAX_CONSECUTIVE_FAILURES: + logger.critical( + "gpu-arbiter: _drain_queue failed %d consecutive " + "times — restarting queue processor after %ds cooldown", + consecutive_failures, COOLDOWN_BACKOFF, + ) + break + await asyncio.sleep(backoff) + # Outer restart loop: wait cooldown, then restart the inner loop + await asyncio.sleep(COOLDOWN_BACKOFF) + logger.warning( + "gpu-arbiter: restarting queue processor after cooldown" + ) except asyncio.CancelledError: raise @@ -559,6 +672,9 @@ async def _drain_queue(self) -> None: drained = False while not self._queue.empty() and not drained: entry = self._queue.get_nowait() + if entry.task.id in self._cancelled_ids: + self._cancelled_ids.discard(entry.task.id) + continue admission = await self._reserve_and_check(entry.task.id, entry.required_vram_mb) if not admission.admitted: # Try eviction-to-make-room for higher-priority queued tasks. @@ -570,11 +686,21 @@ async def _drain_queue(self) -> None: if evicted > 0: admission = await self._reserve_and_check(entry.task.id, entry.required_vram_mb) if admission.admitted: + # Re-check cancelled_ids after the await window above + # (reserve_and_check / eviction). cancel_op may have fired + # concurrently while we were suspended, removing the entry + # from _queued_entries and cancelling its future. Without + # this re-check the cancelled task is still admitted and + # runs to completion despite the cancellation. + if entry.task.id in self._cancelled_ids: + self._cancelled_ids.discard(entry.task.id) + self._release_reservation(entry.task.id) + continue future = getattr(entry.task, "_arbiter_future", None) # Spawn as background task so drain doesn't block and # eviction-to-make-room stays responsive on subsequent ticks. t = asyncio.create_task( - self._run_gpu_task(entry.task, entry.required_vram_mb, entry.evictable, None), + self._run_gpu_task(entry.task, entry.required_vram_mb, entry.evictable, entry.resource_id), name=f"gpu-arbiter-drain-{entry.task.id}", ) @@ -594,6 +720,7 @@ def _propagate(ct: asyncio.Task, f: asyncio.Future = future) -> None: t.add_done_callback(_propagate) + self._queued_entries.pop(entry.task.id, None) # promoted → _running drained = True else: # Not admitted — release reservation (idempotent) and retry later. @@ -601,8 +728,22 @@ def _propagate(ct: asyncio.Task, f: asyncio.Future = future) -> None: retry.append(entry) # Re-queue tasks that still can't be admitted for entry in retry: + # Skip entries cancelled during the admission-await window + # (same race: cancel_op fired while we awaited reserve_and_check). + if entry.task.id in self._cancelled_ids: + self._cancelled_ids.discard(entry.task.id) + continue if not self._queue.full(): self._queue.put_nowait(entry) + self._queued_entries[entry.task.id] = entry + # Re-check _cancelled_ids after insertion — cancel_op may have + # fired between the check above and the put_nowait+dict insert, + # which would resurrect the shadow entry and leak a queue slot. + if entry.task.id in self._cancelled_ids: + self._queued_entries.pop(entry.task.id, None) + self._cancelled_ids.discard(entry.task.id) + # The physical PriorityQueue entry can't be removed but + # _drain_queue skips cancelled entries at dequeue time. else: self._dropped += 1 future = getattr(entry.task, "_arbiter_future", None) @@ -631,20 +772,50 @@ async def running_tasks(self) -> list[dict]: for tid, (task, lid, pri, vram) in self._running.items() ] - def queue_snapshot(self) -> list[dict]: - items: list[_QueuedGpuTask] = [] - while not self._queue.empty(): - try: - items.append(self._queue.get_nowait()) - except asyncio.QueueEmpty: + def _ordered_queued(self) -> list[_QueuedGpuTask]: + return sorted(self._queued_entries.values(), key=lambda e: (e.priority, e.seq)) + + def queue_position(self, task_id: str) -> int | None: + entry = self._queued_entries.get(task_id) + if entry is None: + return None + ahead = 0 + for other in self._ordered_queued(): + if other.task.id == task_id: break - result = [ + if entry.op == "inference": + if other.model == entry.model: + ahead += 1 + else: + ahead += 1 + return ahead + 1 + + async def cancel_op(self, task_id: str) -> bool: + entry = self._queued_entries.pop(task_id, None) + if entry is not None: + self._cancelled_ids.add(task_id) + # Release the VRAM reservation immediately so capacity is freed + # rather than leaking until _drain_queue eventually skips the entry. + # _release_reservation is idempotent — safe if no reservation exists. + self._release_reservation(task_id) + future = getattr(entry.task, "_arbiter_future", None) + if future is not None and not future.done(): + future.cancel() + return True + async with self._running_lock: + running = task_id in self._running + if running: + return bool(await self._evict_task(task_id)) + return False + + def queue_snapshot(self) -> list[dict]: + now = time.time() + return [ {"task_id": e.task.id, "capability": e.task.capability.value, - "priority": e.priority, "vram_mb": e.required_vram_mb, - "queued_seconds": time.time() - e.queued_at} - for e in items + "op": e.op, "model": e.model, "backend_name": e.backend_name, + "submitter": e.task.submitter, "priority": e.priority, + "vram_mb": e.required_vram_mb, + "queued_seconds": now - e.queued_at, + "position": self.queue_position(e.task.id)} + for e in self._ordered_queued() ] - for e in items: - if not self._queue.full(): - self._queue.put_nowait(e) - return result