feat(scheduler): GPU arbiter — VRAM-accounted admission + queue + eviction#1683
feat(scheduler): GPU arbiter — VRAM-accounted admission + queue + eviction#1683hognek wants to merge 6 commits into
Conversation
|
Caution Review failedThe pull request is closed. ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Pro Plus Run ID: 📒 Files selected for processing (3)
📝 WalkthroughWalkthroughThis PR adds a GPU arbiter for VRAM-based admission, queueing, and eviction, wires it into scheduler discovery and app startup, and extends task metadata plus tests for reservation and eviction behavior. ChangesGPU Arbiter Feature
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant App
participant GpuArbiter
participant ClusterManager
participant Scheduler
App->>GpuArbiter: construct(scheduler, cluster_manager, max_queue_size=100, eviction_enabled=True)
App->>GpuArbiter: start()
App->>ClusterManager: _gpu_arbiter = arbiter
GpuArbiter->>GpuArbiter: submit_gpu(task, required_vram_mb)
GpuArbiter->>GpuArbiter: _check_admission()
alt VRAM available
GpuArbiter->>ClusterManager: claim_lease(resource_id)
GpuArbiter->>Scheduler: submit(task) / payload(None)
Scheduler-->>GpuArbiter: result
GpuArbiter->>ClusterManager: release_lease()
else queued
GpuArbiter->>GpuArbiter: enqueue task
GpuArbiter->>GpuArbiter: _process_queue() retries admission
GpuArbiter->>GpuArbiter: evict_lowest_priority()
end
Possibly related PRs
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
|
👋 Thanks for the PR! This one targets See CONTRIBUTING.md for the branch model. |
|
Done — retargeted to |
|
|
Qodo reviews are paused for this user.Troubleshooting steps vary by plan Learn more → On a Teams plan? Using GitHub Enterprise Server, GitLab Self-Managed, or Bitbucket Data Center? |
![kilo-code-bot[bot]](https://avatars.githubusercontent.com/in/2193792?s=60&v=4){kind=small}
| return 0, 0 | ||
|
|
||
|
|
||
| def _probe_nvidia_vram() -> tuple[int, int]: |
There was a problem hiding this comment.
CRITICAL: _probe_nvidia_vram only reads the first GPU line (stdout.split("\n")[0]), so on a multi-GPU host it always reports GPU 0's free VRAM regardless of which gpu-cuda-N resource is being probed. When gpu_count > 1 (the TODO says Phase 2, but the probe is already shared) every GPU will claim GPU 0's free VRAM, which breaks admission for GPUs 1+ and can cause OOM/Xid 62 crashes on those devices.
| def _probe_nvidia_vram() -> tuple[int, int]: | |
| free_lines = [l for l in free_raw.stdout.strip().split("\n") if l.strip()] | |
| total_lines = [l for l in total_raw.stdout.strip().split("\n") if l.strip()] | |
| return int(free_lines[0]), int(total_lines[0]) |
Also consider indexing by CUDA_VISIBLE_DEVICES or accepting a gpu index parameter so per-resource probes are correct.
Reply with @kilocode-bot fix it to have Kilo Code address this issue.
| victim_priority, victim_id = pri, tid | ||
| if victim_id is None: | ||
| return 0 | ||
| return self._evict_task(victim_id) |
There was a problem hiding this comment.
CRITICAL: _evict_task calls future.cancel() on task._arbiter_future, but that attribute is only set on tasks that were queued via submit_gpu and parked on a future. Tasks that were directly admitted and are now in _running (i.e. the common case for evict_lowest_priority callers) have no _arbiter_future, so getattr(..., None) returns None and the cancel is silently a no-op. The scheduler task that is actually doing GPU work keeps running and the lease has already been released, so the next admission for that VRAM will succeed against freed memory that is still in use — exactly the Xid 62 crash the arbiter is supposed to prevent.
You need to track the live asyncio.Task running the payload (e.g. set a _running_task field when calling self._scheduler.submit(task) in _run_gpu_task) and cancel that task from _evict_task instead of the (possibly absent) future.
Reply with @kilocode-bot fix it to have Kilo Code address this issue.
| 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", | ||
| ) |
There was a problem hiding this comment.
WARNING: _drain_queue breaks as soon as it admits and starts one task, even though the loop condition is while not self._queue.empty(). With the 2-second processor interval, draining N admitted tasks from a backlog takes ~2N seconds — under load the queue depth grows faster than it drains. Either keep admitting while capacity remains (drop the break) or at least document that the arbiter admits at most one task per tick.
Reply with @kilocode-bot fix it to have Kilo Code address this issue.
| 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( |
There was a problem hiding this comment.
WARNING: evict_lowest_priority selects the victim by pri > victim_priority (keeping the highest value seen), but Priority is defined as lower integer wins (INTERACTIVE_USER=10, BATCH=40). The variable name victim_priority is misleading — the chosen victim is actually the least important task. Consider renaming to worst_priority / victim_priority_value and adding a comment so future readers don't assume this evicts the highest priority.
Reply with @kilocode-bot fix it to have Kilo Code address this issue.
| self._submitted += 1 | ||
| if required_vram_mb > 0: | ||
| admission = self._check_admission(task, required_vram_mb) | ||
| if not admission.admitted: |
There was a problem hiding this comment.
SUGGESTION: When self._scheduler is None you fall through to await task.payload(None). The Task dataclass types payload as Callable[["Resource"], Awaitable[Any]] — payload authors will reasonably assume a Resource is always passed. A Resource-less None can crash payloads that dereference it (e.g. resource.backend_url_for(...)). Either always run through the scheduler (refuse to start the arbiter without one), or skip the GPU path entirely when no scheduler is configured.
Reply with @kilocode-bot fix it to have Kilo Code address this issue.
| runtime_version="", | ||
| ) | ||
|
|
||
| def _gpu_vram_probe() -> int: |
There was a problem hiding this comment.
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.
| @@ -0,0 +1,216 @@ | |||
| """Tests for the GPU VRAM arbiter (taOS #894 Slice 2).""" | |||
There was a problem hiding this comment.
CRITICAL: The PR description claims "19 unit tests pass", but this file cannot be collected: it imports VramAllocation from tinyagentos.scheduler.gpu_arbiter (only GpuArbiter and GpuAdmission are defined), instantiates Resource(..., gpu_arbiter=arbiter) (Resource.__init__ has no such kwarg — it only knows about memory_probe), and calls arbiter.reserve(...), arbiter.find_eviction_candidates(...), arbiter.evict_and_reserve(...), arbiter.wait_for_vram(), and arbiter.allocations, none of which exist on GpuArbiter (the actual API is submit_gpu + private _check_admission/_run_gpu_task). Either the tests need to be rewritten against the real API or the module needs the missing public surface — as shipped, the test suite fails at import.
Reply with @kilocode-bot fix it to have Kilo Code address this issue.
| payload=lambda r: asyncio.sleep(0), preferred_resources=[], | ||
| priority=priority, estimated_vram_mb=vram_mb, | ||
| ) | ||
|
|
There was a problem hiding this comment.
CRITICAL: _make_resource builds a Resource with gpu_arbiter=arbiter, but the current Resource constructor (tinyagentos/scheduler/resource.py) has no such parameter — it accepts name, signature, concurrency, get_capabilities, backend_lookup, tier, potential_capabilities, score_lookup, memory_probe. The PR neither adds VRAM accounting to Resource.can_admit nor wires the arbiter in, so the TestResourceIntegration block is testing a feature that does not exist in this PR. If VRAM-aware Resource.can_admit was intended, it is missing; if not, the test is dead.
Reply with @kilocode-bot fix it to have Kilo Code address this issue.
| 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 |
There was a problem hiding this comment.
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.
| submitter: str = "unknown" | ||
| estimated_seconds: float = 1.0 | ||
| estimated_memory_mb: int = 0 | ||
| estimated_vram_mb: int = 0 |
There was a problem hiding this comment.
SUGGESTION: Task.estimated_vram_mb defaults to 0, which GpuArbiter._check_admission treats as "no VRAM requirement" (skips admission entirely). That is convenient but also means any new Task constructor that forgets to set it will silently bypass GPU admission control. Consider logging a warning when a GPU-routed task arrives with estimated_vram_mb == 0, or refuse to route zero-VRAM tasks to GPU resources by default.
Reply with @kilocode-bot fix it to have Kilo Code address this issue.
Code Review SummaryStatus: 1 Issue Found | Recommendation: Address before merge Overview
Issue Details (click to expand)WARNING
Files Reviewed (3 changed files in incremental diff vs. d51641a)
Prior findings outside this incremental diff (still active)
Fix these issues in Kilo Cloud Previous Review Summaries (5 snapshots, latest commit d51641a)Current summary above is authoritative. Previous snapshots are kept for context only. Previous review (commit d51641a)Status: 1 Issue Found | Recommendation: Address before merge Overview
Issue Details (click to expand)WARNING
Files Reviewed (2 changed files in incremental diff vs. 74b1185)
Pre-existing findings from the previous review ( Fix these issues in Kilo Cloud Previous review (commit 74b1185)Status: 3 Issues Found | Recommendation: Address before merge Overview
Issue Details (click to expand)CRITICAL
WARNING
SUGGESTION
Files Reviewed (6 changed files at HEAD 74b1185)
The Fix these issues in Kilo Cloud Previous review (commit d193b46)Status: 3 Issues Found | Recommendation: Address before merge Overview
Issue Details (click to expand)CRITICAL
WARNING
SUGGESTION
Files Reviewed (1 changed file in this incremental commit)
Fix these issues in Kilo Cloud Previous review (commit 52c6003)Status: No New Issues Found | Recommendation: Re-review after addressing prior findings Overview
Issue Details (click to expand)CRITICAL
WARNING
SUGGESTION
Prior findings still unfixed (carry-over, already commented)The fix commit (preemption via asyncio.Task cancellation) correctly addresses the
Files Reviewed (4 changed files in this fix commit)
Previous review (commit 1532fef)Status: 3 New Issues Found | Recommendation: Address before merge Overview
Incremental review against
All 10 previously-reported findings (4 CRITICAL on Issue Details (click to expand)CRITICAL
WARNING
SUGGESTION
Files Reviewed (2 changed files in this commit)
Reviewed by hy3-20260706:free · Input: 84.8K · Output: 17.7K · Cached: 189.9K |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 7
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
tests/test_gpu_arbiter.py (1)
1-217: 🗄️ Data Integrity & Integration | 🔴 Critical | 🏗️ Heavy liftAlign this suite with the current
GpuArbitercontract.
tinyagentos/scheduler/gpu_arbiter.pyexposessubmit_gpu/_check_admission/evict_lowest_priorityand astats()shape withoutallocations; it does not provide the VRAM-ledger API this file assumes (reserve,release,can_admit,find_eviction_candidates,evict_and_reserve,wait_for_vram, orVramAllocation). This file will fail at collection, andResource.can_admit()here also doesn't apply a VRAM gate.🤖 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 `@tests/test_gpu_arbiter.py` around lines 1 - 217, The test suite is written against an older VRAM-ledger API that `GpuArbiter` no longer exposes, so update `tests/test_gpu_arbiter.py` to match the current contract. Replace uses of `reserve`, `release`, `can_admit`, `find_eviction_candidates`, `evict_and_reserve`, `wait_for_vram`, and `VramAllocation` with assertions around `submit_gpu`, `_check_admission`, `evict_lowest_priority`, and `stats()`. Also remove or rewrite the `Resource.can_admit()` VRAM-specific expectations in `TestResourceIntegration`, since that path no longer performs a VRAM gate in this contract.
🤖 Prompt for all review comments with 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.
Inline comments:
In `@tinyagentos/app.py`:
- Around line 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.
- 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.
In `@tinyagentos/scheduler/discovery.py`:
- Around line 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.
- Around line 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.
In `@tinyagentos/scheduler/gpu_arbiter.py`:
- Around line 143-145: The eviction counter is being incremented twice for the
same queued task: once in _evict_task and again in submit_gpu’s
asyncio.CancelledError handler. Update submit_gpu (and the eviction path around
_evict_task) so only one of these locations records an eviction, and keep the
CancelledError block limited to cleanup/re-raise without bumping self._evicted
again.
- Around line 151-158: In `_check_admission()` within `GpuArbiter`, the current
`free_vram > 0` check conflates a saturated GPU with no probe data, allowing
admission when `free_vram == 0`; use the `_vram_probe()` total (or another
explicit presence signal) to distinguish an absent/unprobed GPU from a full GPU,
and only admit when a real GPU is present and VRAM is sufficient, otherwise
return a non-admitted `GpuAdmission` with the appropriate reason.
- Around line 206-230: Cancel the actual running coroutine when a task is
evicted, not just the queued submission future. Update evict_lowest_priority()
and _evict_task() in gpu_arbiter.py so that, after removing the task from
_running and releasing the lease, they also cancel the task object or its
running handle created by _run_gpu_task(), since _arbiter_future only covers
queued submissions. Make sure the eviction path handles both queued and
already-admitted tasks consistently.
---
Outside diff comments:
In `@tests/test_gpu_arbiter.py`:
- Around line 1-217: The test suite is written against an older VRAM-ledger API
that `GpuArbiter` no longer exposes, so update `tests/test_gpu_arbiter.py` to
match the current contract. Replace uses of `reserve`, `release`, `can_admit`,
`find_eviction_candidates`, `evict_and_reserve`, `wait_for_vram`, and
`VramAllocation` with assertions around `submit_gpu`, `_check_admission`,
`evict_lowest_priority`, and `stats()`. Also remove or rewrite the
`Resource.can_admit()` VRAM-specific expectations in `TestResourceIntegration`,
since that path no longer performs a VRAM gate in this contract.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro Plus
Run ID: a6cd67dd-4e34-47a5-b87e-3841f3ec9391
📒 Files selected for processing (6)
tests/test_gpu_arbiter.pytinyagentos/app.pytinyagentos/scheduler/__init__.pytinyagentos/scheduler/discovery.pytinyagentos/scheduler/gpu_arbiter.pytinyagentos/scheduler/types.py
|
|
||
| # 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 |
There was a problem hiding this comment.
🩺 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.
| eviction_enabled=True, | ||
| ) | ||
| await gpu_arbiter.start() | ||
| app.state.gpu_arbiter = gpu_arbiter |
There was a problem hiding this comment.
🩺 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.
| # 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 |
There was a problem hiding this comment.
🎯 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.
| def _gpu_vram_probe() -> int: | ||
| free, _total = _probe_nvidia_vram() | ||
| return free if free > 0 else 999_999 # optimistic |
There was a problem hiding this comment.
🩺 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]}")
PYRepository: 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]}")
PYRepository: 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 -SRepository: 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" . -SRepository: 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]}")
PYRepository: 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]}")
PYRepository: 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]}")
PYRepository: 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.
|
Thanks Hogne — read through the whole arbiter. Strong foundation: clean VRAM accounting, admission via the atomic cluster leases (#1680), a priority queue with a FIFO tiebreak, and good unit coverage of the primitives. All the bots passed it, but I dug into the concurrency and there are a few real gaps I'd want sorted before this goes on a live path. Note it's currently dormant — it's wired into 1. Eviction doesn't actually stop running GPU work (the core one)
2. Local-VRAM admission is TOCTOUOn the no-cluster path, 3. Double lease-releaseWhen a task is evicted while also completing, 4. Queue processor serializes on the running task
5. Minor: inconsistent locking on
|
|
Thanks for the thorough review — exactly the kind of feedback that makes this worth doing. Created 5 follow-up tasks for all five points:
They'll get picked up by the worker pipeline and pushed as new commits on this branch. One question on workflow while we're here: we've been keeping PRs as draft until CI passes, then marking ready — but the CI matrix never actually runs because the fork needs maintainer workflow approval (the action_required gate). Is that the intended flow, or should we be doing something differently? We want to make sure we're not adding friction on your side. |
|
Perfect, thanks for splitting all five into tasks — that's exactly right. On the CI question: you've spotted a real friction point. Because your PRs come from a fork, GitHub gates the workflow runs behind maintainer approval (the
No changes needed on your side beyond flipping ready when the code's done. Appreciate you flagging it rather than letting it sit. |
|
All contributors have signed the CLA ✍️ ✅ |
|
I have read the CLA Document and I hereby sign the CLA |
|
recheck |
Review feedback fix: eviction preemption (#894)The latest commit (95876da) addresses the review feedback at #1683#issuecomment-4897343798. Problem:
Fix: Tests: 9 new tests in CI is awaiting approval (first-time contributor policy). |
| return self._evict_task(victim_id) | ||
|
|
||
| def _evict_task(self, task_id: str) -> int: | ||
| if task_id not in self._running: |
There was a problem hiding this comment.
CRITICAL: Race condition: _evict_task checks task_id not in self._running (line 232) and then self._running.pop(task_id) (line 234) without acquiring self._running_lock. Meanwhile _run_gpu_task's finally block DOES hold the lock for the same pop. If a task completes naturally between line 232 and line 234, the .pop() here will raise KeyError, crashing eviction. The new preemption logic added in this commit (asyncio task cancel + lease release) widens the window where this race can fire. Fix by wrapping the check+pop in async with self._running_lock: (and making _evict_task async, or using a sync lock for _running).
Reply with @kilocode-bot fix it to have Kilo Code address this issue.
| assert "t-noevict" in arbiter._running | ||
|
|
||
| # Clean up | ||
| arbiter._running_tasks.get("t-noevict", None) |
There was a problem hiding this comment.
WARNING: Dead code — arbiter._running_tasks.get("t-noevict", None) returns a value that is immediately discarded. It does not cancel the task. Cleanup actually happens on the next line via submit_coro.cancel(). This line is misleading (looks like it should be submit_coro = ... or arbiter._running_tasks["t-noevict"].cancel()) and should either be removed or made to actually cancel the tracked asyncio task.
Reply with @kilocode-bot fix it to have Kilo Code address this issue.
| evicted = arbiter.evict_lowest_priority() | ||
| assert evicted == 1 | ||
|
|
||
| # Low-priority (BATCH=50) should be evicted, high (INTERACTIVE_USER=10) stays |
There was a problem hiding this comment.
SUGGESTION: Comment typo — BATCH=50 should be BATCH=40 per the Priority enum in tinyagentos/scheduler/types.py:30.
Reply with @kilocode-bot fix it to have Kilo Code address this issue.
|
I have read the CLA Document and I hereby sign the CLA |
|
recheck |
| f"GPU lease claim failed for {resource_id} (task {task.id})" | ||
| ) | ||
| lease_id = lease.lease_id | ||
| current = asyncio.current_task() |
There was a problem hiding this comment.
CRITICAL: current = asyncio.current_task() captures the wrong task for the queued-admission path.
When _run_gpu_task is awaited from _drain_queue (line 272), which itself is awaited from _process_queue, asyncio.current_task() returns the _process_queue task — not a task dedicated to running the user's payload. Calling asyncio_task.cancel() on it in _evict_task (line 248) cancels the entire queue processor, killing all future task admission and draining.
For a queued task, eviction of t-queued should stop the user payload (so its _arbiter_future propagates CancelledError to the submitter), but the _process_queue loop must keep running to admit the next queued task. Instead, the queue processor dies and no further queued work is processed for the lifetime of the GpuArbiter.
The fix must capture a task that is specific to running the user's work (e.g. a freshly created asyncio.create_task that wraps task.payload(...) / _scheduler.submit(task)), not the caller's task.
The test test_evict_queued_admitted_task_cancels_submitter in tests/test_gpu_arbiter_894.py does not exercise the real _drain_queue path — it calls arbiter._run_gpu_task(...) directly, so the captured task is the test's own runner and this bug is hidden.
Reply with @kilocode-bot fix it to have Kilo Code address this issue.
| if lease_id is not None and self._cluster_manager is not None: | ||
| self._cluster_manager.release_lease(lease_id) | ||
| # Cancel the running asyncio Task — stops _run_gpu_task and frees VRAM | ||
| asyncio_task = self._running_tasks.pop(task_id, None) |
There was a problem hiding this comment.
WARNING: _evict_task mutates self._running_tasks and calls asyncio_task.cancel() without holding self._running_lock.
_run_gpu_task's finally block (line 209) registers writes/deletes to both self._running and self._running_tasks under _running_lock. _evict_task pops _running atomically (good, fixes the prior TOCTOU) but then does self._running_tasks.pop(task_id, None) and asyncio_task.cancel() outside the lock.
While CPython dict ops are atomic under the GIL, the invariant — that task_id maps to the same task under both views for the duration of the eviction — is no longer enforced. A subtle interleaving with a concurrent _run_gpu_task for a re-used task.id could cancel a task that belongs to a different logical submission. Holding the lock around the read/cancel (or moving the cancel inside the finally-pop in _run_gpu_task via a sentinel) would keep the invariant explicit.
Reply with @kilocode-bot fix it to have Kilo Code address this issue.
| self._admitted += 1 | ||
| return result | ||
| except asyncio.CancelledError: | ||
| logger.info("gpu-arbiter: task %s preempted via CancelledError (pri=%d, vram=%d)", |
There was a problem hiding this comment.
SUGGESTION: Log message "preempted via CancelledError" is misleading — CancelledError is not eviction-specific.
This handler is reached for any CancelledError raised inside _run_gpu_task's try, including:
- shutdown via
await self.stop(), - the parent task (e.g. the arbiter's caller) being cancelled for unrelated reasons,
- asyncio task-group cancellations.
In all of these cases the log line will read as if an eviction happened, but _evicted was not incremented by _evict_task and no lease was released via the eviction path. Recommend distinguishing the eviction case (e.g. a flag set by _evict_task before asyncio_task.cancel()) from generic cancellation in the log message.
Reply with @kilocode-bot fix it to have Kilo Code address this issue.
…iter _drain_queue no longer blocks on _run_gpu_task. The admitted task is spawned as a background asyncio Task with a done-callback that propagates the result/exception to the submitter's _arbiter_future. This keeps the drain loop responsive on subsequent ticks so eviction-to-make-room can kick in when higher-priority tasks arrive while VRAM is full. Also adds eviction-to-make-room: when admission fails for a queued task the drain now calls evict_lowest_priority with the task's priority as the floor, then re-checks admission. Lower-priority running tasks are evicted to make room; higher-priority runners are left alone. Queue processing remains intentionally serial — only one task is admitted per drain cycle to avoid flooding the GPU with concurrent loads. PR jaylfc#1683 review feedback jaylfc#4.
…iter _drain_queue no longer blocks on _run_gpu_task. The admitted task is spawned as a background asyncio Task with a done-callback that propagates the result/exception to the submitter's _arbiter_future. This keeps the drain loop responsive on subsequent ticks so eviction-to-make-room can kick in when higher-priority tasks arrive while VRAM is full. Also adds eviction-to-make-room: when admission fails for a queued task the drain now calls evict_lowest_priority with the task's priority as the floor, then re-checks admission. Lower-priority running tasks are evicted to make room; higher-priority runners are left alone. Queue processing remains intentionally serial — only one task is admitted per drain cycle to avoid flooding the GPU with concurrent loads. PR jaylfc#1683 review feedback jaylfc#4.
|
Status check on the arbiter package: all five review findings are now covered across this PR plus #1705 (TOCTOU), #1706 (non-blocking drain + eviction), and #1707 (running-lock). Nice work, especially the double-release guard. One thing before I land the set: #1706 and #1707 currently conflict with dev. Could you rebase them onto the current #1683/#1705 stack so the whole thing applies as one clean sequence? I want to merge the arbiter as a single coherent package rather than piecemeal, so once those two are clean I'll review the set end to end and merge together. |
|
On the arbiter and cluster epic (#1683, #1689, #1690, #1705, #1706, #1707): I did a deep pass across all of these and they overlap heavily. #1689 re-adds the same gpu_arbiter.py as this PR and shares the #1706/#1707 file footprint, and #1690's drain/pause edits the same manager.py methods as #1689. Before I merge any of it, can we consolidate into one ordered stack? Concretely: this PR plus the #1705/#1706/#1707 fixes should own the arbiter core; close or rebase #1689 on top rather than have two PRs adding gpu_arbiter.py; then sequence #1690 after so its drain/pause builds on the arbiter's lease model. Two things to fix while consolidating: cluster_manager._gpu_arbiter is never wired in app.py, so eviction is a no-op in production (it only works in tests because they pass the kwarg), and #1690's drain uses startswith(name) lease matching, which releases other workers' leases (the collision we fixed before, so use an exact worker-name compare). Once it is one clean sequence I will review the whole package end to end and merge together. |
|
Agreed on all of it. Plan: make #1683 + #1705/#1706/#1707 the arbiter core as one ordered stack, evaluate #1689 and either close it or rebase only its unique pause/resume+detach bits on top [not a second gpu_arbiter.py], and sequence #1690 last so its drain/pause builds on the lease model. Will also fix the two real ones you caught — wiring cluster_manager._gpu_arbiter in app.py so eviction is not a prod no-op, and switching #1690 drain to exact worker-name lease matching instead of startswith. Ill ping here when the whole sequence is clean for your end-to-end pass. |
…iter _drain_queue no longer blocks on _run_gpu_task. The admitted task is spawned as a background asyncio Task with a done-callback that propagates the result/exception to the submitter's _arbiter_future. This keeps the drain loop responsive on subsequent ticks so eviction-to-make-room can kick in when higher-priority tasks arrive while VRAM is full. Also adds eviction-to-make-room: when admission fails for a queued task the drain now calls evict_lowest_priority with the task's priority as the floor, then re-checks admission. Lower-priority running tasks are evicted to make room; higher-priority runners are left alone. Queue processing remains intentionally serial — only one task is admitted per drain cycle to avoid flooding the GPU with concurrent loads. PR jaylfc#1683 review feedback jaylfc#4.
…M reservation _addmission_check read live nvidia-smi free VRAM with no reservation, so two concurrent submit_gpu calls could both pass before either model loaded — the exact concurrent-load crash (NVIDIA Xid 62) the arbiter exists to prevent. Add _reserved_vram_mb counter and _pending_reservations dict, guarded by _reservation_lock. _reserve_and_check() atomically checks admission against (probe - _reserved_vram_mb) and reserves on success. Reservations are released in _run_gpu_task's finally block and in _evict_task. _stats() now includes reserved_vram_mb and pending_reservations for observability. Fixes: PR jaylfc#1683 review feedback #2 Tests: 168 pass (25 arbiter + 143 cluster/worker)
…iter _drain_queue no longer blocks on _run_gpu_task. The admitted task is spawned as a background asyncio Task with a done-callback that propagates the result/exception to the submitter's _arbiter_future. This keeps the drain loop responsive on subsequent ticks so eviction-to-make-room can kick in when higher-priority tasks arrive while VRAM is full. Also adds eviction-to-make-room: when admission fails for a queued task the drain now calls evict_lowest_priority with the task's priority as the floor, then re-checks admission. Lower-priority running tasks are evicted to make room; higher-priority runners are left alone. Queue processing remains intentionally serial — only one task is admitted per drain cycle to avoid flooding the GPU with concurrent loads. PR jaylfc#1683 review feedback jaylfc#4.
…ction 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, jaylfc#893 leases). Task: t_d7208884. Fixes jaylfc#894.
… futures _evict_task previously only cancelled the _arbiter_future on the Task object. Directly-admitted tasks never have one set, and for queued-then- admitted tasks the future is resolved by _drain_queue before _run_gpu_task completes. In both cases the running GPU work keeps executing and VRAM is never freed until natural completion. Now _run_gpu_task registers its asyncio.Task in _running_tasks, and _evict_task cancels that Task directly. CancelledError propagates to the payload, the finally block cleans up VRAM/leases, and eviction actually stops running GPU work. - Added _running_tasks dict to track asyncio.Tasks - _run_gpu_task registers current task and handles CancelledError - _evict_task cancels the asyncio.Task (in addition to _arbiter_future) - _drain_queue catches CancelledError from evicted tasks to keep queue processor alive - Finally block checks if task was already evicted (entry popped) before releasing lease — prevents double-release Fixes jaylfc#894.
_evict_task now uses pop(task_id, None) atomically instead of the check-then-pop pattern, eliminating the TOCTOU KeyError that could occur when _run_gpu_task's finally block pops the entry between the existence check and the pop. Ownership is now explicit: whoever pops self._running first releases the lease. _evict_task's pop is the sole lease releaser for evicted tasks; _run_gpu_task's finally only releases on normal completion (when it still finds the entry). Comments document the invariant. Refs jaylfc#894. Task: t_0c86d08e.
…ncompatible with stack The file was incorrectly brought in from the jaylfc#1689 version and references VramAllocation, total_vram_mb, headroom_mb, and evict_callback that don't exist in the stack's GpuArbiter. test_gpu_arbiter_894.py covers the same territory with the correct API.
…M reservation _addmission_check read live nvidia-smi free VRAM with no reservation, so two concurrent submit_gpu calls could both pass before either model loaded — the exact concurrent-load crash (NVIDIA Xid 62) the arbiter exists to prevent. Add _reserved_vram_mb counter and _pending_reservations dict, guarded by _reservation_lock. _reserve_and_check() atomically checks admission against (probe - _reserved_vram_mb) and reserves on success. Reservations are released in _run_gpu_task's finally block and in _evict_task. _stats() now includes reserved_vram_mb and pending_reservations for observability. Fixes: PR jaylfc#1683 review feedback #2 Tests: 168 pass (25 arbiter + 143 cluster/worker)
…iter _drain_queue no longer blocks on _run_gpu_task. The admitted task is spawned as a background asyncio Task with a done-callback that propagates the result/exception to the submitter's _arbiter_future. This keeps the drain loop responsive on subsequent ticks so eviction-to-make-room can kick in when higher-priority tasks arrive while VRAM is full. Also adds eviction-to-make-room: when admission fails for a queued task the drain now calls evict_lowest_priority with the task's priority as the floor, then re-checks admission. Lower-priority running tasks are evicted to make room; higher-priority runners are left alone. Queue processing remains intentionally serial — only one task is admitted per drain cycle to avoid flooding the GPU with concurrent loads. PR jaylfc#1683 review feedback jaylfc#4.
|
Update on the arbiter package: #1689 is closed [it duplicated the gpu_arbiter core; its concurrency fixes are covered by #1706/#1707]. The remaining stack PRs — #1683, #1705, #1706, #1707, #1690 — are all individually MERGEABLE vs dev again [the #1706/#1707 dev conflicts are resolved]. Were landing your two flagged fixes now: wiring cluster_manager._gpu_arbiter in app.py [so eviction is not a prod no-op] and switching #1690 drain to exact worker-name lease matching. On merging as one coherent package: fully retargeting all the bases into a single sequence keeps exceeding our automation limits — since the PRs are individually clean, would you rather merge them in order [#1683 → #1705 → #1706 → #1707 → #1690] yourself, or should we formally retarget each PRs base onto the one below it? Either works on our end. |
…t a prod no-op Initialize _gpu_arbiter = None in ClusterManager.__init__ and wire cluster_manager._gpu_arbiter = gpu_arbiter in app.py after arbiter creation, matching the _capabilities pattern. Without this, the arbiter is only reachable via app.state.gpu_arbiter and prod eviction paths that go through the cluster manager are dead.
| ) | ||
| await gpu_arbiter.start() | ||
| app.state.gpu_arbiter = gpu_arbiter | ||
| cluster_manager._gpu_arbiter = gpu_arbiter # wire for eviction paths |
There was a problem hiding this comment.
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.
There was a problem hiding this comment.
Actionable comments posted: 5
🧹 Nitpick comments (2)
tests/test_gpu_arbiter_894.py (2)
188-208: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winSimulated "queued" test never exercises the real dequeue path.
The docstring says "queued-then-admitted", but the test calls
arbiter._run_gpu_taskdirectly rather than going throughsubmit_gpu+ the actual_queue/_drain_queueconsumption. This validates eviction on a task that was never truly queued, so preemption behavior for tasks dequeued via the background drain loop remains unverified by this suite.Consider driving this test through
submit_gpuwith VRAM pressure so the task is genuinely queued and then admitted by_drain_queue, to close the coverage gap on the real production path.
[recommended_refactor]🤖 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 `@tests/test_gpu_arbiter_894.py` around lines 188 - 208, The “queued-then-admitted” test is bypassing the real queueing flow by calling arbiter._run_gpu_task directly, so it does not cover the actual submit_gpu -> _queue -> _drain_queue path. Update the test to enqueue the task through submit_gpu under VRAM pressure and let _drain_queue admit it, then verify _evict_task behavior and cancellation on that real queued task using the existing arbiter, runner_task, and payload_cancelled checks.
195-201: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick winFixed
sleep(0.1)for synchronization is flaky.Waiting a fixed 100ms for the task to register in
_running(vs. using anasyncio.Eventas other tests in this file do) can be flaky under CI load, causing intermittent false failures.♻️ Suggested fix: poll instead of a fixed sleep
- runner_task = asyncio.create_task(runner()) - - # Give it a moment to register in _running - await asyncio.sleep(0.1) - - assert "t-queued" in arbiter._running + runner_task = asyncio.create_task(runner()) + + # Poll until registered in _running (avoids fixed-sleep flakiness) + for _ in range(50): + if "t-queued" in arbiter._running: + break + await asyncio.sleep(0.01) + assert "t-queued" in arbiter._running🤖 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 `@tests/test_gpu_arbiter_894.py` around lines 195 - 201, The test in the runner setup is relying on a fixed asyncio.sleep(0.1) to wait for the task to appear in arbiter._running, which is flaky under load. Replace this synchronization in the runner/running-state test with a deterministic wait strategy, such as polling until "t-queued" is present or using an asyncio.Event pattern like the other tests in this file. Keep the assertion against arbiter._running, but ensure the test only proceeds once runner() has actually registered its task.
🤖 Prompt for all review comments with 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.
Inline comments:
In `@tinyagentos/scheduler/gpu_arbiter.py`:
- Around line 180-201: The GPU task bookkeeping in _run_gpu_task and the
eviction path does not preserve the evictable flag, so evict_lowest_priority()
can still cancel non-evictable work. Store evictable in the _running entry for
each task, update the running-task tuple shape consistently wherever it is read
or written, and change evict_lowest_priority() to consider only idle or
evictable tasks as valid victims. Make sure the same TaskTracker/_running
bookkeeping is used everywhere so submit_gpu()’s evictable argument is enforced
end to end.
- Around line 118-159: The admission path in submit_gpu and _check_admission is
not accounting for VRAM already reserved by other queued/running tasks, so two
submissions can be admitted against the same free budget. Update the GPU arbiter
to make admission and reservation atomic by protecting the free-VRAM check plus
reservation update with a single lock, and subtract any reserved VRAM before
deciding whether to admit. Use the existing submit_gpu, _check_admission, and
_running reservation bookkeeping to ensure reservations are recorded before
returning admission.
- Around line 47-55: The GPU queue entry in _QueuedGpuTask currently ignores
waiting time, so low-priority tasks can starve and there is no max_wait_ms
enforcement. Update the scheduling logic in gpu_arbiter.py to use queued_at for
aging and check max_wait_ms when selecting from the queue, likely in the
queue-pop/dispatch path around the GPU arbiter methods that consume
_QueuedGpuTask. Ensure the priority comparison still works, but boost or bypass
entries that have waited too long so older tasks eventually run, and
reject/expire tasks whose wait exceeds max_wait_ms.
- Around line 143-146: The queued GPU work is not being dropped when the
submitter’s waiting future is already cancelled, so stale entries can still be
admitted and run later; update the queue admission path in the arbiter logic
around _drain_queue and the submit/done handling to check the stored
_arbiter_future before scheduling work, and skip/remove entries whose future is
already cancelled or done. Apply the same guard in the related submit/queue code
referenced by the later block so cancelled callers do not leave executable work
behind.
- Around line 160-191: The cluster GPU admission result is selecting a worker,
but that chosen resource is not being preserved into the execution path, so
queued work reaches _run_gpu_task with resource_id unset and never claims a
lease. Update the admission/scheduling flow around the GPU arbiter so the
selected worker/resource_id is returned or carried forward from the admission
check, and ensure the queued path passes that same resource_id into
_run_gpu_task when invoking the lease claim logic. Verify the resource is
threaded through any helper used by the cluster queue path so claim_lease()
always receives the selected cluster resource.
---
Nitpick comments:
In `@tests/test_gpu_arbiter_894.py`:
- Around line 188-208: The “queued-then-admitted” test is bypassing the real
queueing flow by calling arbiter._run_gpu_task directly, so it does not cover
the actual submit_gpu -> _queue -> _drain_queue path. Update the test to enqueue
the task through submit_gpu under VRAM pressure and let _drain_queue admit it,
then verify _evict_task behavior and cancellation on that real queued task using
the existing arbiter, runner_task, and payload_cancelled checks.
- Around line 195-201: The test in the runner setup is relying on a fixed
asyncio.sleep(0.1) to wait for the task to appear in arbiter._running, which is
flaky under load. Replace this synchronization in the runner/running-state test
with a deterministic wait strategy, such as polling until "t-queued" is present
or using an asyncio.Event pattern like the other tests in this file. Keep the
assertion against arbiter._running, but ensure the test only proceeds once
runner() has actually registered its task.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro Plus
Run ID: 684d79bb-ffa2-441a-869f-d7c21f79aa19
📒 Files selected for processing (7)
tests/test_gpu_arbiter_894.pytinyagentos/app.pytinyagentos/cluster/manager.pytinyagentos/scheduler/__init__.pytinyagentos/scheduler/discovery.pytinyagentos/scheduler/gpu_arbiter.pytinyagentos/scheduler/types.py
✅ Files skipped from review due to trivial changes (1)
- tinyagentos/cluster/manager.py
🚧 Files skipped from review as they are similar to previous changes (4)
- tinyagentos/scheduler/init.py
- tinyagentos/scheduler/types.py
- tinyagentos/scheduler/discovery.py
- tinyagentos/app.py
| @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) |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | 🏗️ Heavy lift
Implement aging and max_wait_ms enforcement.
The queue orders only by (priority, seq); queued_at is informational and there is no max_wait_ms path. Low-priority entries can starve indefinitely, which misses the Phase 2 acceptance requirements.
Also applies to: 79-86, 266-290
🤖 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/gpu_arbiter.py` around lines 47 - 55, The GPU queue
entry in _QueuedGpuTask currently ignores waiting time, so low-priority tasks
can starve and there is no max_wait_ms enforcement. Update the scheduling logic
in gpu_arbiter.py to use queued_at for aging and check max_wait_ms when
selecting from the queue, likely in the queue-pop/dispatch path around the GPU
arbiter methods that consume _QueuedGpuTask. Ensure the priority comparison
still works, but boost or bypass entries that have waited too long so older
tasks eventually run, and reject/expire tasks whose wait exceeds max_wait_ms.
| 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) |
There was a problem hiding this comment.
🩺 Stability & Availability | 🔴 Critical | 🏗️ Heavy lift
Make admission and local VRAM reservation atomic.
_check_admission() admits from probed free VRAM only; _running reservations are recorded later and never subtracted. Two GPU submissions can both pass admission against the same free VRAM budget before the driver reflects usage, defeating the arbiter’s OOM prevention goal. Track reserved VRAM under one lock and subtract it during admission/registration.
Also applies to: 192-196
🤖 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/gpu_arbiter.py` around lines 118 - 159, The admission
path in submit_gpu and _check_admission is not accounting for VRAM already
reserved by other queued/running tasks, so two submissions can be admitted
against the same free budget. Update the GPU arbiter to make admission and
reservation atomic by protecting the free-VRAM check plus reservation update
with a single lock, and subtract any reserved VRAM before deciding whether to
admit. Use the existing submit_gpu, _check_admission, and _running reservation
bookkeeping to ensure reservations are recorded before returning admission.
| return await done | ||
| except asyncio.CancelledError: | ||
| self._evicted += 1 | ||
| raise |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
Skip queued work after the submitter is cancelled.
If the caller awaiting done is cancelled, the queue entry remains and _drain_queue() can later execute the GPU task even though no submitter can receive the result. Check the stored _arbiter_future before admission and drop entries that are already cancelled/done.
Also applies to: 266-275
🤖 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/gpu_arbiter.py` around lines 143 - 146, The queued GPU
work is not being dropped when the submitter’s waiting future is already
cancelled, so stale entries can still be admitted and run later; update the
queue admission path in the arbiter logic around _drain_queue and the
submit/done handling to check the stored _arbiter_future before scheduling work,
and skip/remove entries whose future is already cancelled or done. Apply the
same guard in the related submit/queue code referenced by the later block so
cancelled callers do not leave executable work behind.
| 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 | ||
| current = asyncio.current_task() | ||
| async with self._running_lock: | ||
| self._running[task.id] = (task, lease_id, int(task.priority), required_vram_mb) | ||
| if current is not None: | ||
| self._running_tasks[task.id] = current | ||
| try: | ||
| if self._scheduler is not None: | ||
| result = await self._scheduler.submit(task) | ||
| else: | ||
| result = await task.payload(None) |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Respect evictable before cancelling running GPU work.
evictable is accepted by submit_gpu() but never stored in _running; evict_lowest_priority() can therefore cancel any running task, including non-evictable foreground work. Store the flag with the running entry and filter victims to idle/evictable work only.
Also applies to: 220-248
🤖 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/gpu_arbiter.py` around lines 180 - 201, The GPU task
bookkeeping in _run_gpu_task and the eviction path does not preserve the
evictable flag, so evict_lowest_priority() can still cancel non-evictable work.
Store evictable in the _running entry for each task, update the running-task
tuple shape consistently wherever it is read or written, and change
evict_lowest_priority() to consider only idle or evictable tasks as valid
victims. Make sure the same TaskTracker/_running bookkeeping is used everywhere
so submit_gpu()’s evictable argument is enforced end to end.
|
Thanks for consolidating. Closing #1689 and confirming the stack is individually mergeable vs dev is exactly what I wanted. Yes, merging in dependency order works for me. Once your two fixes land (the app.py cluster_manager._gpu_arbiter wiring, and the exact worker-name lease matching in #1690), I will review the whole set end to end (#1683, #1705, #1706, #1707, #1690) and then merge them in that order. I want the full-set review first so I can confirm all five findings are covered and consistent before anything lands, but the strict-order merge itself is fine and avoids the base-retarget automation limits. Ping me when the two fixes are pushed. |
|
Both fixes are pushed and ready for your end-to-end review: [1] cluster_manager._gpu_arbiter is now wired in app.py [commit on #1683: fix(cluster): wire gpu_arbiter into cluster_manager], so eviction is no longer a prod no-op; [2] #1690 drain now uses an exact worker-name compare instead of startswith. The full set — #1683, #1705, #1706, #1707, #1690 — is individually mergeable vs dev and ready for the whole-set review + strict-order merge whenever you have a window. Thanks! |
|
Did the full end-to-end review of the stack (#1683, #1705, #1706, #1707, plus #1690). Solid foundation and the five findings are all mechanically addressed, but the composed state has consistency bugs that block merging as a package. CI being green is misleading here: CodeRabbit never actually reviewed #1705/#1706/#1707 (rate-limited), so Kilo is the only real reviewer and it flagged most of these. Blockers, per PR: #1707 (the big one): evict_lowest_priority became async here but the call site in _drain_queue (gpu_arbiter.py around :378 and :383) still calls it synchronously. At runtime that returns a coroutine, #1706: the post-eviction re-admit calls _check_admission instead of _reserve_and_check, so it spawns a task with no reservation entry and re-opens the TOCTOU window plus drifts the _reserved_vram_mb accounting. #1705: the lease claim raises before the try/finally, so a failed claim never releases the reservation and _reserved_vram_mb leaks until restart. Move the claim inside the try/finally. Also the cluster-mode admission path does not subtract reserved VRAM, so finding #2 is only closed for the local path. #1690: the new /drain, /cancel-drain and /update endpoints have no auth, unlike the sibling cluster routes. Any LAN client can drain or offline a worker or trigger a remote worker update/restart. Please gate them with _require_admin (and the manual-claim rate limiter). Also the force-release path mutates self._leases outside _lease_lock while the sweep runs under it, which can race. Cross-cutting (most important for the epic gate): wiring cluster_manager._gpu_arbiter in app.py is in, but nothing actually calls it. drain_worker(graceful=False) and the monitor stale-drain branch force-release the cluster lease but never ask the arbiter to cancel the running task, and nothing anywhere calls submit_gpu, so the arbiter is constructed, started, and idle. Until the drain path notifies the arbiter and a real admission path routes through submit_gpu, eviction is still a prod no-op regardless of the app.py wiring. Minimum to unblock the merge-in-order: the #1707 awaits, the #1690 auth, the #1705 reservation-release, and the drain to arbiter wiring. Once those are in I will re-review the composed state and merge the stack in order. Happy to pair on the wiring piece if useful. |
…iter _drain_queue no longer blocks on _run_gpu_task. The admitted task is spawned as a background asyncio Task with a done-callback that propagates the result/exception to the submitter's _arbiter_future. This keeps the drain loop responsive on subsequent ticks so eviction-to-make-room can kick in when higher-priority tasks arrive while VRAM is full. Also adds eviction-to-make-room: when admission fails for a queued task the drain now calls evict_lowest_priority with the task's priority as the floor, then re-checks admission. Lower-priority running tasks are evicted to make room; higher-priority runners are left alone. Queue processing remains intentionally serial — only one task is admitted per drain cycle to avoid flooding the GPU with concurrent loads. PR jaylfc#1683 review feedback jaylfc#4.
|
Update on the arbiter blockers — all addressed, but a heads-up on structure. The fixes [#1707 async drain awaits, #1705 reservation-release inside try/finally, #1706 TOCTOU re-admit via _reserve_and_check, #1690 auth + _lease_lock on force-release, and the cross-cutting drain->arbiter task-cancellation wiring] landed CONSOLIDATED on one rebased branch — PR #1718 [feat/worker-auto-update-rebased, currently mergeable, Kilo findings addressed] — rather than distributed across the individual stack branches. That does not match your merge-the-stack-in-order plan. Your call: |
… GpuArbiter Add VramReservationManager as the arbiter's INTERNAL reservation layer. Replaces ad-hoc vram_probe with atomic check-and-reserve under an asyncio.Lock so the arbiter is the single VRAM authority. - VramReservationManager: atomic reserve()/release(), sync available() for best-effort admission checks, stats() for observability - GpuArbiter._check_admission now reads against reserved VRAM - GpuArbiter._run_gpu_task atomically reserves local VRAM before execution and releases in finally (normal completion) or via _evict_task (eviction path) - Cluster lease path unchanged (VRAM tracked by lease system) - Exported via scheduler.__init__ lazy loading - 21 new tests: 11 unit (VramReservationManager), 9 integration (arbiter lifecycle), 1 lazy-export test - 59/59 related tests pass (arbiter + cluster + scheduler + vram)
| # read; the actual reserve happens atomically in _run_gpu_task). | ||
| free_vram = self._vram_reservations.available("local") | ||
| _free_raw, total = self._vram_probe() | ||
| if total > 0: |
There was a problem hiding this comment.
WARNING: Admission no longer falls through to the cluster path when the local GPU is fully reserved
Previously the gate was if free_vram > 0; it is now if total > 0. With a local GPU present (total > 0) but fully reserved (free == 0), a VRAM-needing task now returns admitted=False and never reaches the self._cluster_manager cluster-worker check below. In a mixed deployment (local GPU + cluster workers) this wrongly rejects work that could run on a cluster worker whenever local VRAM is exhausted. If cluster fallback is intended, gate on whether the task targets local vs cluster (e.g. resource_id) or restore the fall-through when local VRAM is insufficient but total > 0.
Reply with @kilocode-bot fix it to have Kilo Code address this issue.
2 participants
Summary
Adds the GPU arbiter (Slice 2 of taOS #894) — VRAM-accounted scheduler admission control with queuing and priority-based eviction.
Builds on:
What's included
New module:
tinyagentos/scheduler/gpu_arbiter.pyGPU resource registration (
discovery.py)gpu-cuda-0resource when GPU backends are healthymemory_probeon the GPU ResourceWiring (
app.py)app.state.gpu_arbiterTests: 19 unit tests pass
tests/test_gpu_arbiter.py: VRAM admission (4), queue (3), eviction (3), stats (3), probes (2), lifecycle (2), dataclass (2)Types
estimated_vram_mb: int = 0field toTaskdataclass (backward-compatible, defaults to 0)Design decisions
evictableflag.Task: t_d7208884. Fixes #894.
Summary by CodeRabbit