feat(cluster): pause/resume GPU queue, graceful worker detach, hardware-aware LLM queuing#1689
feat(cluster): pause/resume GPU queue, graceful worker detach, hardware-aware LLM queuing#1689hognek wants to merge 4 commits into
Conversation
|
Warning Review limit reached
Next review available in: 42 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 (12)
✨ 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 |
|
All contributors have signed the CLA ✍️ ✅ |
|
|
hognek
force-pushed
the
feat/cluster-pause-resume-detach
branch
from
1d02c5e to
3fa5d53
Compare
|
I have read the CLA Document and I hereby sign the CLA |
|
recheck |
hognek
force-pushed
the
feat/cluster-pause-resume-detach
branch
from
03c1c97 to
3fa5d53
Compare
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? |
…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.
taOS jaylfc#894 Slice 2 — builds on Slice 1 (VRAM endpoint in worker heartbeat). Adds VramTracker: per-GPU VRAM admission, reservation, priority-based eviction, and wait-for-free signalling. Integrates into Resource via optional vram_tracker parameter — can_admit() checks VRAM budget and run() reserves VRAM with eviction + wait loop. Prevents concurrent-load driver crashes (NVIDIA Xid 62) by ensuring GPU VRAM is accounted for before admitting inference tasks. New: - tinyagentos/scheduler/vram_tracker.py: VramAllocation + VramTracker - Resource: vram_tracker param, VRAM-aware can_admit(), VRAM-aware run() - tests/test_gpu_arbiter.py: 22 tests covering accounting, admission, eviction, wait signalling, and Resource integration
…re-aware LLM queuing Adds three taOS jaylfc#796 capabilities on top of the GPU arbiter + lease system: 1. Pause/resume queues — GpuArbiter.pause()/resume() halt/restart queue processing without rejecting new submissions. Running tasks finish; queued tasks wait. Exposed via POST /api/cluster/gpu-queue/pause and /api/cluster/gpu-queue/resume. 2. Graceful worker detach — ClusterManager.drain_worker() enters a 'draining' state where no new tasks route to the worker but existing leases run to completion. The monitor loop auto-completes the drain when all leases are released. Force-drain (graceful=False) releases all leases and evicts GPU tasks immediately. Cancel-drain restores online status. Exposed via /api/cluster/workers/{name}/drain and /api/cluster/workers/{name}/cancel-drain. 3. Hardware-aware LLM queuing — GpuArbiter.submit_gpu() accepts required_gpu_arch (e.g. 'sm_86') and checks cluster workers for compatible GPU hardware before admission. Uses both model name and compute_cap field from worker hardware info. Tests: 102 pass (64 existing + 38 new)
hognek
force-pushed
the
feat/cluster-pause-resume-detach
branch
from
3fa5d53 to
b9b02ab
Compare
claim_lease became async in upstream dev; the test files added by this branch were calling it synchronously. Added @pytest.mark.asyncio + async def to the 5 affected test methods and await at each call site.
![kilo-code-bot[bot]](https://avatars.githubusercontent.com/in/2193792?s=60&v=4){kind=small}
| worker_name, _ = parsed | ||
| worker = self._workers.get(worker_name) | ||
| if worker is None or worker.status != "online": | ||
| if worker is None or worker.status not in ("online",): |
| # Force-release all leases for this worker | ||
| lids = [ | ||
| lid for lid, lease in self._leases.items() | ||
| if lease.resource_id.startswith(name + ":") |
| gpu_model = gpu_info.get("model", "") or "" | ||
| # Check both the model string and the compute_cap field if present | ||
| cc = gpu_info.get("compute_cap", "") or "" | ||
| if required_gpu_arch in gpu_model or required_gpu_arch in cc: |
| # Build the GPU arbiter — wraps the resource scheduler with VRAM-accounted | ||
| # admission control, queuing, and eviction for GPU-bound workloads. | ||
| try: | ||
| gpu_arbiter = GpuArbiter( |
| logger.debug("Lease %s released — worker %s went offline", lid, worker.name) | ||
|
|
||
| # Handle draining workers: auto-complete if no active leases | ||
| elif worker.status == "draining": |
| ) | ||
| elif (now - worker.last_heartbeat) > HEARTBEAT_TIMEOUT: | ||
| # Draining worker went stale — force-finish the drain | ||
| async with self._lease_lock: |
| is marked offline immediately. | ||
| """ | ||
| cluster = request.app.state.cluster_manager | ||
| # Read graceful flag from query param or JSON body |
| body = await request.json() | ||
| if isinstance(body, dict): | ||
| graceful = body.get("graceful", True) | ||
| except Exception: |
|
|
||
| async def wait_for_vram(self) -> None: | ||
| self._vram_freed.clear(); await self._vram_freed.wait() | ||
|
|
| model_id=getattr(task, 'model_id', ''), | ||
| priority=int(task.priority)) | ||
| if reserved: break | ||
| await tracker.wait_for_vram() |
|
|
||
| def _gpu_vram_probe() -> int: | ||
| free, _total = _probe_nvidia_vram() | ||
| return free if free > 0 else 999_999 # optimistic |
| continue | ||
| if self._cluster_manager is not None: | ||
| lease = getattr(self._cluster_manager, "_leases", {}).get(lid) | ||
| if lease is not None and lease.resource_id.startswith(worker_name + ":"): |
| gpu_model = gpu_info.get("model", "") or "" | ||
| # Check both the model string and the compute_cap field if present | ||
| 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.
CRITICAL: Substring matching for GPU architecture is unsafe. required_gpu_arch in gpu_model will false-positive on any model name that contains the substring — e.g. requesting sm_86 against an NVIDIA A100 (sm_80) whose later replacement SKU is sm_86x would still match; conversely sm_8 would match every sm_8x. Same risk for required_gpu_arch in cc. Use an exact (case-insensitive) comparison against a normalized arch field, or parse a structured compute_cap field. Also: this branch treats draining workers as valid for admission even though draining workers are excluded from routing — pick one policy and apply it consistently.
Reply with @kilocode-bot fix it to have Kilo Code address this issue.
| # 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 but never wired into cluster_manager._gpu_arbiter. As a result ClusterManager.drain_worker(graceful=False) will fail its if self._gpu_arbiter is not None check (manager.py ~line 433) and skip GPU task eviction — the headlining feature of this PR will silently no-op for server-side initiated drains. Add cluster_manager._gpu_arbiter = gpu_arbiter (and ideally inject it via the new ClusterManager(gpu_arbiter=...) kwarg used in tests).
Reply with @kilocode-bot fix it to have Kilo Code address this issue.
| # Force-release all leases for this worker | ||
| lids = [ | ||
| lid for lid, lease in self._leases.items() | ||
| if lease.resource_id.startswith(name + ":") |
There was a problem hiding this comment.
CRITICAL: Force-drain mutates self._leases without holding self._lease_lock. The codebase serialises all lease mutations through _lease_lock (see claim_lease, _sweep_expired_leases, and the new monitor-loop timeout branch). Popping entries here while a concurrent claim or sweep runs can leave the lease table in an inconsistent state. Make drain_worker accept the lock, or use _lease_lock from a synchronous helper — see the timeout branch below for the async pattern.
Reply with @kilocode-bot fix it to have Kilo Code address this issue.
| logger.debug("Lease %s released — worker %s went offline", lid, worker.name) | ||
|
|
||
| # Handle draining workers: auto-complete if no active leases | ||
| elif worker.status == "draining": |
There was a problem hiding this comment.
WARNING: _worker_for_resource was tightened to exclude draining workers (good), but this monitor-loop code path still inspects draining workers' lease maps and promotes them to offline. Combined with cancel_drain blindly flipping status back to "online" (no heartbeat check, no in-flight task consideration), a draining worker whose heartbeat has briefly gone stale will be marked offline even though its leases are still healthy — and conversely a worker the operator just cancelled draining on could be restored to online while its leases were already force-collected. Verify heartbeat and lease state under _lease_lock before mutating worker.status.
Reply with @kilocode-bot fix it to have Kilo Code address this issue.
| ) | ||
| elif (now - worker.last_heartbeat) > HEARTBEAT_TIMEOUT: | ||
| # Draining worker went stale — force-finish the drain | ||
| async with self._lease_lock: |
There was a problem hiding this comment.
WARNING: lids is computed before async with self._lease_lock: is acquired, so new leases for this worker can be granted (or expired) between the snapshot and the pop. Move the list-comprehension inside the async with block (compute and pop atomically) or re-validate each lid inside the lock. The drain-timed-out branch in drain_worker (line ~435) has the same hazard.
Reply with @kilocode-bot fix it to have Kilo Code address this issue.
| body = await request.json() | ||
| if isinstance(body, dict): | ||
| graceful = body.get("graceful", True) | ||
| except Exception: |
There was a problem hiding this comment.
SUGGESTION: cancel_drain returns HTTP 404 for "worker exists but is not draining". That conflates 'not found' with 'state conflict' — return 409 Conflict instead so callers can distinguish the two cases.
Reply with @kilocode-bot fix it to have Kilo Code address this issue.
|
|
||
| async def wait_for_vram(self) -> None: | ||
| self._vram_freed.clear(); await self._vram_freed.wait() | ||
|
|
There was a problem hiding this comment.
CRITICAL: wait_for_vram does self._vram_freed.clear() and then await self._vram_freed.wait(). If release() is called between clear() and wait() (e.g. by another task on the same loop tick), the set() from release is lost and the waiter hangs forever. Use asyncio.Condition (or check-then-wait while holding _lock) instead of a plain Event.
Reply with @kilocode-bot fix it to have Kilo Code address this issue.
| model_id=getattr(task, 'model_id', ''), | ||
| priority=int(task.priority)) | ||
| if reserved: break | ||
| await tracker.wait_for_vram() |
There was a problem hiding this comment.
WARNING: Resource.run() now loops evict_and_reserve → wait_for_vram forever. If no eviction candidate ever frees enough VRAM (e.g. all running tasks are non-evictable / higher priority), this task will block indefinitely and the caller will never see an error. Add a max-iteration or timeout and surface NoResourceAvailableError so timeouts are observable in upstream code and tests.
Reply with @kilocode-bot fix it to have Kilo Code address this issue.
|
|
||
| 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.
WARNING: _gpu_vram_probe falls back to 999_999 MiB when nvidia-smi fails or returns 0. That optimistic default will silently admit GPU tasks that exceed actual VRAM, exactly the Xid-62 driver-crash scenario the PR exists to prevent. Use a conservative constant (e.g. min(free, known_total/2)) and/or treat the unknown probe as 0 so tasks queue instead of running.
Reply with @kilocode-bot fix it to have Kilo Code address this issue.
| continue | ||
| if self._cluster_manager is not None: | ||
| lease = getattr(self._cluster_manager, "_leases", {}).get(lid) | ||
| if lease is not None and lease.resource_id.startswith(worker_name + ":"): |
There was a problem hiding this comment.
WARNING: release_tasks_for_worker reads self._cluster_manager._leases directly (private attribute) and never acquires ClusterManager._lease_lock. Concurrent mutation by the monitor loop or a drain_worker(graceful=False) can cause the snapshot to miss leases or read stale entries. Add a public method on ClusterManager (e.g. get_leases_for_worker(name)) that takes the lock internally.
Reply with @kilocode-bot fix it to have Kilo Code address this issue.
Code Review SummaryStatus: 11 Issues Found | Recommendation: Address before merge Overview
Issue Details (click to expand)CRITICAL
WARNING
SUGGESTION
Files Reviewed (5 files)
Fix these issues in Kilo Cloud Reviewed by minimax-m3 · Input: 114.9K · Output: 27.8K · Cached: 6.6M |
|
Heads up: this overlaps the arbiter epic (#1683) almost entirely. It re-adds the same gpu_arbiter.py and shares #1706/#1707's file footprint, and its drain/pause collides with #1690 on the same manager.py methods. I left a consolidation note on #1683. Suggest we either close this in favour of that stack or rebase it to sit on top, rather than land two PRs that both add the arbiter core. Not merging this standalone. It also carries a few real concurrency issues (lockless lease mutation in drain_worker, wait_for_vram lost-wakeup, substring GPU-arch match) that #1706/#1707 look like they already address, so folding it into that stack avoids double work. |
…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.
taOS jaylfc#796 features extracted from jaylfc#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.
2 participants
taOS #796: Cluster pause/resume, graceful worker detach, hardware-aware LLM queuing
Three capabilities for stable GPU usage by Skald embedding batches, built on top of the GPU arbiter + lease system.
1. Pause/resume GPU queues
2. Graceful worker detach
3. Hardware-aware LLM queuing
Tests
Fixes #796
PRBODY 2>&1