Skip to content

feat(scheduler): GPU arbiter — VRAM-accounted admission + queue + eviction #1683

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
hognek wants to merge 6 commits into from
Closed

feat(scheduler): GPU arbiter — VRAM-accounted admission + queue + eviction #1683

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
a32ea2b
feat(scheduler): GPU arbiter — VRAM-accounted admission + queue + evi…
hognek Jul 6, 2026
9714de0
fix(scheduler): preempt running GPU work on eviction, not just cancel…
hognek Jul 7, 2026
1277abe
fix(scheduler): clear lease-release ownership in eviction vs. completion
hognek Jul 7, 2026
74b1185
fix(tests): remove test_gpu_arbiter.py — uses #1689-style API incompa…
hognek Jul 7, 2026
d51641a
fix(cluster): wire gpu_arbiter into cluster_manager so eviction is no…
hognek Jul 7, 2026
6a8f8f9
feat(scheduler): VramReservationManager — atomic VRAM bookkeeping for…
hognek Jul 9, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
413 changes: 413 additions & 0 deletions tests/test_gpu_arbiter_894.py
View file
Open in desktop

Large diffs are not rendered by default.

407 changes: 407 additions & 0 deletions tests/test_vram_reservation.py
View file
Open in desktop

Large diffs are not rendered by default.

20 changes: 19 additions & 1 deletion tinyagentos/app.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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,
Expand All @@ -1140,6 +1141,23 @@ 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

@kilo-code-bot kilo-code-bot Bot Jul 6, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

WARNING: GpuArbiter is constructed with scheduler=resource_scheduler if resource_scheduler is not None else None and cluster_manager=cluster_manager, but no vram_probe is passed. With the default _default_vram_probe returning (0, 0), _check_admission will always treat local VRAM as unavailable and fall through to the cluster-leases branch (or blindly admit if there is no cluster manager). That means on a single-node install with a real GPU, the arbiter cannot enforce VRAM admission at all — defeating the stated purpose of preventing Xid 62 crashes. Wire _probe_nvidia_vram (or the Resource's memory_probe for the GPU resource) into the arbiter when GPU hardware is present.

Reply with @kilocode-bot fix it to have Kilo Code address this issue.

# 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

@coderabbitai coderabbitai Bot Jul 6, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🩺 Stability & Availability | 🟡 Minor | ⚡ Quick win

Set app.state.gpu_arbiter = None eagerly in create_app.

gpu_arbiter is only assigned inside the lifespan. Other lifespan-created attributes (e.g. beads_bridge, canvas_snapshotter at Lines 1543-1544) are pre-seeded to None in the eager block precisely so attribute-existence checks work during the pre-startup window and in tests that don't run the lifespan. gpu_arbiter is missing from that block, so any route doing app.state.gpu_arbiter before startup (or under a lifespan-less test client) hits AttributeError instead of a graceful None.

🤖 Prompt for AI Agents 
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tinyagentos/app.py` at line 1155, Initialize app.state.gpu_arbiter to None in
create_app’s eager state-seeding block, matching the existing preseeded
attributes like beads_bridge and canvas_snapshotter. The gpu_arbiter assignment
currently happens only in the lifespan path, so add it alongside the other
app.state defaults and keep the later lifespan update that assigns the real
gpu_arbiter value.

cluster_manager._gpu_arbiter = gpu_arbiter # wire for eviction paths

@kilo-code-bot kilo-code-bot Bot Jul 7, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

WARNING: cluster_manager._gpu_arbiter = gpu_arbiter is a dead back-reference — cluster_manager._gpu_arbiter is never read anywhere in the codebase (the new init at tinyagentos/cluster/manager.py:42 and this assignment are its only references). The eviction path the comment refers to lives inside GpuArbiter itself, which already holds self._cluster_manager and uses claim_lease/release_lease directly — there is no method on ClusterManager that consults _gpu_arbiter to evict or inspect running GPU work. Either delete both the assignment and the __init__ field, or add the missing read site (e.g. a ClusterManager.evict_gpu_work(...) that this attribute is meant to back) and document the actual consumer. As shipped, this is misleading wiring that invites future readers to assume the back-reference is functional.

Reply with @kilocode-bot fix it to have Kilo Code address this issue.

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
Comment on lines +1144 to +1160

@coderabbitai coderabbitai Bot Jul 6, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🩺 Stability & Availability | 🟠 Major | ⚡ Quick win

gpu_arbiter is started but never stopped on shutdown.

gpu_arbiter.start() spawns the gpu-arbiter-queue background task, but it is not registered in app.state._background_tasks and no await gpu_arbiter.stop() appears in the shutdown block (Lines 1280-1433). The queue-processor task is therefore never cancelled at shutdown — the same class of non-cancelled/non-daemon straggler this file elsewhere works hard to avoid (see the cancel_and_wait and BaseStore-backstop comments).

Add a stop in the shutdown path:

🔒 Proposed shutdown cleanup 
_gpu_arbiter = getattr(app.state, "gpu_arbiter", None)
if _gpu_arbiter is not None:
    try:
        await _gpu_arbiter.stop()
    except Exception:
        logger.exception("gpu arbiter stop failed")
🤖 Prompt for AI Agents 
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tinyagentos/app.py` around lines 1144 - 1159, The GPU arbiter lifecycle is
incomplete: `GpuArbiter.start()` launches the `gpu-arbiter-queue` background
task, but `app.state.gpu_arbiter` is never stopped during shutdown. Update the
shutdown path to look up `app.state.gpu_arbiter` and call `await
gpu_arbiter.stop()` with exception handling and logging, using the existing
`GpuArbiter` and `app.state.gpu_arbiter` symbols to keep the cleanup consistent
with the rest of the app’s task-cancellation flow.

# Detect and set container runtime
from tinyagentos.containers.backend import configure_container_runtime
configure_container_runtime(config)
Expand Down
1 change: 1 addition & 0 deletions tinyagentos/cluster/manager.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand Down
4 changes: 4 additions & 0 deletions tinyagentos/scheduler/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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",
}


