feat(scheduler): GPU arbiter with drain→arbiter wiring (taOS #1707)#1727
feat(scheduler): GPU arbiter with drain→arbiter wiring (taOS #1707)#1727hognek wants to merge 1 commit into
Conversation
) - Add GpuArbiter: VRAM-accounted admission, queuing, priority-based eviction - drain→arbiter coordination: arbiter notifies scheduler drain before eviction via drain_notify_fn, waits for on_drain_complete callback before proceeding - Add drain_worker/cancel_drain to ClusterManager with arbiter cross-wiring (graceful and force drain modes, lease release + task cancellation) - Exclude draining workers from capability routing - Monitor loop auto-completes graceful drains when no active leases remain - 34 new tests: 25 gpu_arbiter + 9 cluster drain - All existing cluster/lease/scheduler tests continue to pass (63 total)
📝 WalkthroughWalkthroughIntroduces a ChangesGPU Arbiter and Cluster Draining
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant Caller
participant ClusterManager
participant GpuArbiter
participant Scheduler
Caller->>ClusterManager: drain_worker(name, graceful=False)
ClusterManager->>ClusterManager: mark worker "draining"
ClusterManager->>ClusterManager: release active leases
ClusterManager->>GpuArbiter: cancel_running_for_leases(lease_ids)
GpuArbiter->>Scheduler: _notify_drain_and_wait(model_id)
Scheduler-->>GpuArbiter: on_drain_complete(model_id)
GpuArbiter->>GpuArbiter: release VRAM reservation, cancel task
GpuArbiter-->>ClusterManager: (cancelled, done) counts
ClusterManager->>ClusterManager: mark worker "offline"
ClusterManager-->>Caller: drain status
sequenceDiagram
participant Caller
participant GpuArbiter
participant Queue
Caller->>GpuArbiter: submit_gpu(task, required_vram_mb)
GpuArbiter->>GpuArbiter: _reserve_and_check(vram)
alt insufficient VRAM
GpuArbiter->>Queue: enqueue _QueuedGpuTask
GpuArbiter-->>Caller: pending future
Queue->>GpuArbiter: _drain_queue admits task
GpuArbiter->>GpuArbiter: evict_lowest_priority (if needed)
else admitted
GpuArbiter->>GpuArbiter: _run_gpu_task(task)
end
GpuArbiter-->>Caller: task result
Possibly related PRs
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 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? |
![kilo-code-bot[bot]](https://avatars.githubusercontent.com/in/2193792?s=60&v=4){kind=small}
| if l.resource_id.startswith(worker.name + ":") | ||
| and l.required_vram_mb > 0 | ||
| ) | ||
| available = worker.free_vram_mb - worker_leases - self._reserved_vram_mb |
There was a problem hiding this comment.
CRITICAL: worker.free_vram_mb can be None, causing a TypeError crash.
WorkerInfo.free_vram_mb defaults to None for workers that never reported VRAM (non-NVIDIA, Apple Silicon, RK3588, CPU-only, or simply unreported — see worker_protocol.py:69). With the default probe _default_vram_probe() returning (0, 0), the local-VRAM branch is skipped (free_vram > 0 is False) and every submit falls into this cluster path, hitting this subtraction. Any None worker raises TypeError: unsupported operand type(s) for -: 'NoneType' and 'int' on essentially every admission check in a default or heterogeneous cluster — making the arbiter unusable.
| available = worker.free_vram_mb - worker_leases - self._reserved_vram_mb | |
| available = (worker.free_vram_mb or 0) - worker_leases - self._reserved_vram_mb |
For a stricter fix, continue when worker.free_vram_mb is None, since the available VRAM is genuinely unknown for that worker.
Reply with @kilocode-bot fix it to have Kilo Code address this issue.
| while True: | ||
| await asyncio.sleep(2) | ||
| if not self._paused: | ||
| await self._drain_queue() |
There was a problem hiding this comment.
WARNING: _process_queue has no exception handling around _drain_queue.
Any exception raised inside _drain_queue (including the None-VRAM TypeError above) propagates out of this loop and permanently kills the background queue-processor task. There is no restart path, so every queued submit_gpu future that is awaiting its _arbiter_future will hang forever, and no further queued tasks will ever be admitted. Wrap the drain call in try/except (log + continue) so one bad cycle cannot take down admission for the whole process.
Reply with @kilocode-bot fix it to have Kilo Code address this issue.
| # ── drain→arbiter wiring (taOS #1707) ────────────────────── | ||
| # Notify the scheduler to drain this model and wait for completion. | ||
| model_id = getattr(task, "model_id", None) or task.id | ||
| await self._notify_drain_and_wait(model_id) |
There was a problem hiding this comment.
WARNING: Force-drain cancellation blocks up to drain_timeout (default 60s) per task.
cancel_running_for_leases (the force-drain path from ClusterManager.drain_worker(graceful=False)) calls _evict_task, which unconditionally awaits _notify_drain_and_wait. That waits up to drain_timeout for an on_drain_complete callback that the force-drain flow never issues for these tasks. So a "force"/immediate drain can block drain_worker for up to 60s * N tasks, and additionally re-invokes drain_notify_fn for models the scheduler did not ask to drain. The force path should bypass drain coordination (e.g. a _evict_task(task_id, notify_drain=False) flag) and cancel immediately.
Reply with @kilocode-bot fix it to have Kilo Code address this issue.
| await self._notify_drain_and_wait(model_id) | ||
|
|
||
| # Release VRAM reservation | ||
| self._release_reservation(task_id) |
There was a problem hiding this comment.
WARNING: VRAM reservation is released before the running task is actually cancelled — re-opens the TOCTOU over-subscription.
_evict_task pops the task from _running, then awaits the up-to-60s drain wait, then releases the reservation (here) and only afterwards calls asyncio_task.cancel(). Until the task is actually cancelled it is still consuming VRAM, yet the reservation is already freed and running count already decremented. A concurrent submit_gpu/_drain_queue can therefore admit a new task into VRAM that is still in use — exactly the TOCTOU race the PR claims to fix. Keep the reservation held (or at least don't decrement accounting) until the underlying task has truly stopped.
Reply with @kilocode-bot fix it to have Kilo Code address this issue.
| await self._queue_processor_task | ||
| except asyncio.CancelledError: | ||
| pass | ||
| self._queue_processor_task = None |
There was a problem hiding this comment.
WARNING: stop() cancels the queue processor but never resolves pending _arbiter_futures.
If any submit_gpu calls are awaiting admission in the queue when stop() is called, their futures are never set or cancelled. The awaiting coroutines hang forever (until GC), which can wedge callers/tests. On stop, drain the queue and resolve each pending future with NoResourceAvailableError (or CancelledError) so no caller is left suspended.
Reply with @kilocode-bot fix it to have Kilo Code address this issue.
Code Review SummaryStatus: 5 Issues Found | Recommendation: Address before merge Overview
Issue Details (click to expand)CRITICAL
WARNING
Files Reviewed (5 files)
Fix these issues in Kilo Cloud Reviewed by hy3-20260706:free · Input: 96.7K · Output: 15.6K · Cached: 425.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: 3
🧹 Nitpick comments (5)
tests/test_cluster.py (2)
364-391: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick winNo test for arbiter-cancel exception path.
drain_workerwraps thecancel_running_for_leasescall in try/except and swallows exceptions (per the providedmanager.pysnippet). Consider adding a test where the mocked arbiter raises, to confirmdrain_workerstill returns a validreleased_leasesresult rather than propagating/crashing.🤖 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_cluster.py` around lines 364 - 391, Add a test for the exception path in drain_worker: the current test only covers a successful arbiter cancel in ClusterManager.drain_worker. Create a case where mock_arbiter.cancel_running_for_leases raises an exception, then verify drain_worker("gpu-box", graceful=False) still completes and returns a valid result with released_leases set correctly instead of propagating the error. Reuse the existing ClusterManager, _make_worker, and mock_arbiter setup so the new test stays close to the current one.
364-391: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueMove
AsyncMockimport to top-level for consistency.
AsyncMockis imported locally inside this test whileMagicMockis presumably imported at module scope for the other tests. Minor stylistic inconsistency.♻️ Suggested cleanup
- async def test_drain_force_calls_arbiter_cancel(self): - from unittest.mock import AsyncMock - - mgr = ClusterManager() + async def test_drain_force_calls_arbiter_cancel(self): + mgr = ClusterManager()(and add
AsyncMockto the top-of-fileunittest.mockimport)Also, the assertion only checks
len(call_args) == 2rather than the actual lease IDs of the two claimed leases — asserting the exact set would make the test stricter against regressions in howlidsis constructed indrain_worker.🤖 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_cluster.py` around lines 364 - 391, The test `test_drain_force_calls_arbiter_cancel` should use the same top-level `unittest.mock` imports as the rest of the module by moving `AsyncMock` out of the function and adding it alongside `MagicMock`. Also strengthen the `drain_worker` assertion by checking the actual lease IDs passed to `mock_arbiter.cancel_running_for_leases` instead of only `len(call_args)`, so the test verifies the exact lids collected from the two `claim_lease` calls.tests/test_gpu_arbiter.py (1)
89-115: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winOptional: add a happy-path integration test for queued-task admission.
All admission tests intentionally avoid starting the queue processor (comment at Line 92-93), so there's no coverage of a task queuing, then getting admitted once VRAM frees up (via
start()/eviction/retry). Consider adding one such test in a follow-up to validate the fullsubmit_gpu→ queue → dequeue → completion path.🤖 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_gpu_arbiter.py` around lines 89 - 115, Add a follow-up happy-path integration test for GpuArbiter that covers a task being queued and later admitted once VRAM becomes available. Use the existing GpuArbiter, submit_gpu, stats, and any start()/queue-processing helpers to exercise the full submit_gpu → queue → dequeue → completion flow, since the current tests only verify queuing and queue-full behavior without running the queue processor.tinyagentos/scheduler/gpu_arbiter.py (1)
580-585: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick winQueue drain latency scales with backlog: fixed 2s poll + one admission per cycle.
_process_queuesleeps 2s unconditionally, and_drain_queuesetsdrained=Trueafter admitting a single task, so a backlog of N queued tasks takes ~2·N seconds to clear even when VRAM frees up immediately. Consider signaling the processor (e.g., anasyncio.Eventset on submit/completion) to wake promptly instead of relying solely on the fixed interval, while keeping the one-task-per-admit throttle.🤖 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 580 - 585, The queue processor in _process_queue and the admission path in _drain_queue are too dependent on a fixed 2-second sleep, which makes backlog clearing scale linearly with queue size. Update the scheduler flow to wake the processor promptly using a signal such as an asyncio.Event that is set when work is submitted or when capacity becomes available, and have _process_queue wait on that signal instead of only polling. Keep the existing one-task-per-admit behavior in _drain_queue, but ensure _process_queue is triggered immediately on queue/completion events so multiple queued tasks can drain without waiting for the next sleep cycle.tinyagentos/cluster/manager.py (1)
325-340: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick winHold a strong reference to the fire-and-forget notification task.
asynciokeeps only weak references to tasks created viacreate_task, so thisworker.drainemit can be garbage-collected mid-flight before it completes. Your__init__already maintainsself._background_tasksfor exactly this reason (see the strong-reference comment and its use inregister_worker). Route this emit through the same set.♻️ Proposed fix
try: - asyncio.get_running_loop().create_task( - self._notifications.emit_event( - "worker.drain", - f"Worker '{name}' draining", - detail, - level="info", - ) - ) + _t = asyncio.get_running_loop().create_task( + self._notifications.emit_event( + "worker.drain", + f"Worker '{name}' draining", + detail, + level="info", + ) + ) + self._background_tasks.add(_t) + _t.add_done_callback(self._background_tasks.discard) except RuntimeError: passNote:
heartbeat()at Line 234 uses the same unguarded pattern; consider aligning it too.🤖 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 325 - 340, The fire-and-forget event emission in `manager.py` is creating a task without keeping a strong reference, so the `worker.drain` notification can be garbage-collected before it finishes. Update the `worker.drain` path in `Manager` to store the task in `self._background_tasks`, matching the pattern already used in `register_worker` and the strong-reference handling in `__init__`. Also check `heartbeat()` for the same `create_task` pattern and align it the same way if it is intended to run detached.
🤖 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/scheduler/gpu_arbiter.py`:
- Around line 36-53: The `_probe_nvidia_vram` helper calls `nvidia-smi` by bare
name, which relies on PATH lookup; update it to use an absolute or resolved
executable path instead. Locate the `subprocess.run` calls in
`_probe_nvidia_vram` and resolve `nvidia-smi` explicitly (for example via
`shutil.which` with a fallback failure path) before invoking it, while keeping
the argv static and not using any incoming input.
- Around line 517-535: `evict_lowest_priority()` is currently choosing victims
from `self._running` without considering the task’s `evictable` flag, so
non-evictable tasks can still be cancelled. Persist `evictable` in the
running-state record when tasks are moved into `self._running`, then update
`evict_lowest_priority()` to skip any non-evictable entries while selecting the
victim; keep `cancel_running_for_leases()` unchanged as the force-drain path.
- Around line 410-431: The cluster fallback in gpu_arbiter.py is applying the
global _reserved_vram_mb to every worker, which makes admission overly
coarse-grained; update the worker loop in the GPU admission logic so
reservations are evaluated per-worker if that is the intended policy, or add a
clear comment/documentation if the global reservation is deliberate. Also rename
the loop variable used in the lease filter from l to lease in the
get_leases/get_workers admission path to satisfy the Ruff E741 warning.
---
Nitpick comments:
In `@tests/test_cluster.py`:
- Around line 364-391: Add a test for the exception path in drain_worker: the
current test only covers a successful arbiter cancel in
ClusterManager.drain_worker. Create a case where
mock_arbiter.cancel_running_for_leases raises an exception, then verify
drain_worker("gpu-box", graceful=False) still completes and returns a valid
result with released_leases set correctly instead of propagating the error.
Reuse the existing ClusterManager, _make_worker, and mock_arbiter setup so the
new test stays close to the current one.
- Around line 364-391: The test `test_drain_force_calls_arbiter_cancel` should
use the same top-level `unittest.mock` imports as the rest of the module by
moving `AsyncMock` out of the function and adding it alongside `MagicMock`. Also
strengthen the `drain_worker` assertion by checking the actual lease IDs passed
to `mock_arbiter.cancel_running_for_leases` instead of only `len(call_args)`, so
the test verifies the exact lids collected from the two `claim_lease` calls.
In `@tests/test_gpu_arbiter.py`:
- Around line 89-115: Add a follow-up happy-path integration test for GpuArbiter
that covers a task being queued and later admitted once VRAM becomes available.
Use the existing GpuArbiter, submit_gpu, stats, and any start()/queue-processing
helpers to exercise the full submit_gpu → queue → dequeue → completion flow,
since the current tests only verify queuing and queue-full behavior without
running the queue processor.
In `@tinyagentos/cluster/manager.py`:
- Around line 325-340: The fire-and-forget event emission in `manager.py` is
creating a task without keeping a strong reference, so the `worker.drain`
notification can be garbage-collected before it finishes. Update the
`worker.drain` path in `Manager` to store the task in `self._background_tasks`,
matching the pattern already used in `register_worker` and the strong-reference
handling in `__init__`. Also check `heartbeat()` for the same `create_task`
pattern and align it the same way if it is intended to run detached.
In `@tinyagentos/scheduler/gpu_arbiter.py`:
- Around line 580-585: The queue processor in _process_queue and the admission
path in _drain_queue are too dependent on a fixed 2-second sleep, which makes
backlog clearing scale linearly with queue size. Update the scheduler flow to
wake the processor promptly using a signal such as an asyncio.Event that is set
when work is submitted or when capacity becomes available, and have
_process_queue wait on that signal instead of only polling. Keep the existing
one-task-per-admit behavior in _drain_queue, but ensure _process_queue is
triggered immediately on queue/completion events so multiple queued tasks can
drain without waiting for the next sleep cycle.
🪄 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: 0462f2d3-bddb-4783-ab73-dec22200e51e
📒 Files selected for processing (5)
tests/test_cluster.pytests/test_gpu_arbiter.pytinyagentos/cluster/manager.pytinyagentos/scheduler/__init__.pytinyagentos/scheduler/gpu_arbiter.py
| def _probe_nvidia_vram() -> tuple[int, int]: | ||
| """Probe free/total VRAM from nvidia-smi. Returns (free_mb, total_mb).""" | ||
| try: | ||
| import subprocess | ||
| free_raw = subprocess.run( | ||
| ["nvidia-smi", "--query-gpu=memory.free", "--format=csv,noheader,nounits"], | ||
| capture_output=True, text=True, timeout=5, | ||
| ) | ||
| total_raw = subprocess.run( | ||
| ["nvidia-smi", "--query-gpu=memory.total", "--format=csv,noheader,nounits"], | ||
| capture_output=True, text=True, timeout=5, | ||
| ) | ||
| return ( | ||
| int(free_raw.stdout.strip().split("\n")[0]), | ||
| int(total_raw.stdout.strip().split("\n")[0]), | ||
| ) | ||
| except Exception: | ||
| return 0, 0 |
There was a problem hiding this comment.
🔒 Security & Privacy | 🟡 Minor | ⚡ Quick win
Use an absolute/resolved path for nvidia-smi; the injection SAST alert is a false positive.
nvidia-smi is invoked by bare name, so resolution depends on PATH (Ruff S607). Resolve it explicitly (or via shutil.which) to harden against PATH hijacking. Note the ast-grep "command from incoming request" finding is a false positive here — the argv is fully static with no external input.
🛡️ Suggested hardening
def _probe_nvidia_vram() -> tuple[int, int]:
"""Probe free/total VRAM from nvidia-smi. Returns (free_mb, total_mb)."""
try:
import subprocess
+ import shutil
+ smi = shutil.which("nvidia-smi")
+ if smi is None:
+ return 0, 0
free_raw = subprocess.run(
- ["nvidia-smi", "--query-gpu=memory.free", "--format=csv,noheader,nounits"],
+ [smi, "--query-gpu=memory.free", "--format=csv,noheader,nounits"],
capture_output=True, text=True, timeout=5,
)
total_raw = subprocess.run(
- ["nvidia-smi", "--query-gpu=memory.total", "--format=csv,noheader,nounits"],
+ [smi, "--query-gpu=memory.total", "--format=csv,noheader,nounits"],
capture_output=True, text=True, timeout=5,
)📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| def _probe_nvidia_vram() -> tuple[int, int]: | |
| """Probe free/total VRAM from nvidia-smi. Returns (free_mb, total_mb).""" | |
| try: | |
| import subprocess | |
| free_raw = subprocess.run( | |
| ["nvidia-smi", "--query-gpu=memory.free", "--format=csv,noheader,nounits"], | |
| capture_output=True, text=True, timeout=5, | |
| ) | |
| total_raw = subprocess.run( | |
| ["nvidia-smi", "--query-gpu=memory.total", "--format=csv,noheader,nounits"], | |
| capture_output=True, text=True, timeout=5, | |
| ) | |
| return ( | |
| int(free_raw.stdout.strip().split("\n")[0]), | |
| int(total_raw.stdout.strip().split("\n")[0]), | |
| ) | |
| except Exception: | |
| return 0, 0 | |
| def _probe_nvidia_vram() -> tuple[int, int]: | |
| """Probe free/total VRAM from nvidia-smi. Returns (free_mb, total_mb).""" | |
| try: | |
| import subprocess | |
| import shutil | |
| smi = shutil.which("nvidia-smi") | |
| if smi is None: | |
| return 0, 0 | |
| free_raw = subprocess.run( | |
| [smi, "--query-gpu=memory.free", "--format=csv,noheader,nounits"], | |
| capture_output=True, text=True, timeout=5, | |
| ) | |
| total_raw = subprocess.run( | |
| [smi, "--query-gpu=memory.total", "--format=csv,noheader,nounits"], | |
| capture_output=True, text=True, timeout=5, | |
| ) | |
| return ( | |
| int(free_raw.stdout.strip().split("\n")[0]), | |
| int(total_raw.stdout.strip().split("\n")[0]), | |
| ) | |
| except Exception: | |
| return 0, 0 |
🧰 Tools
🪛 ast-grep (0.44.1)
[error] 39-42: Command coming from incoming request
Context: subprocess.run(
["nvidia-smi", "--query-gpu=memory.free", "--format=csv,noheader,nounits"],
capture_output=True, text=True, timeout=5,
)
Note: [CWE-78] Improper Neutralization of Special Elements used in an OS Command ('OS Command Injection').
(subprocess-from-request)
[error] 43-46: Command coming from incoming request
Context: subprocess.run(
["nvidia-smi", "--query-gpu=memory.total", "--format=csv,noheader,nounits"],
capture_output=True, text=True, timeout=5,
)
Note: [CWE-78] Improper Neutralization of Special Elements used in an OS Command ('OS Command Injection').
(subprocess-from-request)
🪛 Ruff (0.15.20)
[error] 41-41: Starting a process with a partial executable path
(S607)
[error] 45-45: Starting a process with a partial executable path
(S607)
[warning] 52-52: Do not catch blind exception: Exception
(BLE001)
🤖 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 36 - 53, The
`_probe_nvidia_vram` helper calls `nvidia-smi` by bare name, which relies on
PATH lookup; update it to use an absolute or resolved executable path instead.
Locate the `subprocess.run` calls in `_probe_nvidia_vram` and resolve
`nvidia-smi` explicitly (for example via `shutil.which` with a fallback failure
path) before invoking it, while keeping the argv static and not using any
incoming input.
Source: Linters/SAST tools
| # No local VRAM probe — check cluster workers | ||
| if self._cluster_manager is not None: | ||
| leases = self._cluster_manager.get_leases() | ||
| for worker in self._cluster_manager.get_workers(): | ||
| if worker.status != "online": | ||
| continue | ||
| worker_leases = sum( | ||
| l.required_vram_mb for l in leases | ||
| if l.resource_id.startswith(worker.name + ":") | ||
| and l.required_vram_mb > 0 | ||
| ) | ||
| available = worker.free_vram_mb - worker_leases - self._reserved_vram_mb | ||
| if available >= required_vram_mb: | ||
| return GpuAdmission( | ||
| admitted=True, | ||
| free_vram_mb=available, | ||
| required_vram_mb=required_vram_mb, | ||
| ) | ||
| return GpuAdmission( | ||
| admitted=False, required_vram_mb=required_vram_mb, | ||
| reason=f"no cluster worker with {required_vram_mb} MiB free VRAM", | ||
| ) |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Map the file structure first.
ast-grep outline tinyagentos/scheduler/gpu_arbiter.py --view expanded
# Inspect the relevant section around the cited lines.
sed -n '340,460p' tinyagentos/scheduler/gpu_arbiter.py
# Find all references to _reserved_vram_mb and lease handling.
rg -n "_reserved_vram_mb|get_leases|get_workers|required_vram_mb|free_vram_mb" tinyagentos/scheduler/gpu_arbiter.py
# Show the local admission path and any reservation accounting nearby.
ast-grep run --lang python --pattern $'if self._cluster_manager is None:\n $$$\n' tinyagentos/scheduler/gpu_arbiter.pyRepository: jaylfc/taOS
Length of output: 10293
🏁 Script executed:
#!/bin/bash
set -euo pipefail
sed -n '300,470p' tinyagentos/scheduler/gpu_arbiter.py
printf '\n--- reservation references ---\n'
rg -n "_reserved_vram_mb|_pending_reservations|_release_reservation|_reserve_and_check|claim_lease|worker_leases" tinyagentos/scheduler/gpu_arbiter.py
printf '\n--- any cluster-manager docs/usages ---\n'
rg -n "cluster_manager|claim_lease|get_leases|get_workers|free_vram_mb" tinyagentos -g '!**/__pycache__/**'Repository: jaylfc/taOS
Length of output: 20231
🏁 Script executed:
#!/bin/bash
set -euo pipefail
sed -n '420,520p' tinyagentos/cluster/manager.py
printf '\n--- worker capacity / free_vram semantics ---\n'
sed -n '90,140p' tinyagentos/cluster/worker_capacity.py
printf '\n--- worker info shape ---\n'
sed -n '1,120p' tinyagentos/cluster/worker_protocol.pyRepository: jaylfc/taOS
Length of output: 10791
🏁 Script executed:
#!/bin/bash
set -euo pipefail
sed -n '470,640p' tinyagentos/scheduler/gpu_arbiter.py
printf '\n--- queue / drain / eviction references ---\n'
rg -n "_release_reservation|_reserve_and_check|_check_admission|resource_id|claim_lease|free_vram_mb" tinyagentos/scheduler/gpu_arbiter.pyRepository: jaylfc/taOS
Length of output: 8724
🏁 Script executed:
#!/bin/bash
set -euo pipefail
rg -n "submit_gpu\(" tinyagentos -g '!**/__pycache__/**'
printf '\n--- resource_id call sites ---\n'
rg -n "resource_id\s*=" tinyagentos -g '!**/__pycache__/**'Repository: jaylfc/taOS
Length of output: 729
🏁 Script executed:
#!/bin/bash
set -euo pipefail
sed -n '160,260p' tinyagentos/cluster/manager.py
printf '\n--- heartbeat / worker updates ---\n'
rg -n "free_vram_mb|used_vram_mb|heartbeat|update_worker|register_worker|report" tinyagentos/cluster/manager.py tinyagentos/worker -g '!**/__pycache__/**'Repository: jaylfc/taOS
Length of output: 9076
Cluster admission is too coarse-grained
_reserved_vram_mbis applied to every candidate worker in the cluster fallback, so unrelated in-flight reservations can make admission fail even when one worker still has room. If that’s intentional, document the policy; otherwise make the check worker-specific.- Rename
ltoleaseto avoid the Ruff E741 nit.
🧰 Tools
🪛 Ruff (0.15.20)
[error] 417-417: Ambiguous variable name: l
(E741)
🤖 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 410 - 431, The cluster
fallback in gpu_arbiter.py is applying the global _reserved_vram_mb to every
worker, which makes admission overly coarse-grained; update the worker loop in
the GPU admission logic so reservations are evaluated per-worker if that is the
intended policy, or add a clear comment/documentation if the global reservation
is deliberate. Also rename the loop variable used in the lease filter from l to
lease in the get_leases/get_workers admission path to satisfy the Ruff E741
warning.
Source: Linters/SAST tools
| async def evict_lowest_priority(self, min_priority: int | None = None) -> int: | ||
| """Evict the running task with the highest numeric priority value. | ||
|
|
||
| Higher numeric priority = lower actual priority = evicted first. | ||
| If *min_priority* is given, only tasks with priority >= min_priority | ||
| are considered. | ||
| """ | ||
| if not self._eviction_enabled: | ||
| return 0 | ||
| async with self._running_lock: | ||
| victim_id, victim_priority = None, -1 | ||
| for tid, (_t, _lid, pri, _vram) in self._running.items(): | ||
| if min_priority is not None and pri < min_priority: | ||
| continue | ||
| if pri > victim_priority: | ||
| victim_priority, victim_id = pri, tid | ||
| if victim_id is None: | ||
| return 0 | ||
| return await self._evict_task(victim_id) |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
set -euo pipefail
echo '--- file outline ---'
ast-grep outline tinyagentos/scheduler/gpu_arbiter.py --view expanded || true
echo '--- search self._running and evictable ---'
rg -n "self\._running|evictable|cancel_running_for_leases|evict_lowest_priority|_evict_task|running_tasks" tinyagentos/scheduler/gpu_arbiter.py
echo '--- relevant slices ---'
sed -n '430,620p' tinyagentos/scheduler/gpu_arbiter.pyRepository: jaylfc/taOS
Length of output: 11949
Honor evictable when selecting eviction victims
evictable is captured on the queued task but dropped from self._running, so evict_lowest_priority() can still cancel default non-evictable tasks under VRAM pressure. Persist that flag in the running state and skip ineligible tasks there; cancel_running_for_leases() can stay as the force-drain path.
🤖 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 517 - 535,
`evict_lowest_priority()` is currently choosing victims from `self._running`
without considering the task’s `evictable` flag, so non-evictable tasks can
still be cancelled. Persist `evictable` in the running-state record when tasks
are moved into `self._running`, then update `evict_lowest_priority()` to skip
any non-evictable entries while selecting the victim; keep
`cancel_running_for_leases()` unchanged as the force-drain path.
|
Closing as requested — superseded by #1718 (the canonical arbiter superset composing #1683/#1705/#1706/#1707/#1690 with the blocker fixes, clean Kilo). #1727 competes on the same gpu_arbiter.py, drops the app.py set_gpu_arbiter wiring and the #1690 worker-update routes, and carries the CRITICALs you flagged (:421 free_vram_mb None → TypeError, :585 bare-except, :190 stop() leaving _arbiter_futures hung). Converging on #1718. |
1 participant
Fixes #1707.
Summary
Adds the GPU arbiter module with admission control, queuing, and priority-based eviction, plus the drain→arbiter coordination wiring that was the focus of #1707.
Changes
New:
tinyagentos/scheduler/gpu_arbiter.pydrain_notify_fncallback — arbiter calls this when deciding to evict, so the scheduler can stop routing new workon_drain_complete(model_id)— scheduler calls this when draining is done, arbiter then proceeds with evictiondrain_timeoutcancel_running_for_leases()— called from cluster drain paths to cancel GPU tasks for force-released leasesModified:
tinyagentos/cluster/manager.pydrain_worker(name, graceful=True): graceful (leases run to completion) and force (immediate lease release + arbiter task cancellation) modescancel_drain(name): return a draining worker to onlineset_gpu_arbiter(arbiter): cross-wiring hook for the application wiring layerget_workers_for_capability()Modified:
tinyagentos/scheduler/__init__.pyGpuArbiterandGpuAdmissionTests
tests/test_gpu_arbiter.py: admission, eviction, drain wiring (notification + callback), lease cancellation, pause/resume, statstests/test_cluster.py: drain_worker (graceful + force), cancel_drain, routing exclusion, arbiter cross-wiringSummary by CodeRabbit
New Features
Bug Fixes
Tests