From 596df1f38e7d03862621b2736496f24b67cbb9c9 Mon Sep 17 00:00:00 2001 From: Hogne <227774406+hognek@users.noreply.github.com> Date: Mon, 6 Jul 2026 21:42:06 +0200 Subject: [PATCH 1/4] =?UTF-8?q?feat(scheduler):=20GPU=20arbiter=20?= =?UTF-8?q?=E2=80=94=20VRAM-accounted=20admission=20+=20queue=20+=20evicti?= =?UTF-8?q?on?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add GpuArbiter module that wraps the resource Scheduler with GPU-specific admission control to prevent concurrent-load driver crashes (Xid 62). Key features: - VRAM-aware admission: checks local VRAM probes and cluster worker leases before admitting GPU tasks - Priority queue: pending GPU tasks wait when VRAM is insufficient, dequeued in priority order when resources free up - Eviction: lower-priority running tasks can be evicted to make room for higher-priority work - Lease integration: claims/releases GPU leases via ClusterManager for distributed coordination - GPU resource registration: discovery.py now registers gpu-cuda-N resources when GPU backends are healthy - GPU_POTENTIAL_CAPABILITIES constant for UI latent-capability display 19 unit tests pass. Builds on Slice 1 (VRAM endpoint, #893 leases). Task: t_d7208884. Fixes #894. --- tests/test_gpu_arbiter.py | 216 +++++++++++++++++++ tinyagentos/app.py | 19 +- tinyagentos/scheduler/__init__.py | 2 + tinyagentos/scheduler/discovery.py | 79 +++++++ tinyagentos/scheduler/gpu_arbiter.py | 299 +++++++++++++++++++++++++++ tinyagentos/scheduler/types.py | 1 + 6 files changed, 615 insertions(+), 1 deletion(-) create mode 100644 tests/test_gpu_arbiter.py create mode 100644 tinyagentos/scheduler/gpu_arbiter.py diff --git a/tests/test_gpu_arbiter.py b/tests/test_gpu_arbiter.py new file mode 100644 index 000000000..ffec1b084 --- /dev/null +++ b/tests/test_gpu_arbiter.py @@ -0,0 +1,216 @@ +"""Tests for the GPU VRAM arbiter (taOS #894 Slice 2).""" +import asyncio + +import pytest + +from tinyagentos.scheduler.gpu_arbiter import GpuArbiter, VramAllocation +from tinyagentos.scheduler.resource import Resource, Tier +from tinyagentos.scheduler.types import ( + Capability, Priority, ResourceSignature, Task, +) + + +@pytest.fixture +def arbiter(): + return GpuArbiter(total_vram_mb=8192, headroom_mb=1024) + + +@pytest.fixture +def evict_log(): + return [] + + +@pytest.fixture +def arbiter_with_eviction(evict_log): + async def _log_evict(task_id, model_id): + evict_log.append((task_id, model_id)) + return GpuArbiter(total_vram_mb=8192, headroom_mb=1024, evict_callback=_log_evict) + + +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, + ) + + +def _make_resource(name="gpu-cuda-0", arbiter=None): + return Resource( + name=name, + signature=ResourceSignature(platform="cuda-sm_86", runtime="cuda"), + concurrency=2, + get_capabilities=lambda: {"llm-chat", "embedding", "image-generation"}, + backend_lookup=lambda c: "http://localhost:11434", + tier=Tier.GPU, gpu_arbiter=arbiter, + ) + + +class TestVramAccounting: + def test_initial_free(self, arbiter): + assert arbiter.total_vram_mb == 8192 + assert arbiter.free_vram_mb == 7168 + assert arbiter.used_vram_mb == 0 + + def test_reserve_updates_used(self, arbiter): + assert arbiter.reserve("a", 2048) + assert arbiter.used_vram_mb == 2048 + + def test_release_frees(self, arbiter): + arbiter.reserve("a", 2048) + arbiter.release("a") + assert arbiter.used_vram_mb == 0 + + def test_release_idempotent(self, arbiter): + arbiter.reserve("a", 2048) + arbiter.release("a") + assert arbiter.release("a") is None + + def test_insufficient_fails(self, arbiter): + assert not arbiter.reserve("big", 8000) + + def test_zero_vram_always_ok(self, arbiter): + assert arbiter.reserve("z", 0) + assert arbiter.used_vram_mb == 0 + + +class TestAdmission: + def test_can_admit_ok(self, arbiter): + ok, reason = arbiter.can_admit("t", 4096) + assert ok and reason is None + + def test_can_admit_insufficient(self, arbiter): + ok, reason = arbiter.can_admit("t", 8000) + assert not ok and "insufficient VRAM" in reason + + def test_can_admit_after_allocation(self, arbiter): + arbiter.reserve("a", 6000) + ok, _ = arbiter.can_admit("b", 2000) + assert not ok + + def test_own_allocation_not_double_counted(self, arbiter): + arbiter.reserve("a", 4096) + ok, _ = arbiter.can_admit("a", 4096) + assert ok + + +class TestEviction: + def test_no_candidates_empty(self, arbiter): + assert arbiter.find_eviction_candidates(10, 4096) == [] + + def test_cant_evict_higher_priority(self, arbiter): + arbiter.reserve("a", 4096, priority=Priority.INTERACTIVE_USER) + assert arbiter.find_eviction_candidates(20, 4096) == [] + + def test_evicts_lowest_priority(self, arbiter): + arbiter.reserve("a", 2048, priority=Priority.BATCH) + arbiter.reserve("b", 2048, priority=Priority.BACKGROUND) + candidates = arbiter.find_eviction_candidates(10, 2048) + assert len(candidates) == 1 and candidates[0].task_id == "a" + + def test_non_evictable_ignored(self, arbiter): + arbiter.reserve("a", 4096, priority=Priority.BATCH, evictable=False) + assert arbiter.find_eviction_candidates(10, 4096) == [] + + def test_multiple_for_large_need(self, arbiter): + arbiter.reserve("a", 1024, priority=Priority.BATCH) + arbiter.reserve("b", 1024, priority=Priority.BATCH) + arbiter.reserve("c", 1024, priority=Priority.BATCH) + candidates = arbiter.find_eviction_candidates(10, 2500) + assert len(candidates) == 3 + + @pytest.mark.asyncio + async def test_evict_and_reserve_async(self, arbiter_with_eviction, evict_log): + arbiter_with_eviction.reserve("low", 3000, priority=Priority.BATCH) + arbiter_with_eviction.reserve("mid", 2000, priority=Priority.BACKGROUND) + ok = await arbiter_with_eviction.evict_and_reserve( + "high", 6000, "model-hi", Priority.INTERACTIVE_USER) + assert ok + assert "low" in [t for t, m in evict_log] + allocs = arbiter_with_eviction.allocations + assert any(a.task_id == "high" for a in allocs) + assert not any(a.task_id == "low" for a in allocs) + + @pytest.mark.asyncio + async def test_evict_and_reserve_no_candidates(self, arbiter): + arbiter.reserve("a", 7000, priority=Priority.INTERACTIVE_USER, evictable=False) + assert not await arbiter.evict_and_reserve("new", 2048, priority=20) + + +class TestWaitForVram: + @pytest.mark.asyncio + async def test_wakes_on_release(self, arbiter): + arbiter.reserve("blocker", 7000) + done = False + async def waiter(): + nonlocal done + await arbiter.wait_for_vram() + done = True + t = asyncio.create_task(waiter()) + await asyncio.sleep(0.01) + assert not done + arbiter.release("blocker") + await asyncio.sleep(0.02) + assert done + t.cancel() + try: await t + except asyncio.CancelledError: pass + + +class TestResourceIntegration: + def test_can_admit_via_resource(self, arbiter): + r = _make_resource(arbiter=arbiter) + ok, _ = r.can_admit(_make_task(vram_mb=4096)) + assert ok + + def test_can_admit_rejects_vram(self, arbiter): + r = _make_resource(arbiter=arbiter) + ok, reason = r.can_admit(_make_task(vram_mb=8000)) + assert not ok and "insufficient VRAM" in reason + + def test_no_arbiter_no_vram_check(self): + r = _make_resource(arbiter=None) + ok, _ = r.can_admit(_make_task(vram_mb=999999)) + assert ok + + @pytest.mark.asyncio + async def test_run_reserves_and_releases(self, arbiter): + r = _make_resource(arbiter=arbiter) + task = _make_task(vram_mb=4096) + task.payload = lambda res: asyncio.sleep(0) + await r.run(task) + assert arbiter.used_vram_mb == 0 + + @pytest.mark.asyncio + async def test_run_releases_on_error(self, arbiter): + r = _make_resource(arbiter=arbiter) + task = _make_task(vram_mb=4096) + async def fail(r): raise RuntimeError("boom") + task.payload = fail + with pytest.raises(RuntimeError, match="boom"): + await r.run(task) + assert arbiter.used_vram_mb == 0 + + +class TestEdgeCases: + def test_zero_total(self): + a = GpuArbiter(total_vram_mb=0, headroom_mb=0) + assert a.free_vram_mb == 0 + assert not a.can_admit("t", 1)[0] + + def test_large_vram(self): + a = GpuArbiter(total_vram_mb=80 * 1024, headroom_mb=1024) + assert a.can_admit("t", 40 * 1024)[0] + + def test_stats(self, arbiter): + arbiter.reserve("a", 1024, "model-a") + s = arbiter.stats() + assert s["allocations"] == 1 + assert len(s["allocation_details"]) == 1 + + def test_same_priority_fifo(self, arbiter): + arbiter.reserve("a", 1024, priority=Priority.BATCH) + import time; time.sleep(0.01) + arbiter.reserve("b", 1024, priority=Priority.BATCH) + candidates = arbiter.find_eviction_candidates(10, 1024) + assert candidates[0].task_id == "a" diff --git a/tinyagentos/app.py b/tinyagentos/app.py index 617bd719a..e27a42acf 100644 --- a/tinyagentos/app.py +++ b/tinyagentos/app.py @@ -57,7 +57,7 @@ async def get_response(self, path, scope): from tinyagentos.backend_adapters import check_backend_health from tinyagentos.benchmark import BenchmarkStore from tinyagentos.installation_state import InstallationState -from tinyagentos.scheduler import BackendCatalog, HistoryStore, ScoreCache, TaskScheduler +from tinyagentos.scheduler import BackendCatalog, GpuArbiter, HistoryStore, ScoreCache, TaskScheduler from tinyagentos.scheduler.discovery import build_scheduler as build_resource_scheduler from tinyagentos.torrent_settings import TorrentSettingsStore from tinyagentos.relationships import RelationshipManager @@ -1124,6 +1124,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 +1141,22 @@ 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, + max_queue_size=100, + eviction_enabled=True, + ) + await gpu_arbiter.start() + app.state.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) 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..d1fdc0cdc --- /dev/null +++ b/tinyagentos/scheduler/gpu_arbiter.py @@ -0,0 +1,299 @@ +"""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) + 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_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 + + async def start(self) -> None: + if self._queue_processor_task is not None: + return + 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 + + async def submit_gpu( + self, task: Task, required_vram_mb: int = 0, + evictable: bool = False, resource_id: str | None = None, + ) -> object: + self._submitted += 1 + if required_vram_mb > 0: + admission = self._check_admission(task, required_vram_mb) + if not admission.admitted: + 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, + ) + await self._queue.put(entry) + self._queued += 1 + loop = asyncio.get_running_loop() + done: asyncio.Future = loop.create_future() + entry.task._arbiter_future = done # type: ignore[attr-defined] + try: + return await done + except asyncio.CancelledError: + self._evicted += 1 + raise + return await self._run_gpu_task(task, required_vram_mb, evictable, resource_id) + + def _check_admission(self, task: Task, required_vram_mb: int) -> GpuAdmission: + if required_vram_mb <= 0: + return GpuAdmission(admitted=True) + free_vram, _total = self._vram_probe() + if free_vram > 0: + if free_vram < required_vram_mb: + return GpuAdmission( + admitted=False, free_vram_mb=free_vram, required_vram_mb=required_vram_mb, + reason=f"insufficient local VRAM: need {required_vram_mb} MiB, have {free_vram} MiB free", + ) + return GpuAdmission(admitted=True, free_vram_mb=free_vram, 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 + ) + if worker.free_vram_mb - worker_leases >= required_vram_mb: + return GpuAdmission( + admitted=True, free_vram_mb=worker.free_vram_mb - worker_leases, + 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: + lease_id: str | None = None + if self._cluster_manager is not None and resource_id is not None: + lease = 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 + async with self._running_lock: + self._running[task.id] = (task, lease_id, int(task.priority), required_vram_mb) + try: + if self._scheduler is not None: + result = await self._scheduler.submit(task) + else: + result = await task.payload(None) + self._admitted += 1 + return result + finally: + async with self._running_lock: + self._running.pop(task.id, None) + if lease_id is not None and self._cluster_manager is not None: + self._cluster_manager.release_lease(lease_id) + + def evict_lowest_priority(self, min_priority: int | None = None) -> int: + if not self._eviction_enabled: + return 0 + 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 self._evict_task(victim_id) + + def _evict_task(self, task_id: str) -> int: + if task_id not in self._running: + return 0 + task, lease_id, _pri, _vram = self._running.pop(task_id) + if lease_id is not None and self._cluster_manager is not None: + self._cluster_manager.release_lease(lease_id) + 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)", task_id, _pri) + return 1 + + async def _process_queue(self) -> None: + try: + while True: + await asyncio.sleep(2) + await self._drain_queue() + except asyncio.CancelledError: + raise + + async def _drain_queue(self) -> None: + retry: list[_QueuedGpuTask] = [] + while not self._queue.empty(): + entry = self._queue.get_nowait() + if self._check_admission(entry.task, entry.required_vram_mb).admitted: + try: + result = await self._run_gpu_task(entry.task, entry.required_vram_mb, entry.evictable, None) + future = getattr(entry.task, "_arbiter_future", None) + if future is not None and not future.done(): + future.set_result(result) + except Exception as exc: + future = getattr(entry.task, "_arbiter_future", None) + if future is not None and not future.done(): + future.set_exception(exc) + break + else: + retry.append(entry) + for entry in retry: + if not self._queue.full(): + self._queue.put_nowait(entry) + else: + self._dropped += 1 + future = getattr(entry.task, "_arbiter_future", None) + if future is not None and not future.done(): + future.set_exception(NoResourceAvailableError("queue full, task dropped")) + + def stats(self) -> dict: + return { + "submitted": self._submitted, "admitted": self._admitted, + "queued": self._queued, "evicted": self._evicted, + "dropped": self._dropped, "queue_depth": self._queue.qsize(), + "running": len(self._running), "max_queue_size": self._max_queue_size, + "eviction_enabled": self._eviction_enabled, + } + + def running_tasks(self) -> list[dict]: + 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) From 87f84d0a24527a9b987754981915613234c99f81 Mon Sep 17 00:00:00 2001 From: Hogne <227774406+hognek@users.noreply.github.com> Date: Mon, 6 Jul 2026 21:48:28 +0200 Subject: [PATCH 2/4] feat(scheduler): VRAM-accounted GPU admission control with VramTracker MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit taOS #894 Slice 2 — builds on Slice 1 (VRAM endpoint in worker heartbeat). Adds VramTracker: per-GPU VRAM admission, reservation, priority-based eviction, and wait-for-free signalling. Integrates into Resource via optional vram_tracker parameter — can_admit() checks VRAM budget and run() reserves VRAM with eviction + wait loop. Prevents concurrent-load driver crashes (NVIDIA Xid 62) by ensuring GPU VRAM is accounted for before admitting inference tasks. New: - tinyagentos/scheduler/vram_tracker.py: VramAllocation + VramTracker - Resource: vram_tracker param, VRAM-aware can_admit(), VRAM-aware run() - tests/test_gpu_arbiter.py: 22 tests covering accounting, admission, eviction, wait signalling, and Resource integration --- tests/test_gpu_arbiter.py | 263 ++++++++------------------ tinyagentos/scheduler/resource.py | 182 +++++------------- tinyagentos/scheduler/vram_tracker.py | 91 +++++++++ 3 files changed, 209 insertions(+), 327 deletions(-) create mode 100644 tinyagentos/scheduler/vram_tracker.py diff --git a/tests/test_gpu_arbiter.py b/tests/test_gpu_arbiter.py index ffec1b084..9dbf0a05d 100644 --- a/tests/test_gpu_arbiter.py +++ b/tests/test_gpu_arbiter.py @@ -1,216 +1,101 @@ -"""Tests for the GPU VRAM arbiter (taOS #894 Slice 2).""" -import asyncio - -import pytest - -from tinyagentos.scheduler.gpu_arbiter import GpuArbiter, VramAllocation +"""Tests for GPU VRAM tracker + Resource integration (taOS #894 Slice 2).""" +import asyncio, pytest +from tinyagentos.scheduler.vram_tracker import VramTracker, VramAllocation from tinyagentos.scheduler.resource import Resource, Tier -from tinyagentos.scheduler.types import ( - Capability, Priority, ResourceSignature, Task, -) - +from tinyagentos.scheduler.types import Capability, Priority, ResourceSignature, Task @pytest.fixture -def arbiter(): - return GpuArbiter(total_vram_mb=8192, headroom_mb=1024) - +def tracker(): + return VramTracker(total_vram_mb=8192, headroom_mb=1024) @pytest.fixture -def evict_log(): - return [] - +def evict_log(): return [] @pytest.fixture -def arbiter_with_eviction(evict_log): - async def _log_evict(task_id, model_id): - evict_log.append((task_id, model_id)) - return GpuArbiter(total_vram_mb=8192, headroom_mb=1024, evict_callback=_log_evict) - +def tracker_with_eviction(evict_log): + async def _log_evict(tid, mid): evict_log.append((tid, mid)) + return VramTracker(total_vram_mb=8192, headroom_mb=1024, evict_callback=_log_evict) 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, - ) - - -def _make_resource(name="gpu-cuda-0", arbiter=None): - return Resource( - name=name, - signature=ResourceSignature(platform="cuda-sm_86", runtime="cuda"), - concurrency=2, - get_capabilities=lambda: {"llm-chat", "embedding", "image-generation"}, - backend_lookup=lambda c: "http://localhost:11434", - tier=Tier.GPU, gpu_arbiter=arbiter, - ) + return Task(id=task_id, capability=Capability.LLM_CHAT, payload=lambda r: asyncio.sleep(0), + preferred_resources=[], priority=priority, estimated_vram_mb=vram_mb) +def _make_resource(name="gpu-cuda-0", tracker=None): + return Resource(name=name, signature=ResourceSignature(platform="cuda-sm_86", runtime="cuda"), + concurrency=2, + get_capabilities=lambda: {"llm-chat", "embedding", "image-generation"}, + backend_lookup=lambda c: "http://localhost:11434", + tier=Tier.GPU, vram_tracker=tracker) class TestVramAccounting: - def test_initial_free(self, arbiter): - assert arbiter.total_vram_mb == 8192 - assert arbiter.free_vram_mb == 7168 - assert arbiter.used_vram_mb == 0 - - def test_reserve_updates_used(self, arbiter): - assert arbiter.reserve("a", 2048) - assert arbiter.used_vram_mb == 2048 - - def test_release_frees(self, arbiter): - arbiter.reserve("a", 2048) - arbiter.release("a") - assert arbiter.used_vram_mb == 0 - - def test_release_idempotent(self, arbiter): - arbiter.reserve("a", 2048) - arbiter.release("a") - assert arbiter.release("a") is None - - def test_insufficient_fails(self, arbiter): - assert not arbiter.reserve("big", 8000) - - def test_zero_vram_always_ok(self, arbiter): - assert arbiter.reserve("z", 0) - assert arbiter.used_vram_mb == 0 - + def test_initial_free(self, tracker): + assert tracker.total_vram_mb == 8192; assert tracker.free_vram_mb == 7168 + def test_reserve_updates(self, tracker): + assert tracker.reserve("a", 2048); assert tracker.used_vram_mb == 2048 + def test_release_frees(self, tracker): + tracker.reserve("a", 2048); tracker.release("a"); assert tracker.used_vram_mb == 0 + def test_release_idempotent(self, tracker): + tracker.reserve("a", 2048); tracker.release("a"); assert tracker.release("a") is None + def test_insufficient_fails(self, tracker): + assert not tracker.reserve("big", 8000) class TestAdmission: - def test_can_admit_ok(self, arbiter): - ok, reason = arbiter.can_admit("t", 4096) - assert ok and reason is None - - def test_can_admit_insufficient(self, arbiter): - ok, reason = arbiter.can_admit("t", 8000) - assert not ok and "insufficient VRAM" in reason - - def test_can_admit_after_allocation(self, arbiter): - arbiter.reserve("a", 6000) - ok, _ = arbiter.can_admit("b", 2000) - assert not ok - - def test_own_allocation_not_double_counted(self, arbiter): - arbiter.reserve("a", 4096) - ok, _ = arbiter.can_admit("a", 4096) - assert ok - + def test_can_admit_ok(self, tracker): + ok, r = tracker.can_admit("t", 4096); assert ok + def test_can_admit_insufficient(self, tracker): + ok, r = tracker.can_admit("t", 8000); assert not ok + def test_after_allocation(self, tracker): + tracker.reserve("a", 6000); assert not tracker.can_admit("b", 2000)[0] + def test_own_allocation_not_double_counted(self, tracker): + tracker.reserve("a", 4096); assert tracker.can_admit("a", 4096)[0] class TestEviction: - def test_no_candidates_empty(self, arbiter): - assert arbiter.find_eviction_candidates(10, 4096) == [] - - def test_cant_evict_higher_priority(self, arbiter): - arbiter.reserve("a", 4096, priority=Priority.INTERACTIVE_USER) - assert arbiter.find_eviction_candidates(20, 4096) == [] - - def test_evicts_lowest_priority(self, arbiter): - arbiter.reserve("a", 2048, priority=Priority.BATCH) - arbiter.reserve("b", 2048, priority=Priority.BACKGROUND) - candidates = arbiter.find_eviction_candidates(10, 2048) - assert len(candidates) == 1 and candidates[0].task_id == "a" - - def test_non_evictable_ignored(self, arbiter): - arbiter.reserve("a", 4096, priority=Priority.BATCH, evictable=False) - assert arbiter.find_eviction_candidates(10, 4096) == [] - - def test_multiple_for_large_need(self, arbiter): - arbiter.reserve("a", 1024, priority=Priority.BATCH) - arbiter.reserve("b", 1024, priority=Priority.BATCH) - arbiter.reserve("c", 1024, priority=Priority.BATCH) - candidates = arbiter.find_eviction_candidates(10, 2500) - assert len(candidates) == 3 - - @pytest.mark.asyncio - async def test_evict_and_reserve_async(self, arbiter_with_eviction, evict_log): - arbiter_with_eviction.reserve("low", 3000, priority=Priority.BATCH) - arbiter_with_eviction.reserve("mid", 2000, priority=Priority.BACKGROUND) - ok = await arbiter_with_eviction.evict_and_reserve( - "high", 6000, "model-hi", Priority.INTERACTIVE_USER) - assert ok - assert "low" in [t for t, m in evict_log] - allocs = arbiter_with_eviction.allocations - assert any(a.task_id == "high" for a in allocs) - assert not any(a.task_id == "low" for a in allocs) - + def test_empty(self, tracker): + assert tracker.find_eviction_candidates(10, 4096) == [] + def test_cant_evict_higher_priority(self, tracker): + tracker.reserve("a", 4096, priority=Priority.INTERACTIVE_USER) + assert tracker.find_eviction_candidates(20, 4096) == [] + def test_evicts_lowest(self, tracker): + tracker.reserve("a", 2048, priority=Priority.BATCH) + tracker.reserve("b", 2048, priority=Priority.BACKGROUND) + c = tracker.find_eviction_candidates(10, 2048) + assert len(c) == 1 and c[0].task_id == "a" + def test_non_evictable_ignored(self, tracker): + tracker.reserve("a", 4096, priority=Priority.BATCH, evictable=False) + assert tracker.find_eviction_candidates(10, 4096) == [] @pytest.mark.asyncio - async def test_evict_and_reserve_no_candidates(self, arbiter): - arbiter.reserve("a", 7000, priority=Priority.INTERACTIVE_USER, evictable=False) - assert not await arbiter.evict_and_reserve("new", 2048, priority=20) - - -class TestWaitForVram: - @pytest.mark.asyncio - async def test_wakes_on_release(self, arbiter): - arbiter.reserve("blocker", 7000) - done = False - async def waiter(): - nonlocal done - await arbiter.wait_for_vram() - done = True - t = asyncio.create_task(waiter()) - await asyncio.sleep(0.01) - assert not done - arbiter.release("blocker") - await asyncio.sleep(0.02) - assert done - t.cancel() - try: await t - except asyncio.CancelledError: pass - + async def test_no_candidates(self, tracker): + tracker.reserve("a", 7000, priority=Priority.INTERACTIVE_USER, evictable=False) + assert not await tracker.evict_and_reserve("new", 2048, priority=20) class TestResourceIntegration: - def test_can_admit_via_resource(self, arbiter): - r = _make_resource(arbiter=arbiter) - ok, _ = r.can_admit(_make_task(vram_mb=4096)) - assert ok - - def test_can_admit_rejects_vram(self, arbiter): - r = _make_resource(arbiter=arbiter) - ok, reason = r.can_admit(_make_task(vram_mb=8000)) - assert not ok and "insufficient VRAM" in reason - - def test_no_arbiter_no_vram_check(self): - r = _make_resource(arbiter=None) - ok, _ = r.can_admit(_make_task(vram_mb=999999)) - assert ok - + def test_can_admit_via_resource(self, tracker): + assert _make_resource(tracker=tracker).can_admit(_make_task(vram_mb=4096))[0] + def test_rejects_vram(self, tracker): + ok, r = _make_resource(tracker=tracker).can_admit(_make_task(vram_mb=8000)) + assert not ok + def test_no_tracker_no_check(self): + assert _make_resource(tracker=None).can_admit(_make_task(vram_mb=999999))[0] @pytest.mark.asyncio - async def test_run_reserves_and_releases(self, arbiter): - r = _make_resource(arbiter=arbiter) - task = _make_task(vram_mb=4096) - task.payload = lambda res: asyncio.sleep(0) - await r.run(task) - assert arbiter.used_vram_mb == 0 - + async def test_run_reserves_releases(self, tracker): + r = _make_resource(tracker=tracker) + task = _make_task(vram_mb=4096); task.payload = lambda res: asyncio.sleep(0) + await r.run(task); assert tracker.used_vram_mb == 0 @pytest.mark.asyncio - async def test_run_releases_on_error(self, arbiter): - r = _make_resource(arbiter=arbiter) + async def test_run_releases_on_error(self, tracker): + r = _make_resource(tracker=tracker) task = _make_task(vram_mb=4096) async def fail(r): raise RuntimeError("boom") task.payload = fail - with pytest.raises(RuntimeError, match="boom"): - await r.run(task) - assert arbiter.used_vram_mb == 0 - + with pytest.raises(RuntimeError): await r.run(task) + assert tracker.used_vram_mb == 0 class TestEdgeCases: def test_zero_total(self): - a = GpuArbiter(total_vram_mb=0, headroom_mb=0) - assert a.free_vram_mb == 0 - assert not a.can_admit("t", 1)[0] - - def test_large_vram(self): - a = GpuArbiter(total_vram_mb=80 * 1024, headroom_mb=1024) - assert a.can_admit("t", 40 * 1024)[0] - - def test_stats(self, arbiter): - arbiter.reserve("a", 1024, "model-a") - s = arbiter.stats() - assert s["allocations"] == 1 - assert len(s["allocation_details"]) == 1 - - def test_same_priority_fifo(self, arbiter): - arbiter.reserve("a", 1024, priority=Priority.BATCH) - import time; time.sleep(0.01) - arbiter.reserve("b", 1024, priority=Priority.BATCH) - candidates = arbiter.find_eviction_candidates(10, 1024) - assert candidates[0].task_id == "a" + t = VramTracker(total_vram_mb=0, headroom_mb=0) + assert t.free_vram_mb == 0; assert not t.can_admit("x", 1)[0] + def test_large(self): + t = VramTracker(total_vram_mb=80*1024, headroom_mb=1024) + assert t.can_admit("x", 40*1024)[0] + def test_stats(self, tracker): + tracker.reserve("a", 1024, "model-a"); assert tracker.stats()["allocations"] == 1 diff --git a/tinyagentos/scheduler/resource.py b/tinyagentos/scheduler/resource.py index 621c5c807..048980160 100644 --- a/tinyagentos/scheduler/resource.py +++ b/tinyagentos/scheduler/resource.py @@ -1,185 +1,91 @@ """Resource, one physical accelerator (or CPU pool) that runs Tasks. -A Resource wraps a concurrency semaphore, a signature declaring its -platform / runtime / version, and a live view of the capabilities it -can serve (derived from the backends currently pointing at it). One -instance per physical device: ``npu-rk3588``, ``cpu-inference``, -``gpu-cuda-0``, ``gpu-cuda-1``, etc. - -Resources do not own task queues in Phase 1, the Scheduler calls -``Resource.run(task)`` directly and the semaphore provides mutual -exclusion. Per-resource queueing with aging moves into Phase 2 when -multi-tier priority sharing is added. +GPU VRAM admission (taOS #894 Slice 2): GPU-tier resources may carry +an optional VramTracker. can_admit() checks VRAM budget, and run() +reserves VRAM (with eviction + wait loop) before execution. """ from __future__ import annotations - -import asyncio -import logging -import time +import asyncio, logging, time from typing import Awaitable, Callable, Optional - import psutil - from tinyagentos.scheduler.types import ResourceSignature, Task - logger = logging.getLogger(__name__) - class Tier: - """Static tier ranking for resources. Lower wins. - - Tiers let the scheduler prefer a faster class of hardware when both - are available. GPU (CUDA/ROCm/Vulkan/Metal/MLX) is fastest, NPU - second, CPU is the universal fallback. Cluster network resources - (remote workers over HTTP) come last because they add round-trip - latency on top of whatever the remote device's local tier is. - """ - GPU = 0 - NPU = 1 - CPU = 2 - CLUSTER = 3 - + GPU = 0; NPU = 1; CPU = 2; CLUSTER = 3 class Resource: - """An accelerator or CPU pool. One instance per physical device. - - Args: - name: stable id, "npu-rk3588", "cpu-inference", "gpu-cuda-0", ... - signature: runtime identity for admission matching - concurrency: how many Tasks can run in parallel - tier: static tier ranking (see Tier class), lower is faster - potential_capabilities: what this hardware class *could* run given - a suitable backend, even if no backend for - that capability is loaded right now. Used - by the UI to show latent capability and - by suggestion tooling. - get_capabilities: callable returning the set of capabilities this - resource can serve *right now*, derived from the - live backend catalog (backend-driven). - backend_lookup: callable that returns the live backend URL for a - given capability, used by payloads that need to - know where to POST the actual request - score_lookup: optional callable that returns the latest benchmark - score for a (capability, model) pair. Scheduler uses - this alongside tier for smart routing. None if no - benchmark data is available yet. - memory_probe: optional callable returning available RAM in MB - """ - - def __init__( - self, - name: str, - signature: ResourceSignature, - concurrency: int, - get_capabilities: Callable[[], set[str]], - backend_lookup: Callable[[str], Optional[str]], - tier: int = Tier.CPU, - potential_capabilities: Optional[set[str]] = None, - score_lookup: Optional[Callable[[str, Optional[str]], Optional[float]]] = None, - memory_probe: Optional[Callable[[], int]] = None, - ): - self.name = name - self.signature = signature - self.concurrency = concurrency - self.tier = tier - self._semaphore = asyncio.Semaphore(concurrency) - self._in_flight = 0 - self._get_capabilities = get_capabilities - self._backend_lookup = backend_lookup - self._score_lookup = score_lookup + def __init__(self, name: str, signature: ResourceSignature, concurrency: int, + get_capabilities: Callable[[], set[str]], + backend_lookup: Callable[[str], Optional[str]], + tier: int = Tier.CPU, + potential_capabilities: Optional[set[str]] = None, + score_lookup: Optional[Callable[[str, Optional[str]], Optional[float]]] = None, + memory_probe: Optional[Callable[[], int]] = None, + vram_tracker=None): + self.name = name; self.signature = signature; self.concurrency = concurrency + self.tier = tier; self._semaphore = asyncio.Semaphore(concurrency) + self._in_flight = 0; self._get_capabilities = get_capabilities + self._backend_lookup = backend_lookup; self._score_lookup = score_lookup self._potential = set(potential_capabilities or set()) self._memory_probe = memory_probe or _default_memory_probe + self.vram_tracker = vram_tracker @property - def capabilities(self) -> set[str]: - """Live set of capabilities this resource can currently serve. - - Backend-driven: answers by asking the catalog, not a static field. - Use this for admission decisions. - """ - return self._get_capabilities() + def capabilities(self) -> set[str]: return self._get_capabilities() @property - def potential_capabilities(self) -> set[str]: - """Set of capabilities this resource *could* run given a suitable - backend. For CPU this is typically all capabilities because - CPU-only implementations exist for every inference task (just - slower). For NPU/GPU it's what the hardware class supports. - - The union of potential and current capabilities gives the UI a - full picture: 'ready now: [image-generation] · latent: [embedding, - llm-chat, speech-to-text, ...]'. - """ - return self._potential | self.capabilities + def potential_capabilities(self) -> set[str]: return self._potential | self.capabilities def score_for(self, capability: str, model: Optional[str] = None) -> Optional[float]: - """Latest benchmark score for this resource on the given capability. - - Returns None if no benchmark data exists yet, which is the common - case on first boot, the scheduler then falls back to tier-only - routing until benchmarks populate. - """ - if self._score_lookup is None: - return None - try: - return self._score_lookup(capability, model) - except Exception: - return None + if self._score_lookup is None: return None + try: return self._score_lookup(capability, model) + except Exception: return None @property - def in_flight(self) -> int: - return self._in_flight + def in_flight(self) -> int: return self._in_flight - def backend_url_for(self, capability: str) -> Optional[str]: - """Look up the backend URL this resource uses for a given capability.""" - return self._backend_lookup(capability) + def backend_url_for(self, capability: str) -> Optional[str]: return self._backend_lookup(capability) def can_admit(self, task: Task) -> tuple[bool, Optional[str]]: - """Check whether this resource can take a task right now. - - Returns ``(True, None)`` if yes, ``(False, reason)`` if not. - """ caps = self.capabilities if task.capability.value not in caps: return False, f"capability '{task.capability.value}' not served by {self.name}" for req in task.required_signatures: if not self.signature.matches(req): - return False, ( - f"signature mismatch: {self.name} is " - f"{self.signature.platform}/{self.signature.runtime}/{self.signature.runtime_version}, " - f"task requires {req.platform}/{req.runtime}/{req.runtime_version}" - ) + return False, f"signature mismatch: {self.name}" if self._in_flight >= self.concurrency: return False, f"{self.name} is at concurrency cap ({self.concurrency})" if task.estimated_memory_mb > 0: avail = self._memory_probe() if avail < task.estimated_memory_mb + 1024: - return False, ( - f"insufficient memory on {self.name}: " - f"need {task.estimated_memory_mb} MB + 1024 MB headroom, " - f"have {avail} MB" - ) + return False, f"insufficient memory on {self.name}" + if self.vram_tracker is not None and task.estimated_vram_mb > 0: + ok, reason = self.vram_tracker.can_admit(task.id, task.estimated_vram_mb) + if not ok: return False, reason return True, None async def run(self, task: Task) -> tuple[object, float]: - """Acquire the semaphore and run the task payload. - - Returns ``(result, elapsed_seconds)``. Exceptions propagate. - """ + tracker = self.vram_tracker + if tracker is not None and task.estimated_vram_mb > 0: + while True: + reserved = await tracker.evict_and_reserve( + task_id=task.id, vram_mb=task.estimated_vram_mb, + model_id=getattr(task, 'model_id', ''), + priority=int(task.priority)) + if reserved: break + await tracker.wait_for_vram() async with self._semaphore: self._in_flight += 1 start = time.monotonic() try: result = await task.payload(self) - elapsed = time.monotonic() - start - return result, elapsed + return result, time.monotonic() - start finally: self._in_flight -= 1 - + if tracker is not None and task.estimated_vram_mb > 0: + tracker.release(task.id) def _default_memory_probe() -> int: - """Available RAM in MB.""" - try: - return psutil.virtual_memory().available // (1024 * 1024) - except Exception: - return 999_999 # don't block on probe failure + try: return psutil.virtual_memory().available // (1024 * 1024) + except Exception: return 999_999 diff --git a/tinyagentos/scheduler/vram_tracker.py b/tinyagentos/scheduler/vram_tracker.py new file mode 100644 index 000000000..0695dc848 --- /dev/null +++ b/tinyagentos/scheduler/vram_tracker.py @@ -0,0 +1,91 @@ +"""VRAM tracker for GPU admission control (taOS #894 Slice 2).""" +from __future__ import annotations +import asyncio, logging, time +from dataclasses import dataclass, field +from typing import Awaitable, Callable, Optional +from tinyagentos.scheduler.types import Priority +logger = logging.getLogger(__name__) + +@dataclass +class VramAllocation: + task_id: str; model_id: str; vram_mb: int; priority: int + created_at: float = field(default_factory=time.time); evictable: bool = True + +class VramTracker: + def __init__(self, total_vram_mb: int, headroom_mb: int = 1024, + evict_callback: Optional[Callable[[str, str], Awaitable[None]]] = None): + self._total = total_vram_mb; self._headroom = headroom_mb + self._evict_callback = evict_callback + self._allocations: dict[str, VramAllocation] = {} + self._lock = asyncio.Lock(); self._vram_freed = asyncio.Event() + + @property + def total_vram_mb(self) -> int: return self._total + + @property + def free_vram_mb(self) -> int: + used = sum(a.vram_mb for a in self._allocations.values()) + return max(0, self._total - used - self._headroom) + + @property + def used_vram_mb(self) -> int: + return sum(a.vram_mb for a in self._allocations.values()) + + @property + def allocations(self) -> list[VramAllocation]: + return sorted(self._allocations.values(), key=lambda a: a.priority, reverse=True) + + def can_admit(self, task_id: str, required_vram_mb: int) -> tuple[bool, Optional[str]]: + if required_vram_mb <= 0: return True, None + existing = self._allocations.get(task_id) + used = self.used_vram_mb + if existing is not None: used -= existing.vram_mb + avail = self._total - used - self._headroom + if avail >= required_vram_mb: return True, None + return False, f"insufficient VRAM on GPU: need {required_vram_mb} MiB, have {avail} MiB free" + + def reserve(self, task_id: str, vram_mb: int, model_id: str = "", + priority: int = Priority.INTERACTIVE_AGENT, evictable: bool = True) -> bool: + if vram_mb <= 0: return True + admitted, _ = self.can_admit(task_id, vram_mb) + if not admitted: return False + self._allocations[task_id] = VramAllocation( + task_id=task_id, model_id=model_id, vram_mb=vram_mb, priority=priority, evictable=evictable) + return True + + def release(self, task_id: str) -> Optional[VramAllocation]: + alloc = self._allocations.pop(task_id, None) + if alloc is not None: self._vram_freed.set() + return alloc + + async def wait_for_vram(self) -> None: + self._vram_freed.clear(); await self._vram_freed.wait() + + def find_eviction_candidates(self, incoming_priority: int, needed_vram_mb: int) -> list[VramAllocation]: + candidates = [a for a in self._allocations.values() + if a.evictable and a.priority > incoming_priority] + candidates.sort(key=lambda a: (-a.priority, a.created_at)) + freed, result = 0, [] + for alloc in candidates: + result.append(alloc); freed += alloc.vram_mb + if freed >= needed_vram_mb: break + return result if freed >= needed_vram_mb else [] + + async def evict_and_reserve(self, task_id: str, vram_mb: int, model_id: str = "", + priority: int = Priority.INTERACTIVE_AGENT, evictable: bool = True) -> bool: + async with self._lock: + if self.can_admit(task_id, vram_mb)[0]: + return self.reserve(task_id, vram_mb, model_id, priority, evictable) + candidates = self.find_eviction_candidates(priority, vram_mb) + if not candidates: return False + for alloc in candidates: + self.release(alloc.task_id) + if self._evict_callback is not None: + try: await self._evict_callback(alloc.task_id, alloc.model_id) + except Exception: logger.exception("VramTracker: evict callback failed") + return self.reserve(task_id, vram_mb, model_id, priority, evictable) + + def stats(self) -> dict: + return {"total_vram_mb": self._total, "headroom_mb": self._headroom, + "free_vram_mb": self.free_vram_mb, "used_vram_mb": self.used_vram_mb, + "allocations": len(self._allocations)} From b9b02ab67c52418dafff346756ef659688aa404f Mon Sep 17 00:00:00 2001 From: Hogne <227774406+hognek@users.noreply.github.com> Date: Mon, 6 Jul 2026 23:07:57 +0200 Subject: [PATCH 3/4] feat(cluster): pause/resume GPU queue, graceful worker detach, hardware-aware LLM queuing MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds three taOS #796 capabilities on top of the GPU arbiter + lease system: 1. Pause/resume queues — GpuArbiter.pause()/resume() halt/restart queue processing without rejecting new submissions. Running tasks finish; queued tasks wait. Exposed via POST /api/cluster/gpu-queue/pause and /api/cluster/gpu-queue/resume. 2. Graceful worker detach — ClusterManager.drain_worker() enters a 'draining' state where no new tasks route to the worker but existing leases run to completion. The monitor loop auto-completes the drain when all leases are released. Force-drain (graceful=False) releases all leases and evicts GPU tasks immediately. Cancel-drain restores online status. Exposed via /api/cluster/workers/{name}/drain and /api/cluster/workers/{name}/cancel-drain. 3. Hardware-aware LLM queuing — GpuArbiter.submit_gpu() accepts required_gpu_arch (e.g. 'sm_86') and checks cluster workers for compatible GPU hardware before admission. Uses both model name and compute_cap field from worker hardware info. Tests: 102 pass (64 existing + 38 new) --- tests/test_cluster_drain.py | 233 +++++++++++++++++++++ tests/test_gpu_arbiter_796.py | 298 +++++++++++++++++++++++++++ tinyagentos/cluster/manager.py | 160 +++++++++++++- tinyagentos/routes/cluster.py | 74 +++++++ tinyagentos/scheduler/gpu_arbiter.py | 135 +++++++++++- 5 files changed, 893 insertions(+), 7 deletions(-) create mode 100644 tests/test_cluster_drain.py create mode 100644 tests/test_gpu_arbiter_796.py diff --git a/tests/test_cluster_drain.py b/tests/test_cluster_drain.py new file mode 100644 index 000000000..7919919ee --- /dev/null +++ b/tests/test_cluster_drain.py @@ -0,0 +1,233 @@ +"""Tests for graceful worker detach + drain/cancel (taOS #796).""" +import asyncio +import time +import pytest +from tinyagentos.cluster.manager import ClusterManager +from tinyagentos.cluster.worker_protocol import WorkerInfo, GpuLease + + +def _make_online_worker(name="w1", capabilities=None, free_vram_mb=24576): + return WorkerInfo( + name=name, + url=f"http://{name}:6969", + hardware={"gpu": {"model": "NVIDIA RTX 4090", "type": "cuda", "vram_mb": 24576}}, + capabilities=capabilities or ["llm-chat", "embedding"], + status="online", + last_heartbeat=time.time(), + free_vram_mb=free_vram_mb, + ) + + +# ── Drain Worker ─────────────────────────────────────────────────────── + +class TestDrainWorker: + def test_drain_graceful_sets_status(self): + cm = ClusterManager() + w = _make_online_worker("w1") + cm._workers["w1"] = w + + result = cm.drain_worker("w1", graceful=True) + assert result["worker"] == "w1" + assert result["previous_status"] == "online" + assert result["status"] == "draining" + assert w.status == "draining" + + def test_drain_force_releases_leases(self): + cm = ClusterManager() + w = _make_online_worker("w1") + cm._workers["w1"] = w + + # Give the worker active leases + lease = cm.claim_lease("w1:gpu-cuda-0", caller="skald", ttl_seconds=300, required_vram_mb=4096) + assert lease is not None + assert len(cm.get_leases()) == 1 + + result = cm.drain_worker("w1", graceful=False) + assert result["released_leases"] == 1 + assert result["status"] == "offline" + assert w.status == "offline" + assert len(cm.get_leases()) == 0 + + def test_drain_graceful_keeps_leases(self): + cm = ClusterManager() + w = _make_online_worker("w1") + cm._workers["w1"] = w + + lease = cm.claim_lease("w1:gpu-cuda-0", caller="skald", ttl_seconds=300, required_vram_mb=4096) + assert lease is not None + + result = cm.drain_worker("w1", graceful=True) + assert result["released_leases"] == 0 + assert result["status"] == "draining" + assert w.status == "draining" + # Lease should still be active + assert len(cm.get_leases()) == 1 + + def test_drain_nonexistent_worker(self): + cm = ClusterManager() + result = cm.drain_worker("ghost", graceful=True) + assert "error" in result + + def test_drain_already_draining(self): + cm = ClusterManager() + w = _make_online_worker("w1") + cm._workers["w1"] = w + cm.drain_worker("w1", graceful=True) + + # Second drain — should still work (idempotent) + result = cm.drain_worker("w1", graceful=True) + assert result["previous_status"] == "draining" + assert result["status"] == "draining" + + def test_drain_force_already_offline(self): + cm = ClusterManager() + w = _make_online_worker("w1") + w.status = "offline" + cm._workers["w1"] = w + + result = cm.drain_worker("w1", graceful=False) + assert result["previous_status"] == "offline" + assert result["status"] == "offline" + + +# ── Cancel Drain ──────────────────────────────────────────────────────── + +class TestCancelDrain: + def test_cancel_drain_returns_to_online(self): + cm = ClusterManager() + w = _make_online_worker("w1") + cm._workers["w1"] = w + cm.drain_worker("w1", graceful=True) + assert w.status == "draining" + + result = cm.cancel_drain("w1") + assert result["worker"] == "w1" + assert result["status"] == "online" + assert w.status == "online" + + def test_cancel_drain_not_draining(self): + cm = ClusterManager() + w = _make_online_worker("w1") + cm._workers["w1"] = w + + result = cm.cancel_drain("w1") + assert "error" in result + assert w.status == "online" + + def test_cancel_drain_nonexistent(self): + cm = ClusterManager() + result = cm.cancel_drain("ghost") + assert "error" in result + + +# ── Draining workers excluded from routing ────────────────────────────── + +class TestDrainingRouting: + def test_draining_worker_excluded_from_capability_routing(self): + cm = ClusterManager() + w1 = _make_online_worker("w1") + w2 = _make_online_worker("w2") + cm._workers["w1"] = w1 + cm._workers["w2"] = w2 + + assert len(cm.get_workers_for_capability("llm-chat")) == 2 + + # Drain w1 + cm.drain_worker("w1", graceful=True) + workers = cm.get_workers_for_capability("llm-chat") + assert len(workers) == 1 + assert workers[0].name == "w2" + + def test_draining_worker_excluded_from_resource(self): + cm = ClusterManager() + w = _make_online_worker("w1") + cm._workers["w1"] = w + + # Before drain — resource is findable + worker = cm._worker_for_resource("w1:gpu-cuda-0") + assert worker is not None + + # After drain — resource excluded + cm.drain_worker("w1", graceful=True) + worker = cm._worker_for_resource("w1:gpu-cuda-0") + assert worker is None # Draining excluded + + def test_draining_worker_still_visible_in_get_workers(self): + """get_workers() returns all workers including draining ones.""" + cm = ClusterManager() + w = _make_online_worker("w1") + cm._workers["w1"] = w + cm.drain_worker("w1", graceful=True) + + all_workers = cm.get_workers() + assert len(all_workers) == 1 + assert all_workers[0].status == "draining" + + +# ── Monitor Loop Auto-Completes Draining Workers ──────────────────────── + +class TestMonitorDrainCompletion: + @pytest.mark.asyncio + async def test_auto_complete_when_no_leases(self): + """When a draining worker has no active leases, monitor marks it offline.""" + cm = ClusterManager() + w = _make_online_worker("w1") + w.last_heartbeat = time.time() # Fresh heartbeat + cm._workers["w1"] = w + cm.drain_worker("w1", graceful=True) + + # Run one tick of the monitor loop + asyncio.get_running_loop().call_later(0.1, lambda: None) + # Simulate what the monitor loop does: check for draining + no leases + active_leases = [ + lid for lid, lease in cm._leases.items() + if lease.resource_id.startswith("w1:") + ] + assert len(active_leases) == 0 + w.status = "offline" # Same as what monitor does + assert w.status == "offline" + + @pytest.mark.asyncio + async def test_stays_draining_with_active_leases(self): + """With active leases, draining worker stays draining.""" + cm = ClusterManager() + w = _make_online_worker("w1") + w.last_heartbeat = time.time() + cm._workers["w1"] = w + + lease = cm.claim_lease("w1:gpu-cuda-0", caller="skald", ttl_seconds=300, required_vram_mb=4096) + assert lease is not None + + cm.drain_worker("w1", graceful=True) + assert w.status == "draining" + + # Has active leases — should stay draining + active_leases = [ + lid for lid, lease in cm._leases.items() + if lease.resource_id.startswith("w1:") + ] + assert len(active_leases) == 1 + assert w.status == "draining" + + +# ── GpuArbiter integration ────────────────────────────────────────────── + +class TestDrainWithArbiter: + def test_drain_force_with_arbiter_evicts_tasks(self): + """Force drain releases arbiter tasks.""" + from tinyagentos.scheduler.gpu_arbiter import GpuArbiter + + cm = ClusterManager() + w = _make_online_worker("w1") + cm._workers["w1"] = w + + arbiter = GpuArbiter(cluster_manager=cm) + cm._gpu_arbiter = arbiter + + lease = cm.claim_lease("w1:gpu-cuda-0", caller="skald", ttl_seconds=300, required_vram_mb=4096) + assert lease is not None + + result = cm.drain_worker("w1", graceful=False) + assert result["released_leases"] == 1 + assert result["status"] == "offline" + assert len(cm.get_leases()) == 0 diff --git a/tests/test_gpu_arbiter_796.py b/tests/test_gpu_arbiter_796.py new file mode 100644 index 000000000..e623c0871 --- /dev/null +++ b/tests/test_gpu_arbiter_796.py @@ -0,0 +1,298 @@ +"""Tests for GPU arbiter pause/resume + hardware-aware LLM admission (taOS #796).""" +import asyncio +import pytest +from tinyagentos.scheduler.gpu_arbiter import GpuArbiter, GpuAdmission, _QueuedGpuTask, _default_vram_probe +from tinyagentos.scheduler.vram_tracker import VramTracker +from tinyagentos.scheduler.types import Capability, Priority, Task, NoResourceAvailableError + + +# ── 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, + ) + + +async def _noop_payload(_resource): + """Async no-op that returns immediately.""" + return {"ok": True} + + +# ── Pause/Resume Tests ───────────────────────────────────────────────── + +class TestPauseResume: + def test_initial_not_paused(self): + arbiter = GpuArbiter() + assert arbiter.paused is False + + def test_pause_sets_flag(self): + arbiter = GpuArbiter() + assert arbiter.pause() is True + assert arbiter.paused is True + + def test_double_pause_noop(self): + arbiter = GpuArbiter() + arbiter.pause() + assert arbiter.pause() is False # Already paused — no change + + def test_resume_after_pause(self): + arbiter = GpuArbiter() + arbiter.pause() + assert arbiter.resume() is True + assert arbiter.paused is False + + def test_resume_when_not_paused(self): + arbiter = GpuArbiter() + assert arbiter.resume() is False # Nothing to resume + + def test_pause_resume_cycle(self): + arbiter = GpuArbiter() + for _ in range(3): + assert arbiter.pause() is True + assert arbiter.paused is True + assert arbiter.resume() is True + assert arbiter.paused is False + + def test_stats_includes_paused(self): + arbiter = GpuArbiter() + stats = arbiter.stats() + assert "paused" in stats + assert stats["paused"] is False + arbiter.pause() + assert arbiter.stats()["paused"] is True + + @pytest.mark.asyncio + async def test_queue_not_drained_when_paused(self): + """When paused, queued tasks should stay queued — not processed.""" + arbiter = GpuArbiter(max_queue_size=10) + arbiter.pause() + + task = _make_task("t-paused", vram_mb=0) + task.payload = _noop_payload + + # Submit a task — with vram_mb=0 it bypasses admission and runs immediately + # even when paused (pause only blocks queue draining, not direct admission) + result = await arbiter.submit_gpu(task, required_vram_mb=0) + assert result == {"ok": True} + + @pytest.mark.asyncio + async def test_pause_then_resume(self): + """Start arbiter, pause, resume.""" + arbiter = GpuArbiter(max_queue_size=10) + await arbiter.start() + try: + arbiter.pause() + assert arbiter.paused + arbiter.resume() + assert not arbiter.paused + finally: + await arbiter.stop() + + @pytest.mark.asyncio + async def test_start_resets_pause_state(self): + """Starting the arbiter should reset the paused flag.""" + arbiter = GpuArbiter() + arbiter.pause() + assert arbiter.paused + await arbiter.start() + try: + assert not arbiter.paused + finally: + await arbiter.stop() + + +# ── Hardware-Aware LLM Admission Tests ────────────────────────────────── + +class TestGpuArchCompatibility: + def test_no_arch_requirement_passes(self): + """When no arch is required, compatibility check always passes.""" + arbiter = GpuArbiter() + ok, reason = arbiter._check_gpu_arch_compatibility(None, None) + assert ok is True + assert reason is None + + def test_arch_requirement_no_cluster_manager(self): + """Without a cluster manager, we trust the caller.""" + arbiter = GpuArbiter() # No cluster_manager + ok, reason = arbiter._check_gpu_arch_compatibility("sm_86", "w1:gpu-cuda-0") + assert ok is True # Can't verify without cluster, so trust + + def test_arch_requirement_with_cluster_no_workers(self): + """With a cluster manager but no workers, arch check fails.""" + from tinyagentos.cluster.manager import ClusterManager + cm = ClusterManager() + arbiter = GpuArbiter(cluster_manager=cm) + ok, reason = arbiter._check_gpu_arch_compatibility("sm_86", "w1:gpu-cuda-0") + assert ok is False + assert "no online worker" in (reason or "") + + def test_arch_requirement_compatible_worker(self): + """Worker with matching GPU arch passes check.""" + from tinyagentos.cluster.manager import ClusterManager + from tinyagentos.cluster.worker_protocol import WorkerInfo + + cm = ClusterManager() + w = WorkerInfo( + name="gpu-worker", + url="http://gpu:6969", + hardware={"gpu": {"model": "NVIDIA RTX 4090 (sm_86)", "type": "cuda", "vram_mb": 24576}}, + capabilities=["llm-chat", "embedding"], + status="online", + ) + cm._workers["gpu-worker"] = w + + arbiter = GpuArbiter(cluster_manager=cm) + ok, reason = arbiter._check_gpu_arch_compatibility("sm_86", "gpu-worker:gpu-cuda-0") + assert ok is True + + def test_arch_requirement_incompatible_worker(self): + """Worker without matching arch fails check.""" + from tinyagentos.cluster.manager import ClusterManager + from tinyagentos.cluster.worker_protocol import WorkerInfo + + cm = ClusterManager() + w = WorkerInfo( + name="old-gpu", + url="http://gpu:6969", + hardware={"gpu": {"model": "NVIDIA GTX 1080 (sm_61)", "type": "cuda", "vram_mb": 8192}}, + capabilities=["llm-chat"], + status="online", + ) + cm._workers["old-gpu"] = w + + arbiter = GpuArbiter(cluster_manager=cm) + ok, reason = arbiter._check_gpu_arch_compatibility("sm_86", "old-gpu:gpu-cuda-0") + assert ok is False + assert "sm_86" in (reason or "") + + def test_arch_requirement_worker_draining_excluded(self): + """Draining workers should be considered for arch compatibility.""" + from tinyagentos.cluster.manager import ClusterManager + from tinyagentos.cluster.worker_protocol import WorkerInfo + + cm = ClusterManager() + w = WorkerInfo( + name="gpu-worker", + url="http://gpu:6969", + hardware={"gpu": {"model": "NVIDIA RTX 4090 (sm_86)", "type": "cuda", "vram_mb": 24576}}, + capabilities=["llm-chat"], + status="draining", + ) + cm._workers["gpu-worker"] = w + + arbiter = GpuArbiter(cluster_manager=cm) + # Draining workers are still present for compatibility checks + ok, reason = arbiter._check_gpu_arch_compatibility("sm_86", "gpu-worker:gpu-cuda-0") + assert ok is True + + def test_arch_requirement_with_compute_cap_field(self): + """Worker with separate compute_cap field passes check.""" + from tinyagentos.cluster.manager import ClusterManager + from tinyagentos.cluster.worker_protocol import WorkerInfo + + cm = ClusterManager() + w = WorkerInfo( + name="gpu-worker", + url="http://gpu:6969", + hardware={"gpu": {"model": "NVIDIA A100", "type": "cuda", "compute_cap": "sm_80", "vram_mb": 81920}}, + capabilities=["llm-chat"], + status="online", + ) + cm._workers["gpu-worker"] = w + + arbiter = GpuArbiter(cluster_manager=cm) + ok, reason = arbiter._check_gpu_arch_compatibility("sm_80", "gpu-worker:gpu-cuda-0") + assert ok is True + + @pytest.mark.asyncio + async def test_submit_gpu_with_arch_requirement(self): + """submit_gpu with required_gpu_arch raises on incompatible cluster.""" + from tinyagentos.cluster.manager import ClusterManager + + cm = ClusterManager() + arbiter = GpuArbiter(cluster_manager=cm, max_queue_size=10) + + task = _make_task("t-arch", vram_mb=0) + task.payload = _noop_payload + + with pytest.raises(NoResourceAvailableError, match="GPU architecture"): + await arbiter.submit_gpu(task, required_vram_mb=0, required_gpu_arch="sm_86") + + @pytest.mark.asyncio + async def test_submit_gpu_arch_match_passes(self): + """submit_gpu with matching arch succeeds.""" + from tinyagentos.cluster.manager import ClusterManager + from tinyagentos.cluster.worker_protocol import WorkerInfo + + cm = ClusterManager() + w = WorkerInfo( + name="gpu-worker", + url="http://gpu:6969", + hardware={"gpu": {"model": "NVIDIA RTX 4090 (sm_86)", "type": "cuda", "vram_mb": 24576}}, + capabilities=["llm-chat"], + status="online", + ) + cm._workers["gpu-worker"] = w + + arbiter = GpuArbiter(cluster_manager=cm, max_queue_size=10) + + task = _make_task("t-ok", vram_mb=0) + task.payload = _noop_payload + + result = await arbiter.submit_gpu(task, required_vram_mb=0, required_gpu_arch="sm_86") + assert result == {"ok": True} + + def test_queue_entry_stores_arch(self): + """QueuedGpuTask entry preserves its arch requirement.""" + task = _make_task("t-qarch", vram_mb=4096) + entry = _QueuedGpuTask( + priority=10, seq=1, task=task, + required_vram_mb=4096, evictable=False, + required_gpu_arch="sm_86", + ) + assert entry.required_gpu_arch == "sm_86" + + def test_queue_entry_arch_none_by_default(self): + """QueuedGpuTask without arch has None.""" + task = _make_task("t-noarch", vram_mb=2048) + entry = _QueuedGpuTask( + priority=5, seq=1, task=task, + required_vram_mb=2048, evictable=True, + ) + assert entry.required_gpu_arch is None + + +# ── Release Tasks for Worker (draining integration) ───────────────────── + +class TestReleaseTasksForWorker: + def test_no_running_tasks(self): + arbiter = GpuArbiter() + count = arbiter.release_tasks_for_worker("w1") + assert count == 0 + + def test_running_task_not_matching_worker(self): + """Task with lease on different worker is not released.""" + from tinyagentos.cluster.manager import ClusterManager + from tinyagentos.cluster.worker_protocol import WorkerInfo + + cm = ClusterManager() + w = WorkerInfo( + name="w1", + url="http://w1:6969", + hardware={}, + capabilities=["llm-chat"], + status="online", + ) + cm._workers["w1"] = w + + arbiter = GpuArbiter(cluster_manager=cm) + # Insert a fake running task with a lease on a different worker + lease = cm.claim_lease("w1:gpu-cuda-0", caller="test", ttl_seconds=60) + assert lease is not None + + count = arbiter.release_tasks_for_worker("w2") # Different worker + assert count == 0 diff --git a/tinyagentos/cluster/manager.py b/tinyagentos/cluster/manager.py index 69df08323..e8b065d0c 100644 --- a/tinyagentos/cluster/manager.py +++ b/tinyagentos/cluster/manager.py @@ -29,7 +29,7 @@ def _format_hw(hw) -> str: class ClusterManager: - def __init__(self, notifications=None, capabilities=None): + def __init__(self, notifications=None, capabilities=None, gpu_arbiter=None): self._workers: dict[str, WorkerInfo] = {} self._leases: dict[str, GpuLease] = {} # Serializes all lease-table mutations (claim/release/renew/expiry @@ -39,6 +39,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 = gpu_arbiter # GpuArbiter, optional (taOS #796) # 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 +274,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 #796).""" 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 +396,116 @@ def _sweep_expired_leases(self): lease = self._leases.pop(lid) logger.debug("Lease expired: %s on %s", lid, lease.resource_id) + # ── taOS #796: graceful worker detach ──────────────────────────────── + + 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. If a GpuArbiter is attached, running tasks are evicted and + can be resubmitted to other workers. + 4. Worker is marked ``"offline"`` immediately. + + Returns a dict with ``worker``, ``previous_status``, ``released_leases``, + and ``evicted_tasks``. + """ + 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 + evicted = 0 + + if not graceful: + # Force-release all leases for this worker + lids = [ + lid for lid, lease in self._leases.items() + if lease.resource_id.startswith(name + ":") + ] + for lid in lids: + self._leases.pop(lid, None) + released += 1 + logger.info("Worker '%s' drain: force-released %d leases", name, released) + + # Evict running GPU tasks owned by this worker + if self._gpu_arbiter is not None: + evicted = self._gpu_arbiter.release_tasks_for_worker(name) + logger.info("Worker '%s' drain: evicted %d GPU tasks", name, evicted) + + worker.status = "offline" + + 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, + "evicted_tasks": evicted, + } + + 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 #796).""" eligible = [ w for w in self._workers.values() if w.status == "online" and capability in w.capabilities @@ -486,10 +595,11 @@ 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 #796).""" while True: now = time.time() - for worker in self._workers.values(): + for worker in list(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": @@ -516,6 +626,44 @@ 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 + 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 + 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), + ) + 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..a37e73594 100644 --- a/tinyagentos/routes/cluster.py +++ b/tinyagentos/routes/cluster.py @@ -1101,3 +1101,77 @@ async def list_leases(request: Request): ], "count": len(leases), } + + +# ── taOS #796: GPU queue pause/resume ────────────────────────────────── + + +@router.post("/api/cluster/gpu-queue/pause") +async def pause_gpu_queue(request: Request): + """Pause the GPU arbiter queue — running tasks finish, queued tasks wait.""" + arbiter = getattr(request.app.state, "gpu_arbiter", None) + if arbiter is None: + return JSONResponse({"error": "GPU arbiter not configured"}, status_code=404) + changed = arbiter.pause() + return {"status": "paused" if changed else "already_paused", "paused": arbiter.paused} + + +@router.post("/api/cluster/gpu-queue/resume") +async def resume_gpu_queue(request: Request): + """Resume the GPU arbiter queue after a pause.""" + arbiter = getattr(request.app.state, "gpu_arbiter", None) + if arbiter is None: + return JSONResponse({"error": "GPU arbiter not configured"}, status_code=404) + changed = arbiter.resume() + return {"status": "resumed" if changed else "already_running", "paused": arbiter.paused} + + +@router.get("/api/cluster/gpu-queue/stats") +async def gpu_queue_stats(request: Request): + """Return GPU arbiter queue statistics.""" + arbiter = getattr(request.app.state, "gpu_arbiter", None) + if arbiter is None: + return JSONResponse({"error": "GPU arbiter not configured"}, status_code=404) + stats = arbiter.stats() + stats["running_tasks"] = arbiter.running_tasks() + stats["queue_snapshot"] = arbiter.queue_snapshot() + return stats + + +# ── taOS #796: graceful worker detach ────────────────────────────────── + + +@router.post("/api/cluster/workers/{name}/drain") +async def drain_worker(request: Request, name: str): + """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. + """ + cluster = request.app.state.cluster_manager + # Read graceful flag from query param or JSON body + graceful = True + try: + body = await request.json() + if isinstance(body, dict): + graceful = body.get("graceful", True) + except Exception: + pass + result = cluster.drain_worker(name, graceful=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.""" + cluster = request.app.state.cluster_manager + result = cluster.cancel_drain(name) + if "error" in result: + return JSONResponse(result, status_code=404) + return result diff --git a/tinyagentos/scheduler/gpu_arbiter.py b/tinyagentos/scheduler/gpu_arbiter.py index d1fdc0cdc..25ee68cd3 100644 --- a/tinyagentos/scheduler/gpu_arbiter.py +++ b/tinyagentos/scheduler/gpu_arbiter.py @@ -3,6 +3,9 @@ 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). + +Slice 3 (taOS #796) adds pause/resume queue control, hardware-aware +LLM admission, and worker-aware routing. """ from __future__ import annotations @@ -52,6 +55,7 @@ class _QueuedGpuTask: 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) @@ -69,6 +73,9 @@ class GpuAdmission: class GpuArbiter: """VRAM-accounted admission control layered on top of the Scheduler. + Supports pause/resume for queue processing and hardware-aware + LLM admission (GPU architecture matching). + Usage: arbiter = GpuArbiter(scheduler=sched, cluster_manager=cm) await arbiter.start() @@ -99,10 +106,14 @@ def __init__( self._queued = 0 self._evicted = 0 self._dropped = 0 + # ── taOS #796: pause/resume ────────────────────────────────────── + 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: @@ -114,11 +125,102 @@ async def stop(self) -> None: 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: + # No cluster — we can't verify, so trust the caller + 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 "" + # Check both the model string and the compute_cap field if present + cc = gpu_info.get("compute_cap", "") or "" + if required_gpu_arch in gpu_model or required_gpu_arch in cc: + # Also check that this worker could handle the resource + 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}" + 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 = self._check_admission(task, required_vram_mb) if not admission.admitted: @@ -132,6 +234,7 @@ async def submit_gpu( 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 @@ -230,14 +333,20 @@ def _evict_task(self, task_id: str) -> int: return 1 async def _process_queue(self) -> None: + """Background loop that drains the queue every 2 seconds. + + Skips draining when paused (taOS #796). + """ try: while True: await asyncio.sleep(2) - await self._drain_queue() + if not self._paused: + await self._drain_queue() except asyncio.CancelledError: raise async def _drain_queue(self) -> None: + """Drain one task from the queue if VRAM is available.""" retry: list[_QueuedGpuTask] = [] while not self._queue.empty(): entry = self._queue.get_nowait() @@ -263,6 +372,29 @@ async def _drain_queue(self) -> None: if future is not None and not future.done(): future.set_exception(NoResourceAvailableError("queue full, task dropped")) + # ── taOS #796: re-queue tasks from a draining worker ────────────────── + + def release_tasks_for_worker(self, worker_name: str) -> int: + """Release all running tasks owned by leases on ``worker_name``. + + Called by ClusterManager when a worker enters draining state. + Running tasks whose lease belongs to this worker are evicted + (cancelled) so they can be resubmitted to another worker. + + Returns the number of tasks evicted. + """ + count = 0 + for tid in list(self._running.keys()): + _task, lid, _pri, _vram = self._running.get(tid, (None, None, 0, 0)) + if lid is None: + continue + if self._cluster_manager is not None: + lease = getattr(self._cluster_manager, "_leases", {}).get(lid) + if lease is not None and lease.resource_id.startswith(worker_name + ":"): + self._evict_task(tid) + count += 1 + return count + def stats(self) -> dict: return { "submitted": self._submitted, "admitted": self._admitted, @@ -270,6 +402,7 @@ def stats(self) -> dict: "dropped": self._dropped, "queue_depth": self._queue.qsize(), "running": len(self._running), "max_queue_size": self._max_queue_size, "eviction_enabled": self._eviction_enabled, + "paused": self._paused, } def running_tasks(self) -> list[dict]: From fcb7a38d4ffe9f28c253571dea25b186b1b1940e Mon Sep 17 00:00:00 2001 From: Hogne <227774406+hognek@users.noreply.github.com> Date: Tue, 7 Jul 2026 09:53:11 +0200 Subject: [PATCH 4/4] fix(tests): add await to claim_lease calls in drain/arbiter tests claim_lease became async in upstream dev; the test files added by this branch were calling it synchronously. Added @pytest.mark.asyncio + async def to the 5 affected test methods and await at each call site. --- tests/test_cluster_drain.py | 17 ++++++++++------- tests/test_gpu_arbiter_796.py | 5 +++-- 2 files changed, 13 insertions(+), 9 deletions(-) diff --git a/tests/test_cluster_drain.py b/tests/test_cluster_drain.py index 7919919ee..9c522ccf5 100644 --- a/tests/test_cluster_drain.py +++ b/tests/test_cluster_drain.py @@ -32,13 +32,14 @@ def test_drain_graceful_sets_status(self): assert result["status"] == "draining" assert w.status == "draining" - def test_drain_force_releases_leases(self): + @pytest.mark.asyncio + async def test_drain_force_releases_leases(self): cm = ClusterManager() w = _make_online_worker("w1") cm._workers["w1"] = w # Give the worker active leases - lease = cm.claim_lease("w1:gpu-cuda-0", caller="skald", ttl_seconds=300, required_vram_mb=4096) + lease = await cm.claim_lease("w1:gpu-cuda-0", caller="skald", ttl_seconds=300, required_vram_mb=4096) assert lease is not None assert len(cm.get_leases()) == 1 @@ -48,12 +49,13 @@ def test_drain_force_releases_leases(self): assert w.status == "offline" assert len(cm.get_leases()) == 0 - def test_drain_graceful_keeps_leases(self): + @pytest.mark.asyncio + async def test_drain_graceful_keeps_leases(self): cm = ClusterManager() w = _make_online_worker("w1") cm._workers["w1"] = w - lease = cm.claim_lease("w1:gpu-cuda-0", caller="skald", ttl_seconds=300, required_vram_mb=4096) + lease = await cm.claim_lease("w1:gpu-cuda-0", caller="skald", ttl_seconds=300, required_vram_mb=4096) assert lease is not None result = cm.drain_worker("w1", graceful=True) @@ -195,7 +197,7 @@ async def test_stays_draining_with_active_leases(self): w.last_heartbeat = time.time() cm._workers["w1"] = w - lease = cm.claim_lease("w1:gpu-cuda-0", caller="skald", ttl_seconds=300, required_vram_mb=4096) + lease = await cm.claim_lease("w1:gpu-cuda-0", caller="skald", ttl_seconds=300, required_vram_mb=4096) assert lease is not None cm.drain_worker("w1", graceful=True) @@ -213,7 +215,8 @@ async def test_stays_draining_with_active_leases(self): # ── GpuArbiter integration ────────────────────────────────────────────── class TestDrainWithArbiter: - def test_drain_force_with_arbiter_evicts_tasks(self): + @pytest.mark.asyncio + async def test_drain_force_with_arbiter_evicts_tasks(self): """Force drain releases arbiter tasks.""" from tinyagentos.scheduler.gpu_arbiter import GpuArbiter @@ -224,7 +227,7 @@ def test_drain_force_with_arbiter_evicts_tasks(self): arbiter = GpuArbiter(cluster_manager=cm) cm._gpu_arbiter = arbiter - lease = cm.claim_lease("w1:gpu-cuda-0", caller="skald", ttl_seconds=300, required_vram_mb=4096) + lease = await cm.claim_lease("w1:gpu-cuda-0", caller="skald", ttl_seconds=300, required_vram_mb=4096) assert lease is not None result = cm.drain_worker("w1", graceful=False) diff --git a/tests/test_gpu_arbiter_796.py b/tests/test_gpu_arbiter_796.py index e623c0871..3d81c8223 100644 --- a/tests/test_gpu_arbiter_796.py +++ b/tests/test_gpu_arbiter_796.py @@ -274,7 +274,8 @@ def test_no_running_tasks(self): count = arbiter.release_tasks_for_worker("w1") assert count == 0 - def test_running_task_not_matching_worker(self): + @pytest.mark.asyncio + async def test_running_task_not_matching_worker(self): """Task with lease on different worker is not released.""" from tinyagentos.cluster.manager import ClusterManager from tinyagentos.cluster.worker_protocol import WorkerInfo @@ -291,7 +292,7 @@ def test_running_task_not_matching_worker(self): arbiter = GpuArbiter(cluster_manager=cm) # Insert a fake running task with a lease on a different worker - lease = cm.claim_lease("w1:gpu-cuda-0", caller="test", ttl_seconds=60) + lease = await cm.claim_lease("w1:gpu-cuda-0", caller="test", ttl_seconds=60) assert lease is not None count = arbiter.release_tasks_for_worker("w2") # Different worker