fix(scheduler): non-blocking drain + eviction-to-make-room in GPU arbiter#1706
Closed
hognek wants to merge 8 commits into
Closed
fix(scheduler): non-blocking drain + eviction-to-make-room in GPU arbiter#1706hognek wants to merge 8 commits into
hognek wants to merge 8 commits into
Conversation
📝 WalkthroughWalkthroughThis PR introduces a ChangesGPU Arbiter Admission Control
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant Caller
participant GpuArbiter
participant VramProbe
participant ClusterManager
participant Task
Caller->>GpuArbiter: submit_gpu(task, required_vram_mb)
GpuArbiter->>GpuArbiter: _reserve_and_check()
GpuArbiter->>VramProbe: probe free VRAM
VramProbe-->>GpuArbiter: free/total VRAM
alt admitted
GpuArbiter->>ClusterManager: claim lease (optional)
GpuArbiter->>Task: run payload via _run_gpu_task
Task-->>GpuArbiter: result/exception
GpuArbiter->>ClusterManager: release lease
GpuArbiter->>GpuArbiter: _release_reservation()
GpuArbiter-->>Caller: result
else not admitted
GpuArbiter->>GpuArbiter: enqueue task, attach _arbiter_future
GpuArbiter-->>Caller: awaits _arbiter_future
end
sequenceDiagram
participant AppLifespan
participant ResourceScheduler
participant GpuArbiter
participant ClusterManager
AppLifespan->>ResourceScheduler: build_scheduler()
AppLifespan->>GpuArbiter: GpuArbiter(resource_scheduler, cluster_manager, max_queue_size=100, eviction_enabled=True)
AppLifespan->>GpuArbiter: start()
GpuArbiter-->>AppLifespan: app.state.gpu_arbiter set
Possibly related PRs
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 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 |
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}
| "All tasks completed; worker detached gracefully.", | ||
| level="info", | ||
| ) | ||
| elif (now - worker.last_heartbeat) > HEARTBEAT_TIMEOUT: |
There was a problem hiding this comment.
CRITICAL: Stale-draining path force-releases leases but never notifies the GPU arbiter to cancel running tasks. The worker's heartbeats are gone, so the asyncio.Tasks tracked in GpuArbiter._running will continue executing against a dead worker indefinitely, holding VRAM and blocking any new admission on that worker until their natural completion. Call self._gpu_arbiter.release_tasks_for_worker(worker.name) (or equivalent) before popping the leases, mirroring the not graceful path in drain_worker().
Reply with @kilocode-bot fix it to have Kilo Code address this issue.
| def _evict_task(self, task_id: str) -> int: | ||
| if task_id not in self._running: | ||
| return 0 | ||
| task, lease_id, _pri, _vram = self._running.pop(task_id) |
There was a problem hiding this comment.
CRITICAL: _evict_task mutates self._running and self._running_tasks without holding self._running_lock, while _run_gpu_task's finally block reads/writes both dicts under the lock (and even calls release_lease outside the lock). Under CPython GIL the dict op is atomic, but the predicate if entry is not None in the finally relies on _evict_task having popped first — if the eviction races with natural completion, the lease can be released twice (once by _evict_task, once by the finally) or in the wrong order. Wrap the body of _evict_task in async with self._running_lock: (making it an async method or holding the lock via the caller's caller) so the eviction and natural completion paths are mutually exclusive.
Reply with @kilocode-bot fix it to have Kilo Code address this issue.
| # whose priority value is *lower* than N (i.e. tasks that are | ||
| # actually higher priority), so only lower-or-equal priority | ||
| # running tasks are candidates for eviction. | ||
| evicted = self.evict_lowest_priority(min_priority=int(entry.task.priority)) |
There was a problem hiding this comment.
WARNING: evict_lowest_priority(min_priority=int(entry.task.priority)) uses the predicate pri >= min_priority to choose victims, so a running task with the SAME priority as the queued task is eligible for eviction. The PR's stated goal is to preserve higher-priority runners and only evict strictly lower-priority work; equal-priority victims are unfair and will swap a freshly-running interactive task for an identical queued one. Either change the comparison to pri > min_priority (strict), or document this equal-priority eviction explicitly and surface it via stats().
Reply with @kilocode-bot fix it to have Kilo Code address this issue.
| future = getattr(entry.task, "_arbiter_future", None) | ||
| # Spawn as background task so drain doesn't block and | ||
| # eviction-to-make-room stays responsive on subsequent ticks. | ||
| t = asyncio.create_task( |
There was a problem hiding this comment.
SUGGESTION: The spawned t = asyncio.create_task(self._run_gpu_task(...)) is only kept alive by add_done_callback, which is correct, but the task has no name observers and logs from inside _run_gpu_task won't include a tag identifying the queue-drain that spawned it. Use name=f"gpu-arbiter-drain-{entry.task.id}" (already set, good) and also store t on the entry so future eviction logic can reach the exact asyncio Task without depending on the _running_tasks[task.id] mapping that the task itself populates after acquiring the running-lock — this avoids a window where _evict_task is called between create_task and the first await inside _run_gpu_task, in which case the running lock has not yet been entered and the cancellation has no effect.
Reply with @kilocode-bot fix it to have Kilo Code address this issue.
| # ── taOS #796: graceful worker detach ────────────────────────────────── | ||
|
|
||
|
|
||
| @router.post("/api/cluster/workers/{name}/drain") |
There was a problem hiding this comment.
WARNING: The docstring promises graceful is read from "query param or JSON body", but only the JSON body is parsed — a ?graceful=false query string is silently ignored. Either implement the query-param branch (request.query_params.get("graceful")) or update the docstring to say JSON-only. Note also that there is no path-level lock; if /drain is called twice concurrently for the same worker, both attempts pass the existence check before either mutates worker.status.
Reply with @kilocode-bot fix it to have Kilo Code address this issue.
| if lid is None: | ||
| continue | ||
| if self._cluster_manager is not None: | ||
| lease = getattr(self._cluster_manager, "_leases", {}).get(lid) |
There was a problem hiding this comment.
WARNING: release_tasks_for_worker reaches into self._cluster_manager._leases (private attribute) and on _evict_task does not verify that the eviction actually stopped the asyncio Task — it relies on _running_tasks being populated by _run_gpu_task after it acquires _running_lock. If the call lands before that registration, the task continues running and the function silently returns a count that overstates the cleanup it achieved. At minimum, check that the popped asyncio_task is non-None in _evict_task and log a warning when the eviction couldn't reach the asyncio layer; ideally, also clear the GpuArbiter's _running_lock ordering with ClusterManager._leases so the cleanup is atomic.
Reply with @kilocode-bot fix it to have Kilo Code address this issue.
Code Review SummaryStatus: No new issues from incremental commit | Recommendation: Address remaining carry-forwards before merge Overview
Incremental Re-Verification (db19abf..b616b84)Commit 7/7 (
No new issues introduced. Carry-Forward Findings (still valid against
|
| Severity | Count |
|---|---|
| CRITICAL | 1 |
| WARNING | 0 |
| SUGGESTION | 0 |
Incremental Re-Verification (4c544f8..db19abf)
No new issues introduced by commits 5/6 (TOCTOU reservation) and 6/6 (non-blocking drain + eviction-to-make-room). All previously flagged changed-code issues remain valid against current HEAD (db19abf85e):
tinyagentos/scheduler/gpu_arbiter.py:389(CRITICAL, TOCTOU post-eviction re-check) — unchanged. Commit 6/6 still uses bare_check_admissionafterevict_lowest_priority, bypassing_reservation_lockand re-opening the same TOCTOU window the rest of the PR closes. A concurrentsubmit_gpucan reserve the freed VRAM between the eviction and this re-check.tinyagentos/scheduler/gpu_arbiter.py:387(WARNING,evict_lowest_priorityevicts equal-priority runners) — unchanged.tinyagentos/scheduler/gpu_arbiter.py:394(SUGGESTION, spawned task not stored on entry / register-vs-cancel race) — unchanged.tinyagentos/scheduler/gpu_arbiter.py:_evict_task(CRITICAL, mutation without_running_lock) — unchanged in semantics. Atomicpop()is used butrelease_leaseandasyncio_task.cancel()still run outside the lock, so the cancellation can land after_run_gpu_task'sfinallyalready ran, with no way for_evict_taskto know the task completed naturally.tinyagentos/cluster/manager.py(stale-draining force-releases leases without notifyingGpuArbiter) — unchanged.tinyagentos/routes/cluster.py(gracefulquery param ignored) — unchanged.tinyagentos/scheduler/gpu_arbiter.py:release_tasks_for_worker(reaches into_cluster_manager._leases; eviction may not reach the asyncio Task) — unchanged.
New Issues
None.
Issue Details (click to expand)
CRITICAL
| File | Line | Issue |
|---|---|---|
tinyagentos/scheduler/gpu_arbiter.py |
389 | Post-eviction admission check bypasses the TOCTOU reservation and races with concurrent submitters. (carry-forward) |
Files Reviewed (incremental, 7 files)
tests/test_gpu_arbiter_894.py- 0 issues (test-only)tests/test_gpu_arbiter_toctou.py- 0 issues (new)tinyagentos/app.py- 0 issues (GPU arbiter wiring)tinyagentos/scheduler/__init__.py- 0 issues (export)tinyagentos/scheduler/discovery.py- 0 issues (GPU resource registration)tinyagentos/scheduler/gpu_arbiter.py- 0 NEW issues (carry-forwards only)tinyagentos/scheduler/types.py- 0 issues (estimated_vram_mbfield)
Previous review (commit 4c544f8)
Status: 1 New Issue Found | Recommendation: Address before merge
Overview
| Severity | Count |
|---|---|
| CRITICAL | 1 |
| WARNING | 0 |
| SUGGESTION | 0 |
Incremental Re-Verification (363bff1..HEAD)
Previous review issues on changed lines:
gpu_arbiter.py337 (CRITICAL, _evict_task mutates_running/_running_taskswithout_running_lock) — partially addressed:_evict_tasknow pops_runningfirst (line 331) and_run_gpu_task'sfinallyskips release when entry is None. Comment in code documents the ownership. Race remains for_running_tasks(both evict and finally call.pop()which is safe, but the asyncio.Task cancel race vsfinallyordering is still present — flag carries forward but not as new).gpu_arbiter.py387 (WARNING,evict_lowest_priority(min_priority=N)evicts equal-priority runners) — unchanged, carries forward.gpu_arbiter.py394 (SUGGESTION, spawn window before_running_lock) — unchanged, carries forward.
New Issues
CRITICAL
| File | Line | Issue |
|---|---|---|
tinyagentos/scheduler/gpu_arbiter.py |
389 | _drain_queue post-eviction admission re-check bypasses the TOCTOU reservation and races with concurrent submitters. |
Previously Resolved
- All taOS Cluster: pause/resume for benchmarks and agent queues, graceful worker detach, hardware-aware LLM call queuing #796 drain / cancel-drain / pause-resume features and
release_tasks_for_workerremoved cleanly; no orphan references. _release_reservationis correctly called from_run_gpu_taskfinally and from_evict_task, making the reservation lifecycle idempotent for the_run_gpu_taskcaller.
Issue Details (click to expand)
CRITICAL
| File | Line | Issue |
|---|---|---|
tinyagentos/scheduler/gpu_arbiter.py |
389 | _drain_queue post-eviction admission re-check bypasses the TOCTOU reservation and races with concurrent submitters. |
Files Reviewed (incremental, 9 files)
tests/test_cluster_drain.py- 0 issues (deleted)tests/test_gpu_arbiter.py- 0 issuestests/test_gpu_arbiter_796.py- 0 issues (deleted)tests/test_gpu_arbiter_toctou.py- 0 issues (new)tinyagentos/cluster/manager.py- 0 issuestinyagentos/routes/cluster.py- 0 issues (drain endpoints removed)tinyagentos/scheduler/gpu_arbiter.py- 1 issuetinyagentos/scheduler/resource.py- 0 issuestinyagentos/scheduler/vram_tracker.py- 0 issues (deleted)
Fix these issues in Kilo Cloud
Previous review (commit 1481206)
Status: 1 New Issue Found | Recommendation: Address before merge
Overview
| Severity | Count |
|---|---|
| CRITICAL | 1 |
| WARNING | 0 |
| SUGGESTION | 0 |
Incremental Re-Verification (363bff1..HEAD)
Previous review issues on changed lines:
gpu_arbiter.py337 (CRITICAL, _evict_task mutates_running/_running_taskswithout_running_lock) — partially addressed:_evict_tasknow pops_runningfirst (line 331) and_run_gpu_task'sfinallyskips release when entry is None. Comment in code documents the ownership. Race remains for_running_tasks(both evict and finally call.pop()which is safe, but the asyncio.Task cancel race vsfinallyordering is still present — flag carries forward but not as new).gpu_arbiter.py387 (WARNING,evict_lowest_priority(min_priority=N)evicts equal-priority runners) — unchanged, carries forward.gpu_arbiter.py394 (SUGGESTION, spawn window before_running_lock) — unchanged, carries forward.
New Issues
CRITICAL
| File | Line | Issue |
|---|---|---|
tinyagentos/scheduler/gpu_arbiter.py |
389 | _drain_queue post-eviction admission re-check bypasses the TOCTOU reservation and races with concurrent submitters. |
Previously Resolved
- All taOS Cluster: pause/resume for benchmarks and agent queues, graceful worker detach, hardware-aware LLM call queuing #796 drain / cancel-drain / pause-resume features and
release_tasks_for_workerremoved cleanly; no orphan references. _release_reservationis correctly called from_run_gpu_taskfinally and from_evict_task, making the reservation lifecycle idempotent for the_run_gpu_taskcaller.
Issue Details (click to expand)
CRITICAL
| File | Line | Issue |
|---|---|---|
tinyagentos/scheduler/gpu_arbiter.py |
389 | _drain_queue post-eviction admission re-check bypasses the TOCTOU reservation and races with concurrent submitters. |
Files Reviewed (incremental, 9 files)
tests/test_cluster_drain.py- 0 issues (deleted)tests/test_gpu_arbiter.py- 0 issuestests/test_gpu_arbiter_796.py- 0 issues (deleted)tests/test_gpu_arbiter_toctou.py- 0 issues (new)tinyagentos/cluster/manager.py- 0 issuestinyagentos/routes/cluster.py- 0 issues (drain endpoints removed)tinyagentos/scheduler/gpu_arbiter.py- 1 issuetinyagentos/scheduler/resource.py- 0 issuestinyagentos/scheduler/vram_tracker.py- 0 issues (deleted)
Fix these issues in Kilo Cloud
Previous review (commit 363bff1)
Status: 6 Issues Found | Recommendation: Address before merge
Overview
| Severity | Count |
|---|---|
| CRITICAL | 2 |
| WARNING | 3 |
| SUGGESTION | 1 |
Issue Details (click to expand)
CRITICAL
| File | Line | Issue |
|---|---|---|
tinyagentos/cluster/manager.py |
609 | Stale-draining path force-releases leases but never calls the GPU arbiter to cancel running tasks — orphan GPU workloads on a dead worker. |
tinyagentos/scheduler/gpu_arbiter.py |
337 | _evict_task mutates _running/_running_tasks without _running_lock, racing _run_gpu_task's finally block and risking double lease release. |
WARNING
| File | Line | Issue |
|---|---|---|
tinyagentos/scheduler/gpu_arbiter.py |
386 | evict_lowest_priority(min_priority=N) evicts running tasks with equal priority as the queued task, contradicting "preserve higher-priority runners". |
tinyagentos/routes/cluster.py |
863 | /drain docstring promises query-param graceful, but only JSON body is parsed; concurrent calls also race on worker.status. |
tinyagentos/scheduler/gpu_arbiter.py |
444 | release_tasks_for_worker reaches into private cluster_manager._leases and silently no-ops when asyncio_task isn't yet registered. |
SUGGESTION
| File | Line | Issue |
|---|---|---|
tinyagentos/scheduler/gpu_arbiter.py |
393 | Background drain task is only kept alive by add_done_callback; there is a window before _running_lock is acquired where eviction cannot reach the asyncio Task. |
Files Reviewed (19 files)
docs/design/a2a-gpu-lease-proposal.md- 0 issuestests/test_cluster_drain.py- 0 issuestests/test_cluster_worker_protocol.py- 0 issuestests/test_gpu_arbiter.py- 0 issuestests/test_gpu_arbiter_796.py- 0 issuestests/test_gpu_arbiter_894.py- 0 issuestests/test_leases.py- 0 issuestinyagentos/app.py- 0 issuestinyagentos/cluster/manager.py- 1 issuetinyagentos/cluster/worker_capacity.py- 0 issuestinyagentos/cluster/worker_protocol.py- 0 issuestinyagentos/routes/cluster.py- 1 issuetinyagentos/scheduler/__init__.py- 0 issuestinyagentos/scheduler/discovery.py- 0 issuestinyagentos/scheduler/gpu_arbiter.py- 4 issuestinyagentos/scheduler/resource.py- 0 issuestinyagentos/scheduler/types.py- 0 issuestinyagentos/scheduler/vram_tracker.py- 0 issuestinyagentos/worker/agent.py- 0 issues
Fix these issues in Kilo Cloud
Previous review (commit 7ba9342)
Status: 6 Issues Found | Recommendation: Address before merge
Overview
| Severity | Count |
|---|---|
| CRITICAL | 2 |
| WARNING | 3 |
| SUGGESTION | 1 |
Issue Details (click to expand)
CRITICAL
| File | Line | Issue |
|---|---|---|
tinyagentos/cluster/manager.py |
609 | Stale-draining path force-releases leases but never calls the GPU arbiter to cancel running tasks — orphan GPU workloads on a dead worker. |
tinyagentos/scheduler/gpu_arbiter.py |
337 | _evict_task mutates _running/_running_tasks without _running_lock, racing _run_gpu_task's finally block and risking double lease release. |
WARNING
| File | Line | Issue |
|---|---|---|
tinyagentos/scheduler/gpu_arbiter.py |
386 | evict_lowest_priority(min_priority=N) evicts running tasks with equal priority as the queued task, contradicting "preserve higher-priority runners". |
tinyagentos/routes/cluster.py |
863 | /drain docstring promises query-param graceful, but only JSON body is parsed; concurrent calls also race on worker.status. |
tinyagentos/scheduler/gpu_arbiter.py |
444 | release_tasks_for_worker reaches into private cluster_manager._leases and silently no-ops when asyncio_task isn't yet registered. |
SUGGESTION
| File | Line | Issue |
|---|---|---|
tinyagentos/scheduler/gpu_arbiter.py |
393 | Background drain task is only kept alive by add_done_callback; there is a window before _running_lock is acquired where eviction cannot reach the asyncio Task. |
Files Reviewed (19 files)
docs/design/a2a-gpu-lease-proposal.md- 0 issuestests/test_cluster_drain.py- 0 issuestests/test_cluster_worker_protocol.py- 0 issuestests/test_gpu_arbiter.py- 0 issuestests/test_gpu_arbiter_796.py- 0 issuestests/test_gpu_arbiter_894.py- 0 issuestests/test_leases.py- 0 issuestinyagentos/app.py- 0 issuestinyagentos/cluster/manager.py- 1 issuetinyagentos/cluster/worker_capacity.py- 0 issuestinyagentos/cluster/worker_protocol.py- 0 issuestinyagentos/routes/cluster.py- 1 issuetinyagentos/scheduler/__init__.py- 0 issuestinyagentos/scheduler/discovery.py- 0 issuestinyagentos/scheduler/gpu_arbiter.py- 4 issuestinyagentos/scheduler/resource.py- 0 issuestinyagentos/scheduler/types.py- 0 issuestinyagentos/scheduler/vram_tracker.py- 0 issuestinyagentos/worker/agent.py- 0 issues
Reviewed by minimax-m3 · Input: 63.4K · Output: 5.7K · Cached: 711.3K
|
|
hognek
force-pushed
the
fix/queue-processor-nonblocking-drain
branch
2 times, most recently
from
715c4cc to
363bff1
Compare
This was referenced Jul 7, 2026
hognek
force-pushed
the
fix/queue-processor-nonblocking-drain
branch
from
363bff1 to
1481206
Compare
| # running tasks are candidates for eviction. | ||
| evicted = self.evict_lowest_priority(min_priority=int(entry.task.priority)) | ||
| if evicted > 0: | ||
| admission = self._check_admission(entry.task, entry.required_vram_mb) |
There was a problem hiding this comment.
CRITICAL: Post-eviction admission check bypasses the TOCTOU reservation and races with concurrent submitters.
When evicted > 0, this re-check uses the bare _check_admission (no _reservation_lock, no _reserved_vram_mb update) before spawning _run_gpu_task. Two problems:
-
Bypassed reservation. The earlier
_reserve_and_checkat line 380 returnedadmitted=False, so no reservation was made forentry.task.id. The subsequent_run_gpu_taskwill run without holding a VRAM reservation, and itsfinallycalls_release_reservationon a never-registered id (no-op). This re-opens the exact TOCTOU window the rest of this PR closes — a concurrentsubmit_gpucan race and both pass admission while models load. -
Lock-free re-check. Between the eviction at line 387 and this call, another concurrent submitter can reserve VRAM via
_reservation_lock. The check here is racy.
Fix: take _reservation_lock and atomically check-and-reserve, e.g.
async with self._reservation_lock:
admission = self._check_admission(entry.task, entry.required_vram_mb)
if admission.admitted:
self._reserved_vram_mb += entry.required_vram_mb
self._pending_reservations[entry.task.id] = entry.required_vram_mband _release_reservation(entry.task.id) in the else branch (line 418) only for entries that were never reserved on this pass.
Reply with @kilocode-bot fix it to have Kilo Code address this issue.
hognek
force-pushed
the
fix/queue-processor-nonblocking-drain
branch
from
1481206 to
4c544f8
Compare
…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)
hognek
force-pushed
the
fix/queue-processor-nonblocking-drain
branch
from
4c544f8 to
db19abf
Compare
…ed VRAM in cluster path taOS jaylfc#1705 — reservation leak fix: - Move lease claim and _running registration inside try block so the finally always releases the VRAM reservation, even on claim failure. - Await claim_lease / release_lease (both are async on ClusterManager). - Make _evict_task and evict_lowest_priority async (they await release_lease). - Subtract _reserved_vram_mb in the cluster-mode admission path, matching the local-path behaviour so in-flight reservations are accounted for. - Update tests: FakeClusterManager methods made async, all _evict_task and evict_lowest_priority calls now awaited. Tests: 16/16 arbiter + 111/111 cluster 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.
taOS jaylfc#1706 — TOCTOU fix: - After eviction-to-make-room frees VRAM, the re-admit in _drain_queue used _check_admission (no reservation) instead of _reserve_and_check (atomic reservation). This spawned tasks with no reservation entry, re-opening the TOCTOU window and causing _reserved_vram_mb drift. - Changed to _reserve_and_check so post-eviction admission is atomic and properly tracked. - Also fixes the missing await on evict_lowest_priority (made async by the base jaylfc#1705 fix) and sync _evict_task calls in new tests. Tests: 22/22 arbiter pass.
hognek
force-pushed
the
fix/queue-processor-nonblocking-drain
branch
from
db19abf to
b616b84
Compare
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 10
🧹 Nitpick comments (2)
tests/test_gpu_arbiter_894.py (2)
122-156: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueDead no-op cleanup statement.
Line 151 calls
.get()and discards the result — it has no effect and appears to be leftover debug code.🧹 Proposed cleanup
assert await arbiter.evict_lowest_priority() == 0 assert "t-noevict" in arbiter._running # Clean up - arbiter._running_tasks.get("t-noevict", None) submit_coro.cancel()🤖 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 122 - 156, The test contains a dead no-op cleanup statement that calls GpuArbiter._running_tasks.get("t-noevict", None) and discards the result, so remove that unused line from test_evict_eviction_disabled_returns_zero. Keep the rest of the cleanup logic intact by cancelling submit_coro and awaiting it, and use the existing symbols GpuArbiter, submit_gpu, and evict_lowest_priority to locate the test.
16-22: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDuplicate helper across test files.
_make_taskis duplicated verbatim intests/test_gpu_arbiter_toctou.py(lines 16-22 there). Consider moving it to a sharedconftest.pyfixture/helper to avoid drift whenTask's constructor signature changes.🤖 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 16 - 22, The `_make_task` helper is duplicated across GPU arbiter tests, so update `test_gpu_arbiter_894.py` to use a shared helper instead of keeping a local copy. Move the common task निर्माण logic into a shared `conftest.py` fixture or test utility and have both `test_gpu_arbiter_894.py` and `test_gpu_arbiter_toctou.py` call that shared helper; keep the `Task` construction centralized so future constructor changes only need one update.
🤖 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 `@tests/test_gpu_arbiter_894.py`:
- Around line 422-476: The cleanup call in
test_drain_returns_immediately_after_spawning_task is missing an await, so the
_evict_task coroutine on GpuArbiter is never executed and can trigger an
unawaited-coroutine warning. Update the cleanup branch in
test_drain_returns_immediately_after_spawning_task to await
arbiter._evict_task("t-drain-fast"), matching the async usage of _evict_task
elsewhere in the test suite.
In `@tinyagentos/scheduler/discovery.py`:
- Around line 194-197: The GPU backend discovery logic is using a weaker
readiness check than the capability lookup, which can advertise backends that
are not actually schedulable. Update the GPU filtering in discovery.py (the
comprehension used for gpu_backends, and the later selection path in
backends_with_capability) to use the same readiness contract as the existing
backend lookup: require enabled, lifecycle_state == "running", and status ==
"ok" before advertising or returning a GPU backend. Keep the behavior consistent
across the code paths that build the GPU backend list and resolve the URL so
stopped or disabled backends are excluded everywhere.
- Around line 216-218: The _gpu_vram_probe helper currently returns an
optimistic 999_999 MB when _probe_nvidia_vram fails or reports no free VRAM,
which can bypass GPU admission checks; change the fallback in _gpu_vram_probe
(and any caller that relies on its result in the scheduler/discovery path) to
fail closed by returning a conservative value or an explicit unavailable signal
so GPU tasks are not over-admitted when VRAM cannot be probed.
In `@tinyagentos/scheduler/gpu_arbiter.py`:
- Around line 199-203: The await-on-done path in GPUArbiter is still admitting
canceled work, so a canceled submitter can leave a stale queue entry that
`_drain_queue()` later executes. Update the relevant `try/except
asyncio.CancelledError` handling in `gpu_arbiter.py` around the `done` await so
it checks whether the future was canceled before admission and skips/removes
that queued item instead of letting it proceed; apply the same guard in the
other matching cancel path mentioned in the diff.
- Around line 246-248: The lease variable name in the VRAM summation expression
is too ambiguous and triggers Ruff E741. Update the comprehension inside
gpu_arbiter.py’s worker lease calculation to use a descriptive name instead of a
single-letter alias, keeping the logic in the same place but making the symbol
clear and lint-compliant.
- Around line 410-414: The _propagate callback in gpu_arbiter.py currently
ignores all task cancellations, which can leave the submitter future unresolved
and block submit_gpu() awaiters. Update _propagate to distinguish cancellations
caused by _evict_task from unexpected cancellations, and in the unexpected case
resolve or cancel f so the awaiting caller is released. Use the existing
_propagate helper and the submit_gpu() future handling paths to keep the
cancellation flow consistent.
- Around line 172-178: The GPU submission path in submit_gpu currently ignores
Task.estimated_vram_mb when required_vram_mb is omitted, so default admission
should use the task’s estimate instead of silently treating it as zero. Update
submit_gpu to derive the effective VRAM requirement from required_vram_mb when
provided, otherwise fall back to task.estimated_vram_mb, and keep the
reservation/admission check in _reserve_and_check using that resolved value so
callers don’t need to pass it explicitly.
- Around line 123-130: The shutdown path in stop() only cancels
_queue_processor_task and leaves items already stored in _queue with their
_arbiter_future unresolved, so submit_gpu() callers can hang. Update
GpuArbiter.stop() to drain any remaining queued submitters and settle their
futures during shutdown, using the existing _queue_processor_task, _queue, and
_arbiter_future handling so every pending submit_gpu() returns or fails promptly
before stop() completes.
- Around line 221-264: In the GPU admission logic inside GpuArbiter’s VRAM
check, the current `if free_vram > 0` gate lets `(0, total)` and probe failures
like `(0, 0)` fall through to an admitted result. Change the local-path handling
to fail closed whenever VRAM is zero or the probe is unknown, returning
`GpuAdmission(admitted=False, ...)` with an explicit reason instead of
defaulting to admission. Keep the existing `effective_free` reservation
subtraction and cluster-manager fallback behavior intact, and ensure
`GpuAdmission` is only admitted on local GPUs when a positive, sufficient
`free_vram` is confirmed.
- Around line 30-43: The _probe_nvidia_vram helper is still relying on PATH
lookup for nvidia-smi and catching all exceptions, which is too broad. Resolve
the executable with shutil.which before running the probe, pass check=True to
the subprocess calls, and limit the exception handling to the expected
probe/parse failures only. Keep the fix localized to _probe_nvidia_vram so the
lookup and error handling are explicit and narrow.
---
Nitpick comments:
In `@tests/test_gpu_arbiter_894.py`:
- Around line 122-156: The test contains a dead no-op cleanup statement that
calls GpuArbiter._running_tasks.get("t-noevict", None) and discards the result,
so remove that unused line from test_evict_eviction_disabled_returns_zero. Keep
the rest of the cleanup logic intact by cancelling submit_coro and awaiting it,
and use the existing symbols GpuArbiter, submit_gpu, and evict_lowest_priority
to locate the test.
- Around line 16-22: The `_make_task` helper is duplicated across GPU arbiter
tests, so update `test_gpu_arbiter_894.py` to use a shared helper instead of
keeping a local copy. Move the common task निर्माण logic into a shared
`conftest.py` fixture or test utility and have both `test_gpu_arbiter_894.py`
and `test_gpu_arbiter_toctou.py` call that shared helper; keep the `Task`
construction centralized so future constructor changes only need one update.
🪄 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: dfd4d429-22c0-4416-a532-e265a8d3692f
📒 Files selected for processing (8)
.gitignoretests/test_gpu_arbiter_894.pytests/test_gpu_arbiter_toctou.pytinyagentos/app.pytinyagentos/scheduler/__init__.pytinyagentos/scheduler/discovery.pytinyagentos/scheduler/gpu_arbiter.pytinyagentos/scheduler/types.py
Comment on lines
+422
to
+476
| async def test_drain_returns_immediately_after_spawning_task(self): | ||
| """_drain_queue should return without waiting for the GPU task.""" | ||
| arbiter = GpuArbiter( | ||
| vram_probe=lambda: (8192, 8192), | ||
| max_queue_size=10, | ||
| ) | ||
|
|
||
| payload_started = asyncio.Event() | ||
| payload_done = asyncio.Event() | ||
|
|
||
| async def slow_payload(_resource): | ||
| payload_started.set() | ||
| await payload_done.wait() # Block until released | ||
| return {"slow": "done"} | ||
|
|
||
| task = _make_task("t-drain-fast", vram_mb=0) | ||
| task.payload = slow_payload | ||
|
|
||
| # Queue the task | ||
| loop = asyncio.get_running_loop() | ||
| arb_future: asyncio.Future = loop.create_future() | ||
| task._arbiter_future = arb_future # type: ignore[attr-defined] | ||
| entry = _QueuedGpuTask( | ||
| priority=int(task.priority), seq=1, task=task, | ||
| required_vram_mb=0, evictable=False, | ||
| ) | ||
| await arbiter._queue.put(entry) | ||
|
|
||
| # Drain — should return immediately after spawning | ||
| await arbiter._drain_queue() | ||
|
|
||
| # Payload was started (spawned as background task) | ||
| await asyncio.wait_for(payload_started.wait(), timeout=5) | ||
|
|
||
| # _drain_queue returned — queue should be empty | ||
| assert arbiter._queue.empty() | ||
|
|
||
| # Task should be in _running (registered by _run_gpu_task) | ||
| assert "t-drain-fast" in arbiter._running | ||
|
|
||
| # Release the payload | ||
| payload_done.set() | ||
|
|
||
| # arbiter_future should be resolved | ||
| result = await asyncio.wait_for(arb_future, timeout=5) | ||
| assert result is not None | ||
|
|
||
| # Cleanup — task should remove itself from _running on completion | ||
| # (the _run_gpu_task finally block handles this) | ||
| await asyncio.sleep(0.1) | ||
| # After completion, _running should be empty (or at least not contain t-drain-fast) | ||
| # Actually, with the new async model, _run_gpu_task removes itself from _running | ||
| # on completion, so let's check: | ||
| if "t-drain-fast" in arbiter._running: | ||
| arbiter._evict_task("t-drain-fast") |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
Missing await on async cleanup call.
_evict_task is a coroutine (awaited everywhere else, e.g. line 67, 203, 317), but here it's invoked without await, so the coroutine is created and immediately discarded, producing an unawaited-coroutine RuntimeWarning and skipping the intended cleanup.
🐛 Proposed fix
if "t-drain-fast" in arbiter._running:
- arbiter._evict_task("t-drain-fast")
+ await arbiter._evict_task("t-drain-fast")📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| async def test_drain_returns_immediately_after_spawning_task(self): | |
| """_drain_queue should return without waiting for the GPU task.""" | |
| arbiter = GpuArbiter( | |
| vram_probe=lambda: (8192, 8192), | |
| max_queue_size=10, | |
| ) | |
| payload_started = asyncio.Event() | |
| payload_done = asyncio.Event() | |
| async def slow_payload(_resource): | |
| payload_started.set() | |
| await payload_done.wait() # Block until released | |
| return {"slow": "done"} | |
| task = _make_task("t-drain-fast", vram_mb=0) | |
| task.payload = slow_payload | |
| # Queue the task | |
| loop = asyncio.get_running_loop() | |
| arb_future: asyncio.Future = loop.create_future() | |
| task._arbiter_future = arb_future # type: ignore[attr-defined] | |
| entry = _QueuedGpuTask( | |
| priority=int(task.priority), seq=1, task=task, | |
| required_vram_mb=0, evictable=False, | |
| ) | |
| await arbiter._queue.put(entry) | |
| # Drain — should return immediately after spawning | |
| await arbiter._drain_queue() | |
| # Payload was started (spawned as background task) | |
| await asyncio.wait_for(payload_started.wait(), timeout=5) | |
| # _drain_queue returned — queue should be empty | |
| assert arbiter._queue.empty() | |
| # Task should be in _running (registered by _run_gpu_task) | |
| assert "t-drain-fast" in arbiter._running | |
| # Release the payload | |
| payload_done.set() | |
| # arbiter_future should be resolved | |
| result = await asyncio.wait_for(arb_future, timeout=5) | |
| assert result is not None | |
| # Cleanup — task should remove itself from _running on completion | |
| # (the _run_gpu_task finally block handles this) | |
| await asyncio.sleep(0.1) | |
| # After completion, _running should be empty (or at least not contain t-drain-fast) | |
| # Actually, with the new async model, _run_gpu_task removes itself from _running | |
| # on completion, so let's check: | |
| if "t-drain-fast" in arbiter._running: | |
| arbiter._evict_task("t-drain-fast") | |
| async def test_drain_returns_immediately_after_spawning_task(self): | |
| """_drain_queue should return without waiting for the GPU task.""" | |
| arbiter = GpuArbiter( | |
| vram_probe=lambda: (8192, 8192), | |
| max_queue_size=10, | |
| ) | |
| payload_started = asyncio.Event() | |
| payload_done = asyncio.Event() | |
| async def slow_payload(_resource): | |
| payload_started.set() | |
| await payload_done.wait() # Block until released | |
| return {"slow": "done"} | |
| task = _make_task("t-drain-fast", vram_mb=0) | |
| task.payload = slow_payload | |
| # Queue the task | |
| loop = asyncio.get_running_loop() | |
| arb_future: asyncio.Future = loop.create_future() | |
| task._arbiter_future = arb_future # type: ignore[attr-defined] | |
| entry = _QueuedGpuTask( | |
| priority=int(task.priority), seq=1, task=task, | |
| required_vram_mb=0, evictable=False, | |
| ) | |
| await arbiter._queue.put(entry) | |
| # Drain — should return immediately after spawning | |
| await arbiter._drain_queue() | |
| # Payload was started (spawned as background task) | |
| await asyncio.wait_for(payload_started.wait(), timeout=5) | |
| # _drain_queue returned — queue should be empty | |
| assert arbiter._queue.empty() | |
| # Task should be in _running (registered by _run_gpu_task) | |
| assert "t-drain-fast" in arbiter._running | |
| # Release the payload | |
| payload_done.set() | |
| # arbiter_future should be resolved | |
| result = await asyncio.wait_for(arb_future, timeout=5) | |
| assert result is not None | |
| # Cleanup — task should remove itself from _running on completion | |
| # (the _run_gpu_task finally block handles this) | |
| await asyncio.sleep(0.1) | |
| # After completion, _running should be empty (or at least not contain t-drain-fast) | |
| # Actually, with the new async model, _run_gpu_task removes itself from _running | |
| # on completion, so let's check: | |
| if "t-drain-fast" in arbiter._running: | |
| await arbiter._evict_task("t-drain-fast") |
🤖 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 422 - 476, The cleanup call in
test_drain_returns_immediately_after_spawning_task is missing an await, so the
_evict_task coroutine on GpuArbiter is never executed and can trigger an
unawaited-coroutine warning. Update the cleanup branch in
test_drain_returns_immediately_after_spawning_task to await
arbiter._evict_task("t-drain-fast"), matching the async usage of _evict_task
elsewhere in the test suite.
Comment on lines
+194
to
+197
| gpu_backends = [ | ||
| b for b in catalog.backends() | ||
| if b.status == "ok" and b.type in gpu_backend_types | ||
| ] |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Use the same readiness contract for GPU capability advertising and backend lookup.
Lines 194 and 223 treat status == "ok" as enough, but backends_with_capability() also requires enabled and lifecycle_state == "running". A stopped/disabled GPU backend can be advertised as schedulable, then Line 230 returns no URL.
Proposed consistency fix
gpu_backend_types = {"vllm", "ollama", "exo", "mlx"}
+
+ def _is_ready_gpu_backend(backend) -> bool:
+ return (
+ backend.type in gpu_backend_types
+ and backend.status == "ok"
+ and backend.enabled
+ and backend.lifecycle_state == "running"
+ )
+
gpu_backends = [
b for b in catalog.backends()
- if b.status == "ok" and b.type in gpu_backend_types
+ if _is_ready_gpu_backend(b)
]
@@
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:
+ if _is_ready_gpu_backend(b):
caps |= b.capabilities
return capsAlso applies to: 220-230
🤖 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 194 - 197, The GPU backend
discovery logic is using a weaker readiness check than the capability lookup,
which can advertise backends that are not actually schedulable. Update the GPU
filtering in discovery.py (the comprehension used for gpu_backends, and the
later selection path in backends_with_capability) to use the same readiness
contract as the existing backend lookup: require enabled, lifecycle_state ==
"running", and status == "ok" before advertising or returning a GPU backend.
Keep the behavior consistent across the code paths that build the GPU backend
list and resolve the URL so stopped or disabled backends are excluded
everywhere.
Comment on lines
+216
to
+218
| 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
Fail closed when VRAM cannot be probed.
Line 218 reports 999_999 MB free when nvidia-smi fails or when the GPU is ROCm/Metal/native. That bypasses the VRAM admission guard and can over-admit GPU tasks.
Proposed safer fallback
def _gpu_vram_probe() -> int:
+ if gpu_type not in ("cuda", "nvidia"):
+ return 0 # no safe probe wired yet; fail closed
free, _total = _probe_nvidia_vram()
- return free if free > 0 else 999_999 # optimistic
+ return free if free > 0 else 0📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| def _gpu_vram_probe() -> int: | |
| free, _total = _probe_nvidia_vram() | |
| return free if free > 0 else 999_999 # optimistic | |
| def _gpu_vram_probe() -> int: | |
| if gpu_type not in ("cuda", "nvidia"): | |
| return 0 # no safe probe wired yet; fail closed | |
| free, _total = _probe_nvidia_vram() | |
| return free if free > 0 else 0 |
🤖 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 returns an optimistic 999_999 MB when
_probe_nvidia_vram fails or reports no free VRAM, which can bypass GPU admission
checks; change the fallback in _gpu_vram_probe (and any caller that relies on
its result in the scheduler/discovery path) to fail closed by returning a
conservative value or an explicit unavailable signal so GPU tasks are not
over-admitted when VRAM cannot be probed.
Comment on lines
+30
to
+43
| def _probe_nvidia_vram() -> tuple[int, int]: | ||
| """Probe free/total VRAM from nvidia-smi. Returns (free_mb, total_mb).""" | ||
| try: | ||
| import subprocess | ||
| free_raw = subprocess.run( | ||
| ["nvidia-smi", "--query-gpu=memory.free", "--format=csv,noheader,nounits"], | ||
| capture_output=True, text=True, timeout=5, | ||
| ) | ||
| total_raw = subprocess.run( | ||
| ["nvidia-smi", "--query-gpu=memory.total", "--format=csv,noheader,nounits"], | ||
| capture_output=True, text=True, timeout=5, | ||
| ) | ||
| return int(free_raw.stdout.strip().split("\n")[0]), int(total_raw.stdout.strip().split("\n")[0]) | ||
| except Exception: |
There was a problem hiding this comment.
🔒 Security & Privacy | 🟠 Major | ⚡ Quick win
Resolve nvidia-smi via an absolute path and narrow probe failures.
The command argv is constant, so this is not request-driven command injection, but Ruff still flags PATH-based executable lookup and blind exception handling here. Resolve nvidia-smi with shutil.which, use check=True, and catch only expected probe/parse failures.
Proposed cleanup
def _probe_nvidia_vram() -> tuple[int, int]:
"""Probe free/total VRAM from nvidia-smi. Returns (free_mb, total_mb)."""
try:
+ import shutil
import subprocess
- free_raw = subprocess.run(
- ["nvidia-smi", "--query-gpu=memory.free", "--format=csv,noheader,nounits"],
- capture_output=True, text=True, timeout=5,
- )
- total_raw = subprocess.run(
- ["nvidia-smi", "--query-gpu=memory.total", "--format=csv,noheader,nounits"],
- capture_output=True, text=True, timeout=5,
+ nvidia_smi = shutil.which("nvidia-smi")
+ if nvidia_smi is None:
+ return 0, 0
+ raw = subprocess.run(
+ [nvidia_smi, "--query-gpu=memory.free,memory.total", "--format=csv,noheader,nounits"],
+ capture_output=True, text=True, timeout=5, check=True,
)
- return int(free_raw.stdout.strip().split("\n")[0]), int(total_raw.stdout.strip().split("\n")[0])
- except Exception:
+ free_mb, total_mb = (int(part.strip()) for part in raw.stdout.strip().splitlines()[0].split(","))
+ return free_mb, total_mb
+ except (OSError, subprocess.SubprocessError, ValueError, IndexError):
return 0, 0📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| def _probe_nvidia_vram() -> tuple[int, int]: | |
| """Probe free/total VRAM from nvidia-smi. Returns (free_mb, total_mb).""" | |
| try: | |
| import subprocess | |
| free_raw = subprocess.run( | |
| ["nvidia-smi", "--query-gpu=memory.free", "--format=csv,noheader,nounits"], | |
| capture_output=True, text=True, timeout=5, | |
| ) | |
| total_raw = subprocess.run( | |
| ["nvidia-smi", "--query-gpu=memory.total", "--format=csv,noheader,nounits"], | |
| capture_output=True, text=True, timeout=5, | |
| ) | |
| return int(free_raw.stdout.strip().split("\n")[0]), int(total_raw.stdout.strip().split("\n")[0]) | |
| except Exception: | |
| def _probe_nvidia_vram() -> tuple[int, int]: | |
| """Probe free/total VRAM from nvidia-smi. Returns (free_mb, total_mb).""" | |
| try: | |
| import shutil | |
| import subprocess | |
| nvidia_smi = shutil.which("nvidia-smi") | |
| if nvidia_smi is None: | |
| return 0, 0 | |
| raw = subprocess.run( | |
| [nvidia_smi, "--query-gpu=memory.free,memory.total", "--format=csv,noheader,nounits"], | |
| capture_output=True, text=True, timeout=5, check=True, | |
| ) | |
| free_mb, total_mb = (int(part.strip()) for part in raw.stdout.strip().splitlines()[0].split(",")) | |
| return free_mb, total_mb | |
| except (OSError, subprocess.SubprocessError, ValueError, IndexError): | |
| return 0, 0 |
🧰 Tools
🪛 ast-grep (0.44.1)
[error] 33-36: Command coming from incoming request
Context: subprocess.run(
["nvidia-smi", "--query-gpu=memory.free", "--format=csv,noheader,nounits"],
capture_output=True, text=True, timeout=5,
)
Note: [CWE-78] Improper Neutralization of Special Elements used in an OS Command ('OS Command Injection').
(subprocess-from-request)
[error] 37-40: Command coming from incoming request
Context: subprocess.run(
["nvidia-smi", "--query-gpu=memory.total", "--format=csv,noheader,nounits"],
capture_output=True, text=True, timeout=5,
)
Note: [CWE-78] Improper Neutralization of Special Elements used in an OS Command ('OS Command Injection').
(subprocess-from-request)
🪛 Ruff (0.15.20)
[error] 35-35: Starting a process with a partial executable path
(S607)
[error] 39-39: Starting a process with a partial executable path
(S607)
[warning] 43-43: Do not catch blind exception: Exception
(BLE001)
🤖 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 30 - 43, The
_probe_nvidia_vram helper is still relying on PATH lookup for nvidia-smi and
catching all exceptions, which is too broad. Resolve the executable with
shutil.which before running the probe, pass check=True to the subprocess calls,
and limit the exception handling to the expected probe/parse failures only. Keep
the fix localized to _probe_nvidia_vram so the lookup and error handling are
explicit and narrow.
Source: Linters/SAST tools
Comment on lines
+123
to
+130
| async def stop(self) -> None: | ||
| if self._queue_processor_task is not None: | ||
| self._queue_processor_task.cancel() | ||
| try: | ||
| await self._queue_processor_task | ||
| except asyncio.CancelledError: | ||
| pass | ||
| self._queue_processor_task = None |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
Resolve queued submitters during shutdown.
stop() cancels only the queue processor; tasks already sitting in _queue keep their _arbiter_future pending, so callers awaiting submit_gpu() can hang after shutdown.
Proposed shutdown handling
async def stop(self) -> None:
if self._queue_processor_task is not None:
self._queue_processor_task.cancel()
try:
await self._queue_processor_task
except asyncio.CancelledError:
pass
self._queue_processor_task = None
+ while not self._queue.empty():
+ try:
+ entry = self._queue.get_nowait()
+ except asyncio.QueueEmpty:
+ break
+ self._release_reservation(entry.task.id)
+ future = getattr(entry.task, "_arbiter_future", None)
+ if future is not None and not future.done():
+ future.cancel()📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| async def stop(self) -> None: | |
| if self._queue_processor_task is not None: | |
| self._queue_processor_task.cancel() | |
| try: | |
| await self._queue_processor_task | |
| except asyncio.CancelledError: | |
| pass | |
| self._queue_processor_task = None | |
| async def stop(self) -> None: | |
| if self._queue_processor_task is not None: | |
| self._queue_processor_task.cancel() | |
| try: | |
| await self._queue_processor_task | |
| except asyncio.CancelledError: | |
| pass | |
| self._queue_processor_task = None | |
| while not self._queue.empty(): | |
| try: | |
| entry = self._queue.get_nowait() | |
| except asyncio.QueueEmpty: | |
| break | |
| self._release_reservation(entry.task.id) | |
| future = getattr(entry.task, "_arbiter_future", None) | |
| if future is not None and not future.done(): | |
| future.cancel() |
🤖 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 123 - 130, The shutdown
path in stop() only cancels _queue_processor_task and leaves items already
stored in _queue with their _arbiter_future unresolved, so submit_gpu() callers
can hang. Update GpuArbiter.stop() to drain any remaining queued submitters and
settle their futures during shutdown, using the existing _queue_processor_task,
_queue, and _arbiter_future handling so every pending submit_gpu() returns or
fails promptly before stop() completes.
Comment on lines
+172
to
+178
| 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 = await self._reserve_and_check(task.id, required_vram_mb) |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
Default GPU admission to Task.estimated_vram_mb.
The new Task.estimated_vram_mb field is ignored unless every caller also passes required_vram_mb; omitted values silently bypass reservation/admission and run as zero-VRAM tasks.
Proposed API wiring
async def submit_gpu(
- self, task: Task, required_vram_mb: int = 0,
+ self, task: Task, required_vram_mb: int | None = None,
evictable: bool = False, resource_id: str | None = None,
) -> object:
+ if required_vram_mb is None:
+ required_vram_mb = task.estimated_vram_mb
self._submitted += 1📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| 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 = await self._reserve_and_check(task.id, required_vram_mb) | |
| async def submit_gpu( | |
| self, task: Task, required_vram_mb: int | None = None, | |
| evictable: bool = False, resource_id: str | None = None, | |
| ) -> object: | |
| if required_vram_mb is None: | |
| required_vram_mb = task.estimated_vram_mb | |
| self._submitted += 1 | |
| if required_vram_mb > 0: | |
| admission = await self._reserve_and_check(task.id, required_vram_mb) |
🤖 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 172 - 178, The GPU
submission path in submit_gpu currently ignores Task.estimated_vram_mb when
required_vram_mb is omitted, so default admission should use the task’s estimate
instead of silently treating it as zero. Update submit_gpu to derive the
effective VRAM requirement from required_vram_mb when provided, otherwise fall
back to task.estimated_vram_mb, and keep the reservation/admission check in
_reserve_and_check using that resolved value so callers don’t need to pass it
explicitly.
Comment on lines
+199
to
+203
| try: | ||
| 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 cancels.
Cancelling the caller while it awaits done leaves the queue entry behind; _drain_queue() can later reserve VRAM and execute abandoned work. Check the future before admission.
Proposed guard
try:
return await done
except asyncio.CancelledError:
+ if not done.done():
+ done.cancel()
self._evicted += 1
raise drained = False
while not self._queue.empty() and not drained:
entry = self._queue.get_nowait()
+ future = getattr(entry.task, "_arbiter_future", None)
+ if future is not None and future.done():
+ continue
admission = await self._reserve_and_check(entry.task.id, entry.required_vram_mb)Also applies to: 386-389
🤖 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 199 - 203, The
await-on-done path in GPUArbiter is still admitting canceled work, so a canceled
submitter can leave a stale queue entry that `_drain_queue()` later executes.
Update the relevant `try/except asyncio.CancelledError` handling in
`gpu_arbiter.py` around the `done` await so it checks whether the future was
canceled before admission and skips/removes that queued item instead of letting
it proceed; apply the same guard in the other matching cancel path mentioned in
the diff.
Comment on lines
+221
to
+264
| free_vram, _total = self._vram_probe() | ||
| # Subtract in-flight reservations from the probe to close the | ||
| # TOCTOU window between two concurrent submit_gpu calls. | ||
| effective_free = max(0, free_vram - self._reserved_vram_mb) | ||
| if free_vram > 0: | ||
| if effective_free < required_vram_mb: | ||
| return GpuAdmission( | ||
| admitted=False, | ||
| free_vram_mb=effective_free, | ||
| required_vram_mb=required_vram_mb, | ||
| reason=( | ||
| f"insufficient local VRAM: need {required_vram_mb} MiB, " | ||
| f"have {effective_free} MiB available " | ||
| f"({free_vram} MiB free - {self._reserved_vram_mb} MiB reserved)" | ||
| ), | ||
| ) | ||
| return GpuAdmission( | ||
| admitted=True, free_vram_mb=effective_free, | ||
| required_vram_mb=required_vram_mb, | ||
| ) | ||
| if self._cluster_manager is not None: | ||
| leases = self._cluster_manager.get_leases() | ||
| for worker in self._cluster_manager.get_workers(): | ||
| if worker.status != "online": | ||
| continue | ||
| worker_leases = sum( | ||
| l.required_vram_mb for l in leases | ||
| if l.resource_id.startswith(worker.name + ":") and l.required_vram_mb > 0 | ||
| ) | ||
| # Subtract in-flight reservations as well — the cluster-mode | ||
| # path must also account for tasks whose models are still | ||
| # loading (taOS #1705). | ||
| available = worker.free_vram_mb - worker_leases - self._reserved_vram_mb | ||
| if available >= required_vram_mb: | ||
| return GpuAdmission( | ||
| admitted=True, | ||
| free_vram_mb=available, | ||
| required_vram_mb=required_vram_mb, | ||
| ) | ||
| return GpuAdmission( | ||
| admitted=False, required_vram_mb=required_vram_mb, | ||
| reason=f"no cluster worker with {required_vram_mb} MiB free VRAM", | ||
| ) | ||
| return GpuAdmission(admitted=True, required_vram_mb=required_vram_mb) |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
Do not admit required-VRAM tasks when free VRAM is zero or unknown.
Line 225 only handles free_vram > 0; a full local GPU (0, total) or failed probe (0, 0) falls through to admitted=True, bypassing the arbiter exactly when admission control is needed.
Proposed fail-closed admission
free_vram, _total = self._vram_probe()
# Subtract in-flight reservations from the probe to close the
# TOCTOU window between two concurrent submit_gpu calls.
effective_free = max(0, free_vram - self._reserved_vram_mb)
- if free_vram > 0:
+ if _total > 0:
if effective_free < required_vram_mb:
return GpuAdmission(
admitted=False,
free_vram_mb=effective_free,
required_vram_mb=required_vram_mb,
@@
admitted=False, required_vram_mb=required_vram_mb,
reason=f"no cluster worker with {required_vram_mb} MiB free VRAM",
)
- return GpuAdmission(admitted=True, required_vram_mb=required_vram_mb)
+ return GpuAdmission(
+ admitted=False,
+ required_vram_mb=required_vram_mb,
+ reason="unable to determine local or cluster VRAM availability",
+ )📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| free_vram, _total = self._vram_probe() | |
| # Subtract in-flight reservations from the probe to close the | |
| # TOCTOU window between two concurrent submit_gpu calls. | |
| effective_free = max(0, free_vram - self._reserved_vram_mb) | |
| if free_vram > 0: | |
| if effective_free < required_vram_mb: | |
| return GpuAdmission( | |
| admitted=False, | |
| free_vram_mb=effective_free, | |
| required_vram_mb=required_vram_mb, | |
| reason=( | |
| f"insufficient local VRAM: need {required_vram_mb} MiB, " | |
| f"have {effective_free} MiB available " | |
| f"({free_vram} MiB free - {self._reserved_vram_mb} MiB reserved)" | |
| ), | |
| ) | |
| return GpuAdmission( | |
| admitted=True, free_vram_mb=effective_free, | |
| required_vram_mb=required_vram_mb, | |
| ) | |
| if self._cluster_manager is not None: | |
| leases = self._cluster_manager.get_leases() | |
| for worker in self._cluster_manager.get_workers(): | |
| if worker.status != "online": | |
| continue | |
| worker_leases = sum( | |
| l.required_vram_mb for l in leases | |
| if l.resource_id.startswith(worker.name + ":") and l.required_vram_mb > 0 | |
| ) | |
| # Subtract in-flight reservations as well — the cluster-mode | |
| # path must also account for tasks whose models are still | |
| # loading (taOS #1705). | |
| available = worker.free_vram_mb - worker_leases - self._reserved_vram_mb | |
| if available >= required_vram_mb: | |
| return GpuAdmission( | |
| admitted=True, | |
| free_vram_mb=available, | |
| required_vram_mb=required_vram_mb, | |
| ) | |
| return GpuAdmission( | |
| admitted=False, required_vram_mb=required_vram_mb, | |
| reason=f"no cluster worker with {required_vram_mb} MiB free VRAM", | |
| ) | |
| return GpuAdmission(admitted=True, required_vram_mb=required_vram_mb) | |
| free_vram, _total = self._vram_probe() | |
| # Subtract in-flight reservations from the probe to close the | |
| # TOCTOU window between two concurrent submit_gpu calls. | |
| effective_free = max(0, free_vram - self._reserved_vram_mb) | |
| if _total > 0: | |
| if effective_free < required_vram_mb: | |
| return GpuAdmission( | |
| admitted=False, | |
| free_vram_mb=effective_free, | |
| required_vram_mb=required_vram_mb, | |
| reason=( | |
| f"insufficient local VRAM: need {required_vram_mb} MiB, " | |
| f"have {effective_free} MiB available " | |
| f"({free_vram} MiB free - {self._reserved_vram_mb} MiB reserved)" | |
| ), | |
| ) | |
| return GpuAdmission( | |
| admitted=True, free_vram_mb=effective_free, | |
| required_vram_mb=required_vram_mb, | |
| ) | |
| if self._cluster_manager is not None: | |
| leases = self._cluster_manager.get_leases() | |
| for worker in self._cluster_manager.get_workers(): | |
| if worker.status != "online": | |
| continue | |
| worker_leases = sum( | |
| l.required_vram_mb for l in leases | |
| if l.resource_id.startswith(worker.name + ":") and l.required_vram_mb > 0 | |
| ) | |
| # Subtract in-flight reservations as well — the cluster-mode | |
| # path must also account for tasks whose models are still | |
| # loading (taOS `#1705`). | |
| available = worker.free_vram_mb - worker_leases - self._reserved_vram_mb | |
| if available >= required_vram_mb: | |
| return GpuAdmission( | |
| admitted=True, | |
| free_vram_mb=available, | |
| required_vram_mb=required_vram_mb, | |
| ) | |
| return GpuAdmission( | |
| admitted=False, required_vram_mb=required_vram_mb, | |
| reason=f"no cluster worker with {required_vram_mb} MiB free VRAM", | |
| ) | |
| return GpuAdmission( | |
| admitted=False, | |
| required_vram_mb=required_vram_mb, | |
| reason="unable to determine local or cluster VRAM availability", | |
| ) |
🧰 Tools
🪛 Ruff (0.15.20)
[error] 247-247: Ambiguous variable name: l
(E741)
🤖 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 221 - 264, In the GPU
admission logic inside GpuArbiter’s VRAM check, the current `if free_vram > 0`
gate lets `(0, total)` and probe failures like `(0, 0)` fall through to an
admitted result. Change the local-path handling to fail closed whenever VRAM is
zero or the probe is unknown, returning `GpuAdmission(admitted=False, ...)` with
an explicit reason instead of defaulting to admission. Keep the existing
`effective_free` reservation subtraction and cluster-manager fallback behavior
intact, and ensure `GpuAdmission` is only admitted on local GPUs when a
positive, sufficient `free_vram` is confirmed.
Comment on lines
+246
to
+248
| worker_leases = sum( | ||
| l.required_vram_mb for l in leases | ||
| if l.resource_id.startswith(worker.name + ":") and l.required_vram_mb > 0 |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Rename the ambiguous lease variable.
Ruff flags Line 247 as E741; use a descriptive name to keep lint clean.
Proposed rename
worker_leases = sum(
- l.required_vram_mb for l in leases
- if l.resource_id.startswith(worker.name + ":") and l.required_vram_mb > 0
+ lease.required_vram_mb for lease in leases
+ if lease.resource_id.startswith(worker.name + ":") and lease.required_vram_mb > 0
)📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| worker_leases = sum( | |
| l.required_vram_mb for l in leases | |
| if l.resource_id.startswith(worker.name + ":") and l.required_vram_mb > 0 | |
| worker_leases = sum( | |
| lease.required_vram_mb for lease in leases | |
| if lease.resource_id.startswith(worker.name + ":") and lease.required_vram_mb > 0 |
🧰 Tools
🪛 Ruff (0.15.20)
[error] 247-247: Ambiguous variable name: l
(E741)
🤖 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 246 - 248, The lease
variable name in the VRAM summation expression is too ambiguous and triggers
Ruff E741. Update the comprehension inside gpu_arbiter.py’s worker lease
calculation to use a descriptive name instead of a single-letter alias, keeping
the logic in the same place but making the symbol clear and lint-compliant.
Source: Linters/SAST tools
Comment on lines
+410
to
+414
| def _propagate(ct: asyncio.Task, f: asyncio.Future = future) -> None: | ||
| if f.done(): | ||
| return | ||
| if ct.cancelled(): | ||
| return # _evict_task already cancelled the future |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
Propagate unexpected task cancellation to the submitter future.
If the background task is cancelled by anything other than _evict_task, this callback returns without resolving f, leaving submit_gpu() awaiters stuck.
Proposed cancellation propagation
def _propagate(ct: asyncio.Task, f: asyncio.Future = future) -> None:
if f.done():
return
if ct.cancelled():
- return # _evict_task already cancelled the future
+ f.cancel()
+ return📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| def _propagate(ct: asyncio.Task, f: asyncio.Future = future) -> None: | |
| if f.done(): | |
| return | |
| if ct.cancelled(): | |
| return # _evict_task already cancelled the future | |
| def _propagate(ct: asyncio.Task, f: asyncio.Future = future) -> None: | |
| if f.done(): | |
| return | |
| if ct.cancelled(): | |
| f.cancel() | |
| return |
🤖 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 410 - 414, The _propagate
callback in gpu_arbiter.py currently ignores all task cancellations, which can
leave the submitter future unresolved and block submit_gpu() awaiters. Update
_propagate to distinguish cancellations caused by _evict_task from unexpected
cancellations, and in the unexpected case resolve or cancel f so the awaiting
caller is released. Use the existing _propagate helper and the submit_gpu()
future handling paths to keep the cancellation flow consistent.
This was referenced Jul 7, 2026
jaylfc
pushed a commit
that referenced
this pull request
Jul 9, 2026
… fix, #1706) (#1725) * feat(models): atomic VRAM check-and-reserve before model load (TOCTOU fix) Adds VramReservationManager — an asyncio.Lock-protected VRAM reservation system that atomically checks free VRAM (via nvidia-smi) and reserves the required amount before a model load begins. Concurrent callers see the in-flight reservation and back off, preventing the TOCTOU race where two model pulls both see enough free VRAM and both try to load, causing OOM. Integration points: - POST /api/models/pull — reserve before ollama pull, release after - POST /api/models/download (rkllama path) — reserve before installer, release in the background task's finally block - GET /api/models/vram-reservations — stats endpoint for monitoring New files: - tinyagentos/vram_reservation.py — atomic reserve/release manager - tests/test_vram_reservation.py — 15 tests including concurrent safety Wired as app.state.vram_reservation in the FastAPI lifespan alongside the cluster manager and resource scheduler. Issue: #1706 * docs: add TODO for per-model VRAM estimate replacement jaylfc review: min_ram_mb heuristic is safe for v1 (over-reserves, fails closed), but should be replaced with a real per-model VRAM estimate on CUDA GGUF to avoid false 503s in multi-model setups. --------- Co-authored-by: Hogne <227774406+hognek@users.noreply.github.com>
Owner
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What
Fixes the queue processor serialization blocking drain (PR #1683 review feedback #4).
Changes
Non-blocking drain:
_drain_queueno longer awaits_run_gpu_task. The admitted task is spawned as a backgroundasyncio.Taskwith a done-callback that propagates the result/exception to the submitter's_arbiter_future. The drain loop stays responsive on subsequent ticks.Eviction-to-make-room: When admission fails for a queued task,
_drain_queuenow callsevict_lowest_priority(min_priority=int(task.priority))and re-checks admission. Lower-priority running tasks are evicted to make room; higher-priority runners are preserved.Documentation: Added a docstring explaining the intentional serialization — only one task is admitted per drain cycle to avoid flooding the GPU with concurrent loads.
Tests
test_gpu_arbiter_894.py:TestDrainQueueNonBlocking: drain returns immediately, done_callback propagates results and exceptionsTestEvictionToMakeRoom: eviction triggered when VRAM full, no eviction of higher-priority runners, eviction disabled guardNot changed
_evict_task/evict_lowest_priority/_run_gpu_taskare unchanged_process_queueSummary by CodeRabbit
New Features
Bug Fixes