fix(cluster): release GPU leases on worker unregister (fixes #1705)#1726
Conversation
|
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 (4)
🚧 Files skipped from review as they are similar to previous changes (3)
📝 WalkthroughWalkthrough
ChangesAsync worker unregistration with lease cleanup
Estimated code review effort: 3 (Moderate) | ~20 minutes Sequence Diagram(s)sequenceDiagram
participant Client
participant ClusterRoute
participant ClusterManager
participant LeaseStorage
Client->>ClusterRoute: DELETE /api/cluster/workers/{name}
ClusterRoute->>ClusterManager: await unregister_worker(name)
ClusterManager->>LeaseStorage: acquire _lease_lock
ClusterManager->>LeaseStorage: remove worker leases
ClusterManager-->>ClusterRoute: return removed
ClusterRoute-->>Client: 200 removed or 404 not found
Possibly related PRs
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
Qodo reviews are paused for this user.Troubleshooting steps vary by plan Learn more → On a Teams plan? Using GitHub Enterprise Server, GitLab Self-Managed, or Bitbucket Data Center? |
| ) | ||
|
|
||
| self._workers.pop(name, None) | ||
| logger.info("Worker '%s' unregistered — %d leases released", name, len(lids)) |
There was a problem hiding this comment.
WARNING: Worker pop happens outside _lease_lock — operation is not visibly atomic w.r.t. concurrent claims
The self._workers.pop(name, None) call runs after the async with self._lease_lock: block has already exited. While no await currently sits between the lock release and the pop (so pure-asyncio preemption won't fire here today), this is a fragile pattern:
- Any future
awaitadded insideunregister_workerbetween these two lines (e.g., for logging, metrics, or a callback) would open a window where a concurrentclaim_leasecan see the worker still registered, create a new lease, and have that lease silently leaked — the exact bug this PR is fixing. - The docstring advertises "remove a worker and ALL of its active GPU leases" — the implementation should match that contract visibly.
Move self._workers.pop(name, None) inside the async with self._lease_lock: block, alongside the lease release.
Reply with @kilocode-bot fix it to have Kilo Code address this issue.
| heartbeat-timeout path (taOS #1705). | ||
| """ | ||
| worker = self._workers.get(name) | ||
| if worker is None: |
There was a problem hiding this comment.
SUGGESTION: Unused local variable worker
The result of self._workers.get(name) is stored in worker but never read. Replace the two-line check with a single membership test:
if name not in self._workers:
return FalseReply with @kilocode-bot fix it to have Kilo Code address this issue.
| assert lease is not None | ||
| assert lease.caller == "second" | ||
|
|
||
| async def test_unregister_releases_leases_atomically(self): |
There was a problem hiding this comment.
SUGGESTION: Test name claims "atomically" but only verifies lock-blocking, not actual atomicity
The test test_unregister_releases_leases_atomically only holds _lease_lock from the test side and asserts that the unregister coroutine is blocked. It does not actually race a concurrent claim_lease against unregister_worker to verify end-to-end atomicity (e.g., that a claim started just before unregister is either fully created-and-removed, or never created and not leaked).
Either rename the test to test_unregister_blocks_on_lease_lock for accuracy, or strengthen it to spawn a claim_lease task concurrently with unregister_worker and assert the post-state is consistent (no leaked lease, worker gone).
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 (4 files)
Fix these issues in Kilo Cloud Previous Review Summary (commit 6fd0bb1)Current summary above is authoritative. Previous snapshots are kept for context only. Previous review (commit 6fd0bb1)Status: 3 Issues Found | Recommendation: Address before merge Overview
Issue Details (click to expand)WARNING
SUGGESTION
Files Reviewed (5 files)
Reviewed by hy3-20260706:free · Input: 78.3K · Output: 10.2K · Cached: 342.4K |
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
tinyagentos/cluster/manager.py (1)
259-271: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDuplicate lease-release logic vs.
_monitor_loop's offline path.This block duplicates the filter/pop/log pattern used for offline-triggered lease release (Lines 538-546). Consider extracting a shared helper, e.g.
_release_leases_for_worker_locked(name, reason), called from bothunregister_workerand_monitor_loopwhile holding_lease_lock, to avoid the two paths drifting apart over time.🤖 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 259 - 271, The lease-release logic in unregister_worker duplicates the offline-path cleanup already used by _monitor_loop, so extract the shared filter/pop/log flow into a helper such as _release_leases_for_worker_locked(name, reason) and call it from both places while holding _lease_lock. Keep the helper responsible for finding matching lids in self._leases, removing them, and emitting the debug log so the two paths stay consistent.
🤖 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/cluster/manager.py`:
- Around line 255-275: Make `unregister_worker` atomic by moving the `_workers`
existence check and `pop()` into the same `_lease_lock` block used for lease
cleanup, so concurrent callers cannot both pass the pre-check and both report
success. Update `ClusterManager.unregister_worker` to check
`self._workers.get(name)` only while holding `_lease_lock`, release matching
leases, remove the worker from `_workers`, and then log the final lease count
from the locked section.
---
Nitpick comments:
In `@tinyagentos/cluster/manager.py`:
- Around line 259-271: The lease-release logic in unregister_worker duplicates
the offline-path cleanup already used by _monitor_loop, so extract the shared
filter/pop/log flow into a helper such as
_release_leases_for_worker_locked(name, reason) and call it from both places
while holding _lease_lock. Keep the helper responsible for finding matching lids
in self._leases, removing them, and emitting the debug log so the two paths stay
consistent.
🪄 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: 9340a7f5-d468-4950-84c5-51f584b3c7a1
📒 Files selected for processing (5)
.gitignoretests/test_cluster.pytests/test_leases.pytinyagentos/cluster/manager.pytinyagentos/routes/cluster.py
| worker = self._workers.get(name) | ||
| if worker is None: | ||
| return False | ||
|
|
||
| # Release all active leases for this worker's resources. | ||
| async with self._lease_lock: | ||
| lids: list[str] = [ | ||
| lid for lid, lease in self._leases.items() | ||
| if (parsed := self._parse_resource_id(lease.resource_id)) | ||
| and parsed[0] == name | ||
| ] | ||
| for lid in lids: | ||
| self._leases.pop(lid, None) | ||
| logger.debug( | ||
| "Lease %s released — worker '%s' unregistered", | ||
| lid, name, | ||
| ) | ||
|
|
||
| self._workers.pop(name, None) | ||
| logger.info("Worker '%s' unregistered — %d leases released", name, len(lids)) | ||
| return True |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Locate the relevant file and surrounding symbols.
git ls-files tinyagentos/cluster/manager.py tests/test_cluster.py tests/test_leases.py
printf '\n== manager.py outline ==\n'
ast-grep outline tinyagentos/cluster/manager.py --view expanded || true
printf '\n== unregister_worker implementation ==\n'
sed -n '230,290p' tinyagentos/cluster/manager.py
printf '\n== lease lock usages ==\n'
rg -n "_lease_lock|unregister_worker|claim_lease|_workers\.pop|_workers\.get" tinyagentos/cluster/manager.py
printf '\n== relevant tests ==\n'
sed -n '1,260p' tests/test_cluster.py
printf '\n--- test_leases ---\n'
sed -n '1,260p' tests/test_leases.pyRepository: jaylfc/taOS
Length of output: 25349
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '== worker lock / mutations ==\n'
rg -n "_workers\.|_worker_lock|asyncio\.Lock|register_worker|unregister_worker" tinyagentos/cluster/manager.py tinyagentos/cluster -g '*.py'
printf '\n== monitor loop lease-release block ==\n'
sed -n '520,560p' tinyagentos/cluster/manager.py
printf '\n== API routes touching unregister ==\n'
rg -n "unregister_worker|/workers" tinyagentos -g '*.py'Repository: jaylfc/taOS
Length of output: 8345
Make unregister_worker atomic with the lease release. _workers isn’t protected anywhere else, so two concurrent calls can both pass the pre-check and both report success. Move the existence check and pop() into the same _lease_lock section as the lease cleanup.
🤖 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 255 - 275, Make
`unregister_worker` atomic by moving the `_workers` existence check and `pop()`
into the same `_lease_lock` block used for lease cleanup, so concurrent callers
cannot both pass the pre-check and both report success. Update
`ClusterManager.unregister_worker` to check `self._workers.get(name)` only while
holding `_lease_lock`, release matching leases, remove the worker from
`_workers`, and then log the final lease count from the locked section.
|
Reviewed this in depth, @hognek — the production code is correct and safe to ship: One thing blocks the merge, and it is a merge artifact, not your logic: since your merge-base, To land it: rebase on Optional nits (non-blocking, Kilo agrees): the |
unregister_worker() now releases all active leases for the worker's
resources before removing the worker from the registry. Previously
only the heartbeat-timeout path in _monitor_loop released leases;
explicit unregister (DELETE /api/cluster/workers/{name}) leaked them.
Changes:
- manager.py: unregister_worker() made async, acquires _lease_lock,
releases all leases for the worker (exact name match), then pops
the worker. Logs lease count released.
- routes/cluster.py: await the now-async unregister_worker() call.
- tests: 6 new tests (5 unit + 1 HTTP integration) covering release
on unregister, noop for unknown worker, worker isolation, reclaim
after unregister, and lock serialization.
6fd0bb1 to
e7c8f8d
Compare
|
Rebased onto dev (b8acce7) with merge conflict resolved. Merge artifact fix: Added Optional nit applied: Moved Tests: All 53 tests in |
|
Thanks Hognek, #1726 itself looks good now: the rebase auth-header fix on the four lease calls and moving the The reason I am not merging this yet is the epic gate, not #1726: this lands as a package with #1718 (and the merged #1725 VRAM manager), and #1718 is still on its 2026-07-07 commit, so the core arbiter findings are unaddressed there: submit_gpu still has no production caller (the arbiter ships inert), #1718 carries a parallel VRAM reservation that needs reworking onto the merged VramReservationManager rather than double-counting, and the eviction ordering frees the reservation before the task is cancelled and the VRAM is reclaimed. Once #1718 reworks onto the merged manager, wires the submit_gpu gate-closer, and fixes the eviction ordering, I will review #1718 + #1726 + #1725 together and merge the set. #1726 is ready and waiting on that. |
jaylfc
left a comment
There was a problem hiding this comment.
Reviewed at the code level. The fix mirrors the existing _monitor_loop offline-path lease release exactly (same exact-name match so gpu-node does not over-release gpu-node-2, same _lease_lock, idempotent pop). The earlier concern about _workers.pop sitting outside the lock is resolved in the head commit (release + removal are now atomic under _lease_lock). Conflict-free against current dev, and the new tests exercise the #1705 regression end-to-end (DELETE endpoint drives lease count 2 to 0). Remaining bot notes are suggestion/nitpick tier (fold the None pre-check into the lock, rename the atomicity test, optionally extract a shared _release_leases_for_worker_locked helper shared with _monitor_loop) and are fine as a follow-up. Thanks hognek.
- _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.
…artbeat fixes (#894/#185, continues #1718) (#1859) * feat(scheduler): GPU arbiter — VRAM-accounted admission + queue + eviction 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. * fix(scheduler): preempt running GPU work on eviction, not just cancel 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. * fix(scheduler): clear lease-release ownership in eviction vs. completion _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. * fix(tests): remove test_gpu_arbiter.py — uses #1689-style API incompatible 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. * fix(scheduler): close TOCTOU race in GPU admission with in-flight VRAM 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) * fix(scheduler): move lease claim inside try/finally + subtract reserved 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. * fix(scheduler): non-blocking drain + eviction-to-make-room in GPU arbiter _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. * fix(scheduler): use _reserve_and_check in post-eviction re-admit path 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. * fix(scheduler): add _running_lock coverage to eviction paths in GpuArbiter 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. * fix(scheduler): wire GpuArbiter with real VRAM probe in app.py - 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__. * fix(scheduler): add _running_lock coverage to eviction paths 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. * feat(cluster): worker auto-update with graceful drain, pause, install, 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. * fix(cluster): use exact worker-name compare in lease drain matching 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). * feat(scheduler): pause/resume GPU queue + hardware-aware LLM admission 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. * fix(tests): add missing await to _evict_task calls in test_gpu_arbiter_894 _evict_task was made async in commit 1789ad0 (running-lock fix). Three cleanup calls were missing await. * fix(cluster): make drain_worker/cancel_drain async, hold _lease_lock 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. * fix(cluster): move worker.status='offline' inside _lease_lock; return (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. * fix(scheduler): exact GPU-arch + resource-on-worker match (kilo review) - _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. --------- Co-authored-by: Hogne <227774406+hognek@users.noreply.github.com>
Summary
Fixes taOS #1705: reservation leak — GPU leases were not released when a worker is explicitly unregistered.
The bug
ClusterManager.unregister_worker()simply popped the worker from_workersbut left all its active GPU leases in_leases. This meant the resource slots remained locked indefinitely — new claims on those resources would get 409 "already leased" even though the worker was gone.The heartbeat-timeout path in
_monitor_loopalready handled the crash/disconnect case correctly (marks worker offline → releases leases). But the explicit disconnect path (DELETE /api/cluster/workers/{name}) leaked leases.The fix
manager.py:unregister_worker()made async, acquires_lease_lock, releases all leases for the worker's resources (exact name match via_parse_resource_id), then pops the worker. Logs count of leases released.routes/cluster.py: Route handler nowawaits the async method.Tests added (6 new)
test_unregister_releases_all_leasestest_unregister_unknown_worker_nooptest_unregister_one_worker_spares_anothertest_reclaim_after_unregistertest_unregister_releases_leases_atomically_lease_lockserializes unregister with concurrent claimstest_unregister_worker_releases_leasesTest results
Summary by CodeRabbit
New Features
Bug Fixes