feat(scheduler): VramReservationManager — atomic VRAM bookkeeping for GpuArbiter [#1683]#1759
feat(scheduler): VramReservationManager — atomic VRAM bookkeeping for GpuArbiter [#1683]#1759hognek wants to merge 8 commits into
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, 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.
…t a prod no-op Initialize _gpu_arbiter = None in ClusterManager.__init__ and wire cluster_manager._gpu_arbiter = gpu_arbiter in app.py after arbiter creation, matching the _capabilities pattern. Without this, the arbiter is only reachable via app.state.gpu_arbiter and prod eviction paths that go through the cluster manager are dead.
… GpuArbiter Add VramReservationManager as the arbiter's INTERNAL reservation layer. Replaces ad-hoc vram_probe with atomic check-and-reserve under an asyncio.Lock so the arbiter is the single VRAM authority. - VramReservationManager: atomic reserve()/release(), sync available() for best-effort admission checks, stats() for observability - GpuArbiter._check_admission now reads against reserved VRAM - GpuArbiter._run_gpu_task atomically reserves local VRAM before execution and releases in finally (normal completion) or via _evict_task (eviction path) - Cluster lease path unchanged (VRAM tracked by lease system) - Exported via scheduler.__init__ lazy loading - 21 new tests: 11 unit (VramReservationManager), 9 integration (arbiter lifecycle), 1 lazy-export test - 59/59 related tests pass (arbiter + cluster + scheduler + vram)
|
Warning Review limit reached
Next review available in: 6 minutes Enable usage-based reviews in Billing to review now. Otherwise, wait until the next included review is available. How can I continue?After more reviews become available, a review can be triggered using the To avoid repeated limits, reduce automatic review volume by pausing incremental auto-reviews earlier, using label-based review opt-in, excluding WIP or generated PR titles, or requesting reviews manually when the PR is ready. If your team needs uninterrupted high-volume reviews, an organization admin can enable usage-based reviews. How do review limits work?CodeRabbit enforces per-developer PR review limits for each organization. Most developers receive the normal plan review availability. For paid Pro and Pro+ PR reviews, CodeRabbit uses adaptive limits for sustained high-volume activity. When a developer's recent PR review activity reaches the 95th percentile or higher among CodeRabbit users, additional reviews become available more gradually as earlier reviews age out of the rolling window. Please refer docs for additional details. Review details⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Pro Plus Run ID: 📒 Files selected for processing (4)
📝 WalkthroughWalkthroughThis PR adds a GPU scheduling arbiter ( ChangesGPU Arbiter Feature
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant Caller
participant GpuArbiter
participant VramReservationManager
participant ClusterManager
Caller->>GpuArbiter: submit_gpu(task, required_vram_mb)
GpuArbiter->>VramReservationManager: check-and-reserve VRAM
alt admitted
GpuArbiter->>ClusterManager: claim lease (if configured)
GpuArbiter->>GpuArbiter: _run_gpu_task(task)
GpuArbiter->>VramReservationManager: release VRAM on completion
else not admitted
GpuArbiter->>GpuArbiter: enqueue task with priority
GpuArbiter->>GpuArbiter: _drain_queue retries admission
end
Caller->>GpuArbiter: evict_lowest_priority()
GpuArbiter->>GpuArbiter: _evict_task cancels running task/future
GpuArbiter->>ClusterManager: release lease
GpuArbiter->>VramReservationManager: release reservation
Possibly related PRs
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Warning Review ran into problems🔥 ProblemsGit: Failed to clone repository. Please run the 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}
| # Build the GPU arbiter — wraps the resource scheduler with VRAM-accounted | ||
| # admission control, queuing, and eviction for GPU-bound workloads. | ||
| try: | ||
| gpu_arbiter = GpuArbiter( |
There was a problem hiding this comment.
CRITICAL: GpuArbiter is constructed here without a vram_probe, so self._vram_probe falls back to _default_vram_probe() which returns (0, 0).
Consequence: _check_admission never enters the local VRAM branch (if total > 0 is always false) and _run_gpu_task always fails the atomic reserve("local", vram) for any required_vram_mb > 0 local task, raising NoResourceAvailableError. The entire VramReservationManager / "single VRAM authority" feature added in this PR is therefore inactive in production — only the tests (which pass an explicit probe) exercise it, giving false confidence that the 59 tests "pass". Wire a real probe (e.g. vram_probe=_probe_nvidia_vram) so local VRAM admission actually functions.
Reply with @kilocode-bot fix it to have Kilo Code address this issue.
| ) | ||
| await gpu_arbiter.start() | ||
| app.state.gpu_arbiter = gpu_arbiter | ||
| cluster_manager._gpu_arbiter = gpu_arbiter # wire for eviction paths |
There was a problem hiding this comment.
WARNING: cluster_manager._gpu_arbiter is assigned here but never read anywhere in the codebase, so the claimed "eviction is not a prod no-op" fix is incomplete — no cluster eviction path actually consults this attribute (grep for _gpu_arbiter shows only the assignment and the None init in manager.py). Either wire the reads (e.g. cluster eviction invoking self._gpu_arbiter.evict_*) or remove the dead wiring to avoid a false sense of correctness.
Reply with @kilocode-bot fix it to have Kilo Code address this issue.
| # Local GPU: check against probed VRAM minus reserved (best-effort | ||
| # read; the actual reserve happens atomically in _run_gpu_task). | ||
| free_vram = self._vram_reservations.available("local") | ||
| _free_raw, total = self._vram_probe() |
There was a problem hiding this comment.
WARNING: Unused variable _free_raw — only total is used on this line; free_vram already comes from self._vram_reservations.available("local").
| _free_raw, total = self._vram_probe() | |
| _, total = self._vram_probe() |
Reply with @kilocode-bot fix it to have Kilo Code address this issue.
| lease_id = lease.lease_id | ||
| elif required_vram_mb > 0: | ||
| # Local GPU path: atomically reserve VRAM through the reservation manager. | ||
| if not await self._vram_reservations.reserve("local", required_vram_mb): |
There was a problem hiding this comment.
WARNING: Admission (_check_admission) is a non-atomic, best-effort check, but the authoritative reserve happens later inside _run_gpu_task. For a directly-admitted (non-queued) task, if reserve then fails, the caller gets a hard NoResourceAvailableError with no queue fallback — unlike queued tasks, which would simply wait. Two near-concurrent submits that both pass admission can leave the second hard-failing instead of queueing. Consider routing directly-admitted tasks that fail reserve back into the queue, or document this failure-mode asymmetry.
Reply with @kilocode-bot fix it to have Kilo Code address this issue.
| # await here (sync method), so schedule the release as a | ||
| # background task that will run promptly on the event loop. | ||
| if _vram > 0: | ||
| asyncio.create_task( |
There was a problem hiding this comment.
SUGGESTION: _evict_task is a synchronous method, yet asyncio.create_task(...) here will raise RuntimeError if there is no running event loop (e.g. called from a thread or during shutdown), and otherwise schedules a task whose exceptions are silently swallowed. Guard with a running-loop check (and consider capturing the returned task) so the VRAM release cannot fail silently.
Reply with @kilocode-bot fix it to have Kilo Code address this issue.
Code Review SummaryStatus: 1 Issue Found | Recommendation: Address before merge Overview
Issue Details (click to expand)CRITICAL(none) WARNING(none) SUGGESTION
Files Reviewed (1 file, incremental)
Prior findings on unchanged lines (e.g. admission race @295, Resolved in this commit (c7fde9d)
Fix these issues in Kilo Cloud Previous Review Summaries (2 snapshots, latest commit 34ed95a)Current summary above is authoritative. Previous snapshots are kept for context only. Previous review (commit 34ed95a)Status: 6 Issues Found | Recommendation: Address before merge Overview
Issue Details (click to expand)CRITICAL(none) WARNING
SUGGESTION
Files Reviewed (2 files)
Fix these issues in Kilo Cloud Previous review (commit 6a8f8f9)Status: 5 Issues Found | Recommendation: Address before merge Overview
Issue Details (click to expand)CRITICAL
WARNING
SUGGESTION
Files Reviewed (8 files)
Reviewed by hy3-20260706:free · Input: 74.7K · Output: 15K · Cached: 244.7K |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (4)
tinyagentos/app.py (1)
1149-1149: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueSimplify redundant conditional.
resource_scheduler if resource_scheduler is not None else Noneis equivalent to justresource_scheduler— the conditional adds no behavior.♻️ Proposed refactor
- scheduler=resource_scheduler if resource_scheduler is not None else None, + scheduler=resource_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` at line 1149, The scheduler assignment in the app setup is using a redundant conditional that does nothing beyond returning the same value. Update the `scheduler=` argument in the relevant configuration block in `app.py` to pass `resource_scheduler` directly, and remove the unnecessary `resource_scheduler if resource_scheduler is not None else None` expression.tests/test_vram_reservation.py (3)
21-27: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueDuplicated
_make_taskhelper across both test files.The
_make_taskfunction is identical intests/test_vram_reservation.pyandtests/test_gpu_arbiter_894.py. Consider extracting it to a sharedconftest.pyor test utility module to avoid divergence.🤖 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_vram_reservation.py` around lines 21 - 27, The `_make_task` helper is duplicated across the VRAM and GPU arbiter tests, so extract the shared task-construction logic into a common test fixture or utility and update `test_vram_reservation` and `test_gpu_arbiter_894` to import and use it; keep the behavior consistent by centralizing the `Task` creation in the shared helper rather than maintaining two copies.
148-164: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winTest doesn't verify VRAM is reserved during execution.
The test name claims to verify both reservation and release, but only checks
total_reserved == 0before and after. Since the payload (asyncio.sleep(0)) completes instantly, there's no window to observe the reservation. Consider using a blocking payload with an event to assert VRAM is held mid-execution, similar totest_eviction_releases_local_vramortest_concurrent_tasks_account_correctly.🤖 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_vram_reservation.py` around lines 148 - 164, The test for GpuArbiter.submit_gpu only checks reservation state before and after completion, so it never verifies VRAM is actually held during execution. Update test_submit_gpu_reserves_and_releases_vram to use a blocking payload (for example, an event-controlled coroutine like in test_eviction_releases_local_vram or test_concurrent_tasks_account_correctly) and assert arbiter._vram_reservations.total_reserved("local") is nonzero while submit_gpu is still running, then confirm it returns to zero after the task finishes.
250-250: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick winTiming-based sleeps risk CI flakiness.
Several tests use fixed-duration
asyncio.sleep()calls (lines 250, 307, 313, 379, 391) to wait for background operations. These are inherently racy under load. Where possible, prefer event-based synchronization orasyncio.wait_foron a condition. For example, line 250 could wait on an event set by the release path instead of sleeping 0.1s.Also applies to: 307-307, 313-313, 379-379, 391-391
🤖 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_vram_reservation.py` at line 250, Replace the fixed asyncio.sleep() waits in the VRAM reservation tests with deterministic synchronization so the tests do not race under CI load. In the affected test helpers and cases in test_vram_reservation.py, use an Event/condition that is set by the relevant release/background path, or wrap the expected state change with asyncio.wait_for instead of sleeping. Update the waits near the release flow and other background-op checks so they wait on the actual signal from the code under test rather than elapsed time.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@tinyagentos/app.py`:
- Around line 1144-1160: Initialize app.state.gpu_arbiter to None in the eager
app-state setup alongside the other lifespan-managed attributes, so routes can
safely read it before startup completes. Add the initialization near the
existing eager assignments for mcp_supervisor, orchestrator, and trace_registry,
and keep the lifespan code in the GPU arbiter startup block in app.py assigning
the real GpuArbiter instance later as it does now.
In `@tinyagentos/scheduler/gpu_arbiter.py`:
- Around line 219-223: The queued-task path in `GpuArbiter.submit` can block
forever on `await done` because `stop()` only cancels `_queue_processor_task`
and leaves pending `_queue` entries unresolved. Update `GpuArbiter.stop()` to
walk any outstanding queued entries and complete each entry’s `_arbiter_future`
with a failure such as `NoResourceAvailableError` or `CancelledError`, so
callers waiting in `submit()` are released when the arbiter shuts down.
---
Nitpick comments:
In `@tests/test_vram_reservation.py`:
- Around line 21-27: The `_make_task` helper is duplicated across the VRAM and
GPU arbiter tests, so extract the shared task-construction logic into a common
test fixture or utility and update `test_vram_reservation` and
`test_gpu_arbiter_894` to import and use it; keep the behavior consistent by
centralizing the `Task` creation in the shared helper rather than maintaining
two copies.
- Around line 148-164: The test for GpuArbiter.submit_gpu only checks
reservation state before and after completion, so it never verifies VRAM is
actually held during execution. Update
test_submit_gpu_reserves_and_releases_vram to use a blocking payload (for
example, an event-controlled coroutine like in test_eviction_releases_local_vram
or test_concurrent_tasks_account_correctly) and assert
arbiter._vram_reservations.total_reserved("local") is nonzero while submit_gpu
is still running, then confirm it returns to zero after the task finishes.
- Line 250: Replace the fixed asyncio.sleep() waits in the VRAM reservation
tests with deterministic synchronization so the tests do not race under CI load.
In the affected test helpers and cases in test_vram_reservation.py, use an
Event/condition that is set by the relevant release/background path, or wrap the
expected state change with asyncio.wait_for instead of sleeping. Update the
waits near the release flow and other background-op checks so they wait on the
actual signal from the code under test rather than elapsed time.
In `@tinyagentos/app.py`:
- Line 1149: The scheduler assignment in the app setup is using a redundant
conditional that does nothing beyond returning the same value. Update the
`scheduler=` argument in the relevant configuration block in `app.py` to pass
`resource_scheduler` directly, and remove the unnecessary `resource_scheduler if
resource_scheduler is not None else None` expression.
🪄 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: 9b188593-2613-4fd5-a6ab-28a7f6acfe5a
📒 Files selected for processing (8)
tests/test_gpu_arbiter_894.pytests/test_vram_reservation.pytinyagentos/app.pytinyagentos/cluster/manager.pytinyagentos/scheduler/__init__.pytinyagentos/scheduler/discovery.pytinyagentos/scheduler/gpu_arbiter.pytinyagentos/scheduler/types.py
|
Thanks for pushing on the VRAM bookkeeping. I did a full code-level read of this against the epic gate (#185) and the blockers from the #1718 set review. The mechanical class itself is reasonable, but as the gate-closer it does not land yet, and there is one bug that CI is hiding. Concrete findings, most important first: 1. The epic gate is still open. 2. Two VRAM authorities, not one. This adds a 3. HIGH, CI is masking it: async lease APIs called without await. In 4. CRITICAL (kilo, confirmed): no vram_probe wired. The arbiter is built without a 5. MEDIUM: 6. MEDIUM: 7. Eviction-ordering TOCTOU (blocker 2) still open. Net: hold this with the rest of the arbiter set rather than merge piecemeal. The unblocking sequence is (3) await the lease APIs + async test double, (4) wire the real vram_probe, then the big one (1)+(2): one VRAM authority routed through |
|
Staging note (direction on #894): we are shipping #1725's model-load VRAM guard first as the single authority, and staging the arbiter (this PR + #1718) as the follow-up that closes #894. So this is held, not rejected. The unblock sequence from my review above stands: await the async lease APIs (with a matching async test double), wire the real vram_probe, then the big one, route a real admission through submit_gpu and reconcile with #1725 so there is one VRAM authority. That last step is what actually closes the epic. |
…ueued callers, eager gpu_arbiter None init, drop unused _free_raw Kilo CRITICAL: GpuArbiter constructed without vram_probe → _default_vram_probe returned (0,0) in production, making VRAM admission entirely inert. Pass _probe_nvidia_vram from app.py so the arbiter sees real GPU VRAM. Kilo WARNING: unused _free_raw in _check_admission — replace with _. CodeRabbit: app.state.gpu_arbiter now eagerly initialized to None so routes can safely read it before the lifespan creates the real instance. CodeRabbit: stop() walks remaining queue entries and cancels their _arbiter_future so submit_gpu callers don't hang forever on shutdown. Tests: 30/30 pass (test_gpu_arbiter_894 + test_vram_reservation)
…_priority async jfc review finding jaylfc#3 (HIGH): claim_lease() and release_lease() are async def in ClusterManager but were called synchronously in _run_gpu_task and _evict_task. In prod, claim_lease returns a coroutine (never None) → AttributeError on lease.lease_id. CI was green only because FakeClusterManager in tests was sync. - _run_gpu_task: await claim_lease + await release_lease in finally - _evict_task: async def, await release_lease on eviction path - evict_lowest_priority: async def, await _evict_task - FakeClusterManager: claim_lease + release_lease now async def - All test callers updated with await Tests: 30/30 pass
| # Release any queued callers blocked on submit_gpu — they would | ||
| # otherwise hang forever waiting for a queue processor that's gone. | ||
| cancelled = 0 | ||
| while not self._queue.empty(): |
There was a problem hiding this comment.
WARNING: stop() drains only entries still in self._queue, but an entry the queue processor has already dequeued and is actively running is no longer in the queue: _drain_queue removes it via get_nowait() (line 393) and then awaits self._run_gpu_task(...). If stop() cancels the processor mid-run, _drain_queue's CancelledError handler (lines 400-401) just logs and breaks without resolving that entry's _arbiter_future, so the submitter's await done in submit_gpu hangs forever — the exact hang this new code is meant to prevent, for a different subset of tasks.
Also walk _running/_running_tasks (or resolve the future in _drain_queue's CancelledError branch) so in-flight admitted entries are released on shutdown.
Reply with @kilocode-bot fix it to have Kilo Code address this issue.
| break | ||
| future = getattr(entry.task, "_arbiter_future", None) | ||
| if future is not None and not future.done(): | ||
| future.cancel() |
There was a problem hiding this comment.
SUGGESTION: Cancelling the future here routes through submit_gpu's except asyncio.CancelledError: self._evicted += 1; raise (line 238), so a normal arbiter shutdown is miscounted as an eviction in stats(). Consider distinguishing shutdown from eviction (e.g. a dedicated counter, or resolve the future with NoResourceAvailableError('arbiter stopped') instead of cancelling) so the metric isn't skewed.
Reply with @kilocode-bot fix it to have Kilo Code address this issue.
| from tinyagentos.installation_state import InstallationState | ||
| from tinyagentos.scheduler import BackendCatalog, HistoryStore, ScoreCache, TaskScheduler | ||
| from tinyagentos.scheduler import BackendCatalog, GpuArbiter, HistoryStore, ScoreCache, TaskScheduler | ||
| from tinyagentos.scheduler.gpu_arbiter import _probe_nvidia_vram |
There was a problem hiding this comment.
SUGGESTION: Importing the private symbol _probe_nvidia_vram (underscore-prefixed) couples app.py to an internal API of gpu_arbiter that may be renamed or removed without notice. Consider exporting a public probe name (e.g. probe_nvidia_vram) or a GpuArbiter factory, and importing that instead.
Reply with @kilocode-bot fix it to have Kilo Code address this issue.
|
Closing as superseded by the merged #1725 ( |
2 participants
Adds VramReservationManager as the arbiter's internal reservation layer with atomic check-and-reserve under asyncio.Lock. Wired into GpuArbiter so the arbiter is the single VRAM authority.
Tests: 21 new (11 unit + 9 integration + 1 export), 59/59 related tests pass.
Task: t_19db5cd3 | Prior: #1718 (drain/auth/eviction — MERGEABLE)
Summary by CodeRabbit
New Features
Bug Fixes