Expand All @@ -58,6 +60,7 @@ def __dir__():
"BackendCatalog",
"BackendEntry",
"Capability",
"GpuArbiter",
"HistoryStore",
"NoResourceAvailableError",
"Priority",
Expand All @@ -70,4 +73,5 @@ def __dir__():
"TaskRecord",
"TaskScheduler",
"TaskStatus",
"VramReservationManager",
]
79 changes: 79 additions & 0 deletions tinyagentos/scheduler/discovery.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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__)


Expand Down Expand Up @@ -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:

@kilo-code-bot kilo-code-bot Bot Jul 6, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

WARNING: _gpu_vram_probe falls back to 999_999 (effectively "unlimited free VRAM") whenever _probe_nvidia_vram() returns free == 0. That happens both when there is genuinely no NVIDIA GPU and when nvidia-smi is missing/failing — exactly the case where admission should be conservative, not optimistic. Returning a huge free value causes Resource.can_admit to accept any estimated_memory_mb against this resource and then OOM at runtime. If the probe can't read VRAM, prefer to report 0 so the GPU is skipped unless the cluster has a healthier worker.

Reply with @kilocode-bot fix it to have Kilo Code address this issue.

free, _total = _probe_nvidia_vram()
return free if free > 0 else 999_999 # optimistic
Comment on lines +216 to +218

@coderabbitai coderabbitai Bot Jul 6, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🩺 Stability & Availability | 🟠 Major | ⚡ Quick win

🧩 Analysis chain

🏁 Script executed:

#!/bin/bash
set -euo pipefail

python3 - <<'PY'
from pathlib import Path
path = Path('tinyagentos/scheduler/discovery.py')
text = path.read_text()
lines = text.splitlines()
for start, end in [(1, 260)]:
    print(f'লines {start}-{min(end, len(lines))}')
    for i in range(start-1, min(end, len(lines))):
        print(f"{i+1:4d}: {lines[i]}")
PY

Repository: jaylfc/taOS

Length of output: 11510

🏁 Script executed:

#!/bin/bash
set -euo pipefail

# Inspect symbols around the GPU probe helpers and any resource registration logic.
python3 - <<'PY'
from pathlib import Path
path = Path('tinyagentos/scheduler/discovery.py')
lines = path.read_text().splitlines()
for i, line in enumerate(lines, 1):
    if 'def _gpu_vram_probe' in line or '_probe_nvidia_vram' in line or 'Resource' in line or 'vram' in line:
        start = max(1, i-25)
        end = min(len(lines), i+40)
        print(f"\n--- context around line {i} ---")
        for j in range(start, end+1):
            print(f"{j:4d}: {lines[j-1]}")
PY

Repository: jaylfc/taOS

Length of output: 48442

🏁 Script executed:

#!/bin/bash
set -euo pipefail

# Search for the NVIDIA VRAM probe implementation and any documentation/comments
# describing what 0 or unknown means.
rg -n "_probe_nvidia_vram|gpu_vram_probe|vram" tinyagentos -S

Repository: jaylfc/taOS

Length of output: 24249

🏁 Script executed:

#!/bin/bash
set -euo pipefail

# Look for tests or docs around GPU resource discovery semantics.
rg -n "999_999|unknown|probe.*gpu|free.*vram|total.*vram|nvidia" . -S

Repository: jaylfc/taOS

Length of output: 50368

🏁 Script executed:

#!/bin/bash
set -euo pipefail

