feat(scheduler): land GPU arbiter on one VRAM authority + eviction/heartbeat fixes (#894/#185, continues #1718)#1859
Merged
Merged
Conversation
…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, #893 leases). Task: t_d7208884. Fixes #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 #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 #894. Task: t_0c86d08e.
…tible with stack The file was incorrectly brought in from the #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 #1683 review feedback #2 Tests: 168 pass (25 arbiter + 143 cluster/worker)
…ed VRAM in cluster path taOS #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 #1683 review feedback #4.
taOS #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 #1705 fix) and sync _evict_task calls in new tests. Tests: 22/22 arbiter pass.
…biter Make evict_lowest_priority, _evict_task, release_tasks_for_worker, stats, and running_tasks async and guard _running dict access with _running_lock. Previously _running was mutated under the lock in _run_gpu_task but read/popped without it in eviction paths — safe under single-thread-no-await but fragile. Lease release and task cancellation happen outside the lock to avoid deadlock. Update all callers: drain_worker → async, route handlers, and tests. 69 targeted tests pass.
- Added vram_probe=_probe_nvidia_vram so the arbiter uses actual GPU VRAM instead of the default (0,0) probe that makes all checks pass. - Wired cluster_manager._gpu_arbiter so eviction paths that access the cluster manager can reach the arbiter. - Added gpu_arbiter.stop() call in shutdown sequence. - Added _gpu_arbiter attribute to ClusterManager.__init__.
taOS #1707 — _running_lock coverage fix: - _evict_task now holds _running_lock during the pop, preventing races with concurrent drain operations. - stats() made async (needed for _running_lock consistency). - All async call sites now properly awaited (evict_lowest_priority, _evict_task, stats, claim_lease, release_lease). Tests: 22/22 arbiter pass.
…, restart Adds drain_worker/cancel_drain to ClusterManager for graceful worker detach without dropping inflight GPU tasks (taOS #890). ClusterManager: - drain_worker(name, graceful=True): worker enters 'draining' status, excluded from routing/catalog/lease claims. When graceful=False, all leases are force-released and worker is marked offline. - cancel_drain(name): returns draining worker to 'online'. - _monitor_loop: auto-completes drain when all leases released; force-finishes stale drains on heartbeat timeout. - _worker_for_resource, get_workers_for_capability, aggregate_catalog: all exclude draining workers. Routes: - POST /api/cluster/workers/{name}/drain — begin graceful/force drain - POST /api/cluster/workers/{name}/cancel-drain — cancel in-progress drain - POST /api/cluster/workers/{name}/update — full auto-update orchestration: drain → deploy update-worker → restart → re-register Tests: 9 new tests (25 total in test_cluster.py), 191 cluster tests pass.
Replace startswith(name + ':') with _parse_resource_id exact match to prevent draining worker 'foo' from releasing leases belonging to worker 'foo-bar' (same prefix-collision class as TOCTOU fixes).
taOS #796 features extracted from #1689 and rebased onto the arbiter stack: 1. pause/resume queue control: GpuArbiter.pause()/resume() halt/restart queue processing. Running tasks finish; queued tasks wait. Paused queue skips _drain_queue in _process_queue loop. 2. Hardware-aware LLM admission: submit_gpu() accepts required_gpu_arch (e.g. 'sm_86'). _check_gpu_arch_compatibility() checks cluster workers for compatible GPU hardware via model name and compute_cap field before admitting tasks.
…r_894 _evict_task was made async in commit 1789ad0 (running-lock fix). Three cleanup calls were missing await.
…on force-release, wire arbiter task cancellation #1690 lock fix: drain_worker force-release and _monitor_loop stale-drain path now hold _lease_lock during self._leases mutation, serialized with claim/release/sweep. drain_worker and cancel_drain are now async (consistent with claim_lease/release_lease — all callers were already in async handlers). Cross-cutting wiring: added GpuArbiter.cancel_running_for_leases() and wired it into drain_worker(graceful=False) and _monitor_loop stale-drain path. When leases are force-released, any running GPU arbiter tasks for those leases are cancelled via _evict_task — eviction is no longer a prod no-op when the arbiter is running. Added TYPE_CHECKING import for GpuArbiter type annotation on ClusterManager._gpu_arbiter.
… (cancelled, already_completed) from cancel_running_for_leases Addresses Kilo review findings on PR #1718: CRITICAL: Race in stale-drain path — worker.status='offline' was set after releasing _lease_lock, allowing duplicate worker.leave on next monitor tick. Now set inside the lock (both drain_worker force-release and _monitor_loop stale-drain paths). WARNING: Same pattern in drain_worker(graceful=False) — fixed. SUGGESTION: cancel_running_for_leases now returns (cancelled, already_completed) so operators can distinguish force-kills from natural completions. Callers log both counts.
Maintainer takeover of the arbiter epic (continues hognek's #1718). Merges current dev and closes the #185 gate so there is ONE VRAM ledger: - The arbiter no longer keeps its own reservation accounting. It reserves against the shared VramReservationManager (#1725) that the model-load path in routes/models.py already uses, so GPU scheduler tasks and model loads never double-count. app.py builds one manager and passes it into the arbiter; reserve_vram/release_vram is the arbiter front door for non-task callers. Folding the manager in also removes the arbiter's own probe: the manager probes via a worker thread (no event-loop stall) and distinguishes no-probe (fail-open, cluster-parity) from known-full. - Eviction ordering (Xid-62): _evict_task now cancels AND awaits the task before releasing its reservation + lease, so a re-admission cannot pass against not-yet-reclaimed VRAM. Regression test added. - evict_lowest_priority no longer preempts an equal-priority running task (strictly-lower only) to avoid eviction thrash. - heartbeat() no longer flips a draining worker back to online (#890), which had defeated graceful drain. Regression test added. - routes/cluster.py /update raises_for_status before reading the worker response so a failed deploy is not reported as status:"updating". VramReservationManager gains an injectable probe for tests. TOCTOU tests updated to assert against the shared ledger. Model-loads-as-evictable- submit_gpu-tasks (fuller queue/eviction semantics) is deferred to a follow-up. Docs-Reviewed: internal GPU-scheduler/cluster admission plumbing; no agent-facing API surface or route module added or removed.
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? |
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Pro Plus Run ID: 📒 Files selected for processing (1)
🚧 Files skipped from review as they are similar to previous changes (1)
📝 WalkthroughWalkthroughAdds a VRAM-aware GPU arbiter with atomic reservation accounting, queued admission, priority eviction, and lease coordination. Integrates GPU discovery and application lifecycle wiring, and adds graceful worker draining, routing exclusion, lease cleanup, and admin update orchestration. ChangesGPU arbitration and scheduler integration
Worker draining and updates
Estimated code review effort: 5 (Critical) | ~120 minutes Sequence Diagram(s)sequenceDiagram
participant Admin
participant ClusterRoutes
participant ClusterManager
participant Worker
participant GpuArbiter
Admin->>ClusterRoutes: POST drain or update request
ClusterRoutes->>ClusterManager: drain_worker(name, graceful)
ClusterManager->>Worker: set draining or offline status
Worker-->>ClusterManager: heartbeat preserves draining status
ClusterManager->>GpuArbiter: cancel tasks for released leases
ClusterRoutes->>Worker: POST update-worker
Worker-->>ClusterRoutes: deployment result
ClusterRoutes->>ClusterManager: cancel_drain on deployment failure
Possibly related PRs
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches📝 Generate docstrings
🧪 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 |
|
|
![kilo-code-bot[bot]](https://avatars.githubusercontent.com/in/2193792?s=60&v=4){kind=small}
| gpu_info = worker.hardware.get("gpu", {}) if isinstance(worker.hardware, dict) else {} | ||
| gpu_model = gpu_info.get("model", "") or "" | ||
| cc = gpu_info.get("compute_cap", "") or "" | ||
| if required_gpu_arch in gpu_model or required_gpu_arch in cc: |
There was a problem hiding this comment.
WARNING: required_gpu_arch in gpu_model (and in cc) is a substring match that accepts the wrong architecture.
gpu_model/compute_cap are strings like "NVIDIA A100" or "8.0,8.6". A caller asking for sm_8 would match sm_80, sm_86, and sm_89 all at once, while sm_86 as a substring would also spuriously match a model containing sm_860. More importantly, this check can never actually succeed for the documented input: a compute capability like sm_86 is not a substring of a real compute_cap value such as 8.6 or a model string, so the feature as written will reject hardware that genuinely satisfies the requirement. Use exact/normalized comparison (parse compute_cap to a float/set, or require the caller to pass a comparable token) rather than in.
| gpu_model = gpu_info.get("model", "") or "" | ||
| cc = gpu_info.get("compute_cap", "") or "" | ||
| if required_gpu_arch in gpu_model or required_gpu_arch in cc: | ||
| if resource_id is None or resource_id.startswith(worker.name + ":"): |
There was a problem hiding this comment.
WARNING: resource_id.startswith(worker.name + ":") reintroduces the worker-name prefix-collision bug this PR explicitly fixed elsewhere (commits 13/16 switched startswith(name + ":") to an exact _parse_resource_id compare to stop foo from matching foo-bar). A resource like foo-bar:gpu-cuda-0 would be matched against worker foo. Use _parse_resource_id(resource_id)[0] == worker.name for an exact-name comparison, consistent with the drain paths in manager.py.
| victim_ids.append(task_id) | ||
| cancelled = 0 | ||
| already_completed = 0 | ||
| for task_id in victim_ids: |
There was a problem hiding this comment.
SUGGESTION: cancel_running_for_leases conflates two distinct cases in already_completed: _evict_task returns 0 both when the task is not in _running (already finished or was never tracked, e.g. a lease whose task was not a GPU-arbiter task) and when eviction did nothing. So already_completed can over-count and the operator-facing log is misleading. Consider distinguishing "task found and evicted" vs "task not found in _running".
Suggested change
| for task_id in victim_ids: | |
| cancelled = 0 | |
| not_found = 0 | |
| for task_id in victim_ids: | |
| if await self._evict_task(task_id): | |
| cancelled += 1 | |
| else: | |
| not_found += 1 | |
| if cancelled or not_found: | |
| logger.info("gpu-arbiter: drain cancelled %d, not-in-running %d (leases=%d)", | |
| cancelled, not_found, len(lease_ids)) | |
| return cancelled, not_found |
Code Review SummaryStatus: 1 New Issue Found | Recommendation: Address suggestion before merge Overview
Issue Details (click to expand)SUGGESTION
Previously reported, still open (unchanged-code, carried forward)
Resolved since previous review (commit 4839501)
Files Reviewed (incremental)
Fix these issues in Kilo Cloud Previous Review Summary (commit 6bf501a)Current summary above is authoritative. Previous snapshots are kept for context only. Previous review (commit 6bf501a)Status: 3 Issues Found | Recommendation: Address before merge Overview
Issue Details (click to expand)WARNING
SUGGESTION
Files Reviewed (11 files)
Reviewed by hy3:free · Input: 83.9K · Output: 7.3K · Cached: 886.4K |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 16
🧹 Nitpick comments (2)
tests/test_cluster.py (1)
75-90: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winCover heartbeats after a drain reaches
"offline".Add force-drain and graceful-completion cases that heartbeat afterward and assert the worker remains excluded. The current test only covers the intermediate
"draining"state.🤖 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_cluster.py` around lines 75 - 90, Add coverage alongside test_heartbeat_does_not_reonline_draining_worker for workers whose drain completes as "offline": create force-drain and graceful-completion cases, heartbeat each worker afterward, and assert heartbeat succeeds without changing the offline status or re-enabling routing. Reuse the existing ClusterManager and worker-drain APIs and preserve the current intermediate "draining" test.tests/test_gpu_arbiter_894.py (1)
360-414: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd an equal-priority non-eviction regression test.
The PR’s strict-priority boundary is important, but the suite only covers higher-versus-lower priority tasks.
🤖 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 360 - 414, Extend TestEvictLowestPriority with an async regression test that starts two running GPU tasks having equal Priority values, invokes evict_lowest_priority, and verifies that no task is evicted and both remain in arbiter._running. Reuse the existing blocking-task setup and ensure both tasks are cleaned up afterward.
🤖 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 541-575: Update the _QueuedGpuTask entry in
test_eviction_triggered_when_vram_full so its evictable value matches the test’s
expected eviction behavior; set evictable to True and preserve the existing
eviction assertions.
In `@tests/test_gpu_arbiter_toctou.py`:
- Around line 29-58: Extend test_concurrent_reserve_and_check_only_one_admitted
to exercise the no-local-probe ClusterManager admission path, configuring a
worker with 8192 MiB capacity and two concurrent 5000 MiB requests. Assert
exactly one request is admitted and the cluster worker’s reserved capacity
reflects only that single reservation.
In `@tinyagentos/app.py`:
- Around line 1183-1196: Update the GPU arbiter startup flow around GpuArbiter
and await gpu_arbiter.start() so it never starts when resource_scheduler is
unavailable. Skip initialization and retain the existing fallback behavior, or
provide the actual vanilla scheduler instance instead of passing None; ensure
_run_gpu_task always receives a valid scheduler.
In `@tinyagentos/cluster/manager.py`:
- Around line 212-218: Update the worker heartbeat logic around the status
assignment to preserve administrative drain intent independently of transient
status, keeping workers offline after both forced drains and completed graceful
drains rather than re-onlining them. In tinyagentos/cluster/manager.py lines
212-218, use the existing drain state or an equivalent persistent marker when
deciding whether to set status to "online"; add heartbeat-after-force-drain and
heartbeat-after-graceful-completion regression coverage in tests/test_cluster.py
lines 75-90.
In `@tinyagentos/routes/cluster.py`:
- Around line 1218-1229: Update the worker update flow around
cluster.drain_worker to wait for graceful drain completion, using an appropriate
timeout, before creating the deploy request to /api/worker/deploy. Preserve the
existing error response handling and only trigger deployment after active leases
have finished draining.
- Around line 1212-1219: Update the worker deployment rollback flow around the
status check and cluster.drain_worker call to record whether the worker was
already draining before this request. On deployment failure, call cancel_drain()
only when this request initiated the drain, and derive drain_cancelled from the
actual cancellation result rather than setting it unconditionally; preserve
pre-existing drains.
In `@tinyagentos/scheduler/discovery.py`:
- Line 18: Remove the direct _probe_nvidia_vram dependency from scheduler
discovery and obtain capacity from an injected shared-manager snapshot instead.
Ensure scheduler admission uses the arbiter ledger’s reservation-aware values
across NVIDIA, ROCm, and Metal, and represents unknown capacity explicitly
rather than converting probe failures into fictitious free VRAM.
- Around line 190-202: Update the GPU discovery logic around gpu_backend_types
and the gpu_backends comprehension to use an explicit backend acceleration
signal from each catalog backend, rather than treating backend type alone as GPU
evidence. Ensure CPU-only Ollama is excluded while CUDA-enabled llama-cpp is
included, and reuse this acceleration-based result for GPU tier registration and
capability lookup.
In `@tinyagentos/scheduler/gpu_arbiter.py`:
- Around line 252-254: Update the reservation lifecycle in the task execution
flow around _vram.reserve and the corresponding cleanup block so capacity holds
remain valid for the entire task lifetime. Use a renewable hold or a
non-expiring runtime reservation, and ensure the reservation is released in
finally for both normal completion and failure paths.
- Around line 338-342: Update the queue-draining logic in _drain_queue to detect
and remove entries whose awaited future was cancelled before admitting or
starting GPU work. Apply the same guard to the corresponding queue-entry
handling near the other referenced block, while preserving the existing _evicted
increment and CancelledError propagation in the submitter path.
- Around line 48-57: Preserve the task’s requested resource_id throughout
queueing and dispatch: add it to _QueuedGpuTask, populate it when creating
queued entries, and pass it when the queued task is later executed instead of
using None. Update the related queue handling paths noted in the comment while
keeping existing lease behavior unchanged.
- Around line 262-264: The cluster admission path in
GpuArbiter._check_cluster_admission must atomically reserve capacity or claim
the worker lease instead of returning a snapshot-based decision; update
tinyagentos/scheduler/gpu_arbiter.py lines 262-264 accordingly. Add concurrent
no-local-probe coverage in tests/test_gpu_arbiter_toctou.py lines 29-58
verifying that only one request acquires the worker’s available VRAM.
- Around line 548-555: The queue’s eviction-to-make-room path ignores each
entry’s eviction policy. In tinyagentos/scheduler/gpu_arbiter.py lines 548-555,
gate the evict_lowest_priority call and subsequent admission retry on
entry.evictable. In tests/test_gpu_arbiter_894.py lines 541-575, set
evictable=True for the eviction scenario and add an assertion covering
evictable=False to verify eviction is skipped.
- Around line 246-252: Update the GPU admission flow around _vram.available_vram
and reserve so ledger access remains on the event-loop thread. Add an async
manager operation that acquires the manager lock, runs only the blocking
hardware probe in a worker thread, and performs the _pending sweep and
reservation atomically under that lock; replace the direct
asyncio.to_thread(self._vram.available_vram) call with this operation while
preserving the no-local-GPU behavior.
- Around line 205-207: Update the worker eligibility check in the scheduler
admission logic to skip workers whose status is "draining", alongside other
non-online workers. Ensure only workers with status "online" can satisfy
architecture admission, while preserving the existing iteration and filtering
behavior.
In `@tinyagentos/vram_reservation.py`:
- Line 77: Import Callable from typing so the probe annotation in the relevant
reservation API resolves correctly for Ruff and typing.get_type_hints().
---
Nitpick comments:
In `@tests/test_cluster.py`:
- Around line 75-90: Add coverage alongside
test_heartbeat_does_not_reonline_draining_worker for workers whose drain
completes as "offline": create force-drain and graceful-completion cases,
heartbeat each worker afterward, and assert heartbeat succeeds without changing
the offline status or re-enabling routing. Reuse the existing ClusterManager and
worker-drain APIs and preserve the current intermediate "draining" test.
In `@tests/test_gpu_arbiter_894.py`:
- Around line 360-414: Extend TestEvictLowestPriority with an async regression
test that starts two running GPU tasks having equal Priority values, invokes
evict_lowest_priority, and verifies that no task is evicted and both remain in
arbiter._running. Reuse the existing blocking-task setup and ensure both tasks
are cleaned up afterward.
🪄 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: dc4c1935-808c-40ec-98d4-ffe7613d7fee
📒 Files selected for processing (11)
tests/test_cluster.pytests/test_gpu_arbiter_894.pytests/test_gpu_arbiter_toctou.pytinyagentos/app.pytinyagentos/cluster/manager.pytinyagentos/routes/cluster.pytinyagentos/scheduler/__init__.pytinyagentos/scheduler/discovery.pytinyagentos/scheduler/gpu_arbiter.pytinyagentos/scheduler/types.pytinyagentos/vram_reservation.py
Comment on lines
+541
to
+575
| async def test_eviction_triggered_when_vram_full(self): | ||
| """If VRAM is full, drain should evict a lower-priority task to make room.""" | ||
| arbiter = GpuArbiter( | ||
| vram_probe=lambda: ( | ||
| (1024, 8192) if len(arbiter_ctx["ref"]._running) > 0 else (8192, 8192) | ||
| ), | ||
| max_queue_size=10, | ||
| ) | ||
| arbiter_ctx: dict[str, object] = {"ref": arbiter} | ||
|
|
||
| # Put a low-priority task in _running (simulating an already-running task) | ||
| low_task = _make_task("low-running", vram_mb=4096, priority=Priority.BATCH) | ||
| low_task.payload = lambda r: asyncio.sleep(60) | ||
| arbiter._running["low-running"] = (low_task, None, int(Priority.BATCH), 4096) | ||
|
|
||
| # Now queue a higher-priority task that needs VRAM | ||
| hi_task = _make_task("hi-queued", vram_mb=4096, priority=Priority.INTERACTIVE_USER) | ||
| hi_started = asyncio.Event() | ||
| hi_release = asyncio.Event() | ||
|
|
||
| async def hi_payload(_resource): | ||
| hi_started.set() | ||
| await hi_release.wait() | ||
| return {"ok": True} | ||
|
|
||
| hi_task.payload = hi_payload | ||
|
|
||
| loop = asyncio.get_running_loop() | ||
| arb_future: asyncio.Future = loop.create_future() | ||
| hi_task._arbiter_future = arb_future # type: ignore[attr-defined] | ||
|
|
||
| entry = _QueuedGpuTask( | ||
| priority=int(hi_task.priority), seq=1, task=hi_task, | ||
| required_vram_mb=4096, evictable=False, | ||
| ) |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
Align this test with the evictable contract.
It sets evictable=False but expects the queued task to trigger eviction. Set it to True, or change the assertion to require no eviction.
🤖 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 541 - 575, Update the
_QueuedGpuTask entry in test_eviction_triggered_when_vram_full so its evictable
value matches the test’s expected eviction behavior; set evictable to True and
preserve the existing eviction assertions.
Comment on lines
+29
to
+58
| async def test_concurrent_reserve_and_check_only_one_admitted(self): | ||
| """Two concurrent _reserve_and_check calls — only first admitted. | ||
|
|
||
| With 8192 MiB free and each needing 5000 MiB, the second | ||
| _reserve_and_check must see the first's reservation in the shared | ||
| ledger and be rejected. The manager's atomic reserve closes the | ||
| TOCTOU gap. | ||
| """ | ||
| arbiter = GpuArbiter( | ||
| vram_probe=lambda: (8192, 8192), | ||
| max_queue_size=10, | ||
| ) | ||
|
|
||
| results: list[bool] = [] | ||
|
|
||
| async def try_reserve(task_id): | ||
| admission = await arbiter._reserve_and_check(task_id, 5000) | ||
| results.append(admission.admitted) | ||
|
|
||
| # Launch both concurrently — this is the TOCTOU window. | ||
| await asyncio.gather( | ||
| try_reserve("t-a"), | ||
| try_reserve("t-b"), | ||
| ) | ||
|
|
||
| admitted_count = sum(results) | ||
| assert admitted_count == 1, ( | ||
| f"Expected exactly 1 admission, got {admitted_count}: {results}" | ||
| ) | ||
| assert arbiter._vram.reserved_vram_mb == 5000 |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟠 Major | ⚡ Quick win
Exercise the cluster admission branch concurrently.
This test only covers the local manager’s atomic reserve path. Add a no-local-probe ClusterManager case proving two concurrent admissions cannot both claim the same worker capacity.
🤖 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_toctou.py` around lines 29 - 58, Extend
test_concurrent_reserve_and_check_only_one_admitted to exercise the
no-local-probe ClusterManager admission path, configuring a worker with 8192 MiB
capacity and two concurrent 5000 MiB requests. Assert exactly one request is
admitted and the cluster worker’s reserved capacity reflects only that single
reservation.
Comment on lines
+1183
to
+1196
| gpu_arbiter = GpuArbiter( | ||
| scheduler=resource_scheduler if resource_scheduler is not None else None, | ||
| cluster_manager=cluster_manager, | ||
| vram_reservation=vram_reservation, | ||
| max_queue_size=100, | ||
| eviction_enabled=True, | ||
| ) | ||
| await gpu_arbiter.start() | ||
| app.state.gpu_arbiter = gpu_arbiter | ||
| cluster_manager._gpu_arbiter = gpu_arbiter | ||
| logger.info("GPU arbiter ready (queue size=100, eviction=enabled, shared VRAM ledger)") | ||
| 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.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Do not start the production arbiter without a resource scheduler.
When scheduler construction fails, this passes None; _run_gpu_task then calls task.payload(None) rather than using the claimed “vanilla scheduler” fallback. Resource-dependent payloads can fail or bypass routing.
Skip arbiter startup in this case or inject the actual fallback scheduler.
🤖 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 1183 - 1196, Update the GPU arbiter startup
flow around GpuArbiter and await gpu_arbiter.start() so it never starts when
resource_scheduler is unavailable. Skip initialization and retain the existing
fallback behavior, or provide the actual vanilla scheduler instance instead of
passing None; ensure _run_gpu_task always receives a valid scheduler.
Comment on lines
+212
to
+218
| # A worker mid graceful-drain keeps heartbeating while it finishes | ||
| # inflight work. Do NOT flip it back to "online" — that would re-enter | ||
| # it into routing/catalog/lease-claims before its leases drain and | ||
| # defeat the drain (taOS #890). Leave the drain to complete or be | ||
| # cancelled explicitly; every other status re-onlines as before. | ||
| if worker.status != "draining": | ||
| worker.status = "online" |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Keep administratively drained workers offline across heartbeats. Completed and forced drains currently lose their intent once status becomes "offline".
tinyagentos/cluster/manager.py#L212-L218: preserve drain intent independently of the transient worker status.tests/test_cluster.py#L75-L90: add heartbeat-after-force-drain and heartbeat-after-graceful-completion regressions.
📍 Affects 2 files
tinyagentos/cluster/manager.py#L212-L218(this comment)tests/test_cluster.py#L75-L90
🤖 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/cluster/manager.py` around lines 212 - 218, Update the worker
heartbeat logic around the status assignment to preserve administrative drain
intent independently of transient status, keeping workers offline after both
forced drains and completed graceful drains rather than re-onlining them. In
tinyagentos/cluster/manager.py lines 212-218, use the existing drain state or an
equivalent persistent marker when deciding whether to set status to "online";
add heartbeat-after-force-drain and heartbeat-after-graceful-completion
regression coverage in tests/test_cluster.py lines 75-90.
Comment on lines
+1212
to
+1219
| if worker.status not in ("online", "draining"): | ||
| return JSONResponse( | ||
| {"error": f"Worker '{name}' is not online (status={worker.status})"}, | ||
| status_code=400, | ||
| ) | ||
|
|
||
| # Step 1: Begin draining | ||
| drain_result = await cluster.drain_worker(name, graceful=True) |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Do not undo a pre-existing drain on deployment failure.
The endpoint accepts an already-draining worker, then unconditionally calls cancel_drain(), re-admitting it after failure. Only roll back drains initiated by this request, and derive drain_cancelled from the actual rollback result.
Also applies to: 1236-1241
🤖 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/routes/cluster.py` around lines 1212 - 1219, Update the worker
deployment rollback flow around the status check and cluster.drain_worker call
to record whether the worker was already draining before this request. On
deployment failure, call cancel_drain() only when this request initiated the
drain, and derive drain_cancelled from the actual cancellation result rather
than setting it unconditionally; preserve pre-existing drains.
Comment on lines
+252
to
+254
| reservation = await self._vram.reserve(required_vram_mb, caller=f"gpu-task:{task_id}") | ||
| if reservation is not None: | ||
| self._pending_reservations[task_id] = reservation.reservation_id |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | 🏗️ Heavy lift
Renew capacity holds for the full task lifetime.
VRAM reservations expire through the manager TTL, while remote leases expire after 300 seconds. Neither is renewed, so a longer task can remain GPU-resident after its accounting has been reclaimed, allowing over-admission.
Use renewable holds or non-expiring runtime reservations that are released in finally.
Also applies to: 394-397
🤖 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 252 - 254, Update the
reservation lifecycle in the task execution flow around _vram.reserve and the
corresponding cleanup block so capacity holds remain valid for the entire task
lifetime. Use a renewable hold or a non-expiring runtime reservation, and ensure
the reservation is released in finally for both normal completion and failure
paths.
Comment on lines
+262
to
+264
| # No local GPU probe: try the cluster. | ||
| if self._cluster_manager is not None: | ||
| return self._check_cluster_admission(required_vram_mb) |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟠 Major | 🏗️ Heavy lift
Cluster admission is not atomic with capacity ownership.
tinyagentos/scheduler/gpu_arbiter.py#L262-L264: reserve target capacity or claim its lease within the admission operation instead of returning a snapshot result.tests/test_gpu_arbiter_toctou.py#L29-L58: add concurrent no-local-probe coverage proving only one request acquires the worker’s available VRAM.
📍 Affects 2 files
tinyagentos/scheduler/gpu_arbiter.py#L262-L264(this comment)tests/test_gpu_arbiter_toctou.py#L29-L58
🤖 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 262 - 264, The cluster
admission path in GpuArbiter._check_cluster_admission must atomically reserve
capacity or claim the worker lease instead of returning a snapshot-based
decision; update tinyagentos/scheduler/gpu_arbiter.py lines 262-264 accordingly.
Add concurrent no-local-probe coverage in tests/test_gpu_arbiter_toctou.py lines
29-58 verifying that only one request acquires the worker’s available VRAM.
Comment on lines
+338
to
+342
| try: | ||
| return await done | ||
| except asyncio.CancelledError: | ||
| self._evicted += 1 | ||
| raise |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
Discard cancelled queue entries before starting GPU work.
Cancelling the submitter cancels its awaited future but leaves the entry queued. _drain_queue later admits and starts that task even though nobody can receive its result.
Proposed guard
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: 545-557
🤖 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 338 - 342, Update the
queue-draining logic in _drain_queue to detect and remove entries whose awaited
future was cancelled before admitting or starting GPU work. Apply the same guard
to the corresponding queue-entry handling near the other referenced block, while
preserving the existing _evicted increment and CancelledError propagation in the
submitter path.
Comment on lines
+548
to
+555
| # Try eviction-to-make-room for higher-priority queued tasks. | ||
| # evict_lowest_priority(min_priority=N) skips running tasks | ||
| # 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 = await self.evict_lowest_priority(min_priority=int(entry.task.priority)) | ||
| if evicted > 0: | ||
| admission = await self._reserve_and_check(entry.task.id, entry.required_vram_mb) |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
The queue ignores the caller’s eviction policy.
tinyagentos/scheduler/gpu_arbiter.py#L548-L555: gate eviction-to-make-room onentry.evictable.tests/test_gpu_arbiter_894.py#L541-L575: useevictable=Truewhen expecting eviction, and add a false-case assertion.
📍 Affects 2 files
tinyagentos/scheduler/gpu_arbiter.py#L548-L555(this comment)tests/test_gpu_arbiter_894.py#L541-L575
🤖 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 548 - 555, The queue’s
eviction-to-make-room path ignores each entry’s eviction policy. In
tinyagentos/scheduler/gpu_arbiter.py lines 548-555, gate the
evict_lowest_priority call and subsequent admission retry on entry.evictable. In
tests/test_gpu_arbiter_894.py lines 541-575, set evictable=True for the eviction
scenario and add an assertion covering evictable=False to verify eviction is
skipped.
| def __init__( | ||
| self, | ||
| ttl_seconds: float = DEFAULT_RESERVATION_TTL_SECONDS, | ||
| probe: "Callable[[], tuple[int, int] | None] | None" = None, |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
file="tinyagentos/vram_reservation.py"
echo "== outline =="
ast-grep outline "$file" --view expanded || true
echo
echo "== top of file =="
sed -n '1,120p' "$file" | cat -n
echo
echo "== typing imports/usages =="
rg -n "from typing|Callable|get_type_hints|__future__|annotations" "$file" || trueRepository: jaylfc/taOS
Length of output: 6351
🏁 Script executed:
#!/bin/bash
set -euo pipefail
sed -n '1,140p' tinyagentos/vram_reservation.py | cat -nRepository: jaylfc/taOS
Length of output: 6623
Import Callable for the probe annotation. Ruff will flag the unresolved name, and typing.get_type_hints() will fail to resolve it.
🧰 Tools
🪛 Ruff (0.15.21)
[error] 77-77: Undefined name Callable
(F821)
🤖 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/vram_reservation.py` at line 77, Import Callable from typing so
the probe annotation in the relevant reservation API resolves correctly for Ruff
and typing.get_type_hints().
Source: Linters/SAST tools
- _check_gpu_arch_compatibility matches compute_cap exactly so sm_8 cannot prefix-match sm_86 (keeps substring match on the freeform model string). - Add _resource_on_worker using the cluster manager's exact _parse_resource_id so gpu-node does not collide with gpu-node-2 (same guard #1726 applied to the lease paths); used by the arch check and cluster admission.
Owner
Author
|
Addressed the Kilo review in 4839501:
|
| #1726 fixed on the lease paths); falls back to a colon-delimited match. | ||
| """ | ||
| if self._cluster_manager is not None: | ||
| parsed = self._cluster_manager._parse_resource_id(resource_id) |
There was a problem hiding this comment.
SUGGESTION: Coupling to a private cluster-manager API
_resource_on_worker calls self._cluster_manager._parse_resource_id(...) — a name-mangled-private method on ClusterManager. This works today, but reaching into another component's private API means a refactor of _parse_resource_id (or its return shape) will silently break arch/resource matching here with no static signal.
Consider promoting parse_resource_id to a public method on ClusterManager (it is already used internally by _worker_for_resource), or exposing a dedicated resource_on_worker(resource_id, worker_name) helper on the manager so the arbiter does not depend on its internals.
Reply with @kilocode-bot fix it to have Kilo Code address this issue.
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.
Maintainer takeover of the GPU arbiter epic (#894), continuing @hognek's work on #1718 (his commits are preserved here). Rebased onto current dev and closes the #185 gate.
The #185 gate-closer: one VRAM authority
The arbiter shipped inert beside #1725's now-merged
VramReservationManager, so there were two reservation ledgers that would double-count oncesubmit_gpugot a caller. This makes the manager the single ledger:_reserved_vram_mbaccounting. It reserves against the sharedVramReservationManagerthat the model-load path (routes/models.py) already uses.app.pybuilds one manager and passes it into the arbiter, so GPU scheduler tasks and model loads draw from one authority and never double-count.GpuArbiter.reserve_vram()/release_vram()is the arbiter front door for non-task callers.claim_leaseparity) from known-full, which is the behaviour feat(models): atomic VRAM check-and-reserve before model load (TOCTOU fix, #1706) #1725 already reviewed and shipped.Model-loads-as-evictable-
submit_gpu-tasks (full queue/eviction semantics for pulls, so a pull that can't fit waits or evicts instead of 503) is intentionally a follow-up, not this PR: eviction of a user's running model is a product/policy decision, and this PR keeps model-pull behaviour identical to today. The one-ledger foundation here is the prerequisite for that upgrade.Correctness fixes
_evict_tasknow cancels AND awaits the evicted task before releasing its reservation + lease, so an evict-to-make-room re-admission cannot pass against not-yet-reclaimed VRAM and start a second concurrent load. Regression test added.evict_lowest_priorityno longer preempts an equal-priority running task (strictly-lower priority only), avoiding eviction thrash between peers.heartbeat()no longer flips adrainingworker back toonline; that had re-entered it into routing before its leases finished and defeated graceful drain. Regression test added./updatedeploy result.routes/cluster.pynowraise_for_status()on the worker deploy response before reading it, so a failed deploy is not reported asstatus: "updating".Tests
VramReservationManagergains an injectable probe for tests. The TOCTOU tests now assert against the shared ledger. Full local run green: arbiter (894 + toctou), cluster, leases, routes/cluster, routes/models, vram_reservation — 167 passed, plus the two new regression tests.Supersedes #1718 (folds in #1690's worker-auto-update + lock fix via the same branch ancestry). #1725 and #1726 are already merged.
Summary by CodeRabbit