From 8377b964a59b9015fb3a285323212fb47103b3ea Mon Sep 17 00:00:00 2001 From: Hogne <227774406+hognek@users.noreply.github.com> Date: Wed, 8 Jul 2026 00:45:50 +0200 Subject: [PATCH] =?UTF-8?q?feat(scheduler):=20GPU=20arbiter=20with=20drain?= =?UTF-8?q?=E2=86=92arbiter=20wiring=20(taOS=20#1707)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Add GpuArbiter: VRAM-accounted admission, queuing, priority-based eviction - drain→arbiter coordination: arbiter notifies scheduler drain before eviction via drain_notify_fn, waits for on_drain_complete callback before proceeding - Add drain_worker/cancel_drain to ClusterManager with arbiter cross-wiring (graceful and force drain modes, lease release + task cancellation) - Exclude draining workers from capability routing - Monitor loop auto-completes graceful drains when no active leases remain - 34 new tests: 25 gpu_arbiter + 9 cluster drain - All existing cluster/lease/scheduler tests continue to pass (63 total) --- tests/test_cluster.py | 130 +++++ tests/test_gpu_arbiter.py | 405 +++++++++++++++ tinyagentos/cluster/manager.py | 145 +++++- tinyagentos/scheduler/__init__.py | 4 + tinyagentos/scheduler/gpu_arbiter.py | 704 +++++++++++++++++++++++++++ 5 files changed, 1387 insertions(+), 1 deletion(-) create mode 100644 tests/test_gpu_arbiter.py create mode 100644 tinyagentos/scheduler/gpu_arbiter.py diff --git a/tests/test_cluster.py b/tests/test_cluster.py index 74875d441..1eb6b8a6e 100644 --- a/tests/test_cluster.py +++ b/tests/test_cluster.py @@ -259,3 +259,133 @@ async def test_router_returns_none_for_no_workers(self): data, worker_name = await router.route_request("chat", "POST", "/v1/chat/completions", {}) assert data is None assert worker_name is None + + +# --------------------------------------------------------------------------- +# ClusterManager — drain_worker / cancel_drain (taOS #1707) +# --------------------------------------------------------------------------- + +@pytest.mark.asyncio +class TestClusterManagerDrain: + async def test_drain_worker_graceful_marks_draining(self): + mgr = ClusterManager() + w = _make_worker("gpu-box") + await mgr.register_worker(w) + + result = await mgr.drain_worker("gpu-box", graceful=True) + assert result["worker"] == "gpu-box" + assert result["previous_status"] == "online" + assert result["status"] == "draining" + assert result["released_leases"] == 0 + + worker = mgr.get_worker("gpu-box") + assert worker.status == "draining" + + async def test_drain_worker_not_found(self): + mgr = ClusterManager() + result = await mgr.drain_worker("nonexistent") + assert result["error"] == "worker not found" + + async def test_drain_worker_force_releases_leases(self): + mgr = ClusterManager() + w = _make_worker("gpu-box") + await mgr.register_worker(w) + + # Create a lease for this worker + lease = await mgr.claim_lease( + resource_id="gpu-box:gpu-cuda-0", + caller="test", + ttl_seconds=300, + required_vram_mb=1024, + ) + assert lease is not None + + result = await mgr.drain_worker("gpu-box", graceful=False) + assert result["worker"] == "gpu-box" + assert result["status"] == "offline" + assert result["released_leases"] == 1 + + worker = mgr.get_worker("gpu-box") + assert worker.status == "offline" + + # Lease should be gone + leases = mgr.get_leases() + assert len(leases) == 0 + + async def test_cancel_drain_puts_worker_back_online(self): + mgr = ClusterManager() + w = _make_worker("gpu-box") + await mgr.register_worker(w) + + await mgr.drain_worker("gpu-box", graceful=True) + assert mgr.get_worker("gpu-box").status == "draining" + + result = await mgr.cancel_drain("gpu-box") + assert result["status"] == "online" + assert mgr.get_worker("gpu-box").status == "online" + + async def test_cancel_drain_not_draining(self): + mgr = ClusterManager() + w = _make_worker("gpu-box") + await mgr.register_worker(w) + + result = await mgr.cancel_drain("gpu-box") + assert "error" in result + assert "not draining" in result["error"] + + async def test_cancel_drain_not_found(self): + mgr = ClusterManager() + result = await mgr.cancel_drain("nonexistent") + assert result["error"] == "worker not found" + + async def test_draining_workers_excluded_from_routing(self): + mgr = ClusterManager() + await mgr.register_worker( + _make_worker("online-w", capabilities=["chat"], load=0.1) + ) + await mgr.register_worker( + _make_worker("draining-w", capabilities=["chat"], load=0.1) + ) + await mgr.drain_worker("draining-w", graceful=True) + + workers = mgr.get_workers_for_capability("chat") + worker_names = [w.name for w in workers] + assert "online-w" in worker_names + assert "draining-w" not in worker_names + + async def test_set_gpu_arbiter_cross_wiring(self): + mgr = ClusterManager() + # set_gpu_arbiter is called by application wiring code; + # verify it doesn't crash and stores the reference. + mock_arbiter = MagicMock() + mgr.set_gpu_arbiter(mock_arbiter) + assert mgr._gpu_arbiter is mock_arbiter + + async def test_drain_force_calls_arbiter_cancel(self): + from unittest.mock import AsyncMock + + mgr = ClusterManager() + w = _make_worker("gpu-box") + await mgr.register_worker(w) + + mock_arbiter = MagicMock() + mock_arbiter.cancel_running_for_leases = AsyncMock( + return_value=(2, 0) + ) + mgr.set_gpu_arbiter(mock_arbiter) + + # Create leases + await mgr.claim_lease( + resource_id="gpu-box:gpu-cuda-0", caller="test", + ttl_seconds=300, required_vram_mb=1024, + ) + await mgr.claim_lease( + resource_id="gpu-box:gpu-cuda-1", caller="test", + ttl_seconds=300, required_vram_mb=2048, + ) + + result = await mgr.drain_worker("gpu-box", graceful=False) + assert result["released_leases"] == 2 + mock_arbiter.cancel_running_for_leases.assert_called_once() + call_args = mock_arbiter.cancel_running_for_leases.call_args[0][0] + assert len(call_args) == 2 # two lease IDs diff --git a/tests/test_gpu_arbiter.py b/tests/test_gpu_arbiter.py new file mode 100644 index 000000000..ee4a97e57 --- /dev/null +++ b/tests/test_gpu_arbiter.py @@ -0,0 +1,405 @@ +"""Tests for the GPU arbiter — admission, queuing, eviction, and +drain→arbiter wiring (taOS #1707).""" +from __future__ import annotations + +import asyncio +from unittest.mock import AsyncMock + +import pytest + +from tinyagentos.scheduler.gpu_arbiter import ( + GpuAdmission, + GpuArbiter, + _DrainState, +) +from tinyagentos.scheduler.types import ( + Capability, + NoResourceAvailableError, + Priority, + ResourceRef, + Task, +) + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + +def _make_task( + task_id: str = "t1", + priority: Priority = Priority.INTERACTIVE_AGENT, +) -> Task: + return Task( + id=task_id, + capability=Capability.LLM_CHAT, + priority=priority, + submitter="test", + payload=AsyncMock(return_value="ok"), + preferred_resources=[], + ) + + +def _vram_probe(free_mb: int, total_mb: int = 8192): + """Return a probe that always returns (free_mb, total_mb).""" + def _probe() -> tuple[int, int]: + return free_mb, total_mb + return _probe + + +# --------------------------------------------------------------------------- +# GpuAdmission +# --------------------------------------------------------------------------- + +class TestGpuAdmission: + def test_admitted_defaults(self): + a = GpuAdmission(admitted=True) + assert a.admitted is True + assert a.reason is None + assert a.free_vram_mb == 0 + + def test_rejected_with_reason(self): + a = GpuAdmission( + admitted=False, free_vram_mb=1024, required_vram_mb=4096, + reason="insufficient VRAM", + ) + assert a.admitted is False + assert a.reason == "insufficient VRAM" + + +# --------------------------------------------------------------------------- +# GpuArbiter — admission +# --------------------------------------------------------------------------- + +class TestGpuArbiterAdmission: + @pytest.mark.asyncio + async def test_zero_vram_always_admitted(self): + arbiter = GpuArbiter(vram_probe=_vram_probe(0)) + task = _make_task() + result = await arbiter.submit_gpu(task, required_vram_mb=0) + assert result == "ok" + + @pytest.mark.asyncio + async def test_sufficient_vram_admitted(self): + arbiter = GpuArbiter(vram_probe=_vram_probe(4096)) + task = _make_task() + result = await arbiter.submit_gpu(task, required_vram_mb=2048) + assert result == "ok" + + @pytest.mark.asyncio + async def test_insufficient_vram_queued(self): + arbiter = GpuArbiter(vram_probe=_vram_probe(1024), max_queue_size=10) + task = _make_task() + # VRAM is insufficient — task will queue. The _arbiter_future + # won't resolve because no queue processor is running. + with pytest.raises(asyncio.TimeoutError): + await asyncio.wait_for( + arbiter.submit_gpu(task, required_vram_mb=4096), + timeout=0.5, + ) + stats = await arbiter.stats() + assert stats["submitted"] == 1 + assert stats["queued"] == 1 + + @pytest.mark.asyncio + async def test_queue_full_raises(self): + arbiter = GpuArbiter(vram_probe=_vram_probe(1024), max_queue_size=1) + t1 = _make_task("t1") + with pytest.raises(asyncio.TimeoutError): + await asyncio.wait_for( + arbiter.submit_gpu(t1, required_vram_mb=4096), + timeout=0.3, + ) + t2 = _make_task("t2") + with pytest.raises(NoResourceAvailableError, match="queue full"): + await arbiter.submit_gpu(t2, required_vram_mb=4096) + + +# --------------------------------------------------------------------------- +# GpuArbiter — eviction +# --------------------------------------------------------------------------- + +class TestGpuArbiterEviction: + @pytest.mark.asyncio + async def test_evict_lowest_priority(self): + arbiter = GpuArbiter(vram_probe=_vram_probe(8192)) + + # Manually seed a running task (bypass admission). + # Priority.BACKGROUND = 30 (higher value = lower priority = evicted first) + task = _make_task("t-victim", priority=Priority.BACKGROUND) + async with arbiter._running_lock: + arbiter._running[task.id] = (task, None, int(task.priority), 2048) + + evicted = await arbiter.evict_lowest_priority() + assert evicted == 1 + stats = await arbiter.stats() + assert stats["evicted"] == 1 + assert stats["running"] == 0 + + @pytest.mark.asyncio + async def test_evict_respects_min_priority(self): + arbiter = GpuArbiter(vram_probe=_vram_probe(8192)) + + # Priority.INTERACTIVE_USER = 10 (low number = high priority) + task = _make_task("t-important", priority=Priority.INTERACTIVE_USER) + async with arbiter._running_lock: + arbiter._running[task.id] = (task, None, int(task.priority), 2048) + + # min_priority=20 means only evict tasks with pri >= 20. + # INTERACTIVE_USER is 10, so it should be skipped. + evicted = await arbiter.evict_lowest_priority(min_priority=20) + assert evicted == 0 + stats = await arbiter.stats() + assert stats["running"] == 1 + + @pytest.mark.asyncio + async def test_evict_highest_priority_value_wins(self): + """When multiple tasks run, the one with highest numeric priority + (lowest actual priority) gets evicted.""" + arbiter = GpuArbiter(vram_probe=_vram_probe(8192)) + + t_low = _make_task("t-low", priority=Priority.BATCH) # pri=40 + t_med = _make_task("t-med", priority=Priority.BACKGROUND) # pri=30 + t_high = _make_task("t-high", priority=Priority.INTERACTIVE_USER) # pri=10 + + async with arbiter._running_lock: + arbiter._running["t-low"] = (t_low, None, 40, 1024) + arbiter._running["t-med"] = (t_med, None, 30, 2048) + arbiter._running["t-high"] = (t_high, None, 10, 512) + + evicted = await arbiter.evict_lowest_priority() + assert evicted == 1 + # t-low (pri=40) should be evicted + stats = await arbiter.stats() + assert stats["evicted"] == 1 + assert stats["running"] == 2 + + @pytest.mark.asyncio + async def test_evict_not_found(self): + arbiter = GpuArbiter() + evicted = await arbiter.evict_lowest_priority() + assert evicted == 0 + + @pytest.mark.asyncio + async def test_eviction_disabled(self): + arbiter = GpuArbiter(eviction_enabled=False) + task = _make_task() + async with arbiter._running_lock: + arbiter._running[task.id] = (task, None, int(task.priority), 2048) + evicted = await arbiter.evict_lowest_priority() + assert evicted == 0 + + +# --------------------------------------------------------------------------- +# GpuArbiter — drain→arbiter wiring (taOS #1707) +# --------------------------------------------------------------------------- + +class TestGpuArbiterDrainWiring: + @pytest.mark.asyncio + async def test_on_drain_complete_returns_true(self): + arbiter = GpuArbiter() + state = _DrainState(model_id="m1") + arbiter._draining["m1"] = state + result = await arbiter.on_drain_complete("m1") + assert result is True + assert state.event.is_set() + assert "m1" not in arbiter._draining + + @pytest.mark.asyncio + async def test_on_drain_complete_unknown_model(self): + arbiter = GpuArbiter() + result = await arbiter.on_drain_complete("unknown") + assert result is False + + @pytest.mark.asyncio + async def test_notify_drain_and_wait_no_fn_proceeds(self): + """Without drain_notify_fn, _notify_drain_and_wait returns True.""" + arbiter = GpuArbiter(drain_notify_fn=None) + result = await arbiter._notify_drain_and_wait("m1") + assert result is True + + @pytest.mark.asyncio + async def test_notify_drain_and_wait_with_fn(self): + """With drain_notify_fn, waits for on_drain_complete.""" + drain_calls: list[str] = [] + + async def drain_fn(model_id: str) -> None: + drain_calls.append(model_id) + + arbiter = GpuArbiter(drain_notify_fn=drain_fn, drain_timeout=5.0) + + async def do_drain() -> bool: + return await arbiter._notify_drain_and_wait("m1") + + drain_task = asyncio.create_task(do_drain()) + await asyncio.sleep(0.05) + + assert drain_calls == ["m1"] + + result = await arbiter.on_drain_complete("m1") + assert result is True + + drain_result = await asyncio.wait_for(drain_task, timeout=1.0) + assert drain_result is True + + @pytest.mark.asyncio + async def test_notify_drain_timeout(self): + """When drain doesn't complete within timeout, returns False.""" + async def slow_drain(model_id: str) -> None: + pass # never calls on_drain_complete + + arbiter = GpuArbiter(drain_notify_fn=slow_drain, drain_timeout=0.1) + result = await arbiter._notify_drain_and_wait("m1") + assert result is False + + @pytest.mark.asyncio + async def test_evict_task_notifies_drain(self): + """_evict_task calls _notify_drain_and_wait before cancelling.""" + drain_calls: list[str] = [] + + async def drain_fn(model_id: str) -> None: + drain_calls.append(model_id) + # Complete drain asynchronously + asyncio.get_running_loop().create_task( + _complete_drain(arbiter, model_id) + ) + + arbiter = GpuArbiter( + drain_notify_fn=drain_fn, drain_timeout=5.0, + vram_probe=_vram_probe(8192), + ) + + task = _make_task("t-drain-victim") + # Store model_id as an attribute (real tasks won't have this; + # the arbiter uses getattr fallback) + task.model_id = "my-model" # type: ignore[attr-defined] + async with arbiter._running_lock: + arbiter._running[task.id] = (task, None, int(task.priority), 2048) + + evicted = await arbiter._evict_task(task.id) + assert evicted == 1 + assert drain_calls == ["my-model"] + stats = await arbiter.stats() + assert stats["evicted"] == 1 + + @pytest.mark.asyncio + async def test_evict_task_falls_back_to_task_id(self): + """When task has no model_id, uses task.id for drain notification.""" + drain_calls: list[str] = [] + + async def drain_fn(model_id: str) -> None: + drain_calls.append(model_id) + asyncio.get_running_loop().create_task( + _complete_drain(arbiter, model_id) + ) + + arbiter = GpuArbiter(drain_notify_fn=drain_fn, drain_timeout=5.0) + + task = _make_task("t-fallback-id") + # No model_id attribute — falls back to task.id + async with arbiter._running_lock: + arbiter._running[task.id] = (task, None, int(task.priority), 2048) + + evicted = await arbiter._evict_task(task.id) + assert evicted == 1 + assert drain_calls == ["t-fallback-id"] + + +async def _complete_drain(arbiter: GpuArbiter, model_id: str) -> None: + """Helper to complete a drain after a short delay.""" + await asyncio.sleep(0.01) + await arbiter.on_drain_complete(model_id) + + +# --------------------------------------------------------------------------- +# GpuArbiter — cancel_running_for_leases +# --------------------------------------------------------------------------- + +class TestGpuArbiterLeaseCancel: + @pytest.mark.asyncio + async def test_cancel_by_lease_ids(self): + arbiter = GpuArbiter() + + t1 = _make_task("t1") + t2 = _make_task("t2") + t3 = _make_task("t3") + + async with arbiter._running_lock: + arbiter._running["t1"] = (t1, "lease-a", 0, 1024) + arbiter._running["t2"] = (t2, "lease-b", 0, 2048) + arbiter._running["t3"] = (t3, "lease-c", 0, 512) + + cancelled, done = await arbiter.cancel_running_for_leases( + {"lease-a", "lease-c"} + ) + assert cancelled == 2 + assert done == 0 + stats = await arbiter.stats() + assert stats["running"] == 1 + + @pytest.mark.asyncio + async def test_cancel_empty_leases(self): + arbiter = GpuArbiter() + cancelled, done = await arbiter.cancel_running_for_leases(set()) + assert cancelled == 0 + assert done == 0 + + @pytest.mark.asyncio + async def test_cancel_no_match(self): + arbiter = GpuArbiter() + t1 = _make_task("t1") + async with arbiter._running_lock: + arbiter._running["t1"] = (t1, "lease-a", 0, 1024) + cancelled, done = await arbiter.cancel_running_for_leases({"lease-x"}) + assert cancelled == 0 + assert done == 0 + + +# --------------------------------------------------------------------------- +# GpuArbiter — pause / resume +# --------------------------------------------------------------------------- + +class TestGpuArbiterPause: + def test_pause_resume(self): + arbiter = GpuArbiter() + assert arbiter.paused is False + + result = arbiter.pause() + assert result is True + assert arbiter.paused is True + + result = arbiter.pause() + assert result is False + + result = arbiter.resume() + assert result is True + assert arbiter.paused is False + + result = arbiter.resume() + assert result is False + + +# --------------------------------------------------------------------------- +# GpuArbiter — stats and observability +# --------------------------------------------------------------------------- + +class TestGpuArbiterStats: + @pytest.mark.asyncio + async def test_initial_stats(self): + arbiter = GpuArbiter() + stats = await arbiter.stats() + assert stats["submitted"] == 0 + assert stats["admitted"] == 0 + assert stats["evicted"] == 0 + assert stats["running"] == 0 + assert stats["active_drains"] == 0 + + @pytest.mark.asyncio + async def test_running_tasks_empty(self): + arbiter = GpuArbiter() + tasks = await arbiter.running_tasks() + assert tasks == [] + + def test_queue_snapshot_empty(self): + arbiter = GpuArbiter() + snap = arbiter.queue_snapshot() + assert snap == [] diff --git a/tinyagentos/cluster/manager.py b/tinyagentos/cluster/manager.py index 69df08323..68f1f8f5d 100644 --- a/tinyagentos/cluster/manager.py +++ b/tinyagentos/cluster/manager.py @@ -252,6 +252,132 @@ def get_workers(self) -> list[WorkerInfo]: def get_worker(self, name: str) -> WorkerInfo | None: return self._workers.get(name) + # ── GPU arbiter cross-wiring ─────────────────────────────────────── + + def set_gpu_arbiter(self, arbiter) -> None: + """Wire a :class:`GpuArbiter` into the cluster manager so drain + operations can cancel running GPU tasks for evicted workers.""" + self._gpu_arbiter = arbiter + + # ── Worker drain (taOS #1707) ────────────────────────────────────── + + async def drain_worker(self, name: str, graceful: bool = True) -> dict: + """Begin draining a worker — gracefully detach without dropping tasks. + + When ``graceful=True`` (default): + 1. Worker enters ``\"draining\"`` — no new tasks routed to it. + 2. Existing leases run to completion (TTL not touched). + 3. The monitor loop sweeps expired leases and marks the worker + ``\"offline\"`` once all leases are released. + + When ``graceful=False``: + 1. Worker enters ``\"draining\"``. + 2. All active leases for this worker are released immediately. + 3. Running GPU arbiter tasks for those leases are cancelled. + 4. Worker is marked ``\"offline\"`` immediately. + + Returns a dict with ``worker``, ``previous_status``, + ``released_leases``, and ``status``. + """ + worker = self._workers.get(name) + if worker is None: + return {"worker": name, "error": "worker not found"} + prev_status = worker.status + worker.status = "draining" + logger.info( + "Worker '%s' entering drain (was %s, graceful=%s)", + name, prev_status, graceful, + ) + + released = 0 + + if not graceful: + async with self._lease_lock: + lids: list[str] = [ + lid for lid, lease in self._leases.items() + if (parsed := self._parse_resource_id(lease.resource_id)) + and parsed[0] == name + ] + for lid in lids: + self._leases.pop(lid, None) + released += 1 + worker.status = "offline" + logger.info( + "Worker '%s' drain: force-released %d leases", name, released, + ) + + # Cancel running GPU arbiter tasks for the released leases + # so eviction isn't a no-op (taOS #1707 cross-wiring). + if released > 0 and getattr(self, "_gpu_arbiter", None) is not None: + try: + cancelled, already_done = ( + await self._gpu_arbiter.cancel_running_for_leases(set(lids)) + ) + logger.info( + "Worker '%s' drain: arbiter cancelled %d tasks, %d already done", + name, cancelled, already_done, + ) + except Exception: + logger.exception( + "gpu-arbiter: cancel for drain of '%s' failed", name, + ) + + if self._notifications: + detail = ( + f"Worker '{name}' is draining. " + f"{'Tasks will complete before detach.' if graceful else 'All leases released immediately.'}" + ) + try: + asyncio.get_running_loop().create_task( + self._notifications.emit_event( + "worker.drain", + f"Worker '{name}' draining", + detail, + level="info", + ) + ) + except RuntimeError: + pass + + return { + "worker": name, + "previous_status": prev_status, + "status": worker.status, + "released_leases": released, + } + + async def cancel_drain(self, name: str) -> dict: + """Cancel an in-progress drain and return the worker to online. + + Only works when the worker is in ``\"draining\"`` status and has + not yet been marked offline. + """ + worker = self._workers.get(name) + if worker is None: + return {"worker": name, "error": "worker not found"} + if worker.status != "draining": + return { + "worker": name, + "error": f"worker is not draining (status={worker.status})", + } + worker.status = "online" + logger.info("Worker '%s' drain cancelled — back online", name) + + if self._notifications: + try: + asyncio.get_running_loop().create_task( + self._notifications.emit_event( + "worker.online", + f"Worker '{name}' drain cancelled", + f"Worker '{name}' is back online after drain was cancelled.", + level="info", + ) + ) + except RuntimeError: + pass + + return {"worker": name, "status": "online"} + def find_worker_by_host_lan_ip(self, host_lan_ip: str) -> WorkerInfo | None: """Return the worker registered for this bare host's LAN IP, or None. @@ -396,7 +522,9 @@ def _sweep_expired_leases(self): # ── Capability routing ───────────────────────────────────────────── def get_workers_for_capability(self, capability: str) -> list[WorkerInfo]: - """Get online workers that support a capability, sorted by priority (lowest load first).""" + """Get online workers that support a capability, sorted by priority (lowest load first). + + Draining workers are excluded (taOS #1707).""" eligible = [ w for w in self._workers.values() if w.status == "online" and capability in w.capabilities @@ -516,6 +644,21 @@ async def _monitor_loop(self): for lid in offline_lids: self._leases.pop(lid, None) logger.debug("Lease %s released — worker %s went offline", lid, worker.name) + # taOS #1707: gracefully drained workers with no remaining + # leases can now be marked fully offline. + if worker.status == "draining": + active_leases = any( + (parsed := self._parse_resource_id(lease.resource_id)) + and parsed[0] == worker.name + for lease in self._leases.values() + if lease.expires_at > now + ) + if not active_leases: + worker.status = "offline" + logger.info( + "Worker '%s' drain complete — marked offline (no active leases)", + worker.name, + ) async with self._lease_lock: self._sweep_expired_leases() await asyncio.sleep(5) diff --git a/tinyagentos/scheduler/__init__.py b/tinyagentos/scheduler/__init__.py index 187914bef..f7c97e6ad 100644 --- a/tinyagentos/scheduler/__init__.py +++ b/tinyagentos/scheduler/__init__.py @@ -32,6 +32,8 @@ _LAZY_EXPORTS = { "BackendCatalog": "backend_catalog", "BackendEntry": "backend_catalog", + "GpuArbiter": "gpu_arbiter", + "GpuAdmission": "gpu_arbiter", "HistoryStore": "history_store", "Resource": "resource", "Scheduler": "scheduler", @@ -58,6 +60,8 @@ def __dir__(): "BackendCatalog", "BackendEntry", "Capability", + "GpuAdmission", + "GpuArbiter", "HistoryStore", "NoResourceAvailableError", "Priority", diff --git a/tinyagentos/scheduler/gpu_arbiter.py b/tinyagentos/scheduler/gpu_arbiter.py new file mode 100644 index 000000000..dc81eac9f --- /dev/null +++ b/tinyagentos/scheduler/gpu_arbiter.py @@ -0,0 +1,704 @@ +"""GPU Arbiter — VRAM-accounted admission + queue + eviction for GPU workloads. + +Provides admission control, queuing, and priority-based eviction for GPU +inference tasks. Coordinates with the scheduler drain mechanism: when +the arbiter decides to evict a model, it notifies the scheduler drain so +it stops routing new work, then waits for a drain-complete callback +before proceeding with the actual eviction. + +taOS #1707: drain→arbiter wiring. +""" +from __future__ import annotations + +import asyncio +import logging +import time +from dataclasses import dataclass, field +from typing import Callable, Optional + +from tinyagentos.scheduler.types import ( + NoResourceAvailableError, + Task, +) + +logger = logging.getLogger(__name__) + + +# --------------------------------------------------------------------------- +# VRAM probe helpers +# --------------------------------------------------------------------------- + +def _default_vram_probe() -> tuple[int, int]: + """No-GPU fallback: returns (0, 0).""" + return 0, 0 + + +def _probe_nvidia_vram() -> tuple[int, int]: + """Probe free/total VRAM from nvidia-smi. Returns (free_mb, total_mb).""" + try: + import subprocess + free_raw = subprocess.run( + ["nvidia-smi", "--query-gpu=memory.free", "--format=csv,noheader,nounits"], + capture_output=True, text=True, timeout=5, + ) + total_raw = subprocess.run( + ["nvidia-smi", "--query-gpu=memory.total", "--format=csv,noheader,nounits"], + capture_output=True, text=True, timeout=5, + ) + return ( + int(free_raw.stdout.strip().split("\n")[0]), + int(total_raw.stdout.strip().split("\n")[0]), + ) + except Exception: + return 0, 0 + + +# --------------------------------------------------------------------------- +# Internal queue entry +# --------------------------------------------------------------------------- + +@dataclass(order=True) +class _QueuedGpuTask: + """Internal queue entry, ordered by (priority, seq).""" + priority: int + seq: int + task: Task = field(compare=False) + required_vram_mb: int = field(compare=False) + 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) + + +# --------------------------------------------------------------------------- +# Admission result +# --------------------------------------------------------------------------- + +@dataclass +class GpuAdmission: + """Result of a GPU admission check.""" + admitted: bool + reason: str | None = None + free_vram_mb: int = 0 + required_vram_mb: int = 0 + + +# --------------------------------------------------------------------------- +# Drain state tracking +# --------------------------------------------------------------------------- + +@dataclass +class _DrainState: + """Tracks an active drain for a model being evicted. + + The arbiter sets this when it decides to evict and calls + drain_notify_fn. The scheduler drain path calls + arbiter.on_drain_complete(model_id) when draining is finished, + which signals the event so the arbiter can proceed with eviction. + """ + model_id: str + event: asyncio.Event = field(default_factory=asyncio.Event) + started_at: float = field(default_factory=time.time) + + +# --------------------------------------------------------------------------- +# GPU Arbiter +# --------------------------------------------------------------------------- + +class GpuArbiter: + """VRAM-accounted admission control with scheduler drain coordination. + + Layered on top of the Scheduler. Provides admission control, queuing, + and priority-based eviction. Coordinates with the scheduler drain + mechanism via drain_notify_fn and on_drain_complete. + + Usage:: + + arbiter = GpuArbiter( + scheduler=sched, + cluster_manager=cm, + drain_notify_fn=sched.drain_model, + ) + await arbiter.start() + result = await arbiter.submit_gpu(task, required_vram_mb=4096) + await arbiter.stop() + """ + + def __init__( + self, + scheduler=None, + cluster_manager=None, + vram_probe: Callable[[], tuple[int, int]] | None = None, + max_queue_size: int = 100, + eviction_enabled: bool = True, + drain_notify_fn: Optional[Callable[[str], object]] = None, + drain_timeout: float = 60.0, + ): + self._scheduler = scheduler + self._cluster_manager = cluster_manager + self._vram_probe = vram_probe or _default_vram_probe + self._max_queue_size = max_queue_size + self._eviction_enabled = eviction_enabled + + # ── drain→arbiter wiring (taOS #1707) ────────────────────────── + # drain_notify_fn(model_id) is called when the arbiter decides to + # evict a model. The scheduler should stop routing new work to the + # model and call arbiter.on_drain_complete(model_id) when done. + self._drain_notify_fn = drain_notify_fn + self._drain_timeout = drain_timeout + # model_id → _DrainState for active drains + self._draining: dict[str, _DrainState] = {} + + self._queue: asyncio.PriorityQueue[_QueuedGpuTask] = ( + asyncio.PriorityQueue(maxsize=max_queue_size) + ) + self._seq = 0 + self._running: dict[str, tuple[Task, str | None, int, int]] = {} + self._running_tasks: dict[str, asyncio.Task] = {} + self._running_lock = asyncio.Lock() + + # TOCTOU-race fix: in-flight VRAM reservation tracking + self._reserved_vram_mb: int = 0 + self._pending_reservations: dict[str, int] = {} + self._reservation_lock = asyncio.Lock() + + self._queue_processor_task: asyncio.Task | None = None + self._submitted = 0 + self._admitted = 0 + self._queued = 0 + self._evicted = 0 + self._dropped = 0 + self._paused: bool = False + self._paused_at: float | None = None + + # ── Lifecycle ────────────────────────────────────────────────────── + + async def start(self) -> None: + if self._queue_processor_task is not None: + return + self._paused = False + self._queue_processor_task = asyncio.create_task( + self._process_queue(), name="gpu-arbiter-queue" + ) + + async def stop(self) -> None: + if self._queue_processor_task is not None: + self._queue_processor_task.cancel() + try: + await self._queue_processor_task + except asyncio.CancelledError: + pass + self._queue_processor_task = None + + # ── Pause / resume ───────────────────────────────────────────────── + + @property + def paused(self) -> bool: + return self._paused + + def pause(self) -> bool: + if self._paused: + return False + self._paused = True + self._paused_at = time.time() + logger.info( + "gpu-arbiter: queue processing paused (queue_depth=%d, running=%d)", + self._queue.qsize(), len(self._running), + ) + return True + + def resume(self) -> bool: + if not self._paused: + return False + self._paused = False + paused_for = time.time() - (self._paused_at or time.time()) + logger.info( + "gpu-arbiter: queue processing resumed (was paused for %.1fs, queue_depth=%d)", + paused_for, self._queue.qsize(), + ) + self._paused_at = None + return True + + # ── drain→arbiter wiring (taOS #1707) ────────────────────────────── + + async def on_drain_complete(self, model_id: str) -> bool: + """Called by the scheduler drain path when draining is finished. + + Signals the waiting arbiter that it can now proceed with eviction + for *model_id*. Returns False if no drain was in progress for + this model (stale / duplicate callback). + """ + state = self._draining.pop(model_id, None) + if state is None: + logger.debug( + "gpu-arbiter: on_drain_complete for %r — no active drain (stale callback)", + model_id, + ) + return False + elapsed = time.time() - state.started_at + logger.info( + "gpu-arbiter: drain complete for %r (waited %.1fs)", + model_id, elapsed, + ) + state.event.set() + return True + + async def _notify_drain_and_wait(self, model_id: str) -> bool: + """Notify the scheduler to drain *model_id* and wait for completion. + + Calls drain_notify_fn(model_id) if configured, then waits for + on_drain_complete(model_id) to be called (up to _drain_timeout). + + Returns True if drain completed, False if timed out or no drain + function was configured. + """ + if self._drain_notify_fn is None: + # No drain coordination configured — proceed immediately. + return True + + # Already draining this model? Wait on the existing drain. + existing = self._draining.get(model_id) + if existing is not None: + logger.debug( + "gpu-arbiter: %r already draining, waiting on existing drain", + model_id, + ) + try: + await asyncio.wait_for( + existing.event.wait(), timeout=self._drain_timeout, + ) + return True + except asyncio.TimeoutError: + logger.warning( + "gpu-arbiter: drain timeout for %r (%.0fs), proceeding with eviction", + model_id, self._drain_timeout, + ) + self._draining.pop(model_id, None) + return False + + state = _DrainState(model_id=model_id) + self._draining[model_id] = state + + logger.info("gpu-arbiter: notifying scheduler drain for %r", model_id) + try: + result = self._drain_notify_fn(model_id) + if asyncio.iscoroutine(result): + await result + except Exception: + logger.exception( + "gpu-arbiter: drain_notify_fn failed for %r, proceeding with eviction", + model_id, + ) + self._draining.pop(model_id, None) + return False + + try: + await asyncio.wait_for( + state.event.wait(), timeout=self._drain_timeout, + ) + return True + except asyncio.TimeoutError: + logger.warning( + "gpu-arbiter: drain timeout for %r (%.0fs), proceeding with eviction anyway", + model_id, self._drain_timeout, + ) + self._draining.pop(model_id, None) + return False + + # ── VRAM reservation helpers ─────────────────────────────────────── + + def _release_reservation(self, task_id: str) -> None: + """Release an in-flight VRAM reservation for *task_id*. + + Idempotent — safe to call on a task that was never reserved. + """ + vram = self._pending_reservations.pop(task_id, None) + if vram is not None: + self._reserved_vram_mb -= vram + logger.debug( + "gpu-arbiter: released reservation %d MiB for task %s", + vram, task_id, + ) + + async def _reserve_and_check( + self, task_id: str, required_vram_mb: int, + ) -> GpuAdmission: + """Atomically check admission and reserve VRAM if admitted.""" + if required_vram_mb <= 0: + return GpuAdmission(admitted=True) + + async with self._reservation_lock: + admission = self._check_admission(required_vram_mb) + if admission.admitted: + self._reserved_vram_mb += required_vram_mb + self._pending_reservations[task_id] = required_vram_mb + logger.debug( + "gpu-arbiter: reserved %d MiB for task %s (total reserved: %d)", + required_vram_mb, task_id, self._reserved_vram_mb, + ) + return admission + + # ── Public API ───────────────────────────────────────────────────── + + 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, + ) -> object: + """Submit a GPU task with optional hardware-architecture requirements.""" + self._submitted += 1 + + if required_vram_mb > 0: + admission = await self._reserve_and_check(task.id, required_vram_mb) + if not admission.admitted: + self._release_reservation(task.id) + if self._queue.full(): + self._dropped += 1 + raise NoResourceAvailableError( + f"GPU arbiter queue full ({self._max_queue_size}), " + f"dropped task {task.id}" + ) + self._seq += 1 + entry = _QueuedGpuTask( + priority=int(task.priority), seq=self._seq, task=task, + required_vram_mb=required_vram_mb, evictable=evictable, + required_gpu_arch=required_gpu_arch, + ) + await self._queue.put(entry) + self._queued += 1 + loop = asyncio.get_running_loop() + done: asyncio.Future = loop.create_future() + entry.task._arbiter_future = done # type: ignore[attr-defined] + try: + return await done + except asyncio.CancelledError: + self._evicted += 1 + raise + return await self._run_gpu_task(task, required_vram_mb, evictable, resource_id) + + def _check_admission(self, required_vram_mb: int) -> GpuAdmission: + """Check whether *required_vram_mb* can be admitted right now. + + Subtracts in-flight reservations (_reserved_vram_mb) from the + hardware probe to close the TOCTOU window. + """ + if required_vram_mb <= 0: + return GpuAdmission(admitted=True) + + free_vram, _total = self._vram_probe() + effective_free = max(0, free_vram - self._reserved_vram_mb) + + if free_vram > 0: + if effective_free < required_vram_mb: + return GpuAdmission( + admitted=False, + free_vram_mb=effective_free, + required_vram_mb=required_vram_mb, + reason=( + f"insufficient local VRAM: need {required_vram_mb} MiB, " + f"have {effective_free} MiB available " + f"({free_vram} MiB free - {self._reserved_vram_mb} MiB reserved)" + ), + ) + return GpuAdmission( + admitted=True, free_vram_mb=effective_free, + required_vram_mb=required_vram_mb, + ) + + # No local VRAM probe — check cluster workers + if self._cluster_manager is not None: + leases = self._cluster_manager.get_leases() + for worker in self._cluster_manager.get_workers(): + if worker.status != "online": + continue + worker_leases = sum( + l.required_vram_mb for l in leases + if l.resource_id.startswith(worker.name + ":") + and l.required_vram_mb > 0 + ) + available = worker.free_vram_mb - worker_leases - self._reserved_vram_mb + if available >= required_vram_mb: + return GpuAdmission( + admitted=True, + free_vram_mb=available, + required_vram_mb=required_vram_mb, + ) + return GpuAdmission( + admitted=False, required_vram_mb=required_vram_mb, + reason=f"no cluster worker with {required_vram_mb} MiB free VRAM", + ) + + return GpuAdmission(admitted=True, required_vram_mb=required_vram_mb) + + # ── Task execution ───────────────────────────────────────────────── + + async def _run_gpu_task( + self, + task: Task, + required_vram_mb: int, + evictable: bool, + resource_id: str | None, + ) -> object: + """Execute a GPU task, holding a VRAM reservation for its duration.""" + lease_id: str | None = None + try: + if self._cluster_manager is not None and resource_id is not None: + lease = await self._cluster_manager.claim_lease( + resource_id=resource_id, caller=task.submitter, + ttl_seconds=300, required_vram_mb=required_vram_mb, + ) + if lease is None: + raise NoResourceAvailableError( + f"GPU lease claim failed for {resource_id} (task {task.id})" + ) + lease_id = lease.lease_id + current = asyncio.current_task() + async with self._running_lock: + self._running[task.id] = ( + task, lease_id, int(task.priority), required_vram_mb, + ) + if current is not None: + self._running_tasks[task.id] = current + if self._scheduler is not None: + result = await self._scheduler.submit(task) + else: + result = await task.payload(None) + self._admitted += 1 + return result + except asyncio.CancelledError: + logger.info( + "gpu-arbiter: task %s preempted via CancelledError (pri=%d, vram=%d)", + task.id, task.priority, required_vram_mb, + ) + raise + finally: + self._release_reservation(task.id) + async with self._running_lock: + entry = self._running.pop(task.id, None) + self._running_tasks.pop(task.id, None) + if entry is not None: + _task, _lid, _pri, _vram = entry + if _lid is not None and self._cluster_manager is not None: + await self._cluster_manager.release_lease(_lid) + + # ── Eviction ─────────────────────────────────────────────────────── + + async def cancel_running_for_leases(self, lease_ids: set[str]) -> tuple[int, int]: + """Cancel all running GPU tasks whose leases are in *lease_ids*. + + Called from ClusterManager drain paths when a worker's leases are + force-released. + + Returns (cancelled, already_completed). + """ + if not lease_ids: + return 0, 0 + victim_ids: list[str] = [] + async with self._running_lock: + for task_id, (_task, lease_id, _pri, _vram) in self._running.items(): + if lease_id in lease_ids: + victim_ids.append(task_id) + cancelled = 0 + already_completed = 0 + for task_id in victim_ids: + if await self._evict_task(task_id): + cancelled += 1 + else: + already_completed += 1 + if cancelled or already_completed: + logger.info( + "gpu-arbiter: drain cancelled %d, already-done %d (leases=%d)", + cancelled, already_completed, len(lease_ids), + ) + return cancelled, already_completed + + async def evict_lowest_priority(self, min_priority: int | None = None) -> int: + """Evict the running task with the highest numeric priority value. + + Higher numeric priority = lower actual priority = evicted first. + If *min_priority* is given, only tasks with priority >= min_priority + are considered. + """ + if not self._eviction_enabled: + return 0 + async with self._running_lock: + victim_id, victim_priority = None, -1 + for tid, (_t, _lid, pri, _vram) in self._running.items(): + if min_priority is not None and pri < min_priority: + continue + if pri > victim_priority: + victim_priority, victim_id = pri, tid + if victim_id is None: + return 0 + return await self._evict_task(victim_id) + + async def _evict_task(self, task_id: str) -> int: + """Evict a running task, coordinating with the scheduler drain. + + taOS #1707: Before cancelling the running task, notifies the + scheduler drain so it stops routing new work, and waits for + the drain-complete callback before proceeding. + """ + async with self._running_lock: + if task_id not in self._running: + return 0 + task, lease_id, _pri, _vram = self._running.pop(task_id) + asyncio_task = self._running_tasks.pop(task_id, None) + + # ── drain→arbiter wiring (taOS #1707) ────────────────────── + # Notify the scheduler to drain this model and wait for completion. + model_id = getattr(task, "model_id", None) or task.id + await self._notify_drain_and_wait(model_id) + + # Release VRAM reservation + self._release_reservation(task_id) + + # Release cluster lease + if lease_id is not None and self._cluster_manager is not None: + await self._cluster_manager.release_lease(lease_id) + + # Cancel the running asyncio Task + if asyncio_task is not None and not asyncio_task.done(): + asyncio_task.cancel() + + # Also cancel the arbiter future for queued tasks + future = getattr(task, "_arbiter_future", None) + if future is not None and not future.done(): + future.cancel() + + self._evicted += 1 + logger.info( + "gpu-arbiter: evicted task %s (pri=%d, vram=%d, task_cancelled=%s)", + task_id, _pri, _vram, asyncio_task is not None, + ) + return 1 + + # ── Queue processing ─────────────────────────────────────────────── + + async def _process_queue(self) -> None: + try: + while True: + await asyncio.sleep(2) + if not self._paused: + await self._drain_queue() + except asyncio.CancelledError: + raise + + async def _drain_queue(self) -> None: + """Drain the admission queue one task at a time. + + Queue processing is intentionally serial — only one queued task is + admitted per drain cycle to avoid flooding the GPU with concurrent + loads. Uses _reserve_and_check so the queue processor doesn't race + with concurrent submit_gpu calls. + """ + retry: list[_QueuedGpuTask] = [] + drained = False + while not self._queue.empty() and not drained: + entry = self._queue.get_nowait() + admission = await self._reserve_and_check( + entry.task.id, entry.required_vram_mb, + ) + if not admission.admitted: + evicted = await self.evict_lowest_priority( + min_priority=int(entry.task.priority), + ) + if evicted > 0: + admission = await self._reserve_and_check( + entry.task.id, entry.required_vram_mb, + ) + if admission.admitted: + future = getattr(entry.task, "_arbiter_future", None) + t = asyncio.create_task( + self._run_gpu_task( + entry.task, entry.required_vram_mb, + entry.evictable, None, + ), + name=f"gpu-arbiter-drain-{entry.task.id}", + ) + if future is not None: + def _propagate(ct: asyncio.Task, f: asyncio.Future = future) -> None: + if f.done(): + return + if ct.cancelled(): + return + exc = ct.exception() + if exc is not None: + f.set_exception(exc) + else: + f.set_result(ct.result()) + t.add_done_callback(_propagate) + drained = True + else: + self._release_reservation(entry.task.id) + retry.append(entry) + + for entry in retry: + if not self._queue.full(): + self._queue.put_nowait(entry) + else: + self._dropped += 1 + future = getattr(entry.task, "_arbiter_future", None) + if future is not None and not future.done(): + future.set_exception( + NoResourceAvailableError("queue full, task dropped") + ) + + # ── Observability ────────────────────────────────────────────────── + + async def stats(self) -> dict: + async with self._running_lock: + running_count = len(self._running) + return { + "submitted": self._submitted, + "admitted": self._admitted, + "queued": self._queued, + "evicted": self._evicted, + "dropped": self._dropped, + "queue_depth": self._queue.qsize(), + "running": running_count, + "max_queue_size": self._max_queue_size, + "eviction_enabled": self._eviction_enabled, + "reserved_vram_mb": self._reserved_vram_mb, + "pending_reservations": len(self._pending_reservations), + "active_drains": len(self._draining), + } + + async def running_tasks(self) -> list[dict]: + async with self._running_lock: + return [ + { + "task_id": tid, + "capability": task.capability.value, + "submitter": task.submitter, + "priority": pri, + "vram_mb": vram, + "lease_id": lid, + } + for tid, (task, lid, pri, vram) in self._running.items() + ] + + def queue_snapshot(self) -> list[dict]: + """Snapshot of queued tasks (non-destructive — re-queues after reading).""" + items: list[_QueuedGpuTask] = [] + while not self._queue.empty(): + try: + items.append(self._queue.get_nowait()) + except asyncio.QueueEmpty: + break + result = [ + { + "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 + ] + for e in items: + if not self._queue.full(): + self._queue.put_nowait(e) + return result