diff --git a/tests/test_cluster.py b/tests/test_cluster.py index 74875d441..9e543b366 100644 --- a/tests/test_cluster.py +++ b/tests/test_cluster.py @@ -259,3 +259,121 @@ async def test_router_returns_none_for_no_workers(self): data, worker_name = await router.route_request("chat", "POST", "/v1/chat/completions", {}) assert data is None assert worker_name is None + + +@pytest.mark.asyncio +class TestWorkerDrain: + """Tests for taOS #890 — worker auto-update with graceful drain.""" + + async def test_drain_worker_graceful_sets_draining_status(self): + """Graceful drain sets status to 'draining', leaves leases intact.""" + mgr = ClusterManager() + await mgr.register_worker(_make_worker("gpu-box", capabilities=["chat"])) + + result = await mgr.drain_worker("gpu-box", graceful=True) + assert result["worker"] == "gpu-box" + assert result["previous_status"] == "online" + assert result["status"] == "draining" + assert result["released_leases"] == 0 + assert mgr.get_worker("gpu-box").status == "draining" + + async def test_drain_worker_force_releases_leases_and_marks_offline(self): + """Force drain releases all leases and marks worker offline.""" + mgr = ClusterManager() + await mgr.register_worker(_make_worker("gpu-box", capabilities=["chat"])) + + # Add a lease for this worker + from tinyagentos.cluster.worker_protocol import GpuLease + import time + lease = GpuLease( + lease_id="l_test1", + resource_id="gpu-box:gpu-cuda-0", + caller="test", + expires_at=time.time() + 60, + required_vram_mb=0, + ) + mgr._leases["l_test1"] = lease + + result = await mgr.drain_worker("gpu-box", graceful=False) + assert result["released_leases"] == 1 + assert result["status"] == "offline" + assert mgr.get_worker("gpu-box").status == "offline" + assert "l_test1" not in mgr._leases + + async def test_drain_worker_unknown_returns_error(self): + """Draining a non-existent worker returns error dict.""" + mgr = ClusterManager() + result = await mgr.drain_worker("nonexistent") + assert "error" in result + assert result["worker"] == "nonexistent" + + async def test_cancel_drain_returns_worker_to_online(self): + """Cancel drain returns a draining worker back to online.""" + mgr = ClusterManager() + await mgr.register_worker(_make_worker("gpu-box", capabilities=["chat"])) + await mgr.drain_worker("gpu-box", graceful=True) + assert mgr.get_worker("gpu-box").status == "draining" + + result = mgr.cancel_drain("gpu-box") + assert result["worker"] == "gpu-box" + assert result["status"] == "online" + assert mgr.get_worker("gpu-box").status == "online" + + async def test_cancel_drain_only_works_on_draining(self): + """Cancel drain returns error for non-draining workers.""" + mgr = ClusterManager() + await mgr.register_worker(_make_worker("gpu-box")) + + result = mgr.cancel_drain("gpu-box") + assert "error" in result + assert "not draining" in result["error"] + + async def test_cancel_drain_unknown_returns_error(self): + """Cancel drain on unknown worker returns error.""" + mgr = ClusterManager() + result = mgr.cancel_drain("nonexistent") + assert "error" in result + + async def test_draining_workers_excluded_from_routing(self): + """Draining workers should not receive new tasks.""" + mgr = ClusterManager() + await mgr.register_worker(_make_worker("online-gpu", capabilities=["chat"], load=0.1)) + await mgr.register_worker(_make_worker("draining-gpu", capabilities=["chat"], load=0.0)) + await mgr.drain_worker("draining-gpu", graceful=True) + + result = mgr.get_workers_for_capability("chat") + assert len(result) == 1 + assert result[0].name == "online-gpu" + + async def test_draining_workers_excluded_from_catalog(self): + """Draining workers are excluded from the aggregate catalog.""" + mgr = ClusterManager() + online = _make_worker("online", capabilities=["chat"]) + online.backends = [{"name": "b1", "type": "ollama", "url": "u", "capabilities": ["chat"], "models": [{"name": "m1"}]}] + draining = _make_worker("draining", capabilities=["chat"]) + draining.backends = [{"name": "b2", "type": "ollama", "url": "u", "capabilities": ["chat"], "models": [{"name": "m2"}]}] + await mgr.register_worker(online) + await mgr.register_worker(draining) + await mgr.drain_worker("draining", graceful=True) + + out = mgr.aggregate_catalog() + assert [w["name"] for w in out["workers"]] == ["online"] + assert [m["name"] for m in out["models"]] == ["m1"] + + async def test_draining_workers_excluded_from_lease_claim(self): + """Lease claims should fail for draining workers.""" + mgr = ClusterManager() + await mgr.register_worker(_make_worker("draining-gpu", url="http://draining:9000")) + + # Set up worker with VRAM info + worker = mgr.get_worker("draining-gpu") + worker.free_vram_mb = 8000 + + await mgr.drain_worker("draining-gpu", graceful=True) + + lease = await mgr.claim_lease( + resource_id="draining-gpu:gpu-cuda-0", + caller="test", + ttl_seconds=30, + ) + assert lease is None diff --git a/tests/test_routes_cluster.py b/tests/test_routes_cluster.py index 8daf0848d..7dd0df6a0 100644 --- a/tests/test_routes_cluster.py +++ b/tests/test_routes_cluster.py @@ -520,3 +520,282 @@ async def fake_remote_list(): assert fedora is not None assert fedora["hardware_known"] is True, fedora assert fedora["tier_id"] not in ("", "unknown"), fedora + + +# ── update_worker endpoint tests (taOS #890) ────────────────────────────── + + +@pytest.mark.asyncio +async def test_update_worker_not_found(client): + """404 when worker does not exist.""" + resp = await client.post("/api/cluster/workers/nonexistent/update") + assert resp.status_code == 404 + assert "not found" in resp.json()["error"] + + +@pytest.mark.asyncio +async def test_update_worker_already_draining(client, app): + """409 Conflict when worker is already draining (idempotency guard).""" + from tinyagentos.cluster.worker_protocol import WorkerInfo + cluster = app.state.cluster_manager + cluster._workers["gpu-box"] = WorkerInfo( # noqa: SLF001 + name="gpu-box", url="http://gpu:9000", status="draining", + ) + resp = await client.post("/api/cluster/workers/gpu-box/update") + assert resp.status_code == 409 + assert "already updating" in resp.json()["error"] + + +@pytest.mark.asyncio +async def test_update_worker_offline(client, app): + """400 when worker is offline.""" + from tinyagentos.cluster.worker_protocol import WorkerInfo + cluster = app.state.cluster_manager + cluster._workers["gpu-box"] = WorkerInfo( # noqa: SLF001 + name="gpu-box", url="http://gpu:9000", status="offline", + ) + resp = await client.post("/api/cluster/workers/gpu-box/update") + assert resp.status_code == 400 + assert "not online" in resp.json()["error"] + + +@pytest.mark.asyncio +async def test_update_worker_success(client, app): + """Happy path: drain succeeds, deploy returns 200.""" + from tinyagentos.cluster.worker_protocol import WorkerInfo + cluster = app.state.cluster_manager + cluster._workers["gpu-box"] = WorkerInfo( # noqa: SLF001 + name="gpu-box", url="http://gpu:9000", status="online", + ) + + from unittest.mock import MagicMock + + mock_resp = MagicMock() + mock_resp.status_code = 200 + mock_resp.json.return_value = {"ok": True, "command": "update-worker"} + mock_resp.raise_for_status = lambda: None + + mock_client = AsyncMock() + mock_client.__aenter__.return_value = mock_client + mock_client.__aexit__.return_value = None + mock_client.post.return_value = mock_resp + + with patch("tinyagentos.routes.cluster.httpx.AsyncClient", return_value=mock_client): + resp = await client.post("/api/cluster/workers/gpu-box/update") + + assert resp.status_code == 200 + data = resp.json() + assert data["worker"] == "gpu-box" + assert data["status"] == "updating" + assert data["drain"]["status"] == "draining" + assert data["deploy"]["ok"] is True + + # Worker should be draining after the call + assert cluster.get_worker("gpu-box").status == "draining" + + +@pytest.mark.asyncio +async def test_update_worker_deploy_http_error(client, app): + """502 when deploy endpoint returns a non-2xx status — drain cancelled.""" + from tinyagentos.cluster.worker_protocol import WorkerInfo + from unittest.mock import MagicMock + import httpx as _httpx + cluster = app.state.cluster_manager + cluster._workers["gpu-box"] = WorkerInfo( # noqa: SLF001 + name="gpu-box", url="http://gpu:9000", status="online", + ) + + mock_resp = MagicMock() + mock_resp.status_code = 500 + mock_resp.raise_for_status = lambda: (_ for _ in ()).throw( + _httpx.HTTPStatusError("Server error", request=MagicMock(), response=mock_resp) + ) + + mock_client = AsyncMock() + mock_client.__aenter__.return_value = mock_client + mock_client.__aexit__.return_value = None + mock_client.post.return_value = mock_resp + + with patch("tinyagentos.routes.cluster.httpx.AsyncClient", return_value=mock_client): + resp = await client.post("/api/cluster/workers/gpu-box/update") + + assert resp.status_code == 502 + data = resp.json() + assert "rejected by worker" in data["error"] + assert data["drain_cancelled"] is True + + # Drain should be cancelled — worker back online + assert cluster.get_worker("gpu-box").status == "online" + + +@pytest.mark.asyncio +async def test_update_worker_deploy_disconnect_during_restart(client, app): + """200 when worker disconnects mid-update (expected restart). Drain + stays active — the worker will re-register and the monitor loop + will auto-complete the drain once all leases are released.""" + from tinyagentos.cluster.worker_protocol import WorkerInfo + import httpx as _httpx + cluster = app.state.cluster_manager + cluster._workers["gpu-box"] = WorkerInfo( # noqa: SLF001 + name="gpu-box", url="http://gpu:9000", status="online", + ) + + mock_client = AsyncMock() + mock_client.__aenter__.return_value = mock_client + mock_client.__aexit__.return_value = None + mock_client.post.side_effect = _httpx.RemoteProtocolError( + "Server disconnected", request=AsyncMock() + ) + + with patch("tinyagentos.routes.cluster.httpx.AsyncClient", return_value=mock_client): + resp = await client.post("/api/cluster/workers/gpu-box/update") + + assert resp.status_code == 200 + data = resp.json() + assert data["worker"] == "gpu-box" + assert data["status"] == "updating" + assert data["deploy"]["status"] == "acknowledged" + assert "restart expected" in data["deploy"]["note"] + + # Drain must NOT be cancelled — worker should still be draining + assert cluster.get_worker("gpu-box").status == "draining" + + +@pytest.mark.asyncio +async def test_update_worker_deploy_unexpected_error(client, app): + """502 on unexpected exception — drain cancelled.""" + from tinyagentos.cluster.worker_protocol import WorkerInfo + cluster = app.state.cluster_manager + cluster._workers["gpu-box"] = WorkerInfo( # noqa: SLF001 + name="gpu-box", url="http://gpu:9000", status="online", + ) + + mock_client = AsyncMock() + mock_client.__aenter__.return_value = mock_client + mock_client.__aexit__.return_value = None + mock_client.post.side_effect = OSError("Network unreachable") + + with patch("tinyagentos.routes.cluster.httpx.AsyncClient", return_value=mock_client): + resp = await client.post("/api/cluster/workers/gpu-box/update") + + assert resp.status_code == 502 + data = resp.json() + assert "deploy failed" in data["error"] + assert data["drain_cancelled"] is True + assert cluster.get_worker("gpu-box").status == "online" + + +# ── update_worker auth gate tests ─────────────────────────────────────── + + +@pytest.mark.asyncio +async def test_drain_requires_admin(app): + """403 when drain endpoint is called without an admin session.""" + from httpx import ASGITransport + + async with AsyncClient( + transport=ASGITransport(app=app), base_url="http://test" + ) as noauth_client: + resp = await noauth_client.post("/api/cluster/workers/gpu-box/drain", json={}) + assert resp.status_code in (401, 403), f"got {resp.status_code}: {resp.text}" + + +@pytest.mark.asyncio +async def test_cancel_drain_requires_admin(app): + """403 when cancel-drain endpoint is called without an admin session.""" + from httpx import ASGITransport + + async with AsyncClient( + transport=ASGITransport(app=app), base_url="http://test" + ) as noauth_client: + resp = await noauth_client.post("/api/cluster/workers/gpu-box/cancel-drain") + assert resp.status_code in (401, 403), f"got {resp.status_code}: {resp.text}" + + +@pytest.mark.asyncio +async def test_update_worker_requires_admin(app): + """403 when update endpoint is called without an admin session.""" + from httpx import ASGITransport + + async with AsyncClient( + transport=ASGITransport(app=app), base_url="http://test" + ) as noauth_client: + resp = await noauth_client.post("/api/cluster/workers/gpu-box/update") + assert resp.status_code in (401, 403), f"got {resp.status_code}: {resp.text}" + + +# ── update_worker disconnect classification tests ──────────────────────── + + +@pytest.mark.asyncio +async def test_update_worker_deploy_connect_error(client, app): + """502 when deploy endpoint raises ConnectError — connection refused means + the worker is genuinely unreachable, not restarting. Drain is cancelled.""" + from tinyagentos.cluster.worker_protocol import WorkerInfo + import httpx as _httpx + cluster = app.state.cluster_manager + cluster._workers["gpu-box"] = WorkerInfo( # noqa: SLF001 + name="gpu-box", url="http://gpu:9000", status="online", + ) + + mock_client = AsyncMock() + mock_client.__aenter__.return_value = mock_client + mock_client.__aexit__.return_value = None + mock_client.post.side_effect = _httpx.ConnectError( + "Connection refused" + ) + + with patch("tinyagentos.routes.cluster.httpx.AsyncClient", return_value=mock_client): + resp = await client.post("/api/cluster/workers/gpu-box/update") + + assert resp.status_code == 502 + data = resp.json() + assert "connection refused" in data["error"].lower() + assert data["drain_cancelled"] is True + # Drain must be cancelled — worker back online + assert cluster.get_worker("gpu-box").status == "online" + + +@pytest.mark.asyncio +async def test_update_worker_deploy_read_timeout(client, app): + """200 when deploy endpoint raises ReadTimeout — worker restarted + mid-request, which is the normal update path. Drain stays active.""" + from tinyagentos.cluster.worker_protocol import WorkerInfo + import httpx as _httpx + cluster = app.state.cluster_manager + cluster._workers["gpu-box"] = WorkerInfo( # noqa: SLF001 + name="gpu-box", url="http://gpu:9000", status="online", + ) + + mock_client = AsyncMock() + mock_client.__aenter__.return_value = mock_client + mock_client.__aexit__.return_value = None + mock_client.post.side_effect = _httpx.ReadTimeout( + "Read timed out", request=AsyncMock() + ) + + with patch("tinyagentos.routes.cluster.httpx.AsyncClient", return_value=mock_client): + resp = await client.post("/api/cluster/workers/gpu-box/update") + + assert resp.status_code == 200 + data = resp.json() + assert data["worker"] == "gpu-box" + assert data["status"] == "updating" + assert data["deploy"]["status"] == "acknowledged" + assert "restart expected" in data["deploy"]["note"] + # Drain must NOT be cancelled — worker should still be draining + assert cluster.get_worker("gpu-box").status == "draining" + + +@pytest.mark.asyncio +async def test_cancel_drain_state_conflict_returns_409(client, app): + """409 when cancel-drain is called on a worker not in draining state.""" + from tinyagentos.cluster.worker_protocol import WorkerInfo + cluster = app.state.cluster_manager + cluster._workers["gpu-box"] = WorkerInfo( # noqa: SLF001 + name="gpu-box", url="http://gpu:9000", status="online", + ) + + resp = await client.post("/api/cluster/workers/gpu-box/cancel-drain") + assert resp.status_code == 409 + assert "not draining" in resp.json()["error"] diff --git a/tinyagentos/cluster/manager.py b/tinyagentos/cluster/manager.py index 69df08323..d3a3a12b3 100644 --- a/tinyagentos/cluster/manager.py +++ b/tinyagentos/cluster/manager.py @@ -273,13 +273,15 @@ def _parse_resource_id(self, resource_id: str) -> tuple[str, str] | None: return worker, rest def _worker_for_resource(self, resource_id: str) -> WorkerInfo | None: - """Return the WorkerInfo for a resource_id, or None.""" + """Return the WorkerInfo for a resource_id, or None. + + Excludes draining and offline workers (taOS #890).""" parsed = self._parse_resource_id(resource_id) if parsed is None: return None worker_name, _ = parsed worker = self._workers.get(worker_name) - if worker is None or worker.status != "online": + if worker is None or worker.status not in ("online",): return None return worker @@ -393,10 +395,114 @@ def _sweep_expired_leases(self): lease = self._leases.pop(lid) logger.debug("Lease expired: %s on %s", lid, lease.resource_id) + # ── taOS #890: graceful worker drain for auto-update ─────────────── + + async 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" + logger.info("Worker '%s' entering drain (was %s, graceful=%s)", + name, prev_status, graceful) + + released = 0 + + if not graceful: + # Force-release all leases for this worker. + # Lease pops must be serialised under _lease_lock so concurrent + # claim/release operations see a consistent view (taOS #1690). + async with self._lease_lock: + lids = [ + 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) + released += 1 + logger.info("Worker '%s' drain: force-released %d leases", name, released) + worker.status = "offline" + + if self._notifications: + detail = ( + f"Worker '{name}' is draining. " + f"{'Tasks will complete before detach.' if graceful else 'All leases released immediately.'}" + ) + try: + task = asyncio.create_task( + self._notifications.emit_event( + "worker.drain", + f"Worker '{name}' draining", + detail, + level="info", + ) + ) + self._background_tasks.add(task) + task.add_done_callback(self._background_tasks.discard) + except RuntimeError: + pass + + return { + "worker": name, + "previous_status": prev_status, + "status": worker.status, + "released_leases": released, + } + + def cancel_drain(self, name: str) -> dict: + """Cancel an in-progress drain and return the worker to online. + + Only works when the worker is in ``"draining"`` status and has + not yet been marked offline. + """ + worker = self._workers.get(name) + if worker is None: + return {"worker": name, "error": "worker not found"} + if worker.status != "draining": + return {"worker": name, "error": f"worker is not draining (status={worker.status})"} + worker.status = "online" + logger.info("Worker '%s' drain cancelled — back online", name) + + if self._notifications: + try: + task = asyncio.create_task( + self._notifications.emit_event( + "worker.online", + f"Worker '{name}' drain cancelled", + f"Worker '{name}' is back online after drain was cancelled.", + level="info", + ) + ) + self._background_tasks.add(task) + task.add_done_callback(self._background_tasks.discard) + except RuntimeError: + pass + + return {"worker": name, "status": "online"} + # ── Capability routing ───────────────────────────────────────────── def get_workers_for_capability(self, capability: str) -> list[WorkerInfo]: - """Get online workers that support a capability, sorted by priority (lowest load first).""" + """Get online workers that support a capability, sorted by priority (lowest load first). + + Draining workers are excluded (taOS #890).""" eligible = [ w for w in self._workers.values() if w.status == "online" and capability in w.capabilities @@ -443,7 +549,7 @@ def aggregate_catalog(self) -> dict: for worker in self._workers.values(): if worker.status != "online": - continue + continue # skip offline and draining workers (taOS #890) worker_caps = set(worker.capabilities or []) all_capabilities |= worker_caps @@ -486,14 +592,17 @@ def aggregate_catalog(self) -> dict: async def _monitor_loop(self): """Monitor worker heartbeats, mark stale workers as offline, and - sweep expired GPU leases.""" + sweep expired GPU leases. Auto-completes draining workers + whose leases have all been released (taOS #890).""" while True: now = time.time() - for worker in self._workers.values(): - # The 'local' worker is the controller itself — it never sends - # heartbeats (it IS the server), so never mark it offline. + for worker in list(self._workers.values()): + # The local worker (controller itself) must never be marked + # offline by the monitor loop — doing so would drop the + # controller's own leases and break local routing. if worker.name == "local": continue + # Handle online workers that haven't heartbeated if worker.status == "online" and (now - worker.last_heartbeat) > HEARTBEAT_TIMEOUT: worker.status = "offline" logger.warning(f"Worker '{worker.name}' marked offline (no heartbeat for {HEARTBEAT_TIMEOUT}s)") @@ -516,6 +625,46 @@ async def _monitor_loop(self): for lid in offline_lids: self._leases.pop(lid, None) logger.debug("Lease %s released — worker %s went offline", lid, worker.name) + + # Handle draining workers: auto-complete if no active leases (taOS #890) + elif worker.status == "draining": + active_leases = [ + lid for lid, lease in self._leases.items() + if (parsed := self._parse_resource_id(lease.resource_id)) + and parsed[0] == worker.name + ] + if not active_leases: + worker.status = "offline" + logger.info( + "Worker '%s' drained — all leases released, marked offline", + worker.name, + ) + if self._notifications: + await self._notifications.emit_event( + "worker.leave", + f"Worker '{worker.name}' drained and went offline", + "All tasks completed; worker detached gracefully.", + level="info", + ) + elif (now - worker.last_heartbeat) > HEARTBEAT_TIMEOUT: + # Draining worker went stale — force-finish the drain. + # Lease pops must be serialised under _lease_lock so + # concurrent claim/release operations see a consistent + # view of the lease table. + async with self._lease_lock: + lids = [ + lid for lid, lease in self._leases.items() + if (parsed := self._parse_resource_id(lease.resource_id)) + and parsed[0] == worker.name + ] + for lid in lids: + self._leases.pop(lid, None) + worker.status = "offline" + logger.warning( + "Worker '%s' drain timed out (no heartbeat for %ds) — " + "force-released %d leases, marked offline", + worker.name, HEARTBEAT_TIMEOUT, len(lids), + ) async with self._lease_lock: self._sweep_expired_leases() await asyncio.sleep(5) diff --git a/tinyagentos/routes/cluster.py b/tinyagentos/routes/cluster.py index ef5d91708..eff507200 100644 --- a/tinyagentos/routes/cluster.py +++ b/tinyagentos/routes/cluster.py @@ -2,17 +2,17 @@ import asyncio import logging +import re import subprocess from dataclasses import asdict +import httpx from fastapi import APIRouter, Request from fastapi.responses import JSONResponse from pydantic import BaseModel logger = logging.getLogger(__name__) -import re - from tinyagentos.cluster.capabilities import hardware_to_targets, potential_capabilities as _potential_capabilities from tinyagentos.cluster.optimiser import ClusterOptimiser from tinyagentos.cluster.worker_auth import _HMACError, require_worker_hmac @@ -900,7 +900,6 @@ async def deploy_backend(request: Request, name: str, body: DeployRequest): status_code=400, ) - import httpx try: async with httpx.AsyncClient(timeout=620) as client: resp = await client.post( @@ -935,7 +934,6 @@ async def worker_remote_command(request: Request, name: str, body: WorkerRemoteR status_code=403, ) - import httpx try: async with httpx.AsyncClient(timeout=body.timeout + 5) as client: resp = await client.post( @@ -1101,3 +1099,143 @@ async def list_leases(request: Request): ], "count": len(leases), } + + +# ── taOS #890: worker auto-update with graceful drain ──────────────────── + + +class DrainRequest(BaseModel): + graceful: bool = True + + +@router.post("/api/cluster/workers/{name}/drain") +async def drain_worker(request: Request, name: str, body: DrainRequest = DrainRequest()): + """Begin draining a worker — gracefully detach without dropping tasks. + + When ``graceful=true`` (the default), the worker enters "draining" status: + no new tasks are routed to it, but existing leases run to completion. + The monitor loop auto-completes the drain when all leases are released. + + When ``graceful=false``, all leases are force-released and the worker + is marked offline immediately. + """ + ok, err = _require_admin(request) + if not ok: + return err + cluster = request.app.state.cluster_manager + result = await cluster.drain_worker(name, graceful=body.graceful) + if "error" in result: + return JSONResponse(result, status_code=404) + return result + + +@router.post("/api/cluster/workers/{name}/cancel-drain") +async def cancel_drain(request: Request, name: str): + """Cancel an in-progress drain and return the worker to online status.""" + ok, err = _require_admin(request) + if not ok: + return err + cluster = request.app.state.cluster_manager + result = cluster.cancel_drain(name) + if "error" in result: + return JSONResponse(result, status_code=409) + return result + + +@router.post("/api/cluster/workers/{name}/update") +async def update_worker(request: Request, name: str): + """Trigger a full auto-update sequence on a worker. + + Orchestrates the graceful update sequence: + 1. Pause incoming tasks — worker enters "draining" status. + 2. Inflight leases run to completion (drain). + 3. Send ``update-worker`` deploy command to the worker. + 4. Worker restarts, re-registers, returns to "online". + 5. Routing resumes automatically once the worker is back online. + """ + ok, err = _require_admin(request) + if not ok: + return err + cluster = request.app.state.cluster_manager + worker = cluster.get_worker(name) + if not worker: + return JSONResponse({"error": f"Worker '{name}' not found"}, status_code=404) + if worker.status == "draining": + return JSONResponse( + {"error": f"Worker '{name}' is already updating (status=draining)"}, + status_code=409, + ) + if worker.status != "online": + return JSONResponse( + {"error": f"Worker '{name}' is not online (status={worker.status})"}, + status_code=400, + ) + + # Step 1: Begin draining + drain_result = await 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. + # The deploy script runs the update and then restarts the worker + # service. The HTTP connection will break when the worker restarts + # — that is expected and must not cancel the drain. + 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"}, + ) + resp.raise_for_status() + deploy_result = resp.json() + except httpx.HTTPStatusError: + # The deploy endpoint returned a non-2xx status — real failure. + cluster.cancel_drain(name) + logger.exception("Worker '%s' update deploy returned HTTP error", name) + return JSONResponse( + {"error": "Worker update deploy rejected by worker", "drain_cancelled": True}, + status_code=502, + ) + except (httpx.RemoteProtocolError, httpx.ReadError, httpx.ReadTimeout): + # Worker restarted mid-request — the connection dropped because + # the worker service stopped. This is the normal update path. + # Keep the drain active; the worker will re-register and the + # monitor loop will auto-complete the drain when all leases + # are released (taOS #890). + logger.info( + "Worker '%s' disconnected during update — restart expected, " + "drain stays active until re-registration", name, + ) + deploy_result = {"status": "acknowledged", "note": "worker disconnected — restart expected"} + except httpx.ConnectError: + # Connection-refused / DNS failure — the worker is genuinely + # unreachable, not restarting. Cancel the drain so the worker + # can be retried. + cluster.cancel_drain(name) + logger.exception("Worker '%s' update deploy failed — connection refused", name) + return JSONResponse( + {"error": "Worker unreachable — connection refused", "drain_cancelled": True}, + status_code=502, + ) + except Exception as exc: + # Unexpected failure — cancel drain so worker can still serve traffic. + cluster.cancel_drain(name) + logger.exception("Worker '%s' update deploy failed unexpectedly", name) + return JSONResponse( + {"error": f"Worker update deploy failed: {exc}", "drain_cancelled": True}, + status_code=502, + ) + + return { + "worker": name, + "status": "updating", + "previous_status": drain_result["previous_status"], + "drain": drain_result, + "deploy": deploy_result, + "message": ( + "Worker is draining and updating. It will restart and re-register. " + "The controller monitor loop will auto-complete the drain once " + "all leases are released and the new worker process heartbeats." + ), + }