feat(scheduler): VramReservationManager — atomic VRAM bookkeeping for GpuArbiter [#1683]

tinyagentos/app.py

@@ -57,7 +57,8 @@ 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.gpu_arbiter import _probe_nvidia_vram
 from tinyagentos.scheduler.discovery import build_scheduler as build_resource_scheduler
 from tinyagentos.torrent_settings import TorrentSettingsStore
 from tinyagentos.relationships import RelationshipManager
@@ -1124,6 +1125,7 @@ async def _reload_llm_proxy_on_catalog_change() -> None:
         # Build the resource scheduler from hardware profile + live catalog.
         # Phase 1: local resources only (NPU + CPU), capability-based routing
         # with fallback and priority. Cluster-aware dispatch is Phase 3.
+        resource_scheduler = None
         try:
             resource_scheduler = build_resource_scheduler(
                 hardware_profile,
@@ -1140,6 +1142,24 @@ async def _reload_llm_proxy_on_catalog_change() -> None:
         except Exception:
             logger.exception("resource scheduler failed to build — routes will use static config")
             app.state.resource_scheduler = None
+
+        # Build the GPU arbiter — wraps the resource scheduler with VRAM-accounted
+        # admission control, queuing, and eviction for GPU-bound workloads.
+        try:
+            gpu_arbiter = GpuArbiter(
+                scheduler=resource_scheduler if resource_scheduler is not None else None,
+                cluster_manager=cluster_manager,
+                vram_probe=_probe_nvidia_vram,
+                max_queue_size=100,
+                eviction_enabled=True,
+            )
+            await gpu_arbiter.start()
+            app.state.gpu_arbiter = gpu_arbiter
+            cluster_manager._gpu_arbiter = gpu_arbiter  # wire for eviction paths
+            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)
@@ -1567,6 +1587,7 @@ async def dispatch(self, request, call_next):
     import platform as _platform
     app.state.host_arch = _platform.machine()
     app.state.trace_registry = None
+    app.state.gpu_arbiter = None
     app.state.otel_emitter = None
     app.state.span_store_registry = None
     app.state.bridge_sessions = None

tinyagentos/cluster/manager.py

@@ -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 = None      # GpuArbiter, wired after creation 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()

tinyagentos/scheduler/__init__.py

@@ -32,11 +32,13 @@
 _LAZY_EXPORTS = {
     "BackendCatalog": "backend_catalog",
     "BackendEntry": "backend_catalog",
+    "GpuArbiter": "gpu_arbiter",
     "HistoryStore": "history_store",
     "Resource": "resource",
     "Scheduler": "scheduler",
     "ScoreCache": "score_cache",
     "TaskScheduler": "task_scheduler",
+    "VramReservationManager": "gpu_arbiter",
 }
 
 
@@ -58,6 +60,7 @@ def __dir__():
     "BackendCatalog",
     "BackendEntry",
     "Capability",
+    "GpuArbiter",
     "HistoryStore",
     "NoResourceAvailableError",
     "Priority",
@@ -70,4 +73,5 @@ def __dir__():
     "TaskRecord",
     "TaskScheduler",
     "TaskStatus",
+    "VramReservationManager",
 ]

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(