feat(gpu-arbiter): event-driven admission wakeup replacing 2s poll tick (#1864 A3)#1986
feat(gpu-arbiter): event-driven admission wakeup replacing 2s poll tick (#1864 A3)#1986hognek wants to merge 1 commit into
Conversation
…ck (jaylfc#1864 A3) Replace the hardcoded asyncio.sleep(2) in _process_queue with an asyncio.Event-based wake path. The drain loop now blocks on asyncio.wait_for(self._wake.wait(), timeout=drain_tick_seconds) so a reservation release or task completion immediately wakes the drain loop. The 2 s constant becomes the configurable fallback timeout (drain_tick_seconds, default 2.0). - __init__: new drain_tick_seconds param, self._wake Event - _signal_capacity(): new method — calls self._wake.set() - _process_queue: event-driven wait + clear, fallback timeout - _release_reservation: calls _signal_capacity on actual release - _run_gpu_task finally: calls _signal_capacity on completion - tests/test_gpu_arbiter_wakeup.py: 2 new tests * test_release_triggers_immediate_drain (60 s tick, < 0.5 s admit) * test_poll_tick_still_drains_without_signal (0.1 s tick)
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? |
📝 WalkthroughWalkthroughGpuArbiter now supports event-driven queue draining after VRAM or task state changes, with a configurable polling fallback. New async tests verify prompt admission after release and continued draining when capacity changes without an explicit signal. ChangesGPU arbiter capacity wakeups
Estimated code review effort: 3 (Moderate) | ~20 minutes Sequence Diagram(s)sequenceDiagram
participant GPU_Task
participant GpuArbiter
participant ReservationManager
GPU_Task->>GpuArbiter: release reservation or finish
GpuArbiter->>GpuArbiter: signal _wake
GpuArbiter->>ReservationManager: process queued admission
ReservationManager-->>GpuArbiter: report capacity
GpuArbiter->>GPU_Task: admit queued task
🚥 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 |
![kilo-code-bot[bot]](https://avatars.githubusercontent.com/in/2193792?s=60&v=4){kind=small}
| self._paused_at: float | None = None | ||
|
|
||
| # ── taOS #1864 A3: event-driven wakeup ─────────────────────────── | ||
| self._drain_tick_seconds = drain_tick_seconds |
There was a problem hiding this comment.
SUGGESTION: drain_tick_seconds is not validated and is passed straight into asyncio.wait_for(..., timeout=self._drain_tick_seconds). asyncio.wait_for raises ValueError if timeout <= 0, so a caller passing 0 or a negative value will crash the entire _process_queue drain loop on startup. Consider clamping/validating the value (e.g. max(drain_tick_seconds, 0.01) or raising a clear error in __init__).
Reply with @kilocode-bot fix it to have Kilo Code address this issue.
| async with self._running_lock: | ||
| entry = self._running.pop(task.id, None) | ||
| self._running_tasks.pop(task.id, None) | ||
| self._signal_capacity() # taOS #1864 A3: wake drain on completion |
There was a problem hiding this comment.
SUGGESTION: _run_gpu_task's finally calls _release_reservation (line 444), which already invokes _signal_capacity() at line 250, and then calls _signal_capacity() again here. The wake fires twice for every normal task completion. Event.set() is idempotent so this is harmless today, but the redundancy is misleading and could break a future wait_for-with-count design. Consider dropping one of the two calls (e.g. let _release_reservation own the signal and remove this line).
Reply with @kilocode-bot fix it to have Kilo Code address this issue.
Code Review SummaryStatus: 2 Issues Found | Recommendation: Address before merge Overview
Issue Details (click to expand)SUGGESTION
Files Reviewed (2 files)
Fix these issues in Kilo Cloud Reviewed by hy3:free · Input: 39.4K · Output: 2.9K · Cached: 107.5K |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
🧹 Nitpick comments (2)
tinyagentos/scheduler/gpu_arbiter.py (2)
543-551: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick winSignal capacity on queue insertion and unpausing to maximize responsiveness.
To fully achieve the PR's goal of admitting contended operations in less than 2 seconds, consider waking the drain loop when new tasks are queued and when the arbiter is resumed. Currently, if a high-priority task is queued (e.g., due to insufficient capacity requiring eviction) or the arbiter is resumed, processing is still delayed by up to the
drain_tick_secondsfallback.Calling
self._signal_capacity()insubmit_gpu(after queue insertion) andresumeensures prompt preemption and processing.🤖 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 543 - 551, Update submit_gpu to call _signal_capacity immediately after inserting a task into the queue, and update resume to call _signal_capacity when the arbiter is unpaused. Preserve the existing _signal_capacity behavior while ensuring both newly queued work and resumed processing wake the drain loop without waiting for drain_tick_seconds.
552-561: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick winCatch and log unexpected exceptions to prevent the queue processor from silently dying.
If
await self._drain_queue()raises an unexpected exception, thewhile Trueloop will break, and the background_queue_processor_taskwill silently terminate. The arbiter will then permanently stop processing queued tasks.Consider wrapping the queue processing step in a generic exception handler to log the error and allow the loop to continue.
🛡️ Proposed fix to ensure resilience
self._wake.clear() if not self._paused: - await self._drain_queue() + try: + await self._drain_queue() + except Exception as e: + logger.exception("gpu-arbiter: unexpected error in queue processor: %s", e) except asyncio.CancelledError: raise🤖 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 552 - 561, Update _process_queue so unexpected exceptions from _drain_queue are caught by a generic exception handler, logged with the existing arbiter logging mechanism, and followed by continuing the while loop. Preserve the current cancellation and normal queue-processing behavior, ensuring _queue_processor_task does not terminate silently after a drain failure.
🤖 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.
Nitpick comments:
In `@tinyagentos/scheduler/gpu_arbiter.py`:
- Around line 543-551: Update submit_gpu to call _signal_capacity immediately
after inserting a task into the queue, and update resume to call
_signal_capacity when the arbiter is unpaused. Preserve the existing
_signal_capacity behavior while ensuring both newly queued work and resumed
processing wake the drain loop without waiting for drain_tick_seconds.
- Around line 552-561: Update _process_queue so unexpected exceptions from
_drain_queue are caught by a generic exception handler, logged with the existing
arbiter logging mechanism, and followed by continuing the while loop. Preserve
the current cancellation and normal queue-processing behavior, ensuring
_queue_processor_task does not terminate silently after a drain failure.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro Plus
Run ID: 77a5f5bb-66b6-4731-b4f7-414c5727291d
📒 Files selected for processing (2)
tests/test_gpu_arbiter_wakeup.pytinyagentos/scheduler/gpu_arbiter.py
1 participant
Summary
Implements Slice A3 of #1864: event-driven admission wakeup (locked decision 7).
Replaces the hardcoded
asyncio.sleep(2)poll tick in_process_queuewith anasyncio.Event-based wake path. A reservation release or task completion now immediately wakes the drain loop, so contended ops admit in<< 2 sinstead of≤ 2 s. The 2 s constant becomes the configurable fallback timeout (drain_tick_seconds, default 2.0).Changes
__init__: newdrain_tick_seconds: float = 2.0param,self._wake: asyncio.Event_signal_capacity(): new internal method —self._wake.set()_process_queue:asyncio.wait_for(self._wake.wait(), timeout=drain_tick_seconds)+ clear_release_reservation: calls_signal_capacity()on actual reservation release_run_gpu_taskfinally: calls_signal_capacity()after_runningpoptests/test_gpu_arbiter_wakeup.py: 2 new testsTest results
Targeted: 25/25 pass (2 new + 23 existing regression)
Canonical gate: running (no failures visible)
Summary by CodeRabbit
New Features
Bug Fixes