From e7c8f8de5f43947633eeb89ab59f8e454ae7ac6e Mon Sep 17 00:00:00 2001 From: Hogne <227774406+hognek@users.noreply.github.com> Date: Wed, 8 Jul 2026 00:34:50 +0200 Subject: [PATCH] fix(cluster): release GPU leases on worker unregister (taOS #1705) unregister_worker() now releases all active leases for the worker's resources before removing the worker from the registry. Previously only the heartbeat-timeout path in _monitor_loop released leases; explicit unregister (DELETE /api/cluster/workers/{name}) leaked them. Changes: - manager.py: unregister_worker() made async, acquires _lease_lock, releases all leases for the worker (exact name match), then pops the worker. Logs lease count released. - routes/cluster.py: await the now-async unregister_worker() call. - tests: 6 new tests (5 unit + 1 HTTP integration) covering release on unregister, noop for unknown worker, worker isolation, reclaim after unregister, and lock serialization. --- tests/test_cluster.py | 4 +- tests/test_leases.py | 124 +++++++++++++++++++++++++++++++++ tinyagentos/cluster/manager.py | 32 ++++++++- tinyagentos/routes/cluster.py | 2 +- 4 files changed, 157 insertions(+), 5 deletions(-) diff --git a/tests/test_cluster.py b/tests/test_cluster.py index 74875d441..2be9cd264 100644 --- a/tests/test_cluster.py +++ b/tests/test_cluster.py @@ -43,9 +43,9 @@ async def test_register_worker(self): async def test_unregister_worker(self): mgr = ClusterManager() await mgr.register_worker(_make_worker("gpu-box")) - assert mgr.unregister_worker("gpu-box") is True + assert await mgr.unregister_worker("gpu-box") is True assert mgr.get_workers() == [] - assert mgr.unregister_worker("gpu-box") is False + assert await mgr.unregister_worker("gpu-box") is False async def test_heartbeat_updates_load_and_status(self): mgr = ClusterManager() diff --git a/tests/test_leases.py b/tests/test_leases.py index 4287d3bc0..c252e08d8 100644 --- a/tests/test_leases.py +++ b/tests/test_leases.py @@ -338,6 +338,52 @@ async def test_claim_rejects_negative_vram(client, app): assert resp.status_code == 422 +# ── taOS #1705: DELETE /worker endpoint leases ───────────────────────── + + +@pytest.mark.asyncio +async def test_unregister_worker_releases_leases(client, app): + """DELETE /api/cluster/workers/{name} releases all GPU leases held + by that worker (taOS #1705).""" + # Register a worker with enough VRAM + key = await _register_worker( + client, app, "gpu-node", "http://10.0.0.1:9000", + capabilities=["llm-chat"], + hardware={"gpu": {"model": "GTX 1080", "vram_mb": 8192}}, + ) + await _heartbeat(client, key, "gpu-node", free_vram_mb=8000, used_vram_mb=0) + + # Claim two leases + claim1 = await client.post("/api/cluster/leases/claim", headers=_auth(app), json={ + "resource_id": "gpu-node:gpu-cuda-0", + "caller": "test", + "ttl_seconds": 60, + }) + assert claim1.status_code == 200 + + claim2 = await client.post("/api/cluster/leases/claim", headers=_auth(app), json={ + "resource_id": "gpu-node:gpu-cuda-1", + "caller": "test", + "ttl_seconds": 60, + }) + assert claim2.status_code == 200 + + # Verify both leases are present + list_resp = await client.get("/api/cluster/leases", headers=_auth(app)) + assert list_resp.json()["count"] == 2 + + # Unregister the worker — this must release its leases + resp = await client.delete("/api/cluster/workers/gpu-node") + assert resp.status_code == 200 + assert resp.json()["status"] == "removed" + + # Verify all leases are gone + list_resp = await client.get("/api/cluster/leases", headers=_auth(app)) + assert list_resp.json()["count"] == 0 + + await app.state.cluster_pairing.close() + + # ── ClusterManager unit tests ────────────────────────────────────────── @@ -530,3 +576,81 @@ async def test_claim_serialized_by_lease_lock(self): result = await task assert result is not None + + # ── taOS #1705: unregister releases leases ───────────────────────── + + async def test_unregister_releases_all_leases(self): + """Unregistering a worker releases all of its active GPU leases.""" + mgr = ClusterManager() + mgr._workers["gpu-node"] = _worker("gpu-node", free_vram=8000) + + # Claim several leases + lease1 = await mgr.claim_lease("gpu-node:gpu-cuda-0", caller="test") + lease2 = await mgr.claim_lease("gpu-node:gpu-cuda-1", caller="test") + assert lease1 is not None + assert lease2 is not None + assert len(mgr.get_leases()) == 2 + + # Unregister — all leases must be gone + assert await mgr.unregister_worker("gpu-node") is True + assert len(mgr.get_leases()) == 0 + assert lease1.lease_id not in mgr._leases + assert lease2.lease_id not in mgr._leases + + async def test_unregister_unknown_worker_noop(self): + """Unregistering an unknown worker returns False and touches no leases.""" + mgr = ClusterManager() + mgr._workers["gpu-node"] = _worker("gpu-node", free_vram=8000) + await mgr.claim_lease("gpu-node:gpu-cuda-0", caller="test") + assert len(mgr.get_leases()) == 1 + + assert await mgr.unregister_worker("ghost") is False + assert len(mgr.get_leases()) == 1 # existing lease untouched + + async def test_unregister_one_worker_spares_another(self): + """Unregistering worker A leaves worker B's leases untouched.""" + mgr = ClusterManager() + mgr._workers["gpu-a"] = _worker("gpu-a", free_vram=8000) + mgr._workers["gpu-b"] = _worker("gpu-b", free_vram=8000) + + lease_a = await mgr.claim_lease("gpu-a:gpu-cuda-0", caller="test") + lease_b = await mgr.claim_lease("gpu-b:gpu-cuda-0", caller="test") + assert lease_a is not None + assert lease_b is not None + + assert await mgr.unregister_worker("gpu-a") is True + assert len(mgr.get_leases()) == 1 + assert lease_b.lease_id in mgr._leases + assert lease_a.lease_id not in mgr._leases + + async def test_reclaim_after_unregister(self): + """After unregister releases leases, a new claim succeeds.""" + mgr = ClusterManager() + mgr._workers["gpu-node"] = _worker("gpu-node", free_vram=8000) + + await mgr.claim_lease("gpu-node:gpu-cuda-0", caller="first") + assert await mgr.unregister_worker("gpu-node") is True + + # Re-register and claim again + mgr._workers["gpu-node"] = _worker("gpu-node", free_vram=8000) + lease = await mgr.claim_lease("gpu-node:gpu-cuda-0", caller="second") + assert lease is not None + assert lease.caller == "second" + + async def test_unregister_releases_leases_atomically(self): + """Unregister holds _lease_lock so concurrent claims see a + consistent view — they cannot race with the lease release.""" + mgr = ClusterManager() + mgr._workers["gpu-node"] = _worker("gpu-node", free_vram=8000) + + # Hold the lock so unregister blocks on it + async with mgr._lease_lock: + unregister_task = asyncio.create_task( + mgr.unregister_worker("gpu-node") + ) + await asyncio.sleep(0.05) + assert not unregister_task.done() + + # Lock released — unregister completes + assert await unregister_task + assert len(mgr.get_leases()) == 0 diff --git a/tinyagentos/cluster/manager.py b/tinyagentos/cluster/manager.py index ffbed738f..50f78babe 100644 --- a/tinyagentos/cluster/manager.py +++ b/tinyagentos/cluster/manager.py @@ -259,8 +259,36 @@ def heartbeat( pass # No running loop (e.g. in sync tests) — skip gracefully return True - def unregister_worker(self, name: str) -> bool: - return self._workers.pop(name, None) is not None + async def unregister_worker(self, name: str) -> bool: + """Remove a worker and all of its active GPU leases. + + When a worker is explicitly unregistered (admin action), every + active lease tied to the worker's resources is released so that + those resource slots become available for new claims immediately. + This mirrors the lease-release logic in ``_monitor_loop`` for the + heartbeat-timeout path (taOS #1705). + """ + worker = self._workers.get(name) + if worker is None: + return False + + # Release all active leases for this worker's resources. + async with self._lease_lock: + lids: list[str] = [ + lid for lid, lease in self._leases.items() + if (parsed := self._parse_resource_id(lease.resource_id)) + and parsed[0] == name + ] + for lid in lids: + self._leases.pop(lid, None) + logger.debug( + "Lease %s released — worker '%s' unregistered", + lid, name, + ) + + self._workers.pop(name, None) + logger.info("Worker '%s' unregistered — %d leases released", name, len(lids)) + return True def get_workers(self) -> list[WorkerInfo]: return list(self._workers.values()) diff --git a/tinyagentos/routes/cluster.py b/tinyagentos/routes/cluster.py index df7fc4742..a559cb7ce 100644 --- a/tinyagentos/routes/cluster.py +++ b/tinyagentos/routes/cluster.py @@ -515,7 +515,7 @@ async def worker_heartbeat(request: Request, body: HeartbeatBody): @router.delete("/api/cluster/workers/{name}") async def unregister_worker(request: Request, name: str): cluster = request.app.state.cluster_manager - removed = cluster.unregister_worker(name) + removed = await cluster.unregister_worker(name) if not removed: return JSONResponse({"error": "Worker not found"}, status_code=404) return {"status": "removed", "name": name}