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/changelog.d/taos-893-a2a-gpu-lease.md b/changelog.d/taos-893-a2a-gpu-lease.md new file mode 100644 index 000000000..b168f5ca7 --- /dev/null +++ b/changelog.d/taos-893-a2a-gpu-lease.md @@ -0,0 +1,43 @@ +### 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. 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 + 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), 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: 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 + +- 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 e2ada7bf2..e3390b65a 100644 --- a/docs/agent-coordination.md +++ b/docs/agent-coordination.md @@ -351,6 +351,112 @@ 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: + +```text +[GPU CLAIM] node= holder=@you vram=~9.4gb reason=... eta=... expires= +[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. + +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 + 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". 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. +- **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`). 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 + 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. + 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 — 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` + or `/renew` while you hold the card; a claim posted by hand without an + `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, + 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 +1559,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_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/tests/test_gpu_lease_protocol.py b/tests/test_gpu_lease_protocol.py new file mode 100644 index 000000000..6eccd138e --- /dev/null +++ b/tests/test_gpu_lease_protocol.py @@ -0,0 +1,487 @@ +"""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 + +import time + +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_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" + + 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 + + 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"): + 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_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"}, + {"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") == [] + + 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): + 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_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", + 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_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 + ) + 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_claim_publishes_an_integer_expiry(self): + line = render_claim("n1", "@a", 6144, expires_at=1783350000.9) + # 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) + + 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..f948a03df --- /dev/null +++ b/tests/test_routes_a2a_gpu_lease.py @@ -0,0 +1,1237 @@ +"""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 asyncio +import math +import time +from collections.abc import Awaitable, Callable + +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 +from tinyagentos.gpu_lease import claims_for_node, open_claims + +_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.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: + """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, 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): + 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 {}) + 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_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( + "/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" + # 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 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): + 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. 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) + 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 + ): + # 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 + ) + 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 [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 + ): + 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 [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 + ): + # 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_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_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 + ): + """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 + # 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. + folded = open_claims(bus.messages) + 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 + ): + """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 + + # 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( + "/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_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 + ): + """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 + 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 [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_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} + ) + 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": 60}, + ) + assert resp.status_code == 200 + # 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_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_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 + ): + 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_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 + ): + # 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_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: + resp = await bare.get( + "/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 [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( + 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": 600}, + 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/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/cluster/manager.py b/tinyagentos/cluster/manager.py index 9e9ee7351..c5dfebaa6 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( @@ -894,16 +896,63 @@ 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 + ) -> 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. + + 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 def get_leases(self) -> list[GpuLease]: """Return a snapshot of active (non-expired) leases.""" 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 new file mode 100644 index 000000000..35c3ad800 --- /dev/null +++ b/tinyagentos/gpu_lease.py @@ -0,0 +1,518 @@ +# 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 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 + +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. + +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 +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 math +import re +import time +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", "expires_at": "expires"} + + +def _clean(text: object) -> str: + """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), 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( + " " 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. + + 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 = _FIELD_ALIASES.get(m.group("key").lower(), m.group("key").lower()) + fields.setdefault(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 + 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: + """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_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, + *, + 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, + expires_at=parse_epoch(fields.get("expires")), + ) + + +def render_claim( + 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 = [ + 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)}") + 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. 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)}" + + +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], *, 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 + 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). + + 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 + 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. + """ + 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: + msg = parse_message(raw) + if msg is None or msg.kind not in (CLAIM, RELEASE): + continue + key = (msg.node.strip().casefold(), 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 None or msg.expired(cutoff): + continue + out.setdefault(key[0], []).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() + 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) + return [] + + +def _is_mine(claim: GpuLeaseMessage, identity: str) -> bool: + """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) +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, + replace_own: bool = False, +) -> 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". + + ``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) + 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 = ( + int(budget) + own_mb if replace_own else 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..cc838f30a --- /dev/null +++ b/tinyagentos/routes/a2a_gpu_lease.py @@ -0,0 +1,986 @@ +# 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 +import time +from dataclasses import dataclass + +import httpx +from fastapi import APIRouter, HTTPException, Request +from fastapi.responses import JSONResponse +from pydantic import BaseModel, Field + +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 + +# 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 + + +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" + # 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: + 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, 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}, + headers=headers or None, + ) + 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, actor: _Actor +) -> dict[str, list[GpuLeaseMessage]]: + """Read the channel and fold it, mapping an unreadable bus to 503.""" + try: + 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( + 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 (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 + 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(): + 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] + 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(_canonical_node(cluster, 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, + 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) + 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=actor.identity, + 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] + 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 = Field( + default=_DEFAULT_TTL_SECONDS, gt=0, le=MAX_LEASE_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 = Field( + default=_DEFAULT_TTL_SECONDS, gt=0, le=MAX_LEASE_TTL_SECONDS + ) + + +def _resource_id(node: str, resource: str) -> str: + res = (resource or _DEFAULT_RESOURCE).strip() or _DEFAULT_RESOURCE + 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) + if existing is None: + return None + 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*. + + 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 == f"a2a:{actor.identity}" + + +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}" + + +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 + + +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 ───────────────────────────────────────────────────────────────── + + +@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) + # 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( + request, + node=node, + required_mb=required_mb or 0, + actor=actor, + 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() + + # 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, + required_mb=vram_mb, + actor=actor, + channel=channel, + resource=body.resource, + replace_own=reclaimer, + ) + if not admission["admitted"]: + return JSONResponse({"status": "denied", **admission}, status_code=409) + + 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( + { + "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). 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, + ) + # 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 + else: + lease = await cluster.claim_lease( + resource_id=resource_id, + caller=caller, + ttl_seconds=float(body.ttl_seconds), + required_vram_mb=vram_mb, + claim_channel=channel, + ) + 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 + created_lease = True + + # 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: + # 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 { + "status": "claimed", + "node": node, + "holder": actor.holder, + "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, + "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 + lease = None + if cluster is not 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(_canonical_node(cluster, node), body.resource), + actor, + ) + 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 + # every peer's fold keeps reading the node as claimed while the local lease + # 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 + # 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, 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, + "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) + # 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) + # 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}, + 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. + 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 + # 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, + 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: + # 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" + 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 keep-alive rolled back (restored=%s): " + "claim repost failed (channel=%s)", + lease.lease_id, + restored, + 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, + }