From 0dd39b4ece3ee62ad64eb97e64737fb0b143322d Mon Sep 17 00:00:00 2001 From: hognek <227774406+hognek@users.noreply.github.com> Date: Sat, 12 Sep 2026 00:11:04 +0200 Subject: [PATCH 01/17] feat(a2a): GPU lease protocol over the coordination bus (#893) Two agents sharing one GPU could silently co-load past its VRAM: taOS FLUX image-gen OOM-killed itself because taOSmd had ~9.4 GB of Ollama models resident on the same 12 GB card after an earlier "free" signal. The controller already had a lease registry and the GPU arbiter, but nothing coordinated across PRODUCTS over the A2A bus. Adds the one-line text protocol (CHECK/CLAIM/RELEASE/REQUEST) as a pure parser/fold in tinyagentos/gpu_lease.py, plus an authenticated route surface in tinyagentos/routes/a2a_gpu_lease.py: - CHECK folds the channel's open [GPU CLAIM]/[GPU RELEASE] claims AND this controller's own cluster leases (invisible on the bus) and nets both against the node's live VRAM. - CLAIM is admission-checked, registers a real cluster lease (TTL keep-alive via /renew) and posts [GPU CLAIM]; the lease is rolled back if the bus post fails, so the bus and the local scheduler never disagree about who holds the node. - RELEASE frees the lease and posts [GPU RELEASE]; REQUEST posts [GPU REQUEST]. - Claiming is denied (409) on another holder's claim or insufficient VRAM, and CHECK/CLAIM fail closed with 503 when the channel cannot be read rather than reporting a node free. A claim's authoritative holder is the bus message author, not the caller- controlled holder= field. An agent posts as its own registry identity (bus `from` = the token's sub, readable handle in holder=) and presents its registry JWT, matching the bus auth contract from #2112. The new paths are added to the agent Bearer allowlist. Tests: 72 (protocol unit tests + route integration against a faked bus and a real ClusterManager). Fixes #893. --- changelog.d/taos-893-a2a-gpu-lease.md | 16 + docs/agent-coordination.md | 72 +++ tests/test_gpu_lease_protocol.py | 323 +++++++++++++ tests/test_routes_a2a_gpu_lease.py | 604 ++++++++++++++++++++++++ tinyagentos/auth_middleware.py | 16 + tinyagentos/gpu_lease.py | 414 ++++++++++++++++ tinyagentos/routes/__init__.py | 3 + tinyagentos/routes/a2a_gpu_lease.py | 647 ++++++++++++++++++++++++++ 8 files changed, 2095 insertions(+) create mode 100644 changelog.d/taos-893-a2a-gpu-lease.md create mode 100644 tests/test_gpu_lease_protocol.py create mode 100644 tests/test_routes_a2a_gpu_lease.py create mode 100644 tinyagentos/gpu_lease.py create mode 100644 tinyagentos/routes/a2a_gpu_lease.py diff --git a/changelog.d/taos-893-a2a-gpu-lease.md b/changelog.d/taos-893-a2a-gpu-lease.md new file mode 100644 index 000000000..4b58c88cb --- /dev/null +++ b/changelog.d/taos-893-a2a-gpu-lease.md @@ -0,0 +1,16 @@ +### Added + +- Shared-GPU coordination over the A2A bus (taOS #893). Agents sharing one card + can `GET /api/a2a/gpu/check` (folds the channel's open `[GPU CLAIM]`/ + `[GPU RELEASE]` messages into the claims still held and subtracts them from + the node's live free VRAM), then `POST /api/a2a/gpu/claim` for an + admission-checked claim that both registers a real cluster lease (TTL, kept + alive with `POST /api/a2a/gpu/renew`) and posts the `[GPU CLAIM]` line peers + read; `POST /api/a2a/gpu/release` frees it and `POST /api/a2a/gpu/request` + asks for a window when blocked. Claiming refuses (409) on another holder's + open claim or insufficient VRAM, and CHECK/CLAIM return 503 rather than + reporting a node free when the bus cannot be read, so two agents can no + longer silently co-load past the card's VRAM. An agent posts as its own + registry identity (scope `a2a_receive` to check, `a2a_send` to act) and the + node label resolves to a cluster worker's heartbeat VRAM or the controller's + own shared VRAM ledger. diff --git a/docs/agent-coordination.md b/docs/agent-coordination.md index e2ada7bf2..a42dd7f9d 100644 --- a/docs/agent-coordination.md +++ b/docs/agent-coordination.md @@ -351,6 +351,70 @@ Read through the controller with your own registry token, not the raw bus port: If the bus is silent, check `channel_known` and your cursor before concluding nobody is talking. A read that returns `200` with nothing is the failure mode that looks like peace. +## Shared-GPU leases (`/api/a2a/gpu/*`) + +Two agents on one host share a physical GPU and must not silently co-load past +its VRAM (taOS #893). Coordinate over the bus with a one-line text protocol, and +use the controller's endpoints so the protocol is admission-checked and backed by +a real lease: + +``` +[GPU CLAIM] node= holder=@you vram=~9.4gb reason=... eta=... +[GPU RELEASE] node= holder=@you +[GPU REQUEST] node= need=~6gb +``` + +| Method | Path | Scope | Purpose | +|--------|------|-------|---------| +| GET | `/api/a2a/gpu/check` | `a2a_receive` | Fold the channel's open claims + the node's live VRAM and answer "may I load?" | +| POST | `/api/a2a/gpu/claim` | `a2a_send` | Admission-checked claim: cluster lease (TTL) + `[GPU CLAIM]` post | +| POST | `/api/a2a/gpu/release` | `a2a_send` | Release the lease + `[GPU RELEASE]` post | +| POST | `/api/a2a/gpu/request` | `a2a_send` | Post `[GPU REQUEST]` when blocked | +| POST | `/api/a2a/gpu/renew` | `a2a_send` | Keep-alive: extend a lease TTL | + +Do not post a claim line directly to the bus and skip `/claim`: only the endpoint +checks admission (another holder's claim, the node's free VRAM, and the cluster's +own lease table) before the line is posted. A line posted by hand is recorded but +enforces nothing. + +Rules that matter when you use it: + +- **CHECK before load, always.** It folds the channel's `[GPU CLAIM]`/`[GPU + RELEASE]` history into the claims still open, folds in this controller's own + GPU leases (which the bus never shows), and subtracts both from the node's + live free VRAM. A node claimed by ANY other holder is blocked even if the card + looks free, because "claimed" means a load is in flight. +- **Holder identity is the bus author, not the `holder=` text.** The body is + caller-controlled; `from` is what the bus authenticated (see *Posting to the + coordination bus*). The `holder=` field is a readable label for humans. +- **An agent always acts as itself.** `from` is the agent's registry canonical + id, the body's `holder=` is its registry handle, and the caller's own registry + JWT is forwarded to the bus exactly as on `/api/a2a/bus/send`. A `holder` + field in the request body is ignored for agent callers. +- **Fail closed.** If the channel cannot be read, `check` and `claim` return + `503` rather than reporting the node free: an unreadable channel looks exactly + like "nobody has claimed anything". +- **Claim is both halves or neither.** The cluster lease is rolled back if the + bus post fails, so a peer that only watches the bus never disagrees with the + local scheduler about who holds the node. +- **Keep-alive is the TTL, not a promise.** A lease expires after + `ttl_seconds` (default 300) unless renewed via `/renew`; a crashed or idle + holder therefore frees the node without anyone releasing it. +- **A claim is only visible inside the channel fold window** (the newest 500 + messages). For a load that outlives the chatter around it, re-POST `/claim` + periodically: it is idempotent (it extends the lease and reposts the line, + which the fold treats as a replacement, never a second claim). +- **`node` labels** resolve to a cluster worker by name, by its URL host, or to + the local controller for `local`/`localhost`/this hostname. A node this + controller does not know is bus-governed only (no local lease), and its CHECK + is reported as `vram_verified: false` rather than as free. +- The channel defaults to `gpu`; point every agent at the same thread with + `TAOS_A2A_GPU_CHANNEL`. + +`check` returns `admitted`, `blockers`, `free_mb`, `capacity_mb`, `claimed_mb`, +`vram_verified`, `reason`, and the `claims` it folded. `claim` returns the +`lease_id` (when the node is a cluster worker) and the exact `line` posted. + ## Bus restarts during a controller update `POST /api/settings/update` on a host that also runs taOSmd locally (config @@ -1453,6 +1517,14 @@ worker lane is therefore refused on `POST` (it lacks the create grant, `403`) and authorised on `GET`. `tests/test_routes_task_checklist.py` pins this scope split directly, not behind an xfail. +Shared-GPU leases (taOS #893). `GET /api/a2a/gpu/check` (scope `a2a_receive`) +and `POST /api/a2a/gpu/{claim,release,request,renew}` (scope `a2a_send`). The +route resolves the acting identity from the token and forces the bus `from` to +the identity the token proves, so an agent can only claim/release GPU capacity +for itself. See *Shared-GPU leases (`/api/a2a/gpu/*`)* above for the protocol, +the admission rules, and why CHECK/CLAIM fail closed when the channel is +unreadable. + Container provisioning request (P1 + P2, agent-container-provisioning spec): - `POST /api/containers/requests` and `POST /api/container-requests` -- an active agent submits a container provisioning request with its own registry JWT. The route resolves the canonical_id from the token (never from the request body) and applies the provisioning policy (per-agent quota + threshold). Under quota the request is auto-approved; over quota it lands in `pending-approval`; over threshold it is escalated to a Decisions-app item for Jay. This is an identity-only check (no scope grant required), matching the scope-request create flow's use of `check_agent_identity`. diff --git a/tests/test_gpu_lease_protocol.py b/tests/test_gpu_lease_protocol.py new file mode 100644 index 000000000..33d97e8fd --- /dev/null +++ b/tests/test_gpu_lease_protocol.py @@ -0,0 +1,323 @@ +"""Unit tests for the A2A GPU lease protocol (taOS #893). + +The protocol is the cross-product contract between agents sharing one GPU: +CHECK before loading, CLAIM before loading, RELEASE when done. These tests pin +the wire format, the fold (which claims are still open), and the admission +rules that stop a silent co-load past the card's VRAM. +""" +from __future__ import annotations + +from tinyagentos.gpu_lease import ( + CLAIM, + RELEASE, + REQUEST, + Admission, + claims_for_node, + evaluate_admission, + format_vram_mb, + open_claims, + parse_message, + parse_vram_mb, + render_check, + render_claim, + render_release, + render_request, + same_holder, +) + + +class TestVramParsing: + def test_gigabytes_convert_to_mib(self): + assert parse_vram_mb("~9.4gb") == 9626 + assert parse_vram_mb("6gb") == 6144 + assert parse_vram_mb("1GB") == 1024 + + def test_megabytes_and_bare_numbers_are_mib(self): + assert parse_vram_mb("4096mb") == 4096 + assert parse_vram_mb("4096") == 4096 + assert parse_vram_mb("512M") == 512 + + def test_unparseable_is_none_not_zero(self): + # None ("no figure given") must stay distinct from a real 0: a claim + # with a garbled vram figure must not be read as "needs nothing". + assert parse_vram_mb(None) is None + assert parse_vram_mb("lots") is None + assert parse_vram_mb("") is None + + def test_format_round_trips_through_parse(self): + for mb in (512, 1024, 6144, 9626): + assert parse_vram_mb(format_vram_mb(mb)) == mb + + def test_format_uses_gb_above_a_gib(self): + assert format_vram_mb(6144) == "~6gb" + assert format_vram_mb(512) == "512mb" + + +class TestParseMessage: + def test_claim_line_full(self): + msg = parse_message( + "[GPU CLAIM] node=linstation holder=@taOSmd vram=~9.4gb " + "reason=ollama-models eta=~10m" + ) + assert msg is not None + assert msg.kind == CLAIM + assert msg.node == "linstation" + assert msg.holder == "@taOSmd" + assert msg.vram_mb == 9626 + assert msg.reason == "ollama-models" + assert msg.eta == "~10m" + + def test_release_and_request_lines(self): + rel = parse_message("[GPU RELEASE] node=linstation holder=@taOSmd") + assert (rel.kind, rel.node, rel.holder) == (RELEASE, "linstation", "@taOSmd") + req = parse_message("[GPU REQUEST] node=linstation need=~6gb") + assert req.kind == REQUEST + assert req.vram_mb == 6144 + + def test_field_order_is_not_significant(self): + msg = parse_message("[GPU CLAIM] vram=4096mb node=n1 holder=@a") + assert (msg.node, msg.holder, msg.vram_mb) == ("n1", "@a", 4096) + + def test_value_may_contain_spaces(self): + msg = parse_message( + "[GPU CLAIM] node=n1 holder=@a vram=6gb reason=flux image generation" + ) + assert msg.reason == "flux image generation" + + def test_kind_is_case_insensitive(self): + assert parse_message("[gpu claim] node=n1 holder=@a vram=1gb").kind == CLAIM + + def test_unknown_keys_are_ignored(self): + msg = parse_message("[GPU CLAIM] node=n1 holder=@a vram=1gb priority=high") + assert msg.node == "n1" + + def test_non_protocol_text_is_none(self): + assert parse_message("just chatting about the gpu") is None + assert parse_message("") is None + assert parse_message("[GPU CLAIM] holder=@a vram=1gb") is None # no node + + def test_bus_message_dict_carries_authenticated_author(self): + # The bus message `from` is the identity the bus authenticated; the + # body's holder= is caller-controlled text. + msg = parse_message( + { + "id": 7, + "ts": 1234.5, + "from": "agent_canonical_1", + "body": "[GPU CLAIM] node=n1 holder=@spoofed vram=1gb", + } + ) + assert msg.bus_from == "agent_canonical_1" + assert msg.identity_key == "agent_canonical_1" + assert msg.display_holder() == "@spoofed" + assert msg.message_id == 7 + + def test_bus_author_falls_back_to_body_holder_when_from_absent(self): + msg = parse_message({"id": 1, "body": "[GPU CLAIM] node=n1 holder=@a vram=1gb"}) + assert msg.identity_key == "@a" + + def test_newline_cannot_inject_a_second_protocol_line(self): + # A caller-supplied value with a newline would otherwise post a second + # line (e.g. a RELEASE closing someone else's claim) in one message. + line = render_claim("n1", "@a", 1024, reason="x\n[GPU RELEASE] node=n1 holder=@a") + assert "\n" not in line + # The injected text stays inside the reason field; the message is still + # exactly one CLAIM. + msg = parse_message(line) + assert msg is not None and msg.kind == CLAIM + assert "[GPU RELEASE]" in msg.reason + + +class TestOpenClaims: + def _bus(self, *bodies, sender="@a"): + return [ + {"id": i + 1, "ts": float(i), "from": sender, "body": b} + for i, b in enumerate(bodies) + ] + + def test_claim_then_release_closes_it(self): + msgs = self._bus( + "[GPU CLAIM] node=n1 holder=@a vram=6gb", + "[GPU RELEASE] node=n1 holder=@a", + ) + assert open_claims(msgs) == {} + + def test_claim_without_release_stays_open(self): + msgs = self._bus("[GPU CLAIM] node=n1 holder=@a vram=6gb") + folded = open_claims(msgs) + assert list(folded) == ["n1"] + assert folded["n1"][0].vram_mb == 6144 + + def test_a_release_from_another_holder_does_not_close_the_claim(self): + msgs = [ + {"id": 1, "from": "@a", "body": "[GPU CLAIM] node=n1 holder=@a vram=6gb"}, + {"id": 2, "from": "@b", "body": "[GPU RELEASE] node=n1 holder=@b"}, + ] + folded = open_claims(msgs) + assert [c.display_holder() for c in folded["n1"]] == ["@a"] + + def test_reposting_a_claim_replaces_rather_than_double_counts(self): + msgs = self._bus( + "[GPU CLAIM] node=n1 holder=@a vram=6gb", + "[GPU CLAIM] node=n1 holder=@a vram=4gb", + ) + folded = open_claims(msgs) + assert len(folded["n1"]) == 1 + assert folded["n1"][0].vram_mb == 4096 + + def test_two_holders_on_one_node_both_stay_open(self): + msgs = [ + {"id": 1, "from": "@a", "body": "[GPU CLAIM] node=n1 holder=@a vram=6gb"}, + {"id": 2, "from": "@b", "body": "[GPU CLAIM] node=n1 holder=@b vram=2gb"}, + ] + folded = open_claims(msgs) + assert {c.display_holder() for c in folded["n1"]} == {"@a", "@b"} + + def test_requests_and_chatter_do_not_create_claims(self): + msgs = self._bus( + "[GPU REQUEST] node=n1 need=6gb", + "morning all", + "[GPU CHECK] node=n1 need=6gb", + ) + assert open_claims(msgs) == {} + + def test_claims_are_grouped_by_node(self): + msgs = [ + {"id": 1, "from": "@a", "body": "[GPU CLAIM] node=n1 holder=@a vram=1gb"}, + {"id": 2, "from": "@a", "body": "[GPU CLAIM] node=n2 holder=@a vram=1gb"}, + ] + folded = open_claims(msgs) + assert set(folded) == {"n1", "n2"} + assert claims_for_node(folded, "N2")[0].node == "n2" + assert claims_for_node(folded, "n3") == [] + + +class TestSameHolder: + def test_alias_and_canonical_id_match(self): + assert same_holder("agent_abc", "@agent_abc") is True + assert same_holder("@TAOSmd", "taosmd") is True + + def test_empty_never_matches(self): + assert same_holder("", "") is False + assert same_holder(None, "@a") is False + assert same_holder("@a", None) is False + + +class TestEvaluateAdmission: + def _claim(self, sender, vram_mb, node="n1", holder=None): + return parse_message( + { + "id": 1, + "from": sender, + "body": f"[GPU CLAIM] node={node} holder={holder or sender} vram={vram_mb}mb", + } + ) + + def test_free_node_is_admitted(self): + d = evaluate_admission( + node="n1", required_mb=6144, identity="@a", claims=[], free_mb=9000 + ) + assert d.admitted is True + assert d.verified is True + assert d.blockers == () + + def test_another_holders_claim_blocks_even_with_vram_to_spare(self): + d = evaluate_admission( + node="n1", + required_mb=1024, + identity="@a", + claims=[self._claim("@taosmd", 6144)], + free_mb=12288, + ) + assert d.admitted is False + assert d.blockers == ("@taosmd",) + assert "already claimed" in d.reason + + def test_own_claim_does_not_block_and_is_subtracted_from_the_budget(self): + d = evaluate_admission( + node="n1", + required_mb=3072, + identity="@a", + claims=[self._claim("@a", 6144)], + free_mb=8192, + ) + # 8192 free - 6144 already promised to ourselves = 2048 < 3072. + assert d.admitted is False + assert d.claimed_mb == 6144 + assert "insufficient VRAM" in d.reason + + def test_own_claim_matches_through_the_registry_canonical_id(self): + claim = parse_message( + {"id": 1, "from": "agent_abc", "body": "[GPU CLAIM] node=n1 holder=@a vram=6144mb"} + ) + d = evaluate_admission( + node="n1", required_mb=1024, identity="agent_abc", claims=[claim], free_mb=8192 + ) + # Recognised as our own claim (not a blocker) and netted out of the + # budget: 8192 - 6144 = 2048 >= 1024. + assert d.admitted is True + assert d.claimed_mb == 6144 + assert d.blockers == () + + def test_insufficient_vram_is_denied(self): + d = evaluate_admission( + node="n1", required_mb=8192, identity="@a", claims=[], free_mb=4096 + ) + assert d.admitted is False + assert d.free_mb == 4096 + assert "need 8192" in d.reason + + def test_capacity_is_used_when_free_is_unknown(self): + d = evaluate_admission( + node="n1", required_mb=8192, identity="@a", claims=[], capacity_mb=12288 + ) + assert d.admitted is True + d2 = evaluate_admission( + node="n1", required_mb=16384, identity="@a", claims=[], capacity_mb=12288 + ) + assert d2.admitted is False + + def test_unknown_vram_admits_but_flags_unverified(self): + d = evaluate_admission(node="n1", required_mb=6144, identity="@a", claims=[]) + assert d.admitted is True + assert d.verified is False + assert "no VRAM figure" in d.reason + + def test_a_claim_blocks_even_a_zero_vram_check(self): + # "Load only if free + unclaimed": unclaimed applies regardless of size. + d = evaluate_admission( + node="n1", + required_mb=0, + identity="@a", + claims=[self._claim("@b", 1024)], + free_mb=12288, + ) + assert d.admitted is False + + def test_admission_dict_shape(self): + d = evaluate_admission(node="n1", required_mb=1, identity="@a", free_mb=2) + assert isinstance(d, Admission) + body = d.as_dict() + assert body["admitted"] is True and body["vram_verified"] is True + + +class TestRender: + def test_render_claim_includes_reason_and_eta(self): + line = render_claim("n1", "@a", 6144, "ollama", "~10m") + assert line == "[GPU CLAIM] node=n1 holder=@a vram=~6gb reason=ollama eta=~10m" + assert parse_message(line).vram_mb == 6144 + + def test_render_claim_omits_empty_optionals(self): + assert render_claim("n1", "@a", 1024) == "[GPU CLAIM] node=n1 holder=@a vram=~1gb" + + def test_render_release_request_check(self): + assert render_release("n1", "@a") == "[GPU RELEASE] node=n1 holder=@a" + assert render_request("n1", 6144, "blocked") == ( + "[GPU REQUEST] node=n1 need=~6gb reason=blocked" + ) + assert render_check("n1", 6144) == "[GPU CHECK] node=n1 need=~6gb" + + def test_render_flattens_whitespace_in_values(self): + assert render_claim("n 1", "@a", 1024, "x\ty") == ( + "[GPU CLAIM] node=n 1 holder=@a vram=~1gb reason=x y" + ) diff --git a/tests/test_routes_a2a_gpu_lease.py b/tests/test_routes_a2a_gpu_lease.py new file mode 100644 index 000000000..541ba5cc4 --- /dev/null +++ b/tests/test_routes_a2a_gpu_lease.py @@ -0,0 +1,604 @@ +"""Route tests for the A2A GPU lease surface (taOS #893). + +The bus is faked (no network): a ``FakeBus`` serves ``GET /a2a/messages`` from +an in-memory list and appends ``POST /a2a/send`` payloads to it, so a claim made +through the route is visible to the next CHECK exactly as it would be in +production. That is what makes the "claim then check" assertions meaningful +rather than tautological. +""" +from __future__ import annotations + +import pytest +import pytest_asyncio +from httpx import ASGITransport, AsyncClient + +from taos_test_csrf import csrf_event_hooks +from tinyagentos.agent_registry_store import mint_registry_token +from tinyagentos.cluster.manager import ClusterManager +from tinyagentos.cluster.worker_protocol import WorkerInfo + +_ROUTE_PATCH = "tinyagentos.routes.a2a_gpu_lease.httpx.AsyncClient" + + +class FakeBus: + """In-memory stand-in for the raw coordination bus on :7900.""" + + def __init__(self) -> None: + self.messages: list[dict] = [] + self.sends: list[dict] = [] + self.fail_get = False + self.fail_post = False + self._id = 0 + + def seed(self, body: str, sender: str = "@peer") -> dict: + """Append a message as if a peer had posted it.""" + self._id += 1 + msg = { + "id": self._id, + "ts": float(self._id), + "from": sender, + "body": body, + "thread": "gpu", + } + self.messages.append(msg) + return msg + + @property + def last_line(self) -> str | None: + return self.sends[-1]["payload"]["body"] if self.sends else None + + @property + def last_from(self) -> str | None: + return self.sends[-1]["payload"]["from"] if self.sends else None + + @property + def last_headers(self) -> dict: + return self.sends[-1]["headers"] or {} if self.sends else {} + + +def _install_fake_bus(monkeypatch, bus: FakeBus) -> None: + class _Resp: + def __init__(self, payload: dict, status: int = 200) -> None: + self._payload = payload + self.status_code = status + + def raise_for_status(self) -> None: + if self.status_code >= 400: + raise RuntimeError(f"bus returned {self.status_code}") + + def json(self) -> dict: + return self._payload + + class _Client: + def __init__(self, *args, **kwargs) -> None: + pass + + async def __aenter__(self): + return self + + async def __aexit__(self, *exc) -> bool: + return False + + async def get(self, url, params=None): + if bus.fail_get: + raise RuntimeError("bus unreachable") + return _Resp({"messages": list(bus.messages)}) + + async def post(self, url, json=None, headers=None): + if bus.fail_post: + raise RuntimeError("bus unreachable") + payload = dict(json or {}) + bus._id += 1 + msg = { + "id": bus._id, + "ts": float(bus._id), + "from": payload.get("from"), + "body": payload.get("body"), + "thread": payload.get("thread"), + } + bus.sends.append({"payload": payload, "headers": headers}) + bus.messages.append(msg) + return _Resp(msg) + + monkeypatch.setattr(_ROUTE_PATCH, _Client) + + +@pytest.fixture +def bus(monkeypatch) -> FakeBus: + fake = FakeBus() + _install_fake_bus(monkeypatch, fake) + return fake + + +class _StubLedger: + """Stands in for the shared VramReservationManager (taOS #185).""" + + def __init__(self, free_mb: int, total_mb: int) -> None: + self.free_mb = free_mb + self.total_mb = total_mb + + def available_vram(self) -> tuple[int, int]: + return self.free_mb, self.total_mb + + +@pytest.fixture +def local_vram(app): + """Set the local node's reported (free, total) VRAM, in MiB.""" + + def _set(free_mb: int, total_mb: int = 12288) -> None: + app.state.vram_reservation = _StubLedger(free_mb, total_mb) + + return _set + + +@pytest_asyncio.fixture +async def lease_client(app, tmp_data_dir): + """Admin client with the agent registry + grant stores initialised. + + Exposes ``._app`` so tests can register agents / mint tokens and drive bare + (cookieless) requests with a Bearer header. + """ + for attr in ("agent_registry", "agent_grants", "metrics"): + store = getattr(app.state, attr) + if store._db is None: + await store.init() + app.state.auth.setup_user("admin", "Test Admin", "", "testpass") + record = app.state.auth.find_user("admin") + token = app.state.auth.create_session( + user_id=record["id"] if record else "", long_lived=True + ) + app.state._startup_complete = True + async with AsyncClient( + transport=ASGITransport(app=app), + base_url="http://test", + cookies={"taos_session": token}, + event_hooks=csrf_event_hooks(), + ) as c: + c._app = app + yield c + for attr in ("agent_registry", "agent_grants", "metrics"): + store = getattr(app.state, attr) + if store._db is not None: + await store.close() + + +def _bare(app) -> AsyncClient: + return AsyncClient( + transport=ASGITransport(app=app), base_url="http://test", + event_hooks=csrf_event_hooks(), + ) + + +async def _agent_token(app, *, scopes=("a2a_send",), handle="@taosmd"): + """Register an agent, grant scopes, return (canonical_id, jwt).""" + registry = app.state.agent_registry + grants = app.state.agent_grants + priv, _pub = app.state.agent_registry_keypair + rec = await registry.register( + framework="taosmd", display_name="Peer", origin="taos-deployed", handle=handle + ) + cid = rec["canonical_id"] + for scope in scopes: + await grants.add_grant(cid, scope) + return cid, mint_registry_token(cid, priv, user_id="u", framework="taosmd") + + +@pytest_asyncio.fixture +async def cluster(app): + """A real ClusterManager with one online GPU worker (linstation).""" + cm = ClusterManager() + ok, reason = await cm.register_worker( + WorkerInfo( + name="linstation", + url="http://10.0.0.9:9000", + status="online", + free_vram_mb=8192, + hardware={"gpu": {"vram_mb": 12288}}, + resources=["gpu-cuda-0"], + ) + ) + assert ok, reason + app.state.cluster_manager = cm + return cm + + +@pytest.mark.asyncio +class TestCheck: + async def test_missing_node_is_400(self, lease_client, bus): + resp = await lease_client.get("/api/a2a/gpu/check") + assert resp.status_code == 400 + assert bus.sends == [] + + async def test_free_node_is_admitted(self, lease_client, bus, local_vram): + local_vram(9000) + resp = await lease_client.get( + "/api/a2a/gpu/check", params={"node": "local", "vram_mb": 6144} + ) + assert resp.status_code == 200 + data = resp.json() + assert data["admitted"] is True + assert data["vram_verified"] is True + assert data["free_mb"] == 9000 + + async def test_peer_claim_blocks_the_node(self, lease_client, bus, local_vram): + local_vram(11000) + bus.seed("[GPU CLAIM] node=local holder=@taosmd vram=~9.4gb", sender="@taosmd") + resp = await lease_client.get( + "/api/a2a/gpu/check", params={"node": "local", "vram_mb": 2048} + ) + assert resp.status_code == 200 + data = resp.json() + assert data["admitted"] is False + assert data["blockers"] == ["@taosmd"] + assert data["claims"][0]["vram_mb"] == 9626 + assert data["claims"][0]["identity"] == "@taosmd" + + async def test_a_spoofed_body_holder_does_not_evade_the_block( + self, lease_client, bus, local_vram + ): + # Blocking is keyed on the bus-authenticated author, never the + # caller-controlled holder= field in the body. + local_vram(11000) + bus.seed( + "[GPU CLAIM] node=local holder=@operator vram=9gb", sender="agent_peer" + ) + resp = await lease_client.post( + "/api/a2a/gpu/claim", json={"node": "local", "vram_mb": 2048} + ) + assert resp.status_code == 409 + assert resp.json()["claims"][0]["identity"] == "agent_peer" + assert bus.sends == [] + + async def test_released_peer_claim_no_longer_blocks(self, lease_client, bus, local_vram): + local_vram(11000) + bus.seed("[GPU CLAIM] node=local holder=@taosmd vram=9gb", sender="@peer") + bus.seed("[GPU RELEASE] node=local holder=@taosmd", sender="@peer") + resp = await lease_client.get( + "/api/a2a/gpu/check", params={"node": "local", "vram_mb": 2048} + ) + assert resp.json()["admitted"] is True + + async def test_insufficient_vram_is_denied(self, lease_client, bus, local_vram): + local_vram(4096) + resp = await lease_client.get( + "/api/a2a/gpu/check", params={"node": "local", "vram_mb": 8192} + ) + data = resp.json() + assert data["admitted"] is False + assert data["blockers"] == [] + + async def test_unknown_node_admits_but_flags_unverified(self, lease_client, bus): + resp = await lease_client.get( + "/api/a2a/gpu/check", params={"node": "ghostbox", "vram_mb": 4096} + ) + data = resp.json() + assert data["admitted"] is True + assert data["vram_verified"] is False + + async def test_bus_unreadable_fails_closed(self, lease_client, bus): + bus.fail_get = True + resp = await lease_client.get( + "/api/a2a/gpu/check", params={"node": "local", "vram_mb": 4096} + ) + assert resp.status_code == 503 + assert "cannot verify" in resp.json()["detail"] + + async def test_worker_node_uses_heartbeat_vram(self, lease_client, bus, cluster): + resp = await lease_client.get( + "/api/a2a/gpu/check", params={"node": "linstation", "vram_mb": 4096} + ) + data = resp.json() + assert data["admitted"] is True + assert data["free_mb"] == 8192 + assert data["capacity_mb"] == 12288 + + +@pytest.mark.asyncio +class TestClaim: + async def test_claim_without_vram_is_400(self, lease_client, bus): + resp = await lease_client.post("/api/a2a/gpu/claim", json={"node": "local"}) + assert resp.status_code == 400 + assert bus.sends == [] + + async def test_claim_posts_the_protocol_line(self, lease_client, bus, local_vram): + local_vram(11000) + resp = await lease_client.post( + "/api/a2a/gpu/claim", + json={"node": "local", "vram_mb": 6144, "reason": "flux", "eta": "~5m"}, + ) + assert resp.status_code == 200 + data = resp.json() + assert data["status"] == "claimed" + assert data["line"] == ( + "[GPU CLAIM] node=local holder=@operator vram=~6gb reason=flux eta=~5m" + ) + assert bus.last_line == data["line"] + + async def test_claim_then_check_sees_the_claim(self, lease_client, bus, local_vram): + local_vram(11000) + await lease_client.post( + "/api/a2a/gpu/claim", json={"node": "local", "vram_mb": 6144} + ) + # Same caller: its own claim must not block it. + mine = await lease_client.get( + "/api/a2a/gpu/check", params={"node": "local", "vram_mb": 2048} + ) + assert mine.json()["admitted"] is True + # A different caller sees a claimed node. + peer = await lease_client.get( + "/api/a2a/gpu/check", params={"node": "local", "vram_mb": 2048, "channel": "gpu"} + ) + assert peer.status_code == 200 + + async def test_peer_owned_claim_is_denied_409(self, lease_client, bus, local_vram): + local_vram(11000) + bus.seed("[GPU CLAIM] node=local holder=@taosmd vram=9gb", sender="agent_peer") + resp = await lease_client.post( + "/api/a2a/gpu/claim", json={"node": "local", "vram_mb": 2048} + ) + assert resp.status_code == 409 + assert resp.json()["blockers"] == ["@taosmd"] + assert bus.sends == [] # denied claims never reach the channel + + async def test_insufficient_vram_is_denied_409(self, lease_client, bus, local_vram): + local_vram(2048) + resp = await lease_client.post( + "/api/a2a/gpu/claim", json={"node": "local", "vram_mb": 8192} + ) + assert resp.status_code == 409 + assert bus.sends == [] + + async def test_bus_failure_does_not_grant_the_claim(self, lease_client, bus, local_vram): + local_vram(11000) + bus.fail_post = True + resp = await lease_client.post( + "/api/a2a/gpu/claim", json={"node": "local", "vram_mb": 4096} + ) + assert resp.status_code == 502 + + +@pytest.mark.asyncio +class TestClusterLeaseIntegration: + async def test_claim_creates_a_real_lease_and_release_frees_it( + self, lease_client, bus, cluster + ): + resp = await lease_client.post( + "/api/a2a/gpu/claim", + json={"node": "linstation", "vram_mb": 6144, "ttl_seconds": 300}, + ) + assert resp.status_code == 200 + data = resp.json() + assert data["lease_id"], data + leases = cluster.get_leases() + assert len(leases) == 1 + assert leases[0].resource_id == "linstation:gpu-cuda-0" + assert leases[0].required_vram_mb == 6144 + + rel = await lease_client.post( + "/api/a2a/gpu/release", + json={"node": "linstation", "lease_id": data["lease_id"]}, + ) + assert rel.status_code == 200 + assert rel.json()["lease_id"] == data["lease_id"] + assert rel.json()["line"] == "[GPU RELEASE] node=linstation holder=@operator" + assert cluster.get_leases() == [] + + async def test_release_without_a_lease_id_releases_the_callers_own( + self, lease_client, bus, cluster + ): + await lease_client.post( + "/api/a2a/gpu/claim", json={"node": "linstation", "vram_mb": 4096} + ) + assert len(cluster.get_leases()) == 1 + resp = await lease_client.post( + "/api/a2a/gpu/release", json={"node": "linstation"} + ) + assert resp.status_code == 200 + assert resp.json()["lease_id"] is not None + assert cluster.get_leases() == [] + + async def test_release_does_not_free_another_holders_lease( + self, lease_client, bus, cluster + ): + # A lease taken by the scheduler (not by this A2A caller) must survive. + foreign = await cluster.claim_lease( + "linstation:gpu-cuda-0", caller="skald-dispatcher", ttl_seconds=300 + ) + assert foreign is not None + resp = await lease_client.post( + "/api/a2a/gpu/release", json={"node": "linstation"} + ) + assert resp.status_code == 200 + assert resp.json()["lease_id"] is None + assert [l.lease_id for l in cluster.get_leases()] == [foreign.lease_id] + + async def test_reclaiming_extends_rather_than_conflicts(self, lease_client, bus, cluster): + first = await lease_client.post( + "/api/a2a/gpu/claim", json={"node": "linstation", "vram_mb": 4096} + ) + second = await lease_client.post( + "/api/a2a/gpu/claim", json={"node": "linstation", "vram_mb": 4096} + ) + assert second.status_code == 200 + assert second.json()["lease_id"] == first.json()["lease_id"] + assert len(cluster.get_leases()) == 1 + + async def test_lease_claim_is_refused_when_the_worker_lacks_the_vram( + self, lease_client, bus, cluster + ): + resp = await lease_client.post( + "/api/a2a/gpu/claim", json={"node": "linstation", "vram_mb": 16384} + ) + assert resp.status_code == 409 + assert cluster.get_leases() == [] + + async def test_claim_rolls_back_the_lease_when_the_bus_post_fails( + self, lease_client, bus, cluster + ): + bus.fail_post = True + resp = await lease_client.post( + "/api/a2a/gpu/claim", json={"node": "linstation", "vram_mb": 4096} + ) + assert resp.status_code == 502 + assert cluster.get_leases() == [] + + async def test_renew_extends_the_ttl(self, lease_client, bus, cluster): + claim = await lease_client.post( + "/api/a2a/gpu/claim", + json={"node": "linstation", "vram_mb": 4096, "ttl_seconds": 60}, + ) + lease_id = claim.json()["lease_id"] + before = cluster.get_leases()[0].expires_at + resp = await lease_client.post( + "/api/a2a/gpu/renew", json={"lease_id": lease_id, "ttl_seconds": 600} + ) + assert resp.status_code == 200 + assert cluster.get_leases()[0].expires_at > before + + async def test_renew_unknown_lease_is_409(self, lease_client, bus, cluster): + resp = await lease_client.post( + "/api/a2a/gpu/renew", json={"lease_id": "l_nope"} + ) + assert resp.status_code == 409 + + async def test_expired_lease_auto_frees_the_node(self, lease_client, bus, cluster): + resp = await lease_client.post( + "/api/a2a/gpu/claim", + json={"node": "linstation", "vram_mb": 4096, "ttl_seconds": -1}, + ) + assert resp.status_code == 200 + # TTL already elapsed: the node is free again with no explicit release. + nxt = await lease_client.post( + "/api/a2a/gpu/claim", json={"node": "linstation", "vram_mb": 4096} + ) + assert nxt.status_code == 200 + + async def test_check_blocks_when_the_scheduler_holds_a_lease( + self, lease_client, bus, cluster + ): + # A lease the A2A layer did not take (here: the Skald dispatcher) is + # invisible on the bus but is still a real reservation. + lease = await cluster.claim_lease( + "linstation:gpu-cuda-0", caller="skald-dispatcher", ttl_seconds=300 + ) + assert lease is not None + resp = await lease_client.get( + "/api/a2a/gpu/check", params={"node": "linstation", "vram_mb": 2048} + ) + data = resp.json() + assert data["admitted"] is False + assert data["blockers"] == ["skald-dispatcher"] + + async def test_claim_is_denied_while_the_scheduler_holds_a_lease( + self, lease_client, bus, cluster + ): + await cluster.claim_lease( + "linstation:gpu-cuda-0", caller="skald-dispatcher", ttl_seconds=300 + ) + resp = await lease_client.post( + "/api/a2a/gpu/claim", json={"node": "linstation", "vram_mb": 2048} + ) + assert resp.status_code == 409 + assert resp.json()["blockers"] == ["skald-dispatcher"] + assert bus.sends == [] + + async def test_check_counts_the_callers_own_lease_as_its_own( + self, lease_client, bus, cluster + ): + await lease_client.post( + "/api/a2a/gpu/claim", json={"node": "linstation", "vram_mb": 4096} + ) + resp = await lease_client.get( + "/api/a2a/gpu/check", params={"node": "linstation", "vram_mb": 2048} + ) + data = resp.json() + # 8192 free - 4096 promised to ourselves = 4096 >= 2048. + assert data["admitted"] is True + assert data["claimed_mb"] == 4096 + + +@pytest.mark.asyncio +class TestRequest: + async def test_request_posts_a_request_line(self, lease_client, bus): + resp = await lease_client.post( + "/api/a2a/gpu/request", + json={"node": "linstation", "need_mb": 6144, "reason": "blocked"}, + ) + assert resp.status_code == 200 + assert resp.json()["line"] == ( + "[GPU REQUEST] node=linstation need=~6gb reason=blocked" + ) + assert bus.last_line == resp.json()["line"] + + async def test_request_accepts_a_gb_string(self, lease_client, bus): + resp = await lease_client.post( + "/api/a2a/gpu/request", json={"node": "n1", "need": "~2.5gb"} + ) + assert resp.status_code == 200 + assert resp.json()["need_mb"] == 2560 + + async def test_request_without_need_is_400(self, lease_client, bus): + resp = await lease_client.post("/api/a2a/gpu/request", json={"node": "n1"}) + assert resp.status_code == 400 + assert bus.sends == [] + + +@pytest.mark.asyncio +class TestAgentToken: + async def test_agent_claims_as_its_registry_identity( + self, lease_client, bus + ): + cid, token = await _agent_token(lease_client._app, scopes=("a2a_send",)) + async with _bare(lease_client._app) as bare: + resp = await bare.post( + "/api/a2a/gpu/claim", + json={"node": "remote-node", "vram_mb": 4096, "holder": "@spoofed"}, + headers={"Authorization": f"Bearer {token}"}, + ) + assert resp.status_code == 200 + # The bus `from` is the identity the token proves, never the body field. + assert bus.last_from == cid + assert "holder=@taosmd" in bus.last_line + assert "@spoofed" not in bus.last_line + # The credential travels with the attribution. + assert bus.last_headers.get("Authorization") == f"Bearer {token}" + + async def test_agent_without_send_scope_cannot_claim(self, lease_client, bus): + _cid, token = await _agent_token(lease_client._app, scopes=("a2a_receive",)) + async with _bare(lease_client._app) as bare: + resp = await bare.post( + "/api/a2a/gpu/claim", + json={"node": "n1", "vram_mb": 4096}, + headers={"Authorization": f"Bearer {token}"}, + ) + assert resp.status_code == 403 + assert bus.sends == [] + + async def test_agent_check_needs_receive_scope(self, lease_client, bus): + _cid, token = await _agent_token(lease_client._app, scopes=("a2a_send",)) + async with _bare(lease_client._app) as bare: + resp = await bare.get( + "/api/a2a/gpu/check", + params={"node": "n1"}, + headers={"Authorization": f"Bearer {token}"}, + ) + assert resp.status_code == 403 + + async def test_agent_with_receive_scope_can_check(self, lease_client, bus): + _cid, token = await _agent_token(lease_client._app, scopes=("a2a_receive",)) + async with _bare(lease_client._app) as bare: + resp = await bare.get( + "/api/a2a/gpu/check", + params={"node": "n1", "vram_mb": 1024}, + headers={"Authorization": f"Bearer {token}"}, + ) + assert resp.status_code == 200 + assert resp.json()["holder"] == "@taosmd" + + async def test_agent_token_is_not_a_skeleton_key(self, lease_client, bus): + _cid, token = await _agent_token(lease_client._app, scopes=("a2a_send",)) + async with _bare(lease_client._app) as bare: + resp = await bare.get( + "/api/cluster/leases", headers={"Authorization": f"Bearer {token}"} + ) + assert resp.status_code in (401, 403) diff --git a/tinyagentos/auth_middleware.py b/tinyagentos/auth_middleware.py index 1610055e8..78e8cef10 100644 --- a/tinyagentos/auth_middleware.py +++ b/tinyagentos/auth_middleware.py @@ -46,6 +46,20 @@ _A2A_BUS_WRITE_PATHS = frozenset({ "/api/a2a/bus/send", }) +# A2A GPU lease coordination (taOS #893). Read (CHECK) needs a2a_receive; the +# mutating actions (CLAIM/RELEASE/REQUEST/RENEW) need a2a_send. The route +# resolves the acting identity from the token itself and forces the bus `from` +# to the identity the token proves, so an agent can only claim/release for +# itself. Same passthrough contract as the bus paths above. +_A2A_GPU_READ_PATHS = frozenset({ + "/api/a2a/gpu/check", +}) +_A2A_GPU_WRITE_PATHS = frozenset({ + "/api/a2a/gpu/claim", + "/api/a2a/gpu/release", + "/api/a2a/gpu/request", + "/api/a2a/gpu/renew", +}) # Observatory routes an agent may reach with its own registry JWT (scope # observatory_control). The route verifies the JWT + grant itself; the # middleware only passes the Bearer through. Admin/local-token is handled @@ -82,6 +96,8 @@ _REGISTRY_FEED_PATHS | _A2A_BUS_READ_PATHS | _A2A_BUS_WRITE_PATHS + | _A2A_GPU_READ_PATHS + | _A2A_GPU_WRITE_PATHS | _OBSERVATORY_PATHS | _CONTAINER_REQUEST_PATHS | frozenset({"/api/agents/me/models", "/api/agents/me/model"}) diff --git a/tinyagentos/gpu_lease.py b/tinyagentos/gpu_lease.py new file mode 100644 index 000000000..9d9ee75b6 --- /dev/null +++ b/tinyagentos/gpu_lease.py @@ -0,0 +1,414 @@ +# tinyagentos/gpu_lease.py +"""A2A GPU lease protocol — CHECK / CLAIM / RELEASE / REQUEST over the bus. + +taOS #893: two agents (@taOS and @taOSmd) share one physical GPU and must not +silently co-load past its VRAM. Found when a FLUX image generation OOM-killed +itself because the other agent had ~9.4 GB of Ollama models loaded on the same +card despite an earlier "free" signal. + +The coordination channel is the A2A bus (see ``routes/a2a_bus.py``) and the wire +format is a single readable line so a human watching the channel can follow it:: + + [GPU CLAIM] node=linstation holder=@taOSmd vram=~9.4gb reason=ollama eta=~10m + [GPU RELEASE] node=linstation holder=@taOSmd + [GPU REQUEST] node=linstation need=~6gb + [GPU CHECK] node=linstation need=~6gb + +This module owns the FORMAT and the FOLD: parsing a line, rendering a line, and +reducing a channel's message history to the claims still open (a CLAIM is closed +by a later RELEASE from the same holder). It is deliberately pure — no FastAPI, +no httpx — so the admission rules can be tested without a bus. + +Identity +-------- +The authoritative holder of a claim is the BUS MESSAGE AUTHOR (``from``), not +the ``holder=`` field in the body. The body is caller-controlled text; ``from`` +is the identity the bus authenticated (taOS #2112, ``routes/a2a_bus.py``). +``holder=`` is kept in the rendered line for human readability, and is also +accepted when matching an identity so the interim handle-spelled posts that +predate bus auth still close correctly. Bus ids are matched first and wins. +""" + +from __future__ import annotations + +import re +from dataclasses import dataclass +from typing import Iterable, Mapping + +# Message kinds. +CLAIM = "CLAIM" +RELEASE = "RELEASE" +REQUEST = "REQUEST" +CHECK = "CHECK" +_KINDS = (CLAIM, RELEASE, REQUEST, CHECK) + +_HEADER_RE = re.compile( + r"^\s*\[\s*GPU\s+(?PCLAIM|RELEASE|REQUEST|CHECK)\s*\]\s*(?P.*)$", + re.IGNORECASE | re.DOTALL, +) + +# A field starts at the beginning of the string or after whitespace. Values can +# contain spaces (`reason=flux image gen`), so a field's value runs until the +# next `key=` token rather than to the next whitespace. +_FIELD_START_RE = re.compile(r"(?:^|\s)(?P[A-Za-z_][A-Za-z0-9_]*)\s*=") + +_VRAM_RE = re.compile( + r"^\s*~?\s*(?P\d+(?:\.\d+)?)\s*(?Pgib|gb|g|mib|mb|m)?\s*$", + re.IGNORECASE, +) + +# Field names understood in a line body. Unknown keys are ignored (forward +# compatibility: a future `priority=` must not break an older reader). +_FIELD_ALIASES = {"needed": "need", "vram_mb": "vram"} + + +def _clean(text: object) -> str: + """Collapse *text* to a single printable line. + + Every field value is caller-supplied text that ends up inside a one-line bus + message. A raw newline would let a caller inject a SECOND protocol line + (e.g. a RELEASE that closes another holder's claim), so all whitespace runs + collapse to a single space and non-printable control characters are dropped. + """ + keep = "".join( + ch if (ch.isprintable() or ch.isspace()) else "" for ch in str(text) + ) + return " ".join(keep.split()) + + +def _split_fields(rest: str) -> dict[str, str]: + """Return ``key=value`` pairs from a line body, values allowed spaces.""" + matches = list(_FIELD_START_RE.finditer(rest)) + fields: dict[str, str] = {} + for i, m in enumerate(matches): + start = m.end() + end = matches[i + 1].start() if i + 1 < len(matches) else len(rest) + key = m.group("key").lower() + fields[_FIELD_ALIASES.get(key, key)] = rest[start:end].strip() + return fields + + +def parse_vram_mb(text: object) -> int | None: + """Parse a VRAM figure (``"~9.4gb"``, ``"4096mb"``, ``"6"``) into MiB. + + A unit-less number is MiB, matching the internal convention. Returns None + for anything unparseable so a caller can tell "no figure given" from "zero". + """ + if text is None: + return None + m = _VRAM_RE.match(str(text)) + if m is None: + return None + num = float(m.group("num")) + unit = (m.group("unit") or "mb").lower() + if unit in ("gib", "gb", "g"): + return int(round(num * 1024)) + return int(round(num)) + + +def format_vram_mb(mb: int | None) -> str: + """Render MiB as a bus-friendly figure (``~9.4gb`` / ``512mb``).""" + if mb is None: + return "?" + if mb >= 1024: + gb = f"{mb / 1024:.1f}".rstrip("0").rstrip(".") + return f"~{gb}gb" + return f"{mb}mb" + + +def same_holder(a: str | None, b: str | None) -> bool: + """True when two identity spellings name the same holder. + + Case-insensitive and ``@``-insensitive so a registry canonical_id and the + readable ``@handle`` alias of the same agent compare equal. Nil/empty + identities never match anything (an unidentified claim must not be + mistaken for the caller's own). + """ + if not a or not b: + return False + return a.strip().lstrip("@").casefold() == b.strip().lstrip("@").casefold() + + +@dataclass(frozen=True) +class GpuLeaseMessage: + """One parsed ``[GPU ...]`` line plus the bus metadata it arrived with.""" + + kind: str + node: str + holder: str = "" + vram_mb: int | None = None + reason: str = "" + eta: str = "" + raw: str = "" + message_id: int | None = None + ts: float | None = None + bus_from: str | None = None + + @property + def identity_key(self) -> str: + """The authenticated author when known, else the body's holder.""" + return (self.bus_from or self.holder or "").strip() + + def display_holder(self) -> str: + """Readable holder for a rendered line (never empty).""" + return (self.holder or self.bus_from or "@unknown").strip() + + def as_dict(self) -> dict: + return { + "kind": self.kind, + "node": self.node, + "holder": self.display_holder(), + "identity": self.identity_key, + "vram_mb": self.vram_mb, + "reason": self.reason, + "eta": self.eta, + "message_id": self.message_id, + "ts": self.ts, + } + + +def parse_message( + raw: object, + *, + sender: str | None = None, + message_id: int | None = None, + ts: float | None = None, +) -> GpuLeaseMessage | None: + """Parse a bus message (dict) or a bare line body into a message. + + Returns None when *raw* is not a GPU lease line or carries no ``node=``. + A dict is expected in the bus shape (``id``/``ts``/``from``/``body``); its + ``from`` becomes the authoritative ``bus_from`` identity. + """ + if isinstance(raw, GpuLeaseMessage): + return raw + + body: str + if isinstance(raw, Mapping): + body = str(raw.get("body") or raw.get("text") or "") + sender = raw.get("from") or sender + message_id = raw.get("id", message_id) + ts = raw.get("ts", ts) + else: + body = str(raw) + + m = _HEADER_RE.match(body) + if m is None: + return None + fields = _split_fields(m.group("rest")) + node = fields.get("node", "").strip() + if not node: + return None + kind = m.group("kind").upper() + vram_raw = fields.get("vram") + if vram_raw is None and kind in (REQUEST, CHECK): + vram_raw = fields.get("need") + return GpuLeaseMessage( + kind=kind, + node=node, + holder=_clean(fields.get("holder", ""))[:64], + vram_mb=parse_vram_mb(vram_raw) if vram_raw is not None else None, + reason=_clean(fields.get("reason", ""))[:200], + eta=_clean(fields.get("eta", ""))[:64], + raw=body.strip(), + message_id=int(message_id) if isinstance(message_id, int) else None, + ts=float(ts) if isinstance(ts, (int, float)) else None, + bus_from=sender.strip() if isinstance(sender, str) and sender.strip() else None, + ) + + +def render_claim( + node: str, holder: str, vram_mb: int, reason: str = "", eta: str = "" +) -> str: + """Render a ``[GPU CLAIM]`` line. Every value is flattened to one line.""" + parts = [ + f"node={_clean(node)}", + f"holder={_clean(holder)}", + f"vram={format_vram_mb(vram_mb)}", + ] + if reason: + parts.append(f"reason={_clean(reason)}") + if eta: + parts.append(f"eta={_clean(eta)}") + return f"[GPU CLAIM] {' '.join(parts)}" + + +def render_release(node: str, holder: str) -> str: + """Render a ``[GPU RELEASE]`` line.""" + return f"[GPU RELEASE] node={_clean(node)} holder={_clean(holder)}" + + +def render_request(node: str, need_mb: int, reason: str = "") -> str: + """Render a ``[GPU REQUEST]`` line.""" + parts = [f"node={_clean(node)}", f"need={format_vram_mb(need_mb)}"] + if reason: + parts.append(f"reason={_clean(reason)}") + return f"[GPU REQUEST] {' '.join(parts)}" + + +def render_check(node: str, need_mb: int | None = None) -> str: + """Render a ``[GPU CHECK]`` line (informational; CHECK is normally local).""" + parts = [f"node={_clean(node)}"] + if need_mb is not None: + parts.append(f"need={format_vram_mb(need_mb)}") + return f"[GPU CHECK] {' '.join(parts)}" + + +def open_claims(messages: Iterable[object]) -> dict[str, list[GpuLeaseMessage]]: + """Fold a channel's messages into the claims that are still open. + + Messages are expected oldest-first (the bus returns them in that order). A + CLAIM opens a slot keyed by ``(node, identity)``; a RELEASE from the same + identity on the same node closes it. A re-CLAIM by the same identity + replaces the previous one (so a reposted claim is never double-counted). + """ + open_by_key: dict[tuple[str, str], GpuLeaseMessage] = {} + order: list[tuple[str, str]] = [] + for raw in messages: + msg = parse_message(raw) + if msg is None or msg.kind not in (CLAIM, RELEASE): + continue + key = (msg.node, msg.identity_key.casefold()) + if msg.kind == CLAIM: + if key not in open_by_key: + order.append(key) + open_by_key[key] = msg + else: + open_by_key.pop(key, None) + + out: dict[str, list[GpuLeaseMessage]] = {} + for key in order: + msg = open_by_key.get(key) + if msg is not None: + out.setdefault(msg.node, []).append(msg) + return out + + +def claims_for_node( + claims: Mapping[str, list[GpuLeaseMessage]], node: str +) -> list[GpuLeaseMessage]: + """Claims for a node, matching the node case-insensitively.""" + wanted = (node or "").strip().casefold() + for name, entries in claims.items(): + if name.strip().casefold() == wanted: + return list(entries) + return [] + + +def _is_mine(claim: GpuLeaseMessage, identity: str) -> bool: + return same_holder(claim.identity_key, identity) or same_holder( + claim.holder, identity + ) + + +@dataclass(frozen=True) +class Admission: + """The outcome of a CHECK: may this holder load on this node?""" + + admitted: bool + node: str + required_mb: int + claimed_mb: int = 0 + free_mb: int | None = None + capacity_mb: int | None = None + blockers: tuple[str, ...] = () + reason: str | None = None + verified: bool = True + + def as_dict(self) -> dict: + return { + "admitted": self.admitted, + "node": self.node, + "required_mb": self.required_mb, + "claimed_mb": self.claimed_mb, + "free_mb": self.free_mb, + "capacity_mb": self.capacity_mb, + "blockers": list(self.blockers), + "reason": self.reason, + "vram_verified": self.verified, + } + + +def evaluate_admission( + *, + node: str, + required_mb: int, + identity: str, + claims: Iterable[GpuLeaseMessage] = (), + free_mb: int | None = None, + capacity_mb: int | None = None, +) -> Admission: + """Decide whether *identity* may use *node* for *required_mb* of VRAM. + + Rules, in order: + + 1. A claim by ANOTHER holder blocks the node outright — the issue's rule is + "load only if free + unclaimed", independent of how much VRAM is free. + 2. With no other holder, the caller's remaining need is checked against the + node's budget: live free VRAM when known, else total capacity. The + caller's own prior claim is subtracted (it is a pending load not yet + reflected in a physical probe) so a re-check cannot over-admit. + 3. When neither figure is known the claim is admitted but flagged + ``verified=False`` with a reason: a node that cannot report VRAM is a + gap in the guarantee, and the caller must probe before loading rather + than be told a silent "free". + """ + required_mb = max(0, int(required_mb)) + entries = list(claims) + others = [c for c in entries if not _is_mine(c, identity)] + own_mb = sum((c.vram_mb or 0) for c in entries if _is_mine(c, identity)) + + if others: + blockers = tuple( + sorted({c.display_holder() for c in others if c.display_holder()}) + ) + return Admission( + admitted=False, + node=node, + required_mb=required_mb, + claimed_mb=own_mb, + free_mb=free_mb, + capacity_mb=capacity_mb, + blockers=blockers, + reason=( + f"{node} is already claimed by {', '.join(blockers)}; " + "release it or post a [GPU REQUEST] to negotiate a window" + ), + ) + + budget = free_mb if free_mb is not None else capacity_mb + if budget is None: + return Admission( + admitted=True, + node=node, + required_mb=required_mb, + claimed_mb=own_mb, + verified=False, + reason=( + f"{node} reports no VRAM figure — claim admitted unverified; " + "probe nvidia-smi before loading" + ), + ) + + available = max(0, int(budget) - own_mb) + if required_mb > available: + return Admission( + admitted=False, + node=node, + required_mb=required_mb, + claimed_mb=own_mb, + free_mb=free_mb, + capacity_mb=capacity_mb, + reason=( + f"insufficient VRAM on {node}: need {required_mb} MiB, " + f"{available} MiB available" + ), + ) + return Admission( + admitted=True, + node=node, + required_mb=required_mb, + claimed_mb=own_mb, + free_mb=free_mb, + capacity_mb=capacity_mb, + ) diff --git a/tinyagentos/routes/__init__.py b/tinyagentos/routes/__init__.py index cf090d7c4..6458b6856 100644 --- a/tinyagentos/routes/__init__.py +++ b/tinyagentos/routes/__init__.py @@ -120,6 +120,9 @@ def register_all_routers(app): from tinyagentos.routes.a2a_bus import router as a2a_bus_router app.include_router(a2a_bus_router, dependencies=_csrf) + from tinyagentos.routes.a2a_gpu_lease import router as a2a_gpu_lease_router + app.include_router(a2a_gpu_lease_router, dependencies=_csrf) + from tinyagentos.routes.scheduler import router as scheduler_router app.include_router(scheduler_router, dependencies=_csrf) diff --git a/tinyagentos/routes/a2a_gpu_lease.py b/tinyagentos/routes/a2a_gpu_lease.py new file mode 100644 index 000000000..228820430 --- /dev/null +++ b/tinyagentos/routes/a2a_gpu_lease.py @@ -0,0 +1,647 @@ +# tinyagentos/routes/a2a_gpu_lease.py +"""A2A GPU lease endpoints — CHECK / CLAIM / RELEASE / REQUEST / RENEW. + +taOS #893. Two agents sharing one GPU coordinate over the A2A bus with the +one-line text protocol owned by ``tinyagentos/gpu_lease.py``. These endpoints +are the authenticated half: + +* ``GET /api/a2a/gpu/check`` — read the channel, fold open claims, combine + them with the node's live VRAM, and answer "may I load?". +* ``POST /api/a2a/gpu/claim`` — admission-checked claim: registers a real + GPU lease in the cluster manager (TTL + auto-expiry, the "keep-alive + auto-eviction" half of the issue) AND posts the ``[GPU CLAIM]`` line so + peers that only see the bus are coordinated too. +* ``POST /api/a2a/gpu/release`` — release the lease and post ``[GPU RELEASE]``. +* ``POST /api/a2a/gpu/request`` — post ``[GPU REQUEST]`` when blocked. +* ``POST /api/a2a/gpu/renew`` — keep-alive for a long-running load. + +Why both halves +--------------- +The bus claim is what a peer product (@taOSmd) can see; the cluster lease is +what THIS controller's scheduler enforces (``GpuArbiter`` already subtracts +``ClusterManager.get_leases()`` when admitting cluster work, taOS #1705). A +claim is only real when both exist, so a failed bus post rolls the lease back +rather than leaving a reservation no peer knows about. + +Identity and credentials +------------------------ +An agent caller posts as its OWN registry identity: the bus ``from`` is the +token's ``sub`` and the caller's registry JWT travels with it, because the bus +verifies the signature and then requires ``token sub == from`` (taOS #2112, +``routes/a2a_bus.py``). The readable ``holder=`` in the body is the agent's +registry handle. An admin session may post as an explicit handle and forwards +no credential. ``from`` is never taken from the request body for an agent. + +Fail closed +----------- +If the channel cannot be read, CHECK and CLAIM return 503 rather than reporting +the node free: an unreadable channel is indistinguishable from "everyone else's +claims are invisible", which is exactly the silent co-load the issue is about. +""" + +from __future__ import annotations + +import asyncio +import logging +import os +import socket +from dataclasses import dataclass + +import httpx +from fastapi import APIRouter, HTTPException, Request +from fastapi.responses import JSONResponse +from pydantic import BaseModel + +from tinyagentos.agent_token_auth import check_agent_scope +from tinyagentos.gpu_lease import ( + CLAIM, + GpuLeaseMessage, + claims_for_node, + evaluate_admission, + open_claims, + parse_vram_mb, + render_claim, + render_release, + render_request, +) +from tinyagentos.routes.a2a_bus import _bus_url, _credential_may_cross + +logger = logging.getLogger(__name__) + +router = APIRouter() + +# Coordination channel. The issue does not name one; "gpu" is the default and an +# operator can point every agent at the same thread with TAOS_A2A_GPU_CHANNEL. +_DEFAULT_CHANNEL = "gpu" +_DEFAULT_RESOURCE = "gpu-cuda-0" + +# The channel is folded on every CHECK/CLAIM, so the read window has to be deep +# enough to contain the CLAIM that is still open. 500 is the largest page the +# proxy exposes; a claim that scrolls out of this window stops being visible, so +# a long-running load re-POSTs /claim (idempotent) to refresh its position. +_CHANNEL_LIMIT = 500 + +# Default lease TTL. Long enough for a model load to start and renew; short +# enough that a crashed holder stops blocking the node (issue: advisory keep-alive). +_DEFAULT_TTL_SECONDS = 300.0 + + +def _channel() -> str: + return os.environ.get("TAOS_A2A_GPU_CHANNEL", _DEFAULT_CHANNEL).strip() or _DEFAULT_CHANNEL + + +def _clean_handle(raw: str | None) -> str: + """Flatten a caller-supplied handle to one printable token.""" + s = "".join(c for c in (raw or "") if c.isprintable()) + return " ".join(s.split())[:64] + + +@dataclass(frozen=True) +class _Actor: + """Who a lease action is attributed to, and the credential proving it.""" + + identity: str # the bus `from` (canonical_id for agents) + holder: str # readable handle for the body's holder= field + credential: str | None = None + is_admin: bool = False + + +def _bearer_token(request: Request) -> str | None: + """Return the caller's raw Bearer credential, verbatim, or None.""" + auth_header = request.headers.get("Authorization", "") + if not auth_header.lower().startswith("bearer "): + return None + return auth_header[7:].strip() or None + + +async def _resolve_actor( + request: Request, body_holder: str | None, scope: str +) -> _Actor: + """Resolve the acting identity, or raise 401/403 (fail closed).""" + if getattr(request.state, "is_admin", False): + holder = _clean_handle(body_holder) or "@operator" + return _Actor(identity=holder, holder=holder, is_admin=True) + + caller = await check_agent_scope(request, scope) + if caller is None: + raise HTTPException(status_code=403, detail="forbidden") + + registry = getattr(request.app.state, "agent_registry", None) + record = await registry.get(caller) if registry is not None else None + handle = _clean_handle((record or {}).get("handle")) + if not handle: + raise HTTPException(status_code=403, detail="agent has no bus handle") + if not handle.startswith("@"): + handle = f"@{handle}" + # The bus `from` must be the identity the token proves, and the token must + # travel with it -- see the module docstring. + return _Actor(identity=caller, holder=handle, credential=_bearer_token(request)) + + +async def _read_channel(channel: str) -> list[dict]: + """Fetch the channel's messages oldest-first. Raises on an unreadable bus.""" + bus = _bus_url() + async with httpx.AsyncClient(timeout=5.0) as client: + resp = await client.get( + f"{bus}/a2a/messages", params={"thread": channel, "limit": _CHANNEL_LIMIT} + ) + resp.raise_for_status() + data = resp.json() + if not isinstance(data, dict): + raise ValueError(f"unexpected bus payload: {type(data).__name__}") + messages = data.get("messages") + if not isinstance(messages, list): + raise ValueError("bus payload has no messages list") + return [m for m in messages if isinstance(m, dict)] + + +async def _post_line(channel: str, actor: _Actor, text: str) -> dict: + """Post a protocol line to the channel. Raises HTTPException(502) on failure.""" + headers: dict[str, str] = {} + if actor.credential: + bus = _bus_url() + if _credential_may_cross(bus): + headers["Authorization"] = f"Bearer {actor.credential}" + else: + logger.warning( + "A2A GPU lease credential withheld for non-loopback http destination %s", + bus, + ) + bus = _bus_url() + payload = {"from": actor.identity, "thread": channel, "body": text} + try: + async with httpx.AsyncClient(timeout=5.0) as client: + resp = await client.post(f"{bus}/a2a/send", json=payload, headers=headers or None) + resp.raise_for_status() + data = resp.json() + except Exception as exc: # noqa: BLE001 + logger.warning("A2A GPU lease post failed (%s): %s", bus, exc) + raise HTTPException(status_code=502, detail="a2a bus unavailable") + return data if isinstance(data, dict) else {} + + +async def _folded_claims(request: Request, channel: str) -> dict[str, list[GpuLeaseMessage]]: + """Read the channel and fold it, mapping an unreadable bus to 503.""" + try: + messages = await _read_channel(channel) + except Exception as exc: # noqa: BLE001 + logger.warning("A2A GPU lease channel read failed (channel=%s): %s", channel, exc) + raise HTTPException( + status_code=503, + detail=( + "a2a bus unavailable — cannot verify GPU claims, so the node " + "cannot be reported free" + ), + ) + return open_claims(messages) + + +def _match_worker(cluster, node: str): + """Return the cluster worker *node* names, or None. + + Accepts the worker name, a URL host, or the local controller's own aliases + ("local", "localhost", this hostname) so a bus node label resolves to the + worker whose VRAM the label refers to. + """ + if cluster is None: + return None + worker = cluster.get_worker(node) + if worker is not None: + return worker + wanted = (node or "").strip().casefold() + if not wanted: + return None + if wanted in ("local", "localhost", socket.gethostname().casefold()): + return cluster.get_worker("local") + for w in cluster.get_workers(): + url = getattr(w, "url", "") or "" + host = url.split("//", 1)[-1].split("/", 1)[0].split(":", 1)[0] + if host and host.casefold() == wanted: + return w + return None + + +def _worker_capacity_mb(worker) -> int | None: + hw = getattr(worker, "hardware", None) + if not isinstance(hw, dict): + return None + gpu = hw.get("gpu") + if not isinstance(gpu, dict): + return None + vram = gpu.get("vram_mb") + try: + vram = int(vram) + except (TypeError, ValueError): + return None + return vram if vram > 0 else None + + +async def _vram_for_node(request: Request, node: str) -> tuple[int | None, int | None]: + """Return ``(free_mb, capacity_mb)`` for *node*; either may be None. + + A cluster worker's last-heartbeat figures win. The local controller uses + the shared VRAM ledger (which nets out in-flight reservations, taOS #185) + when it is available, falling back to a live nvidia-smi probe. + """ + cluster = getattr(request.app.state, "cluster_manager", None) + worker = _match_worker(cluster, node) + if worker is not None: + free = getattr(worker, "free_vram_mb", None) + capacity = _worker_capacity_mb(worker) + if getattr(worker, "name", "") == "local": + local_free, local_total = await _local_vram(request) + return (local_free if local_free is not None else free, capacity or local_total) + return free, capacity + + if (node or "").strip().casefold() in ( + "local", + "localhost", + socket.gethostname().casefold(), + ): + return await _local_vram(request) + return None, None + + +async def _local_vram(request: Request) -> tuple[int | None, int | None]: + """Probe this host's VRAM: (free_mb, total_mb), either possibly None.""" + ledger = getattr(request.app.state, "vram_reservation", None) + if ledger is not None: + try: + free, total = await asyncio.to_thread(ledger.available_vram) + return (free if total and total > 0 else None, total if total and total > 0 else None) + except Exception: # noqa: BLE001 (fall through to a raw probe) + logger.debug("gpu-lease: ledger probe failed", exc_info=True) + try: + from tinyagentos.system_stats import read_nvidia_vram + + pair = await asyncio.to_thread(read_nvidia_vram) + if pair is not None: + used, total = pair + return max(0, total - used), total + except Exception: # noqa: BLE001 + logger.debug("gpu-lease: nvidia-smi probe failed", exc_info=True) + return None, None + + +def _resolve_vram_mb(vram_mb: int | None, vram: str | None) -> int | None: + """Resolve an explicit MiB figure or a protocol-style ``vram`` string.""" + if vram_mb is not None: + return max(0, int(vram_mb)) + if vram: + return parse_vram_mb(vram) + return None + + +def _cluster_lease_claims(cluster, node: str, resource: str) -> list[GpuLeaseMessage]: + """Represent this controller's own GPU leases on *node* as open claims. + + The bus only carries what peers posted; a lease taken by the local + scheduler (``skald-dispatcher``, the GPU arbiter) is invisible there but is + just as real a reservation. Folding the two together is what makes CHECK + account for pending local loads whose VRAM a heartbeat has not yet seen + (the gap ``GpuArbiter._check_cluster_admission`` closes for the scheduler, + taOS #1705). + + ``caller`` is stripped of the ``a2a:`` prefix our own claims add so a + caller's own lease matches its identity instead of blocking it. + + A holder that already has an open BUS claim is skipped: that claim and its + lease are the same reservation seen twice, and charging both would double + the VRAM it holds. + """ + if cluster is None: + return [] + resource_id = _resource_id(node, resource) + out: list[GpuLeaseMessage] = [] + seen: set[str] = set() + for lease in cluster.get_leases(): + if lease.resource_id != resource_id: + continue + caller = lease.caller or "" + who = caller[4:] if caller.startswith("a2a:") else caller + if who.casefold() in seen: + continue + seen.add(who.casefold()) + out.append( + GpuLeaseMessage( + kind=CLAIM, + node=node, + holder=who, + vram_mb=(lease.required_vram_mb or None), + reason="cluster lease", + bus_from=who, + ) + ) + return out + + +async def _check_node( + request: Request, + *, + node: str, + required_mb: int, + identity: str, + channel: str, + resource: str = _DEFAULT_RESOURCE, +) -> tuple[dict, list[GpuLeaseMessage]]: + """Run the full CHECK for a node; returns (admission dict, node claims).""" + folded = await _folded_claims(request, channel) + cluster = getattr(request.app.state, "cluster_manager", None) + bus_claims = claims_for_node(folded, node) + # An A2A claim and the cluster lease it created are the same reservation; + # charge it once (the bus line is the holder's own declared figure). + on_bus = set() + for c in bus_claims: + on_bus.add(c.identity_key.casefold()) + on_bus.add(c.holder.casefold()) + node_claims = bus_claims + [ + c + for c in _cluster_lease_claims(cluster, node, resource) + if c.holder.casefold() not in on_bus + ] + free_mb, capacity_mb = await _vram_for_node(request, node) + admission = evaluate_admission( + node=node, + required_mb=required_mb, + identity=identity, + claims=node_claims, + free_mb=free_mb, + capacity_mb=capacity_mb, + ) + body = admission.as_dict() + body["claims"] = [c.as_dict() for c in node_claims] + body["channel"] = channel + return body, node_claims + + +# ── request bodies ──────────────────────────────────────────────────────────── + + +class _LeaseBody(BaseModel): + node: str + channel: str | None = None + resource: str = _DEFAULT_RESOURCE + + +class ClaimBody(_LeaseBody): + vram_mb: int | None = None + vram: str | None = None + reason: str = "" + eta: str = "" + ttl_seconds: float = _DEFAULT_TTL_SECONDS + holder: str | None = None # honored for admin callers only + + +class ReleaseBody(_LeaseBody): + lease_id: str | None = None + holder: str | None = None + + +class RequestBody(_LeaseBody): + need_mb: int | None = None + need: str | None = None + reason: str = "" + holder: str | None = None + + +class RenewBody(BaseModel): + lease_id: str + ttl_seconds: float = _DEFAULT_TTL_SECONDS + + +def _resource_id(node: str, resource: str) -> str: + res = (resource or _DEFAULT_RESOURCE).strip() or _DEFAULT_RESOURCE + return f"{node}:{res}" + + +def _lease_for_actor(cluster, resource_id: str, actor: _Actor): + """The actor's own active lease on *resource_id*, if any.""" + existing = cluster.find_existing_lease(resource_id) + if existing is None: + return None + caller = f"a2a:{actor.identity}" + if existing.caller == caller or existing.caller == f"a2a:{actor.holder}": + return existing + return None + + +# ── endpoints ───────────────────────────────────────────────────────────────── + + +@router.get("/api/a2a/gpu/check") +async def gpu_check(request: Request): + """May the caller load on *node*? Reads the channel, folds claims, probes VRAM. + + Authorized readers: an admin session / host local token, or an active agent + registry JWT holding ``a2a_receive``. Fails closed with 503 when the channel + cannot be read (an unreadable channel must never read as "free"). + """ + actor = await _resolve_actor(request, None, "a2a_receive") + node = (request.query_params.get("node") or "").strip() + if not node: + return JSONResponse({"error": "node required"}, status_code=400) + required_mb = _resolve_vram_mb( + None, + request.query_params.get("vram_mb") or request.query_params.get("vram"), + ) + channel = (request.query_params.get("channel") or "").strip() or _channel() + resource = (request.query_params.get("resource") or _DEFAULT_RESOURCE).strip() + body, _claims = await _check_node( + request, + node=node, + required_mb=required_mb or 0, + identity=actor.identity, + channel=channel, + resource=resource or _DEFAULT_RESOURCE, + ) + body["holder"] = actor.holder + return body + + +@router.post("/api/a2a/gpu/claim") +async def gpu_claim(request: Request, body: ClaimBody): + """Claim *node* for this caller: admission check + cluster lease + bus post. + + Returns 409 with the admission detail when the node is claimed by someone + else or lacks the VRAM; 503 when the channel cannot be read. The cluster + lease is rolled back if the bus post fails, so a claim is never half-made. + """ + node = (body.node or "").strip() + if not node: + return JSONResponse({"error": "node required"}, status_code=400) + actor = await _resolve_actor(request, body.holder, "a2a_send") + + vram_mb = _resolve_vram_mb(body.vram_mb, body.vram) + if vram_mb is None or vram_mb <= 0: + return JSONResponse( + {"error": "vram_mb (or vram, e.g. '~6gb') required and must be > 0"}, + status_code=400, + ) + + channel = (body.channel or "").strip() or _channel() + admission, _claims = await _check_node( + request, + node=node, + required_mb=vram_mb, + identity=actor.identity, + channel=channel, + resource=body.resource, + ) + if not admission["admitted"]: + return JSONResponse({"status": "denied", **admission}, status_code=409) + + # Productized half: a real lease the scheduler enforces. Only for a node + # this controller knows as a worker; an external node is bus-governed. + cluster = getattr(request.app.state, "cluster_manager", None) + lease_id: str | None = None + lease = None + if cluster is not None and _match_worker(cluster, node) is not None: + resource_id = _resource_id(node, body.resource) + caller = f"a2a:{actor.identity}" + existing = cluster.find_existing_lease(resource_id) + if existing is not None and existing.caller != caller: + return JSONResponse( + { + "status": "denied", + **admission, + "reason": f"{resource_id} is leased by {existing.caller}", + "blockers": [existing.caller], + }, + status_code=409, + ) + if existing is not None: + # Idempotent re-claim: extend our own lease rather than failing on + # the "already leased" guard (which the manager applies to every + # caller, including the current holder). + lease = await cluster.renew_lease( + existing.lease_id, ttl_seconds=float(body.ttl_seconds) + ) + lease_id = existing.lease_id if lease is not None else None + else: + lease = await cluster.claim_lease( + resource_id=resource_id, + caller=caller, + ttl_seconds=float(body.ttl_seconds), + required_vram_mb=vram_mb, + ) + if lease is None: + return JSONResponse( + { + "status": "denied", + **admission, + "reason": f"cluster refused a lease on {resource_id}", + }, + status_code=409, + ) + lease_id = lease.lease_id + + line = render_claim(node, actor.holder, vram_mb, body.reason, body.eta) + try: + posted = await _post_line(channel, actor, line) + except HTTPException: + if lease_id is not None and cluster is not None: + await cluster.release_lease(lease_id) + raise + + return { + "status": "claimed", + "node": node, + "holder": actor.holder, + "vram_mb": vram_mb, + "lease_id": lease_id, + "expires_at": getattr(lease, "expires_at", None), + "line": line, + "channel": channel, + "message": posted, + "admission": admission, + } + + +@router.post("/api/a2a/gpu/release") +async def gpu_release(request: Request, body: ReleaseBody): + """Release this caller's claim on *node*: cluster lease + ``[GPU RELEASE]``. + + Idempotent: the release line is posted even when no local lease is found, so + a peer holding the bus-side claim learns the node is free. + """ + node = (body.node or "").strip() + if not node: + return JSONResponse({"error": "node required"}, status_code=400) + actor = await _resolve_actor(request, body.holder, "a2a_send") + channel = (body.channel or "").strip() or _channel() + + cluster = getattr(request.app.state, "cluster_manager", None) + released_id = body.lease_id + if cluster is not None: + if released_id is None: + lease = _lease_for_actor(cluster, _resource_id(node, body.resource), actor) + released_id = lease.lease_id if lease is not None else None + if released_id is not None: + # release_lease is idempotent and returns False for an unknown id. + await cluster.release_lease(released_id) + + line = render_release(node, actor.holder) + posted = await _post_line(channel, actor, line) + return { + "status": "released", + "node": node, + "holder": actor.holder, + "lease_id": released_id, + "line": line, + "channel": channel, + "message": posted, + } + + +@router.post("/api/a2a/gpu/request") +async def gpu_request(request: Request, body: RequestBody): + """Post ``[GPU REQUEST]`` — the caller is blocked and asks for a window.""" + node = (body.node or "").strip() + if not node: + return JSONResponse({"error": "node required"}, status_code=400) + actor = await _resolve_actor(request, body.holder, "a2a_send") + need_mb = _resolve_vram_mb(body.need_mb, body.need) + if need_mb is None or need_mb <= 0: + return JSONResponse( + {"error": "need_mb (or need, e.g. '~6gb') required and must be > 0"}, + status_code=400, + ) + channel = (body.channel or "").strip() or _channel() + line = render_request(node, need_mb, body.reason) + posted = await _post_line(channel, actor, line) + return { + "status": "requested", + "node": node, + "holder": actor.holder, + "need_mb": need_mb, + "line": line, + "channel": channel, + "message": posted, + } + + +@router.post("/api/a2a/gpu/renew") +async def gpu_renew(request: Request, body: RenewBody): + """Keep-alive: extend the TTL of a lease taken through :func:`gpu_claim`. + + The lease expires on its own if the holder stops renewing, which is how a + crashed or idle consumer frees the node without anyone releasing it. + """ + actor = await _resolve_actor(request, None, "a2a_send") + cluster = getattr(request.app.state, "cluster_manager", None) + if cluster is None: + return JSONResponse({"error": "cluster manager unavailable"}, status_code=503) + lease = await cluster.renew_lease(body.lease_id, ttl_seconds=float(body.ttl_seconds)) + if lease is None: + return JSONResponse( + {"error": "lease not found or expired", "lease_id": body.lease_id}, + status_code=409, + ) + if lease.caller not in (f"a2a:{actor.identity}", f"a2a:{actor.holder}") and not actor.is_admin: + return JSONResponse({"error": "not the lease holder"}, status_code=403) + return { + "status": "renewed", + "lease_id": lease.lease_id, + "resource_id": lease.resource_id, + "expires_at": lease.expires_at, + } From b41fb7fe981c3455e754a80eefe26ceb73cb5a27 Mon Sep 17 00:00:00 2001 From: hognek <227774406+hognek@users.noreply.github.com> Date: Sat, 12 Sep 2026 00:18:48 +0200 Subject: [PATCH 02/17] fix(a2a): GPU lease ownership checks, 400 on unparseable VRAM Self-review hardening on the A2A GPU lease surface: - A caller-supplied lease_id could previously release or renew ANY holder's lease. `_may_act_on` now requires the lease to be the caller's own (or an operator for the explicit-id path, mirroring POST /api/cluster/leases/*); the node-scoped path stays strictly ownership-matched. - The renew route extended the lease BEFORE the ownership check, so a rejected renewal had already moved another holder's expiry. - `GET /api/a2a/gpu/check` treated an unparseable vram figure as 0, i.e. "check the claims, not the VRAM" -- a silent downgrade indistinguishable from a real VRAM check. It is now a 400, matching the other endpoints. Tests: 77 (5 new: unparseable-vram 400, agent release/renew of a foreign lease, agent releasing its own lease, admin explicit-id override, node-scoped release ownership). --- tests/test_routes_a2a_gpu_lease.py | 81 ++++++++++++++++++++++++++++- tinyagentos/routes/a2a_gpu_lease.py | 72 +++++++++++++++++++++---- 2 files changed, 143 insertions(+), 10 deletions(-) diff --git a/tests/test_routes_a2a_gpu_lease.py b/tests/test_routes_a2a_gpu_lease.py index 541ba5cc4..53231bb7d 100644 --- a/tests/test_routes_a2a_gpu_lease.py +++ b/tests/test_routes_a2a_gpu_lease.py @@ -209,6 +209,12 @@ async def test_missing_node_is_400(self, lease_client, bus): assert resp.status_code == 400 assert bus.sends == [] + async def test_unparseable_vram_is_400_not_a_silent_zero(self, lease_client, bus): + resp = await lease_client.get( + "/api/a2a/gpu/check", params={"node": "local", "vram": "six gigabytes"} + ) + assert resp.status_code == 400 + async def test_free_node_is_admitted(self, lease_client, bus, local_vram): local_vram(9000) resp = await lease_client.get( @@ -400,7 +406,8 @@ async def test_release_without_a_lease_id_releases_the_callers_own( async def test_release_does_not_free_another_holders_lease( self, lease_client, bus, cluster ): - # A lease taken by the scheduler (not by this A2A caller) must survive. + # The node-scoped release frees only the caller's OWN claim; a lease + # taken by the scheduler (not by this A2A caller) must survive. foreign = await cluster.claim_lease( "linstation:gpu-cuda-0", caller="skald-dispatcher", ttl_seconds=300 ) @@ -412,6 +419,21 @@ async def test_release_does_not_free_another_holders_lease( assert resp.json()["lease_id"] is None assert [l.lease_id for l in cluster.get_leases()] == [foreign.lease_id] + async def test_admin_may_release_an_explicit_lease_id( + self, lease_client, bus, cluster + ): + # Operator override, mirroring POST /api/cluster/leases/release. + foreign = await cluster.claim_lease( + "linstation:gpu-cuda-0", caller="skald-dispatcher", ttl_seconds=300 + ) + assert foreign is not None + resp = await lease_client.post( + "/api/a2a/gpu/release", + json={"node": "linstation", "lease_id": foreign.lease_id}, + ) + assert resp.status_code == 200 + assert cluster.get_leases() == [] + async def test_reclaiming_extends_rather_than_conflicts(self, lease_client, bus, cluster): first = await lease_client.post( "/api/a2a/gpu/claim", json={"node": "linstation", "vram_mb": 4096} @@ -602,3 +624,60 @@ async def test_agent_token_is_not_a_skeleton_key(self, lease_client, bus): "/api/cluster/leases", headers={"Authorization": f"Bearer {token}"} ) assert resp.status_code in (401, 403) + + async def test_agent_releases_its_own_lease(self, lease_client, bus, cluster): + cid, token = await _agent_token(lease_client._app, scopes=("a2a_send",)) + async with _bare(lease_client._app) as bare: + claimed = await bare.post( + "/api/a2a/gpu/claim", + json={"node": "linstation", "vram_mb": 4096}, + headers={"Authorization": f"Bearer {token}"}, + ) + assert claimed.status_code == 200 + lease_id = claimed.json()["lease_id"] + assert lease_id is not None + async with _bare(lease_client._app) as bare: + released = await bare.post( + "/api/a2a/gpu/release", + json={"node": "linstation", "lease_id": lease_id}, + headers={"Authorization": f"Bearer {token}"}, + ) + assert released.status_code == 200 + assert cluster.get_leases() == [] + + async def test_agent_cannot_release_another_holders_lease( + self, lease_client, bus, cluster + ): + foreign = await cluster.claim_lease( + "linstation:gpu-cuda-0", caller="skald-dispatcher", ttl_seconds=300 + ) + assert foreign is not None + _cid, token = await _agent_token(lease_client._app, scopes=("a2a_send",)) + async with _bare(lease_client._app) as bare: + resp = await bare.post( + "/api/a2a/gpu/release", + json={"node": "linstation", "lease_id": foreign.lease_id}, + headers={"Authorization": f"Bearer {token}"}, + ) + assert resp.status_code == 403 + assert [l.lease_id for l in cluster.get_leases()] == [foreign.lease_id] + assert bus.sends == [] # and no [GPU RELEASE] line clears the peer claim + + async def test_agent_cannot_renew_another_holders_lease( + self, lease_client, bus, cluster + ): + foreign = await cluster.claim_lease( + "linstation:gpu-cuda-0", caller="skald-dispatcher", ttl_seconds=300 + ) + assert foreign is not None + before = cluster.get_leases()[0].expires_at + _cid, token = await _agent_token(lease_client._app, scopes=("a2a_send",)) + async with _bare(lease_client._app) as bare: + resp = await bare.post( + "/api/a2a/gpu/renew", + json={"lease_id": foreign.lease_id, "ttl_seconds": 6000}, + headers={"Authorization": f"Bearer {token}"}, + ) + assert resp.status_code == 403 + # The rejection must not have extended the lease first. + assert cluster.get_leases()[0].expires_at == before diff --git a/tinyagentos/routes/a2a_gpu_lease.py b/tinyagentos/routes/a2a_gpu_lease.py index 228820430..b939133c0 100644 --- a/tinyagentos/routes/a2a_gpu_lease.py +++ b/tinyagentos/routes/a2a_gpu_lease.py @@ -419,12 +419,34 @@ def _lease_for_actor(cluster, resource_id: str, actor: _Actor): existing = cluster.find_existing_lease(resource_id) if existing is None: return None - caller = f"a2a:{actor.identity}" - if existing.caller == caller or existing.caller == f"a2a:{actor.holder}": + if _lease_owned_by(existing, actor): return existing return None +def _lease_owned_by(lease, actor: _Actor) -> bool: + """True when *lease* was taken by *actor*. + + Strict caller match: the node-scoped release/renew paths must not let an + admin's session free a lease it did not take (that is what the explicit-id + paths and the cluster lease API are for). + """ + return lease.caller in (f"a2a:{actor.identity}", f"a2a:{actor.holder}") + + +def _may_act_on(lease, actor: _Actor) -> bool: + """Ownership for an EXPLICIT lease id: the holder, or an operator.""" + return _lease_owned_by(lease, actor) or actor.is_admin + + +def _find_lease(cluster, lease_id: str): + """Look up an active lease by id, or None.""" + for lease in cluster.get_leases(): + if lease.lease_id == lease_id: + return lease + return None + + # ── endpoints ───────────────────────────────────────────────────────────────── @@ -440,10 +462,22 @@ async def gpu_check(request: Request): node = (request.query_params.get("node") or "").strip() if not node: return JSONResponse({"error": "node required"}, status_code=400) - required_mb = _resolve_vram_mb( - None, - request.query_params.get("vram_mb") or request.query_params.get("vram"), - ) + # A figure that cannot be parsed must be a 400, not a silent 0: a CHECK that + # quietly downgrades to "how many claims are open" is indistinguishable from + # one that checked the VRAM, which is the failure mode this surface exists + # to remove. + raw_vram = request.query_params.get("vram_mb") or request.query_params.get("vram") + required_mb = _resolve_vram_mb(None, raw_vram) + if raw_vram is not None and required_mb is None: + return JSONResponse( + { + "error": ( + "vram_mb must be an integer number of MiB " + "(or use vram='~6gb')" + ) + }, + status_code=400, + ) channel = (request.query_params.get("channel") or "").strip() or _channel() resource = (request.query_params.get("resource") or _DEFAULT_RESOURCE).strip() body, _claims = await _check_node( @@ -573,7 +607,19 @@ async def gpu_release(request: Request, body: ReleaseBody): cluster = getattr(request.app.state, "cluster_manager", None) released_id = body.lease_id if cluster is not None: - if released_id is None: + if released_id is not None: + # A caller-supplied id must belong to the caller: releasing another + # holder's lease (and posting the RELEASE that clears its bus claim) + # would hand any agent with a2a_send the power to free someone + # else's GPU. An id that names no ACTIVE lease is a no-op (it has + # already expired), so it falls through to the idempotent post. + lease = _find_lease(cluster, released_id) + if lease is not None and not _may_act_on(lease, actor): + return JSONResponse( + {"error": "not the lease holder", "lease_id": released_id}, + status_code=403, + ) + else: lease = _lease_for_actor(cluster, _resource_id(node, body.resource), actor) released_id = lease.lease_id if lease is not None else None if released_id is not None: @@ -631,14 +677,22 @@ async def gpu_renew(request: Request, body: RenewBody): cluster = getattr(request.app.state, "cluster_manager", None) if cluster is None: return JSONResponse({"error": "cluster manager unavailable"}, status_code=503) + # Ownership is checked BEFORE the renewal: renewing first and rejecting + # afterwards would already have extended another holder's lease. + existing = _find_lease(cluster, body.lease_id) + if existing is None: + return JSONResponse( + {"error": "lease not found or expired", "lease_id": body.lease_id}, + status_code=409, + ) + if not _may_act_on(existing, actor): + return JSONResponse({"error": "not the lease holder"}, status_code=403) lease = await cluster.renew_lease(body.lease_id, ttl_seconds=float(body.ttl_seconds)) if lease is None: return JSONResponse( {"error": "lease not found or expired", "lease_id": body.lease_id}, status_code=409, ) - if lease.caller not in (f"a2a:{actor.identity}", f"a2a:{actor.holder}") and not actor.is_admin: - return JSONResponse({"error": "not the lease holder"}, status_code=403) return { "status": "renewed", "lease_id": lease.lease_id, From 5df3950d29d17080321f0284a37f3dfaec692e7b Mon Sep 17 00:00:00 2001 From: hognek <227774406+hognek@users.noreply.github.com> Date: Sat, 12 Sep 2026 00:22:05 +0200 Subject: [PATCH 03/17] fix(a2a): post the RELEASE line before freeing the lease Kilo review of PR #2988 (WARNING): gpu_release freed the cluster lease first and posted [GPU RELEASE] afterwards, so a failed bus post left the controller with no lease while peers still read an open claim -- the next CHECK then blocks a node that is actually free. The line is now posted first and the lease released only on success; a failed post changes nothing and the caller can retry. `release_lease` is an idempotent in-memory pop, so the halves cannot be left disagreeing the other way. (The review's CRITICAL finding -- a caller-supplied lease_id bypassing the ownership check in gpu_release/gpu_renew -- was already fixed in b41fb7fe.) Tests: 78 (1 new: a failed post leaves the lease intact). --- docs/agent-coordination.md | 4 +++- tests/test_routes_a2a_gpu_lease.py | 17 +++++++++++++++++ tinyagentos/routes/a2a_gpu_lease.py | 10 +++++++--- 3 files changed, 27 insertions(+), 4 deletions(-) diff --git a/docs/agent-coordination.md b/docs/agent-coordination.md index a42dd7f9d..edce9db87 100644 --- a/docs/agent-coordination.md +++ b/docs/agent-coordination.md @@ -396,7 +396,9 @@ Rules that matter when you use it: like "nobody has claimed anything". - **Claim is both halves or neither.** The cluster lease is rolled back if the bus post fails, so a peer that only watches the bus never disagrees with the - local scheduler about who holds the node. + local scheduler about who holds the node. Release is ordered the same way: the + line is posted BEFORE the local lease is freed, so a failed post leaves the + lease intact rather than freeing a node peers still see as claimed. - **Keep-alive is the TTL, not a promise.** A lease expires after `ttl_seconds` (default 300) unless renewed via `/renew`; a crashed or idle holder therefore frees the node without anyone releasing it. diff --git a/tests/test_routes_a2a_gpu_lease.py b/tests/test_routes_a2a_gpu_lease.py index 53231bb7d..0a7d0339b 100644 --- a/tests/test_routes_a2a_gpu_lease.py +++ b/tests/test_routes_a2a_gpu_lease.py @@ -419,6 +419,23 @@ async def test_release_does_not_free_another_holders_lease( assert resp.json()["lease_id"] is None assert [l.lease_id for l in cluster.get_leases()] == [foreign.lease_id] + async def test_release_keeps_the_lease_when_the_bus_post_fails( + self, lease_client, bus, cluster + ): + claim = await lease_client.post( + "/api/a2a/gpu/claim", json={"node": "linstation", "vram_mb": 4096} + ) + lease_id = claim.json()["lease_id"] + assert lease_id is not None + bus.fail_post = True + resp = await lease_client.post( + "/api/a2a/gpu/release", json={"node": "linstation", "lease_id": lease_id} + ) + assert resp.status_code == 502 + # Nothing has changed: the node is still reserved locally, so a retry + # cannot hand the same GPU to two holders. + assert [l.lease_id for l in cluster.get_leases()] == [lease_id] + async def test_admin_may_release_an_explicit_lease_id( self, lease_client, bus, cluster ): diff --git a/tinyagentos/routes/a2a_gpu_lease.py b/tinyagentos/routes/a2a_gpu_lease.py index b939133c0..f70aabb80 100644 --- a/tinyagentos/routes/a2a_gpu_lease.py +++ b/tinyagentos/routes/a2a_gpu_lease.py @@ -622,12 +622,16 @@ async def gpu_release(request: Request, body: ReleaseBody): else: lease = _lease_for_actor(cluster, _resource_id(node, body.resource), actor) released_id = lease.lease_id if lease is not None else None - if released_id is not None: - # release_lease is idempotent and returns False for an unknown id. - await cluster.release_lease(released_id) + # Post BEFORE releasing the local lease, so a bus failure changes nothing + # and the caller can retry. Releasing first would free the node here while + # peers still read an open claim, i.e. block a node that is actually free + # (Kilo review of #2988). release_lease itself is an idempotent in-memory + # pop, so the two halves cannot be left disagreeing in the other direction. line = render_release(node, actor.holder) posted = await _post_line(channel, actor, line) + if released_id is not None and cluster is not None: + await cluster.release_lease(released_id) return { "status": "released", "node": node, From 3736b11d31979777a0eb7b89880f9b400176f0c7 Mon Sep 17 00:00:00 2001 From: hognek <227774406+hognek@users.noreply.github.com> Date: Sat, 12 Sep 2026 00:42:35 +0200 Subject: [PATCH 04/17] fix(a2a): address CodeRabbit review on the GPU lease surface Findings from the CodeRabbit review of PR #2988, all verified against the code: - **Admission ownership trusted a spoofable field** (CWE-290). `_is_mine` accepted the body's `holder=` even when the bus authenticated the author, so `from=@attacker holder=@victim` made the victim's own admission treat the attacker's claim as its own and load anyway. The authenticated `from` now decides; `holder=` is only a fallback for pre-bus-auth posts. - **Field injection through a value.** `reason=a node=ghost` re-parsed as a new `node`. `_clean` neutralises `=` in rendered values and `_split_fields` now keeps the FIRST value for a repeated key. - **Node spelling split the fold.** `Linstation` and `linstation` produced two groups (so one holder's claim was invisible) and a differently-spelled RELEASE never closed its claim. The fold key is case-folded, groups merge, and `_match_worker` resolves a worker name case-insensitively so both spellings reach the same lease. - **TTL was unbounded.** `ttl_seconds: 1e9` removed the auto-expiry the whole mechanism rests on; claim/renew now bound it to (0, 3600]. - **Claim rollback freed a pre-existing lease.** A re-claim renews the caller's own lease; a failed repost then deleted it. Rollback now applies only to a lease this call created. - **Bus reads sent no credential.** A bus that gates reads answered 401, which surfaced as a 503 "unreadable channel". The caller's registry JWT is now presented on the read too (same loopback/HTTPS guard as posting). - Docs: protocol fence language, the known-worker scope of the local lease and its TTL, the read credential. Changelog updated for the scope note. Tests: 89 (11 new). --- changelog.d/taos-893-a2a-gpu-lease.md | 5 +- data/agent_registry_signing.pem.lock | 0 docs/agent-coordination.md | 21 ++++-- tests/test_gpu_lease_protocol.py | 64 ++++++++++++++++++ tests/test_routes_a2a_gpu_lease.py | 94 ++++++++++++++++++++++++-- tinyagentos/gpu_lease.py | 55 +++++++++++---- tinyagentos/routes/a2a_gpu_lease.py | 96 +++++++++++++++++++++------ 7 files changed, 290 insertions(+), 45 deletions(-) create mode 100644 data/agent_registry_signing.pem.lock diff --git a/changelog.d/taos-893-a2a-gpu-lease.md b/changelog.d/taos-893-a2a-gpu-lease.md index 4b58c88cb..afdc19807 100644 --- a/changelog.d/taos-893-a2a-gpu-lease.md +++ b/changelog.d/taos-893-a2a-gpu-lease.md @@ -7,7 +7,10 @@ admission-checked claim that both registers a real cluster lease (TTL, kept alive with `POST /api/a2a/gpu/renew`) and posts the `[GPU CLAIM]` line peers read; `POST /api/a2a/gpu/release` frees it and `POST /api/a2a/gpu/request` - asks for a window when blocked. Claiming refuses (409) on another holder's + asks for a window when blocked. The cluster lease applies to nodes the + controller knows as cluster workers; a node it does not know is coordinated + over the bus alone (admission-checked and posted, with no local TTL). + Claiming refuses (409) on another holder's open claim or insufficient VRAM, and CHECK/CLAIM return 503 rather than reporting a node free when the bus cannot be read, so two agents can no longer silently co-load past the card's VRAM. An agent posts as its own diff --git a/data/agent_registry_signing.pem.lock b/data/agent_registry_signing.pem.lock new file mode 100644 index 000000000..e69de29bb diff --git a/docs/agent-coordination.md b/docs/agent-coordination.md index edce9db87..ee45cc50e 100644 --- a/docs/agent-coordination.md +++ b/docs/agent-coordination.md @@ -358,7 +358,7 @@ its VRAM (taOS #893). Coordinate over the bus with a one-line text protocol, and use the controller's endpoints so the protocol is admission-checked and backed by a real lease: -``` +```text [GPU CLAIM] node= holder=@you vram=~9.4gb reason=... eta=... [GPU RELEASE] node= holder=@you [GPU REQUEST] node= need=~6gb @@ -377,6 +377,14 @@ checks admission (another holder's claim, the node's free VRAM, and the cluster' own lease table) before the line is posted. A line posted by hand is recorded but enforces nothing. +The **local** half of a claim — the cluster lease with its TTL, renewal and +rollback — applies to nodes this controller knows as cluster workers. A node that +resolves to no worker is coordinated over the bus alone: CLAIM still admission- +checks the channel and posts the line, but there is no local reservation to +expire, so the bus-side `[GPU RELEASE]` is the only thing that frees it. The +endpoint reports a node it cannot measure as `vram_verified: false` rather than +as free. + Rules that matter when you use it: - **CHECK before load, always.** It folds the channel's `[GPU CLAIM]`/`[GPU @@ -393,15 +401,18 @@ Rules that matter when you use it: field in the request body is ignored for agent callers. - **Fail closed.** If the channel cannot be read, `check` and `claim` return `503` rather than reporting the node free: an unreadable channel looks exactly - like "nobody has claimed anything". + like "nobody has claimed anything". An agent's registry JWT is presented on + that read too, so a bus that gates reads does not look like a dead channel. - **Claim is both halves or neither.** The cluster lease is rolled back if the bus post fails, so a peer that only watches the bus never disagrees with the local scheduler about who holds the node. Release is ordered the same way: the line is posted BEFORE the local lease is freed, so a failed post leaves the lease intact rather than freeing a node peers still see as claimed. -- **Keep-alive is the TTL, not a promise.** A lease expires after - `ttl_seconds` (default 300) unless renewed via `/renew`; a crashed or idle - holder therefore frees the node without anyone releasing it. +- **Keep-alive is the TTL, not a promise.** A cluster-worker lease expires after + `ttl_seconds` (default 300, capped at 3600) unless renewed via `/renew`; a + crashed or idle holder therefore frees the node without anyone releasing it. + An unbounded TTL would let one agent take the shared GPU permanently, so the + cap is enforced by the request model. - **A claim is only visible inside the channel fold window** (the newest 500 messages). For a load that outlives the chatter around it, re-POST `/claim` periodically: it is idempotent (it extends the lease and reposts the line, diff --git a/tests/test_gpu_lease_protocol.py b/tests/test_gpu_lease_protocol.py index 33d97e8fd..1283c7ba6 100644 --- a/tests/test_gpu_lease_protocol.py +++ b/tests/test_gpu_lease_protocol.py @@ -127,6 +127,23 @@ def test_newline_cannot_inject_a_second_protocol_line(self): assert msg is not None and msg.kind == CLAIM assert "[GPU RELEASE]" in msg.reason + def test_a_value_cannot_inject_a_field(self): + # `=` is neutralised in rendered values, so a reason cannot smuggle a + # `node=` that a reader re-parses as the real node. + line = render_claim("n1", "@a", 1024, reason="x node=ghost") + msg = parse_message(line) + assert msg is not None + assert msg.node == "n1" + assert "ghost" in msg.reason + + def test_a_hand_written_duplicate_key_keeps_the_first_value(self): + # First-wins: a trailing `node=` inside a value cannot override the + # structural field a renderer emits first. + msg = parse_message( + "[GPU CLAIM] node=n1 holder=@a vram=1gb reason=y node=ghost" + ) + assert msg.node == "n1" + class TestOpenClaims: def _bus(self, *bodies, sender="@a"): @@ -191,6 +208,26 @@ def test_claims_are_grouped_by_node(self): assert claims_for_node(folded, "N2")[0].node == "n2" assert claims_for_node(folded, "n3") == [] + def test_node_spelling_is_case_insensitive_in_the_fold(self): + # Two spellings of one hostname must land in ONE group: otherwise + # claims_for_node returns only one of them and the other holder's claim + # is invisible to admission. + msgs = [ + {"id": 1, "from": "@a", "body": "[GPU CLAIM] node=Linstation holder=@a vram=6gb"}, + {"id": 2, "from": "@b", "body": "[GPU CLAIM] node=linstation holder=@b vram=2gb"}, + ] + folded = open_claims(msgs) + assert list(folded) == ["linstation"] + assert len(folded["linstation"]) == 2 + assert len(claims_for_node(folded, "LINSTATION")) == 2 + + def test_release_closes_a_claim_spelled_with_a_different_case(self): + msgs = [ + {"id": 1, "from": "@a", "body": "[GPU CLAIM] node=Linstation holder=@a vram=6gb"}, + {"id": 2, "from": "@a", "body": "[GPU RELEASE] node=linstation holder=@a"}, + ] + assert open_claims(msgs) == {} + class TestSameHolder: def test_alias_and_canonical_id_match(self): @@ -233,6 +270,33 @@ def test_another_holders_claim_blocks_even_with_vram_to_spare(self): assert d.blockers == ("@taosmd",) assert "already claimed" in d.reason + def test_a_spoofed_holder_does_not_make_a_claim_mine(self): + # `from` is the authenticated author; `holder=` is caller-controlled. + # Accepting the holder for ownership would let an attacker post + # from=@attacker holder=@victim and have the victim's own admission + # treat the attacker's claim as its own (i.e. not a blocker, CWE-290). + claim = parse_message( + { + "id": 1, + "from": "@attacker", + "body": "[GPU CLAIM] node=n1 holder=@victim vram=6gb", + } + ) + d = evaluate_admission( + node="n1", required_mb=1024, identity="@victim", claims=[claim], + free_mb=12288, + ) + assert d.admitted is False + assert d.blockers == ("@victim",) # the readable label is still shown + + def test_interim_claim_without_an_authenticated_author_matches_on_holder(self): + # Posts that predate bus auth carry no `from`; the holder is all we have. + claim = parse_message("[GPU CLAIM] node=n1 holder=@a vram=1gb") + d = evaluate_admission( + node="n1", required_mb=0, identity="@a", claims=[claim], free_mb=4096 + ) + assert d.admitted is True + def test_own_claim_does_not_block_and_is_subtracted_from_the_budget(self): d = evaluate_admission( node="n1", diff --git a/tests/test_routes_a2a_gpu_lease.py b/tests/test_routes_a2a_gpu_lease.py index 0a7d0339b..563287ccf 100644 --- a/tests/test_routes_a2a_gpu_lease.py +++ b/tests/test_routes_a2a_gpu_lease.py @@ -8,6 +8,8 @@ """ from __future__ import annotations +import time + import pytest import pytest_asyncio from httpx import ASGITransport, AsyncClient @@ -26,6 +28,7 @@ class FakeBus: def __init__(self) -> None: self.messages: list[dict] = [] self.sends: list[dict] = [] + self.gets: list[dict] = [] self.fail_get = False self.fail_post = False self._id = 0 @@ -79,9 +82,10 @@ async def __aenter__(self): async def __aexit__(self, *exc) -> bool: return False - async def get(self, url, params=None): + async def get(self, url, params=None, headers=None): if bus.fail_get: raise RuntimeError("bus unreachable") + bus.gets.append({"params": params, "headers": headers}) return _Resp({"messages": list(bus.messages)}) async def post(self, url, json=None, headers=None): @@ -417,7 +421,7 @@ async def test_release_does_not_free_another_holders_lease( ) assert resp.status_code == 200 assert resp.json()["lease_id"] is None - assert [l.lease_id for l in cluster.get_leases()] == [foreign.lease_id] + assert [lease.lease_id for lease in cluster.get_leases()] == [foreign.lease_id] async def test_release_keeps_the_lease_when_the_bus_post_fails( self, lease_client, bus, cluster @@ -434,7 +438,7 @@ async def test_release_keeps_the_lease_when_the_bus_post_fails( assert resp.status_code == 502 # Nothing has changed: the node is still reserved locally, so a retry # cannot hand the same GPU to two holders. - assert [l.lease_id for l in cluster.get_leases()] == [lease_id] + assert [lease.lease_id for lease in cluster.get_leases()] == [lease_id] async def test_admin_may_release_an_explicit_lease_id( self, lease_client, bus, cluster @@ -503,15 +507,69 @@ async def test_renew_unknown_lease_is_409(self, lease_client, bus, cluster): async def test_expired_lease_auto_frees_the_node(self, lease_client, bus, cluster): resp = await lease_client.post( "/api/a2a/gpu/claim", - json={"node": "linstation", "vram_mb": 4096, "ttl_seconds": -1}, + json={"node": "linstation", "vram_mb": 4096, "ttl_seconds": 60}, ) assert resp.status_code == 200 - # TTL already elapsed: the node is free again with no explicit release. + # Age the lease past its TTL (deterministic): the node must be free + # again with no explicit release, which is the keep-alive guarantee. + cluster.get_leases()[0].expires_at = time.time() - 1 nxt = await lease_client.post( "/api/a2a/gpu/claim", json={"node": "linstation", "vram_mb": 4096} ) assert nxt.status_code == 200 + async def test_ttl_is_bounded_on_claim_and_renew(self, lease_client, bus, cluster): + # An unbounded TTL would let one agent hold the shared GPU forever and + # remove the auto-expiry the mechanism rests on. + too_long = await lease_client.post( + "/api/a2a/gpu/claim", + json={"node": "linstation", "vram_mb": 4096, "ttl_seconds": 1e9}, + ) + assert too_long.status_code == 422 + assert cluster.get_leases() == [] + + claim = await lease_client.post( + "/api/a2a/gpu/claim", json={"node": "linstation", "vram_mb": 4096} + ) + lease_id = claim.json()["lease_id"] + long_renew = await lease_client.post( + "/api/a2a/gpu/renew", json={"lease_id": lease_id, "ttl_seconds": 99999} + ) + assert long_renew.status_code == 422 + zero = await lease_client.post( + "/api/a2a/gpu/renew", json={"lease_id": lease_id, "ttl_seconds": 0} + ) + assert zero.status_code == 422 + + async def test_node_spelling_resolves_to_one_lease(self, lease_client, bus, cluster): + # `Linstation` and `linstation` name the same worker; they must not take + # two leases on the one GPU. + first = await lease_client.post( + "/api/a2a/gpu/claim", json={"node": "linstation", "vram_mb": 4096} + ) + second = await lease_client.post( + "/api/a2a/gpu/claim", json={"node": "Linstation", "vram_mb": 4096} + ) + assert second.status_code == 200 + assert second.json()["lease_id"] == first.json()["lease_id"] + assert len(cluster.get_leases()) == 1 + + async def test_reclaim_rollback_does_not_drop_the_existing_lease( + self, lease_client, bus, cluster + ): + first = await lease_client.post( + "/api/a2a/gpu/claim", json={"node": "linstation", "vram_mb": 4096} + ) + lease_id = first.json()["lease_id"] + # A re-claim renews the caller's own lease; a failed repost must not + # free a reservation the holder still believes it owns. + bus.fail_post = True + retry = await lease_client.post( + "/api/a2a/gpu/claim", json={"node": "linstation", "vram_mb": 4096} + ) + assert retry.status_code == 502 + assert [lease.lease_id for lease in cluster.get_leases()] == [lease_id] + async def test_check_blocks_when_the_scheduler_holds_a_lease( self, lease_client, bus, cluster ): @@ -634,6 +692,28 @@ async def test_agent_with_receive_scope_can_check(self, lease_client, bus): assert resp.status_code == 200 assert resp.json()["holder"] == "@taosmd" + async def test_agent_check_presents_its_credential_to_the_bus( + self, lease_client, bus + ): + # A bus that gates reads fails a credential-less GET with 401, which + # this route would report as an unreadable channel (503). + _cid, token = await _agent_token(lease_client._app, scopes=("a2a_receive",)) + async with _bare(lease_client._app) as bare: + resp = await bare.get( + "/api/a2a/gpu/check", + params={"node": "n1", "vram_mb": 1024}, + headers={"Authorization": f"Bearer {token}"}, + ) + assert resp.status_code == 200 + assert bus.gets[-1]["headers"].get("Authorization") == f"Bearer {token}" + + async def test_admin_check_sends_no_credential_to_the_bus(self, lease_client, bus): + resp = await lease_client.get( + "/api/a2a/gpu/check", params={"node": "n1", "vram_mb": 1024} + ) + assert resp.status_code == 200 + assert bus.gets[-1]["headers"] is None + async def test_agent_token_is_not_a_skeleton_key(self, lease_client, bus): _cid, token = await _agent_token(lease_client._app, scopes=("a2a_send",)) async with _bare(lease_client._app) as bare: @@ -677,7 +757,7 @@ async def test_agent_cannot_release_another_holders_lease( headers={"Authorization": f"Bearer {token}"}, ) assert resp.status_code == 403 - assert [l.lease_id for l in cluster.get_leases()] == [foreign.lease_id] + assert [lease.lease_id for lease in cluster.get_leases()] == [foreign.lease_id] assert bus.sends == [] # and no [GPU RELEASE] line clears the peer claim async def test_agent_cannot_renew_another_holders_lease( @@ -692,7 +772,7 @@ async def test_agent_cannot_renew_another_holders_lease( async with _bare(lease_client._app) as bare: resp = await bare.post( "/api/a2a/gpu/renew", - json={"lease_id": foreign.lease_id, "ttl_seconds": 6000}, + json={"lease_id": foreign.lease_id, "ttl_seconds": 600}, headers={"Authorization": f"Bearer {token}"}, ) assert resp.status_code == 403 diff --git a/tinyagentos/gpu_lease.py b/tinyagentos/gpu_lease.py index 9d9ee75b6..98ea5f79e 100644 --- a/tinyagentos/gpu_lease.py +++ b/tinyagentos/gpu_lease.py @@ -63,28 +63,39 @@ def _clean(text: object) -> str: - """Collapse *text* to a single printable line. + """Collapse *text* to a single printable, structural-free line. Every field value is caller-supplied text that ends up inside a one-line bus message. A raw newline would let a caller inject a SECOND protocol line - (e.g. a RELEASE that closes another holder's claim), so all whitespace runs - collapse to a single space and non-printable control characters are dropped. + (e.g. a RELEASE that closes another holder's claim), and an unescaped ``=`` + would let it inject a `key=` INSIDE a value that a reader re-parses as a new + field (`reason=a node=ghost` moved the node). So: whitespace runs collapse + to one space, non-printable characters are dropped, and ``=`` is neutralised + by turning it into a space (a value is never structural). """ keep = "".join( - ch if (ch.isprintable() or ch.isspace()) else "" for ch in str(text) + " " if ch == "=" or ch.isspace() else (ch if ch.isprintable() else "") + for ch in str(text) ) return " ".join(keep.split()) def _split_fields(rest: str) -> dict[str, str]: - """Return ``key=value`` pairs from a line body, values allowed spaces.""" + """Return ``key=value`` pairs from a line body, values allowed spaces. + + A key seen twice keeps its FIRST value. Values may legitimately contain the + text of a later field, so last-wins would let a trailing ``node=`` inside a + ``reason`` overwrite the real node (CWE-290 by value injection); the fields + this module renders are always emitted node/holder/vram first, so first-wins + keeps the structural fields stable even for a hand-written line. + """ matches = list(_FIELD_START_RE.finditer(rest)) fields: dict[str, str] = {} for i, m in enumerate(matches): start = m.end() end = matches[i + 1].start() if i + 1 < len(matches) else len(rest) - key = m.group("key").lower() - fields[_FIELD_ALIASES.get(key, key)] = rest[start:end].strip() + key = _FIELD_ALIASES.get(m.group("key").lower(), m.group("key").lower()) + fields.setdefault(key, rest[start:end].strip()) return fields @@ -261,6 +272,13 @@ def open_claims(messages: Iterable[object]) -> dict[str, list[GpuLeaseMessage]]: CLAIM opens a slot keyed by ``(node, identity)``; a RELEASE from the same identity on the same node closes it. A re-CLAIM by the same identity replaces the previous one (so a reposted claim is never double-counted). + + Both the node and the identity are matched case-insensitively: peers do not + agree on the spelling of a hostname, and a `node=Linstation` claim whose + release says `node=linstation` would otherwise stay open forever while a + differently-spelled claim hid in a second group. The returned mapping is + therefore keyed by the CASE-FOLDED node name; each message keeps the + spelling it arrived with. """ open_by_key: dict[tuple[str, str], GpuLeaseMessage] = {} order: list[tuple[str, str]] = [] @@ -268,7 +286,7 @@ def open_claims(messages: Iterable[object]) -> dict[str, list[GpuLeaseMessage]]: msg = parse_message(raw) if msg is None or msg.kind not in (CLAIM, RELEASE): continue - key = (msg.node, msg.identity_key.casefold()) + key = (msg.node.strip().casefold(), msg.identity_key.casefold()) if msg.kind == CLAIM: if key not in open_by_key: order.append(key) @@ -280,7 +298,7 @@ def open_claims(messages: Iterable[object]) -> dict[str, list[GpuLeaseMessage]]: for key in order: msg = open_by_key.get(key) if msg is not None: - out.setdefault(msg.node, []).append(msg) + out.setdefault(key[0], []).append(msg) return out @@ -289,6 +307,10 @@ def claims_for_node( ) -> list[GpuLeaseMessage]: """Claims for a node, matching the node case-insensitively.""" wanted = (node or "").strip().casefold() + entries = claims.get(wanted) + if entries: + return list(entries) + # Tolerate a caller-built mapping whose keys were not case-folded. for name, entries in claims.items(): if name.strip().casefold() == wanted: return list(entries) @@ -296,9 +318,18 @@ def claims_for_node( def _is_mine(claim: GpuLeaseMessage, identity: str) -> bool: - return same_holder(claim.identity_key, identity) or same_holder( - claim.holder, identity - ) + """True when *claim* was made by *identity*. + + When the bus authenticated the author (``bus_from``), ONLY that value is + compared. The body's ``holder=`` is caller-controlled, so accepting it would + let an attacker post ``from=@attacker holder=@victim`` and make the victim's + own admission treat the attacker's claim as the victim's, i.e. not a blocker + (CWE-290). The readable holder is only a fallback for the interim posts that + predate bus auth and therefore have no authenticated author. + """ + if claim.bus_from: + return same_holder(claim.bus_from, identity) + return same_holder(claim.holder, identity) @dataclass(frozen=True) diff --git a/tinyagentos/routes/a2a_gpu_lease.py b/tinyagentos/routes/a2a_gpu_lease.py index f70aabb80..a7b2080d7 100644 --- a/tinyagentos/routes/a2a_gpu_lease.py +++ b/tinyagentos/routes/a2a_gpu_lease.py @@ -50,7 +50,7 @@ import httpx from fastapi import APIRouter, HTTPException, Request from fastapi.responses import JSONResponse -from pydantic import BaseModel +from pydantic import BaseModel, Field from tinyagentos.agent_token_auth import check_agent_scope from tinyagentos.gpu_lease import ( @@ -85,6 +85,12 @@ # enough that a crashed holder stops blocking the node (issue: advisory keep-alive). _DEFAULT_TTL_SECONDS = 300.0 +# Upper bound on a caller-supplied TTL. Without it an authenticated agent could +# take the shared GPU with `ttl_seconds: 1e9` and remove the auto-expiry the +# whole mechanism rests on. A longer load is kept alive by renewing, not by one +# enormous lease. +MAX_LEASE_TTL_SECONDS = 3600.0 + def _channel() -> str: return os.environ.get("TAOS_A2A_GPU_CHANNEL", _DEFAULT_CHANNEL).strip() or _DEFAULT_CHANNEL @@ -138,12 +144,30 @@ async def _resolve_actor( return _Actor(identity=caller, holder=handle, credential=_bearer_token(request)) -async def _read_channel(channel: str) -> list[dict]: - """Fetch the channel's messages oldest-first. Raises on an unreadable bus.""" +async def _read_channel(channel: str, actor: _Actor) -> list[dict]: + """Fetch the channel's messages oldest-first. Raises on an unreadable bus. + + The caller's registry credential is presented here too: a bus that gates + reads fails a credential-less GET with 401, which this route would surface + as an unreadable channel (503). Forwarding it is subject to the same + loopback/HTTPS guard as the post path. + """ + headers: dict[str, str] = {} + if actor.credential: + bus = _bus_url() + if _credential_may_cross(bus): + headers["Authorization"] = f"Bearer {actor.credential}" + else: + logger.warning( + "A2A GPU lease read credential withheld for non-loopback http destination %s", + bus, + ) bus = _bus_url() async with httpx.AsyncClient(timeout=5.0) as client: resp = await client.get( - f"{bus}/a2a/messages", params={"thread": channel, "limit": _CHANNEL_LIMIT} + f"{bus}/a2a/messages", + params={"thread": channel, "limit": _CHANNEL_LIMIT}, + headers=headers or None, ) resp.raise_for_status() data = resp.json() @@ -180,10 +204,12 @@ async def _post_line(channel: str, actor: _Actor, text: str) -> dict: return data if isinstance(data, dict) else {} -async def _folded_claims(request: Request, channel: str) -> dict[str, list[GpuLeaseMessage]]: +async def _folded_claims( + request: Request, channel: str, actor: _Actor +) -> dict[str, list[GpuLeaseMessage]]: """Read the channel and fold it, mapping an unreadable bus to 503.""" try: - messages = await _read_channel(channel) + messages = await _read_channel(channel, actor) except Exception as exc: # noqa: BLE001 logger.warning("A2A GPU lease channel read failed (channel=%s): %s", channel, exc) raise HTTPException( @@ -199,9 +225,11 @@ async def _folded_claims(request: Request, channel: str) -> dict[str, list[GpuLe def _match_worker(cluster, node: str): """Return the cluster worker *node* names, or None. - Accepts the worker name, a URL host, or the local controller's own aliases - ("local", "localhost", this hostname) so a bus node label resolves to the - worker whose VRAM the label refers to. + Accepts the worker name (case-insensitively, since a bus label is free + text), a URL host, or the local controller's own aliases ("local", + "localhost", this hostname) so a bus node label resolves to the worker whose + VRAM the label refers to. Failing to resolve is a real outcome: an unknown + node is coordinated over the bus alone, with no local lease. """ if cluster is None: return None @@ -213,6 +241,10 @@ def _match_worker(cluster, node: str): return None if wanted in ("local", "localhost", socket.gethostname().casefold()): return cluster.get_worker("local") + for w in cluster.get_workers(): + name = getattr(w, "name", "") or "" + if name and name.casefold() == wanted: + return w for w in cluster.get_workers(): url = getattr(w, "url", "") or "" host = url.split("//", 1)[-1].split("/", 1)[0].split(":", 1)[0] @@ -311,7 +343,7 @@ def _cluster_lease_claims(cluster, node: str, resource: str) -> list[GpuLeaseMes """ if cluster is None: return [] - resource_id = _resource_id(node, resource) + resource_id = _resource_id(_canonical_node(cluster, node), resource) out: list[GpuLeaseMessage] = [] seen: set[str] = set() for lease in cluster.get_leases(): @@ -340,12 +372,12 @@ async def _check_node( *, node: str, required_mb: int, - identity: str, + actor: _Actor, channel: str, resource: str = _DEFAULT_RESOURCE, ) -> tuple[dict, list[GpuLeaseMessage]]: """Run the full CHECK for a node; returns (admission dict, node claims).""" - folded = await _folded_claims(request, channel) + folded = await _folded_claims(request, channel, actor) cluster = getattr(request.app.state, "cluster_manager", None) bus_claims = claims_for_node(folded, node) # An A2A claim and the cluster lease it created are the same reservation; @@ -363,7 +395,7 @@ async def _check_node( admission = evaluate_admission( node=node, required_mb=required_mb, - identity=identity, + identity=actor.identity, claims=node_claims, free_mb=free_mb, capacity_mb=capacity_mb, @@ -388,7 +420,9 @@ class ClaimBody(_LeaseBody): vram: str | None = None reason: str = "" eta: str = "" - ttl_seconds: float = _DEFAULT_TTL_SECONDS + ttl_seconds: float = Field( + default=_DEFAULT_TTL_SECONDS, gt=0, le=MAX_LEASE_TTL_SECONDS + ) holder: str | None = None # honored for admin callers only @@ -406,7 +440,9 @@ class RequestBody(_LeaseBody): class RenewBody(BaseModel): lease_id: str - ttl_seconds: float = _DEFAULT_TTL_SECONDS + ttl_seconds: float = Field( + default=_DEFAULT_TTL_SECONDS, gt=0, le=MAX_LEASE_TTL_SECONDS + ) def _resource_id(node: str, resource: str) -> str: @@ -414,6 +450,17 @@ def _resource_id(node: str, resource: str) -> str: return f"{node}:{res}" +def _canonical_node(cluster, node: str) -> str: + """Return the cluster worker's own name for *node*, else *node*. + + The bus label is free text ("Linstation", "linstation", a URL host) while + leases are keyed on the worker's registered name, so resolving through the + manager keeps two spellings of one host from taking two leases on one GPU. + """ + worker = _match_worker(cluster, node) + return getattr(worker, "name", None) or node + + def _lease_for_actor(cluster, resource_id: str, actor: _Actor): """The actor's own active lease on *resource_id*, if any.""" existing = cluster.find_existing_lease(resource_id) @@ -484,7 +531,7 @@ async def gpu_check(request: Request): request, node=node, required_mb=required_mb or 0, - identity=actor.identity, + actor=actor, channel=channel, resource=resource or _DEFAULT_RESOURCE, ) @@ -517,7 +564,7 @@ async def gpu_claim(request: Request, body: ClaimBody): request, node=node, required_mb=vram_mb, - identity=actor.identity, + actor=actor, channel=channel, resource=body.resource, ) @@ -529,8 +576,9 @@ async def gpu_claim(request: Request, body: ClaimBody): cluster = getattr(request.app.state, "cluster_manager", None) lease_id: str | None = None lease = None + created_lease = False if cluster is not None and _match_worker(cluster, node) is not None: - resource_id = _resource_id(node, body.resource) + resource_id = _resource_id(_canonical_node(cluster, node), body.resource) caller = f"a2a:{actor.identity}" existing = cluster.find_existing_lease(resource_id) if existing is not None and existing.caller != caller: @@ -568,12 +616,16 @@ async def gpu_claim(request: Request, body: ClaimBody): status_code=409, ) lease_id = lease.lease_id + created_lease = True line = render_claim(node, actor.holder, vram_mb, body.reason, body.eta) try: posted = await _post_line(channel, actor, line) except HTTPException: - if lease_id is not None and cluster is not None: + # Roll back only a lease THIS call created. A re-claim renews the + # caller's own pre-existing lease, and freeing that on a transient bus + # failure would drop a reservation the holder still believes it owns. + if created_lease and lease_id is not None and cluster is not None: await cluster.release_lease(lease_id) raise @@ -620,7 +672,11 @@ async def gpu_release(request: Request, body: ReleaseBody): status_code=403, ) else: - lease = _lease_for_actor(cluster, _resource_id(node, body.resource), actor) + lease = _lease_for_actor( + cluster, + _resource_id(_canonical_node(cluster, node), body.resource), + actor, + ) released_id = lease.lease_id if lease is not None else None # Post BEFORE releasing the local lease, so a bus failure changes nothing From 12ade1998ad366d0adc4b4727defd482f7df7558 Mon Sep 17 00:00:00 2001 From: hognek <227774406+hognek@users.noreply.github.com> Date: Sat, 12 Sep 2026 01:01:56 +0200 Subject: [PATCH 05/17] chore(git): stop tracking the registry signing-key lock sidecar `data/agent_registry_signing.pem.lock` is a 0-byte flock sidecar the registry key writer creates beside the key. `data/*.pem` already ignores the key, but not the lock, so a stray `git add` swept it into the previous commit on this branch - the same accident the neighbouring note in .gitignore describes (PR #2540). Untrack it and enumerate the pattern; verified against origin/dev that no tracked file becomes ignored. --- .gitignore | 8 ++++++++ data/agent_registry_signing.pem.lock | 0 2 files changed, 8 insertions(+) delete mode 100644 data/agent_registry_signing.pem.lock diff --git a/.gitignore b/.gitignore index 722fc2b6d..3645397c8 100644 --- a/.gitignore +++ b/.gitignore @@ -203,3 +203,11 @@ data/*.sqlite3-wal data/apps/ data/.install_id data/store_popularity.json + +# Registry signing-key lock sidecar. Added 2026-09-12: `data/*.pem` (above) +# covers the key itself, but not the 0-byte `*.pem.lock` flock sidecar the +# registry key writer creates beside it - so it rode into PR #2988 on a stray +# `git add`, the same accident the note above describes. Enumerated, not a +# blanket data/ rule, for the reason given there. +# Verified against every tracked file on origin/dev: 0 become ignored. +data/*.pem.lock diff --git a/data/agent_registry_signing.pem.lock b/data/agent_registry_signing.pem.lock deleted file mode 100644 index e69de29bb..000000000 From 4829cf580fba35450be11b21c175fa1d78a8b57c Mon Sep 17 00:00:00 2001 From: hognek <227774406+hognek@users.noreply.github.com> Date: Sat, 12 Sep 2026 01:02:01 +0200 Subject: [PATCH 06/17] fix(a2a): attribute an operator's GPU lease release to the freed holder CodeRabbit on #2988: releasing another holder's lease by explicit `lease_id` posted the `[GPU RELEASE]` as `@operator`. A bus claim is keyed on its AUTHOR (the fold's identity_key is bus_from or holder), so that line closed nothing: the local lease was gone while every peer's fold still read the node as claimed, i.e. a GPU that was actually free stayed blocked - the exact local/peer disagreement the endpoint exists to remove. The line is now attributed to the freed lease's bus identity (an admin session may post with an explicit `from`), with the readable handle resolved from the agent registry; `released_holder` reports whose claim the line closes while `holder` keeps reporting who acted. A lease with no bus claim behind it (a non-`a2a:` caller such as skald-dispatcher) is unchanged, and a bus that authenticates senders refuses the substitution: the post fails before the local lease is freed, so the override cannot half-apply. Tests: 47 in tests/test_routes_a2a_gpu_lease.py (1 new), 296 with the neighbouring suites. The new test fails on the pre-fix code. --- changelog.d/taos-893-a2a-gpu-lease.md | 10 ++++++ docs/agent-coordination.md | 10 ++++++ tests/test_routes_a2a_gpu_lease.py | 47 +++++++++++++++++++++++++ tinyagentos/routes/a2a_gpu_lease.py | 50 +++++++++++++++++++++++++-- 4 files changed, 115 insertions(+), 2 deletions(-) diff --git a/changelog.d/taos-893-a2a-gpu-lease.md b/changelog.d/taos-893-a2a-gpu-lease.md index afdc19807..363d4ff98 100644 --- a/changelog.d/taos-893-a2a-gpu-lease.md +++ b/changelog.d/taos-893-a2a-gpu-lease.md @@ -17,3 +17,13 @@ registry identity (scope `a2a_receive` to check, `a2a_send` to act) and the node label resolves to a cluster worker's heartbeat VRAM or the controller's own shared VRAM ledger. + +### Fixed + +- Freeing another holder's GPU lease by explicit `lease_id` (the operator + override, `POST /api/a2a/gpu/release`) posts the `[GPU RELEASE]` line as the + **freed holder** rather than as the operator. A claim is keyed on its bus + author, so the old line cleared nothing: the local lease was gone while every + peer's fold still read the node as claimed, blocking a GPU that was actually + free. The response now distinguishes `holder` (who acted) from + `released_holder` (whose claim the line closes). diff --git a/docs/agent-coordination.md b/docs/agent-coordination.md index ee45cc50e..c354448b8 100644 --- a/docs/agent-coordination.md +++ b/docs/agent-coordination.md @@ -408,6 +408,16 @@ Rules that matter when you use it: local scheduler about who holds the node. Release is ordered the same way: the line is posted BEFORE the local lease is freed, so a failed post leaves the lease intact rather than freeing a node peers still see as claimed. +- **An operator's release is attributed to the holder whose claim it closes.** + A lease taken through `/claim` can also be freed by an explicit `lease_id` — + by its holder, or by an operator (`_may_act_on`). Since a bus claim is keyed + on its **author**, the operator's `[GPU RELEASE]` is posted as the freed + holder, not as the operator (an admin session may set an explicit `from`, see + *Posting to the coordination bus*); the response reports both, `holder` (who + acted) and `released_holder` (whose claim the line closes). A bus that + authenticates senders refuses the substitution, and because the line is still + posted before the local lease is freed, the override then fails loudly + (`502`) with the lease intact rather than leaving the two views disagreeing. - **Keep-alive is the TTL, not a promise.** A cluster-worker lease expires after `ttl_seconds` (default 300, capped at 3600) unless renewed via `/renew`; a crashed or idle holder therefore frees the node without anyone releasing it. diff --git a/tests/test_routes_a2a_gpu_lease.py b/tests/test_routes_a2a_gpu_lease.py index 563287ccf..4e956ce0f 100644 --- a/tests/test_routes_a2a_gpu_lease.py +++ b/tests/test_routes_a2a_gpu_lease.py @@ -455,6 +455,53 @@ async def test_admin_may_release_an_explicit_lease_id( assert resp.status_code == 200 assert cluster.get_leases() == [] + async def test_operator_release_closes_the_holders_bus_claim( + self, lease_client, bus, cluster + ): + """An operator override must close the claim it frees (CR on #2988). + + The bus claim is keyed on its AUTHOR, so a RELEASE posted as @operator + would free the local lease while every peer's fold kept reading the node + as claimed -- the local/peer disagreement this surface exists to remove. + """ + cid, token = await _agent_token(lease_client._app, scopes=("a2a_send",)) + async with _bare(lease_client._app) as bare: + claimed = await bare.post( + "/api/a2a/gpu/claim", + json={"node": "linstation", "vram_mb": 4096}, + headers={"Authorization": f"Bearer {token}"}, + ) + assert claimed.status_code == 200 + lease_id = claimed.json()["lease_id"] + assert lease_id is not None + + # An admin session (the operator) frees the agent's lease by explicit id. + released = await lease_client.post( + "/api/a2a/gpu/release", + json={"node": "linstation", "lease_id": lease_id}, + ) + assert released.status_code == 200 + assert released.json()["holder"] == "@operator" + assert released.json()["released_holder"] == "@taosmd" + # ... attributed to the holder whose claim it closes, not to the operator. + assert bus.last_from == cid + assert bus.last_line == "[GPU RELEASE] node=linstation holder=@taosmd" + assert cluster.get_leases() == [] + + # So another agent's CHECK no longer reports the claim. + _other, other = await _agent_token( + lease_client._app, scopes=("a2a_receive",), handle="@taos" + ) + async with _bare(lease_client._app) as bare: + checked = await bare.get( + "/api/a2a/gpu/check", + params={"node": "linstation", "vram_mb": 1024}, + headers={"Authorization": f"Bearer {other}"}, + ) + assert checked.status_code == 200 + assert checked.json()["admitted"] is True + assert checked.json()["blockers"] == [] + async def test_reclaiming_extends_rather_than_conflicts(self, lease_client, bus, cluster): first = await lease_client.post( "/api/a2a/gpu/claim", json={"node": "linstation", "vram_mb": 4096} diff --git a/tinyagentos/routes/a2a_gpu_lease.py b/tinyagentos/routes/a2a_gpu_lease.py index a7b2080d7..eb8b4a223 100644 --- a/tinyagentos/routes/a2a_gpu_lease.py +++ b/tinyagentos/routes/a2a_gpu_lease.py @@ -481,6 +481,30 @@ def _lease_owned_by(lease, actor: _Actor) -> bool: return lease.caller in (f"a2a:{actor.identity}", f"a2a:{actor.holder}") +def _lease_bus_identity(lease) -> str | None: + """The bus identity *lease* was taken under, or None. + + ``gpu_claim`` records its leases as ``a2a:``. A lease taken + through the cluster API directly (``skald-dispatcher``) has no bus claim + behind it, so there is no identity to attribute a RELEASE to either. + """ + caller = (getattr(lease, "caller", "") or "").strip() + prefix = "a2a:" + if caller.startswith(prefix): + return caller[len(prefix):].strip() or None + return None + + +async def _holder_for(request: Request, identity: str) -> str: + """The readable ``@handle`` for a bus *identity*, else the identity itself.""" + registry = getattr(request.app.state, "agent_registry", None) + record = await registry.get(identity) if registry is not None else None + handle = _clean_handle((record or {}).get("handle")) + if not handle: + return identity + return handle if handle.startswith("@") else f"@{handle}" + + def _may_act_on(lease, actor: _Actor) -> bool: """Ownership for an EXPLICIT lease id: the holder, or an operator.""" return _lease_owned_by(lease, actor) or actor.is_admin @@ -658,6 +682,7 @@ async def gpu_release(request: Request, body: ReleaseBody): cluster = getattr(request.app.state, "cluster_manager", None) released_id = body.lease_id + lease = None if cluster is not None: if released_id is not None: # A caller-supplied id must belong to the caller: releasing another @@ -679,19 +704,40 @@ async def gpu_release(request: Request, body: ReleaseBody): ) released_id = lease.lease_id if lease is not None else None + # Whose claim the [GPU RELEASE] closes. An operator freeing another holder's + # lease by explicit id must attribute the line to THAT holder: a claim is + # keyed on its bus AUTHOR, so a line posted as @operator clears nothing and + # every peer's fold keeps reading the node as claimed while the local lease + # is already gone (CodeRabbit on #2988). An admin session may post with an + # explicit `from` (docs/agent-coordination.md, *Posting to the coordination + # bus*); a bus that authenticates senders refuses the substitution, and the + # post-before-release ordering below then leaves the local lease intact, so + # the override cannot half-apply. + line_actor = actor + if lease is not None and not _lease_owned_by(lease, actor): + owner = _lease_bus_identity(lease) + if owner is not None: + line_actor = _Actor( + identity=owner, + holder=await _holder_for(request, owner), + credential=actor.credential, + is_admin=actor.is_admin, + ) + # Post BEFORE releasing the local lease, so a bus failure changes nothing # and the caller can retry. Releasing first would free the node here while # peers still read an open claim, i.e. block a node that is actually free # (Kilo review of #2988). release_lease itself is an idempotent in-memory # pop, so the two halves cannot be left disagreeing in the other direction. - line = render_release(node, actor.holder) - posted = await _post_line(channel, actor, line) + line = render_release(node, line_actor.holder) + posted = await _post_line(channel, line_actor, line) if released_id is not None and cluster is not None: await cluster.release_lease(released_id) return { "status": "released", "node": node, "holder": actor.holder, + "released_holder": line_actor.holder, "lease_id": released_id, "line": line, "channel": channel, From eb923915d52e08563655161062dff93a814e1e33 Mon Sep 17 00:00:00 2001 From: hognek <227774406+hognek@users.noreply.github.com> Date: Sat, 12 Sep 2026 01:17:06 +0200 Subject: [PATCH 07/17] fix(a2a): expire the bus claim when the holder stops keeping alive CodeRabbit on #2988: `ttl_seconds` bounded only the cluster lease. A CLAIM published no expiry and the fold kept it open until a RELEASE arrived, so after a holder crashed (or went idle) another identity was denied the shared card until its claim aged out of the fold window - a GPU held by nobody, which is the mirror image of the co-load this surface exists to prevent. - `render_claim` can publish `expires=`, and `open_claims` drops a claim whose published expiry has passed, exactly as the cluster lease's TTL frees the reservation behind it. A claim posted without an expiry is unchanged: bound by a RELEASE alone, which is what the interim hand-posted lines rely on. - `gpu_claim` publishes the backing lease's expiry (or now + the requested TTL on a bus-only node) and reports it as `claim_expires_at`. - `gpu_renew` reposts the claim line with the new expiry as it extends the lease, so keeping the lease alive keeps the claim alive; a failed repost is logged and reported (`bus_claim_refreshed: false`) rather than failing a renewal that did happen locally. `RenewBody.channel` names the thread the claim lives on. Tests: 101 in tests/test_gpu_lease_protocol.py + tests/test_routes_a2a_gpu_lease.py (12 new), 296 with the neighbouring suites. The new tests fail without the change (a crashed holder's node becomes claimable only once the published expiry passes). --- changelog.d/taos-893-a2a-gpu-lease.md | 8 ++ docs/agent-coordination.md | 12 ++- tests/test_gpu_lease_protocol.py | 66 +++++++++++++++ tests/test_routes_a2a_gpu_lease.py | 116 +++++++++++++++++++++++++- tinyagentos/gpu_lease.py | 70 ++++++++++++++-- tinyagentos/routes/a2a_gpu_lease.py | 105 +++++++++++++++++++---- 6 files changed, 352 insertions(+), 25 deletions(-) diff --git a/changelog.d/taos-893-a2a-gpu-lease.md b/changelog.d/taos-893-a2a-gpu-lease.md index 363d4ff98..c5043cc0e 100644 --- a/changelog.d/taos-893-a2a-gpu-lease.md +++ b/changelog.d/taos-893-a2a-gpu-lease.md @@ -17,6 +17,14 @@ registry identity (scope `a2a_receive` to check, `a2a_send` to act) and the node label resolves to a cluster worker's heartbeat VRAM or the controller's own shared VRAM ledger. +- Claims carry their own expiry. `POST /api/a2a/gpu/claim` publishes + `expires=` on the `[GPU CLAIM]` line (the backing cluster lease's + expiry, or the requested TTL for a bus-only node) and `POST /api/a2a/gpu/renew` + reposts that line as it extends the lease, so a holder that keeps its lease + alive keeps its claim alive. The fold drops a claim whose published expiry has + passed, so a holder that crashed or stopped keeping alive no longer blocks the + shared card until its claim ages out of the fold window. A claim posted by hand + without an `expires=` is unchanged: bounded by a RELEASE alone. ### Fixed diff --git a/docs/agent-coordination.md b/docs/agent-coordination.md index c354448b8..ddaaf86c5 100644 --- a/docs/agent-coordination.md +++ b/docs/agent-coordination.md @@ -359,7 +359,7 @@ use the controller's endpoints so the protocol is admission-checked and backed b a real lease: ```text -[GPU CLAIM] node= holder=@you vram=~9.4gb reason=... eta=... +[GPU CLAIM] node= holder=@you vram=~9.4gb reason=... eta=... expires= [GPU RELEASE] node= holder=@you [GPU REQUEST] node= need=~6gb ``` @@ -423,6 +423,16 @@ Rules that matter when you use it: crashed or idle holder therefore frees the node without anyone releasing it. An unbounded TTL would let one agent take the shared GPU permanently, so the cap is enforced by the request model. +- **A claim carries its own expiry, and the fold honours it.** `/claim` publishes + `expires=` on the line — the backing lease's expiry when the node is a + cluster worker, else the TTL it was asked for — and `/renew` reposts the claim + as it extends the lease. A fold drops a claim whose published expiry has + passed, exactly as the cluster lease's TTL frees its reservation, so a holder + that crashed or stopped keeping alive no longer blocks the card until its + claim ages out of the fold window. Keep-alive therefore means re-POST `/claim` + or `/renew` while you hold the card; a claim posted by hand without an + `expires=` never expires (it is bounded by a RELEASE alone), which is what the + interim protocol in #893 relies on. - **A claim is only visible inside the channel fold window** (the newest 500 messages). For a load that outlives the chatter around it, re-POST `/claim` periodically: it is idempotent (it extends the lease and reposts the line, diff --git a/tests/test_gpu_lease_protocol.py b/tests/test_gpu_lease_protocol.py index 1283c7ba6..cf9d960be 100644 --- a/tests/test_gpu_lease_protocol.py +++ b/tests/test_gpu_lease_protocol.py @@ -7,6 +7,8 @@ """ from __future__ import annotations +import time + from tinyagentos.gpu_lease import ( CLAIM, RELEASE, @@ -87,6 +89,21 @@ def test_value_may_contain_spaces(self): def test_kind_is_case_insensitive(self): assert parse_message("[gpu claim] node=n1 holder=@a vram=1gb").kind == CLAIM + def test_expires_is_parsed_as_an_absolute_instant(self): + msg = parse_message( + "[GPU CLAIM] node=n1 holder=@a vram=6gb expires=1783350000" + ) + assert msg.expires_at == 1783350000.0 + assert msg.expired(1783350001) is True + assert msg.expired(1783349999) is False + + def test_an_unparseable_expiry_means_no_expiry_not_expiry_at_zero(self): + # "expires=" must not read as expired-since-1970, which would + # silently turn a live claim into a free card. + msg = parse_message("[GPU CLAIM] node=n1 holder=@a vram=6gb expires=soon") + assert msg.expires_at is None + assert msg.expired(2**31) is False + def test_unknown_keys_are_ignored(self): msg = parse_message("[GPU CLAIM] node=n1 holder=@a vram=1gb priority=high") assert msg.node == "n1" @@ -165,6 +182,44 @@ def test_claim_without_release_stays_open(self): assert list(folded) == ["n1"] assert folded["n1"][0].vram_mb == 6144 + def test_a_claim_past_its_published_expiry_is_not_open(self): + # taOS #893 / CR on #2988: a holder that crashed (no RELEASE, no + # keep-alive) must not block the card forever. + msgs = self._bus("[GPU CLAIM] node=n1 holder=@a vram=6gb expires=1000") + assert open_claims(msgs, now=1001) == {} + assert list(open_claims(msgs, now=999)) == ["n1"] + # The boundary counts as expired, matching the cluster lease's TTL. + assert open_claims(msgs, now=1000) == {} + + def test_a_claim_without_a_published_expiry_never_expires(self): + # Interim hand-posted lines stay bounded by RELEASE / the fold window. + msgs = self._bus("[GPU CLAIM] node=n1 holder=@a vram=6gb") + assert list(open_claims(msgs, now=time.time() + 10**9)) == ["n1"] + + def test_a_refreshed_claim_carries_its_new_expiry(self): + msgs = self._bus( + "[GPU CLAIM] node=n1 holder=@a vram=6gb expires=1000", + "[GPU CLAIM] node=n1 holder=@a vram=6gb expires=2000", + ) + folded = open_claims(msgs, now=1500) + assert folded["n1"][0].expires_at == 2000 + + def test_one_holders_expiry_does_not_drop_another_holders_claim(self): + msgs = [ + { + "id": 1, + "from": "@a", + "body": "[GPU CLAIM] node=n1 holder=@a vram=6gb expires=1000", + }, + { + "id": 2, + "from": "@b", + "body": "[GPU CLAIM] node=n1 holder=@b vram=2gb expires=9999", + }, + ] + folded = open_claims(msgs, now=2000) + assert [c.display_holder() for c in folded["n1"]] == ["@b"] + def test_a_release_from_another_holder_does_not_close_the_claim(self): msgs = [ {"id": 1, "from": "@a", "body": "[GPU CLAIM] node=n1 holder=@a vram=6gb"}, @@ -374,6 +429,17 @@ def test_render_claim_includes_reason_and_eta(self): def test_render_claim_omits_empty_optionals(self): assert render_claim("n1", "@a", 1024) == "[GPU CLAIM] node=n1 holder=@a vram=~1gb" + def test_render_claim_publishes_an_integer_expiry(self): + line = render_claim("n1", "@a", 6144, expires_at=1783350000.9) + assert line == ( + "[GPU CLAIM] node=n1 holder=@a vram=~6gb expires=1783350000" + ) + # It round-trips as an absolute instant, which is what the fold compares. + assert parse_message(line).expires_at == 1783350000.0 + + def test_render_claim_omits_the_expiry_when_it_is_unknown(self): + assert "expires=" not in render_claim("n1", "@a", 6144) + def test_render_release_request_check(self): assert render_release("n1", "@a") == "[GPU RELEASE] node=n1 holder=@a" assert render_request("n1", 6144, "blocked") == ( diff --git a/tests/test_routes_a2a_gpu_lease.py b/tests/test_routes_a2a_gpu_lease.py index 4e956ce0f..bcc3f88a4 100644 --- a/tests/test_routes_a2a_gpu_lease.py +++ b/tests/test_routes_a2a_gpu_lease.py @@ -8,6 +8,7 @@ """ from __future__ import annotations +import asyncio import time import pytest @@ -18,6 +19,7 @@ from tinyagentos.agent_registry_store import mint_registry_token from tinyagentos.cluster.manager import ClusterManager from tinyagentos.cluster.worker_protocol import WorkerInfo +from tinyagentos.gpu_lease import claims_for_node, open_claims _ROUTE_PATCH = "tinyagentos.routes.a2a_gpu_lease.httpx.AsyncClient" @@ -319,9 +321,16 @@ async def test_claim_posts_the_protocol_line(self, lease_client, bus, local_vram assert resp.status_code == 200 data = resp.json() assert data["status"] == "claimed" - assert data["line"] == ( - "[GPU CLAIM] node=local holder=@operator vram=~6gb reason=flux eta=~5m" - ) + # The line carries the holder's own expiry: an integer instant, so a + # peer folds it against its own clock. This node has no cluster lease + # behind it, so it is the TTL this call asked for (the 300s default). + assert data["line"].startswith( + "[GPU CLAIM] node=local holder=@operator vram=~6gb " + "reason=flux eta=~5m expires=" + ) + published = int(data["line"].rsplit("expires=", 1)[1]) + assert abs(published - (time.time() + 300)) <= 5 + assert int(data["claim_expires_at"]) == published assert bus.last_line == data["line"] async def test_claim_then_check_sees_the_claim(self, lease_client, bus, local_vram): @@ -502,6 +511,107 @@ async def test_operator_release_closes_the_holders_bus_claim( assert checked.json()["admitted"] is True assert checked.json()["blockers"] == [] + async def test_claim_publishes_a_bus_expiry_from_its_lease( + self, lease_client, bus, cluster + ): + """A peer must be able to tell when a claim lapses (CR on #2988).""" + resp = await lease_client.post( + "/api/a2a/gpu/claim", json={"node": "linstation", "vram_mb": 4096} + ) + assert resp.status_code == 200 + lease = cluster.get_leases()[0] + assert resp.json()["claim_expires_at"] == lease.expires_at + assert f"expires={int(lease.expires_at)}" in bus.last_line + # The fold reads it, so the claim is bounded even without a RELEASE. + # (The wire carries whole seconds; the lease keeps its float.) + folded = open_claims(bus.messages) + assert claims_for_node(folded, "linstation")[0].expires_at == int( + lease.expires_at + ) + + async def test_a_claim_that_stops_being_kept_alive_frees_the_node( + self, lease_client, bus, cluster + ): + """CR on #2988: a crashed holder must not block the card forever. + + No RELEASE and no keep-alive: once the published expiry passes, another + identity can take the node. Otherwise one crashed agent denies the + shared GPU to everyone until its claim ages out of the fold window. + """ + _cid, token = await _agent_token(lease_client._app, scopes=("a2a_send",)) + async with _bare(lease_client._app) as bare: + claimed = await bare.post( + "/api/a2a/gpu/claim", + json={"node": "linstation", "vram_mb": 4096, "ttl_seconds": 0.5}, + headers={"Authorization": f"Bearer {token}"}, + ) + assert claimed.status_code == 200 + + _other, other = await _agent_token( + lease_client._app, scopes=("a2a_receive", "a2a_send"), handle="@taos" + ) + headers = {"Authorization": f"Bearer {other}"} + + # While the holder's claim is live, the node is blocked for the peer. + async with _bare(lease_client._app) as bare: + blocked = await bare.get( + "/api/a2a/gpu/check", + params={"node": "linstation", "vram_mb": 1024}, + headers=headers, + ) + assert blocked.status_code == 200 + assert blocked.json()["admitted"] is False + + await asyncio.sleep(0.8) + + async with _bare(lease_client._app) as bare: + freed = await bare.get( + "/api/a2a/gpu/check", + params={"node": "linstation", "vram_mb": 1024}, + headers=headers, + ) + took = await bare.post( + "/api/a2a/gpu/claim", + json={"node": "linstation", "vram_mb": 4096}, + headers=headers, + ) + assert freed.json()["admitted"] is True + assert freed.json()["blockers"] == [] + assert took.status_code == 200 + + async def test_renew_republishes_the_claim_keep_alive( + self, lease_client, bus, cluster + ): + """A renewed lease must keep its bus claim alive (CR on #2988). + + The fold drops a claim once its published expiry passes, so a holder + that only renewed the lease would let its claim lapse while still + holding the card - and a peer would read the node as free. + """ + claimed = await lease_client.post( + "/api/a2a/gpu/claim", json={"node": "linstation", "vram_mb": 4096} + ) + lease_id = claimed.json()["lease_id"] + before = cluster.get_leases()[0].expires_at + + renewed = await lease_client.post( + "/api/a2a/gpu/renew", + json={"lease_id": lease_id, "ttl_seconds": 600}, + ) + assert renewed.status_code == 200 + assert renewed.json()["bus_claim_refreshed"] is True + after = cluster.get_leases()[0].expires_at + assert after > before + assert bus.last_line == renewed.json()["line"] + assert bus.last_line.startswith("[GPU CLAIM] node=linstation holder=@operator") + assert "reason=keep-alive" in bus.last_line + assert f"expires={int(after)}" in bus.last_line + # The repost replaces the original claim rather than double-counting it. + folded = open_claims(bus.messages) + assert [int(c.expires_at) for c in claims_for_node(folded, "linstation")] == [ + int(after) + ] + async def test_reclaiming_extends_rather_than_conflicts(self, lease_client, bus, cluster): first = await lease_client.post( "/api/a2a/gpu/claim", json={"node": "linstation", "vram_mb": 4096} diff --git a/tinyagentos/gpu_lease.py b/tinyagentos/gpu_lease.py index 98ea5f79e..8af9a1398 100644 --- a/tinyagentos/gpu_lease.py +++ b/tinyagentos/gpu_lease.py @@ -10,6 +10,7 @@ format is a single readable line so a human watching the channel can follow it:: [GPU CLAIM] node=linstation holder=@taOSmd vram=~9.4gb reason=ollama eta=~10m + [GPU CLAIM] node=linstation holder=@taOSmd vram=~9.4gb expires=1783350000 [GPU RELEASE] node=linstation holder=@taOSmd [GPU REQUEST] node=linstation need=~6gb [GPU CHECK] node=linstation need=~6gb @@ -19,6 +20,16 @@ by a later RELEASE from the same holder). It is deliberately pure — no FastAPI, no httpx — so the admission rules can be tested without a bus. +Expiry +------ +A rendered CLAIM may carry the holder's own ``expires=`` (a unix timestamp, +taken from the backing cluster lease's TTL). The fold drops a claim whose +expiry has passed, exactly as the cluster lease's TTL frees the reservation it +backs: a crashed or idle holder must not block the card forever. A claim +published without ``expires=`` never expires here — it is bounded only by a +RELEASE or by ageing out of the fold window, which keeps the interim +hand-posted lines working unchanged. + Identity -------- The authoritative holder of a claim is the BUS MESSAGE AUTHOR (``from``), not @@ -31,7 +42,9 @@ from __future__ import annotations +import math import re +import time from dataclasses import dataclass from typing import Iterable, Mapping @@ -59,7 +72,7 @@ # Field names understood in a line body. Unknown keys are ignored (forward # compatibility: a future `priority=` must not break an older reader). -_FIELD_ALIASES = {"needed": "need", "vram_mb": "vram"} +_FIELD_ALIASES = {"needed": "need", "vram_mb": "vram", "expires_at": "expires"} def _clean(text: object) -> str: @@ -154,6 +167,16 @@ class GpuLeaseMessage: message_id: int | None = None ts: float | None = None bus_from: str | None = None + expires_at: float | None = None + + def expired(self, now: float) -> bool: + """True when the holder's published expiry has passed. + + A claim with no ``expires=`` (an interim hand-posted line) never + expires here: it is bounded only by a RELEASE or by ageing out of the + fold window. + """ + return self.expires_at is not None and self.expires_at <= now @property def identity_key(self) -> str: @@ -178,6 +201,23 @@ def as_dict(self) -> dict: } +def parse_epoch(text: object) -> float | None: + """Parse a unix-timestamp field (``expires=1783350000``) into seconds. + + Returns None for anything unparseable (or non-finite) so a caller can tell + "no expiry published" from "expired at zero", which would read as expired + since 1970. + """ + raw = str(text).strip() if text is not None else "" + if not raw: + return None + try: + value = float(raw) + except (TypeError, ValueError): + return None + return value if math.isfinite(value) else None + + def parse_message( raw: object, *, @@ -225,11 +265,17 @@ def parse_message( message_id=int(message_id) if isinstance(message_id, int) else None, ts=float(ts) if isinstance(ts, (int, float)) else None, bus_from=sender.strip() if isinstance(sender, str) and sender.strip() else None, + expires_at=parse_epoch(fields.get("expires")), ) def render_claim( - node: str, holder: str, vram_mb: int, reason: str = "", eta: str = "" + node: str, + holder: str, + vram_mb: int, + reason: str = "", + eta: str = "", + expires_at: float | None = None, ) -> str: """Render a ``[GPU CLAIM]`` line. Every value is flattened to one line.""" parts = [ @@ -241,6 +287,10 @@ def render_claim( parts.append(f"reason={_clean(reason)}") if eta: parts.append(f"eta={_clean(eta)}") + if expires_at is not None: + # An integer unix timestamp: the fold compares it against its own clock, + # so it must be an absolute instant, never a duration. + parts.append(f"expires={int(expires_at)}") return f"[GPU CLAIM] {' '.join(parts)}" @@ -265,7 +315,9 @@ def render_check(node: str, need_mb: int | None = None) -> str: return f"[GPU CHECK] {' '.join(parts)}" -def open_claims(messages: Iterable[object]) -> dict[str, list[GpuLeaseMessage]]: +def open_claims( + messages: Iterable[object], *, now: float | None = None +) -> dict[str, list[GpuLeaseMessage]]: """Fold a channel's messages into the claims that are still open. Messages are expected oldest-first (the bus returns them in that order). A @@ -273,6 +325,12 @@ def open_claims(messages: Iterable[object]) -> dict[str, list[GpuLeaseMessage]]: identity on the same node closes it. A re-CLAIM by the same identity replaces the previous one (so a reposted claim is never double-counted). + A claim that published an ``expires=`` in the past is NOT open: it is + dropped here exactly as the cluster lease it backs is dropped by its TTL, + so a holder that crashed (or stopped keeping alive) cannot block the card + forever. A claim with no published expiry is unaffected - it stays open + until a RELEASE closes it or it ages out of the fold window. + Both the node and the identity are matched case-insensitively: peers do not agree on the spelling of a hostname, and a `node=Linstation` claim whose release says `node=linstation` would otherwise stay open forever while a @@ -280,6 +338,7 @@ def open_claims(messages: Iterable[object]) -> dict[str, list[GpuLeaseMessage]]: therefore keyed by the CASE-FOLDED node name; each message keeps the spelling it arrived with. """ + cutoff = time.time() if now is None else float(now) open_by_key: dict[tuple[str, str], GpuLeaseMessage] = {} order: list[tuple[str, str]] = [] for raw in messages: @@ -297,8 +356,9 @@ def open_claims(messages: Iterable[object]) -> dict[str, list[GpuLeaseMessage]]: out: dict[str, list[GpuLeaseMessage]] = {} for key in order: msg = open_by_key.get(key) - if msg is not None: - out.setdefault(key[0], []).append(msg) + if msg is None or msg.expired(cutoff): + continue + out.setdefault(key[0], []).append(msg) return out diff --git a/tinyagentos/routes/a2a_gpu_lease.py b/tinyagentos/routes/a2a_gpu_lease.py index eb8b4a223..afe39cb15 100644 --- a/tinyagentos/routes/a2a_gpu_lease.py +++ b/tinyagentos/routes/a2a_gpu_lease.py @@ -45,6 +45,7 @@ import logging import os import socket +import time from dataclasses import dataclass import httpx @@ -443,6 +444,9 @@ class RenewBody(BaseModel): ttl_seconds: float = Field( default=_DEFAULT_TTL_SECONDS, gt=0, le=MAX_LEASE_TTL_SECONDS ) + # The channel the claim was posted on, so the keep-alive repost refreshes + # the same line instead of leaking one onto the default thread. + channel: str | None = None def _resource_id(node: str, resource: str) -> str: @@ -505,6 +509,30 @@ async def _holder_for(request: Request, identity: str) -> str: return handle if handle.startswith("@") else f"@{handle}" +async def _claim_holder_actor(request: Request, lease, actor: _Actor) -> _Actor: + """The actor a line ABOUT *lease* must be attributed to. + + A claim is keyed on its bus author, so an operator (or any caller acting on + a lease it does not hold) must post as the holder whose claim the line + closes - otherwise the post clears nothing and peers keep reading the node + as claimed while the local lease is already gone (an admin session may post + with an explicit ``from``, docs/agent-coordination.md). A lease with no bus + claim behind it (a non-``a2a:`` caller, e.g. ``skald-dispatcher``) has no + holder to attribute to, so the acting identity stands. + """ + if _lease_owned_by(lease, actor): + return actor + owner = _lease_bus_identity(lease) + if owner is None: + return actor + return _Actor( + identity=owner, + holder=await _holder_for(request, owner), + credential=actor.credential, + is_admin=actor.is_admin, + ) + + def _may_act_on(lease, actor: _Actor) -> bool: """Ownership for an EXPLICIT lease id: the holder, or an operator.""" return _lease_owned_by(lease, actor) or actor.is_admin @@ -642,7 +670,24 @@ async def gpu_claim(request: Request, body: ClaimBody): lease_id = lease.lease_id created_lease = True - line = render_claim(node, actor.holder, vram_mb, body.reason, body.eta) + # Publish the holder's own expiry on the line, so a peer's fold can drop a + # claim whose holder crashed or stopped keeping alive instead of blocking + # the card forever. It is the backing cluster lease's expiry when there is + # one, else the TTL this call asked for (a bus-only node is governed by the + # bus alone, so the line is the only thing that can free it). + claim_expires_at = ( + getattr(lease, "expires_at", None) + if lease is not None + else time.time() + float(body.ttl_seconds) + ) + line = render_claim( + node, + actor.holder, + vram_mb, + body.reason, + body.eta, + expires_at=claim_expires_at, + ) try: posted = await _post_line(channel, actor, line) except HTTPException: @@ -660,6 +705,7 @@ async def gpu_claim(request: Request, body: ClaimBody): "vram_mb": vram_mb, "lease_id": lease_id, "expires_at": getattr(lease, "expires_at", None), + "claim_expires_at": claim_expires_at, "line": line, "channel": channel, "message": posted, @@ -708,21 +754,12 @@ async def gpu_release(request: Request, body: ReleaseBody): # lease by explicit id must attribute the line to THAT holder: a claim is # keyed on its bus AUTHOR, so a line posted as @operator clears nothing and # every peer's fold keeps reading the node as claimed while the local lease - # is already gone (CodeRabbit on #2988). An admin session may post with an - # explicit `from` (docs/agent-coordination.md, *Posting to the coordination - # bus*); a bus that authenticates senders refuses the substitution, and the - # post-before-release ordering below then leaves the local lease intact, so - # the override cannot half-apply. - line_actor = actor - if lease is not None and not _lease_owned_by(lease, actor): - owner = _lease_bus_identity(lease) - if owner is not None: - line_actor = _Actor( - identity=owner, - holder=await _holder_for(request, owner), - credential=actor.credential, - is_admin=actor.is_admin, - ) + # is already gone (CodeRabbit on #2988). A bus that authenticates senders + # refuses the substitution, and the post-before-release ordering below then + # leaves the local lease intact, so the override cannot half-apply. + line_actor = ( + actor if lease is None else await _claim_holder_actor(request, lease, actor) + ) # Post BEFORE releasing the local lease, so a bus failure changes nothing # and the caller can retry. Releasing first would free the node here while @@ -799,9 +836,45 @@ async def gpu_renew(request: Request, body: RenewBody): {"error": "lease not found or expired", "lease_id": body.lease_id}, status_code=409, ) + # Keep the BUS claim alive too. A fold drops a claim whose published expiry + # has passed, so a holder that only renewed its lease would let the claim + # lapse while it still holds the card - and a peer's fold would then read + # the node as free, which is the co-load this surface exists to prevent. + # The repost is the claim line shape, so the fold replaces the previous + # claim (same node + identity) rather than double-counting it. + node, _, _resource = (lease.resource_id or "").partition(":") + hold = await _claim_holder_actor(request, lease, actor) + channel = (body.channel or "").strip() or _channel() + line = render_claim( + node, + hold.holder, + lease.required_vram_mb or 0, + reason="keep-alive", + expires_at=lease.expires_at, + ) + bus_claim_refreshed = False + refresh_error: str | None = None + if node: + try: + await _post_line(channel, hold, line) + bus_claim_refreshed = True + except HTTPException: + # The local lease IS renewed, so this is reported rather than fatal: + # the caller must know the claim did not reach the bus (peers will + # treat the node as free once the old expiry passes) and re-claim. + refresh_error = "a2a bus unavailable" + logger.warning( + "A2A GPU lease %s renewed but the claim repost failed (channel=%s)", + lease.lease_id, + channel, + ) return { "status": "renewed", "lease_id": lease.lease_id, "resource_id": lease.resource_id, "expires_at": lease.expires_at, + "line": line, + "channel": channel, + "bus_claim_refreshed": bus_claim_refreshed, + "bus_refresh_error": refresh_error, } From 71a167a30e6621315fc56b78bb33e46e7ac74ac5 Mon Sep 17 00:00:00 2001 From: hognek <227774406+hognek@users.noreply.github.com> Date: Sat, 12 Sep 2026 01:44:02 +0200 Subject: [PATCH 08/17] fix(a2a): round the claim expiry up, pin its channel, roll back a failed keep-alive CodeRabbit round 4 on #2988, each finding verified against the code: - **Truncation published an expiry BEFORE the lease ended.** `int(expires_at)` turned a lease ending at 1000.9 into `expires=1000`, so a peer could be admitted in the gap while the reservation was still live. The rendered instant is now rounded UP (`math.ceil`): a published expiry never precedes the reservation it describes. - **A renewal could refresh a different channel than the claim.** `RenewBody` accepted a caller-supplied `channel`, so a keep-alive could repost onto another thread while the original claim expired. The resolved channel is stored on the lease (`GpuLease.claim_channel`, set by `gpu_claim`) and reused by `gpu_renew`; `RenewBody.channel` is gone, since the channel is an input to the *claim* and never to its renewal. - **A failed repost left a renewed lease that peers had already seen lapse.** Half a renewal is no renewal: on a failed claim repost the local extension is rolled back to the previous instant (`ClusterManager.restore_lease_expiry`), so both views agree and the caller retries, with `bus_claim_refreshed: false` and `bus_refresh_error` reporting why. Tests: 103 in the two suites (2 new: the channel pin and the rollback). --- changelog.d/taos-893-a2a-gpu-lease.md | 15 +++-- docs/agent-coordination.md | 8 ++- tests/test_gpu_lease_protocol.py | 10 ++-- tests/test_routes_a2a_gpu_lease.py | 76 ++++++++++++++++++++++---- tinyagentos/cluster/manager.py | 21 +++++++ tinyagentos/cluster/worker_protocol.py | 6 ++ tinyagentos/gpu_lease.py | 7 ++- tinyagentos/routes/a2a_gpu_lease.py | 25 ++++++--- 8 files changed, 135 insertions(+), 33 deletions(-) diff --git a/changelog.d/taos-893-a2a-gpu-lease.md b/changelog.d/taos-893-a2a-gpu-lease.md index c5043cc0e..8e9aca501 100644 --- a/changelog.d/taos-893-a2a-gpu-lease.md +++ b/changelog.d/taos-893-a2a-gpu-lease.md @@ -19,12 +19,15 @@ own shared VRAM ledger. - Claims carry their own expiry. `POST /api/a2a/gpu/claim` publishes `expires=` on the `[GPU CLAIM]` line (the backing cluster lease's - expiry, or the requested TTL for a bus-only node) and `POST /api/a2a/gpu/renew` - reposts that line as it extends the lease, so a holder that keeps its lease - alive keeps its claim alive. The fold drops a claim whose published expiry has - passed, so a holder that crashed or stopped keeping alive no longer blocks the - shared card until its claim ages out of the fold window. A claim posted by hand - without an `expires=` is unchanged: bounded by a RELEASE alone. + expiry, or the requested TTL for a bus-only node), rounded up so a published + expiry can never precede the lease it describes, and `POST /api/a2a/gpu/renew` + reposts that line as it extends the lease — on the channel the claim was made + on, since the channel is an input to the claim and never to its renewal. If the + repost fails the local extension is rolled back rather than leaving a lease + that peers have already seen lapse. The fold drops a claim whose published + expiry has passed, so a holder that crashed or stopped keeping alive no longer + blocks the shared card until its claim ages out of the fold window. A claim + posted by hand without an `expires=` is unchanged: bounded by a RELEASE alone. ### Fixed diff --git a/docs/agent-coordination.md b/docs/agent-coordination.md index ddaaf86c5..81c572169 100644 --- a/docs/agent-coordination.md +++ b/docs/agent-coordination.md @@ -425,8 +425,12 @@ Rules that matter when you use it: cap is enforced by the request model. - **A claim carries its own expiry, and the fold honours it.** `/claim` publishes `expires=` on the line — the backing lease's expiry when the node is a - cluster worker, else the TTL it was asked for — and `/renew` reposts the claim - as it extends the lease. A fold drops a claim whose published expiry has + cluster worker, else the TTL it was asked for — rounded up, so a published + expiry can never precede the reservation it describes. `/renew` reposts the + claim as it extends the lease, on the channel the claim was made on (the + channel is an input to the *claim*, never to its renewal); if that repost + fails, the local extension is rolled back, so peers and this controller still + agree and the holder can retry. A fold drops a claim whose published expiry has passed, exactly as the cluster lease's TTL frees its reservation, so a holder that crashed or stopped keeping alive no longer blocks the card until its claim ages out of the fold window. Keep-alive therefore means re-POST `/claim` diff --git a/tests/test_gpu_lease_protocol.py b/tests/test_gpu_lease_protocol.py index cf9d960be..d78c81d68 100644 --- a/tests/test_gpu_lease_protocol.py +++ b/tests/test_gpu_lease_protocol.py @@ -431,11 +431,11 @@ def test_render_claim_omits_empty_optionals(self): def test_render_claim_publishes_an_integer_expiry(self): line = render_claim("n1", "@a", 6144, expires_at=1783350000.9) - assert line == ( - "[GPU CLAIM] node=n1 holder=@a vram=~6gb expires=1783350000" - ) - # It round-trips as an absolute instant, which is what the fold compares. - assert parse_message(line).expires_at == 1783350000.0 + # Rounded UP: a published expiry must never precede the reservation it + # describes, or a peer could admit itself into the truncation gap. + assert line == "[GPU CLAIM] node=n1 holder=@a vram=~6gb expires=1783350001" + assert parse_message(line).expires_at == 1783350001.0 + assert parse_message(line).expires_at >= 1783350000.9 def test_render_claim_omits_the_expiry_when_it_is_unknown(self): assert "expires=" not in render_claim("n1", "@a", 6144) diff --git a/tests/test_routes_a2a_gpu_lease.py b/tests/test_routes_a2a_gpu_lease.py index bcc3f88a4..87bbcc0f5 100644 --- a/tests/test_routes_a2a_gpu_lease.py +++ b/tests/test_routes_a2a_gpu_lease.py @@ -9,6 +9,7 @@ from __future__ import annotations import asyncio +import math import time import pytest @@ -330,7 +331,7 @@ async def test_claim_posts_the_protocol_line(self, lease_client, bus, local_vram ) published = int(data["line"].rsplit("expires=", 1)[1]) assert abs(published - (time.time() + 300)) <= 5 - assert int(data["claim_expires_at"]) == published + assert math.ceil(data["claim_expires_at"]) == published assert bus.last_line == data["line"] async def test_claim_then_check_sees_the_claim(self, lease_client, bus, local_vram): @@ -521,13 +522,13 @@ async def test_claim_publishes_a_bus_expiry_from_its_lease( assert resp.status_code == 200 lease = cluster.get_leases()[0] assert resp.json()["claim_expires_at"] == lease.expires_at - assert f"expires={int(lease.expires_at)}" in bus.last_line + # The wire carries whole seconds, rounded UP so a published expiry can + # never precede the lease it describes. + published = int(bus.last_line.rsplit("expires=", 1)[1]) + assert lease.expires_at <= published < lease.expires_at + 1 # The fold reads it, so the claim is bounded even without a RELEASE. - # (The wire carries whole seconds; the lease keeps its float.) folded = open_claims(bus.messages) - assert claims_for_node(folded, "linstation")[0].expires_at == int( - lease.expires_at - ) + assert claims_for_node(folded, "linstation")[0].expires_at == float(published) async def test_a_claim_that_stops_being_kept_alive_frees_the_node( self, lease_client, bus, cluster @@ -562,7 +563,10 @@ async def test_a_claim_that_stops_being_kept_alive_frees_the_node( assert blocked.status_code == 200 assert blocked.json()["admitted"] is False - await asyncio.sleep(0.8) + # The published expiry is the lease's TTL rounded UP to whole seconds, + # so wait for the instant the peers will actually fold it as lapsed. + published = int(claimed.json()["line"].rsplit("expires=", 1)[1]) + await asyncio.sleep(max(0.5, published - time.time() + 0.1)) async with _bare(lease_client._app) as bare: freed = await bare.get( @@ -605,13 +609,65 @@ async def test_renew_republishes_the_claim_keep_alive( assert bus.last_line == renewed.json()["line"] assert bus.last_line.startswith("[GPU CLAIM] node=linstation holder=@operator") assert "reason=keep-alive" in bus.last_line - assert f"expires={int(after)}" in bus.last_line + published = int(bus.last_line.rsplit("expires=", 1)[1]) + assert after <= published < after + 1 # The repost replaces the original claim rather than double-counting it. folded = open_claims(bus.messages) - assert [int(c.expires_at) for c in claims_for_node(folded, "linstation")] == [ - int(after) + assert [c.expires_at for c in claims_for_node(folded, "linstation")] == [ + float(published) ] + async def test_renew_refreshes_the_channel_the_claim_was_made_on( + self, lease_client, bus, cluster + ): + """The channel is an input to the CLAIM, not to its renewal (CR #2988). + + Refreshing onto a different thread would leave the original claim to + expire while this lease is still held. + """ + claimed = await lease_client.post( + "/api/a2a/gpu/claim", + json={"node": "linstation", "vram_mb": 4096, "channel": "gpu-lab"}, + ) + assert claimed.status_code == 200 + assert claimed.json()["channel"] == "gpu-lab" + assert bus.sends[-1]["payload"]["thread"] == "gpu-lab" + + renewed = await lease_client.post( + "/api/a2a/gpu/renew", + json={"lease_id": claimed.json()["lease_id"], "ttl_seconds": 600}, + ) + assert renewed.status_code == 200 + assert renewed.json()["channel"] == "gpu-lab" + assert bus.sends[-1]["payload"]["thread"] == "gpu-lab" + + async def test_a_failed_keep_alive_repost_rolls_the_renewal_back( + self, lease_client, bus, cluster + ): + """Half a renewal is no renewal (CR #2988). + + If the claim cannot be refreshed, the local extension must not stand: + peers would free the card at the expiry they still hold while this + controller believes it is reserved. + """ + claimed = await lease_client.post( + "/api/a2a/gpu/claim", json={"node": "linstation", "vram_mb": 4096} + ) + lease_id = claimed.json()["lease_id"] + before = cluster.get_leases()[0].expires_at + + bus.fail_post = True + renewed = await lease_client.post( + "/api/a2a/gpu/renew", + json={"lease_id": lease_id, "ttl_seconds": 600}, + ) + assert renewed.status_code == 200 + assert renewed.json()["bus_claim_refreshed"] is False + assert renewed.json()["bus_refresh_error"] == "a2a bus unavailable" + # Rolled back: the lease still ends when the published claim does. + assert renewed.json()["expires_at"] == before + assert cluster.get_leases()[0].expires_at == before + async def test_reclaiming_extends_rather_than_conflicts(self, lease_client, bus, cluster): first = await lease_client.post( "/api/a2a/gpu/claim", json={"node": "linstation", "vram_mb": 4096} diff --git a/tinyagentos/cluster/manager.py b/tinyagentos/cluster/manager.py index 9e9ee7351..54c4ab79d 100644 --- a/tinyagentos/cluster/manager.py +++ b/tinyagentos/cluster/manager.py @@ -806,6 +806,7 @@ async def claim_lease( caller: str = "", ttl_seconds: float = 30, required_vram_mb: int = 0, + claim_channel: str = "", ) -> GpuLease | None: """Attempt to claim a GPU lease on ``resource_id``. @@ -875,6 +876,7 @@ async def claim_lease( caller=caller, expires_at=time.time() + ttl_seconds, required_vram_mb=required_vram_mb, + claim_channel=claim_channel, ) self._leases[lease_id] = lease logger.info( @@ -905,6 +907,25 @@ async def renew_lease(self, lease_id: str, ttl_seconds: float = 30) -> GpuLease lease.expires_at = now + ttl_seconds return lease + async def restore_lease_expiry(self, lease_id: str, expires_at: float) -> bool: + """Put a lease's expiry back after a keep-alive's other half failed. + + Renewal has two halves: the local reservation and the peer-visible + claim published on the bus. When the second cannot be refreshed, the + first must not stay extended - peers would then free the card at the + expiry they still hold while this controller believes it is reserved. + Restoring the previous instant makes the two views agree again, so the + caller can retry rather than sit on a renewal nobody else can see. + + Returns False when the lease is already gone (nothing to restore). + """ + async with self._lease_lock: + lease = self._leases.get(lease_id) + if lease is None: + return False + lease.expires_at = expires_at + return True + def get_leases(self) -> list[GpuLease]: """Return a snapshot of active (non-expired) leases.""" now = time.time() diff --git a/tinyagentos/cluster/worker_protocol.py b/tinyagentos/cluster/worker_protocol.py index f1890540a..253dde2aa 100644 --- a/tinyagentos/cluster/worker_protocol.py +++ b/tinyagentos/cluster/worker_protocol.py @@ -96,9 +96,15 @@ class GpuLease: required_vram_mb: How many MiB of VRAM the caller declared it needs. Used by the pre-claim check to refuse a claim when the worker's ``free_vram_mb`` is too low. + claim_channel: The coordination thread this lease's peer-visible + claim was published on (empty for a lease with no bus claim, + e.g. one taken by the dispatcher through the cluster API). + Kept so a keep-alive refreshes the SAME thread: the channel is + an input to the claim, never to its renewal. """ lease_id: str resource_id: str caller: str = "" expires_at: float = 0.0 required_vram_mb: int = 0 + claim_channel: str = "" diff --git a/tinyagentos/gpu_lease.py b/tinyagentos/gpu_lease.py index 8af9a1398..29d12dad8 100644 --- a/tinyagentos/gpu_lease.py +++ b/tinyagentos/gpu_lease.py @@ -289,8 +289,11 @@ def render_claim( parts.append(f"eta={_clean(eta)}") if expires_at is not None: # An integer unix timestamp: the fold compares it against its own clock, - # so it must be an absolute instant, never a duration. - parts.append(f"expires={int(expires_at)}") + # so it must be an absolute instant, never a duration. Rounded UP, so a + # published expiry can never precede the reservation it describes - a + # lease ending at 1000.9 must not publish expires=1000, or a peer could + # admit itself into the truncation gap (CR on #2988). + parts.append(f"expires={math.ceil(expires_at)}") return f"[GPU CLAIM] {' '.join(parts)}" diff --git a/tinyagentos/routes/a2a_gpu_lease.py b/tinyagentos/routes/a2a_gpu_lease.py index afe39cb15..290377728 100644 --- a/tinyagentos/routes/a2a_gpu_lease.py +++ b/tinyagentos/routes/a2a_gpu_lease.py @@ -444,9 +444,6 @@ class RenewBody(BaseModel): ttl_seconds: float = Field( default=_DEFAULT_TTL_SECONDS, gt=0, le=MAX_LEASE_TTL_SECONDS ) - # The channel the claim was posted on, so the keep-alive repost refreshes - # the same line instead of leaking one onto the default thread. - channel: str | None = None def _resource_id(node: str, resource: str) -> str: @@ -657,6 +654,7 @@ async def gpu_claim(request: Request, body: ClaimBody): caller=caller, ttl_seconds=float(body.ttl_seconds), required_vram_mb=vram_mb, + claim_channel=channel, ) if lease is None: return JSONResponse( @@ -830,6 +828,11 @@ async def gpu_renew(request: Request, body: RenewBody): ) if not _may_act_on(existing, actor): return JSONResponse({"error": "not the lease holder"}, status_code=403) + # The renewal's other half is the bus claim. Keep the instant before the + # extension: if the claim cannot be refreshed the local expiry is rolled + # back to it, so the two views agree rather than this controller holding a + # reservation every peer has already seen lapse. + previous_expiry = existing.expires_at lease = await cluster.renew_lease(body.lease_id, ttl_seconds=float(body.ttl_seconds)) if lease is None: return JSONResponse( @@ -844,7 +847,10 @@ async def gpu_renew(request: Request, body: RenewBody): # claim (same node + identity) rather than double-counting it. node, _, _resource = (lease.resource_id or "").partition(":") hold = await _claim_holder_actor(request, lease, actor) - channel = (body.channel or "").strip() or _channel() + # The channel is an input to the CLAIM, never to its renewal: refreshing + # onto a different thread would leave the original claim to expire while + # this lease is still held (CR on #2988). + channel = getattr(lease, "claim_channel", "") or _channel() line = render_claim( node, hold.holder, @@ -859,12 +865,15 @@ async def gpu_renew(request: Request, body: RenewBody): await _post_line(channel, hold, line) bus_claim_refreshed = True except HTTPException: - # The local lease IS renewed, so this is reported rather than fatal: - # the caller must know the claim did not reach the bus (peers will - # treat the node as free once the old expiry passes) and re-claim. + # Half a renewal is no renewal: the claim never reached the bus, so + # peers free the card at the expiry they still hold. Undo the local + # extension and report it, rather than hold a reservation nobody + # else can see. refresh_error = "a2a bus unavailable" + if not await cluster.restore_lease_expiry(lease.lease_id, previous_expiry): + refresh_error = "a2a bus unavailable; lease already expired" logger.warning( - "A2A GPU lease %s renewed but the claim repost failed (channel=%s)", + "A2A GPU lease %s renewal rolled back: claim repost failed (channel=%s)", lease.lease_id, channel, ) From 4bbe24d79665184bb980de6c67d03a643bed0194 Mon Sep 17 00:00:00 2001 From: hognek <227774406+hognek@users.noreply.github.com> Date: Sat, 12 Sep 2026 02:03:30 +0200 Subject: [PATCH 09/17] fix(cluster): keep a lease rollback from clobbering a newer renewal Kilo (CRITICAL) and CodeRabbit (Major) on #2988 raised the same race: `gpu_renew` posts the claim outside `_lease_lock`, so another renewal can extend the lease while the first one's bus post is in flight. `restore_lease_expiry` then restored unconditionally, so a failing renewal's rollback would silently undo the renewal that succeeded - whose response and bus claim still describe an expiry the lease no longer has. The rollback is now a compare-and-set: it restores only while the lease still carries the expiry THIS caller attempted, and the route reports the difference (`a2a bus unavailable; a newer renewal stands`). Tests: 1 new manager test pinning the supersede guard; 278 passed across the cluster, lease and a2a GPU route suites. --- tests/test_cluster.py | 43 +++++++++++++++++++++++++++++ tinyagentos/cluster/manager.py | 12 ++++++-- tinyagentos/routes/a2a_gpu_lease.py | 14 ++++++++-- 3 files changed, 64 insertions(+), 5 deletions(-) diff --git a/tests/test_cluster.py b/tests/test_cluster.py index 4e6bd9980..5879c1cb2 100644 --- a/tests/test_cluster.py +++ b/tests/test_cluster.py @@ -422,6 +422,49 @@ async def test_draining_workers_excluded_from_lease_claim(self): ) assert lease is None + async def test_restore_lease_expiry_refuses_to_clobber_a_newer_renewal(self): + """A rollback must not undo a renewal that landed in the meantime. + + `gpu_renew` posts to the bus outside `_lease_lock`, so another renewal + can extend the lease while the first one's post is in flight. The + rollback is therefore a compare-and-set on the expiry THIS caller + attempted (Kilo/CR review of #2988). + """ + mgr = ClusterManager() + await mgr.register_worker(_make_worker("gpu-box", url="http://gpu-box:9000")) + mgr.get_worker("gpu-box").free_vram_mb = 8000 + lease = await mgr.claim_lease( + resource_id="gpu-box:gpu-cuda-0", caller="a2a:@peer", ttl_seconds=300 + ) + assert lease is not None + previous = lease.expires_at + + ours = await mgr.renew_lease(lease.lease_id, ttl_seconds=600) + attempted = ours.expires_at + # Another renewal moves the expiry on while our bus post is in flight. + later = await mgr.renew_lease(lease.lease_id, ttl_seconds=900) + assert later.expires_at > attempted + + restored = await mgr.restore_lease_expiry( + lease.lease_id, previous, attempted_expiry=attempted + ) + assert restored is False + assert mgr.get_leases()[0].expires_at == later.expires_at + + # With nothing superseding us, the rollback applies. + restored = await mgr.restore_lease_expiry( + lease.lease_id, previous, attempted_expiry=later.expires_at + ) + assert restored is True + assert mgr.get_leases()[0].expires_at == previous + + # A lease that is gone needs no restore. + await mgr.release_lease(lease.lease_id) + restored = await mgr.restore_lease_expiry( + lease.lease_id, previous, attempted_expiry=previous + ) + assert restored is False + # ── Worker-initiated drain (taOS #890 C2) ────────────────────────── async def test_heartbeat_status_draining_triggers_worker_self_drain(self): diff --git a/tinyagentos/cluster/manager.py b/tinyagentos/cluster/manager.py index 54c4ab79d..467b33bd8 100644 --- a/tinyagentos/cluster/manager.py +++ b/tinyagentos/cluster/manager.py @@ -907,7 +907,9 @@ async def renew_lease(self, lease_id: str, ttl_seconds: float = 30) -> GpuLease lease.expires_at = now + ttl_seconds return lease - async def restore_lease_expiry(self, lease_id: str, expires_at: float) -> bool: + async def restore_lease_expiry( + self, lease_id: str, expires_at: float, *, attempted_expiry: float + ) -> bool: """Put a lease's expiry back after a keep-alive's other half failed. Renewal has two halves: the local reservation and the peer-visible @@ -917,12 +919,18 @@ async def restore_lease_expiry(self, lease_id: str, expires_at: float) -> bool: Restoring the previous instant makes the two views agree again, so the caller can retry rather than sit on a renewal nobody else can see. - Returns False when the lease is already gone (nothing to restore). + Restores only while the lease still carries *attempted_expiry*, the + extension this caller made. The bus post happens outside ``_lease_lock``, + so a renewal can land in between; that newer expiry owns the lease and + must not be clobbered by this rollback. Returns False when the lease is + gone or superseded - i.e. when there is nothing of ours to undo. """ async with self._lease_lock: lease = self._leases.get(lease_id) if lease is None: return False + if lease.expires_at != attempted_expiry: + return False lease.expires_at = expires_at return True diff --git a/tinyagentos/routes/a2a_gpu_lease.py b/tinyagentos/routes/a2a_gpu_lease.py index 290377728..1b5cc1e3c 100644 --- a/tinyagentos/routes/a2a_gpu_lease.py +++ b/tinyagentos/routes/a2a_gpu_lease.py @@ -845,6 +845,7 @@ async def gpu_renew(request: Request, body: RenewBody): # the node as free, which is the co-load this surface exists to prevent. # The repost is the claim line shape, so the fold replaces the previous # claim (same node + identity) rather than double-counting it. + attempted_expiry = lease.expires_at node, _, _resource = (lease.resource_id or "").partition(":") hold = await _claim_holder_actor(request, lease, actor) # The channel is an input to the CLAIM, never to its renewal: refreshing @@ -870,11 +871,18 @@ async def gpu_renew(request: Request, body: RenewBody): # extension and report it, rather than hold a reservation nobody # else can see. refresh_error = "a2a bus unavailable" - if not await cluster.restore_lease_expiry(lease.lease_id, previous_expiry): - refresh_error = "a2a bus unavailable; lease already expired" + restored = await cluster.restore_lease_expiry( + lease.lease_id, previous_expiry, attempted_expiry=attempted_expiry + ) + if not restored: + # The lease went away, or another renewal landed while this + # one's bus post was in flight: that newer expiry stands. + refresh_error = "a2a bus unavailable; a newer renewal stands" logger.warning( - "A2A GPU lease %s renewal rolled back: claim repost failed (channel=%s)", + "A2A GPU lease %s keep-alive rolled back (restored=%s): " + "claim repost failed (channel=%s)", lease.lease_id, + restored, channel, ) return { From abfe524be678dfaad9930cc159ef15dd05350a53 Mon Sep 17 00:00:00 2001 From: hognek <227774406+hognek@users.noreply.github.com> Date: Sat, 12 Sep 2026 02:31:06 +0200 Subject: [PATCH 10/17] fix(a2a): stop reading the body's holder as the admin's identity CodeRabbit on #2988: `_resolve_actor` built the session-admin actor with `identity = body.holder or "@operator"`, and `_lease_owned_by` matched a lease against both the identity and the holder. An admin session presenting an `a2a:`-looking `holder` therefore satisfied the ownership check on the node-scoped release/renew path (which takes no lease id) and could free a lease it did not hold. A session admin now acts as the fixed `@operator` principal and `holder` is display data only; ownership is an identity match and nothing else. An operator still frees any lease by EXPLICIT id - that is the operator path - and the freed holder is still the identity the RELEASE line is posted under. Tests: 104 in the two suites (1 new: an admin cannot take ownership by holder). --- docs/agent-coordination.md | 5 +++- tests/test_routes_a2a_gpu_lease.py | 38 +++++++++++++++++++++++++++++ tinyagentos/routes/a2a_gpu_lease.py | 20 +++++++++++---- 3 files changed, 57 insertions(+), 6 deletions(-) diff --git a/docs/agent-coordination.md b/docs/agent-coordination.md index 81c572169..7f556d517 100644 --- a/docs/agent-coordination.md +++ b/docs/agent-coordination.md @@ -410,7 +410,10 @@ Rules that matter when you use it: lease intact rather than freeing a node peers still see as claimed. - **An operator's release is attributed to the holder whose claim it closes.** A lease taken through `/claim` can also be freed by an explicit `lease_id` — - by its holder, or by an operator (`_may_act_on`). Since a bus claim is keyed + by its holder, or by an operator (`_may_act_on`). The node-scoped form (no + `lease_id`) only ever selects the caller's OWN lease: a body `holder` is + display text and never an identity, so it cannot be used to select someone + else's lease. Since a bus claim is keyed on its **author**, the operator's `[GPU RELEASE]` is posted as the freed holder, not as the operator (an admin session may set an explicit `from`, see *Posting to the coordination bus*); the response reports both, `holder` (who diff --git a/tests/test_routes_a2a_gpu_lease.py b/tests/test_routes_a2a_gpu_lease.py index 87bbcc0f5..c4a929f93 100644 --- a/tests/test_routes_a2a_gpu_lease.py +++ b/tests/test_routes_a2a_gpu_lease.py @@ -668,6 +668,44 @@ async def test_a_failed_keep_alive_repost_rolls_the_renewal_back( assert renewed.json()["expires_at"] == before assert cluster.get_leases()[0].expires_at == before + async def test_an_admin_cannot_take_ownership_of_an_agent_lease_by_holder( + self, lease_client, bus, cluster + ): + """`holder` is display data, never an identity (CR on #2988). + + A session admin acts as a fixed principal. If the body's `holder` were + read as that identity, a holder spelled like an agent's `a2a:` lease + would satisfy the ownership check on the node-scoped release path (which + takes no lease id) and free a lease the admin does not hold. + """ + cid, token = await _agent_token(lease_client._app, scopes=("a2a_send",)) + async with _bare(lease_client._app) as bare: + claimed = await bare.post( + "/api/a2a/gpu/claim", + json={"node": "linstation", "vram_mb": 4096}, + headers={"Authorization": f"Bearer {token}"}, + ) + assert claimed.status_code == 200 + held = cluster.get_leases()[0].lease_id + + # Node-scoped release (no id) while presenting the agent's identity. + spoofed = await lease_client.post( + "/api/a2a/gpu/release", + json={"node": "linstation", "holder": cid}, + ) + assert spoofed.status_code == 200 + assert spoofed.json()["lease_id"] is None + assert [lease.lease_id for lease in cluster.get_leases()] == [held] + + # The operator path is the explicit id, attributed to the holder. + released = await lease_client.post( + "/api/a2a/gpu/release", + json={"node": "linstation", "lease_id": held}, + ) + assert released.status_code == 200 + assert released.json()["released_holder"] == "@taosmd" + assert cluster.get_leases() == [] + async def test_reclaiming_extends_rather_than_conflicts(self, lease_client, bus, cluster): first = await lease_client.post( "/api/a2a/gpu/claim", json={"node": "linstation", "vram_mb": 4096} diff --git a/tinyagentos/routes/a2a_gpu_lease.py b/tinyagentos/routes/a2a_gpu_lease.py index 1b5cc1e3c..374ef2ddb 100644 --- a/tinyagentos/routes/a2a_gpu_lease.py +++ b/tinyagentos/routes/a2a_gpu_lease.py @@ -127,7 +127,14 @@ async def _resolve_actor( """Resolve the acting identity, or raise 401/403 (fail closed).""" if getattr(request.state, "is_admin", False): holder = _clean_handle(body_holder) or "@operator" - return _Actor(identity=holder, holder=holder, is_admin=True) + # A session admin acts as a FIXED verified principal; `holder` is the + # readable label only. Reading the body's holder as the admin's IDENTITY + # would let a `holder=` spelled like another holder's lease satisfy + # `_lease_owned_by` - i.e. take ownership of an `a2a:` lease the admin + # does not hold, on the node-scoped release/renew paths that take no + # lease id (CR on #2988). An operator still frees any lease by EXPLICIT + # id, which is the operator path. + return _Actor(identity="@operator", holder=holder, is_admin=True) caller = await check_agent_scope(request, scope) if caller is None: @@ -475,11 +482,14 @@ def _lease_for_actor(cluster, resource_id: str, actor: _Actor): def _lease_owned_by(lease, actor: _Actor) -> bool: """True when *lease* was taken by *actor*. - Strict caller match: the node-scoped release/renew paths must not let an - admin's session free a lease it did not take (that is what the explicit-id - paths and the cluster lease API are for). + Identity only - never the body's ``holder=``. That field is caller-supplied + display text, so matching on it would let any caller claim another holder's + ``a2a:`` lease as its own (CR on #2988). The node-scoped release/renew paths + must not let an admin's session free a lease it did not take either (that is + what the explicit-id paths and the cluster lease API are for), which the + fixed ``@operator`` principal enforces. """ - return lease.caller in (f"a2a:{actor.identity}", f"a2a:{actor.holder}") + return lease.caller == f"a2a:{actor.identity}" def _lease_bus_identity(lease) -> str | None: From f1cfcca46a3588a28c19ad51f337d93c253705ca Mon Sep 17 00:00:00 2001 From: hognek <227774406+hognek@users.noreply.github.com> Date: Sat, 12 Sep 2026 11:06:56 +0200 Subject: [PATCH 11/17] fix(cluster): read a renewal's previous expiry under the lease lock CR on #2988 (tinyagentos/cluster/manager.py:912). The keep-alive route captured `previous_expiry` from a lease object read OUTSIDE `_lease_lock` and only then called `renew_lease`. A renewal completing in between made that value stale, so if this request's bus post then failed, its rollback restored the stale expiry and clobbered the intervening renewal - even though the compare-and-set guard (4bbe24d7) matched this request's own attempted expiry. `renew_lease_with_previous` returns `(lease, previous_expiry)` from the same critical section that writes the new expiry, so the value a rollback restores is exactly the one this renewal replaced. `renew_lease` keeps its existing contract for every other caller (routes/cluster.py, the GPU arbiter) by delegating. The keep-alive route now uses the paired call. Tests (tests/test_routes_a2a_gpu_lease.py): - test_renewal_reports_the_expiry_it_actually_replaced pins the lock-held capture, and that a superseded renewal's rollback is refused while the owner's still applies. - test_a_failed_keep_alive_does_not_clobber_a_newer_renewal lands a second renewal while the bus post is in flight (new FakeBus.on_post hook) and asserts the newer expiry survives the failed request's rollback. --- tests/test_routes_a2a_gpu_lease.py | 83 +++++++++++++++++++++++++++++ tinyagentos/cluster/manager.py | 26 +++++++-- tinyagentos/routes/a2a_gpu_lease.py | 15 +++--- 3 files changed, 115 insertions(+), 9 deletions(-) diff --git a/tests/test_routes_a2a_gpu_lease.py b/tests/test_routes_a2a_gpu_lease.py index c4a929f93..d17076744 100644 --- a/tests/test_routes_a2a_gpu_lease.py +++ b/tests/test_routes_a2a_gpu_lease.py @@ -11,6 +11,7 @@ import asyncio import math import time +from collections.abc import Awaitable, Callable import pytest import pytest_asyncio @@ -34,6 +35,8 @@ def __init__(self) -> None: self.gets: list[dict] = [] self.fail_get = False self.fail_post = False + # Optional async hook run while a post is "in flight" (see _install_fake_bus). + self.on_post: Callable[[dict], Awaitable[None]] | None = None self._id = 0 def seed(self, body: str, sender: str = "@peer") -> dict: @@ -92,6 +95,10 @@ async def get(self, url, params=None, headers=None): return _Resp({"messages": list(bus.messages)}) async def post(self, url, json=None, headers=None): + if bus.on_post is not None: + # A hook for tests that need something to happen while a post is + # in flight (e.g. a concurrent renewal landing before a failure). + await bus.on_post(dict(json or {})) if bus.fail_post: raise RuntimeError("bus unreachable") payload = dict(json or {}) @@ -583,6 +590,82 @@ async def test_a_claim_that_stops_being_kept_alive_frees_the_node( assert freed.json()["blockers"] == [] assert took.status_code == 200 + async def test_renewal_reports_the_expiry_it_actually_replaced(self, cluster): + """`previous_expiry` must come from inside the locked renewal (CR #2988). + + Captured before the lock it can be stale: a renewal that completed in + between would then be rolled back by this one's failure, clobbering an + expiry that owns the lease. + """ + lease = await cluster.claim_lease( + "linstation:gpu-cuda-0", caller="a2a:@a", ttl_seconds=60 + ) + assert lease is not None + claimed_expiry = lease.expires_at + + first, previous = await cluster.renew_lease_with_previous( + lease.lease_id, ttl_seconds=120 + ) + assert previous == claimed_expiry + # The lease object is mutated in place, so take the attempted expiry now. + first_attempted = first.expires_at + second, replaced = await cluster.renew_lease_with_previous( + lease.lease_id, ttl_seconds=300 + ) + # The second renewal replaced the FIRST renewal's expiry, never the + # original: that is what makes the rollback below safe. + assert replaced == first_attempted + + # A rollback for the superseded renewal must not clobber the newer one. + assert ( + await cluster.restore_lease_expiry( + lease.lease_id, previous, attempted_expiry=first_attempted + ) + is False + ) + assert cluster.get_leases()[0].expires_at == second.expires_at + # ...while the renewal that owns the lease can still roll itself back. + assert ( + await cluster.restore_lease_expiry( + lease.lease_id, replaced, attempted_expiry=second.expires_at + ) + is True + ) + assert cluster.get_leases()[0].expires_at == replaced + + async def test_a_failed_keep_alive_does_not_clobber_a_newer_renewal( + self, lease_client, bus, cluster + ): + """A rollback restores only the expiry this request replaced (CR #2988). + + The bus post runs outside the manager's lock, so a concurrent renewal + can land between the extension and its rollback. That newer expiry owns + the lease and must survive. + """ + claimed = await lease_client.post( + "/api/a2a/gpu/claim", json={"node": "linstation", "vram_mb": 4096} + ) + lease_id = claimed.json()["lease_id"] + assert lease_id is not None + + async def _newer_renewal(_payload): + # A second /renew lands while this request's bus post is in flight. + await cluster.renew_lease(lease_id, ttl_seconds=900) + bus.fail_post = True + + bus.on_post = _newer_renewal + renewed = await lease_client.post( + "/api/a2a/gpu/renew", json={"lease_id": lease_id, "ttl_seconds": 600} + ) + assert renewed.status_code == 200 + assert renewed.json()["bus_claim_refreshed"] is False + assert renewed.json()["bus_refresh_error"] == ( + "a2a bus unavailable; a newer renewal stands" + ) + # The newer (900s) renewal still stands: the failed one's rollback did + # not restore the stale expiry. + assert cluster.get_leases()[0].expires_at > time.time() + 700 + async def test_renew_republishes_the_claim_keep_alive( self, lease_client, bus, cluster ): diff --git a/tinyagentos/cluster/manager.py b/tinyagentos/cluster/manager.py index 467b33bd8..c5dfebaa6 100644 --- a/tinyagentos/cluster/manager.py +++ b/tinyagentos/cluster/manager.py @@ -896,16 +896,36 @@ async def release_lease(self, lease_id: str) -> bool: async def renew_lease(self, lease_id: str, ttl_seconds: float = 30) -> GpuLease | None: """Extend a lease's TTL. Returns the lease, or None if expired/unknown.""" + lease, _previous_expiry = await self.renew_lease_with_previous( + lease_id, ttl_seconds=ttl_seconds + ) + return lease + + async def renew_lease_with_previous( + self, lease_id: str, ttl_seconds: float = 30 + ) -> tuple[GpuLease | None, float | None]: + """Extend a lease's TTL, returning ``(lease, previous_expiry)``. + + ``previous_expiry`` is the expiry this renewal actually REPLACED, read + under ``_lease_lock`` in the same critical section that writes the new + one. A caller that wants to undo a failed keep-alive needs exactly that + value: capturing the expiry before the lock (or from a lease object read + outside it) can be stale - another renewal may complete in between - and + restoring a stale expiry would clobber a newer renewal that owns the + lease (CR on #2988). ``None`` for both when the lease is unknown or + already expired, matching :meth:`renew_lease`'s contract. + """ async with self._lease_lock: lease = self._leases.get(lease_id) if lease is None: - return None + return None, None now = time.time() if lease.expires_at <= now: self._leases.pop(lease_id, None) - return None + return None, None + previous_expiry = lease.expires_at lease.expires_at = now + ttl_seconds - return lease + return lease, previous_expiry async def restore_lease_expiry( self, lease_id: str, expires_at: float, *, attempted_expiry: float diff --git a/tinyagentos/routes/a2a_gpu_lease.py b/tinyagentos/routes/a2a_gpu_lease.py index 374ef2ddb..2ea491b0f 100644 --- a/tinyagentos/routes/a2a_gpu_lease.py +++ b/tinyagentos/routes/a2a_gpu_lease.py @@ -838,12 +838,15 @@ async def gpu_renew(request: Request, body: RenewBody): ) if not _may_act_on(existing, actor): return JSONResponse({"error": "not the lease holder"}, status_code=403) - # The renewal's other half is the bus claim. Keep the instant before the - # extension: if the claim cannot be refreshed the local expiry is rolled - # back to it, so the two views agree rather than this controller holding a - # reservation every peer has already seen lapse. - previous_expiry = existing.expires_at - lease = await cluster.renew_lease(body.lease_id, ttl_seconds=float(body.ttl_seconds)) + # The renewal's other half is the bus claim. The manager returns the expiry + # this renewal actually REPLACED (read under its lock): if the claim cannot + # be refreshed the local expiry is rolled back to it, so the two views agree + # rather than this controller holding a reservation every peer has already + # seen lapse. Reading the expiry here instead would be stale as soon as a + # concurrent renewal lands (CR on #2988). + lease, previous_expiry = await cluster.renew_lease_with_previous( + body.lease_id, ttl_seconds=float(body.ttl_seconds) + ) if lease is None: return JSONResponse( {"error": "lease not found or expired", "lease_id": body.lease_id}, From 102fa59be14a281e2665e8e02e3b12499a5af880 Mon Sep 17 00:00:00 2001 From: hognek <227774406+hognek@users.noreply.github.com> Date: Sat, 12 Sep 2026 11:08:26 +0200 Subject: [PATCH 12/17] fix(a2a): count a re-claimed reservation once, not twice CR on #2988 (tinyagentos/gpu_lease.py:453). Admission subtracted the caller's OWN open claim from the node's live free VRAM. Once the model was loaded that claim is already reflected in free_mb, so a 12-GiB card with 6 GiB free and a same-holder 6-GiB claim computed zero available VRAM and the next claim/check was denied - which broke the idempotent re-POST the 500-message fold window requires. `evaluate_admission(replace_own=True)` treats the requirement as a REPLACEMENT of the caller's own reservation (its reservation comes back into the budget) rather than an addition to it. The default stays False, so CHECK - where an own claim really is a pending load - is unchanged. The claim route passes it only when the caller already owns the lease on that resource; the lease id / resource are resolved before admission to know that. Tests: - tests/test_gpu_lease_protocol.py: the finding's exact scenario is admitted under replace_own, and a replacement still cannot exceed the caller's own reservation plus the free figure. - tests/test_routes_a2a_gpu_lease.py: re-claiming the same node after its free VRAM has dropped to the loaded state renews the same lease instead of 409-ing. --- tests/test_gpu_lease_protocol.py | 34 ++++++++++++++++++++++++++++ tests/test_routes_a2a_gpu_lease.py | 23 +++++++++++++++++++ tinyagentos/gpu_lease.py | 12 +++++++++- tinyagentos/routes/a2a_gpu_lease.py | 35 +++++++++++++++++++++++------ 4 files changed, 96 insertions(+), 8 deletions(-) diff --git a/tests/test_gpu_lease_protocol.py b/tests/test_gpu_lease_protocol.py index d78c81d68..6eccd138e 100644 --- a/tests/test_gpu_lease_protocol.py +++ b/tests/test_gpu_lease_protocol.py @@ -378,6 +378,40 @@ def test_own_claim_matches_through_the_registry_canonical_id(self): assert d.claimed_mb == 6144 assert d.blockers == () + def test_a_repeated_claim_replaces_the_own_reservation_rather_than_stacking( + self, + ): + # A loaded model is already reflected in the live free figure, so + # subtracting the caller's own claim a second time would read a 12-GiB + # card with 6 GiB free as full and deny the idempotent re-claim + # (CR on #2988). + d = evaluate_admission( + node="n1", + required_mb=6144, + identity="@a", + claims=[self._claim("@a", 6144)], + free_mb=6144, + capacity_mb=12288, + replace_own=True, + ) + assert d.admitted is True + assert d.claimed_mb == 6144 + + def test_replacing_the_own_reservation_does_not_conjure_vram(self): + # The caller may use its own reservation plus what is actually free - + # no more. + d = evaluate_admission( + node="n1", + required_mb=16384, + identity="@a", + claims=[self._claim("@a", 6144)], + free_mb=6144, + capacity_mb=12288, + replace_own=True, + ) + assert d.admitted is False + assert "insufficient VRAM" in d.reason + def test_insufficient_vram_is_denied(self): d = evaluate_admission( node="n1", required_mb=8192, identity="@a", claims=[], free_mb=4096 diff --git a/tests/test_routes_a2a_gpu_lease.py b/tests/test_routes_a2a_gpu_lease.py index d17076744..02a4c2d15 100644 --- a/tests/test_routes_a2a_gpu_lease.py +++ b/tests/test_routes_a2a_gpu_lease.py @@ -888,6 +888,29 @@ async def test_node_spelling_resolves_to_one_lease(self, lease_client, bus, clus assert second.json()["lease_id"] == first.json()["lease_id"] assert len(cluster.get_leases()) == 1 + async def test_reclaiming_a_loaded_card_is_admitted(self, lease_client, bus, cluster): + """A re-claim replaces the caller's reservation; it never stacks on it. + + Once the model is loaded the node's free figure already reflects it, so + charging the caller's own claim again would read the card as full and + deny the idempotent re-POST the fold window needs (CR on #2988). + """ + first = await lease_client.post( + "/api/a2a/gpu/claim", json={"node": "linstation", "vram_mb": 6144} + ) + assert first.status_code == 200 + lease_id = first.json()["lease_id"] + assert lease_id is not None + + # The load is now on the card: 2 of the worker's 12 GiB remain free. + cluster.get_worker("linstation").free_vram_mb = 2048 + again = await lease_client.post( + "/api/a2a/gpu/claim", json={"node": "linstation", "vram_mb": 6144} + ) + assert again.status_code == 200, again.json() + assert again.json()["lease_id"] == lease_id + assert len(cluster.get_leases()) == 1 + async def test_reclaim_rollback_does_not_drop_the_existing_lease( self, lease_client, bus, cluster ): diff --git a/tinyagentos/gpu_lease.py b/tinyagentos/gpu_lease.py index 29d12dad8..35c3ad800 100644 --- a/tinyagentos/gpu_lease.py +++ b/tinyagentos/gpu_lease.py @@ -431,6 +431,7 @@ def evaluate_admission( claims: Iterable[GpuLeaseMessage] = (), free_mb: int | None = None, capacity_mb: int | None = None, + replace_own: bool = False, ) -> Admission: """Decide whether *identity* may use *node* for *required_mb* of VRAM. @@ -446,6 +447,13 @@ def evaluate_admission( ``verified=False`` with a reason: a node that cannot report VRAM is a gap in the guarantee, and the caller must probe before loading rather than be told a silent "free". + + ``replace_own=True`` says the request REPLACES the caller's own claim + rather than adding to it (an idempotent re-claim of the same lease). Its + own reservation then comes back into the budget instead of being charged + again: a loaded model can already be reflected in the live ``free_mb``, so + subtracting the claim a second time would read a reserved card as full and + deny the very re-POST the fold window needs (CR on #2988). """ required_mb = max(0, int(required_mb)) entries = list(claims) @@ -484,7 +492,9 @@ def evaluate_admission( ), ) - available = max(0, int(budget) - own_mb) + available = ( + int(budget) + own_mb if replace_own else max(0, int(budget) - own_mb) + ) if required_mb > available: return Admission( admitted=False, diff --git a/tinyagentos/routes/a2a_gpu_lease.py b/tinyagentos/routes/a2a_gpu_lease.py index 2ea491b0f..e0dc19dbe 100644 --- a/tinyagentos/routes/a2a_gpu_lease.py +++ b/tinyagentos/routes/a2a_gpu_lease.py @@ -383,6 +383,7 @@ async def _check_node( actor: _Actor, channel: str, resource: str = _DEFAULT_RESOURCE, + replace_own: bool = False, ) -> tuple[dict, list[GpuLeaseMessage]]: """Run the full CHECK for a node; returns (admission dict, node claims).""" folded = await _folded_claims(request, channel, actor) @@ -407,6 +408,7 @@ async def _check_node( claims=node_claims, free_mb=free_mb, capacity_mb=capacity_mb, + replace_own=replace_own, ) body = admission.as_dict() body["claims"] = [c.as_dict() for c in node_claims] @@ -619,6 +621,30 @@ async def gpu_claim(request: Request, body: ClaimBody): ) channel = (body.channel or "").strip() or _channel() + + # Productized half: a real lease the scheduler enforces. Only for a node + # this controller knows as a worker; an external node is bus-governed. The + # lease identity/resource are resolved BEFORE admission so a re-claim of + # the caller's OWN lease is admitted as a REPLACEMENT of it (see below). + cluster = getattr(request.app.state, "cluster_manager", None) + lease_node = cluster is not None and _match_worker(cluster, node) is not None + caller = f"a2a:{actor.identity}" + resource_id = ( + _resource_id(_canonical_node(cluster, node), body.resource) + if lease_node and cluster is not None + else None + ) + existing = ( + cluster.find_existing_lease(resource_id) + if cluster is not None and resource_id is not None + else None + ) + # A re-claim of our own lease replaces the reservation rather than adding + # a second one. Charging the caller's own claim again would read a card + # whose load is already reflected in the live free VRAM as full and deny + # the idempotent re-POST the fold window needs (CR on #2988). + reclaimer = existing is not None and existing.caller == caller + admission, _claims = await _check_node( request, node=node, @@ -626,20 +652,15 @@ async def gpu_claim(request: Request, body: ClaimBody): actor=actor, channel=channel, resource=body.resource, + replace_own=reclaimer, ) if not admission["admitted"]: return JSONResponse({"status": "denied", **admission}, status_code=409) - # Productized half: a real lease the scheduler enforces. Only for a node - # this controller knows as a worker; an external node is bus-governed. - cluster = getattr(request.app.state, "cluster_manager", None) lease_id: str | None = None lease = None created_lease = False - if cluster is not None and _match_worker(cluster, node) is not None: - resource_id = _resource_id(_canonical_node(cluster, node), body.resource) - caller = f"a2a:{actor.identity}" - existing = cluster.find_existing_lease(resource_id) + if cluster is not None and resource_id is not None: if existing is not None and existing.caller != caller: return JSONResponse( { From acd9f722093981e9fcfb58600df36dd2cad31256 Mon Sep 17 00:00:00 2001 From: hognek <227774406+hognek@users.noreply.github.com> Date: Sat, 12 Sep 2026 11:09:09 +0200 Subject: [PATCH 13/17] fix(a2a): keep a repeated claim to the lease's original contract CR on #2988 (tinyagentos/routes/a2a_gpu_lease.py:660). The idempotent re-claim path only extends `expires_at`, then published the NEW request's `vram_mb` and `channel`. A changed vram split the GpuLease from the bus claim, and a changed channel reposted the line on a thread the original claim was never on - leaving peers watching the first channel with no visible renewal while the local lease stayed held. A re-claim whose vram or channel differs from the held lease is now rejected (409) with the held contract in the reason; the way to change the shape of a reservation is to release it first. Test: test_reclaiming_with_different_parameters_is_rejected asserts both the larger-vram and the other-channel re-claims are refused and that the lease's required_vram_mb / claim_channel and the channel's messages are untouched. --- tests/test_routes_a2a_gpu_lease.py | 35 +++++++++++++++++++++++++++++ tinyagentos/routes/a2a_gpu_lease.py | 23 ++++++++++++++++++- 2 files changed, 57 insertions(+), 1 deletion(-) diff --git a/tests/test_routes_a2a_gpu_lease.py b/tests/test_routes_a2a_gpu_lease.py index 02a4c2d15..95a75d77d 100644 --- a/tests/test_routes_a2a_gpu_lease.py +++ b/tests/test_routes_a2a_gpu_lease.py @@ -911,6 +911,41 @@ async def test_reclaiming_a_loaded_card_is_admitted(self, lease_client, bus, clu assert again.json()["lease_id"] == lease_id assert len(cluster.get_leases()) == 1 + async def test_reclaiming_with_different_parameters_is_rejected( + self, lease_client, bus, cluster + ): + """A re-claim keeps the contract the lease was taken with (CR on #2988). + + `renew_lease` only moves `expires_at`, so republishing a different vram + figure or onto a different channel would leave the GpuLease and the bus + claim disagreeing - and peers on the original channel would never see + the renewal. + """ + first = await lease_client.post( + "/api/a2a/gpu/claim", json={"node": "linstation", "vram_mb": 4096} + ) + assert first.status_code == 200 + lease_id = first.json()["lease_id"] + sends = len(bus.sends) + + bigger = await lease_client.post( + "/api/a2a/gpu/claim", json={"node": "linstation", "vram_mb": 6144} + ) + assert bigger.status_code == 409 + assert "release it before re-claiming" in bigger.json()["reason"] + + elsewhere = await lease_client.post( + "/api/a2a/gpu/claim", + json={"node": "linstation", "vram_mb": 4096, "channel": "gpu-lab"}, + ) + assert elsewhere.status_code == 409 + + # Nothing moved: same lease, same contract, no line on the other thread. + assert [lease.lease_id for lease in cluster.get_leases()] == [lease_id] + assert cluster.get_leases()[0].required_vram_mb == 4096 + assert cluster.get_leases()[0].claim_channel == "gpu" + assert len(bus.sends) == sends + async def test_reclaim_rollback_does_not_drop_the_existing_lease( self, lease_client, bus, cluster ): diff --git a/tinyagentos/routes/a2a_gpu_lease.py b/tinyagentos/routes/a2a_gpu_lease.py index e0dc19dbe..f5993727e 100644 --- a/tinyagentos/routes/a2a_gpu_lease.py +++ b/tinyagentos/routes/a2a_gpu_lease.py @@ -674,7 +674,28 @@ async def gpu_claim(request: Request, body: ClaimBody): if existing is not None: # Idempotent re-claim: extend our own lease rather than failing on # the "already leased" guard (which the manager applies to every - # caller, including the current holder). + # caller, including the current holder). The lease keeps the + # contract it was taken with: `renew_lease` only moves `expires_at`, + # so publishing a different vram figure - or onto a different + # channel - would leave the GpuLease and the bus claim disagreeing, + # and the original channel's claim unaware of the renewal. Releasing + # first is how a reservation changes shape (CR on #2988). + held_channel = getattr(existing, "claim_channel", "") or "" + if (existing.required_vram_mb or 0) != vram_mb or held_channel != channel: + return JSONResponse( + { + "status": "denied", + **admission, + "reason": ( + f"{resource_id} is already leased by {caller} with " + f"vram_mb={existing.required_vram_mb or 0} on channel " + f"{held_channel or _DEFAULT_CHANNEL}; release it " + "before re-claiming with different parameters" + ), + "blockers": [existing.caller], + }, + status_code=409, + ) lease = await cluster.renew_lease( existing.lease_id, ttl_seconds=float(body.ttl_seconds) ) From 8d59a6e76ab269dcaac29c3cfda9adb1af8d43dc Mon Sep 17 00:00:00 2001 From: hognek <227774406+hognek@users.noreply.github.com> Date: Sat, 12 Sep 2026 11:10:15 +0200 Subject: [PATCH 14/17] fix(a2a): undo a re-claim's extension when its repost fails CR on #2988 (tinyagentos/routes/a2a_gpu_lease.py:706). The route renewed the caller's own lease BEFORE posting the claim line, but the HTTPException rollback only handled leases this call CREATED. A failed repost on the idempotent re-claim path therefore left the local lease extended while the bus kept the claim's OLD expiry - after that older instant peers free the card and can co-load while this controller still believes it is reserved. That is the same half-renewal `gpu_renew` already guards against. The re-claim now takes its previous expiry from the locked `renew_lease_with_previous` and, on a failed post, restores it with the same compare-and-set rollback, so an intervening renewal still owns the lease. A lease this call created is still released, as before. Test: test_reclaim_rollback_restores_the_extended_expiry re-claims with a longer TTL while the bus is down and asserts the expiry is back to the value the bus claim still carries. --- tests/test_routes_a2a_gpu_lease.py | 26 +++++++++++++++++++++ tinyagentos/routes/a2a_gpu_lease.py | 36 ++++++++++++++++++++++++----- 2 files changed, 56 insertions(+), 6 deletions(-) diff --git a/tests/test_routes_a2a_gpu_lease.py b/tests/test_routes_a2a_gpu_lease.py index 95a75d77d..efaf607a3 100644 --- a/tests/test_routes_a2a_gpu_lease.py +++ b/tests/test_routes_a2a_gpu_lease.py @@ -962,6 +962,32 @@ async def test_reclaim_rollback_does_not_drop_the_existing_lease( assert retry.status_code == 502 assert [lease.lease_id for lease in cluster.get_leases()] == [lease_id] + async def test_reclaim_rollback_restores_the_extended_expiry( + self, lease_client, bus, cluster + ): + """A failed re-claim repost must undo its own extension (CR on #2988). + + The lease is renewed before the bus post; if the line never lands, the + bus claim keeps its OLD expiry - so leaving the extension standing + would let peers free the card at that older instant while this + controller still held it. + """ + first = await lease_client.post( + "/api/a2a/gpu/claim", json={"node": "linstation", "vram_mb": 4096} + ) + lease_id = first.json()["lease_id"] + before = cluster.get_leases()[0].expires_at + + bus.fail_post = True + retry = await lease_client.post( + "/api/a2a/gpu/claim", + json={"node": "linstation", "vram_mb": 4096, "ttl_seconds": 900}, + ) + assert retry.status_code == 502 + # Rolled back to the expiry the claim on the bus still carries. + assert [lease.lease_id for lease in cluster.get_leases()] == [lease_id] + assert cluster.get_leases()[0].expires_at == before + async def test_check_blocks_when_the_scheduler_holds_a_lease( self, lease_client, bus, cluster ): diff --git a/tinyagentos/routes/a2a_gpu_lease.py b/tinyagentos/routes/a2a_gpu_lease.py index f5993727e..d80886040 100644 --- a/tinyagentos/routes/a2a_gpu_lease.py +++ b/tinyagentos/routes/a2a_gpu_lease.py @@ -660,6 +660,7 @@ async def gpu_claim(request: Request, body: ClaimBody): lease_id: str | None = None lease = None created_lease = False + previous_expiry: float | None = None if cluster is not None and resource_id is not None: if existing is not None and existing.caller != caller: return JSONResponse( @@ -696,7 +697,10 @@ async def gpu_claim(request: Request, body: ClaimBody): }, status_code=409, ) - lease = await cluster.renew_lease( + # The expiry this renewal replaces, read under the manager's lock: + # a failed repost must undo exactly this extension (and a renewal + # that landed meanwhile owns the lease - see the rollback below). + lease, previous_expiry = await cluster.renew_lease_with_previous( existing.lease_id, ttl_seconds=float(body.ttl_seconds) ) lease_id = existing.lease_id if lease is not None else None @@ -741,11 +745,31 @@ async def gpu_claim(request: Request, body: ClaimBody): try: posted = await _post_line(channel, actor, line) except HTTPException: - # Roll back only a lease THIS call created. A re-claim renews the - # caller's own pre-existing lease, and freeing that on a transient bus - # failure would drop a reservation the holder still believes it owns. - if created_lease and lease_id is not None and cluster is not None: - await cluster.release_lease(lease_id) + # Undo the LOCAL half of the lease this call touched. A call that + # created the lease frees it; one that re-claimed only extended the + # caller's own lease, and freeing that would drop a reservation the + # holder still believes it owns - but leaving the extension standing + # would be just as wrong the other way: the bus claim keeps its OLD + # expiry, so after that older instant peers free the card while this + # controller still holds it. Roll the extension back to the expiry it + # replaced, compare-and-set so a concurrent renewal is never clobbered + # (CR on #2988). + if cluster is not None and lease is not None: + if created_lease and lease_id is not None: + await cluster.release_lease(lease_id) + elif previous_expiry is not None: + restored = await cluster.restore_lease_expiry( + lease.lease_id, + previous_expiry, + attempted_expiry=lease.expires_at, + ) + logger.warning( + "A2A GPU lease %s re-claim rolled back (restored=%s): " + "claim repost failed (channel=%s)", + lease.lease_id, + restored, + channel, + ) raise return { From b4615b8ebfb1e823c58fff16ba2e4fd39063984b Mon Sep 17 00:00:00 2001 From: hognek <227774406+hognek@users.noreply.github.com> Date: Sat, 12 Sep 2026 11:11:20 +0200 Subject: [PATCH 15/17] fix(a2a): release the identified lease's node and channel, not the request's CR on #2988 (tinyagentos/routes/a2a_gpu_lease.py:735). `gpu_release` found a lease by explicit id but still rendered the RELEASE from the request's `node` and `channel`. A caller - including an operator - could therefore post `[GPU RELEASE] node=` on a thread the claim was never on, and then delete the identified local lease: the real claim stays open on the bus while the reservation is gone, i.e. peers keep reading the node as claimed. A found lease now supplies both: the node from its `resource_id` and the channel from `claim_channel`. Request values still apply to the idempotent bus-only release (no matching local lease), which is what tells peers a free node is free. Test: test_release_uses_the_leases_own_node_and_channel claims on node=linstation/channel=gpu-lab, releases it naming node=local and another channel, and asserts the lease's node/channel are what reach the bus. --- tests/test_routes_a2a_gpu_lease.py | 29 +++++++++++++++++++++++++++++ tinyagentos/routes/a2a_gpu_lease.py | 10 ++++++++++ 2 files changed, 39 insertions(+) diff --git a/tests/test_routes_a2a_gpu_lease.py b/tests/test_routes_a2a_gpu_lease.py index efaf607a3..611941182 100644 --- a/tests/test_routes_a2a_gpu_lease.py +++ b/tests/test_routes_a2a_gpu_lease.py @@ -519,6 +519,35 @@ async def test_operator_release_closes_the_holders_bus_claim( assert checked.json()["admitted"] is True assert checked.json()["blockers"] == [] + async def test_release_uses_the_leases_own_node_and_channel( + self, lease_client, bus, cluster + ): + """The release line belongs to the lease, not to the request (CR #2988). + + A request may name any node and channel; posting the RELEASE from those + would announce a different resource - on a thread the claim was never + on - and then delete the identified local lease anyway. + """ + claimed = await lease_client.post( + "/api/a2a/gpu/claim", + json={"node": "linstation", "vram_mb": 4096, "channel": "gpu-lab"}, + ) + assert claimed.status_code == 200 + lease_id = claimed.json()["lease_id"] + assert lease_id is not None + + released = await lease_client.post( + "/api/a2a/gpu/release", + json={"node": "local", "lease_id": lease_id, "channel": "gpu-other"}, + ) + assert released.status_code == 200 + # The lease's node/channel win over the request's. + assert released.json()["node"] == "linstation" + assert released.json()["channel"] == "gpu-lab" + assert bus.sends[-1]["payload"]["thread"] == "gpu-lab" + assert bus.last_line == "[GPU RELEASE] node=linstation holder=@operator" + assert cluster.get_leases() == [] + async def test_claim_publishes_a_bus_expiry_from_its_lease( self, lease_client, bus, cluster ): diff --git a/tinyagentos/routes/a2a_gpu_lease.py b/tinyagentos/routes/a2a_gpu_lease.py index d80886040..cc838f30a 100644 --- a/tinyagentos/routes/a2a_gpu_lease.py +++ b/tinyagentos/routes/a2a_gpu_lease.py @@ -824,6 +824,16 @@ async def gpu_release(request: Request, body: ReleaseBody): ) released_id = lease.lease_id if lease is not None else None + # A lease found by id OWNS the release's node and channel: the request's + # `node`/`channel` are free text, so honouring them would post a RELEASE + # for a different resource - on a thread the claim was never on - and only + # then delete the identified local lease, leaving the real claim open + # (CR on #2988). Request values stand for a bus-only release with no + # matching lease, which is the idempotent "the node is free" post. + if lease is not None: + node = (lease.resource_id or "").partition(":")[0] or node + channel = getattr(lease, "claim_channel", "") or channel + # Whose claim the [GPU RELEASE] closes. An operator freeing another holder's # lease by explicit id must attribute the line to THAT holder: a claim is # keyed on its bus AUTHOR, so a line posted as @operator clears nothing and From 35358e63a6c9c52d3e2dac9d08ce31a6bfb79a83 Mon Sep 17 00:00:00 2001 From: hognek <227774406+hognek@users.noreply.github.com> Date: Sat, 12 Sep 2026 11:11:42 +0200 Subject: [PATCH 16/17] test(a2a): make the "different caller is blocked" block test a different caller CR on #2988 (tests/test_routes_a2a_gpu_lease.py:351). The block claimed to show a different caller seeing a claimed node, but it reused `lease_client` - the same admin identity that had just made the claim - so the only assertion was `status_code == 200` and the block proved nothing about blocking. It now checks from a second, agent-token identity: the peer is admitted `False` and is blocked by `@operator`, the claim's bus author. --- tests/test_routes_a2a_gpu_lease.py | 16 +++++++++++++--- 1 file changed, 13 insertions(+), 3 deletions(-) diff --git a/tests/test_routes_a2a_gpu_lease.py b/tests/test_routes_a2a_gpu_lease.py index 611941182..f948a03df 100644 --- a/tests/test_routes_a2a_gpu_lease.py +++ b/tests/test_routes_a2a_gpu_lease.py @@ -351,11 +351,21 @@ async def test_claim_then_check_sees_the_claim(self, lease_client, bus, local_vr "/api/a2a/gpu/check", params={"node": "local", "vram_mb": 2048} ) assert mine.json()["admitted"] is True - # A different caller sees a claimed node. - peer = await lease_client.get( - "/api/a2a/gpu/check", params={"node": "local", "vram_mb": 2048, "channel": "gpu"} + # A DIFFERENT caller sees a claimed node. The admin's own identity is + # not a second caller: the block is only meaningful from another bus + # identity (CR on #2988). + _cid, token = await _agent_token( + lease_client._app, scopes=("a2a_receive",), handle="@taos" ) + async with _bare(lease_client._app) as bare: + peer = await bare.get( + "/api/a2a/gpu/check", + params={"node": "local", "vram_mb": 2048}, + headers={"Authorization": f"Bearer {token}"}, + ) assert peer.status_code == 200 + assert peer.json()["admitted"] is False + assert peer.json()["blockers"] == ["@operator"] async def test_peer_owned_claim_is_denied_409(self, lease_client, bus, local_vram): local_vram(11000) From 4f00a2892cf5a642430de14ffac8da9f7953a95e Mon Sep 17 00:00:00 2001 From: hognek <227774406+hognek@users.noreply.github.com> Date: Sat, 12 Sep 2026 11:12:03 +0200 Subject: [PATCH 17/17] docs(a2a): state both lifetime limits of a hand-posted claim CR on #2988 (docs/agent-coordination.md:441, changelog.d line 30). "Never expires (bounded by a RELEASE alone)" was wrong in one direction: the protocol also limits a claim's visibility to the newest 500 messages, so a hand-posted claim without `expires=` can still age out of the fold and has to be reposted periodically. Say that, instead of implying it outlives the channel. --- changelog.d/taos-893-a2a-gpu-lease.md | 5 ++++- docs/agent-coordination.md | 6 ++++-- 2 files changed, 8 insertions(+), 3 deletions(-) diff --git a/changelog.d/taos-893-a2a-gpu-lease.md b/changelog.d/taos-893-a2a-gpu-lease.md index 8e9aca501..b168f5ca7 100644 --- a/changelog.d/taos-893-a2a-gpu-lease.md +++ b/changelog.d/taos-893-a2a-gpu-lease.md @@ -27,7 +27,10 @@ that peers have already seen lapse. The fold drops a claim whose published expiry has passed, so a holder that crashed or stopped keeping alive no longer blocks the shared card until its claim ages out of the fold window. A claim - posted by hand without an `expires=` is unchanged: bounded by a RELEASE alone. + posted by hand without an `expires=` is unchanged: it has no time-based + expiry, so only a RELEASE closes it - though it is still limited by the fold + window (the newest 500 messages), so a long-lived one is reposted + periodically. ### Fixed diff --git a/docs/agent-coordination.md b/docs/agent-coordination.md index 7f556d517..e3390b65a 100644 --- a/docs/agent-coordination.md +++ b/docs/agent-coordination.md @@ -438,8 +438,10 @@ Rules that matter when you use it: that crashed or stopped keeping alive no longer blocks the card until its claim ages out of the fold window. Keep-alive therefore means re-POST `/claim` or `/renew` while you hold the card; a claim posted by hand without an - `expires=` never expires (it is bounded by a RELEASE alone), which is what the - interim protocol in #893 relies on. + `expires=` has no TIME-based expiry — it is closed by a RELEASE — but it is + still subject to the fold window below, so a long-lived hand-posted claim has + to be reposted periodically like any other. That is what the interim protocol + in #893 relies on. - **A claim is only visible inside the channel fold window** (the newest 500 messages). For a load that outlives the chatter around it, re-POST `/claim` periodically: it is idempotent (it extends the lease and reposts the line,