feat(cluster): worker auto-update with graceful drain, pause, install, restart (taOS #890)#1690
feat(cluster): worker auto-update with graceful drain, pause, install, restart (taOS #890)#1690hognek wants to merge 4 commits into
Conversation
…, restart Adds drain_worker/cancel_drain to ClusterManager for graceful worker detach without dropping inflight GPU tasks (taOS jaylfc#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.
📝 WalkthroughWalkthroughAdds graceful and forced worker draining, excludes draining workers from cluster selection, auto-completes drains in the monitor loop, exposes drain and update endpoints, and tests lifecycle and deployment failure behavior. ChangesWorker Drain Lifecycle
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant Client
participant ClusterRoute
participant ClusterManager
participant RemoteWorker
participant MonitorLoop
Client->>ClusterRoute: POST worker update
ClusterRoute->>ClusterManager: drain_worker
ClusterManager-->>ClusterRoute: draining status
ClusterRoute->>RemoteWorker: deploy update-worker
alt deploy failure
RemoteWorker-->>ClusterRoute: error
ClusterRoute->>ClusterManager: cancel_drain
ClusterManager-->>ClusterRoute: online status
ClusterRoute-->>Client: 502 failure
else deploy succeeds or restart disconnect
RemoteWorker-->>ClusterRoute: deployment result
ClusterRoute-->>Client: update response
MonitorLoop->>ClusterManager: inspect draining worker
ClusterManager-->>MonitorLoop: complete or timeout drain
end
Possibly related issues
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? |
![kilo-code-bot[bot]](https://avatars.githubusercontent.com/in/2193792?s=60&v=4){kind=small}
| # Force-release all leases for this worker | ||
| lids = [ | ||
| lid for lid, lease in self._leases.items() | ||
| if lease.resource_id.startswith(name + ":") |
There was a problem hiding this comment.
CRITICAL: Worker name prefix-match collision
lease.resource_id.startswith(name + ":") reintroduces the exact bug the existing monitor-loop code (lines 604-610) was fixed to avoid — see the comment immediately above this hunk: "Match on the exact worker name (not a resource_id prefix) so 'gpu-node' and 'gpu-node-2' don't collide." A force-drain of gpu-node will release leases belonging to gpu-node-2. Use _parse_resource_id() and compare parsed[0] == name like the offline path does.
| if lease.resource_id.startswith(name + ":") | |
| if (parsed := self._parse_resource_id(lease.resource_id)) and parsed[0] == name |
Reply with @kilocode-bot fix it to have Kilo Code address this issue.
| if lease.resource_id.startswith(name + ":") | ||
| ] | ||
| for lid in lids: | ||
| self._leases.pop(lid, None) |
There was a problem hiding this comment.
WARNING: Lease mutations bypass _lease_lock
The class docstring (lines 35-38) states all lease-table mutations are serialized via self._lease_lock. This sync method mutates self._leases directly without acquiring it, while _sweep_expired_leases, _monitor_loop, and other lease-mutating paths all hold the lock. Concurrent force-drain during a claim_lease or sweep can corrupt iteration. Wrap the release loop in async with self._lease_lock: — or make drain_worker async and await it from the route.
Reply with @kilocode-bot fix it to have Kilo Code address this issue.
| f"{'Tasks will complete before detach.' if graceful else 'All leases released immediately.'}" | ||
| ) | ||
| try: | ||
| asyncio.get_running_loop().create_task( |
There was a problem hiding this comment.
WARNING: Fire-and-forget notification task not retained
asyncio.get_running_loop().create_task(...) returns a task object that this code drops on the floor. The class already maintains self._background_tasks: set[asyncio.Task] (line 46) for exactly this purpose — add the task to it so the notification coroutine isn't GC'd mid-flight if the event loop runs a GC cycle before it completes.
Reply with @kilocode-bot fix it to have Kilo Code address this issue.
| elif worker.status == "draining": | ||
| active_leases = [ | ||
| lid for lid, lease in self._leases.items() | ||
| if lease.resource_id.startswith(worker.name + ":") |
There was a problem hiding this comment.
WARNING: Same prefix-collision bug in monitor drain path
Line 620 (and 639) repeat the lease.resource_id.startswith(worker.name + ":") pattern. A draining gpu-node whose leases actually belong to gpu-node-2 will incorrectly auto-complete the drain (line 622) and then mark gpu-node-2's owner offline when the drain timer fires (line 635-642). Use _parse_resource_id and compare parsed[0] == worker.name, matching the offline path above.
Reply with @kilocode-bot fix it to have Kilo Code address this issue.
| return result | ||
|
|
||
|
|
||
| @router.post("/api/cluster/workers/{name}/update") |
There was a problem hiding this comment.
CRITICAL: Destructive endpoints have no authentication/authorization
/drain, /cancel-drain, and /update can take down arbitrary cluster workers and force-release their inflight GPU leases. If the rest of the cluster API requires auth (look at neighboring routes in this file), these endpoints must too. At minimum, restrict them to an admin role and add rate limiting — an attacker or a buggy automation hitting /update in a loop will churn the whole cluster.
Reply with @kilocode-bot fix it to have Kilo Code address this issue.
| worker = cluster.get_worker(name) | ||
| if not worker: | ||
| return JSONResponse({"error": f"Worker '{name}' not found"}, status_code=404) | ||
| if worker.status not in ("online", "draining"): |
There was a problem hiding this comment.
WARNING: No idempotency / re-entry guard on update
Allowing update when worker.status == "draining" means a second concurrent call (or a retry) silently kicks off another update-worker deploy while the first is still in flight. Combined with the cancel-on-failure logic at line 1179, two simultaneous calls can race: one cancels the other's drain, or both POST update-worker to a worker that is about to restart. Track an in-flight updating flag (or an update_token) on the worker and 409 on duplicate.
Reply with @kilocode-bot fix it to have Kilo Code address this issue.
| f"{worker.url}/api/worker/deploy", | ||
| json={"command": "update-worker"}, | ||
| ) | ||
| deploy_result = resp.json() |
There was a problem hiding this comment.
CRITICAL: HTTP response status unchecked before parsing
resp.json() is called regardless of resp.status_code. If the worker returns 5xx with an HTML or plain-text error page (very likely during a deploy that's about to kill the process), the JSON parser raises, the broad except Exception fires, and the route cancels the drain — leaving a worker that actually accepted the update-worker command and is mid-restart. The handler returns 502 to the caller even though the update is proceeding. Check resp.raise_for_status() first, and on failure don't cancel the drain if status_code suggests the request was received (e.g. 5xx after the deploy was triggered).
Reply with @kilocode-bot fix it to have Kilo Code address this issue.
| json={"command": "update-worker"}, | ||
| ) | ||
| deploy_result = resp.json() | ||
| except Exception as exc: |
There was a problem hiding this comment.
WARNING: except Exception blindly cancels drain
The handler catches every exception (network errors, JSON parse errors, timeouts) and unconditionally calls cluster.cancel_drain(name). But the drain may have been started independently by a prior /drain call from another operator — cancelling it here wipes their intent. Cancel only when this route initiated the drain (track that via a flag on the worker or a return value from drain_worker). Also narrow the except to httpx.HTTPError so programming bugs surface.
Reply with @kilocode-bot fix it to have Kilo Code address this issue.
| cluster = request.app.state.cluster_manager | ||
| result = cluster.cancel_drain(name) | ||
| if "error" in result: | ||
| return JSONResponse(result, status_code=404) |
There was a problem hiding this comment.
SUGGESTION: Wrong status code for cancel-drain error
cancel_drain returns {"error": "worker is not draining (status=...)"} — a state-conflict error, not a missing resource. The route returns 404, but 409 Conflict is more accurate (and 400 is acceptable). The 404 leaks the existence of the worker name.
Reply with @kilocode-bot fix it to have Kilo Code address this issue.
| return JSONResponse(drain_result, status_code=404) | ||
|
|
||
| # Step 2: Trigger the update-worker deploy command on the worker | ||
| import httpx |
There was a problem hiding this comment.
SUGGESTION: Inline import of httpx
import httpx is inside the handler body. Move it to the module-level imports at the top of the file for consistency with the rest of the codebase and to avoid the import cost on every call.
Reply with @kilocode-bot fix it to have Kilo Code address this issue.
Code Review SummaryStatus: 6 Issues Found | Recommendation: Address before merge Overview
Issue Details (click to expand)CRITICAL
WARNING
SUGGESTION
Resolved in HEAD (incremental commit)
Files Reviewed (3 files)
Fix these issues in Kilo Cloud Previous Review Summaries (3 snapshots, latest commit a46b09b)Current summary above is authoritative. Previous snapshots are kept for context only. Previous review (commit a46b09b)Status: 6 Issues Found | Recommendation: Address before merge Overview
Issue Details (click to expand)CRITICAL
WARNING
SUGGESTION
Resolved in HEAD (incremental commit)
Files Reviewed (3 files)
Fix these issues in Kilo Cloud Previous review (commit 6903c00)Status: 8 Issues Found | Recommendation: Address before merge Overview
Issue Details (click to expand)CRITICAL
WARNING
SUGGESTION
Resolved in HEAD (incremental commit)
Files Reviewed (3 files)
Fix these issues in Kilo Cloud Previous review (commit c6d9a98)Status: 10 Issues Found | Recommendation: Address before merge Overview
Issue Details (click to expand)CRITICAL
WARNING
SUGGESTION
Files Reviewed (3 files)
Reviewed by hy3-20260706:free · Input: 78.5K · Output: 11.7K · Cached: 185.1K |
Replace startswith(name + ':') with _parse_resource_id + exact parsed[0] == name compare in three drain code paths to prevent accidental cross-worker lease release when worker names share a prefix (e.g. 'worker-a' matching 'worker-a-pro' leases). Fixes jaylfc#1690
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (2)
tests/test_cluster.py (1)
363-379: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winAdd a heartbeat regression for draining workers.
The suite verifies lease claims fail after
drain_worker(), but it does not cover the normal heartbeat path that can flip a draining worker back to"online".🧪 Proposed test
+ async def test_heartbeat_does_not_cancel_drain(self): + """Heartbeats should refresh metadata without making a draining worker routable.""" + mgr = ClusterManager() + await mgr.register_worker(_make_worker("draining-gpu", capabilities=["chat"])) + + mgr.drain_worker("draining-gpu", graceful=True) + assert mgr.heartbeat("draining-gpu", load=0.2) is True + + assert mgr.get_worker("draining-gpu").status == "draining" + assert mgr.get_workers_for_capability("chat") == []🤖 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 363 - 379, Add a regression test covering the heartbeat path for draining workers, since `test_draining_workers_excluded_from_lease_claim` only checks `claim_lease()` after `drain_worker()`. Extend the `ClusterManager`/worker heartbeat flow to verify that a worker marked draining does not get flipped back to "online" by a normal heartbeat, and then assert `claim_lease()` still returns `None` for that worker after the heartbeat. Use the existing `ClusterManager`, `drain_worker()`, and heartbeat handling logic to place the test near the current draining-worker lease test.tinyagentos/routes/cluster.py (1)
1114-1114: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAvoid constructing
DrainRequestin the function default.Move the default construction into the handler so the endpoint satisfies Ruff B008 and avoids shared default-object surprises.
Proposed fix
-async def drain_worker(request: Request, name: str, body: DrainRequest = DrainRequest()): +async def drain_worker(request: Request, name: str, body: DrainRequest | None = None): @@ """ + body = body or DrainRequest() cluster = request.app.state.cluster_manager🤖 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/routes/cluster.py` at line 1114, The drain_worker handler is constructing a DrainRequest as a function default, which triggers Ruff B008 and can create a shared mutable default. Update drain_worker so the body parameter defaults to None, then instantiate DrainRequest inside the handler when no body is provided; keep the change localized to the drain_worker endpoint and preserve the existing request handling behavior.Source: Linters/SAST tools
🤖 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 400-421: The `drain_worker()` flow in `ClusterManager` sets a
worker to `"draining"`, but `heartbeat()` later overwrites that status back to
`"online"` unconditionally. Update `heartbeat()` to preserve `"draining"`
workers as draining (and avoid re-enabling routing) until the drain completes or
the worker is explicitly taken offline, using the existing
`drain_worker()`/`heartbeat()` status handling logic.
In `@tinyagentos/routes/cluster.py`:
- Around line 1162-1176: The restart flow in cluster.py is triggering the
update-worker deploy too early after drain_worker(..., graceful=True), which
only starts draining and does not guarantee leases are finished. Update the
logic around the drain_result and the httpx.AsyncClient POST to
/api/worker/deploy so the deploy only happens after the worker is confirmed
lease-free/offline, or split it into a drain-start response plus a
background/monitor step that sends the deploy once drain completion is reached.
Use the existing drain_worker and deploy_result flow to locate and restructure
this sequence.
---
Nitpick comments:
In `@tests/test_cluster.py`:
- Around line 363-379: Add a regression test covering the heartbeat path for
draining workers, since `test_draining_workers_excluded_from_lease_claim` only
checks `claim_lease()` after `drain_worker()`. Extend the
`ClusterManager`/worker heartbeat flow to verify that a worker marked draining
does not get flipped back to "online" by a normal heartbeat, and then assert
`claim_lease()` still returns `None` for that worker after the heartbeat. Use
the existing `ClusterManager`, `drain_worker()`, and heartbeat handling logic to
place the test near the current draining-worker lease test.
In `@tinyagentos/routes/cluster.py`:
- Line 1114: The drain_worker handler is constructing a DrainRequest as a
function default, which triggers Ruff B008 and can create a shared mutable
default. Update drain_worker so the body parameter defaults to None, then
instantiate DrainRequest inside the handler when no body is provided; keep the
change localized to the drain_worker endpoint and preserve the existing request
handling behavior.
🪄 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: a3cfade6-53b5-44c7-9806-8c0d64329b54
📒 Files selected for processing (3)
tests/test_cluster.pytinyagentos/cluster/manager.pytinyagentos/routes/cluster.py
| def drain_worker(self, name: str, graceful: bool = True) -> dict: | ||
| """Begin draining a worker — gracefully detach without dropping tasks. | ||
|
|
||
| When ``graceful=True`` (the default): | ||
| 1. Worker enters ``"draining"`` status — no new tasks routed to it. | ||
| 2. Existing leases are allowed to run to completion (TTL not touched). | ||
| 3. The monitor loop will eventually sweep expired leases and then | ||
| mark the worker ``"offline"`` once all leases are released. | ||
|
|
||
| When ``graceful=False``: | ||
| 1. Worker enters ``"draining"`` status. | ||
| 2. All active leases for this worker are released immediately. | ||
| 3. Worker is marked ``"offline"`` immediately. | ||
|
|
||
| Returns a dict with ``worker``, ``previous_status``, ``released_leases``, | ||
| and ``status``. | ||
| """ | ||
| worker = self._workers.get(name) | ||
| if worker is None: | ||
| return {"worker": name, "error": "worker not found"} | ||
| prev_status = worker.status | ||
| worker.status = "draining" |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Preserve draining across heartbeats.
drain_worker() only changes worker.status, but heartbeat() unconditionally sets that same worker back to "online" on the next heartbeat, so a draining worker can re-enter routing before leases finish.
🐛 Proposed fix
- worker.status = "online"
+ if worker.status != "draining":
+ worker.status = "online"🤖 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 400 - 421, The `drain_worker()`
flow in `ClusterManager` sets a worker to `"draining"`, but `heartbeat()` later
overwrites that status back to `"online"` unconditionally. Update `heartbeat()`
to preserve `"draining"` workers as draining (and avoid re-enabling routing)
until the drain completes or the worker is explicitly taken offline, using the
existing `drain_worker()`/`heartbeat()` status handling logic.
| # Step 1: Begin draining | ||
| drain_result = cluster.drain_worker(name, graceful=True) | ||
| if "error" in drain_result: | ||
| return JSONResponse(drain_result, status_code=404) | ||
|
|
||
| # Step 2: Trigger the update-worker deploy command on the worker | ||
| import httpx | ||
| deploy_result = None | ||
| try: | ||
| async with httpx.AsyncClient(timeout=620) as client: | ||
| resp = await client.post( | ||
| f"{worker.url}/api/worker/deploy", | ||
| json={"command": "update-worker"}, | ||
| ) | ||
| deploy_result = resp.json() |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟠 Major | 🏗️ Heavy lift
Wait for graceful drain completion before triggering the restart.
drain_worker(..., graceful=True) only marks the worker as draining; existing leases remain active. Posting update-worker immediately can restart the worker before inflight work finishes, contradicting the documented sequence in Lines 1145-1149. Gate the deploy on the worker becoming lease-free/offline, or split this into a drain-start response plus a monitor/background step that sends the deploy after drain completion.
🤖 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/routes/cluster.py` around lines 1162 - 1176, The restart flow in
cluster.py is triggering the update-worker deploy too early after
drain_worker(..., graceful=True), which only starts draining and does not
guarantee leases are finished. Update the logic around the drain_result and the
httpx.AsyncClient POST to /api/worker/deploy so the deploy only happens after
the worker is confirmed lease-free/offline, or split it into a drain-start
response plus a background/monitor step that sends the deploy once drain
completion is reached. Use the existing drain_worker and deploy_result flow to
locate and restructure this sequence.
…on force-release, wire arbiter task cancellation jaylfc#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.
|
Deep-reviewed this one, @hognek — the drain/pause/routing/catalog/lease-exclusion design is genuinely good and the unit tests for it are solid. Two blocking items before I merge, plus a couple worth folding: BLOCKING 1 — the BLOCKING 2 — the update flow reports failure on success and cancels a drain the worker accepted (Kilo CRITICAL). In Worth folding (non-blocking):
Once 1 + 2 are in (and ideally an |
…#1690 - Restore local worker guard in _monitor_loop (prevents controller from marking itself offline and dropping its own leases). - Fix update_worker error handling: use raise_for_status() for real HTTP errors; catch expected disconnect during worker restart without cancelling drain; keep drain active until re-registration. - Move force-path lease pops under _lease_lock for consistency with the online->offline path. - Add idempotency guard: /update returns 409 if worker is already draining. - Move httpx import to module level in routes/cluster.py (matches codebase convention). - Add 7 tests for update_worker covering all error paths.
|
Addressed all review blockers from jaylfc's deep review: BLOCKING 1 — Restored BLOCKING 2 — Fixed
Non-blocking:
Tests: 7/7 new tests pass; all cluster manager + routes_cluster tests pass (44/44). |
| {"error": f"Worker update deploy rejected by worker", "drain_cancelled": True}, | ||
| status_code=502, | ||
| ) | ||
| except (httpx.ConnectError, httpx.RemoteProtocolError, httpx.ReadError): |
There was a problem hiding this comment.
WARNING: httpx.ConnectError is misclassified as an expected mid-restart disconnect.
ConnectError is raised when the TCP connection can never be established (connection refused, DNS failure, worker down/unreachable) — meaning the update-worker deploy request was never delivered. Treating it as "restart expected" makes the endpoint return 200 {status: acknowledged} and keep the worker draining even though the update never started. The worker then stays out of rotation until the monitor's HEARTBEAT_TIMEOUT force-finishes the drain, silently removing a healthy/unreachable worker from capacity without ever updating it.
A true mid-restart disconnect surfaces as RemoteProtocolError/ReadError (connection already established, server drops it). Consider dropping ConnectError from this branch, or only treating it as "restart expected" after the request is known to have been sent.
Additionally, a restart that stalls will raise httpx.ReadTimeout/TimeoutException (not in this tuple), falling through to the generic except Exception at line 1202 and wrongly cancelling the drain. Timeouts during a restart should arguably be handled here too.
Reply with @kilocode-bot fix it to have Kilo Code address this issue.
| cluster.cancel_drain(name) | ||
| logger.exception("Worker '%s' update deploy returned HTTP error", name) | ||
| return JSONResponse( | ||
| {"error": f"Worker update deploy rejected by worker", "drain_cancelled": True}, |
There was a problem hiding this comment.
SUGGESTION: f-string has no placeholders — drop the f prefix.
| {"error": f"Worker update deploy rejected by worker", "drain_cancelled": True}, | |
| {"error": "Worker update deploy rejected by worker", "drain_cancelled": True}, |
Reply with @kilocode-bot fix it to have Kilo Code address this issue.
|
Thanks Hognek, the blocker fixes look right: restoring the Kilo flagged a few things on the new revision worth folding before this merges, one of them important: CRITICAL (routes/cluster.py ~1139): the destructive endpoints ( WARNING (routes/cluster.py ~1191): the disconnect classification is not quite right yet. WARNING (cluster/manager.py ~435): there is a second force path in WARNING (cluster/manager.py ~446): the fire-and-forget notification Minor: Once the auth gate plus the two manager.py WARNINGs land I will re-review. This is a standalone worker-auto-update change (it does not touch the VRAM arbiter code), so it can go in on its own once these clear. Nice progress on the blockers. |
…lfc#1690 CRITICAL: Add _require_admin auth gate to /drain, /cancel-drain, /update endpoints — all three were missing the admin-session check that every other mutating cluster route uses. WARNING — disconnect classification: Remove ConnectError from the expected-restart tuple (connection-refused means unreachable, not restarting) and add ReadTimeout (the real mid-restart signal). ConnectError now cancels the drain and returns 502. WARNING — lease-lock: Make drain_worker async and wrap the force-release lease pops in async with self._lease_lock so concurrent claim/release operations see a consistent view. WARNING — task retention: Capture notification create_task results in drain_worker and cancel_drain, retain in self._background_tasks with add_done_callback(discard). Minor — cancel_drain returns 409 (not 404) on state conflict. Minor — bare f-string fix in update_worker error message. Tests: 22/22 drain+update tests pass (9 unit + 13 HTTP), including 6 new tests (3 auth-gate, ConnectError-as-failure, ReadTimeout-as-success, cancel_drain 409).
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (2)
tests/test_routes_cluster.py (2)
691-724: 🔒 Security & Privacy | 🔵 Trivial | ⚡ Quick winAssert the exact authorization response.
Accepting 401 allows a middleware response to pass even if the endpoint-level admin guard regresses.
_require_admindefines a 403 response; assert that status and theforbiddenbody.Proposed assertions
- assert resp.status_code in (401, 403), f"got {resp.status_code}: {resp.text}" + assert resp.status_code == 403 + assert resp.json() == {"error": "forbidden"}Apply this to all three authorization tests.
🤖 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_routes_cluster.py` around lines 691 - 724, Update test_drain_requires_admin, test_cancel_drain_requires_admin, and test_update_worker_requires_admin to assert exactly HTTP 403 and verify the response body contains the expected "forbidden" value, replacing the current (401, 403) status checks.
573-591: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winVerify the outgoing deployment contract.
The mock succeeds regardless of URL or payload, so a regression sending the wrong command would pass.
Proposed assertion
with patch("tinyagentos.routes.cluster.httpx.AsyncClient", return_value=mock_client): resp = await client.post("/api/cluster/workers/gpu-box/update") + mock_client.post.assert_awaited_once_with( + "http://gpu:9000/api/worker/deploy", + json={"command": "update-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_routes_cluster.py` around lines 573 - 591, Strengthen the test for the worker update route by asserting the mocked httpx.AsyncClient.post call in the test setup around the request to `/api/cluster/workers/gpu-box/update`. Verify the outgoing deployment URL and payload include the expected worker identifier and `"update-worker"` command, while retaining the existing response assertions.
🤖 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/routes/cluster.py`:
- Around line 1139-1141: Update the cancel-drain flow around the
cluster.cancel_drain call and the worker-draining logic near the referenced
update handling to track a distinct active-update state or token; while that
state is active, reject ordinary cancellation with an appropriate conflict
response instead of reopening routing, and clear the state only when the update
completes or is safely aborted.
---
Nitpick comments:
In `@tests/test_routes_cluster.py`:
- Around line 691-724: Update test_drain_requires_admin,
test_cancel_drain_requires_admin, and test_update_worker_requires_admin to
assert exactly HTTP 403 and verify the response body contains the expected
"forbidden" value, replacing the current (401, 403) status checks.
- Around line 573-591: Strengthen the test for the worker update route by
asserting the mocked httpx.AsyncClient.post call in the test setup around the
request to `/api/cluster/workers/gpu-box/update`. Verify the outgoing deployment
URL and payload include the expected worker identifier and `"update-worker"`
command, while retaining the existing response assertions.
🪄 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: fdca2baa-579f-4f29-932e-ec4367206e64
📒 Files selected for processing (4)
tests/test_cluster.pytests/test_routes_cluster.pytinyagentos/cluster/manager.pytinyagentos/routes/cluster.py
🚧 Files skipped from review as they are similar to previous changes (2)
- tests/test_cluster.py
- tinyagentos/cluster/manager.py
| result = cluster.cancel_drain(name) | ||
| if "error" in result: | ||
| return JSONResponse(result, status_code=409) |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | 🏗️ Heavy lift
Prevent cancel-drain from reopening routing during an active update.
After Line 1175 marks the worker as draining, the deployment request can remain pending for up to 620 seconds. During that window, /cancel-drain can restore it to online, allowing new leases immediately before the worker restarts. Track a distinct update state/token and reject ordinary cancellation while it is active.
Also applies to: 1174-1186
🤖 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/routes/cluster.py` around lines 1139 - 1141, Update the
cancel-drain flow around the cluster.cancel_drain call and the worker-draining
logic near the referenced update handling to track a distinct active-update
state or token; while that state is active, reject ordinary cancellation with an
appropriate conflict response instead of reopening routing, and clear the state
only when the update completes or is safely aborted.
2 participants
Summary
Implements worker auto-update with graceful pause/drain/restart for taOS cluster workers (taOS #890).
Skald workers on taOS nodes can now update without dropping inflight embedding or extraction jobs. The controller orchestrates the sequence: pause incoming routing → drain inflight leases → trigger deploy update-worker → worker restarts → re-registers → routing resumes.
Changes
ClusterManager (
tinyagentos/cluster/manager.py)drain_worker(name, graceful=True): Worker enters"draining"status, excluded from routing, catalog, and lease claims. Whengraceful=False, all leases are force-released and worker is marked offline immediately.cancel_drain(name): Returns a draining worker back to"online".get_workers_for_capability,aggregate_catalog, and_worker_for_resourceall exclude draining workers.Routes (
tinyagentos/routes/cluster.py)POST /api/cluster/workers/{name}/drain— Begin graceful or force drainPOST /api/cluster/workers/{name}/cancel-drain— Cancel an in-progress drainPOST /api/cluster/workers/{name}/update— Full orchestrated auto-update: drain → deploy update-worker → (worker restarts, re-registers, resumes routing)Tests (
tests/test_cluster.py)Flow
Summary by CodeRabbit
New Features
Bug Fixes