# Read only the relevant file sections with line numbers.
python3 - <<'PY'
from pathlib import Path
path = Path('tinyagentos/scheduler/discovery.py')
lines = path.read_text().splitlines()
for start, end in [(1, 260)]:
    for i in range(start-1, min(end, len(lines))):
        print(f"{i+1:4d}: {lines[i]}")
PY

Repository: jaylfc/taOS

Length of output: 11498

🏁 Script executed:

#!/bin/bash
set -euo pipefail

python3 - <<'PY'
from pathlib import Path
path = Path('tinyagentos/scheduler/discovery.py')
lines = path.read_text().splitlines()

targets = ['def _probe_nvidia_vram', 'def _gpu_vram_probe', 'Resource(', 'vram']
for idx, line in enumerate(lines, 1):
    if any(t in line for t in targets):
        start = max(1, idx-20)
        end = min(len(lines), idx+35)
        print(f"\n### around line {idx}")
        for j in range(start, end+1):
            print(f"{j:4d}: {lines[j-1]}")
PY

Repository: jaylfc/taOS

Length of output: 16576

🏁 Script executed:

#!/bin/bash
set -euo pipefail

# Inspect the NVIDIA VRAM probe and how Resource.memory_probe is used.
python3 - <<'PY'
from pathlib import Path

for file_path, ranges in [
    ('tinyagentos/scheduler/gpu_arbiter.py', [(1, 90), (140, 180)]),
    ('tinyagentos/scheduler/resource.py', [(1, 260)]),
]:
    path = Path(file_path)
    lines = path.read_text().splitlines()
    print(f"\n### {file_path} ({len(lines)} lines)")
    for start, end in ranges:
        print(f"\n--- lines {start}-{min(end, len(lines))} ---")
        for i in range(start-1, min(end, len(lines))):
            print(f"{i+1:4d}: {lines[i]}")
PY

Repository: jaylfc/taOS

Length of output: 14635

Separate probe failure from real 0-free VRAM. tinyagentos/scheduler/discovery.py:216-218 maps both a saturated GPU and any probe failure (_probe_nvidia_vram() returns 0, 0) to 999_999, so the resource looks effectively unbounded under pressure. Return 0 for actual zero-free cases and only use the optimistic fallback when the probe is unavailable.

🤖 Prompt for AI Agents 
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tinyagentos/scheduler/discovery.py` around lines 216 - 218, The
_gpu_vram_probe helper currently treats both a real 0-free GPU and a probe
failure from _probe_nvidia_vram() as the same optimistic 999_999 value. Update
_gpu_vram_probe so it preserves an actual zero-free result as 0, and only
returns the optimistic fallback when the probe is unavailable or invalid,
keeping the behavior localized to the GPU discovery path in
tinyagentos/scheduler/discovery.py.


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
Comment on lines +190 to +231

@coderabbitai coderabbitai Bot Jul 6, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win

gpu_backend_types omits the image-generation backends named in the comment.

The comment (Lines 190-192) and the PR objective both call out sd-cpp/sd-gpu (and CUDA llama-cpp) as GPU backends, but gpu_backend_types = {"vllm", "ollama", "exo", "mlx"} excludes them. As a result _gpu_capabilities/_gpu_backend_for will never surface image-generation on the GPU resource even though it is in GPU_POTENTIAL_CAPABILITIES, so GPU image-gen won't route through the admission path.

Confirm whether the image-gen backend types should be included here.

🤖 Prompt for AI Agents 
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tinyagentos/scheduler/discovery.py` around lines 190 - 231, The GPU backend
filter in discovery is too narrow and misses backends explicitly called out by
the comment and admission logic. Update gpu_backend_types in the GPU detection
block inside the discovery routine so it includes the image-generation backend
types sd-cpp and sd-gpu, and also the CUDA llama-cpp variant if it should
participate in GPU scheduling; then make sure _gpu_capabilities and
_gpu_backend_for continue to recognize these backends when building the
ResourceSignature and routing by capability.


for gpu_idx in range(gpu_count):
gpu_name = f"gpu-cuda-{gpu_idx}"

@kilo-code-bot kilo-code-bot Bot Jul 6, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

WARNING: The GPU resource is registered as gpu-cuda-{gpu_idx} unconditionally, even when the detected platform is ROCm, Metal, Vulkan, or generic "gpu". On a ROCm host you'll register gpu-cuda-0, on Apple Silicon gpu-cuda-0, etc. — misleading for logs, metrics, UI, and any downstream code that switches behavior based on the name prefix. Build the name from the platform, e.g. f"gpu-{gpu_signature.platform}-{gpu_idx}".

Reply with @kilocode-bot fix it to have Kilo Code address this issue.

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(
Expand Down
Loading
Loading