diff --git a/.gitignore b/.gitignore index 8c5f276c3..ab07259a1 100644 --- a/.gitignore +++ b/.gitignore @@ -124,3 +124,4 @@ docs/STATUS.md data/pending-restart.json # Stray dev screenshot artifacts desktop-initial.png +venv/ diff --git a/tests/test_cluster.py b/tests/test_cluster.py index 74875d441..70eec27ee 100644 --- a/tests/test_cluster.py +++ b/tests/test_cluster.py @@ -259,3 +259,121 @@ 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 + + +@pytest.mark.asyncio +class TestWorkerDrain: + """Tests for taOS #890 — worker auto-update with graceful drain.""" + + async def test_drain_worker_graceful_sets_draining_status(self): + """Graceful drain sets status to 'draining', leaves leases intact.""" + mgr = ClusterManager() + await mgr.register_worker(_make_worker("gpu-box", capabilities=["chat"])) + + 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 + assert mgr.get_worker("gpu-box").status == "draining" + + async def test_drain_worker_force_releases_leases_and_marks_offline(self): + """Force drain releases all leases and marks worker offline.""" + mgr = ClusterManager() + await mgr.register_worker(_make_worker("gpu-box", capabilities=["chat"])) + + # Add a lease for this worker + from tinyagentos.cluster.worker_protocol import GpuLease + import time + lease = GpuLease( + lease_id="l_test1", + resource_id="gpu-box:gpu-cuda-0", + caller="test", + expires_at=time.time() + 60, + required_vram_mb=0, + ) + mgr._leases["l_test1"] = lease + + result = await mgr.drain_worker("gpu-box", graceful=False) + assert result["released_leases"] == 1 + assert result["status"] == "offline" + assert mgr.get_worker("gpu-box").status == "offline" + assert "l_test1" not in mgr._leases + + async def test_drain_worker_unknown_returns_error(self): + """Draining a non-existent worker returns error dict.""" + mgr = ClusterManager() + result = await mgr.drain_worker("nonexistent") + assert "error" in result + assert result["worker"] == "nonexistent" + + async def test_cancel_drain_returns_worker_to_online(self): + """Cancel drain returns a draining worker back to online.""" + mgr = ClusterManager() + await mgr.register_worker(_make_worker("gpu-box", capabilities=["chat"])) + 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["worker"] == "gpu-box" + assert result["status"] == "online" + assert mgr.get_worker("gpu-box").status == "online" + + async def test_cancel_drain_only_works_on_draining(self): + """Cancel drain returns error for non-draining workers.""" + mgr = ClusterManager() + await mgr.register_worker(_make_worker("gpu-box")) + + result = await mgr.cancel_drain("gpu-box") + assert "error" in result + assert "not draining" in result["error"] + + async def test_cancel_drain_unknown_returns_error(self): + """Cancel drain on unknown worker returns error.""" + mgr = ClusterManager() + result = await mgr.cancel_drain("nonexistent") + assert "error" in result + + async def test_draining_workers_excluded_from_routing(self): + """Draining workers should not receive new tasks.""" + mgr = ClusterManager() + await mgr.register_worker(_make_worker("online-gpu", capabilities=["chat"], load=0.1)) + await mgr.register_worker(_make_worker("draining-gpu", capabilities=["chat"], load=0.0)) + await mgr.drain_worker("draining-gpu", graceful=True) + + result = mgr.get_workers_for_capability("chat") + assert len(result) == 1 + assert result[0].name == "online-gpu" + + async def test_draining_workers_excluded_from_catalog(self): + """Draining workers are excluded from the aggregate catalog.""" + mgr = ClusterManager() + online = _make_worker("online", capabilities=["chat"]) + online.backends = [{"name": "b1", "type": "ollama", "url": "u", "capabilities": ["chat"], "models": [{"name": "m1"}]}] + draining = _make_worker("draining", capabilities=["chat"]) + draining.backends = [{"name": "b2", "type": "ollama", "url": "u", "capabilities": ["chat"], "models": [{"name": "m2"}]}] + await mgr.register_worker(online) + await mgr.register_worker(draining) + await mgr.drain_worker("draining", graceful=True) + + out = mgr.aggregate_catalog() + assert [w["name"] for w in out["workers"]] == ["online"] + assert [m["name"] for m in out["models"]] == ["m1"] + + async def test_draining_workers_excluded_from_lease_claim(self): + """Lease claims should fail for draining workers.""" + mgr = ClusterManager() + await mgr.register_worker(_make_worker("draining-gpu", url="http://draining:9000")) + + # Set up worker with VRAM info + worker = mgr.get_worker("draining-gpu") + worker.free_vram_mb = 8000 + + await mgr.drain_worker("draining-gpu", graceful=True) + + lease = await mgr.claim_lease( + resource_id="draining-gpu:gpu-cuda-0", + caller="test", + ttl_seconds=30, + ) + assert lease is None diff --git a/tests/test_gpu_arbiter_894.py b/tests/test_gpu_arbiter_894.py new file mode 100644 index 000000000..92e83dc4a --- /dev/null +++ b/tests/test_gpu_arbiter_894.py @@ -0,0 +1,675 @@ +"""Tests for GPU arbiter eviction preemption (taOS #894 fix). + +Verifies that _evict_task actually stops running GPU work by cancelling +the asyncio Task, not just the _arbiter_future. Covers both the direct- +admission path and the queued-then-admitted path. +""" + +import asyncio +import pytest +from tinyagentos.scheduler.gpu_arbiter import GpuArbiter, _QueuedGpuTask +from tinyagentos.scheduler.types import Capability, Priority, Task + + +# ── Helpers ───────────────────────────────────────────────────────────── + +def _make_task(task_id="t1", vram_mb=0, priority=Priority.INTERACTIVE_AGENT): + return Task( + id=task_id, capability=Capability.LLM_CHAT, + payload=lambda r: asyncio.sleep(0), + preferred_resources=[], priority=priority, + estimated_vram_mb=vram_mb, + ) + + +# ── Direct-admission eviction ────────────────────────────────────────── + +class TestDirectAdmissionEviction: + """Eviction of tasks admitted directly (no queue).""" + + @pytest.mark.asyncio + async def test_evict_directly_admitted_task_cancels_submitter(self): + """A directly-admitted task should be preempted by eviction.""" + arbiter = GpuArbiter( + vram_probe=lambda: (8192, 8192), # 8 GiB free + max_queue_size=10, + ) + + # Long-running payload that never finishes + payload_started = asyncio.Event() + payload_cancelled = asyncio.Event() + + async def long_running(_resource): + payload_started.set() + try: + await asyncio.sleep(60) # won't actually finish + except asyncio.CancelledError: + payload_cancelled.set() + raise + + task = _make_task("t-direct", vram_mb=0) + task.payload = long_running + + # Submit in background so we can evict while it runs + async def submitter(): + try: + await arbiter.submit_gpu(task, required_vram_mb=0) + except asyncio.CancelledError: + return "cancelled" + return "completed" + + submit_coro = asyncio.create_task(submitter()) + + # Wait for payload to actually start + await asyncio.wait_for(payload_started.wait(), timeout=5) + + # Evict the running task + evicted = await arbiter._evict_task("t-direct") + assert evicted == 1 + + # Submitter should get CancelledError + result = await submit_coro + assert result == "cancelled" + + # Payload should have been cancelled + assert payload_cancelled.is_set() + + # Task should not be in _running anymore + assert "t-direct" not in arbiter._running + assert "t-direct" not in arbiter._running_tasks + + @pytest.mark.asyncio + async def test_evict_frees_running_slot(self): + """After eviction, the task is removed from _running.""" + arbiter = GpuArbiter( + vram_probe=lambda: (8192, 8192), + max_queue_size=10, + ) + + started = asyncio.Event() + + async def blocking(_resource): + started.set() + await asyncio.sleep(60) + + task = _make_task("t-slot", vram_mb=0) + task.payload = blocking + + async def submitter(): + try: + await arbiter.submit_gpu(task, required_vram_mb=0) + except asyncio.CancelledError: + pass + + submit_coro = asyncio.create_task(submitter()) + await asyncio.wait_for(started.wait(), timeout=5) + + assert "t-slot" in arbiter._running + assert len(arbiter._running) == 1 + + await arbiter._evict_task("t-slot") + + await submit_coro + assert "t-slot" not in arbiter._running + + @pytest.mark.asyncio + async def test_evict_nonexistent_task_returns_zero(self): + """Evicting a non-existent task returns 0.""" + arbiter = GpuArbiter() + assert (await arbiter._evict_task("nonexistent")) == 0 + + @pytest.mark.asyncio + async def test_evict_eviction_disabled_returns_zero(self): + """evict_lowest_priority returns 0 when eviction is disabled.""" + arbiter = GpuArbiter( + vram_probe=lambda: (8192, 8192), + eviction_enabled=False, + ) + + started = asyncio.Event() + + async def blocking(_resource): + started.set() + await asyncio.sleep(60) + + task = _make_task("t-noevict", vram_mb=0) + task.payload = blocking + + async def submitter(): + try: + await arbiter.submit_gpu(task, required_vram_mb=0) + except asyncio.CancelledError: + pass + + submit_coro = asyncio.create_task(submitter()) + await asyncio.wait_for(started.wait(), timeout=5) + + assert await arbiter.evict_lowest_priority() == 0 + assert "t-noevict" in arbiter._running + + # Clean up + arbiter._running_tasks.get("t-noevict", None) + submit_coro.cancel() + try: + await submit_coro + except asyncio.CancelledError: + pass + + +# ── Queued-task eviction ──────────────────────────────────────────────── + +class TestQueuedTaskEviction: + """Eviction of tasks that were queued, then admitted via _drain_queue.""" + + @pytest.mark.asyncio + async def test_evict_queued_admitted_task_cancels_submitter(self): + """A queued-then-admitted task should be preempted by eviction. + + Simulated: we manually call _run_gpu_task via _drain_queue path + and evict while it's running. + """ + arbiter = GpuArbiter( + vram_probe=lambda: (8192, 8192), + max_queue_size=10, + ) + + payload_cancelled = asyncio.Event() + + async def long_running(_resource): + try: + await asyncio.sleep(60) + except asyncio.CancelledError: + payload_cancelled.set() + raise + + task = _make_task("t-queued", vram_mb=0) + task.payload = long_running + + # Manually queue the task and run it like _drain_queue would + async def runner(): + try: + return await arbiter._run_gpu_task(task, 0, False, None) + except asyncio.CancelledError: + return "cancelled" + + runner_task = asyncio.create_task(runner()) + + # Give it a moment to register in _running + await asyncio.sleep(0.1) + + assert "t-queued" in arbiter._running + + # Evict + evicted = await arbiter._evict_task("t-queued") + assert evicted == 1 + + result = await runner_task + assert result == "cancelled" + assert payload_cancelled.is_set() + + @pytest.mark.asyncio + async def test_evict_with_arbiter_future_cancels_future(self): + """_evict_task cancels _arbiter_future if present on the task.""" + arbiter = GpuArbiter( + vram_probe=lambda: (8192, 8192), + max_queue_size=10, + ) + + started = asyncio.Event() + + async def blocking(_resource): + started.set() + await asyncio.sleep(60) + + task = _make_task("t-future", vram_mb=0) + task.payload = blocking + + # Attach an _arbiter_future (as submit_gpu does for queued tasks) + loop = asyncio.get_running_loop() + arb_future: asyncio.Future = loop.create_future() + task._arbiter_future = arb_future # type: ignore[attr-defined] + + async def submitter(): + try: + await arbiter._run_gpu_task(task, 0, False, None) + except asyncio.CancelledError: + pass + + runner_task = asyncio.create_task(submitter()) + await asyncio.wait_for(started.wait(), timeout=5) + + # Evict + await arbiter._evict_task("t-future") + + # The arbiter future should be cancelled + assert arb_future.cancelled() is True + assert arb_future.done() is True + + await runner_task + + +# ── Lease handling during eviction ────────────────────────────────────── + +class FakeClusterManager: + """Minimal fake for testing lease claim/release during eviction.""" + + def __init__(self): + self._leases: dict[str, object] = {} + self.release_calls: list[str] = [] + self.claim_calls: list[str] = [] + + class FakeLease: + def __init__(self, lease_id: str, resource_id: str): + self.lease_id = lease_id + self.resource_id = resource_id + + async def claim_lease(self, resource_id, caller, ttl_seconds, required_vram_mb): + lease_id = f"lease-{resource_id}" + lease = self.FakeLease(lease_id, resource_id) + self._leases[lease_id] = lease + self.claim_calls.append(resource_id) + return lease + + async def release_lease(self, lease_id): + self._leases.pop(lease_id, None) + self.release_calls.append(lease_id) + + def get_leases(self): + return list(self._leases.values()) + + def get_workers(self): + return [] + + +class TestLeaseHandling: + """Lease claim/release behaviour during eviction.""" + + @pytest.mark.asyncio + async def test_evict_releases_lease_once(self): + """Eviction releases the lease exactly once, not double-released.""" + cm = FakeClusterManager() + arbiter = GpuArbiter( + vram_probe=lambda: (8192, 8192), + cluster_manager=cm, + max_queue_size=10, + ) + + started = asyncio.Event() + + async def blocking(_resource): + started.set() + await asyncio.sleep(60) + + task = _make_task("t-lease", vram_mb=0) + task.payload = blocking + + async def runner(): + try: + await arbiter._run_gpu_task(task, 0, False, "gpu-cuda-0:t-lease") + except asyncio.CancelledError: + pass + + runner_task = asyncio.create_task(runner()) + await asyncio.wait_for(started.wait(), timeout=5) + + assert len(cm.claim_calls) == 1 + + await arbiter._evict_task("t-lease") + + # Lease should be released exactly once by eviction + # (the finally block in _run_gpu_task should NOT double-release + # because the entry was already popped from _running) + assert len(cm.release_calls) == 1, f"Expected 1 release, got {len(cm.release_calls)}" + + await runner_task + + @pytest.mark.asyncio + async def test_evict_without_lease_no_double_free(self): + """Eviction without a lease shouldn't cause errors.""" + arbiter = GpuArbiter( + vram_probe=lambda: (8192, 8192), + cluster_manager=None, # No cluster manager + max_queue_size=10, + ) + + started = asyncio.Event() + + async def blocking(_resource): + started.set() + await asyncio.sleep(60) + + task = _make_task("t-nolease", vram_mb=0) + task.payload = blocking + + async def runner(): + try: + await arbiter._run_gpu_task(task, 0, False, None) + except asyncio.CancelledError: + pass + + runner_task = asyncio.create_task(runner()) + await asyncio.wait_for(started.wait(), timeout=5) + + # Should not raise + await arbiter._evict_task("t-nolease") + await runner_task + + +# ── Low-priority eviction ────────────────────────────────────────────── + +class TestEvictLowestPriority: + """evict_lowest_priority selection logic.""" + + @pytest.mark.asyncio + async def test_evicts_lowest_priority_task(self): + """evict_lowest_priority picks the highest priority value (lowest pri).""" + arbiter = GpuArbiter( + vram_probe=lambda: (8192, 8192), + max_queue_size=10, + ) + + started_high = asyncio.Event() + started_low = asyncio.Event() + + async def blocking_high(_resource): + started_high.set() + await asyncio.sleep(60) + + async def blocking_low(_resource): + started_low.set() + await asyncio.sleep(60) + + task_high = _make_task("high", vram_mb=0, priority=Priority.INTERACTIVE_USER) + task_high.payload = blocking_high + task_low = _make_task("low", vram_mb=0, priority=Priority.BATCH) + task_low.payload = blocking_low + + async def runner(task): + try: + await arbiter._run_gpu_task(task, 0, False, None) + except asyncio.CancelledError: + pass + + runner_high = asyncio.create_task(runner(task_high)) + runner_low = asyncio.create_task(runner(task_low)) + + await asyncio.wait_for(started_high.wait(), timeout=5) + await asyncio.wait_for(started_low.wait(), timeout=5) + + # Both running + assert len(arbiter._running) == 2 + + # Evict lowest priority (BATCH > INTERACTIVE_USER in numeric value) + evicted = await arbiter.evict_lowest_priority() + assert evicted == 1 + + # Low-priority (BATCH=50) should be evicted, high (INTERACTIVE_USER=10) stays + assert "low" not in arbiter._running + assert "high" in arbiter._running + + # Cleanup + await arbiter._evict_task("high") + await runner_high + await runner_low + + +# ── Non-blocking drain + eviction-to-make-room ───────────────────────── + +class TestDrainQueueNonBlocking: + """_drain_queue spawns background tasks; doesn't block on _run_gpu_task.""" + + @pytest.mark.asyncio + async def test_drain_returns_immediately_after_spawning_task(self): + """_drain_queue should return without waiting for the GPU task.""" + arbiter = GpuArbiter( + vram_probe=lambda: (8192, 8192), + max_queue_size=10, + ) + + payload_started = asyncio.Event() + payload_done = asyncio.Event() + + async def slow_payload(_resource): + payload_started.set() + await payload_done.wait() # Block until released + return {"slow": "done"} + + task = _make_task("t-drain-fast", vram_mb=0) + task.payload = slow_payload + + # Queue the task + loop = asyncio.get_running_loop() + arb_future: asyncio.Future = loop.create_future() + task._arbiter_future = arb_future # type: ignore[attr-defined] + entry = _QueuedGpuTask( + priority=int(task.priority), seq=1, task=task, + required_vram_mb=0, evictable=False, + ) + await arbiter._queue.put(entry) + + # Drain — should return immediately after spawning + await arbiter._drain_queue() + + # Payload was started (spawned as background task) + await asyncio.wait_for(payload_started.wait(), timeout=5) + + # _drain_queue returned — queue should be empty + assert arbiter._queue.empty() + + # Task should be in _running (registered by _run_gpu_task) + assert "t-drain-fast" in arbiter._running + + # Release the payload + payload_done.set() + + # arbiter_future should be resolved + result = await asyncio.wait_for(arb_future, timeout=5) + assert result is not None + + # Cleanup — task should remove itself from _running on completion + # (the _run_gpu_task finally block handles this) + await asyncio.sleep(0.1) + # After completion, _running should be empty (or at least not contain t-drain-fast) + # Actually, with the new async model, _run_gpu_task removes itself from _running + # on completion, so let's check: + if "t-drain-fast" in arbiter._running: + await arbiter._evict_task("t-drain-fast") + + @pytest.mark.asyncio + async def test_done_callback_sets_arbiter_future_result(self): + """When the background task completes, _arbiter_future gets the result.""" + arbiter = GpuArbiter( + vram_probe=lambda: (8192, 8192), + max_queue_size=10, + ) + + async def quick_payload(_resource): + return {"answer": 42} + + task = _make_task("t-done-cb", vram_mb=0) + task.payload = quick_payload + + loop = asyncio.get_running_loop() + arb_future: asyncio.Future = loop.create_future() + task._arbiter_future = arb_future # type: ignore[attr-defined] + + entry = _QueuedGpuTask( + priority=int(task.priority), seq=1, task=task, + required_vram_mb=0, evictable=False, + ) + await arbiter._queue.put(entry) + + await arbiter._drain_queue() + + result = await asyncio.wait_for(arb_future, timeout=5) + assert result == {"answer": 42} + + @pytest.mark.asyncio + async def test_done_callback_propagates_exception(self): + """When the background task raises, the exception propagates to the future.""" + arbiter = GpuArbiter( + vram_probe=lambda: (8192, 8192), + max_queue_size=10, + ) + + async def failing_payload(_resource): + raise ValueError("boom") + + task = _make_task("t-fail-cb", vram_mb=0) + task.payload = failing_payload + + loop = asyncio.get_running_loop() + arb_future: asyncio.Future = loop.create_future() + task._arbiter_future = arb_future # type: ignore[attr-defined] + + entry = _QueuedGpuTask( + priority=int(task.priority), seq=1, task=task, + required_vram_mb=0, evictable=False, + ) + await arbiter._queue.put(entry) + + await arbiter._drain_queue() + + with pytest.raises(ValueError, match="boom"): + await asyncio.wait_for(arb_future, timeout=5) + + +class TestEvictionToMakeRoom: + """When admission fails, _drain_queue tries eviction before re-queuing.""" + + @pytest.mark.asyncio + async def test_eviction_triggered_when_vram_full(self): + """If VRAM is full, drain should evict a lower-priority task to make room.""" + arbiter = GpuArbiter( + vram_probe=lambda: ( + (1024, 8192) if len(arbiter_ctx["ref"]._running) > 0 else (8192, 8192) + ), + max_queue_size=10, + ) + arbiter_ctx: dict[str, object] = {"ref": arbiter} + + # Put a low-priority task in _running (simulating an already-running task) + low_task = _make_task("low-running", vram_mb=4096, priority=Priority.BATCH) + low_task.payload = lambda r: asyncio.sleep(60) + arbiter._running["low-running"] = (low_task, None, int(Priority.BATCH), 4096) + + # Now queue a higher-priority task that needs VRAM + hi_task = _make_task("hi-queued", vram_mb=4096, priority=Priority.INTERACTIVE_USER) + hi_started = asyncio.Event() + hi_release = asyncio.Event() + + async def hi_payload(_resource): + hi_started.set() + await hi_release.wait() + return {"ok": True} + + hi_task.payload = hi_payload + + loop = asyncio.get_running_loop() + arb_future: asyncio.Future = loop.create_future() + hi_task._arbiter_future = arb_future # type: ignore[attr-defined] + + entry = _QueuedGpuTask( + priority=int(hi_task.priority), seq=1, task=hi_task, + required_vram_mb=4096, evictable=False, + ) + await arbiter._queue.put(entry) + + # Drain — should evict "low-running" and admit "hi-queued" + await arbiter._drain_queue() + + # Yield so the background task can register in _running + await asyncio.sleep(0) + + # Low-priority task should be evicted + assert "low-running" not in arbiter._running, ( + f"Expected low-priority task to be evicted, but _running={list(arbiter._running)}" + ) + + # High-priority task should be running + assert "hi-queued" in arbiter._running + + # Release the payload and wait for result + hi_release.set() + result = await asyncio.wait_for(arb_future, timeout=5) + assert result == {"ok": True} + + @pytest.mark.asyncio + async def test_no_eviction_when_all_running_higher_priority(self): + """Don't evict tasks that are higher priority than the queued task.""" + arbiter = GpuArbiter( + vram_probe=lambda: (1024, 8192), # Some free VRAM but not enough + max_queue_size=10, + ) + + # Put a high-priority task in _running + hi_task = _make_task("hi-running", vram_mb=4096, priority=Priority.INTERACTIVE_USER) + hi_task.payload = lambda r: asyncio.sleep(60) + arbiter._running["hi-running"] = (hi_task, None, int(Priority.INTERACTIVE_USER), 4096) + + # Queue a lower-priority task + lo_task = _make_task("lo-queued", vram_mb=4096, priority=Priority.BATCH) + + loop = asyncio.get_running_loop() + arb_future: asyncio.Future = loop.create_future() + lo_task._arbiter_future = arb_future # type: ignore[attr-defined] + lo_task.payload = lambda r: asyncio.sleep(0) + + entry = _QueuedGpuTask( + priority=int(lo_task.priority), seq=1, task=lo_task, + required_vram_mb=4096, evictable=False, + ) + await arbiter._queue.put(entry) + + # Drain — should NOT evict the higher-priority running task + await arbiter._drain_queue() + + # High-priority task should still be running + assert "hi-running" in arbiter._running + + # Low-priority task should be back in the queue (couldn't be admitted) + assert not arbiter._queue.empty() + re_queued = arbiter._queue.get_nowait() + assert re_queued.task.id == "lo-queued" + + # Cleanup + await arbiter._evict_task("hi-running") + + @pytest.mark.asyncio + async def test_eviction_disabled_no_eviction(self): + """When eviction is disabled, don't evict even if VRAM is full.""" + arbiter = GpuArbiter( + vram_probe=lambda: (1024, 8192), # Some free VRAM but not enough + max_queue_size=10, + eviction_enabled=False, + ) + + # Put a task in _running + run_task = _make_task("running", vram_mb=4096, priority=Priority.BATCH) + run_task.payload = lambda r: asyncio.sleep(60) + arbiter._running["running"] = (run_task, None, int(Priority.BATCH), 4096) + + # Queue a higher-priority task + hi_task = _make_task("hi-queued", vram_mb=4096, priority=Priority.INTERACTIVE_USER) + hi_task.payload = lambda r: asyncio.sleep(0) + + loop = asyncio.get_running_loop() + arb_future: asyncio.Future = loop.create_future() + hi_task._arbiter_future = arb_future # type: ignore[attr-defined] + + entry = _QueuedGpuTask( + priority=int(hi_task.priority), seq=1, task=hi_task, + required_vram_mb=4096, evictable=False, + ) + await arbiter._queue.put(entry) + + await arbiter._drain_queue() + + # Running task should NOT be evicted + assert "running" in arbiter._running + + # Queued task should be back in the queue + assert not arbiter._queue.empty() + + # Cleanup + await arbiter._evict_task("running") diff --git a/tests/test_gpu_arbiter_toctou.py b/tests/test_gpu_arbiter_toctou.py new file mode 100644 index 000000000..0ff4734c4 --- /dev/null +++ b/tests/test_gpu_arbiter_toctou.py @@ -0,0 +1,210 @@ +"""Tests for TOCTOU VRAM reservation fix (taOS #894 review feedback #2). + +Verifies that two concurrent admissions cannot both pass when the combined +required VRAM exceeds the available free VRAM. The fix subtracts in-flight +reservations from the nvidia-smi probe so the second caller sees the +capacity already promised to the first. +""" +import asyncio + +import pytest + +from tinyagentos.scheduler.gpu_arbiter import GpuArbiter +from tinyagentos.scheduler.types import Capability, Priority, Task + + +def _make_task(task_id="t1", vram_mb=0, priority=Priority.INTERACTIVE_AGENT): + return Task( + id=task_id, capability=Capability.LLM_CHAT, + payload=lambda r: asyncio.sleep(0), + preferred_resources=[], priority=priority, + estimated_vram_mb=vram_mb, + ) + + +class TestToctouReservation: + """Concurrent admission doesn't over-commit VRAM.""" + + @pytest.mark.asyncio + async def test_concurrent_reserve_and_check_only_one_admitted(self): + """Two concurrent _reserve_and_check calls — only first admitted. + + With 8192 MiB free and each needing 5000 MiB, the second + _reserve_and_check must see the first's reservation and be + rejected. The atomic lock ensures no TOCTOU gap. + """ + arbiter = GpuArbiter( + vram_probe=lambda: (8192, 8192), + max_queue_size=10, + ) + + results: list[bool] = [] + + async def try_reserve(task_id): + admission = await arbiter._reserve_and_check(task_id, 5000) + results.append(admission.admitted) + + # Launch both concurrently — this is the TOCTOU window. + await asyncio.gather( + try_reserve("t-a"), + try_reserve("t-b"), + ) + + admitted_count = sum(results) + assert admitted_count == 1, ( + f"Expected exactly 1 admission, got {admitted_count}: {results}" + ) + assert arbiter._reserved_vram_mb == 5000 + + @pytest.mark.asyncio + async def test_submit_gpu_reserves_and_releases(self): + """submit_gpu reserves VRAM on admission, releases after completion.""" + arbiter = GpuArbiter( + vram_probe=lambda: (8192, 8192), + max_queue_size=10, + ) + + task = _make_task("t-release", vram_mb=0) + task.payload = lambda r: asyncio.sleep(0) + + assert arbiter._reserved_vram_mb == 0 + assert len(arbiter._pending_reservations) == 0 + + await arbiter.submit_gpu(task, required_vram_mb=2048) + + # After completion, reservation should be released. + assert arbiter._reserved_vram_mb == 0 + assert len(arbiter._pending_reservations) == 0 + + @pytest.mark.asyncio + async def test_reservation_released_on_eviction(self): + """Evicting a running task releases its VRAM reservation.""" + arbiter = GpuArbiter( + vram_probe=lambda: (8192, 8192), + max_queue_size=10, + ) + + started = asyncio.Event() + + async def blocking(_resource): + started.set() + await asyncio.sleep(60) + + task = _make_task("t-evict", vram_mb=0) + task.payload = blocking + + async def runner(): + try: + await arbiter._run_gpu_task(task, 4096, False, None) + except asyncio.CancelledError: + pass + + runner_task = asyncio.create_task(runner()) + await asyncio.wait_for(started.wait(), timeout=5) + + # Simulate what _reserve_and_check would have done when the task + # was admitted. + arbiter._reserved_vram_mb += 4096 + arbiter._pending_reservations["t-evict"] = 4096 + assert arbiter._reserved_vram_mb == 4096 + + # Evict + await arbiter._evict_task("t-evict") + await runner_task + + # After eviction, reservation should be released. + assert arbiter._reserved_vram_mb == 0 + assert "t-evict" not in arbiter._pending_reservations + + @pytest.mark.asyncio + async def test_reservation_subtracted_from_admission_check(self): + """_check_admission sees reduced capacity when VRAM is reserved.""" + arbiter = GpuArbiter( + vram_probe=lambda: (8192, 8192), + max_queue_size=10, + ) + + # Initially 8192 free — 6000 should be okay. + admission = arbiter._check_admission(None, 6000) + assert admission.admitted + assert admission.free_vram_mb == 8192 # effective free (0 reserved) + + # Reserve 4000 MiB. + arbiter._reserved_vram_mb = 4000 + arbiter._pending_reservations["t-holder"] = 4000 + + # Now effective free = 8192 - 4000 = 4192. + # 6000 should be rejected. + admission = arbiter._check_admission(None, 6000) + assert not admission.admitted + assert admission.free_vram_mb == 4192 + + # 4000 should still be admitted. + admission = arbiter._check_admission(None, 4000) + assert admission.admitted + + @pytest.mark.asyncio + async def test_atomic_reserve_and_check(self): + """_reserve_and_check atomically checks and reserves.""" + arbiter = GpuArbiter( + vram_probe=lambda: (8192, 8192), + max_queue_size=10, + ) + + # First call — should admit and reserve. + admission = await arbiter._reserve_and_check("t1", 5000) + assert admission.admitted + assert arbiter._reserved_vram_mb == 5000 + assert arbiter._pending_reservations == {"t1": 5000} + + # Second call — should see the reservation and fail. + admission2 = await arbiter._reserve_and_check("t2", 5000) + assert not admission2.admitted + # Reservation was NOT made for failed admission. + assert arbiter._reserved_vram_mb == 5000 + assert "t2" not in arbiter._pending_reservations + + # Release t1 + arbiter._release_reservation("t1") + assert arbiter._reserved_vram_mb == 0 + + @pytest.mark.asyncio + async def test_release_reservation_idempotent(self): + """_release_reservation is safe to call multiple times.""" + arbiter = GpuArbiter( + vram_probe=lambda: (8192, 8192), + max_queue_size=10, + ) + + arbiter._reserved_vram_mb = 1000 + arbiter._pending_reservations["t-idem"] = 1000 + + arbiter._release_reservation("t-idem") + assert arbiter._reserved_vram_mb == 0 + + # Second call — no-op, no crash. + arbiter._release_reservation("t-idem") + assert arbiter._reserved_vram_mb == 0 + + # Releasing a task that never had a reservation — no-op. + arbiter._release_reservation("non-existent") + assert arbiter._reserved_vram_mb == 0 + + @pytest.mark.asyncio + async def test_reserved_vram_in_stats(self): + """stats() includes reservation info.""" + arbiter = GpuArbiter( + vram_probe=lambda: (8192, 8192), + max_queue_size=10, + ) + + stats = await arbiter.stats() + assert stats["reserved_vram_mb"] == 0 + assert stats["pending_reservations"] == 0 + + arbiter._reserved_vram_mb = 2048 + arbiter._pending_reservations["t-stats"] = 2048 + + stats = await arbiter.stats() + assert stats["reserved_vram_mb"] == 2048 + assert stats["pending_reservations"] == 1 diff --git a/tinyagentos/app.py b/tinyagentos/app.py index 617bd719a..f5e55ae8a 100644 --- a/tinyagentos/app.py +++ b/tinyagentos/app.py @@ -59,6 +59,7 @@ async def get_response(self, path, scope): from tinyagentos.installation_state import InstallationState from tinyagentos.scheduler import BackendCatalog, HistoryStore, ScoreCache, TaskScheduler from tinyagentos.scheduler.discovery import build_scheduler as build_resource_scheduler +from tinyagentos.scheduler.gpu_arbiter import GpuArbiter, _probe_nvidia_vram from tinyagentos.torrent_settings import TorrentSettingsStore from tinyagentos.relationships import RelationshipManager from tinyagentos.github_identities import GitHubIdentitiesStore @@ -1124,6 +1125,7 @@ async def _reload_llm_proxy_on_catalog_change() -> None: # Build the resource scheduler from hardware profile + live catalog. # Phase 1: local resources only (NPU + CPU), capability-based routing # with fallback and priority. Cluster-aware dispatch is Phase 3. + resource_scheduler = None try: resource_scheduler = build_resource_scheduler( hardware_profile, @@ -1140,6 +1142,24 @@ async def _reload_llm_proxy_on_catalog_change() -> None: except Exception: logger.exception("resource scheduler failed to build — routes will use static config") app.state.resource_scheduler = None + + # Build the GPU arbiter — wraps the resource scheduler with VRAM-accounted + # admission control, queuing, and eviction for GPU-bound workloads. + try: + gpu_arbiter = GpuArbiter( + scheduler=resource_scheduler if resource_scheduler is not None else None, + cluster_manager=cluster_manager, + vram_probe=_probe_nvidia_vram, + max_queue_size=100, + eviction_enabled=True, + ) + await gpu_arbiter.start() + app.state.gpu_arbiter = gpu_arbiter + cluster_manager._gpu_arbiter = gpu_arbiter + logger.info("GPU arbiter ready (queue size=100, eviction=enabled)") + except Exception: + logger.exception("GPU arbiter failed to start — GPU tasks will use vanilla scheduler") + app.state.gpu_arbiter = None # Detect and set container runtime from tinyagentos.containers.backend import configure_container_runtime configure_container_runtime(config) @@ -1294,6 +1314,8 @@ async def _notify_emitter(row: dict) -> None: # local_heartbeat_task is cancelled+awaited above under the bounded # cancel_and_wait budget alongside the supervised background tasks. await cluster_manager.stop() + if app.state.gpu_arbiter is not None: + await app.state.gpu_arbiter.stop() llm_proxy.stop() try: from tinyagentos.taos_agent_runtime import stop_taos_opencode_server diff --git a/tinyagentos/cluster/manager.py b/tinyagentos/cluster/manager.py index 69df08323..60fa049d3 100644 --- a/tinyagentos/cluster/manager.py +++ b/tinyagentos/cluster/manager.py @@ -3,8 +3,13 @@ import logging import secrets import time +from typing import TYPE_CHECKING + from tinyagentos.cluster.worker_protocol import GpuLease, WorkerInfo +if TYPE_CHECKING: + from tinyagentos.scheduler.gpu_arbiter import GpuArbiter + logger = logging.getLogger(__name__) HEARTBEAT_TIMEOUT = 30 # seconds before marking worker offline @@ -39,6 +44,7 @@ def __init__(self, notifications=None, capabilities=None): self._monitor_task: asyncio.Task | None = None self._notifications = notifications # NotificationStore, optional self._capabilities = capabilities # CapabilityChecker, optional + self._gpu_arbiter: GpuArbiter | None = None # GpuArbiter, wired in app.py # Track worker names seen at least once so we only fire worker.join # on the very first appearance within this process lifetime. self._ever_seen: set[str] = set() @@ -273,13 +279,15 @@ def _parse_resource_id(self, resource_id: str) -> tuple[str, str] | None: return worker, rest def _worker_for_resource(self, resource_id: str) -> WorkerInfo | None: - """Return the WorkerInfo for a resource_id, or None.""" + """Return the WorkerInfo for a resource_id, or None. + + Excludes draining and offline workers (taOS #890).""" parsed = self._parse_resource_id(resource_id) if parsed is None: return None worker_name, _ = parsed worker = self._workers.get(worker_name) - if worker is None or worker.status != "online": + if worker is None or worker.status not in ("online",): return None return worker @@ -393,10 +401,121 @@ def _sweep_expired_leases(self): lease = self._leases.pop(lid) logger.debug("Lease expired: %s on %s", lid, lease.resource_id) + # ── taOS #890: graceful worker drain for auto-update ─────────────── + + async def drain_worker(self, name: str, graceful: bool = True) -> dict: + """Begin draining a worker — gracefully detach without dropping tasks. + + When ``graceful=True`` (the default): + 1. Worker enters ``"draining"`` status — no new tasks routed to it. + 2. Existing leases are allowed to run to completion (TTL not touched). + 3. The monitor loop will eventually sweep expired leases and then + mark the worker ``"offline"`` once all leases are released. + + When ``graceful=False``: + 1. Worker enters ``"draining"`` status. + 2. All active leases for this worker are released immediately. + 3. Any running GPU arbiter tasks 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: + # Force-release all leases for this worker, holding _lease_lock + # so the mutation is serialized with claim/release/sweep. + 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 any running GPU arbiter tasks for the released leases + # so eviction isn't a prod no-op (taOS cross-cutting wiring). + if released > 0 and self._gpu_arbiter 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"} + # ── 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 #890).""" eligible = [ w for w in self._workers.values() if w.status == "online" and capability in w.capabilities @@ -443,7 +562,7 @@ def aggregate_catalog(self) -> dict: for worker in self._workers.values(): if worker.status != "online": - continue + continue # skip offline and draining workers (taOS #890) worker_caps = set(worker.capabilities or []) all_capabilities |= worker_caps @@ -486,14 +605,12 @@ def aggregate_catalog(self) -> dict: async def _monitor_loop(self): """Monitor worker heartbeats, mark stale workers as offline, and - sweep expired GPU leases.""" + sweep expired GPU leases. Auto-completes draining workers + whose leases have all been released (taOS #890).""" while True: now = time.time() - for worker in self._workers.values(): - # The 'local' worker is the controller itself — it never sends - # heartbeats (it IS the server), so never mark it offline. - if worker.name == "local": - continue + for worker in list(self._workers.values()): + # Handle online workers that haven't heartbeated if worker.status == "online" and (now - worker.last_heartbeat) > HEARTBEAT_TIMEOUT: worker.status = "offline" logger.warning(f"Worker '{worker.name}' marked offline (no heartbeat for {HEARTBEAT_TIMEOUT}s)") @@ -516,6 +633,57 @@ 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) + + # Handle draining workers: auto-complete if no active leases (taOS #890) + elif worker.status == "draining": + active_leases = [ + lid for lid, lease in self._leases.items() + if (parsed := self._parse_resource_id(lease.resource_id)) + and parsed[0] == worker.name + ] + if not active_leases: + worker.status = "offline" + logger.info( + "Worker '%s' drained — all leases released, marked offline", + worker.name, + ) + if self._notifications: + await self._notifications.emit_event( + "worker.leave", + f"Worker '{worker.name}' drained and went offline", + "All tasks completed; worker detached gracefully.", + level="info", + ) + elif (now - worker.last_heartbeat) > HEARTBEAT_TIMEOUT: + # Draining worker went stale — force-finish the drain. + # Hold _lease_lock so mutation is serialized with + # claim/release/sweep (taOS #1690 lock fix). + async with self._lease_lock: + lids = [ + lid for lid, lease in self._leases.items() + if (parsed := self._parse_resource_id(lease.resource_id)) + and parsed[0] == worker.name + ] + for lid in lids: + self._leases.pop(lid, None) + worker.status = "offline" + logger.warning( + "Worker '%s' drain timed out (no heartbeat for %ds) — " + "force-released %d leases, marked offline", + worker.name, HEARTBEAT_TIMEOUT, len(lids), + ) + # Cancel any running GPU arbiter tasks for the + # released leases (taOS cross-cutting wiring). + if lids and self._gpu_arbiter is not None: + try: + cancelled, already_done = await self._gpu_arbiter.cancel_running_for_leases(set(lids)) + logger.info( + "Worker '%s' stale-drain: arbiter cancelled %d tasks, %d already done", + worker.name, cancelled, already_done) + except Exception: + logger.exception( + "gpu-arbiter: cancel for stale-drain of '%s' failed", + worker.name) async with self._lease_lock: self._sweep_expired_leases() await asyncio.sleep(5) diff --git a/tinyagentos/routes/cluster.py b/tinyagentos/routes/cluster.py index ef5d91708..3f9fb901a 100644 --- a/tinyagentos/routes/cluster.py +++ b/tinyagentos/routes/cluster.py @@ -1101,3 +1101,105 @@ async def list_leases(request: Request): ], "count": len(leases), } + + +# ── taOS #890: worker auto-update with graceful drain ──────────────────── + + +class DrainRequest(BaseModel): + graceful: bool = True + + +@router.post("/api/cluster/workers/{name}/drain") +async def drain_worker(request: Request, name: str, body: DrainRequest = DrainRequest()): + """Begin draining a worker — gracefully detach without dropping tasks. + + When ``graceful=true`` (the default), the worker enters "draining" status: + no new tasks are routed to it, but existing leases run to completion. + The monitor loop auto-completes the drain when all leases are released. + + When ``graceful=false``, all leases are force-released and the worker + is marked offline immediately. + """ + ok, err = _require_admin(request) + if not ok: + return err + cluster = request.app.state.cluster_manager + result = await cluster.drain_worker(name, graceful=body.graceful) + if "error" in result: + return JSONResponse(result, status_code=404) + return result + + +@router.post("/api/cluster/workers/{name}/cancel-drain") +async def cancel_drain(request: Request, name: str): + """Cancel an in-progress drain and return the worker to online status.""" + ok, err = _require_admin(request) + if not ok: + return err + cluster = request.app.state.cluster_manager + result = await cluster.cancel_drain(name) + if "error" in result: + return JSONResponse(result, status_code=404) + return result + + +@router.post("/api/cluster/workers/{name}/update") +async def update_worker(request: Request, name: str): + """Trigger a full auto-update sequence on a worker. + + Orchestrates the graceful update sequence: + 1. Pause incoming tasks — worker enters "draining" status. + 2. Inflight leases run to completion (drain). + 3. Send ``update-worker`` deploy command to the worker. + 4. Worker restarts, re-registers, returns to "online". + 5. Routing resumes automatically once the worker is back online. + """ + ok, err = _require_admin(request) + if not ok: + return err + cluster = request.app.state.cluster_manager + worker = cluster.get_worker(name) + if not worker: + return JSONResponse({"error": f"Worker '{name}' not found"}, status_code=404) + if worker.status not in ("online", "draining"): + return JSONResponse( + {"error": f"Worker '{name}' is not online (status={worker.status})"}, + status_code=400, + ) + + # Step 1: Begin draining + drain_result = await cluster.drain_worker(name, graceful=True) + if "error" in drain_result: + return JSONResponse(drain_result, status_code=404) + + # Step 2: Trigger the update-worker deploy command on the worker + import httpx + deploy_result = None + try: + async with httpx.AsyncClient(timeout=620) as client: + resp = await client.post( + f"{worker.url}/api/worker/deploy", + json={"command": "update-worker"}, + ) + deploy_result = resp.json() + except Exception as exc: + # Deploy failed — cancel drain so worker can still serve traffic + await cluster.cancel_drain(name) + return JSONResponse( + {"error": f"Worker update deploy failed: {exc}", "drain_cancelled": True}, + status_code=502, + ) + + return { + "worker": name, + "status": "updating", + "previous_status": drain_result["previous_status"], + "drain": drain_result, + "deploy": deploy_result, + "message": ( + "Worker is draining and updating. It will restart and re-register. " + "The controller monitor loop will auto-complete the drain once " + "all leases are released and the new worker process heartbeats." + ), + } diff --git a/tinyagentos/scheduler/__init__.py b/tinyagentos/scheduler/__init__.py index 187914bef..def900993 100644 --- a/tinyagentos/scheduler/__init__.py +++ b/tinyagentos/scheduler/__init__.py @@ -32,6 +32,7 @@ _LAZY_EXPORTS = { "BackendCatalog": "backend_catalog", "BackendEntry": "backend_catalog", + "GpuArbiter": "gpu_arbiter", "HistoryStore": "history_store", "Resource": "resource", "Scheduler": "scheduler", @@ -58,6 +59,7 @@ def __dir__(): "BackendCatalog", "BackendEntry", "Capability", + "GpuArbiter", "HistoryStore", "NoResourceAvailableError", "Priority", diff --git a/tinyagentos/scheduler/discovery.py b/tinyagentos/scheduler/discovery.py index 1f4f4ff2e..641192853 100644 --- a/tinyagentos/scheduler/discovery.py +++ b/tinyagentos/scheduler/discovery.py @@ -15,6 +15,7 @@ from tinyagentos.scheduler.backend_catalog import BackendCatalog from tinyagentos.scheduler.history_store import HistoryStore +from tinyagentos.scheduler.gpu_arbiter import _probe_nvidia_vram from tinyagentos.scheduler.resource import Resource, Tier from tinyagentos.scheduler.scheduler import Scheduler from tinyagentos.scheduler.score_cache import ScoreCache @@ -51,6 +52,20 @@ "vision", } +# Every capability a GPU can run given the right backend. GPU is the +# fastest tier; the potential set mirrors CPU/NPU because any inference +# task that can run on CPU can also run (faster) on GPU. Live capabilities +# are still filtered to what backends actually have loaded right now. +GPU_POTENTIAL_CAPABILITIES: set[str] = { + "llm-chat", + "embedding", + "reranking", + "image-generation", + "speech-to-text", + "text-to-speech", + "vision", +} + logger = logging.getLogger(__name__) @@ -172,6 +187,70 @@ def _npu_backend_for(capability: str) -> Optional[str]: ) ) + # GPU (CUDA/ROCm/Vulkan/Metal), only if a healthy GPU-capable backend exists. + # GPU backends: vllm, llama-cpp (when built with CUDA), ollama (GPU mode), + # exo, mlx (Apple Silicon GPU). Also sd-cpp/sd-gpu for image-generation. + gpu_backend_types = {"vllm", "ollama", "exo", "mlx"} + gpu_backends = [ + b for b in catalog.backends() + if b.status == "ok" and b.type in gpu_backend_types + ] + gpu_info = getattr(hardware_profile, "gpu", None) + gpu_type = getattr(gpu_info, "type", None) if gpu_info else None + has_gpu_hardware = gpu_type not in (None, "", "none") + + if gpu_backends or has_gpu_hardware: + gpu_count = 1 # Default single-GPU; multi-GPU is Phase 2 + gpu_signature = ResourceSignature( + platform="cuda" if gpu_type in ("cuda", "nvidia") else ( + "rocm" if gpu_type == "rocm" else ( + "metal" if gpu_type == "apple" else "gpu" + ) + ), + runtime="cuda" if gpu_type in ("cuda", "nvidia") else ( + "rocm" if gpu_type == "rocm" else "native" + ), + runtime_version="", + ) + + def _gpu_vram_probe() -> int: + free, _total = _probe_nvidia_vram() + return free if free > 0 else 999_999 # optimistic + + def _gpu_capabilities() -> set[str]: + caps: set[str] = set() + for b in catalog.backends(): + if b.status == "ok" and b.type in gpu_backend_types: + caps |= b.capabilities + return caps + + def _gpu_backend_for(capability: str): + for b in catalog.backends_with_capability(capability): + if b.type in gpu_backend_types: + return b.url + return None + + for gpu_idx in range(gpu_count): + gpu_name = f"gpu-cuda-{gpu_idx}" + scheduler.register( + Resource( + name=gpu_name, + signature=gpu_signature, + concurrency=1, + tier=Tier.GPU, + potential_capabilities=GPU_POTENTIAL_CAPABILITIES, + get_capabilities=_gpu_capabilities, + backend_lookup=_gpu_backend_for, + score_lookup=_make_score_lookup(gpu_name), + memory_probe=_gpu_vram_probe, + ) + ) + logger.info( + "discovery: registered GPU resource %s (%s, concurrency=1)", + gpu_name, + gpu_signature.platform, + ) + # CPU inference, always register. Backend-driven: only advertises the # capabilities that some CPU backend currently serves (sd-cpp, llama-cpp, etc.) cpu_signature = ResourceSignature( diff --git a/tinyagentos/scheduler/gpu_arbiter.py b/tinyagentos/scheduler/gpu_arbiter.py new file mode 100644 index 000000000..17b4a075e --- /dev/null +++ b/tinyagentos/scheduler/gpu_arbiter.py @@ -0,0 +1,596 @@ +"""GPU Arbiter — VRAM-accounted admission + queue + eviction for GPU workloads. + +Slice 2 of taOS #894 — builds on Slice 1 (VRAM endpoint) and the lease +system (#893). Provides admission control, queuing, and priority-based +eviction to prevent concurrent-load driver crashes (NVIDIA Xid 62). +""" +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 ( + Capability, + NoResourceAvailableError, + Priority, + Task, +) + +logger = logging.getLogger(__name__) + + +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 + + +@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) + + +@dataclass +class GpuAdmission: + """Result of a GPU admission check.""" + admitted: bool + reason: str | None = None + existing_lease_id: str | None = None + existing_lease_holder: str | None = None + free_vram_mb: int = 0 + required_vram_mb: int = 0 + + +class GpuArbiter: + """VRAM-accounted admission control layered on top of the Scheduler. + + Usage: + arbiter = GpuArbiter(scheduler=sched, cluster_manager=cm) + 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, + ): + 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 + 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 (taOS #894 review feedback #2) --- + # _check_admission reads live nvidia-smi free VRAM with no reservation. + # Two concurrent submit_gpu calls can both pass before either model loads + # — the exact concurrent-load crash the arbiter exists to prevent. + # _reserved_vram_mb tracks in-flight admissions that haven't yet + # consumed physical VRAM (model still loading). _reserve_and_check() + # atomically checks admission against (probe - _reserved_vram_mb) and + # reserves the VRAM if admitted. _release_reservation() is called in + # _run_gpu_task's finally and in _evict_task. + 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 + # ── taOS #796: pause/resume queue control ──────────────────────── + self._paused: bool = False + self._paused_at: float | None = None + + 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 + + # ── taOS #796: pause/resume queue control ────────────────────────── + + @property + def paused(self) -> bool: + """Whether the queue processor is currently paused.""" + return self._paused + + def pause(self) -> bool: + """Pause queue processing. New tasks still queue; running tasks finish. + + Returns True if paused, False if already paused. + """ + 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: + """Resume queue processing after a pause. + + Returns True if resumed, False if not paused. + """ + 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 + + # ── taOS #796: hardware-aware LLM admission ──────────────────────── + + def _check_gpu_arch_compatibility( + self, required_gpu_arch: str | None, resource_id: str | None = None, + ) -> tuple[bool, str | None]: + """Check if a worker's GPU architecture satisfies the requirement. + + ``required_gpu_arch`` is a CUDA compute-capability string like + ``"sm_86"`` or ``"sm_75"``. When the caller specifies this, the + arbiter checks that at least one online cluster worker has a + compatible GPU. Without a cluster_manager, the check is skipped + (standalone mode trusts the local GPU). + + Returns ``(True, None)`` if compatible or no arch requirement, + ``(False, reason)`` if incompatible. + """ + if not required_gpu_arch: + return True, None + if self._cluster_manager is None: + return True, None + + for worker in self._cluster_manager.get_workers(): + if worker.status not in ("online", "draining"): + continue + gpu_info = worker.hardware.get("gpu", {}) if isinstance(worker.hardware, dict) else {} + gpu_model = gpu_info.get("model", "") or "" + cc = gpu_info.get("compute_cap", "") or "" + if required_gpu_arch in gpu_model or required_gpu_arch in cc: + if resource_id is None or resource_id.startswith(worker.name + ":"): + return True, None + return False, f"no online worker with GPU architecture {required_gpu_arch}" + + # ── Reservation helpers (TOCTOU fix) ────────────────────────────── + + 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 (e.g. + a zero-VRAM task or eviction of a queued-but-not-yet-admitted entry). + """ + 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. + + Acquires _reservation_lock, calls _check_admission (which subtracts + _reserved_vram_mb from the hardware probe), and if admitted + immediately adds required_vram_mb to _reserved_vram_mb so + concurrent callers see the reduced capacity. + + Returns the GpuAdmission. The caller MUST call _release_reservation + when the task finishes or is evicted. + """ + if required_vram_mb <= 0: + return GpuAdmission(admitted=True) + + async with self._reservation_lock: + # Build a synthetic task just for the admission check. The + # reservation-aware check uses _reserved_vram_mb internally. + admission = self._check_admission(None, 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. + + Args: + task: The Task to run. + required_vram_mb: VRAM needed in MiB (0 = no VRAM check). + 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"``). + """ + self._submitted += 1 + + # Hardware architecture check (taOS #796) + if required_gpu_arch: + arch_ok, arch_reason = self._check_gpu_arch_compatibility( + required_gpu_arch, resource_id, + ) + if not arch_ok: + raise NoResourceAvailableError( + f"GPU architecture requirement not met: {arch_reason} " + f"(task {task.id})" + ) + + if required_vram_mb > 0: + admission = await self._reserve_and_check(task.id, required_vram_mb) + if not admission.admitted: + # Not admitted — release the reservation (belt-and-suspenders: + # _reserve_and_check only reserves on admit, but be safe). + 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 + # Admitted — reservation already held by _reserve_and_check, so + # _run_gpu_task can proceed safely. + return await self._run_gpu_task(task, required_vram_mb, evictable, resource_id) + + def _check_admission(self, task: Task | None, 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 so that concurrent admission attempts see the + capacity already promised to other tasks whose models are still + loading. + + *task* may be None when called from _reserve_and_check (which + doesn't need the full Task — just the VRAM budget). + """ + if required_vram_mb <= 0: + return GpuAdmission(admitted=True) + free_vram, _total = self._vram_probe() + # Subtract in-flight reservations from the probe to close the + # TOCTOU window between two concurrent submit_gpu calls. + 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, + ) + 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 + ) + # Subtract in-flight reservations as well — the cluster-mode + # path must also account for tasks whose models are still + # loading (taOS #1705). + 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) + + 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. + + The caller (submit_gpu or _drain_queue) must have already reserved + VRAM via _reserve_and_check. This method releases the reservation + in its finally block. + + The lease claim and _running registration are inside the try block + so that the finally always releases the reservation, even when + claim_lease fails (taOS #1705 — reservation leak fix). + """ + 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: + # Release the VRAM reservation whether we completed, errored, + # or were cancelled. _evict_task handles its own reservation + # release so idempotency matters. + self._release_reservation(task.id) + async with self._running_lock: + entry = self._running.pop(task.id, None) + self._running_tasks.pop(task.id, None) + # 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. + 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) + + 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 — the running GPU task must be evicted so its VRAM + and asyncio resources are freed. + + Returns (cancelled, already_completed) so operators can distinguish + force-kills from natural completions. + """ + 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: + 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: + 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) + # Release the VRAM reservation — the task is being preempted and + # its reserved VRAM should be returned to the pool immediately. + self._release_reservation(task_id) + 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 — stops _run_gpu_task and frees VRAM + if asyncio_task is not None and not asyncio_task.done(): + asyncio_task.cancel() + # Also cancel the arbiter future for queued tasks (belt-and-suspenders) + 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 + + async def _process_queue(self) -> None: + try: + while True: + await asyncio.sleep(2) + if not self._paused: # taOS #796: skip drain while 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. The admitted task is spawned as a background asyncio Task so + the drain loop can continue on the next tick and handle + eviction-to-make-room for higher-priority arrivals. + + Uses _reserve_and_check so the queue processor doesn't race with + concurrent submit_gpu calls — the reservation is atomic with the + admission check (TOCTOU-safe). + """ + 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: + # Try eviction-to-make-room for higher-priority queued tasks. + # evict_lowest_priority(min_priority=N) skips running tasks + # whose priority value is *lower* than N (i.e. tasks that are + # actually higher priority), so only lower-or-equal priority + # running tasks are candidates for eviction. + 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) + # 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), + name=f"gpu-arbiter-drain-{entry.task.id}", + ) + + # Wire result / exception from the background task back to the + # submitter's _arbiter_future once it finishes (or fails). + if future is not None: + def _propagate(ct: asyncio.Task, f: asyncio.Future = future) -> None: + if f.done(): + return + if ct.cancelled(): + return # _evict_task already cancelled the future + 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: + # Not admitted — release reservation (idempotent) and retry later. + self._release_reservation(entry.task.id) + retry.append(entry) + # Re-queue tasks that still can't be admitted + 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")) + + 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), + } + + 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]: + 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 diff --git a/tinyagentos/scheduler/types.py b/tinyagentos/scheduler/types.py index cf4f279c0..eb77c6cb5 100644 --- a/tinyagentos/scheduler/types.py +++ b/tinyagentos/scheduler/types.py @@ -87,6 +87,7 @@ class Task: submitter: str = "unknown" estimated_seconds: float = 1.0 estimated_memory_mb: int = 0 + estimated_vram_mb: int = 0 required_signatures: list[ResourceSignature] = field(default_factory=list) id: str = field(default_factory=lambda: uuid.uuid4().hex[:12]) submitted_at: float = field(default_factory=time.time)