diff --git a/tests/test_cluster.py b/tests/test_cluster.py index 2be9cd264..7c228653b 100644 --- a/tests/test_cluster.py +++ b/tests/test_cluster.py @@ -72,6 +72,22 @@ async def test_heartbeat_revives_offline_worker(self): mgr.heartbeat("gpu-box", load=0.1) assert mgr.get_worker("gpu-box").status == "online" + async def test_heartbeat_does_not_reonline_draining_worker(self): + """A worker mid graceful-drain keeps heartbeating; heartbeat must NOT + flip it back to online or it re-enters routing before leases finish + (taOS #890).""" + mgr = ClusterManager() + w = _make_worker("gpu-box") + await mgr.register_worker(w) + w.status = "draining" + + ok = mgr.heartbeat("gpu-box", load=0.1) + assert ok is True + # Load/last_heartbeat still update, but status stays draining. + updated = mgr.get_worker("gpu-box") + assert updated.status == "draining" + assert updated.load == 0.1 + async def test_heartbeat_timeout_marks_offline(self): mgr = ClusterManager() w = _make_worker("gpu-box") @@ -259,3 +275,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..f1c68a347 --- /dev/null +++ b/tests/test_gpu_arbiter_894.py @@ -0,0 +1,717 @@ +"""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") + + +class TestEvictionOrdering: + """Xid-62 fix: _evict_task awaits the cancelled task so its VRAM is + physically reclaimed BEFORE the reservation is freed and re-admission is + allowed (taOS #894).""" + + @pytest.mark.asyncio + async def test_evict_awaits_task_before_releasing_reservation(self): + arbiter = GpuArbiter(vram_probe=lambda: (8192, 8192), max_queue_size=10) + + started = asyncio.Event() + unloaded = asyncio.Event() + + async def payload(_resource): + started.set() + try: + await asyncio.sleep(60) + finally: + # Stands in for the model's physical unload during teardown. + unloaded.set() + + task = _make_task("t-order", vram_mb=0) + task.payload = payload + + # Reserve against the shared ledger exactly as _reserve_and_check would. + reservation = await arbiter._vram.reserve(4096, caller="gpu-task:t-order") + arbiter._pending_reservations["t-order"] = reservation.reservation_id + assert arbiter._vram.reserved_vram_mb == 4096 + + bg = asyncio.create_task(arbiter._run_gpu_task(task, 4096, True, None)) + await asyncio.wait_for(started.wait(), timeout=5) + + await arbiter._evict_task("t-order") + + # The fix: the task coroutine has fully unwound (unload ran) and the + # reservation is released by the time _evict_task returns — so a + # re-admission cannot pass against not-yet-reclaimed VRAM. + assert unloaded.is_set() + assert bg.done() + assert arbiter._vram.reserved_vram_mb == 0 + assert "t-order" not in arbiter._pending_reservations diff --git a/tests/test_gpu_arbiter_toctou.py b/tests/test_gpu_arbiter_toctou.py new file mode 100644 index 000000000..00fa14ee8 --- /dev/null +++ b/tests/test_gpu_arbiter_toctou.py @@ -0,0 +1,209 @@ +"""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 reservation ledger is the +shared VramReservationManager (taOS #185): it does atomic check-and-reserve 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 in the shared + ledger and be rejected. The manager's atomic reserve closes the + 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._vram.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._vram.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._vram.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 + + # Reserve against the shared ledger exactly as _reserve_and_check would + # before _run_gpu_task is entered. + reservation = await arbiter._vram.reserve(4096, caller="gpu-task:t-evict") + assert reservation is not None + arbiter._pending_reservations["t-evict"] = reservation.reservation_id + assert arbiter._vram.reserved_vram_mb == 4096 + + 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) + + # Evict — cancels the task, awaits it, then releases the reservation. + await arbiter._evict_task("t-evict") + await runner_task + + # After eviction, reservation should be released. + assert arbiter._vram.reserved_vram_mb == 0 + assert "t-evict" not in arbiter._pending_reservations + + @pytest.mark.asyncio + async def test_reservation_subtracted_from_admission_check(self): + """_reserve_and_check sees reduced capacity when VRAM is reserved.""" + arbiter = GpuArbiter( + vram_probe=lambda: (8192, 8192), + max_queue_size=10, + ) + + # Reserve 4000 MiB. + admission = await arbiter._reserve_and_check("t-holder", 4000) + assert admission.admitted + assert admission.free_vram_mb == 8192 # effective free before this reserve + assert arbiter._vram.reserved_vram_mb == 4000 + + # Now effective free = 8192 - 4000 = 4192, so 6000 is rejected. + admission = await arbiter._reserve_and_check("t-big", 6000) + assert not admission.admitted + assert admission.free_vram_mb == 4192 + assert "t-big" not in arbiter._pending_reservations + + # 4000 still fits in the remaining 4192. + admission = await arbiter._reserve_and_check("t-fit", 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._vram.reserved_vram_mb == 5000 + assert "t1" in arbiter._pending_reservations + + # 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._vram.reserved_vram_mb == 5000 + assert "t2" not in arbiter._pending_reservations + + # Release t1 + arbiter._release_reservation("t1") + assert arbiter._vram.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, + ) + + reservation = await arbiter._vram.reserve(1000, caller="gpu-task:t-idem") + arbiter._pending_reservations["t-idem"] = reservation.reservation_id + + arbiter._release_reservation("t-idem") + assert arbiter._vram.reserved_vram_mb == 0 + + # Second call — no-op, no crash. + arbiter._release_reservation("t-idem") + assert arbiter._vram.reserved_vram_mb == 0 + + # Releasing a task that never had a reservation — no-op. + arbiter._release_reservation("non-existent") + assert arbiter._vram.reserved_vram_mb == 0 + + @pytest.mark.asyncio + async def test_reserved_vram_in_stats(self): + """stats() includes reservation info from the shared ledger.""" + 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 + + reservation = await arbiter._vram.reserve(2048, caller="gpu-task:t-stats") + arbiter._pending_reservations["t-stats"] = reservation.reservation_id + + 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 4fd1dca28..8eb50b6f6 100644 --- a/tinyagentos/app.py +++ b/tinyagentos/app.py @@ -61,6 +61,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 from tinyagentos.torrent_settings import TorrentSettingsStore from tinyagentos.relationships import RelationshipManager from tinyagentos.github_identities import GitHubIdentitiesStore @@ -1146,6 +1147,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, @@ -1163,10 +1165,35 @@ async def _reload_llm_proxy_on_catalog_change() -> None: logger.exception("resource scheduler failed to build — routes will use static config") app.state.resource_scheduler = None - # VRAM reservation manager — atomic check-and-reserve so two - # concurrent model loads cannot both pass a VRAM check before - # either consumes physical VRAM (TOCTOU race, taOS #1706). - app.state.vram_reservation = VramReservationManager() + # Single VRAM authority (taOS #185). One VramReservationManager is the + # sole ledger: it does atomic check-and-reserve so two concurrent model + # loads cannot both pass a VRAM check before either consumes physical + # VRAM (TOCTOU, #1706). The GPU arbiter reserves against THIS SAME + # instance for scheduler tasks, and the model-load path + # (routes/models.py) reserves against it too, so there is exactly one + # ledger and no double-counting. + vram_reservation = VramReservationManager() + app.state.vram_reservation = vram_reservation + + # Build the GPU arbiter — VRAM-accounted admission control, queuing, and + # eviction for GPU-bound scheduler workloads, layered on the resource + # scheduler. It shares the reservation ledger above (does not keep its + # own), so admission for tasks and model loads draws from one authority. + try: + gpu_arbiter = GpuArbiter( + scheduler=resource_scheduler if resource_scheduler is not None else None, + cluster_manager=cluster_manager, + vram_reservation=vram_reservation, + 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, shared VRAM ledger)") + 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 @@ -1347,6 +1374,8 @@ async def _web_push_sender(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 50f78babe..a7d9fd87c 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() @@ -203,7 +209,13 @@ def heartbeat( prev_status = worker.status worker.last_heartbeat = time.time() worker.load = load - worker.status = "online" + # A worker mid graceful-drain keeps heartbeating while it finishes + # inflight work. Do NOT flip it back to "online" — that would re-enter + # it into routing/catalog/lease-claims before its leases drain and + # defeat the drain (taOS #890). Leave the drain to complete or be + # cancelled explicitly; every other status re-onlines as before. + if worker.status != "draining": + worker.status = "online" if models is not None: worker.models = models if backends is not None: @@ -317,13 +329,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 @@ -441,10 +455,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 @@ -491,7 +616,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 @@ -534,14 +659,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)") @@ -564,6 +687,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 a559cb7ce..21c2312ff 100644 --- a/tinyagentos/routes/cluster.py +++ b/tinyagentos/routes/cluster.py @@ -1148,3 +1148,108 @@ 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"}, + ) + # Surface a non-2xx worker response as a failure instead of + # reporting status:"updating" with the error body as deploy_result. + resp.raise_for_status() + 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..a01a1bb10 --- /dev/null +++ b/tinyagentos/scheduler/gpu_arbiter.py @@ -0,0 +1,650 @@ +"""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, +) +from tinyagentos.vram_reservation import VramReservationManager + +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, + vram_reservation: VramReservationManager | None = None, + max_queue_size: int = 100, + eviction_enabled: bool = True, + ): + self._scheduler = scheduler + self._cluster_manager = cluster_manager + 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() + + # --- Single VRAM authority (taOS #185) --- + # The arbiter does NOT keep its own VRAM ledger. It reserves against a + # shared VramReservationManager — the same instance the model-load path + # (routes/models.py) uses — so GPU tasks and model loads never + # double-count. The manager does atomic check-and-reserve with an + # nvidia-smi probe run in a worker thread (no event-loop stall) plus a + # TTL sweep for hung reservations, which is why folding it in also + # closes the arbiter's own probe-under-lock and stale-hold gaps. + # + # When a legacy ``vram_probe`` (free, total) is supplied (tests/standalone), + # adapt it to the manager's ``(free, total) | None`` convention: a total + # of 0 means "no local GPU" so admission routes to the cluster instead + # of fail-open-admitting locally. + if vram_reservation is not None: + self._vram = vram_reservation + else: + probe_adapter: Callable[[], tuple[int, int] | None] | None = None + if vram_probe is not None: + def probe_adapter() -> tuple[int, int] | None: + free_mb, total_mb = vram_probe() + return (free_mb, total_mb) if total_mb > 0 else None + self._vram = VramReservationManager(probe=probe_adapter) + # task_id -> reservation_id held in the shared ledger (None-safe pop). + self._pending_reservations: dict[str, str] = {} + + 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 "" + # compute_cap is the authoritative token (e.g. "sm_86") — match it + # exactly so "sm_8" cannot prefix-match "sm_86"; the freeform model + # string keeps a substring match for vendor-name requirements. + if required_gpu_arch == cc or required_gpu_arch in gpu_model: + if resource_id is None or self._resource_on_worker(resource_id, worker.name): + return True, None + return False, f"no online worker with GPU architecture {required_gpu_arch}" + + def _resource_on_worker(self, resource_id: str, worker_name: str) -> bool: + """True when *resource_id* names a resource on *worker_name*. + + Uses the cluster manager's exact parse when available so that + ``gpu-node`` does not match ``gpu-node-2`` (the same prefix-collision + #1726 fixed on the lease paths); falls back to a colon-delimited match. + """ + if self._cluster_manager is not None: + parsed = self._cluster_manager._parse_resource_id(resource_id) + if parsed is not None: + return parsed[0] == worker_name + return resource_id.startswith(worker_name + ":") + + # ── 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, an eviction of a queued-but-not-yet-admitted entry, + or a task whose _run_gpu_task finally already released it). + """ + reservation_id = self._pending_reservations.pop(task_id, None) + if reservation_id is not None: + self._vram.release(reservation_id) + logger.debug("gpu-arbiter: released reservation %s for task %s", reservation_id, task_id) + + async def _reserve_and_check(self, task_id: str, required_vram_mb: int) -> GpuAdmission: + """Atomically check admission and reserve VRAM if admitted. + + Reserves against the shared VramReservationManager (the single VRAM + authority, taOS #185), whose ``reserve`` is itself atomic + check-and-reserve, so two concurrent callers cannot both pass before + either consumes physical VRAM. Local admission is decided by the + manager; when the host has no local GPU probe the request routes to + the cluster instead. + + 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) + + # available_vram() reports (effective_free, total); total==0 means no + # local GPU probe. Run it in a thread — nvidia-smi is a blocking + # subprocess and this is on the admission hot path. + effective_free, total = await asyncio.to_thread(self._vram.available_vram) + if total > 0: + # Local GPU present — atomic check-and-reserve on the shared ledger. + reservation = await self._vram.reserve(required_vram_mb, caller=f"gpu-task:{task_id}") + if reservation is not None: + self._pending_reservations[task_id] = reservation.reservation_id + return GpuAdmission(admitted=True, free_vram_mb=effective_free, + required_vram_mb=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"), + ) + # No local GPU probe: try the cluster. + if self._cluster_manager is not None: + return self._check_cluster_admission(required_vram_mb) + # No local GPU and no cluster: fail open with bookkeeping, matching the + # reservation manager's no-probe convention (cluster claim_lease parity). + reservation = await self._vram.reserve(required_vram_mb, caller=f"gpu-task:{task_id}") + if reservation is not None: + self._pending_reservations[task_id] = reservation.reservation_id + return GpuAdmission(admitted=True, required_vram_mb=required_vram_mb) + + async def reserve_vram(self, vram_mb: int, caller: str = ""): + """Front-door VRAM reservation for non-task callers (e.g. model loads). + + Reserves against the arbiter's single shared ledger so model loads and + GPU scheduler tasks draw from ONE authority and never double-count + (taOS #185). Returns a ``VramReservation`` or ``None`` when there is + not enough VRAM; the caller MUST call :meth:`release_vram` when done. + """ + return await self._vram.reserve(vram_mb, caller=caller) + + def release_vram(self, reservation_id: str) -> bool: + """Release a reservation taken via :meth:`reserve_vram`. Idempotent.""" + return self._vram.release(reservation_id) + + # ── 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_cluster_admission(self, required_vram_mb: int) -> GpuAdmission: + """Check whether a cluster worker can host *required_vram_mb*. + + Used only when the local host has no GPU probe. Subtracts each + worker's active GPU leases and the shared ledger's in-flight + reservations so a task whose model is still loading is accounted for + (taOS #1705). + """ + if self._cluster_manager is None: + return GpuAdmission(admitted=True, required_vram_mb=required_vram_mb) + leases = self._cluster_manager.get_leases() + reserved = self._vram.reserved_vram_mb + 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 self._resource_on_worker(l.resource_id, worker.name) and l.required_vram_mb > 0 + ) + available = worker.free_vram_mb - worker_leases - reserved + 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", + ) + + 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(): + # Only STRICTLY lower-priority running tasks are eviction + # candidates. With "lower int = higher priority", skip anything + # whose priority value is <= the requester's so an equal-priority + # task is never preempted (avoids eviction thrash between peers). + 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) + # ORDERING MATTERS (Xid-62 fix, taOS #894): cancel the running task and + # AWAIT it to completion BEFORE releasing the VRAM reservation and lease. + # cancel() is cooperative — the model is still physically resident when + # cancel() returns. If we freed the reservation first, an evict-to-make- + # room re-admission could pass against not-yet-reclaimed VRAM and start a + # second concurrent load (the exact driver crash the arbiter prevents). + # Awaiting lets the task's coroutine unwind (its finally releases its own + # reservation) so the capacity is genuinely free before re-admission. + if (asyncio_task is not None and not asyncio_task.done() + and asyncio_task is not asyncio.current_task()): + asyncio_task.cancel() + try: + await asyncio_task + except (asyncio.CancelledError, Exception): + # Swallow the cancellation and any teardown error — the task is + # being force-evicted; its result/exception is intentionally + # discarded and surfaced to the submitter via _arbiter_future. + pass + # Now the VRAM is physically reclaimed. Release the reservation + # (idempotent — the task's own finally may already have) and the lease + # (the task skipped it because we popped _running first). + 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) + # 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._vram.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) diff --git a/tinyagentos/vram_reservation.py b/tinyagentos/vram_reservation.py index e248aafd7..f6e481368 100644 --- a/tinyagentos/vram_reservation.py +++ b/tinyagentos/vram_reservation.py @@ -74,11 +74,16 @@ class VramReservationManager: def __init__( self, ttl_seconds: float = DEFAULT_RESERVATION_TTL_SECONDS, + probe: "Callable[[], tuple[int, int] | None] | None" = None, ) -> None: self._lock = asyncio.Lock() self._reserved_vram_mb: int = 0 self._pending: dict[str, VramReservation] = {} self._ttl_seconds = float(ttl_seconds) + # An injectable probe overrides nvidia-smi (used by the GPU arbiter to + # share one ledger with a test-supplied probe, taOS #185/#894). It must + # return ``(free_mb, total_mb)`` or ``None`` when VRAM cannot be read. + self._probe_override = probe # ── public API ────────────────────────────────────────────────── @@ -264,15 +269,16 @@ def _sweep_stale_unlocked(self, now: float | None = None) -> int: self._reserved_vram_mb = 0 return len(stale_ids) - @staticmethod - def _probe_vram() -> tuple[int, int] | None: - """Probe real-time free/total VRAM via nvidia-smi. + def _probe_vram(self) -> tuple[int, int] | None: + """Probe real-time free/total VRAM via nvidia-smi (or an injected probe). Returns ``(free_mb, total_mb)``, or ``None`` when no nvidia-smi is available or the probe fails. ``None`` (cannot know) is deliberately distinct from ``(0, 0)`` (known-full): admission fails open on ``None``, matching cluster claim_lease semantics. """ + if self._probe_override is not None: + return self._probe_override() try: from tinyagentos.system_stats import read_nvidia_vram