Skip to content

fix(cluster): release GPU leases on worker unregister (fixes #1705)#1726

Merged
jaylfc merged 1 commit into
jaylfc:devfrom
hognek:fix/unregister-worker-lease-leak
Jul 16, 2026
Merged

fix(cluster): release GPU leases on worker unregister (fixes #1705)#1726
jaylfc merged 1 commit into
jaylfc:devfrom
hognek:fix/unregister-worker-lease-leak

Conversation

@hognek

@hognek hognek commented Jul 7, 2026

Copy link
Copy Markdown
Contributor

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 _workers but 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_loop already 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 now awaits the async method.

Tests added (6 new)

Test What it covers
test_unregister_releases_all_leases Unregistering releases all active leases
test_unregister_unknown_worker_noop Unknown worker → False, no lease side effects
test_unregister_one_worker_spares_another Unregistering A doesn't touch B's leases
test_reclaim_after_unregister New claim succeeds on re-registered worker
test_unregister_releases_leases_atomically _lease_lock serializes unregister with concurrent claims
test_unregister_worker_releases_leases HTTP integration: DELETE /workers/{name} releases leases

Test results

tests/test_leases.py ........... 34 passed
tests/test_cluster.py ......... 16 passed
Total: 50 passed

Summary by CodeRabbit

  • New Features

    • Removing a worker now automatically releases all GPU leases assigned to it.
    • Workers can be re-added and claim GPU leases again after removal.
  • Bug Fixes

    • Prevented stale leases from remaining after worker removal.
    • Preserved leases belonging to other workers during removal.
    • Improved handling of repeated removal requests with clear success or not-found responses.

@coderabbitai

coderabbitai Bot commented Jul 7, 2026

Copy link
Copy Markdown

Review Change Stack

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info
⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro Plus

Run ID: c2c58a46-77d5-42a6-a5e4-a52b968d6aef

📥 Commits

Reviewing files that changed from the base of the PR and between 6fd0bb1 and e7c8f8d.

📒 Files selected for processing (4)
  • tests/test_cluster.py
  • tests/test_leases.py
  • tinyagentos/cluster/manager.py
  • tinyagentos/routes/cluster.py
🚧 Files skipped from review as they are similar to previous changes (3)
  • tests/test_cluster.py
  • tinyagentos/routes/cluster.py
  • tinyagentos/cluster/manager.py

📝 Walkthrough

Walkthrough

unregister_worker is now asynchronous and removes the worker’s active leases under _lease_lock. The DELETE route awaits it, and unit/API tests cover cleanup, concurrency, repeated removal, worker isolation, and re-registration.

Changes

Async worker unregistration with lease cleanup

Layer / File(s) Summary
Async unregister and lease cleanup
tinyagentos/cluster/manager.py
unregister_worker becomes async, removes matching leases while holding _lease_lock, removes the worker, and returns whether removal occurred.
Route integration
tinyagentos/routes/cluster.py
The DELETE worker endpoint awaits unregister_worker while retaining its existing response branching.
Behavior and API coverage
tests/test_cluster.py, tests/test_leases.py
Tests verify async results, lease cleanup, worker isolation, re-registration, lock serialization, and DELETE endpoint behavior.

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
Loading

Possibly related PRs

  • jaylfc/taOS#1687: Introduces the broader GPU lease coordination feature used by this lease cleanup logic.
🚥 Pre-merge checks | ✅ 5
✅ Passed checks (5 passed)
Check name Status Explanation
Description Check ✅ Passed Check skipped - CodeRabbit’s high-level summary is enabled.
Title check ✅ Passed The title clearly states the main change: releasing GPU leases when a worker is unregistered.
Docstring Coverage ✅ Passed No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check ✅ Passed Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check ✅ Passed Check skipped because no linked issues were found for this pull request.
✨ Finishing Touches
🧪 Generate unit tests (beta)
  • Create PR with unit tests

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.

❤️ Share

Comment @coderabbitai help to get the list of available commands.

@hognek
hognek marked this pull request as ready for review July 7, 2026 22:35
@qodo-code-review

Copy link
Copy Markdown

Qodo reviews are paused for this user.

Troubleshooting steps vary by plan Learn more →

On a Teams plan?
Reviews resume once this user has a paid seat and their Git account is linked in Qodo.
Link Git account →

Using GitHub Enterprise Server, GitLab Self-Managed, or Bitbucket Data Center?
These require an Enterprise plan - Contact us
Contact us →

@gitar-bot

gitar-bot Bot commented Jul 7, 2026

Copy link
Copy Markdown

Gitar is working

Gitar

)

self._workers.pop(name, None)
logger.info("Worker '%s' unregistered — %d leases released", name, len(lids))

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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 await added inside unregister_worker between these two lines (e.g., for logging, metrics, or a callback) would open a window where a concurrent claim_lease can 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:

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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 False

Reply with @kilocode-bot fix it to have Kilo Code address this issue.

Comment thread tests/test_leases.py
assert lease is not None
assert lease.caller == "second"

async def test_unregister_releases_leases_atomically(self):

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

@kilo-code-bot

kilo-code-bot Bot commented Jul 7, 2026

Copy link
Copy Markdown

Code Review Summary

Status: 2 Issues Found | Recommendation: Address before merge

Overview

Severity Count
CRITICAL 0
WARNING 0
SUGGESTION 2
Issue Details (click to expand)

SUGGESTION

File Line Issue
tinyagentos/cluster/manager.py 271 worker = self._workers.get(name) is only read to test for None; a direct membership check (if name not in self._workers: return False) is clearer.
tests/test_leases.py 640 test_unregister_releases_leases_atomically only verifies that unregister blocks on _lease_lock from outside; it does not race a concurrent claim_lease. Rename to reflect what it tests or strengthen it.
Files Reviewed (4 files)
  • tinyagentos/cluster/manager.py - 1 issue
  • tinyagentos/routes/cluster.py - 0 issues
  • tests/test_cluster.py - 0 issues
  • tests/test_leases.py - 1 issue

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

Severity Count
CRITICAL 0
WARNING 1
SUGGESTION 2
Issue Details (click to expand)

WARNING

File Line Issue
tinyagentos/cluster/manager.py 274 self._workers.pop(name, None) runs after the _lease_lock block exits. No await sits between them today, so pure-asyncio preemption doesn't fire here, but any future await (logging, metrics, callback) between the lock release and the pop would let a concurrent claim_lease create a lease that is then leaked — the exact bug this PR fixes. Move the pop inside the async with self._lease_lock: block.

SUGGESTION

File Line Issue
tinyagentos/cluster/manager.py 256 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 False).
tests/test_leases.py 583 test_unregister_releases_leases_atomically only verifies that unregister blocks on _lease_lock from outside; it does not actually race a concurrent claim_lease against unregister_worker to verify end-to-end atomicity. Either rename to reflect what it tests (test_unregister_blocks_on_lease_lock) or strengthen it to spawn a concurrent claim task and assert the post-state is consistent.
Files Reviewed (5 files)
  • .gitignore - 0 issues
  • tests/test_cluster.py - 0 issues
  • tests/test_leases.py - 1 issue
  • tinyagentos/cluster/manager.py - 2 issues
  • tinyagentos/routes/cluster.py - 0 issues

Fix these issues in Kilo Cloud


Reviewed by hy3-20260706:free · Input: 78.3K · Output: 10.2K · Cached: 342.4K

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Actionable comments posted: 1

🧹 Nitpick comments (1)
tinyagentos/cluster/manager.py (1)

259-271: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Duplicate 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 both unregister_worker and _monitor_loop while 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

📥 Commits

Reviewing files that changed from the base of the PR and between 0cff19c and 6fd0bb1.

📒 Files selected for processing (5)
  • .gitignore
  • tests/test_cluster.py
  • tests/test_leases.py
  • tinyagentos/cluster/manager.py
  • tinyagentos/routes/cluster.py

Comment on lines +255 to +275
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

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🩺 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.py

Repository: 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.

@jaylfc

jaylfc commented Jul 9, 2026

Copy link
Copy Markdown
Owner

Reviewed this in depth, @hognek — the production code is correct and safe to ship: unregister_worker releases leases under _lease_lock (matching the monitor offline path), it is idempotent (unknown worker → False, double-unregister → False, pop(lid, None) cannot KeyError), and the DELETE route correctly awaits the now-async call. The 4 ClusterManager unit tests genuinely exercise the path. Nice fix.

One thing blocks the merge, and it is a merge artifact, not your logic: since your merge-base, dev added auth enforcement on the lease endpoints (_require_lease_access, routes/cluster.py:1029) and an _auth(app) test helper, adding headers=_auth(app) to every lease call in tests/test_leases.py. That is the textual conflict, and it also means your new integration test test_unregister_worker_releases_leases will start returning 401/403 post-merge — it calls POST /api/cluster/leases/claim (x2) and GET /api/cluster/leases without auth headers. CI is green today only because the auth gate lives on dev, not in your branch.

To land it: rebase on dev, keep dev's _auth helper + header edits, and add headers=_auth(app) to the two leases/claim POSTs and the GET /leases in your new test. Then it is a clean merge and I will take it.

Optional nits (non-blocking, Kilo agrees): the self._workers.pop(name, None) runs just after the async with self._lease_lock block — safe today (no await between release and pop) but you could move it inside the lock to keep it future-proof; and test_unregister_releases_leases_atomically proves the call blocks on the lock but does not race a real concurrent claim. Both optional.

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.
@hognek
hognek force-pushed the fix/unregister-worker-lease-leak branch from 6fd0bb1 to e7c8f8d Compare July 10, 2026 21:30
@hognek

hognek commented Jul 10, 2026

Copy link
Copy Markdown
Contributor Author

Rebased onto dev (b8acce7) with merge conflict resolved.

Merge artifact fix: Added headers=_auth(app) to all four lease API calls in test_unregister_worker_releases_leases (two POST /leases/claim + two GET /leases). This keeps the test passing post-rebase against dev's auth gate (PR #1711).

Optional nit applied: Moved self._workers.pop(name, None) inside the async with self._lease_lock: block for future-proof atomicity.

Tests: All 53 tests in tests/test_leases.py + tests/test_cluster.py pass locally.

@jaylfc

jaylfc commented Jul 10, 2026

Copy link
Copy Markdown
Owner

Thanks Hognek, #1726 itself looks good now: the rebase auth-header fix on the four lease calls and moving the pop inside _lease_lock are exactly right, and Kilo is down to two non-blocking suggestions (the membership-check readability at manager.py:271 and the test name/strength at test_leases.py:640). Fold those if you like, but they are not blockers.

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 jaylfc left a comment

Copy link
Copy Markdown
Owner

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

@jaylfc
jaylfc merged commit 2d72bfb into jaylfc:dev Jul 16, 2026
4 checks passed
jaylfc added a commit that referenced this pull request Jul 16, 2026
- _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.
jaylfc added a commit that referenced this pull request Jul 16, 2026
…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>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants