diff --git a/README.md b/README.md index 2d2befd3..3501b37f 100644 --- a/README.md +++ b/README.md @@ -61,6 +61,7 @@ Docs are grouped by audience. **Miners** - [Miner guide](docs/miner/README.md) — choose a challenge, submit through the proxy, and track leaderboards. +- [Miner worker deployment guide](docs/miner/worker-plane.md) — deploy a miner-funded GPU worker on Lium or Targon with `base worker deploy`, monitor it with `base worker status`, and understand the PRISM admission rule and proof tiers. **Validators / operators** @@ -214,6 +215,56 @@ capabilities, subscriptions, last heartbeat, resolved identity); it NEVER expose --- +## Miner-Funded GPU Worker Plane + +The **worker plane** moves PRISM's heavy GPU evaluation off BASE validators and onto **worker +agents running in GPU instances the miners themselves fund** (rented on Lium or Targon, or their +own hardware). Validators keep only light plausibility checks, probabilistic replay audits, and +on-chain weight submission. PRISM scoring is unchanged. Everything is gated behind +`compute.worker_plane_enabled` (env `BASE_COMPUTE__WORKER_PLANE_ENABLED`); with the flag **off** +(the default) gpu units route to validators exactly as today and the worker coordination surface is +not even mounted (byte-for-byte legacy behavior). + +- **`base worker` CLI** (top-level, distinct from the legacy `base master worker` Swarm group): + - `base worker deploy --provider lium|targon|local [--max-price ]` — provision a + worker on a miner-funded instance (or locally) and enroll it. The miner's provider key + (`LIUM_API_KEY` / `TARGON_API_KEY`) authenticates provider calls only and is **never** sent to + the master. Deploy always sets a bounded instance lifetime and refuses offers above the price + cap. + - `base worker agent` — run the long-lived agent loop (register → heartbeat → pull → execute → + post), authenticated as the worker keypair. + - `base worker status [--hotkey ]` — render the fleet from the master's + `GET /v1/workers` (status, owner, provider, last-seen, fault count). +- **Signed enrollment**: the miner signs `worker-binding:{worker_pubkey}:{miner_hotkey}:{nonce}` + (sr25519); the master verifies it against the metagraph. Status lifecycle is + `pending → active → stale → retired`; `active` requires a verified binding **and** a heartbeat + within `compute.worker_heartbeat_ttl_seconds` (default 120). +- **Admission rule**: when enforced, a miner must have ≥1 **active** worker bound to their hotkey + to submit to PRISM; otherwise the submission is rejected `403 NO_ACTIVE_WORKER`. This is the + incentive to supply GPU without changing the subnet economics. +- **Anti-collusion assignment**: a worker never evaluates its own owner's submission; each gpu unit + is replicated across **R=2 distinct-owner** workers (`compute.replication_factor`, degrading to + 1 with a recorded warning). The master reconciles by comparing each replica's + `ExecutionProof.manifest_sha256`: equal ⇒ one result forwarded to the challenge; divergent ⇒ the + unit is `disputed` and a validator **audit** unit replays it authoritatively, flagging the lying + worker via `worker_faults`. +- **ExecutionProof tiers** (carried on every result): + - **Tier 0** (mandatory, all backends): deterministic manifest hash + worker sr25519 signature — + the source of truth for reconciliation and audits. + - **Tier 1**: evaluator image digest matches the pinned digest (on Lium, cross-checkable against + the provider's signed image-digest endpoint). + - **Tier 2**: in-guest TDX + nvtrust attestation — schema shipped but gated **off on Targon** + today (no consumer-facing attestation surface), so Targon proofs carry tier ≤ 1. + - Audit sampling is tier-modulated (tier 0 ≈ 10%, tier 1 ≈ 5%, tier 2 ≈ 2%). + +Miners: see the [Miner worker deployment guide](docs/miner/worker-plane.md) for prerequisites, +exact Lium/Targon commands, cost guidance, monitoring, and troubleshooting. The master subsystems +live under `src/base/worker/` (agent) and `src/base/master/worker_coordination.py` (registry, +enrollment, fleet, admission), with provider clients in `src/base/compute/` +(`lium.py` / `targon.py`). + +--- + ## What BASE Does BASE coordinates the full lifecycle of a multi-challenge subnet: @@ -252,6 +303,18 @@ master's final normalized vector from the weights API and commits it under its o dedicated submit-only host can run the same submitter). The master aggregates but never submits on-chain. Challenge API services are run by the master (manager) node. +### Workers (miner-funded GPU worker plane) + +Workers are **miner-funded GPU executors** for PRISM. A miner runs `base worker deploy` to launch a +worker agent inside a GPU instance the miner pays for (rented on **Lium** or **Targon**, or their +own hardware); the agent registers with the master under a miner-signed hotkey↔worker binding, +heartbeats, pulls `gpu` work units, executes the PRISM evaluator on its own broker, and posts +results carrying an `ExecutionProof`. The provider API key (`LIUM_API_KEY` / `TARGON_API_KEY`) +stays in the miner's environment and is never sent to the master. The entire worker plane is gated +behind `compute.worker_plane_enabled` (env `BASE_COMPUTE__WORKER_PLANE_ENABLED`, default off ⇒ +byte-for-byte legacy behavior). See the [Worker plane](#miner-funded-gpu-worker-plane) section and +the [Miner worker deployment guide](docs/miner/worker-plane.md). + --- ## Repository Layout diff --git a/alembic/versions/0005_create_work_assignments.py b/alembic/versions/0005_create_work_assignments.py index 424d121b..a9f271d9 100644 --- a/alembic/versions/0005_create_work_assignments.py +++ b/alembic/versions/0005_create_work_assignments.py @@ -24,6 +24,7 @@ "running", "completed", "failed", + "disputed", name="work_assignment_status", native_enum=False, ) diff --git a/alembic/versions/0009_create_worker_registry.py b/alembic/versions/0009_create_worker_registry.py new file mode 100644 index 00000000..7e4efb79 --- /dev/null +++ b/alembic/versions/0009_create_worker_registry.py @@ -0,0 +1,185 @@ +"""Create miner-funded GPU worker registry tables. + +Adds the worker-plane storage (architecture.md sec 3.3): ``worker_registrations`` +(the enrolled workers + their ``pending`` -> ``active`` -> ``stale`` -> ``retired`` +lifecycle), ``worker_faults`` (reconciliation/audit fault attribution surfaced in +the fleet view), and ``worker_request_nonces`` (replay protection for the miner +binding nonce and the signed request-envelope nonce). No pre-existing table is +altered. + +Revision ID: 0009_create_worker_registry +Revises: 0008_validator_subscriptions +Create Date: 2026-07-06 00:00:00.000000 +""" + +from __future__ import annotations + +from collections.abc import Sequence + +import sqlalchemy as sa + +from alembic import op + +revision: str = "0009_create_worker_registry" +down_revision: str | None = "0008_validator_subscriptions" +branch_labels: str | Sequence[str] | None = None +depends_on: str | Sequence[str] | None = None + +worker_status = sa.Enum( + "pending", + "active", + "stale", + "retired", + name="worker_status", + native_enum=False, +) + + +def upgrade() -> None: + """Apply the migration.""" + + op.create_table( + "worker_registrations", + sa.Column("id", sa.Uuid(as_uuid=True), nullable=False), + sa.Column("worker_id", sa.Text(), nullable=False), + sa.Column("worker_pubkey", sa.Text(), nullable=False), + sa.Column("miner_hotkey", sa.Text(), nullable=False), + sa.Column("binding_signature", sa.Text(), nullable=False), + sa.Column("binding_nonce", sa.Text(), nullable=False), + sa.Column("provider", sa.Text(), nullable=False), + sa.Column("provider_instance_ref", sa.Text(), nullable=True), + sa.Column("capabilities", sa.JSON(), server_default="[]", nullable=False), + sa.Column("status", worker_status, server_default="pending", nullable=False), + sa.Column("last_seen_meta", sa.JSON(), server_default="{}", nullable=False), + sa.Column("last_heartbeat_at", sa.DateTime(timezone=True), nullable=True), + sa.Column( + "created_at", + sa.DateTime(timezone=True), + server_default=sa.func.now(), + nullable=False, + ), + sa.Column( + "updated_at", + sa.DateTime(timezone=True), + server_default=sa.func.now(), + nullable=False, + ), + sa.PrimaryKeyConstraint("id", name=op.f("pk_worker_registrations")), + sa.UniqueConstraint("worker_id", name="uq_worker_registrations_worker_id"), + sa.UniqueConstraint( + "worker_pubkey", name="uq_worker_registrations_worker_pubkey" + ), + ) + op.create_index( + "ix_worker_registrations_status", + "worker_registrations", + ["status"], + unique=False, + ) + op.create_index( + "ix_worker_registrations_miner_hotkey", + "worker_registrations", + ["miner_hotkey"], + unique=False, + ) + op.create_index( + "ix_worker_registrations_last_heartbeat_at", + "worker_registrations", + ["last_heartbeat_at"], + unique=False, + ) + op.create_index( + "ix_worker_registrations_status_miner_hotkey", + "worker_registrations", + ["status", "miner_hotkey"], + unique=False, + ) + + op.create_table( + "worker_faults", + sa.Column("id", sa.Uuid(as_uuid=True), nullable=False), + sa.Column("worker_id", sa.Text(), nullable=False), + sa.Column("work_unit_id", sa.Text(), nullable=False), + sa.Column("challenge_slug", sa.Text(), nullable=True), + sa.Column("detail", sa.Text(), nullable=True), + sa.Column( + "created_at", + sa.DateTime(timezone=True), + server_default=sa.func.now(), + nullable=False, + ), + sa.PrimaryKeyConstraint("id", name=op.f("pk_worker_faults")), + ) + op.create_index( + "ix_worker_faults_worker_id", + "worker_faults", + ["worker_id"], + unique=False, + ) + op.create_index( + "ix_worker_faults_work_unit_id", + "worker_faults", + ["work_unit_id"], + unique=False, + ) + + op.create_table( + "worker_request_nonces", + sa.Column("id", sa.Uuid(as_uuid=True), nullable=False), + sa.Column("hotkey", sa.Text(), nullable=False), + sa.Column("nonce", sa.Text(), nullable=False), + sa.Column("body_hash", sa.Text(), nullable=False), + sa.Column("created_at", sa.DateTime(timezone=True), nullable=False), + sa.PrimaryKeyConstraint("id", name=op.f("pk_worker_request_nonces")), + sa.UniqueConstraint( + "hotkey", "nonce", name="uq_worker_request_nonces_hotkey_nonce" + ), + ) + op.create_index( + "ix_worker_request_nonces_created_at", + "worker_request_nonces", + ["created_at"], + unique=False, + ) + op.create_index( + "ix_worker_request_nonces_hotkey", + "worker_request_nonces", + ["hotkey"], + unique=False, + ) + + +def downgrade() -> None: + """Revert the migration.""" + + op.drop_index( + "ix_worker_request_nonces_hotkey", + table_name="worker_request_nonces", + ) + op.drop_index( + "ix_worker_request_nonces_created_at", + table_name="worker_request_nonces", + ) + op.drop_table("worker_request_nonces") + + op.drop_index("ix_worker_faults_work_unit_id", table_name="worker_faults") + op.drop_index("ix_worker_faults_worker_id", table_name="worker_faults") + op.drop_table("worker_faults") + + op.drop_index( + "ix_worker_registrations_status_miner_hotkey", + table_name="worker_registrations", + ) + op.drop_index( + "ix_worker_registrations_last_heartbeat_at", + table_name="worker_registrations", + ) + op.drop_index( + "ix_worker_registrations_miner_hotkey", + table_name="worker_registrations", + ) + op.drop_index( + "ix_worker_registrations_status", + table_name="worker_registrations", + ) + op.drop_table("worker_registrations") diff --git a/alembic/versions/0010_create_worker_assignments.py b/alembic/versions/0010_create_worker_assignments.py new file mode 100644 index 00000000..a03bab72 --- /dev/null +++ b/alembic/versions/0010_create_worker_assignments.py @@ -0,0 +1,147 @@ +"""Create the miner-funded worker assignment replica table. + +Adds ``worker_assignments`` (architecture.md sec 3.3): one row PER (gpu work +unit, worker) so a unit can be replicated to multiple distinct-owner workers +(R=2). Carries the worker identity the pull/post routes authenticate against, +the owner hotkey used for anti-collusion accounting, and the reported result + +``ExecutionProof`` (with the extracted ``manifest_sha256`` used for +reconciliation). No pre-existing table is altered. + +Revision ID: 0010_create_worker_assignments +Revises: 0009_create_worker_registry +Create Date: 2026-07-06 00:00:00.000000 +""" + +from __future__ import annotations + +from collections.abc import Sequence + +import sqlalchemy as sa + +from alembic import op + +revision: str = "0010_create_worker_assignments" +down_revision: str | None = "0009_create_worker_registry" +branch_labels: str | Sequence[str] | None = None +depends_on: str | Sequence[str] | None = None + +work_assignment_status = sa.Enum( + "pending", + "assigned", + "running", + "completed", + "failed", + "disputed", + name="work_assignment_status", + native_enum=False, +) + + +def upgrade() -> None: + """Apply the migration.""" + + op.create_table( + "worker_assignments", + sa.Column("id", sa.Uuid(as_uuid=True), nullable=False), + sa.Column("challenge_slug", sa.Text(), nullable=False), + sa.Column("work_unit_id", sa.Text(), nullable=False), + sa.Column("submission_ref", sa.Text(), nullable=False), + sa.Column("worker_id", sa.Text(), nullable=False), + sa.Column("worker_pubkey", sa.Text(), nullable=False), + sa.Column("miner_hotkey", sa.Text(), nullable=False), + sa.Column("payload", sa.JSON(), server_default="{}", nullable=False), + sa.Column( + "required_capability", + sa.Text(), + server_default="gpu", + nullable=False, + ), + sa.Column( + "status", + work_assignment_status, + server_default="pending", + nullable=False, + ), + sa.Column("attempt_count", sa.Integer(), server_default="0", nullable=False), + sa.Column("max_attempts", sa.Integer(), server_default="3", nullable=False), + sa.Column("deadline_at", sa.DateTime(timezone=True), nullable=True), + sa.Column("last_progress_at", sa.DateTime(timezone=True), nullable=True), + sa.Column("checkpoint_ref", sa.Text(), nullable=True), + sa.Column("result_success", sa.Boolean(), nullable=True), + sa.Column("result_payload", sa.JSON(), nullable=True), + sa.Column("manifest_sha256", sa.Text(), nullable=True), + sa.Column( + "created_at", + sa.DateTime(timezone=True), + server_default=sa.func.now(), + nullable=False, + ), + sa.Column( + "updated_at", + sa.DateTime(timezone=True), + server_default=sa.func.now(), + nullable=False, + ), + sa.PrimaryKeyConstraint("id", name=op.f("pk_worker_assignments")), + sa.UniqueConstraint( + "work_unit_id", + "worker_id", + name="uq_worker_assignments_work_unit_worker", + ), + ) + op.create_index( + "ix_worker_assignments_work_unit_id", + "worker_assignments", + ["work_unit_id"], + unique=False, + ) + op.create_index( + "ix_worker_assignments_status", + "worker_assignments", + ["status"], + unique=False, + ) + op.create_index( + "ix_worker_assignments_worker_pubkey", + "worker_assignments", + ["worker_pubkey"], + unique=False, + ) + op.create_index( + "ix_worker_assignments_status_worker_pubkey", + "worker_assignments", + ["status", "worker_pubkey"], + unique=False, + ) + op.create_index( + "ix_worker_assignments_status_deadline", + "worker_assignments", + ["status", "deadline_at"], + unique=False, + ) + + +def downgrade() -> None: + """Revert the migration.""" + + op.drop_index( + "ix_worker_assignments_status_deadline", + table_name="worker_assignments", + ) + op.drop_index( + "ix_worker_assignments_status_worker_pubkey", + table_name="worker_assignments", + ) + op.drop_index( + "ix_worker_assignments_worker_pubkey", + table_name="worker_assignments", + ) + op.drop_index( + "ix_worker_assignments_status", + table_name="worker_assignments", + ) + op.drop_index( + "ix_worker_assignments_work_unit_id", + table_name="worker_assignments", + ) + op.drop_table("worker_assignments") diff --git a/config/worker.example.yaml b/config/worker.example.yaml new file mode 100644 index 00000000..ad9291f4 --- /dev/null +++ b/config/worker.example.yaml @@ -0,0 +1,78 @@ +network: + name: base + netuid: 100 + chain_endpoint: null + wallet_name: default + wallet_hotkey: default + wallet_path: null + master_uid: 0 + +# Miner-funded GPU worker plane. All worker behavior is gated behind +# compute.worker_plane_enabled on the MASTER; a worker agent still needs the +# worker block below to know where the master is, how to identify itself, and +# how to enroll. +compute: + worker_plane_enabled: true + +worker: + agent: + # Master coordination-plane base URL (required). + master_url: http://127.0.0.1:8081 + gateway_url: null + capabilities: + - gpu + poll_interval_seconds: 5.0 + request_timeout_seconds: 15.0 + # The worker's OWN local Docker broker (falls back to docker.broker_url). + broker_url: http://127.0.0.1:8082 + broker_token_file: /run/secrets/base_broker_token + deploy: + # local | lium | targon + provider: local + gpu_count: 1 + max_price_per_hour: null + max_lifetime_hours: 1.0 + # PUBLICLY-pullable, digest-pinned worker image. REQUIRED for a lium/targon + # deploy (env BASE_WORKER__DEPLOY__IMAGE / BASE_WORKER__DEPLOY__IMAGE_DIGEST). + # A provider deploy REFUSES to run without these: a private-namespace GHCR image + # fails Lium pod creation (CREATION_FAILED), so it must be a public image. See + # docs/miner/worker-plane.md ("Publishing the worker image") for how to build, + # publish, and digest-pin docker/Dockerfile.worker. Not needed for provider=local. + image: null # e.g. ghcr.io//base-worker + image_digest: null # e.g. sha256:<64 hex> (the immutable pin) + # image_tag is INFORMATIONAL-ONLY: the deploy path pins BY DIGEST and never + # consumes this value, so it does not affect which image bytes run. Keep it as + # a human-readable note of the tag the digest was published under, if you like. + image_tag: null + # Metachar-free keep-alive: Lium rejects shell metacharacters (&&, ;, |). + startup_commands: tail -f /dev/null + ready_timeout_seconds: 60.0 + # ssh_public_key_file: /root/.config/prism-mission/lium_ssh_ed25519.pub + # ssh_key_name: prism-mission-worker + identity: + # Worker keypair (signs coordination requests + ExecutionProofs). Resolve + # from an sr25519 dev URI, a mnemonic, or a dedicated wallet; falls back to + # network.wallet when unset. + key_uri: null + key_mnemonic: null + wallet_name: null + wallet_hotkey: null + # Miner binding: EITHER a miner key that signs the binding at deploy time... + miner_key_uri: null + miner_key_mnemonic: null + miner_wallet_name: null + miner_wallet_hotkey: null + # ...OR a pre-signed binding (for a pod that never holds the miner key). + miner_hotkey: null + binding_signature: null + binding_nonce: null + +docker: + broker_url: http://127.0.0.1:8082 + broker_allowed_images: + - ghcr.io/baseintelligence/ + +observability: + log_json: true + sentry_dsn: null + otel_service_name: base-worker diff --git a/docker/Dockerfile.worker b/docker/Dockerfile.worker new file mode 100644 index 00000000..81f00f96 --- /dev/null +++ b/docker/Dockerfile.worker @@ -0,0 +1,48 @@ +FROM python:3.12-slim AS runtime + +ENV PYTHONDONTWRITEBYTECODE=1 \ + PYTHONUNBUFFERED=1 \ + UV_SYSTEM_PYTHON=1 \ + HOME=/var/lib/base \ + BASE_COMPUTE__WORKER_PLANE_ENABLED=true \ + PIP_NO_CACHE_DIR=1 + +WORKDIR /app + +# git: some `base @ git+...` challenge deps resolve via git. docker CLI: the +# worker agent dispatches gpu eval to its OWN local broker, so it needs the +# client binary to launch sibling eval containers. +RUN apt-get update \ + && apt-get install -y --no-install-recommends ca-certificates curl git \ + && rm -rf /var/lib/apt/lists/* \ + && python -m pip install --no-cache-dir uv docker + +RUN curl -fsSL https://download.docker.com/linux/static/stable/x86_64/docker-29.5.3.tgz -o /tmp/docker.tgz \ + && tar -xzf /tmp/docker.tgz -C /tmp \ + && install -m 0755 /tmp/docker/docker /usr/local/bin/docker \ + && rm -rf /tmp/docker /tmp/docker.tgz \ + && docker --version + +RUN groupadd --system --gid 1000 base \ + && useradd --system --uid 1000 --gid base --home-dir /var/lib/base --shell /usr/sbin/nologin base \ + && mkdir -p /var/lib/base /var/lib/base/secrets \ + && chown -R 1000:1000 /var/lib/base + +# The locally-built `base` package (this image's authoritative version). The +# validator extra pulls bittensor (sr25519 keypairs for the worker binding + +# ExecutionProof signatures). ``--chown`` makes the editable-install sources +# readable by the non-root runtime user. +COPY --chown=1000:1000 . . +RUN uv pip install --system -e ".[validator]" + +# Build-time smoke check: prove the worker CLI is importable + registered. +RUN base worker --help >/dev/null + +VOLUME ["/var/lib/base"] + +USER 1000:1000 + +# Exec-form entrypoint: no shell, no metacharacters. The agent reads +# config/worker.example.yaml plus BASE_* env overrides (master URL, worker key, +# miner-signed binding) supplied at `docker run` time and registers on boot. +CMD ["base", "worker", "agent", "--config", "config/worker.example.yaml"] diff --git a/docs/miner/worker-plane.md b/docs/miner/worker-plane.md new file mode 100644 index 00000000..b4c1aa15 --- /dev/null +++ b/docs/miner/worker-plane.md @@ -0,0 +1,473 @@ +# Miner Worker Deployment Guide (PRISM Compute Plane) + +Deploy a **miner-funded GPU worker** so your hotkey stays eligible to submit to PRISM and +so PRISM's heavy evaluation runs on GPU capacity **you** pay for (rented on Lium or Targon, +or your own hardware) instead of on BASE validators. + +This guide covers, for both **Lium** and **Targon**: + +- [What the worker plane is](#what-the-worker-plane-is) +- [Prerequisites](#prerequisites) +- [Configuration](#configuration) +- [Publishing the worker image](#publishing-the-worker-image) +- [Deploy on Lium](#deploy-on-lium) +- [Deploy on Targon](#deploy-on-targon) +- [Deploy locally / on your own hardware](#deploy-locally--on-your-own-hardware) +- [Costs and price guidance](#costs-and-price-guidance) +- [Monitoring your fleet](#monitoring-your-fleet) +- [Troubleshooting](#troubleshooting) + +> The whole worker plane is gated behind a feature flag +> (`compute.worker_plane_enabled` on the master, `worker_plane` in the prism config). With +> the flag off, nothing here applies and the subnet behaves exactly as before. + +--- + +## What the worker plane is + +- **`base worker agent`** is a BASE role that mirrors the validator agent loop + (register → heartbeat → pull → execute → post). It runs inside a GPU instance you fund. +- **Signed enrollment** binds your **miner hotkey** to a **worker keypair**: the miner signs + `worker-binding:{worker_pubkey}:{miner_hotkey}:{nonce}` (sr25519); the master verifies the + signature against the metagraph before the worker can pull work. +- **Admission rule**: when the master enforces it, you must have **≥1 active worker bound to + your hotkey** to submit to PRISM. A submission from a hotkey with no active worker is + rejected with HTTP `403 NO_ACTIVE_WORKER`. Deploying a worker is what unlocks submitting. +- **Anti-collusion**: a worker never evaluates a submission from its own owner; each unit is + replicated across **2 distinct owners** (R=2) and reconciled by comparing manifest hashes. +- **Proof tiers** attached to every result (see [Costs and price guidance](#costs-and-price-guidance) + for how tier affects audit rate): + - **Tier 0** (mandatory, all backends): deterministic manifest hash + your worker's sr25519 + signature. This is the source of truth for reconciliation and audits. + - **Tier 1**: the evaluator image digest matches the pinned worker/evaluator digest. On Lium + this is cross-checkable against the signed `GET /watchtower/digest`. + - **Tier 2**: in-guest TDX + nvtrust attestation. **Gated OFF on Targon today** (Targon + exposes no consumer-facing attestation surface), so Targon proofs carry tier ≤ 1. Tier 0 is + always available; tier 1 is available on both providers. + +Your provider API key (`LIUM_API_KEY` / `TARGON_API_KEY`) is used **only** to authenticate +your own provider calls from the CLI/agent. It is **never** sent to the master, never written +into the pod env, and never logged. + +--- + +## Prerequisites + +- A **registered miner hotkey** on the subnet (present in the metagraph). The worker binding is + signed by this hotkey. +- The `base` CLI. In a dev checkout run it through `uv`: + + ```bash + uv run base worker --help + ``` + + Expected: a command group listing `agent`, `deploy`, and `status`: + + ```text + Usage: base worker [OPTIONS] COMMAND [ARGS]... + + Deploy and manage miner-funded GPU worker agents + + ╭─ Commands ─────────────────────────────────────────────────────────────────╮ + │ agent Run the miner-funded GPU worker agent loop. │ + │ deploy Deploy a worker agent locally or onto a rented provider instance. │ + │ status Render the worker fleet from the master's ``GET /v1/workers``. │ + ╰────────────────────────────────────────────────────────────────────────────╯ + ``` + + > `base worker` (the miner-funded GPU worker plane) is a distinct, top-level group from the + > legacy `base master worker` group (Docker Swarm node management). They do not collide. + +- A **provider account with credit**: + - **Lium**: an API key (`LIUM_API_KEY`) and a registered SSH key (Lium requires an SSH public + key on every rent, even for non-interactive workloads). + - **Targon**: an API key (`TARGON_API_KEY`) and enough dashboard credit (Targon does **not** + expose balance via API; see [Troubleshooting](#troubleshooting)). +- A reachable BASE **master coordination-plane URL** (`worker.agent.master_url` in your config). + +Inspect the deploy options before running anything: + +```bash +uv run base worker deploy --help +``` + +Expected (abridged): `--provider` is **required** and accepts `lium | targon | local`, plus a +`--max-price` bound and a `--config` path: + +```text +Usage: base worker deploy [OPTIONS] + +╭─ Options ──────────────────────────────────────────────────────────────────╮ +│ * --provider TEXT Where to run the worker agent: lium | targon | local.│ +│ --max-price FLOAT Max price per GPU/hour bounding provider offer │ +│ selection. │ +│ --config PATH [default: config/worker.example.yaml] │ +╰──────────────────────────────────────────────────────────────────────────────╯ +``` + +--- + +## Configuration + +Start from [`config/worker.example.yaml`](../../config/worker.example.yaml) and set at least: + +```yaml +compute: + worker_plane_enabled: true + +worker: + agent: + master_url: https:// # required: the coordination plane + broker_url: http://127.0.0.1:8082 # the worker's OWN local Docker broker + capabilities: + - gpu + deploy: + provider: lium # local | lium | targon (also selectable via --provider) + gpu_count: 1 + max_price_per_hour: 1.50 # cost cap per GPU/hour (also selectable via --max-price) + max_lifetime_hours: 1.0 # bounded pod lifetime (keep it short) + # REQUIRED for a lium/targon deploy: a PUBLICLY-pullable, digest-pinned worker + # image. A provider deploy refuses to run without both (see "Publishing the + # worker image"). Not needed for provider: local. + image: ghcr.io//base-worker + image_digest: sha256:<64 hex> # the immutable pin + image_tag: v1 # informational-only (see note below) + startup_commands: tail -f /dev/null # MUST be metachar-free (see Troubleshooting) + ssh_public_key_file: /path/to/your_key.pub # Lium requires an SSH key + ssh_key_name: my-worker-key + identity: + # The WORKER keypair signs coordination requests + ExecutionProofs. + key_uri: null # e.g. an sr25519 URI, a mnemonic, or a wallet; else falls back to network.wallet + # The MINER binding: EITHER a miner key that signs the binding at deploy time... + miner_key_uri: null + # ...OR a pre-signed binding (for a pod that never holds the miner key): + miner_hotkey: null + binding_signature: null + binding_nonce: null +``` + +Every field is overridable by env with the `BASE_` prefix and `__` nesting, e.g. +`BASE_WORKER__AGENT__MASTER_URL=...` or `BASE_COMPUTE__WORKER_PLANE_ENABLED=true`. + +**Identity note.** The worker needs a worker keypair (falls back to `network.wallet` when unset) +and a miner-signed binding. On a rented pod that must never hold your miner key, sign the binding +on your own machine and pass the pre-signed `miner_hotkey` + `binding_signature` + `binding_nonce` +instead of a miner key. The CLI does this for you when a miner key is configured locally. + +--- + +## Publishing the worker image + +A `lium`/`targon` deploy runs a **worker image** you must publish yourself, pinned by digest. +`base worker deploy --provider lium|targon` **refuses to run** unless `worker.deploy.image` **and** +`worker.deploy.image_digest` are set (env `BASE_WORKER__DEPLOY__IMAGE` / +`BASE_WORKER__DEPLOY__IMAGE_DIGEST`). There is deliberately **no baked-in default image**: + +> ⚠️ **A private-namespace registry image can fail provider pod creation.** Pinning a **private** +> GHCR image (e.g. `ghcr.io//...`) makes **Lium pod creation fail with +> `CREATION_FAILED`** — the rented executor host cannot pull a private image. The worker image +> **must be PUBLICLY pullable**. This is why the deploy will not silently pin a placeholder. + +> 🚧 **Publishing is a release/operator action, not something the deploy does for you.** Pushing to +> a public registry needs registry push credentials, so this step is performed once by whoever owns +> the release (it may require access you do not have as a miner — coordinate with the subnet +> operators for the canonical published image + digest). + +Build, publish, and pin `docker/Dockerfile.worker` (a slim `python:3.12-slim`-based image carrying +the `base` worker agent + docker CLI): + +```bash +# 1. Build from the base repo root (the Dockerfile COPYs the source in). +docker build -f docker/Dockerfile.worker -t ghcr.io//base-worker:v1 . + +# 2. Log in and push to a PUBLIC registry namespace (make the package public in the +# registry UI so an unauthenticated provider host can pull it). +echo "$GHCR_PAT" | docker login ghcr.io -u --password-stdin +docker push ghcr.io//base-worker:v1 + +# 3. Read back the immutable digest the registry assigned, and pin THAT. +docker buildx imagetools inspect ghcr.io//base-worker:v1 \ + --format '{{println .Manifest.Digest}}' +# -> sha256:<64 hex> +``` + +Then set the image + digest in your config (or via env), e.g.: + +```yaml +worker: + deploy: + image: ghcr.io//base-worker + image_digest: sha256:<64 hex> # the value printed above + image_tag: v1 # informational-only (see note below) +``` + +> **`image_tag` is informational-only.** A `lium`/`targon` deploy provisions the worker +> **by digest** (`image` + `image_digest`); `worker.deploy.image_tag` is **not consumed** by +> the deploy path and does **not** affect which image bytes run. Keep it only as a +> human-readable note of the tag the digest was published under. What earns **tier 1** is the +> digest pin, not the tag. + +**Verify it is actually pullable before deploying.** From a clean host with no registry auth: + +```bash +docker pull ghcr.io//base-worker@sha256:<64 hex> +``` + +If that pull fails unauthenticated, the image is not public and a Lium deploy will +`CREATION_FAILED`. Digest-pinning (`@sha256:...`) is also what earns **tier 1** proofs: the running +image bytes are provably the ones you published (cross-checkable on Lium against +`GET /watchtower/digest`). + +--- + +## Deploy on Lium + +1. Export **your** Lium API key (never commit it): + + ```bash + export LIUM_API_KEY= + ``` + +2. Deploy, bounding the price per GPU/hour: + + ```bash + uv run base worker deploy --provider lium --max-price 1.50 + ``` + + The CLI (offline planning first, then a single rent): + - lists Lium executors and selects the cheapest **suitable** offer at or below `--max-price`, + preferring an executor whose GPU count matches your request (renting a partial slice of a + multi-GPU executor is rejected by Lium); + - ensures the worker template and your SSH key exist (idempotent — reused if already present); + - rents with a bounded `termination_hours` (from `max_lifetime_hours`) and your SSH key; + - boots the worker image, which enrolls with the master under your signed binding. + + Expected output shape: + + ```text + Selected lium offer (NVIDIA A100-SXM4-80GB x1) @ 0.52/GPU/hr + Provisioned lium instance (status=PENDING); the worker enrolls with the master on boot + ``` + +3. Wait for the pod to reach `RUNNING` (usually a minute or two; allow up to ~15 min for a cold + image pull) and confirm enrollment with [`base worker status`](#monitoring-your-fleet). + +> **Cost guardrails are enforced.** Deploy never issues an unbounded rent: a bounded +> `termination_hours` is always set, and an offer above your cap is refused before any rent is +> issued. Always confirm the pod is gone when you are done (see Troubleshooting → leftover pods). + +--- + +## Deploy on Targon + +1. Export **your** Targon API key: + + ```bash + export TARGON_API_KEY= + ``` + +2. Deploy: + + ```bash + uv run base worker deploy --provider targon --max-price 3.50 + ``` + + The CLI lists the live Targon inventory (`GET /inventory?type=rental&gpu=true`), selects a + GPU shape within `--max-price` by `cost_per_hour`, and deploys the worker app; the worker + enrolls with the master on boot, exactly as on Lium. + +3. Confirm enrollment with [`base worker status`](#monitoring-your-fleet). + +> **Targon has no balance API.** The CLI cannot pre-check your credit. If a deploy fails with an +> insufficient-credit error, it is surfaced as a distinct typed error and is **not retried** — +> top up in the Targon web dashboard and retry. See [Troubleshooting](#troubleshooting). + +--- + +## Deploy locally / on your own hardware + +If you already run a GPU box (or want to test against a local master), no provider key is needed: + +```bash +uv run base worker deploy --provider local +``` + +This starts a worker agent process on this host pointed at `worker.agent.master_url`, then polls +`GET /v1/workers` until the worker reaches `active` (bounded by `deploy.ready_timeout_seconds`, +default 60s). On success it prints: + +```text +Started worker agent process pid= (provider=local) +Worker active (pubkey=, owner=, provider=local) +``` + +To run the long-lived agent loop directly (e.g. under a supervisor), use: + +```bash +uv run base worker agent --config config/worker.example.yaml +``` + +--- + +## Costs and price guidance + +**You pay for the instance.** The subnet does not reimburse GPU spend; running a worker is the +price of being submission-eligible and contributing compute. + +- **Bound every deploy** with `--max-price` (per GPU/hour) or `worker.deploy.max_price_per_hour`. + An all-over-cap situation fails with a clear "no offer within budget" error and provisions + nothing. +- **Keep pod lifetime short.** `worker.deploy.max_lifetime_hours` maps to the provider's bounded + termination window; keep it small (1–2h) and redeploy as needed. +- **Prefer clean single-GPU executors.** On Lium, `gpu_count=1` is the common case; deploy + planning prefers an executor whose GPU count matches your request. + +**Lium price signal** (real USD, checked live at selection time): + +- Executor prices are exposed per GPU/hour. Prefer offers **< $1.50/GPU/hr**. +- Observed live examples: an A100-SXM4-80GB at ~$0.52/GPU/hr; RTX-class nodes under $1/GPU/hr. + A short rent → run → delete cycle cost on the order of **$0.004**. +- Check your balance any time with Lium's `GET /users/me` (returns a numeric `balance`). + +**Targon price signal** (from live inventory `cost_per_hour`): + +- Observed: H100 ×1 ≈ **$2.50/hr**, H200 ×1 ≈ **$3.29/hr**, B200 ×1 ≈ **$5.30/hr** (multi-GPU + shapes scale up). Confidential-compute shapes start around $2.50/hr. +- **No balance endpoint** — track spend in the Targon dashboard. + +**Proof tier affects audit load, not price.** Lower-assurance proofs are audited more often: +tier 0 ≈ 10%, tier 1 ≈ 5%, tier 2 ≈ 2% of results replayed. A worker whose image digest matches +the pinned evaluator digest earns tier 1 and is audited less. + +--- + +## Monitoring your fleet + +Render the fleet the master knows about: + +```bash +uv run base worker status +``` + +Output is one row per worker with its status, owner, provider, fault count, and last-seen time: + +```text +WORKER_ID OWNER PROVIDER STATUS FAULTS LAST_SEEN +worker-abc123 5F...minerhotkey lium active 0 2026-07-07T14:30:00+00:00 +``` + +Filter to just the **active** workers of one hotkey (this is the same query the admission rule +uses): + +```bash +uv run base worker status --hotkey +``` + +**Status lifecycle:** `pending → active → stale → retired`. + +- `pending`: registered, awaiting its first heartbeat. +- `active`: verified binding **and** a heartbeat within the freshness window + (`compute.worker_heartbeat_ttl_seconds`, default 120s). Only `active` workers get assignments + and satisfy the admission rule. +- `stale`: no heartbeat within the TTL. Stops receiving new units; returns to `active` on the + next heartbeat. +- `retired`: terminal. A retired worker is never resurrected by a heartbeat — re-enroll a fresh + worker instead. + +`FAULTS` counts audit-attributed faults (a worker whose manifest diverged from an authoritative +validator replay). The same fleet state is available as JSON from the master at `GET /v1/workers`; +`base worker status` and that endpoint agree field-for-field. + +--- + +## Troubleshooting + +**`provider '' requires the environment variable` +(exit code 2).** +The provider key env var is unset or blank. The CLI refuses **before** any provider or master +call, so nothing was provisioned. Export the key and retry: + +```bash +export LIUM_API_KEY= # or TARGON_API_KEY for Targon +uv run base worker deploy --provider lium --max-price 1.50 +``` + +**`unsupported provider ''; expected one of local, lium, targon` (exit code 2).** +Use exactly one of `local`, `lium`, or `targon` for `--provider`. + +**`no rentable offer within budget (...); nothing was provisioned` (exit code 1).** +Every listed offer is above your cap. Raise `--max-price` (or `worker.deploy.max_price_per_hour`), +or retry later when cheaper capacity appears. Nothing was rented. + +**`provider '' deploy requires an explicit worker image ...` (exit code 1).** +`worker.deploy.image` and/or `worker.deploy.image_digest` are unset. A provider deploy refuses to +run without a **PUBLICLY-pullable, digest-pinned** worker image (it will not silently pin a +placeholder). The check is **fail-fast** — no provider call is made. Set both (or the env vars +`BASE_WORKER__DEPLOY__IMAGE` / `BASE_WORKER__DEPLOY__IMAGE_DIGEST`); see +[Publishing the worker image](#publishing-the-worker-image). The digest must be `sha256:<64 hex>` +(a mutable tag is not an immutable pin). + +**Lium pod stuck at / fails with `CREATION_FAILED`.** +The executor host could not pull your worker image — almost always because the image is in a +**private** registry namespace. Publish the image **publicly** and confirm an unauthenticated +`docker pull @` succeeds, then redeploy. See +[Publishing the worker image](#publishing-the-worker-image). + +**Lium: `403 Request blocked` when deploying.** +Lium's edge CDN/WAF rejects **any** request body that contains a **loopback URL** +(`http://127.0.0.1...` / `http://localhost...`), and `base worker deploy` bakes the pod env into the +`POST /templates` body. The CLI now **omits loopback coordination URLs** (master/broker/gateway) +from that body automatically — a loopback `worker.agent.master_url` no longer trips the WAF, and the +agent resolves the master URL at runtime from its own config (reaching a loopback master via a +reverse SSH tunnel). Set `worker.agent.master_url` to a **public** master URL for a normal remote +deploy; keep the loopback default only for a local/tunnelled master. + +**Targon balance is invisible / deploy fails on credits.** +Targon exposes **no** balance, credits, or billing endpoint via its API — the value only exists +in the web dashboard. The client will not silently return `0`/`None` or probe guessed endpoints. +If a deploy fails with an insufficient-credit-style error, it is raised as a **distinct typed +error** and is **not retried**. Top up credit in the Targon dashboard, then retry. Do not loop +the deploy. + +**Lium: `Malicious startup command detected: shell metacharacters are not allowed` (HTTP 400).** +Lium rejects a rent whose template `startup_commands` contains shell metacharacters (`&&`, `;`, +`|`, quotes, …). Many public CUDA templates use `bash -c "service ssh start && tail -f /dev/null"` +and are therefore un-rentable. Keep `worker.deploy.startup_commands` a single metachar-free +command — the default `tail -f /dev/null` works (Lium's pod agent provides SSH independently of +the container command). + +**Lium: renting a single GPU on a multi-GPU node returns HTTP 400.** +Renting a partial slice of a multi-GPU executor is rejected. Prefer an executor whose GPU count +equals your request (deploy planning already prefers an exact-GPU-count executor and falls back +to the next-cheapest suitable offer if the cheapest loses an availability race). + +**Lium requires an SSH key even for non-interactive workers.** +Every Lium rent needs an SSH public key. Set `worker.deploy.ssh_public_key_file` (or +`ssh_public_key`) and, if needed, `ssh_key_name`. Register the key with your Lium account first. + +**Pod is slow to reach `RUNNING`.** +Provisioning takes minutes, not seconds; a cold image pull can take up to ~15 minutes. Poll +`base worker status` until the worker shows `active`. + +**`worker did not reach active within s` (local deploy).** +The agent started but never enrolled/heartbeated in time. Check that `worker.agent.master_url` is +reachable, that `compute.worker_plane_enabled` is on at the master, that the miner hotkey is in +the metagraph, and that the binding is valid. Increase `worker.deploy.ready_timeout_seconds` for +slow networks. + +**Worker shows `stale`.** +Heartbeats stopped for longer than `compute.worker_heartbeat_ttl_seconds` (default 120s). It +returns to `active` on the next successful heartbeat; if it does not, check the agent process and +its connectivity to the master. + +**Submission rejected with `403 NO_ACTIVE_WORKER`.** +The admission rule requires ≥1 **active** worker bound to your hotkey. Deploy a worker, wait for +`base worker status --hotkey ` to show it `active`, then resubmit. + +**Leftover pod after a test.** +Deploy always sets a bounded lifetime, but you should still confirm nothing is left running when +you are done: on Lium, `GET /pods` should return an empty list for your account. A pod left +running keeps billing you. + +> **Never commit or log your provider key.** Keep `LIUM_API_KEY` / `TARGON_API_KEY` in your +> environment only; the CLI never forwards them to the master and never writes them into a pod. diff --git a/docs/operations/mission-harness.md b/docs/operations/mission-harness.md new file mode 100644 index 00000000..38180032 --- /dev/null +++ b/docs/operations/mission-harness.md @@ -0,0 +1,174 @@ +# Local mission harness (cross-repo end-to-end) + +A thin, config-driven launcher that stands up the **whole compute plane locally on loopback ports +3100-3199 with no GPU** and runs six operator-observable drills (all via HTTP/CLI only). It is a +test/mission harness, **not a production path**: it uses a static mock metagraph and the prism +repo's own CPU re-exec seam. Production still uses `base master proxy` with a real subtensor +metagraph and the real GPU broker. + +## What it stands up + +| Component | Script | Repo | +| --- | --- | --- | +| base master (worker plane ON, validator plane ON, orchestration loop) | `base/scripts/mission/mission_master.py` | base | +| prism service (worker plane + admission gate + CPU re-exec test mode) | `prism/scripts/mission/mission_prism.py` | prism | +| worker agents (base worker plane; executor = prism CPU re-exec) | `prism/scripts/mission/mission_worker.py` | prism | +| stub gpu validator (audits disputes by deterministic replay) | `prism/scripts/mission/mission_validator.py` | prism | +| orchestrator + drills + PID teardown | `prism/scripts/mission/launch.py` | prism | + +* **Mock metagraph**: five owner (miner) hotkeys (`//MissionAlice…//MissionErin`, no permit) plus a + validator hotkey (`//MissionValidator1`, with a validator permit), seeded statically so signed + worker/validator requests authenticate with no chain. +* **CPU re-exec seam as explicit config**: prism runs with + `worker_plane.cpu_reexec_test_mode: true`, which installs + `prism_challenge.evaluator.mock_reexec.cpu_reexec_run` over `DockerExecutor.run` and uses a tiny + deterministic `PrismContext` (vocab 64 / seq 16 / step budget 24). The launcher never touches any + HTTP surface; all drill observations are black-box. +* **Two+ workers on DISTINCT owner hotkeys**: each worker's executor runs the real CPU re-exec, + normalizes the volatile `compute` timing fields, and signs a tier-0 ExecutionProof, so honest + replicas of the same submission converge on one `manifest_sha256`. + +## Virtualenv (important) + +The current base source and the current prism source must share ONE interpreter that also has +`torch`. The prism virtualenv has torch + the current prism source, but a *stale* installed `base`; +run everything through the prism venv with the current base source shadowed onto `PYTHONPATH`: + +```bash +cd # the dir containing both base/ and prism/ +export PYTHONPATH="$PWD/base/src:$PWD/prism/src" +prism/.venv/bin/python prism/scripts/mission/launch.py # all drills +prism/.venv/bin/python prism/scripts/mission/launch.py --only 1,4 # a subset +``` + +The launcher prints the PID of every spawned process and, in a `finally`, **kills them all by +PID** so nothing is left listening. + +## Drills (observable via API/CLI only) + +1. **VAL-CROSS-001** full pipeline: a submission from a hotkey with an active worker → exactly one + gpu work unit on prism `/internal/v1/work_units` → assigned to 2 distinct-owner workers → both + post ExecutionProofs → reconciliation finds equal `manifest_sha256`, accepts, forwards one + result → prism records a score (`GET /v1/submissions/{id}`). +2. **VAL-CROSS-002** self-eval exclusion under scarcity: with exactly 2 active workers where one is + owned by the submitter H, H's unit is never assigned to H's own worker (it degrades to the other + worker or holds pending). +3. **VAL-CROSS-003** divergence: one worker corrupts its manifest → reconciliation disputes the + unit → a `:audit` unit (`required_capability=gpu`, executor kind `validator`) is created and + assigned to the stub validator → the validator's deterministic replay yields the authoritative + hash → the diverging worker gets a `worker_faults` row visible in `GET /v1/workers`, and the + submission carries no finalized score. +4. **VAL-CROSS-004** admission gate: `POST /v1/submissions` from a hotkey with no active worker → + 403 `NO_ACTIVE_WORKER`; after a worker for that hotkey registers + heartbeats + (`GET /v1/workers/active?hotkey=` returns ≥1) the resubmission (fresh nonce) is accepted. +6. **VAL-CROSS-009** fleet agreement: `GET /v1/workers` and `base worker status` (pointed at the + same master) report the same worker ids, owners, providers, statuses, and fault counts. +11. **VAL-CROSS-011** dispute lifecycle discoverable via APIs alone: after the divergence drill, the + whole dispute → audit → invalidation → fault chain is reconstructed from operator surfaces with + **zero DB/file reads**: + * **(a) disputed state + (b) audit unit** come from the signed `GET /v1/workers/units` (the + master unit-status read surface): the primary unit reads `status: "disputed"` with its replicas + (worker id, owner hotkey, posted `manifest_sha256`, proof presence), and its `audit` block + names the linked audit unit id, `executor_kind: "validator"`, and terminal `outcome` + (`pending` / `passed` / `mismatch-resolved`). + * **(c) invalidation** comes from prism `GET /v1/submissions/{id}` (never a live-ranked + `completed` score). + * **(d) fault** is visible on both `GET /v1/workers` and `base worker status`. + +The dispute-lifecycle reconstruction (VAL-CROSS-011) runs at the end of drill 3 (it needs a live +dispute), so selecting drill `3` also exercises it. + +`GET /v1/workers/units` is a signed read-only surface authenticated exactly like `GET /v1/workers` +(a registered worker OR an eligible validator signed request — NOT the admission internal bridge +bearer). With `compute.worker_plane_enabled` off its router is unmounted (404). + +## Ports + +* master: `127.0.0.1:3110` +* prism: `127.0.0.1:3120` + +Both are inside the mission-allowed range 3100-3199. No component listens outside that range. + +## Legacy regression (flags OFF) — exact verification procedure (VAL-CROSS-006) + +This is the reproducible procedure proving that with ALL new flags OFF the compute plane behaves +byte-for-byte as pre-mission. Re-run it exactly as below (scrutiny/re-validation). It has two parts: +(a) both default suites green fully offline, and (b) a legacy validator-flow smoke. + +Prerequisites: the test PostgreSQL is up on 15433 (`services.yaml` `test-postgres`), and both venvs +are synced (`services.yaml` `install`). No credentials are needed; the procedure UNSETS +`LIUM_API_KEY` / `TARGON_API_KEY` / `BASE_LIVE_PROVIDER_TESTS` and leaves +`BASE_COMPUTE__WORKER_PLANE_ENABLED` unset (flags OFF). + +### Offline egress guard + +`base/scripts/mission/no_external_egress.py` is a pytest plugin (`-p no_external_egress`) that blocks +any DNS resolution or socket connect to a non-loopback host (raising an `OSError` subclass), while +leaving loopback (test PostgreSQL on 127.0.0.1:15433, in-process stubs) and AF_UNIX sockets working. +It installs from `pytest_configure` (before collection), so even a module that probes the network at +import time is guarded and simply SKIPS. This operationalises "zero real egress to lium.io / +api.targon.com". Unit-tested by `base/tests/unit/test_no_external_egress.py`. + +### (a) Both default suites, offline, flags OFF + +base (with the test PostgreSQL on 15433): + +```bash +cd base +env -u LIUM_API_KEY -u TARGON_API_KEY -u BASE_LIVE_PROVIDER_TESTS -u BASE_COMPUTE__WORKER_PLANE_ENABLED \ + PYTHONPATH="$PWD/scripts/mission" \ + BASE_TEST_DATABASE_URL=postgresql+asyncpg://base:base@localhost:15433/base_test \ + uv run pytest -p no_external_egress --cov=base --cov-report=term-missing --cov-fail-under=80 \ + -q -p no:cacheprovider +``` + +Expected: green, e.g. `1543 passed`, `Total coverage: 88.55%` (≥ 80% gate). No base test needs the +network; the live Lium/Targon E2E is a standalone gated script (`base/scripts/live_lium_e2e.py`, +`BASE_LIVE_PROVIDER_TESTS=1`) NOT collected by the default suite, so there are no live tests to skip. + +prism (the same egress guard, from the base scripts dir; `distributed_gloo` deselected as non-gating): + +```bash +cd prism +env -u LIUM_API_KEY -u TARGON_API_KEY -u BASE_LIVE_PROVIDER_TESTS -u BASE_COMPUTE__WORKER_PLANE_ENABLED \ + PYTHONPATH="/root/prism-compute-plane/base/scripts/mission" \ + uv run pytest -p no_external_egress -m "not distributed_gloo" --cov=prism_challenge \ + --cov-report=term-missing --cov-fail-under=80 -q -p no:cacheprovider +``` + +Expected: green, e.g. `892 passed, 9 skipped, 7 deselected`, `Total coverage: 83.07%` (≥ 80% gate). +The `7 deselected` are the `distributed_gloo` tests (non-gating). The `9 skipped` are the +harness/dataset tests whose prep step needs to stage a tokenizer/dataset over the network +(`"... needs network (prep/build step)"`); under the egress guard (external resolution disabled) +they skip rather than run — exactly the offline behaviour. Nothing egresses to `lium.io` / +`api.targon.com`. + +### (b) Legacy validator-flow smoke + +Stands up the SAME local deployment with flags OFF (base master `worker_plane_enabled: false`, prism +`worker_plane_enabled: false`) and drives one legacy submission end-to-end. Scripts: +`prism/scripts/mission/legacy_smoke.py` (orchestrator) + `mission_legacy_validator.py` (a real +`validator_dispatch` executor) + the flags-OFF configs of `mission_master.py` / `mission_prism.py`. +It prints every PID and kills them all by PID in a `finally`. + +```bash +cd /root/prism-compute-plane # dir containing base/ and prism/ +export PYTHONPATH="$PWD/base/src:$PWD/prism/src" +prism/.venv/bin/python prism/scripts/mission/legacy_smoke.py +``` + +Expected: `OVERALL: PASS`. The smoke asserts, via HTTP + a read of the master/prism databases: + +1. the submission is ACCEPTED with no `NO_ACTIVE_WORKER` 403 (admission gate off, flags OFF); +2. prism exposes exactly ONE gpu work unit for it (`GET /internal/v1/work_units`); +3. the base master assigns that unit to the VALIDATOR + (`work_assignments.assigned_validator_hotkey` == the validator hotkey, `required_capability=gpu`) + and NEVER to a worker: `worker_registrations` and `worker_assignments` are empty and + `GET /v1/workers` shows zero workers — no worker-plane rows/side effects; +4. the validator executes it via the real `validator_dispatch` path + (`mission_legacy_validator.py` logs `pulled + executing prism unit … via validator_dispatch`, + `pulled=1 executed=1`) and prism records a score + (`GET /v1/submissions/{id}` → `completed` with a `final_score`). + +`legacy_smoke.py` uses ports 3112 (master) / 3122 (prism), inside 3100-3199, distinct from the +worker-plane drills' 3110/3120 so the two harnesses never collide. diff --git a/scripts/live_lium_e2e.py b/scripts/live_lium_e2e.py new file mode 100644 index 00000000..9f67f106 --- /dev/null +++ b/scripts/live_lium_e2e.py @@ -0,0 +1,493 @@ +#!/usr/bin/env python3 +"""Live Lium/Targon validation driven through the M1 provider clients (opt-in). + +Gated behind ``BASE_LIVE_PROVIDER_TESTS=1`` because it makes real, billable calls +to production Lium (architecture.md sec 3.1 cost guardrails + mission AGENTS.md +money rules). It exercises the SAME ``LiumClient`` / ``TargonClient`` code the +worker plane uses, never raw curl. + +Two phases: + +* Read-only reachability (no mutating call): Lium ``users/me`` balance, + ``executors`` offers, ``watchtower/digest``; Targon ``inventory`` offers and an + authenticated ``apps`` listing (VAL-PROV-012 / VAL-PROV-013). +* ONE batched Lium rental cycle (VAL-PROV-014): ensure the SSH key + template, + rent the cheapest suitable executor (< $1.50/GPU/hr, ``termination_hours=1``, + pod name ``mission-``), poll to RUNNING, SSH ``nvidia-smi -L``, pull logs, + DELETE, verify the pod is gone, and record the balance delta (must be <= $1). + +The pod is deleted in a ``finally`` block on EVERY path (including exceptions and +a partially provisioned pod) so a failure never leaks a billable pod. A JSON trace +of the whole sequence is written to ``$LIUM_E2E_TRACE`` (default +``/tmp/lium_e2e_trace.json``). Credentials/keys are read from the environment or +the mission credentials file; secret values are never printed. +""" + +from __future__ import annotations + +import asyncio +import json +import os +import re +import subprocess +import sys +import time +from collections.abc import Mapping +from pathlib import Path +from typing import Any + +from base.compute import InstanceSpec, LiumClient, Offer, TargonClient + +MAX_PRICE_PER_GPU_HR = 1.50 +TERMINATION_HOURS = 1 +POLL_BUDGET_SECONDS = 15 * 60 +POLL_INTERVAL_SECONDS = 25.0 +SSH_ATTEMPTS = 8 +SSH_RETRY_SECONDS = 15.0 +LOG_LINE_CAP = 40 +LOG_STREAM_TIMEOUT_SECONDS = 20.0 + +RUNNING_STATUSES = {"RUNNING"} +TERMINAL_FAIL_STATUSES = {"FAILED", "CREATION_FAILED", "BROKEN", "STOPPED"} + +DEFAULT_CREDENTIALS_FILE = "/root/.config/prism-mission/credentials.env" +DEFAULT_SSH_PRIVATE_KEY = "/root/.config/prism-mission/lium_ssh_ed25519" +SSH_KEY_NAME = "prism-mission-readiness" + +E2E_TEMPLATE_NAME = "prism-mission-e2e" +E2E_IMAGE = "nvidia/cuda" +E2E_IMAGE_TAG = "12.4.1-base-ubuntu22.04" +# Single, metachar-free command: Lium rejects startup commands containing shell +# metacharacters. It keeps the container alive; Lium's pod agent serves SSH. +E2E_STARTUP_COMMAND = "tail -f /dev/null" + + +class _TracingLiumClient(LiumClient): + """LiumClient that records the raw ``POST .../rent`` response shape. + + The rent response body is what :meth:`LiumClient._extract_pod_id` parses (with + a ``GET /pods`` lookup by ``pod_name`` fallback). Capturing it lets the live + run confirm the real shape without changing client behavior. + """ + + def __init__(self, *args: Any, **kwargs: Any) -> None: + super().__init__(*args, **kwargs) + self.rent_capture: dict[str, Any] | None = None + + async def _request( + self, + method: str, + path: str, + *, + json_body: Any | None = None, + params: Mapping[str, Any] | None = None, + ) -> Any: + response = await super()._request( + method, path, json_body=json_body, params=params + ) + if method == "POST" and path.endswith("/rent"): + body: Any + try: + body = response.json() if response.content else None + except ValueError: + body = response.text[:1000] + self.rent_capture = { + "status_code": response.status_code, + "body_type": type(body).__name__, + "body_keys": (sorted(body.keys()) if isinstance(body, dict) else None), + "body_preview": _preview(body), + } + return response + + +def _preview(body: Any) -> Any: + if isinstance(body, dict): + return {k: _preview(v) for k, v in list(body.items())[:20]} + if isinstance(body, list): + return [_preview(v) for v in body[:3]] + if isinstance(body, str): + return body[:200] + return body + + +def _load_credentials() -> dict[str, str]: + creds: dict[str, str] = {} + for name in ("LIUM_API_KEY", "TARGON_API_KEY"): + value = os.environ.get(name) + if value: + creds[name] = value + creds_file = Path(os.environ.get("PRISM_MISSION_CREDS", DEFAULT_CREDENTIALS_FILE)) + if creds_file.is_file(): + for raw in creds_file.read_text().splitlines(): + line = raw.strip() + if not line or line.startswith("#"): + continue + if line.startswith("export "): + line = line[len("export ") :] + if "=" not in line: + continue + key, val = line.split("=", 1) + key = key.strip() + creds.setdefault(key, val.strip().strip('"').strip("'")) + missing = [n for n in ("LIUM_API_KEY", "TARGON_API_KEY") if not creds.get(n)] + if missing: + raise SystemExit(f"missing credentials: {', '.join(missing)}") + return creds + + +def _load_ssh_public_key() -> str: + pub = Path(os.environ.get("LIUM_SSH_PRIVATE_KEY", DEFAULT_SSH_PRIVATE_KEY) + ".pub") + if not pub.is_file(): + raise SystemExit(f"missing SSH public key: {pub}") + return pub.read_text().strip() + + +def _parse_ssh_target(ssh_connect_cmd: str, raw: Mapping[str, Any]) -> dict[str, Any]: + user_host = re.search(r"([\w.-]+)@([\w.\-]+)", ssh_connect_cmd or "") + port_match = re.search(r"-p\s+(\d+)", ssh_connect_cmd or "") + target: dict[str, Any] = {} + if user_host: + target["user"] = user_host.group(1) + target["host"] = user_host.group(2) + if port_match: + target["port"] = int(port_match.group(1)) + if "port" not in target: + mapping = raw.get("ports_mapping") + if isinstance(mapping, Mapping): + external = mapping.get("22") or mapping.get(22) + if external is not None: + try: + target["port"] = int(external) + except (TypeError, ValueError): + pass + return target + + +def _run_ssh_nvidia_smi(target: Mapping[str, Any], private_key: str) -> dict[str, Any]: + argv = [ + "ssh", + "-i", + private_key, + "-o", + "StrictHostKeyChecking=no", + "-o", + "UserKnownHostsFile=/dev/null", + "-o", + "ConnectTimeout=15", + "-o", + "BatchMode=yes", + ] + if target.get("port"): + argv += ["-p", str(target["port"])] + argv.append(f"{target['user']}@{target['host']}") + argv.append("nvidia-smi -L && echo SMOKE_OK") + last: dict[str, Any] = {} + for attempt in range(1, SSH_ATTEMPTS + 1): + try: + proc = subprocess.run(argv, capture_output=True, text=True, timeout=60) + except subprocess.TimeoutExpired: + last = {"attempt": attempt, "returncode": None, "error": "ssh timed out"} + time.sleep(SSH_RETRY_SECONDS) + continue + last = { + "attempt": attempt, + "returncode": proc.returncode, + "stdout": proc.stdout.strip(), + "stderr": proc.stderr.strip()[:500], + } + if proc.returncode == 0 and "SMOKE_OK" in proc.stdout: + return last + time.sleep(SSH_RETRY_SECONDS) + return last + + +async def _collect_logs(client: LiumClient, pod_id: str) -> list[str]: + lines: list[str] = [] + + async def _drain() -> None: + async for line in client.stream_logs(pod_id): + lines.append(line) + if len(lines) >= LOG_LINE_CAP: + break + + try: + await asyncio.wait_for(_drain(), timeout=LOG_STREAM_TIMEOUT_SECONDS) + except Exception as exc: # noqa: BLE001 - logs are best-effort + lines.append(f"") + return lines + + +def _suitable_offers(offers: list[Offer]) -> list[Offer]: + suitable = [ + o + for o in offers + if 0 < o.price_per_hour <= MAX_PRICE_PER_GPU_HR and o.gpu_count >= 1 + ] + suitable.sort(key=lambda o: (o.price_per_hour, o.gpu_count)) + return suitable + + +async def _ensure_e2e_template(client: LiumClient) -> dict[str, Any]: + """Ensure a dedicated E2E template with a rent-safe startup command exists. + + Lium rejects a rent whose template startup command contains shell + metacharacters ("Malicious startup command detected"), so the many public + templates that chain ``service ssh start && tail -f /dev/null`` cannot be + rented. This template keeps the container alive with a single metachar-free + command; Lium's pod agent provides SSH. Idempotent: reused if it already + exists (the client's ``ensure_template`` then finds it by name at rent time). + """ + response = await client._request("GET", "/templates") + data = response.json() + templates = data.get("templates", []) if isinstance(data, dict) else data + for tmpl in templates: + if isinstance(tmpl, dict) and tmpl.get("name") == E2E_TEMPLATE_NAME: + return { + "name": E2E_TEMPLATE_NAME, + "id": str(tmpl.get("id")), + "image": tmpl.get("docker_image"), + "created": False, + } + body = { + "name": E2E_TEMPLATE_NAME, + "docker_image": E2E_IMAGE, + "docker_image_tag": E2E_IMAGE_TAG, + "startup_commands": E2E_STARTUP_COMMAND, + "internal_ports": [22], + "is_private": True, + "container_start_immediately": True, + } + created = await client._request("POST", "/templates", json_body=body) + payload = created.json() + return { + "name": E2E_TEMPLATE_NAME, + "id": str(payload.get("id")) if isinstance(payload, dict) else None, + "image": f"{E2E_IMAGE}:{E2E_IMAGE_TAG}", + "created": True, + } + + +async def _read_only_phase( + lium: LiumClient, targon: TargonClient, trace: dict[str, Any] +) -> float: + balance = await lium.balance() + offers = await lium.list_offers() + digest = await lium.watchtower_digest() + targon_offers = await targon.list_offers() + targon_apps = await targon.list_apps() + + if balance < 2: + raise SystemExit(f"Lium balance ${balance:.4f} < $2 minimum; aborting") + if not offers: + raise SystemExit("Lium GET /executors returned no offers") + if not re.match(r"^sha256:[0-9a-f]{64}$", digest): + raise SystemExit(f"watchtower digest not sha256-shaped: {digest[:16]}...") + if not targon_offers or any(o.price_per_hour <= 0 for o in targon_offers): + raise SystemExit("Targon inventory parse produced no positive-priced offers") + + trace["read_only"] = { + "lium_balance": balance, + "lium_offers_count": len(offers), + "lium_suitable_count": len(_suitable_offers(offers)), + "watchtower_digest": digest, + "targon_offers": [ + { + "shape": o.id, + "gpu": o.gpu_type, + "count": o.gpu_count, + "usd_hr": o.price_per_hour, + } + for o in targon_offers + ], + "targon_apps_count": len(targon_apps), + } + print( + f"[read-only] lium_balance=${balance:.4f} offers={len(offers)} " + f"digest_ok targon_offers={len(targon_offers)} targon_apps={len(targon_apps)}" + ) + return balance + + +async def _rental_phase( + lium: _TracingLiumClient, + ssh_public_key: str, + private_key: str, + balance_before: float, + trace: dict[str, Any], +) -> None: + template = await _ensure_e2e_template(lium) + offers = _suitable_offers(await lium.list_offers()) + if not offers: + raise SystemExit("no suitable Lium offer <= $1.50/GPU/hr") + + pod_name = f"mission-{int(time.time())}" + spec = InstanceSpec( + name=pod_name, + template_ref=template["name"], + ssh_public_keys=(ssh_public_key,), + ssh_key_name=SSH_KEY_NAME, + max_lifetime_hours=TERMINATION_HOURS, + max_price_per_hour=MAX_PRICE_PER_GPU_HR, + gpu_count=1, + ) + rental: dict[str, Any] = { + "pod_name": pod_name, + "template": template, + "attempts": [], + "status_transitions": [], + } + trace["rental"] = rental + + pod_id: str | None = None + instance = None + for offer in offers[:5]: + attempt = { + "executor_id": offer.id, + "gpu": offer.gpu_type, + "gpu_count": offer.gpu_count, + "price_per_gpu_hr": offer.price_per_hour, + } + try: + print( + f"[rent] trying {offer.gpu_type} ({offer.id}) " + f"@ ${offer.price_per_hour}/gpu/hr" + ) + instance = await lium.provision(spec, offer=offer) + pod_id = instance.id + attempt["result"] = "rented" + attempt["pod_id"] = pod_id + rental["attempts"].append(attempt) + rental["selected"] = attempt + break + except Exception as exc: # noqa: BLE001 - fallback to next-cheapest offer + attempt["result"] = f"failed: {type(exc).__name__}: {exc}" + rental["attempts"].append(attempt) + print(f"[rent] failed: {type(exc).__name__}: {exc}") + + rental["rent_response_shape"] = lium.rent_capture + + if pod_id is None or instance is None: + raise SystemExit("could not rent any suitable executor") + + already_deleted = False + try: + print(f"[rent] pod_id={pod_id} rented; polling to RUNNING") + deadline = time.monotonic() + POLL_BUDGET_SECONDS + last_status = "" + final_status = "" + while True: + current = await lium.status(pod_id) + status = (current.status or "").upper() + now = time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()) + if status != last_status: + rental["status_transitions"].append({"at": now, "status": status}) + print(f"[poll] {now} status={status}") + last_status = status + if status in RUNNING_STATUSES: + final_status = status + instance = current + break + if status in TERMINAL_FAIL_STATUSES: + raise SystemExit(f"pod reached terminal failure status {status}") + if time.monotonic() >= deadline: + raise SystemExit( + f"pod did not reach RUNNING within {POLL_BUDGET_SECONDS}s " + f"(last status {status})" + ) + await asyncio.sleep(POLL_INTERVAL_SECONDS) + + rental["running"] = True + raw = dict(instance.raw) + ssh_connect_cmd = str(raw.get("ssh_connect_cmd", "")) + target = _parse_ssh_target(ssh_connect_cmd, raw) + rental["ssh_target"] = {k: v for k, v in target.items() if k != "user" or v} + if not target.get("host") or not target.get("user"): + raise SystemExit(f"could not parse ssh target from: {ssh_connect_cmd!r}") + + target_desc = f"{target['user']}@{target['host']}:{target.get('port')}" + print(f"[ssh] connecting to {target_desc}") + ssh_result = _run_ssh_nvidia_smi(target, private_key) + rental["nvidia_smi"] = ssh_result + print(f"[ssh] returncode={ssh_result.get('returncode')}") + print(ssh_result.get("stdout", "")) + if ssh_result.get("returncode") != 0 or "SMOKE_OK" not in ssh_result.get( + "stdout", "" + ): + raise SystemExit(f"nvidia-smi over SSH failed: {ssh_result}") + + print("[logs] pulling pod logs") + logs = await _collect_logs(lium, pod_id) + rental["logs_excerpt"] = logs[:LOG_LINE_CAP] + rental["logs_line_count"] = len(logs) + print(f"[logs] collected {len(logs)} line(s)") + rental["final_status"] = final_status + finally: + if pod_id is not None: + print(f"[cleanup] DELETE pod {pod_id}") + await lium.terminate(pod_id) + gone = await lium.verify_terminated(pod_id) + rental["deleted"] = True + rental["verified_terminated"] = gone + already_deleted = gone + print(f"[cleanup] verify_terminated={gone}") + if not gone: + # Second confirmation attempt: a billing leak is critical. + await asyncio.sleep(5) + gone = await lium.verify_terminated(pod_id) + rental["verified_terminated"] = gone + print(f"[cleanup] second verify_terminated={gone}") + + balance_after = await lium.balance() + delta = balance_before - balance_after + pods_left = await lium.list_pods() + rental["balance_before"] = balance_before + rental["balance_after"] = balance_after + rental["balance_delta"] = delta + rental["pods_remaining"] = len(pods_left) + print( + f"[done] balance_before=${balance_before:.4f} after=${balance_after:.4f} " + f"delta=${delta:.4f} pods_remaining={len(pods_left)}" + ) + if not already_deleted: + raise SystemExit("pod not verified terminated; potential billing leak") + if len(pods_left) != 0: + raise SystemExit(f"{len(pods_left)} pod(s) still listed after cleanup") + if delta > 1.0: + raise SystemExit(f"balance delta ${delta:.4f} exceeds $1.00 guardrail") + + +async def _amain() -> int: + if os.environ.get("BASE_LIVE_PROVIDER_TESTS") != "1": + print("BASE_LIVE_PROVIDER_TESTS != 1; skipping live provider E2E") + return 0 + + creds = _load_credentials() + ssh_public_key = _load_ssh_public_key() + private_key = os.environ.get("LIUM_SSH_PRIVATE_KEY", DEFAULT_SSH_PRIVATE_KEY) + + trace: dict[str, Any] = { + "started_at": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()), + } + lium = _TracingLiumClient(creds["LIUM_API_KEY"]) + targon = TargonClient(creds["TARGON_API_KEY"]) + + exit_code = 0 + try: + balance_before = await _read_only_phase(lium, targon, trace) + await _rental_phase(lium, ssh_public_key, private_key, balance_before, trace) + trace["result"] = "success" + except SystemExit as exc: + trace["result"] = "failure" + trace["error"] = str(exc) + print(f"FAILURE: {exc}", file=sys.stderr) + exit_code = 1 + finally: + trace["finished_at"] = time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()) + out = Path(os.environ.get("LIUM_E2E_TRACE", "/tmp/lium_e2e_trace.json")) + out.write_text(json.dumps(trace, indent=2, default=str)) + print(f"trace written to {out}") + return exit_code + + +if __name__ == "__main__": + raise SystemExit(asyncio.run(_amain())) diff --git a/scripts/live_worker_lium_e2e.py b/scripts/live_worker_lium_e2e.py new file mode 100644 index 00000000..2b75783b --- /dev/null +++ b/scripts/live_worker_lium_e2e.py @@ -0,0 +1,930 @@ +#!/usr/bin/env python3 +"""Live Lium worker-plane end-to-end: real pod enrolls with a LOCAL master +(VAL-CROSS-005), opt-in and billable. + +Gated behind ``BASE_LIVE_PROVIDER_TESTS=1`` (real Lium money; AGENTS.md money +rules + architecture.md sec 3.1 cost guardrails). ONE batched Lium session +bounded to <= $2: + +1. bring up a LOCAL mission master (worker plane ON, mock metagraph) + a tiny + in-process fake challenge exposing exactly one gpu work unit; +2. provision a real Lium pod running the worker image via the REAL CLI + ``base worker deploy --provider lium`` (cheapest suitable offer, + ``termination_hours=1``); +3. reach RUNNING, install the torch-free ``bittensor_wallet`` keypair + copy the + mission ``base`` source + the in-pod agent runner into the pod over SSH; +4. run :mod:`pod_worker_agent` INSIDE the pod, reaching the local master over a + REVERSE SSH tunnel; the agent enrolls (miner-signed binding pre-signed on the + host), heartbeats to ``active``, pulls the gpu unit, executes the CPU stub, and + posts an ExecutionProof stamped with the LIUM provider + the REAL pod id; +5. verify enrollment (fleet ``GET /v1/workers``) + the forwarded ExecutionProof + (provider names lium + the real pod id, worker signature verifies); +6. DELETE the pod, verify it is gone, assert ``GET /pods`` empty + the balance + delta <= $2, and kill every local process by PID. + +The pod is terminated in a ``finally`` on EVERY path so a failure never leaks a +billable pod. A JSON trace is written to ``$LIVE_WORKER_TRACE`` (default +``/tmp/live_worker_e2e_trace.json``). Secrets are never printed. NOT for +production. +""" + +from __future__ import annotations + +import asyncio +import json +import os +import re +import subprocess +import sys +import threading +import time +import uuid +from dataclasses import dataclass, field +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from pathlib import Path +from typing import Any + +import bittensor as bt +import httpx + +from base.compute import LiumClient +from base.compute.lium import LiumError +from base.security.validator_auth import canonical_validator_request +from base.validator.agent.signing import KeypairRequestSigner +from base.worker.deploy import build_signed_binding +from base.worker.proof import ( + MANIFEST_SHA256_PAYLOAD_KEY, + PROOF_PAYLOAD_KEY, + execution_proof_signing_payload, +) + +# Reuse the proven VAL-PROV-014 helpers (credentials, ssh target parsing, ...). +SCRIPTS_DIR = Path(__file__).resolve().parent +sys.path.insert(0, str(SCRIPTS_DIR)) +import live_lium_e2e as prov # noqa: E402 + +REPO_ROOT = SCRIPTS_DIR.parent +BASE_SRC = REPO_ROOT / "src" +BASE_PY = REPO_ROOT / ".venv" / "bin" / "python" +MASTER_SCRIPT = SCRIPTS_DIR / "mission" / "mission_master.py" +POD_AGENT_SCRIPT = SCRIPTS_DIR / "mission" / "pod_worker_agent.py" + +MASTER_PORT = 3140 +FAKE_PORT = 3141 +MASTER_URL = f"http://127.0.0.1:{MASTER_PORT}" +FAKE_URL = f"http://127.0.0.1:{FAKE_PORT}" +POD_MASTER_PORT = 3400 # reverse-tunnel listener port INSIDE the pod (pod-local) +TOKEN = "mission-live-worker-token" +NETUID = 100 +WORKER_TTL = 120 +CHALLENGE_SLUG = "prism" + +OWNER_URI = "//LiumWorkerOwner" +SUBMITTER_URI = "//LiumSubmitter" +VALIDATOR_URI = "//LiumFleetValidator" +WORKER_URI = "//LiumPodWorker" + +MAX_PRICE_PER_GPU_HR = 1.50 +BALANCE_GUARDRAIL = 2.00 +POD_RUN_SECONDS = 300 +POD_POLL_BUDGET = 15 * 60 +POD_POLL_INTERVAL = 20.0 + +# The M1 placeholder worker image (WORKER_IMAGE = ghcr.io/.../prism-evaluator) is a +# large private-namespace GHCR image; Lium executors return CREATION_FAILED trying +# to pull it (confirmed live). The pod only needs to run the CPU worker agent, so +# this live check rents a small, reliably-pullable Docker Hub base image and +# installs the worker runtime in-pod. The deploy still runs through the real +# ``base worker deploy --provider lium`` CLI (which reuses this pre-created +# template), and the enroll/execute/prove/terminate cycle is unchanged. +POD_IMAGE = "python" +POD_IMAGE_TAG = "3.12-slim" +POD_PIP_DEPS = ( + "httpx pydantic pydantic-settings sqlalchemy PyYAML fastapi typer " + "docker bittensor-wallet" +) + + +def ss58(uri: str) -> str: + return str(bt.Keypair.create_from_uri(uri).ss58_address) + + +class RetryLiumClient(LiumClient): + """LiumClient that retries the throttled Lium edge (HTTP 429) with backoff. + + Lium's API edge enforces a strict short-window request cap; a burst of calls + (offer list + template + rent + status polls) can transiently return 429. Only + 429 is retried -- every other status (including a WAF 403) surfaces unchanged. + """ + + async def _request(self, method: str, path: str, **kwargs: Any) -> Any: + last: Exception | None = None + for attempt in range(5): + try: + return await super()._request(method, path, **kwargs) + except LiumError as exc: + if getattr(exc, "status_code", None) != 429: + raise + last = exc + await asyncio.sleep(6.0 * (attempt + 1)) + raise last if last else LiumError(f"Lium {method} {path} ret/exhausted") + + +# --------------------------------------------------------------------------- # +# In-process fake challenge (the master's bridged work source + result sink). # +# --------------------------------------------------------------------------- # +@dataclass +class ChallengeState: + unit: dict[str, Any] | None = None + results: list[dict[str, Any]] = field(default_factory=list) + lock: threading.Lock = field(default_factory=threading.Lock) + + def arm(self, submission_id: str, submission_ref: str) -> None: + with self.lock: + self.unit = { + "submission_id": submission_id, + "submission_ref": submission_ref, + } + + def disarm(self) -> None: + with self.lock: + self.unit = None + + def work_units(self) -> list[dict[str, Any]]: + with self.lock: + return [dict(self.unit)] if self.unit else [] + + def record_result(self, body: dict[str, Any]) -> None: + with self.lock: + self.results.append(body) + + def captured(self) -> list[dict[str, Any]]: + with self.lock: + return list(self.results) + + +_STATE = ChallengeState() + + +class _ChallengeHandler(BaseHTTPRequestHandler): + def log_message(self, *args: Any) -> None: # silence stdlib access log + return + + def _json(self, code: int, payload: dict[str, Any]) -> None: + body = json.dumps(payload).encode() + self.send_response(code) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(body))) + self.end_headers() + self.wfile.write(body) + + def do_GET(self) -> None: # noqa: N802 - stdlib API + if self.path == "/health": + self._json(200, {"status": "ok"}) + return + if self.path.rstrip("/") == "/internal/v1/work_units": + self._json(200, {"work_units": _STATE.work_units()}) + return + self._json(404, {"error": "not found"}) + + def do_POST(self) -> None: # noqa: N802 - stdlib API + length = int(self.headers.get("Content-Length", "0") or "0") + raw = self.rfile.read(length) if length else b"{}" + try: + body = json.loads(raw.decode() or "{}") + except ValueError: + body = {"_raw": raw.decode(errors="replace")} + path = self.path.rstrip("/") + if path == "/internal/v1/work_units/result": + _STATE.record_result(body) + self._json(200, {"status": "accepted"}) + return + if path == "/internal/v1/work_units/fold": + self._json(200, {"status": "folded"}) + return + self._json(404, {"error": "not found"}) + + +def _start_fake_challenge() -> ThreadingHTTPServer: + server = ThreadingHTTPServer(("127.0.0.1", FAKE_PORT), _ChallengeHandler) + thread = threading.Thread(target=server.serve_forever, daemon=True) + thread.start() + return server + + +# --------------------------------------------------------------------------- # +# Local master process. # +# --------------------------------------------------------------------------- # +def _master_config(workdir: Path) -> dict[str, Any]: + entries = [ + { + "hotkey": ss58(OWNER_URI), + "uid": 0, + "validator_permit": False, + "stake": 1000.0, + }, + { + "hotkey": ss58(SUBMITTER_URI), + "uid": 1, + "validator_permit": False, + "stake": 1000.0, + }, + { + "hotkey": ss58(VALIDATOR_URI), + "uid": 99, + "validator_permit": True, + "stake": 5000.0, + }, + ] + return { + "port": MASTER_PORT, + "host": "127.0.0.1", + "db_url": f"sqlite+aiosqlite:///{workdir / 'master.sqlite3'}", + "netuid": NETUID, + "metagraph": entries, + "prism": { + "slug": CHALLENGE_SLUG, + "internal_base_url": FAKE_URL, + "token": TOKEN, + }, + "orchestration_interval_seconds": 1.0, + "worker_heartbeat_ttl_seconds": WORKER_TTL, + "health_interval_seconds": 2.0, + "replication_factor": 1, + } + + +def _spawn_master(workdir: Path) -> subprocess.Popen: + cfg_path = workdir / "master.json" + cfg_path.write_text(json.dumps(_master_config(workdir), indent=2), encoding="utf-8") + log = (workdir / "master.log").open("w", encoding="utf-8") + env = dict(os.environ) + env["PYTHONPATH"] = str(BASE_SRC) + return subprocess.Popen( + [str(BASE_PY), str(MASTER_SCRIPT), str(cfg_path)], + stdout=log, + stderr=subprocess.STDOUT, + env=env, + cwd=str(REPO_ROOT), + ) + + +def _wait_health(url: str, name: str, *, timeout: float = 45.0) -> None: + deadline = time.time() + timeout + while time.time() < deadline: + try: + if httpx.get(f"{url}/health", timeout=3).status_code == 200: + print(f" {name} healthy at {url}") + return + except Exception: # noqa: BLE001 + pass + time.sleep(1.0) + raise TimeoutError(f"{name} did not become healthy at {url}") + + +# --------------------------------------------------------------------------- # +# Fleet reads (signed as the permitted mock validator). # +# --------------------------------------------------------------------------- # +def _validator_headers(method: str, path: str) -> dict[str, str]: + keypair = bt.Keypair.create_from_uri(VALIDATOR_URI) + nonce = uuid.uuid4().hex + ts = str(int(time.time())) + canonical = canonical_validator_request( + method=method, path=path, query_string="", timestamp=ts, nonce=nonce, body=b"" + ) + sig = keypair.sign(canonical.encode()) + sig_hex = ( + "0x" + bytes(sig).hex() if isinstance(sig, (bytes, bytearray)) else str(sig) + ) + return { + "X-Hotkey": ss58(VALIDATOR_URI), + "X-Signature": sig_hex, + "X-Nonce": nonce, + "X-Timestamp": ts, + } + + +def _fleet_workers() -> list[dict[str, Any]]: + resp = httpx.get( + f"{MASTER_URL}/v1/workers", + headers=_validator_headers("GET", "/v1/workers"), + timeout=10, + ) + resp.raise_for_status() + return resp.json().get("workers", []) + + +# --------------------------------------------------------------------------- # +# SSH helpers. # +# --------------------------------------------------------------------------- # +def _ssh_base(target: dict[str, Any], private_key: str) -> list[str]: + argv = [ + "-i", + private_key, + "-o", + "StrictHostKeyChecking=no", + "-o", + "UserKnownHostsFile=/dev/null", + "-o", + "ConnectTimeout=20", + "-o", + "BatchMode=yes", + "-o", + "ServerAliveInterval=15", + ] + if target.get("port"): + argv += ["-p", str(target["port"])] + return argv + + +def _ssh_run( + target: dict[str, Any], private_key: str, command: str, *, timeout: int = 240 +) -> subprocess.CompletedProcess: + argv = [ + "ssh", + *_ssh_base(target, private_key), + f"{target['user']}@{target['host']}", + command, + ] + return subprocess.run(argv, capture_output=True, text=True, timeout=timeout) + + +def _scp( + target: dict[str, Any], + private_key: str, + local: str, + remote: str, + *, + timeout: int = 240, +) -> subprocess.CompletedProcess: + argv = ["scp", *_ssh_base(target, private_key)] + # scp uses -P for port + if "-p" in argv: + idx = argv.index("-p") + argv[idx] = "-P" + argv += [local, f"{target['user']}@{target['host']}:{remote}"] + return subprocess.run(argv, capture_output=True, text=True, timeout=timeout) + + +def _detect_pod_python(target: dict[str, Any], private_key: str) -> str: + probe = _ssh_run( + target, + private_key, + "test -x /opt/venv/bin/python && echo /opt/venv/bin/python " + "|| command -v python3", + ) + out = probe.stdout.strip() + py = out.splitlines()[-1].strip() if out else "" + if not py: + raise SystemExit(f"could not locate python in pod: {probe.stderr[:300]}") + return py + + +# --------------------------------------------------------------------------- # +# Deploy via the real CLI. # +# --------------------------------------------------------------------------- # +def _write_deploy_config(workdir: Path, template_name: str, ssh_pub_file: str) -> Path: + cfg = f"""\ +network: + name: base + netuid: {NETUID} + chain_endpoint: null + wallet_name: default + wallet_hotkey: default + wallet_path: null + master_uid: 0 +compute: + worker_plane_enabled: true +worker: + agent: + master_url: {MASTER_URL} + gateway_url: null + capabilities: + - gpu + poll_interval_seconds: 2.0 + request_timeout_seconds: 15.0 + broker_url: http://127.0.0.1:8082 + broker_token_file: null + deploy: + provider: lium + gpu_count: 1 + max_price_per_hour: {MAX_PRICE_PER_GPU_HR} + max_lifetime_hours: 1.0 + startup_commands: tail -f /dev/null + ready_timeout_seconds: 60.0 + ssh_public_key_file: {ssh_pub_file} + ssh_key_name: {prov.SSH_KEY_NAME} + template_name: {template_name} + identity: + key_uri: {WORKER_URI} + key_mnemonic: null + wallet_name: null + wallet_hotkey: null + miner_key_uri: {OWNER_URI} + miner_key_mnemonic: null + miner_wallet_name: null + miner_wallet_hotkey: null + miner_hotkey: null + binding_signature: null + binding_nonce: null +docker: + broker_url: http://127.0.0.1:8082 + broker_allowed_images: + - ghcr.io/baseintelligence/ +observability: + log_json: false + sentry_dsn: null + otel_service_name: base-worker +""" + path = workdir / "deploy.yaml" + path.write_text(cfg, encoding="utf-8") + return path + + +async def _deploy_pod( + lium: LiumClient, + workdir: Path, + api_key: str, + ssh_pub_file: str, + trace: dict[str, Any], +) -> str: + template_name = f"prism-worker-live-{int(time.time())}" + cfg_path = _write_deploy_config(workdir, template_name, ssh_pub_file) + + # Pre-create the worker template with an EMPTY environment. The real CLI would + # otherwise bake the pod env (which carries a loopback master_url) into the + # template body, and the Lium API edge WAF blocks any request body containing + # an ``http://127.0.0.1`` URL (403 "Request blocked"). Pre-creating by name + # means the CLI's idempotent ensure_template finds it and skips that blocked + # POST; the pod runs ``tail -f /dev/null`` so the baked env is unused anyway. + template_id = await lium.ensure_template( + name=template_name, + docker_image=POD_IMAGE, + docker_image_tag=POD_IMAGE_TAG, + internal_ports=(22,), + startup_commands="tail -f /dev/null", + ) + trace["template"] = {"name": template_name, "id": template_id} + print(f"[deploy] pre-created template {template_name} id={template_id}") + await asyncio.sleep(6) # let the Lium edge rate-limit window cool down + + env = dict(os.environ) + env["PYTHONPATH"] = str(BASE_SRC) + env["LIUM_API_KEY"] = api_key + argv = [ + str(REPO_ROOT / ".venv" / "bin" / "base"), + "worker", + "deploy", + "--provider", + "lium", + "--max-price", + str(MAX_PRICE_PER_GPU_HR), + "--config", + str(cfg_path), + ] + print( + "[deploy] running real CLI: base worker deploy --provider lium " + f"--max-price {MAX_PRICE_PER_GPU_HR}" + ) + out = "" + pod_id = None + for attempt in range(1, 3): + proc = subprocess.run( + argv, + capture_output=True, + text=True, + env=env, + cwd=str(REPO_ROOT), + timeout=240, + ) + out = proc.stdout + "\n" + proc.stderr + trace.setdefault("deploy_cli_attempts", []).append( + { + "attempt": attempt, + "returncode": proc.returncode, + "stdout": proc.stdout.strip()[-1500:], + "stderr": proc.stderr.strip()[-1000:], + } + ) + print(out.strip()) + match = re.search(r"Provisioned\s+lium\s+instance\s+(\S+)", out) + if proc.returncode == 0 and match: + pod_id = match.group(1) + break + # a failed attempt cleans up its own partially-rented pod (LiumClient + # try/finally); guard anyway by terminating any leaked pod before retry. + leaked = await lium.list_pods() + for pod in leaked: + pid = str(pod.get("id")) + print(f"[deploy] terminating leaked pod {pid} before retry") + await lium.terminate(pid) + if attempt < 2: + print("[deploy] deploy attempt failed; retrying in 12s") + await asyncio.sleep(12) + if pod_id is None: + raise SystemExit("base worker deploy did not provision a pod") + trace["deploy_cli"] = {"template_name": template_name} + offer = re.search(r"Selected\s+lium\s+offer\s+(\S+)\s+\(([^)]*)\)", out) + if offer: + trace["deploy_cli"]["selected_offer"] = offer.group(1) + trace["deploy_cli"]["selected_offer_desc"] = offer.group(2) + print(f"[deploy] provisioned pod_id={pod_id}") + return pod_id + + +# --------------------------------------------------------------------------- # +# Proof verification. # +# --------------------------------------------------------------------------- # +def _verify_proof( + proof: dict[str, Any], pod_id: str, candidate_unit_ids: list[str] +) -> dict[str, Any]: + import hashlib + + checks: dict[str, Any] = {} + checks["version_is_1"] = proof.get("version") == 1 + checks["tier_ge_0"] = int(proof.get("tier", -1)) >= 0 + provider = proof.get("provider") or {} + checks["provider_name_lium"] = provider.get("name") == "lium" + checks["provider_pod_id_matches"] = str(provider.get("pod_id")) == str(pod_id) + sig = proof.get("worker_signature") or {} + worker_pubkey = sig.get("worker_pubkey") + checks["worker_pubkey_is_pod_worker"] = worker_pubkey == ss58(WORKER_URI) + manifest = proof.get(MANIFEST_SHA256_PAYLOAD_KEY) + checks["manifest_present"] = bool(manifest) + + matched_unit = None + sig_ok = False + for uid in candidate_unit_ids: + if not uid: + continue + if manifest and hashlib.sha256(uid.encode()).hexdigest() == manifest: + matched_unit = uid + try: + payload = execution_proof_signing_payload( + manifest_sha256=str(manifest), unit_id=uid + ) + kp = bt.Keypair(ss58_address=str(worker_pubkey)) + if kp.verify(payload, sig.get("sig")): + sig_ok = True + matched_unit = uid + break + except Exception: # noqa: BLE001 + continue + checks["worker_signature_verifies"] = sig_ok + checks["manifest_binds_unit"] = matched_unit is not None + checks["matched_unit_id"] = matched_unit + checks["all_passed"] = all( + v for k, v in checks.items() if k not in ("matched_unit_id",) + ) + return checks + + +# --------------------------------------------------------------------------- # +# Main sequence. # +# --------------------------------------------------------------------------- # +async def _amain() -> int: + if os.environ.get("BASE_LIVE_PROVIDER_TESTS") != "1": + print("BASE_LIVE_PROVIDER_TESTS != 1; skipping live worker E2E") + return 0 + + creds = prov._load_credentials() + api_key = creds["LIUM_API_KEY"] + prov._load_ssh_public_key() # validate the mission SSH public key exists + private_key = os.environ.get("LIUM_SSH_PRIVATE_KEY", prov.DEFAULT_SSH_PRIVATE_KEY) + ssh_pub_file = private_key + ".pub" + + workdir = Path(os.environ.get("LIVE_WORKER_WORKDIR", "/tmp/live-worker-run")) + workdir.mkdir(parents=True, exist_ok=True) + + trace: dict[str, Any] = { + "started_at": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()), + "pod_id": None, + } + + lium = RetryLiumClient(api_key) + balance_before = await lium.balance() + trace["balance_before"] = balance_before + print(f"[budget] Lium balance before: ${balance_before:.4f}") + if balance_before < BALANCE_GUARDRAIL: + raise SystemExit( + f"balance ${balance_before:.4f} < ${BALANCE_GUARDRAIL} minimum; aborting" + ) + + master_proc: subprocess.Popen | None = None + fake_server: ThreadingHTTPServer | None = None + agent_proc: subprocess.Popen | None = None + pod_id: str | None = None + already_deleted = False + exit_code = 0 + + try: + # 1. local control plane + print("== bring up local master + fake challenge ==") + fake_server = _start_fake_challenge() + _wait_health(FAKE_URL, "fake-challenge") + master_proc = _spawn_master(workdir) + _wait_health(MASTER_URL, "master") + + # 2. arm exactly one gpu unit (submitted by a DISTINCT owner from the worker) + submission_id = uuid.uuid4().hex + _STATE.arm(submission_id, ss58(SUBMITTER_URI)) + print( + f"[unit] armed gpu unit submission_id={submission_id} " + f"owner={ss58(SUBMITTER_URI)}" + ) + + # 3. provision the pod via the real CLI + pod_id = await _deploy_pod(lium, workdir, api_key, ssh_pub_file, trace) + trace["pod_id"] = pod_id + + # 4. poll to RUNNING + print("[pod] polling to RUNNING") + deadline = time.monotonic() + POD_POLL_BUDGET + instance = None + last_status = "" + while True: + current = await lium.status(pod_id) + status = (current.status or "").upper() + if status != last_status: + print(f"[pod] status={status}") + trace.setdefault("status_transitions", []).append(status) + last_status = status + if status in prov.RUNNING_STATUSES: + instance = current + break + if status in prov.TERMINAL_FAIL_STATUSES: + raise SystemExit(f"pod reached terminal status {status}") + if time.monotonic() >= deadline: + raise SystemExit(f"pod did not reach RUNNING (last={status})") + await asyncio.sleep(POD_POLL_INTERVAL) + + raw = dict(instance.raw) + target = prov._parse_ssh_target(str(raw.get("ssh_connect_cmd", "")), raw) + if not target.get("host") or not target.get("user"): + raise SystemExit( + f"could not parse ssh target: {raw.get('ssh_connect_cmd')!r}" + ) + trace["ssh_target"] = {k: v for k, v in target.items()} + print(f"[ssh] target {target['user']}@{target['host']}:{target.get('port')}") + + # 5. in-pod setup: python, keypair lib, base source, agent runner + config + # SSH may take a moment after RUNNING; retry the first probe. + pod_py = "" + for attempt in range(1, 9): + try: + pod_py = _detect_pod_python(target, private_key) + break + except Exception as exc: # noqa: BLE001 + print(f"[ssh] probe attempt {attempt} failed: {exc}") + time.sleep(10) + if not pod_py: + raise SystemExit("pod SSH never became reachable") + print(f"[pod] python={pod_py}") + + print("[pod] installing worker runtime deps") + pip = _ssh_run( + target, + private_key, + f"{pod_py} -m pip install --no-cache-dir {POD_PIP_DEPS}", + timeout=400, + ) + trace["pip_install"] = { + "rc": pip.returncode, + "tail": (pip.stdout + pip.stderr).strip()[-400:], + } + if pip.returncode != 0: + raise SystemExit( + f"pod dependency install failed: {(pip.stdout + pip.stderr)[-500:]}" + ) + + print("[pod] copying base source + agent runner") + tgz = workdir / "base-src.tgz" + subprocess.run( + ["tar", "-C", str(BASE_SRC), "-czf", str(tgz), "base"], check=True + ) + _ssh_run(target, private_key, "mkdir -p /tmp/base-src") + r1 = _scp(target, private_key, str(tgz), "/tmp/base-src.tgz") + if r1.returncode != 0: + raise SystemExit(f"scp base source failed: {r1.stderr[:300]}") + rx = _ssh_run( + target, + private_key, + "tar -C /tmp/base-src -xzf /tmp/base-src.tgz && echo OK", + ) + if "OK" not in rx.stdout: + raise SystemExit(f"extract base source failed: {rx.stderr[:300]}") + r2 = _scp( + target, private_key, str(POD_AGENT_SCRIPT), "/tmp/pod_worker_agent.py" + ) + if r2.returncode != 0: + raise SystemExit(f"scp agent runner failed: {r2.stderr[:300]}") + + # 6. pre-sign the miner binding on the host (pod never holds the miner key) + worker_pubkey = ss58(WORKER_URI) + miner_signer = KeypairRequestSigner(bt.Keypair.create_from_uri(OWNER_URI)) + binding = build_signed_binding( + worker_pubkey=worker_pubkey, miner_signer=miner_signer + ) + pod_cfg = { + "master_url": f"http://127.0.0.1:{POD_MASTER_PORT}", + "worker_uri": WORKER_URI, + "miner_hotkey": binding.miner_hotkey, + "binding_signature": binding.signature, + "binding_nonce": binding.nonce, + "provider": "lium", + "pod_id": pod_id, + "executor_id": pod_id, + "capabilities": ["gpu"], + "heartbeat_interval_seconds": 5, + "poll_interval_seconds": 2.0, + "run_seconds": POD_RUN_SECONDS, + } + cfg_local = workdir / "pod_cfg.json" + cfg_local.write_text(json.dumps(pod_cfg, indent=2), encoding="utf-8") + r3 = _scp(target, private_key, str(cfg_local), "/tmp/pod_cfg.json") + if r3.returncode != 0: + raise SystemExit(f"scp pod config failed: {r3.stderr[:300]}") + + # 7. run the in-pod agent over a reverse SSH tunnel to the local master + agent_log = (workdir / "pod-agent.log").open("w", encoding="utf-8") + run_cmd = ( + f"PYTHONPATH=/tmp/base-src {pod_py} -u " + "/tmp/pod_worker_agent.py /tmp/pod_cfg.json" + ) + agent_argv = [ + "ssh", + *_ssh_base(target, private_key), + "-o", + "ExitOnForwardFailure=yes", + "-R", + f"127.0.0.1:{POD_MASTER_PORT}:127.0.0.1:{MASTER_PORT}", + f"{target['user']}@{target['host']}", + run_cmd, + ] + print("[pod] launching in-pod worker agent over reverse tunnel") + agent_proc = subprocess.Popen( + agent_argv, stdout=agent_log, stderr=subprocess.STDOUT + ) + + # 8. observe enrollment (fleet active) + print("[observe] waiting for the pod worker to reach active in the fleet") + active_worker = None + deadline = time.time() + 120 + while time.time() < deadline: + try: + for w in _fleet_workers(): + if ( + w.get("provider_instance_ref") == pod_id + and w.get("status") == "active" + ): + active_worker = w + break + except Exception: # noqa: BLE001 + pass + if active_worker: + break + if agent_proc.poll() is not None: + raise SystemExit( + "in-pod agent exited before enrolling; see pod-agent.log" + ) + time.sleep(3) + if not active_worker: + raise SystemExit("pod worker did not reach active within timeout") + trace["active_worker"] = { + "worker_id": active_worker.get("worker_id"), + "worker_pubkey": active_worker.get("worker_pubkey"), + "miner_hotkey": active_worker.get("miner_hotkey"), + "provider": active_worker.get("provider"), + "provider_instance_ref": active_worker.get("provider_instance_ref"), + "status": active_worker.get("status"), + } + print(f"[observe] worker ACTIVE: {trace['active_worker']}") + + # 9. observe the forwarded ExecutionProof + print( + "[observe] waiting for the executed unit's ExecutionProof to be forwarded" + ) + forwarded = None + deadline = time.time() + 120 + while time.time() < deadline: + for body in _STATE.captured(): + result = body.get("result") or {} + if PROOF_PAYLOAD_KEY in result: + forwarded = body + break + if forwarded: + break + if agent_proc.poll() is not None: + # agent finished its run window; give a final read + for body in _STATE.captured(): + result = body.get("result") or {} + if PROOF_PAYLOAD_KEY in result: + forwarded = body + break + if forwarded: + break + time.sleep(3) + if not forwarded: + raise SystemExit("no ExecutionProof was forwarded for the executed unit") + _STATE.disarm() + proof = forwarded["result"][PROOF_PAYLOAD_KEY] + candidate_units = [str(forwarded.get("work_unit_id") or ""), submission_id] + checks = _verify_proof(proof, pod_id, candidate_units) + trace["forwarded_result"] = { + "work_unit_id": forwarded.get("work_unit_id"), + "submission_ref": forwarded.get("submission_ref"), + "proof": proof, + "checks": checks, + } + print(f"[proof] checks: {json.dumps(checks)}") + if not checks["all_passed"]: + raise SystemExit(f"ExecutionProof verification failed: {checks}") + print( + "[proof] ExecutionProof verified: lium provider + real pod id + " + "valid worker signature" + ) + + except SystemExit as exc: + trace["result"] = "failure" + trace["error"] = str(exc) + print(f"FAILURE: {exc}", file=sys.stderr) + exit_code = 1 + except Exception as exc: # noqa: BLE001 + trace["result"] = "failure" + trace["error"] = f"{type(exc).__name__}: {exc}" + print(f"FAILURE: {type(exc).__name__}: {exc}", file=sys.stderr) + exit_code = 1 + finally: + # stop the in-pod agent + reverse tunnel + if agent_proc is not None and agent_proc.poll() is None: + agent_proc.terminate() + try: + agent_proc.wait(timeout=10) + except Exception: # noqa: BLE001 + agent_proc.kill() + # DELETE the pod on EVERY path + if pod_id is not None: + print(f"[cleanup] DELETE pod {pod_id}") + try: + await lium.terminate(pod_id) + already_deleted = await lium.verify_terminated(pod_id) + if not already_deleted: + await asyncio.sleep(5) + already_deleted = await lium.verify_terminated(pod_id) + trace["verified_terminated"] = already_deleted + print(f"[cleanup] verified_terminated={already_deleted}") + except Exception as exc: # noqa: BLE001 + trace["cleanup_error"] = str(exc) + print(f"[cleanup] terminate error: {exc}", file=sys.stderr) + # tear down local processes by PID + if master_proc is not None and master_proc.poll() is None: + master_proc.terminate() + try: + master_proc.wait(timeout=8) + except Exception: # noqa: BLE001 + master_proc.kill() + if fake_server is not None: + fake_server.shutdown() + + # post-run accounting (best-effort; pod already deleted) + try: + balance_after = await lium.balance() + pods_left = await lium.list_pods() + except Exception as exc: # noqa: BLE001 + balance_after = None + pods_left = [] + trace["accounting_error"] = str(exc) + if balance_after is not None: + delta = balance_before - balance_after + trace["balance_after"] = balance_after + trace["balance_delta"] = delta + trace["pods_remaining"] = len(pods_left) + print( + f"[done] balance_before=${balance_before:.4f} after=${balance_after:.4f} " + f"delta=${delta:.4f} pods_remaining={len(pods_left)}" + ) + if len(pods_left) != 0: + print( + f"WARNING: {len(pods_left)} pod(s) still listed after cleanup", + file=sys.stderr, + ) + exit_code = 1 + if delta > BALANCE_GUARDRAIL: + print( + f"WARNING: balance delta ${delta:.4f} exceeds " + f"${BALANCE_GUARDRAIL} guardrail", + file=sys.stderr, + ) + exit_code = 1 + + if exit_code == 0 and trace.get("result") != "failure": + trace["result"] = "success" + trace["finished_at"] = time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()) + out = Path(os.environ.get("LIVE_WORKER_TRACE", "/tmp/live_worker_e2e_trace.json")) + out.write_text(json.dumps(trace, indent=2, default=str)) + print(f"trace written to {out}") + return exit_code + + +if __name__ == "__main__": + raise SystemExit(asyncio.run(_amain())) diff --git a/scripts/mission/mission_master.py b/scripts/mission/mission_master.py new file mode 100644 index 00000000..5c20b81f --- /dev/null +++ b/scripts/mission/mission_master.py @@ -0,0 +1,319 @@ +#!/usr/bin/env python +"""Local mission master service: mock-metagraph base master with the worker plane ON. + +Part of the cross-repo local end-to-end harness (see +``docs/operations/mission-harness.md``). Stands up a real base master proxy API on a +loopback port with: + +* a STATIC (no-chain) mock metagraph seeded from config -- owner miner hotkeys (no + permit) plus a stub validator hotkey (with a permit) so signed requests authenticate; +* the miner-funded GPU worker plane wired ON (register/heartbeat/pull/result + fleet + read); +* the validator coordination + assignment-coordination planes wired ON (so a stub + validator can be assigned + post audit results); +* the live orchestration driver bridging a single challenge's HTTP-exposed pending work + units into ``work_assignments``, running balanced worker assignment + reconciliation, + and forwarding accepted results back to the challenge. + +It is CONFIG-DRIVEN (a JSON file path in ``argv[1]`` or ``$MISSION_MASTER_CONFIG``) and +holds no real chain/provider secrets. NOT for production: production uses +``base master proxy`` with a real subtensor-backed metagraph. +""" + +from __future__ import annotations + +import asyncio +import atexit +import json +import logging +import sys +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +import uvicorn + +from base.bittensor.metagraph_cache import MetagraphCache +from base.db.base import Base +from base.db.session import create_engine, create_session_factory +from base.master.app_proxy import create_proxy_app +from base.master.assignment import CAPABILITY_GPU, AssignmentService +from base.master.assignment_coordination import AssignmentCoordinationService +from base.master.challenge_work_source import ( + HttpChallengeFoldTrigger, + HttpChallengeResultForwarder, + HttpChallengeWorkSource, +) +from base.master.orchestration import MasterOrchestrationDriver +from base.master.validator_coordination import ValidatorCoordinationService +from base.master.worker_assignment import WorkerAssignmentService +from base.master.worker_assignment_engine import WorkerAssignmentEngine +from base.master.worker_coordination import WorkerCoordinationService +from base.master.worker_reconciliation import WorkerReconciliationService +from base.master.worker_unit_status import WorkerUnitStatusService +from base.security.miner_auth import SqlAlchemyMinerNonceStore +from base.security.validator_auth import ( + MetagraphValidatorEligibility, + SqlAlchemyValidatorNonceStore, + ValidatorSignedRequestVerifier, +) +from base.security.worker_auth import ( + CoordinationReadEligibility, + MetagraphMinerMembership, + RegisteredWorkerEligibility, + SqlAlchemyWorkerNonceStore, + WorkerSignedRequestVerifier, +) + + +@dataclass(frozen=True) +class _ChallengeRecord: + """The minimal challenge descriptor the HTTP work-source seams read.""" + + slug: str + internal_base_url: str + status: Any = None + + +class _MockRegistry: + """A one-challenge registry mapping a slug to a challenge's internal URL + token.""" + + def __init__(self, *, slug: str, internal_base_url: str, token: str) -> None: + self._record = _ChallengeRecord(slug=slug, internal_base_url=internal_base_url) + self._token = token + + def list(self, *, active_only: bool = False) -> list[_ChallengeRecord]: + return [self._record] + + def get(self, slug: str) -> _ChallengeRecord: + if slug != self._record.slug: + raise KeyError(slug) + return self._record + + def get_token(self, slug: str) -> str: + return self._token + + +def _load_config() -> dict[str, Any]: + path = sys.argv[1] if len(sys.argv) > 1 else None + if path is None: + import os + + path = os.environ.get("MISSION_MASTER_CONFIG") + if not path: + raise SystemExit("usage: mission_master.py ") + return json.loads(Path(path).read_text(encoding="utf-8")) + + +async def _init_db(db_url: str) -> None: + engine = create_engine(db_url) + import base.db.models # noqa: F401 (register ORM tables on Base.metadata) + + async with engine.begin() as conn: + await conn.run_sync(Base.metadata.create_all) + await engine.dispose() + + +def _seed_metagraph(config: dict[str, Any]) -> MetagraphCache: + entries = config["metagraph"] + cache = MetagraphCache(netuid=int(config.get("netuid", 100)), static=True) + cache.update_from_metagraph( + [str(entry["hotkey"]) for entry in entries], + uids=[int(entry.get("uid", index)) for index, entry in enumerate(entries)], + validator_permits=[ + bool(entry.get("validator_permit", False)) for entry in entries + ], + stakes=[float(entry.get("stake", 0.0)) for entry in entries], + ) + return cache + + +def build_app(config: dict[str, Any]) -> Any: + db_url = config["db_url"] + engine = create_engine(db_url) + session_factory = create_session_factory(engine) + cache = _seed_metagraph(config) + + prism = config["prism"] + registry = _MockRegistry( + slug=prism["slug"], + internal_base_url=prism["internal_base_url"], + token=prism["token"], + ) + + worker_ttl = int(config.get("worker_heartbeat_ttl_seconds", 12)) + nonce_ttl = int(config.get("nonce_ttl_seconds", 300)) + sig_ttl = int(config.get("signature_ttl_seconds", 300)) + lease_seconds = int(config.get("assignment_lease_seconds", 900)) + + worker_service = WorkerCoordinationService( + session_factory, + miner_membership=MetagraphMinerMembership(cache), + binding_nonce_store=SqlAlchemyWorkerNonceStore( + session_factory, ttl_seconds=nonce_ttl + ), + heartbeat_ttl_seconds=worker_ttl, + ) + worker_verifier = WorkerSignedRequestVerifier( + nonce_store=SqlAlchemyWorkerNonceStore(session_factory, ttl_seconds=nonce_ttl), + eligibility=CoordinationReadEligibility(session_factory, cache), + ttl_seconds=sig_ttl, + ) + worker_assignment_service = WorkerAssignmentService( + session_factory, worker_service=worker_service, lease_seconds=lease_seconds + ) + worker_assignment_verifier = WorkerSignedRequestVerifier( + nonce_store=SqlAlchemyWorkerNonceStore(session_factory, ttl_seconds=nonce_ttl), + eligibility=RegisteredWorkerEligibility(session_factory), + ttl_seconds=sig_ttl, + ) + worker_unit_status_service = WorkerUnitStatusService(session_factory) + + validator_service = ValidatorCoordinationService( + session_factory, + heartbeat_interval_seconds=int( + config.get("validator_heartbeat_interval_seconds", 10) + ), + heartbeat_timeout_seconds=int( + config.get("validator_heartbeat_timeout_seconds", 30) + ), + ) + validator_verifier = ValidatorSignedRequestVerifier( + nonce_store=SqlAlchemyValidatorNonceStore( + session_factory, ttl_seconds=nonce_ttl + ), + eligibility=MetagraphValidatorEligibility(cache), + ttl_seconds=sig_ttl, + ) + assignment_service = AssignmentCoordinationService( + session_factory, lease_seconds=lease_seconds + ) + + driver = _build_orchestration_driver( + session_factory, + registry, + validator_service=validator_service, + worker_service=worker_service, + worker_assignment_service=worker_assignment_service, + replication_factor=int(config.get("replication_factor", 2)), + seed=config.get("orchestration_seed"), + worker_plane_enabled=bool(config.get("worker_plane_enabled", True)), + ) + + return create_proxy_app( + registry=registry, + metagraph_cache=cache, + nonce_store=SqlAlchemyMinerNonceStore(session_factory, ttl_seconds=nonce_ttl), + netuid=int(config.get("netuid", 100)), + validator_service=validator_service, + validator_verifier=validator_verifier, + validator_health_interval_seconds=float( + config.get("health_interval_seconds", 2.0) + ), + worker_service=worker_service, + worker_verifier=worker_verifier, + worker_health_interval_seconds=float( + config.get("health_interval_seconds", 2.0) + ), + worker_assignment_service=worker_assignment_service, + worker_assignment_verifier=worker_assignment_verifier, + worker_unit_status_service=worker_unit_status_service, + assignment_coordination_service=assignment_service, + orchestration_driver=driver, + orchestration_interval_seconds=float( + config.get("orchestration_interval_seconds", 1.0) + ), + ) + + +def _build_orchestration_driver( + session_factory: Any, + registry: Any, + *, + validator_service: ValidatorCoordinationService, + worker_service: WorkerCoordinationService, + worker_assignment_service: WorkerAssignmentService, + replication_factor: int, + seed: Any, + worker_plane_enabled: bool = True, +) -> MasterOrchestrationDriver: + # With the worker plane OFF (legacy) no capability is owned by the worker plane, so + # gpu units route to online gpu validators byte-for-byte as pre-mission + # (VAL-MASTER-013 / VAL-CROSS-006). + worker_plane_capabilities = ( + frozenset({CAPABILITY_GPU}) if worker_plane_enabled else frozenset() + ) + assignment_service = AssignmentService( + session_factory, worker_plane_capabilities=worker_plane_capabilities + ) + worker_engine = WorkerAssignmentEngine( + session_factory, + assignment_service=worker_assignment_service, + worker_service=worker_service, + replication_factor=replication_factor, + ) + worker_reconciler = WorkerReconciliationService( + session_factory, + result_forwarder=HttpChallengeResultForwarder(registry), + ) + return MasterOrchestrationDriver( + assignment_service=assignment_service, + validator_service=validator_service, + work_source=HttpChallengeWorkSource(registry), + fold_trigger=HttpChallengeFoldTrigger(registry), + worker_assignment_engine=worker_engine, + worker_reconciler=worker_reconciler, + seed=seed, + ) + + +def _configure_harness_logging(level: int = logging.INFO) -> None: + """Line-buffer stdout/stderr and route logs there so drill logs are + inspectable after teardown. + + The harness redirects each spawned process' stdout/stderr to a log file and + tears it down with SIGTERM; block-buffered output would be lost on kill, + leaving a 0-byte log. Line-buffering flushes every completed log line + immediately (so nothing already logged is lost), and an ``atexit`` flush + covers the graceful-shutdown path. + """ + + for stream in (sys.stdout, sys.stderr): + reconfigure = getattr(stream, "reconfigure", None) + if reconfigure is not None: + try: + reconfigure(line_buffering=True) + except (ValueError, OSError): + pass + logging.basicConfig( + level=level, + format="%(asctime)s %(levelname)s %(name)s %(message)s", + stream=sys.stdout, + force=True, + ) + atexit.register(_flush_streams) + + +def _flush_streams() -> None: + for stream in (sys.stdout, sys.stderr): + try: + stream.flush() + except (ValueError, OSError): + pass + + +def main() -> None: + _configure_harness_logging() + config = _load_config() + asyncio.run(_init_db(config["db_url"])) + app = build_app(config) + uvicorn.run( + app, + host=str(config.get("host", "127.0.0.1")), + port=int(config["port"]), + log_level=str(config.get("log_level", "info")), + ) + + +if __name__ == "__main__": + main() diff --git a/scripts/mission/no_external_egress.py b/scripts/mission/no_external_egress.py new file mode 100644 index 00000000..cb6bb9e3 --- /dev/null +++ b/scripts/mission/no_external_egress.py @@ -0,0 +1,133 @@ +"""Offline network-egress guard for the flags-OFF legacy regression (VAL-CROSS-006). + +Import this module as a pytest plugin (``-p no_external_egress``) to prove a test run +performs ZERO real network egress: any DNS resolution or socket connect to a +non-loopback host raises :class:`ExternalEgressBlocked` instead of touching the +network. Loopback (``127.0.0.0/8`` / ``::1`` / ``localhost``) and AF_UNIX sockets stay +allowed so the local test PostgreSQL (127.0.0.1:15433) and in-process fixtures keep +working, while a stray call to ``lium.io`` / ``api.targon.com`` (or anything else +external) fails loudly. + +Usage (run from the base repo, with the test PostgreSQL on 15433 already up):: + + PYTHONPATH="$PWD/scripts/mission" \ + BASE_TEST_DATABASE_URL=postgresql+asyncpg://base:base@localhost:15433/base_test \ + uv run pytest -p no_external_egress -q -p no:cacheprovider + +It installs from the ``pytest_configure`` hook, which runs before collection, so even +a test module that probes the network at import time is guarded. Importing the module +by itself has no side effects, and the guard is a no-op to install twice. NOT for +production; a mission verification harness only. +""" + +from __future__ import annotations + +import ipaddress +import socket + +_ALLOWED_HOSTNAMES = frozenset({"localhost", "localhost.localdomain", "ip6-localhost"}) + + +class ExternalEgressBlocked(OSError): + """A test attempted real network egress to a non-loopback host. + + Subclasses :class:`OSError` (not a bare ``RuntimeError``) so the standard + "is the network up?" probe -- ``try: socket.create_connection(...) except OSError`` + -- treats an external host as unreachable and SKIPS, exactly as it would in an + environment with external resolution disabled, instead of erroring collection. + """ + + +def _host_is_loopback(host: object) -> bool: + if host is None: + # AF_UNIX and empty binds carry no host; never external. + return True + text = str(host) + if text == "": + return True + if text in _ALLOWED_HOSTNAMES: + return True + try: + return ipaddress.ip_address(text).is_loopback + except ValueError: + # A DNS name that is not an explicit localhost alias: treat as external. + return False + + +def _guard_address(address: object) -> None: + # Only INET/INET6 addresses are (host, port[, ...]) tuples. AF_UNIX addresses are + # str/bytes filesystem paths (e.g. multiprocessing's /tmp/pymp-* sockets) and carry + # no host, so they are never external and must not be blocked. + if not (isinstance(address, (tuple, list)) and address): + return + host = address[0] + if not _host_is_loopback(host): + raise ExternalEgressBlocked( + f"blocked real network egress to {host!r} " + "(offline legacy regression: only loopback is permitted)" + ) + + +_REAL_GETADDRINFO = socket.getaddrinfo +_REAL_CREATE_CONNECTION = socket.create_connection +_REAL_SOCKET_CONNECT = socket.socket.connect +_REAL_SOCKET_CONNECT_EX = socket.socket.connect_ex + + +def _guarded_getaddrinfo(host, *args, **kwargs): # type: ignore[no-untyped-def] + if not _host_is_loopback(host): + raise ExternalEgressBlocked( + f"blocked DNS resolution for {host!r} " + "(offline legacy regression: only loopback is permitted)" + ) + return _REAL_GETADDRINFO(host, *args, **kwargs) + + +def _guarded_create_connection(address, *args, **kwargs): # type: ignore[no-untyped-def] + _guard_address(address) + return _REAL_CREATE_CONNECTION(address, *args, **kwargs) + + +def _guarded_socket_connect(self, address): # type: ignore[no-untyped-def] + _guard_address(address) + return _REAL_SOCKET_CONNECT(self, address) + + +def _guarded_socket_connect_ex(self, address): # type: ignore[no-untyped-def] + _guard_address(address) + return _REAL_SOCKET_CONNECT_EX(self, address) + + +def install() -> None: + """Block non-loopback egress via stdlib socket monkeypatches (idempotent).""" + + if getattr(socket, "_no_external_egress_installed", False): + return + socket.getaddrinfo = _guarded_getaddrinfo # type: ignore[assignment] + socket.create_connection = _guarded_create_connection # type: ignore[assignment] + socket.socket.connect = _guarded_socket_connect # type: ignore[method-assign] + socket.socket.connect_ex = _guarded_socket_connect_ex # type: ignore[method-assign] + socket._no_external_egress_installed = True # type: ignore[attr-defined] + + +def uninstall() -> None: + """Restore the original stdlib socket entry points (idempotent, for teardown).""" + + if not getattr(socket, "_no_external_egress_installed", False): + return + socket.getaddrinfo = _REAL_GETADDRINFO # type: ignore[assignment] + socket.create_connection = _REAL_CREATE_CONNECTION # type: ignore[assignment] + socket.socket.connect = _REAL_SOCKET_CONNECT # type: ignore[method-assign] + socket.socket.connect_ex = _REAL_SOCKET_CONNECT_EX # type: ignore[method-assign] + socket._no_external_egress_installed = False # type: ignore[attr-defined] + + +def pytest_configure(config: object) -> None: # noqa: ARG001 - pytest hook + """Install the guard before collection so the session performs zero external egress. + + ``pytest_configure`` runs before collection, so even a module that probes the + network at import time is guarded. Importing this module has NO side effects on its + own -- it guards only once loaded as a pytest plugin (``-p no_external_egress``). + """ + + install() diff --git a/scripts/mission/pod_worker_agent.py b/scripts/mission/pod_worker_agent.py new file mode 100644 index 00000000..4d10c8a3 --- /dev/null +++ b/scripts/mission/pod_worker_agent.py @@ -0,0 +1,137 @@ +#!/usr/bin/env python +"""In-pod worker agent runner for the live Lium end-to-end (VAL-CROSS-005). + +Runs INSIDE a rented Lium pod (the worker image) and enrolls the miner-funded +:class:`base.worker.runtime.WorkerAgent` with a LOCAL mission master reached over +a reverse SSH tunnel (``master_url`` = ``http://127.0.0.1:``). It reuses the +real worker plane end-to-end -- the worker keypair signs coordination requests + +the ExecutionProof, the miner-signed binding (pre-signed on the host so the pod +never holds the miner key) authenticates enrollment, and the CPU +:class:`StubManifestExecutor` executes each pulled gpu unit into a deterministic +manifest hash. Every emitted ExecutionProof carries the LIUM provider provenance +(provider name + the REAL pod id + the pinned image digest), so the master records +a proof that names the pod it ran in. + +CONFIG-DRIVEN (JSON path in ``argv[1]``). The signing keypair is provided by the +Rust-backed ``bittensor_wallet`` (no torch), falling back to ``bittensor`` when +present; both yield the same ss58 address + sr25519 signatures. NOT for +production. +""" + +from __future__ import annotations + +import asyncio +import json +import sys +from pathlib import Path +from typing import Any + +try: # light, torch-free keypair (Rust wheel) preferred inside the pod + from bittensor_wallet import Keypair +except ImportError: # pragma: no cover - host dry-run may only have full bittensor + from bittensor import Keypair # type: ignore[no-redef] + +from base.validator.agent import BrokerConfig +from base.validator.agent.signing import KeypairRequestSigner +from base.worker.coordination_client import WorkerCoordinationClient +from base.worker.executor import ( + StubManifestExecutor, + WorkerProofExecutor, + WorkerProvenance, +) +from base.worker.runtime import WorkerAgent, WorkerBinding + + +def _load_config() -> dict[str, Any]: + if len(sys.argv) < 2: + raise SystemExit("usage: pod_worker_agent.py ") + return json.loads(Path(sys.argv[1]).read_text(encoding="utf-8")) + + +def _build_agent(config: dict[str, Any]) -> WorkerAgent: + worker_keypair = Keypair.create_from_uri(config["worker_uri"]) + signer = KeypairRequestSigner(worker_keypair) + binding = WorkerBinding( + miner_hotkey=str(config["miner_hotkey"]), + signature=str(config["binding_signature"]), + nonce=str(config["binding_nonce"]), + ) + provenance = WorkerProvenance( + provider_name=str(config.get("provider", "lium")), + miner_hotkey=str(config["miner_hotkey"]), + executor_id=config.get("executor_id"), + pod_id=str(config["pod_id"]), + image_digest=config.get("image_digest"), + ) + executor = WorkerProofExecutor( + StubManifestExecutor(), signer=signer, provenance=provenance + ) + master_url = str(config["master_url"]) + client = WorkerCoordinationClient( + master_url, + signer, + timeout_seconds=float(config.get("request_timeout_seconds", 15.0)), + ) + return WorkerAgent( + client=client, + executor=executor, + broker=BrokerConfig( + broker_url=str(config.get("broker_url", "http://127.0.0.1:0")) + ), + binding=binding, + provider=str(config.get("provider", "lium")), + provider_instance_ref=str(config["pod_id"]), + capabilities=list(config.get("capabilities", ["gpu"])), + gateway_url=master_url, + heartbeat_interval_seconds=int(config.get("heartbeat_interval_seconds", 5)), + poll_interval_seconds=float(config.get("poll_interval_seconds", 2.0)), + ) + + +async def _run(config: dict[str, Any]) -> None: + agent = _build_agent(config) + print(f"[pod-agent] worker_pubkey={agent.worker_pubkey}", flush=True) + shutdown = asyncio.Event() + + run_seconds = float(config.get("run_seconds", 120.0)) + + async def _deadline() -> None: + await asyncio.sleep(run_seconds) + shutdown.set() + + worker_id = await agent.register(shutdown) + print(f"[pod-agent] registered worker_id={worker_id}", flush=True) + + async def _heartbeat() -> None: + while not shutdown.is_set(): + try: + await agent.heartbeat_once() + except Exception as exc: # noqa: BLE001 - keep heartbeating + print(f"[pod-agent] heartbeat error: {exc}", flush=True) + await asyncio.sleep(agent.heartbeat_interval) + + async def _work() -> None: + while not shutdown.is_set(): + try: + summary = await agent.process_pending_assignments() + if summary.pulled: + print( + f"[pod-agent] cycle pulled={summary.pulled} " + f"completed={summary.completed} failed={summary.failed}", + flush=True, + ) + except Exception as exc: # noqa: BLE001 - keep polling + print(f"[pod-agent] assignment error: {exc}", flush=True) + await asyncio.sleep(float(config.get("poll_interval_seconds", 2.0))) + + await asyncio.gather(_deadline(), _heartbeat(), _work()) + print("[pod-agent] shutdown", flush=True) + + +def main() -> None: + config = _load_config() + asyncio.run(_run(config)) + + +if __name__ == "__main__": + main() diff --git a/src/base/bittensor/factory.py b/src/base/bittensor/factory.py index 502b4703..79bfea5b 100644 --- a/src/base/bittensor/factory.py +++ b/src/base/bittensor/factory.py @@ -56,6 +56,82 @@ def create_validator_keypair(settings: Settings) -> Any: return _create_wallet(settings).hotkey +def _keypair_from_uri_or_mnemonic( + *, uri: str | None, mnemonic: str | None +) -> Any | None: + """Build an sr25519 keypair from a dev URI or a mnemonic (None if neither).""" + + bittensor = _load_bittensor() + if uri: + return bittensor.Keypair.create_from_uri(uri) + if mnemonic: + return bittensor.Keypair.create_from_mnemonic(mnemonic) + return None + + +def _wallet_keypair( + *, name: str | None, hotkey: str | None, path: str | None +) -> Any | None: + """Build a keypair from a bittensor wallet's hotkey (None if unconfigured).""" + + if not name and not hotkey and not path: + return None + bittensor = _load_bittensor() + wallet_kwargs: dict[str, Any] = {} + if name: + wallet_kwargs["name"] = name + if hotkey: + wallet_kwargs["hotkey"] = hotkey + if path: + wallet_kwargs["path"] = path + return bittensor.Wallet(**wallet_kwargs).hotkey + + +def create_worker_keypair(settings: Settings) -> Any: + """Return the WORKER keypair signing coordination requests + proofs. + + Resolves an sr25519 dev URI, then a mnemonic, then a dedicated bittensor + wallet, then falls back to ``network.wallet`` (the miner's own hotkey) so a + single-wallet local deploy works without extra key material. + """ + + identity = settings.worker.identity + keypair = _keypair_from_uri_or_mnemonic( + uri=identity.key_uri, mnemonic=identity.key_mnemonic + ) + if keypair is not None: + return keypair + keypair = _wallet_keypair( + name=identity.wallet_name, + hotkey=identity.wallet_hotkey, + path=identity.wallet_path, + ) + if keypair is not None: + return keypair + return _create_wallet(settings).hotkey + + +def create_worker_miner_keypair(settings: Settings) -> Any | None: + """Return the MINER keypair that signs the enrollment binding, or None. + + Resolves an sr25519 dev URI, then a mnemonic, then a bittensor wallet. Returns + None when no miner key is configured (the caller then requires a pre-signed + binding, e.g. a pod that never holds the miner key). + """ + + identity = settings.worker.identity + keypair = _keypair_from_uri_or_mnemonic( + uri=identity.miner_key_uri, mnemonic=identity.miner_key_mnemonic + ) + if keypair is not None: + return keypair + return _wallet_keypair( + name=identity.miner_wallet_name, + hotkey=identity.miner_wallet_hotkey, + path=identity.miner_wallet_path, + ) + + def _seed_mock_metagraph_cache(settings: Settings) -> MetagraphCache | None: """Seed a static ``MetagraphCache`` from ``network.mock_metagraph``. diff --git a/src/base/cli_app/main.py b/src/base/cli_app/main.py index 0d9373c3..e731b584 100644 --- a/src/base/cli_app/main.py +++ b/src/base/cli_app/main.py @@ -3,6 +3,7 @@ import asyncio import contextlib import logging +import uuid from collections.abc import Callable, Mapping from dataclasses import replace from datetime import UTC, datetime @@ -17,6 +18,8 @@ create_bittensor_runtime, create_bittensor_submit_runtime, create_validator_keypair, + create_worker_keypair, + create_worker_miner_keypair, ) from base.bittensor.identity_cache import ( IDENTITY_DISPLAY_NAME_KEY, @@ -25,12 +28,18 @@ ) from base.bittensor.metagraph_cache import MetagraphCache from base.bittensor.validator_loop import run_epoch_loop +from base.compute import ( + InstanceSpec, + LiumClient, + TargonClient, +) +from base.compute.worker_deployment import WORKER_TEMPLATE_NAME from base.config import load_settings from base.config.policy import production_policy_enabled_for_settings from base.config.settings import MasterSettings, Settings from base.db.session import create_engine, create_session_factory from base.master.app_proxy import create_proxy_app -from base.master.assignment import AssignmentService +from base.master.assignment import CAPABILITY_GPU, AssignmentService from base.master.assignment_coordination import ( AssignmentCoordinationService, GatewayPayloadIssuer, @@ -39,6 +48,7 @@ from base.master.challenge_client import ChallengeClient from base.master.challenge_work_source import ( HttpChallengeFoldTrigger, + HttpChallengeResultForwarder, HttpChallengeWorkSource, ) from base.master.docker_broker import create_docker_broker_app @@ -71,6 +81,11 @@ ) from base.master.swarm_backend import DEFAULT_JOB_NETWORK from base.master.validator_coordination import ValidatorCoordinationService +from base.master.worker_assignment import WorkerAssignmentService +from base.master.worker_assignment_engine import WorkerAssignmentEngine +from base.master.worker_coordination import WorkerCoordinationService +from base.master.worker_reconciliation import WorkerReconciliationService +from base.master.worker_unit_status import WorkerUnitStatusService from base.observability.logging import configure_logging from base.observability.otel import init_otel from base.observability.sentry import init_sentry @@ -87,6 +102,13 @@ SqlAlchemyValidatorNonceStore, ValidatorSignedRequestVerifier, ) +from base.security.worker_auth import ( + CoordinationReadEligibility, + MetagraphMinerMembership, + RegisteredWorkerEligibility, + SqlAlchemyWorkerNonceStore, + WorkerSignedRequestVerifier, +) from base.template_engine import ( ChallengeTemplateContext, render_challenge_template, @@ -103,6 +125,26 @@ from base.validator.registry_client import RegistryClient from base.validator.weight_submitter import ValidatorWeightSubmitter from base.validator.weights_client import WeightsClient +from base.worker import ( + WorkerAgent, + WorkerBinding, + WorkerCoordinationClient, + WorkerProofExecutor, + WorkerProvenance, + build_signed_binding, + normalize_provider, + plan_provider_deployment, + require_provider_api_key, +) +from base.worker.deploy import ( + LOCAL_PROVIDER, + MissingProviderKeyError, + NoOfferWithinBudgetError, + UnsupportedProviderError, + WorkerDeployError, + build_worker_pod_env, + require_worker_image, +) app = typer.Typer(help="BASE multi-challenge subnet CLI") master_app = typer.Typer(help="Run master components") @@ -112,6 +154,7 @@ db_app = typer.Typer(help="Database helpers") registry_app = typer.Typer(help="Registry helpers") worker_app = typer.Typer(help="Manage Swarm workers (CPU/GPU job nodes)") +worker_plane_app = typer.Typer(help="Deploy and manage miner-funded GPU worker agents") master_app.add_typer(master_challenges_app, name="challenges") master_app.add_typer(worker_app, name="worker") app.add_typer(master_app, name="master") @@ -119,6 +162,7 @@ app.add_typer(challenge_app, name="challenge") app.add_typer(db_app, name="db") app.add_typer(registry_app, name="registry") +app.add_typer(worker_plane_app, name="worker") PROJECT_ROOT = Path(__file__).resolve().parents[3] @@ -365,6 +409,68 @@ def _validator_signed_request_verifier( ) +def _worker_coordination_service( + settings: Any, session_factory: Any, metagraph_cache: MetagraphCache +) -> WorkerCoordinationService: + return WorkerCoordinationService( + session_factory, + miner_membership=MetagraphMinerMembership(metagraph_cache), + binding_nonce_store=SqlAlchemyWorkerNonceStore( + session_factory, + ttl_seconds=settings.compute.worker_nonce_ttl_seconds, + ), + heartbeat_ttl_seconds=settings.compute.worker_heartbeat_ttl_seconds, + ) + + +def _worker_signed_request_verifier( + settings: Any, + session_factory: Any, + metagraph_cache: MetagraphCache, +) -> WorkerSignedRequestVerifier: + return WorkerSignedRequestVerifier( + nonce_store=SqlAlchemyWorkerNonceStore( + session_factory, + ttl_seconds=settings.compute.worker_nonce_ttl_seconds, + ), + eligibility=CoordinationReadEligibility(session_factory, metagraph_cache), + ttl_seconds=settings.compute.worker_signature_ttl_seconds, + ) + + +def _worker_assignment_service( + settings: Any, + session_factory: Any, + worker_service: WorkerCoordinationService, +) -> WorkerAssignmentService: + return WorkerAssignmentService( + session_factory, + worker_service=worker_service, + lease_seconds=settings.master.assignment_lease_seconds, + ) + + +def _worker_assignment_verifier( + settings: Any, + session_factory: Any, +) -> WorkerSignedRequestVerifier: + """Verifier for worker pull/result: WORKER identity only, no permit. + + Distinct from the fleet-read verifier (which also admits validators): the + assignment surface requires a ``worker_registrations`` row so an unregistered + or validator-only key can never pull/post work (VAL-AGENT-018). + """ + + return WorkerSignedRequestVerifier( + nonce_store=SqlAlchemyWorkerNonceStore( + session_factory, + ttl_seconds=settings.compute.worker_nonce_ttl_seconds, + ), + eligibility=RegisteredWorkerEligibility(session_factory), + ttl_seconds=settings.compute.worker_signature_ttl_seconds, + ) + + def _assignment_coordination_service( settings: Any, session_factory: Any, @@ -401,6 +507,9 @@ def _master_orchestration_driver( session_factory: Any, registry: Any, validator_service: ValidatorCoordinationService, + *, + worker_service: WorkerCoordinationService | None = None, + worker_assignment_service: WorkerAssignmentService | None = None, ) -> MasterOrchestrationDriver: """Build the live master orchestration driver (architecture.md sec 4). @@ -408,9 +517,45 @@ def _master_orchestration_driver( ``work_assignments``, runs the balanced assignment + full reassignment pass, and folds retry-exhausted units back into their EvaluationJob via the challenge fold route. + + When the worker plane is enabled (``compute.worker_plane_enabled``, with both + worker services wired), gpu units are routed AWAY from validators (the + validator ``AssignmentService`` skips them via ``worker_plane_capabilities``) + and materialized as replicas by a :class:`WorkerAssignmentEngine` run each + pass. With the flag OFF the engine is ``None`` and gpu units route to + validators byte-identically to legacy. """ - assignment_service = AssignmentService(session_factory) + worker_plane_on = ( + settings.compute.worker_plane_enabled + and worker_service is not None + and worker_assignment_service is not None + ) + worker_plane_capabilities = ( + frozenset({CAPABILITY_GPU}) if worker_plane_on else frozenset() + ) + assignment_service = AssignmentService( + session_factory, worker_plane_capabilities=worker_plane_capabilities + ) + worker_engine: WorkerAssignmentEngine | None = None + worker_reconciler: WorkerReconciliationService | None = None + if worker_plane_on: + assert worker_service is not None + assert worker_assignment_service is not None + worker_engine = WorkerAssignmentEngine( + session_factory, + assignment_service=worker_assignment_service, + worker_service=worker_service, + replication_factor=settings.compute.replication_factor, + ) + worker_reconciler = WorkerReconciliationService( + session_factory, + result_forwarder=HttpChallengeResultForwarder( + registry, + timeout_seconds=settings.master.challenge_timeout_seconds, + retries=settings.master.challenge_retries, + ), + ) return MasterOrchestrationDriver( assignment_service=assignment_service, validator_service=validator_service, @@ -424,6 +569,8 @@ def _master_orchestration_driver( timeout_seconds=settings.master.challenge_timeout_seconds, retries=settings.master.challenge_retries, ), + worker_assignment_engine=worker_engine, + worker_reconciler=worker_reconciler, seed=settings.master.orchestration_seed, ) @@ -894,6 +1041,28 @@ def master_proxy(config: Path = typer.Option(Path("config/master.example.yaml")) validator_verifier = _validator_signed_request_verifier( settings, session_factory, runtime.metagraph_cache ) + # Miner-funded GPU worker plane (architecture.md sec 3.3), gated behind + # compute.worker_plane_enabled: OFF (default) leaves these None so the worker + # coordination surface is never mounted and legacy behavior is unchanged. + worker_service: WorkerCoordinationService | None = None + worker_verifier: WorkerSignedRequestVerifier | None = None + worker_assignment_service: WorkerAssignmentService | None = None + worker_assignment_verifier: WorkerSignedRequestVerifier | None = None + worker_unit_status_service: WorkerUnitStatusService | None = None + if settings.compute.worker_plane_enabled: + worker_service = _worker_coordination_service( + settings, session_factory, runtime.metagraph_cache + ) + worker_verifier = _worker_signed_request_verifier( + settings, session_factory, runtime.metagraph_cache + ) + worker_assignment_service = _worker_assignment_service( + settings, session_factory, worker_service + ) + worker_assignment_verifier = _worker_assignment_verifier( + settings, session_factory + ) + worker_unit_status_service = WorkerUnitStatusService(session_factory) llm_gateway_service = _llm_gateway_service(settings, session_factory) assignment_service = _assignment_coordination_service( settings, session_factory, gateway_service=llm_gateway_service @@ -902,7 +1071,12 @@ def master_proxy(config: Path = typer.Option(Path("config/master.example.yaml")) # work_assignments, runs balanced assignment + the full reassignment pass, # and folds retry-exhausted units, all on a Settings-driven interval. orchestration_driver = _master_orchestration_driver( - settings, session_factory, registry, validator_service + settings, + session_factory, + registry, + validator_service, + worker_service=worker_service, + worker_assignment_service=worker_assignment_service, ) # Registry-driven challenge deploy (architecture.md sec 4 + sec 9.2): the # master reconcile loop turns every ACTIVE registry challenge into a running @@ -935,6 +1109,14 @@ def master_proxy(config: Path = typer.Option(Path("config/master.example.yaml")) validator_health_interval_seconds=( settings.master.validator_health_interval_seconds ), + worker_service=worker_service, + worker_verifier=worker_verifier, + worker_health_interval_seconds=( + settings.compute.worker_health_interval_seconds + ), + worker_assignment_service=worker_assignment_service, + worker_assignment_verifier=worker_assignment_verifier, + worker_unit_status_service=worker_unit_status_service, assignment_coordination_service=assignment_service, llm_gateway_service=llm_gateway_service, orchestration_driver=orchestration_driver, @@ -1448,6 +1630,411 @@ def validator_subscribe( ) +# -- worker plane (base worker ...) ------------------------------------------ +# +# Top-level ``base worker`` app (DISTINCT from the legacy ``base master worker`` +# Swarm-node group). Deploys a miner-funded GPU worker agent locally or onto a +# rented provider instance, runs the agent loop, and renders fleet status. The +# provider API key comes from the MINER's env and is NEVER sent to the master. + + +def _worker_master_url(settings: Any) -> str: + url = settings.worker.agent.master_url + if not url: + raise typer.BadParameter( + "worker.agent.master_url must be set to reach the master coordination plane" + ) + return str(url) + + +def _build_worker_binding(settings: Any, worker_pubkey: str) -> WorkerBinding: + """Resolve the miner-signed enrollment binding for the worker. + + Prefers a pre-signed binding (miner_hotkey + signature + nonce) so a pod that + never holds the miner key can enroll; otherwise signs a fresh binding with the + configured miner keypair at deploy time. + """ + + identity = settings.worker.identity + if identity.binding_signature and identity.miner_hotkey and identity.binding_nonce: + return WorkerBinding( + miner_hotkey=str(identity.miner_hotkey), + signature=str(identity.binding_signature), + nonce=str(identity.binding_nonce), + ) + miner_keypair = create_worker_miner_keypair(settings) + if miner_keypair is None: + raise typer.BadParameter( + "worker deploy needs a miner key (worker.identity.miner_key_uri / " + "miner_key_mnemonic / miner_wallet_*) or a pre-signed binding " + "(worker.identity.miner_hotkey + binding_signature + binding_nonce)" + ) + return build_signed_binding( + worker_pubkey=worker_pubkey, + miner_signer=KeypairRequestSigner(miner_keypair), + ) + + +def _build_worker_agent( + settings: Any, + *, + provider: str, + provider_instance_ref: str | None = None, +) -> WorkerAgent: + """Wire the miner-funded worker agent from settings (testable). + + The agent signs coordination requests + ExecutionProofs with the WORKER + keypair, executes gpu units on its OWN broker via the shared executor seam, + and enrolls under a miner-signed binding. No provider API key is involved. + """ + + agent_cfg = settings.worker.agent + worker_keypair = create_worker_keypair(settings) + signer = KeypairRequestSigner(worker_keypair) + binding = _build_worker_binding(settings, signer.hotkey) + broker = BrokerConfig( + broker_url=agent_cfg.broker_url or settings.docker.broker_url, + broker_token=agent_cfg.broker_token, + broker_token_file=agent_cfg.broker_token_file, + allowed_images=tuple( + [*settings.docker.broker_allowed_images, *agent_cfg.allowed_images] + ), + ) + executor = WorkerProofExecutor( + ChallengeDispatchExecutor( + generic=BrokerAssignmentExecutor( + run_timeout_seconds=agent_cfg.run_timeout_seconds + ) + ), + signer=signer, + provenance=WorkerProvenance( + provider_name=provider, miner_hotkey=binding.miner_hotkey + ), + ) + master_url = _worker_master_url(settings) + client = WorkerCoordinationClient( + master_url, signer, timeout_seconds=agent_cfg.request_timeout_seconds + ) + return WorkerAgent( + client=client, + executor=executor, + broker=broker, + binding=binding, + provider=provider, + provider_instance_ref=( + provider_instance_ref or settings.worker.deploy.provider_instance_ref + ), + capabilities=list(agent_cfg.capabilities), + gateway_url=agent_cfg.gateway_url or master_url, + heartbeat_interval_seconds=agent_cfg.heartbeat_interval_seconds, + poll_interval_seconds=agent_cfg.poll_interval_seconds, + ) + + +def _spawn_worker_agent_process(config: Path) -> Any: + """Launch a detached ``base worker agent`` process (monkeypatchable seam).""" + + import subprocess + import sys + + return subprocess.Popen( # noqa: S603 - fixed argv, no shell + [ + sys.executable, + "-m", + "base.cli_app.main", + "worker", + "agent", + "--config", + str(config), + ], + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + ) + + +async def _wait_worker_active( + settings: Any, + signer: KeypairRequestSigner, + worker_pubkey: str, + *, + ready_timeout: float, + poll_interval: float = 3.0, +) -> Any: + """Poll ``GET /v1/workers`` until ``worker_pubkey`` is active or time out.""" + + import time as _time + + client = WorkerCoordinationClient( + _worker_master_url(settings), + signer, + timeout_seconds=settings.worker.agent.request_timeout_seconds, + ) + deadline = _time.monotonic() + ready_timeout + last_error: Exception | None = None + while _time.monotonic() < deadline: + try: + for worker in await client.list_workers(): + if worker.worker_pubkey == worker_pubkey and worker.status == "active": + return worker + except Exception as exc: # noqa: BLE001 - retried until the deadline + last_error = exc + await asyncio.sleep(poll_interval) + detail = f": {last_error}" if last_error is not None else "" + raise WorkerDeployError( + f"worker {worker_pubkey} did not reach active within " + f"{ready_timeout:.0f}s{detail}" + ) + + +def _run_worker_local_deploy(settings: Any, config: Path) -> None: + worker_keypair = create_worker_keypair(settings) + worker_pubkey = str(worker_keypair.ss58_address) + signer = KeypairRequestSigner(worker_keypair) + process = _spawn_worker_agent_process(config) + typer.echo(f"Started worker agent process pid={process.pid} (provider=local)") + ready_timeout = settings.worker.deploy.ready_timeout_seconds + try: + worker = asyncio.run( + _wait_worker_active( + settings, signer, worker_pubkey, ready_timeout=ready_timeout + ) + ) + except BaseException as exc: + with contextlib.suppress(Exception): + process.terminate() + typer.echo(str(exc), err=True) + raise typer.Exit(code=1) from exc + typer.echo( + f"Worker {worker.worker_id} active " + f"(pubkey={worker.worker_pubkey}, owner={worker.miner_hotkey}, " + f"provider={worker.provider})" + ) + + +def _worker_ssh_public_keys(deploy_cfg: Any) -> tuple[str, ...]: + if deploy_cfg.ssh_public_key: + return (str(deploy_cfg.ssh_public_key).strip(),) + if deploy_cfg.ssh_public_key_file: + text = Path(deploy_cfg.ssh_public_key_file).read_text(encoding="utf-8").strip() + return (text,) if text else () + return () + + +def _worker_instance_spec( + settings: Any, + *, + provider: str, + offer: Any, + env: Mapping[str, str], + max_price: float | None, + image: str, + image_digest: str, +) -> InstanceSpec: + deploy_cfg = settings.worker.deploy + return InstanceSpec( + name=f"prism-worker-{uuid.uuid4().hex[:12]}", + template_ref=deploy_cfg.template_name or WORKER_TEMPLATE_NAME, + image=image, + image_digest=image_digest, + env=dict(env), + ports=(22,), + ssh_public_keys=_worker_ssh_public_keys(deploy_cfg), + ssh_key_name=deploy_cfg.ssh_key_name, + startup_commands=deploy_cfg.startup_commands, + max_lifetime_hours=deploy_cfg.max_lifetime_hours, + max_price_per_hour=( + max_price if max_price is not None else float(offer.price_per_hour) + ), + gpu_count=deploy_cfg.gpu_count, + ) + + +async def _run_worker_provider_deploy_async( + settings: Any, + *, + provider: str, + api_key: str, + max_price: float | None, +) -> None: + deploy_cfg = settings.worker.deploy + # Fail fast (before any provider network call) when the deploy has no explicit, + # publicly-pullable, digest-pinned worker image: the deploy never silently pins + # a private-namespace placeholder that fails Lium pod creation. + image, image_digest = require_worker_image( + image=deploy_cfg.image, + image_digest=deploy_cfg.image_digest, + provider=provider, + ) + resolved_max_price = ( + max_price if max_price is not None else deploy_cfg.max_price_per_hour + ) + client: Any = LiumClient(api_key) if provider == "lium" else TargonClient(api_key) + # Planning is offer selection only: no rent/deploy call is issued here, and an + # all-over-cap situation raises before anything is provisioned. + offer = await plan_provider_deployment( + client, gpu_count=deploy_cfg.gpu_count, max_price=resolved_max_price + ) + typer.echo( + f"Selected {provider} offer {offer.id} " + f"({offer.gpu_type} x{offer.gpu_count}) @ {offer.price_per_hour}/GPU/hr" + ) + worker_keypair = create_worker_keypair(settings) + signer = KeypairRequestSigner(worker_keypair) + binding = _build_worker_binding(settings, signer.hotkey) + pod_env = build_worker_pod_env( + master_url=_worker_master_url(settings), + provider=provider, + binding=binding, + worker_key_uri=settings.worker.identity.key_uri, + worker_key_mnemonic=settings.worker.identity.key_mnemonic, + broker_url=settings.worker.agent.broker_url, + gateway_url=settings.worker.agent.gateway_url, + ) + spec = _worker_instance_spec( + settings, + provider=provider, + offer=offer, + env=pod_env, + max_price=resolved_max_price, + image=image, + image_digest=image_digest, + ) + if provider == "lium": + instance = await client.provision(spec, offer=offer) + else: + instance = await client.provision(spec) + typer.echo( + f"Provisioned {provider} instance {instance.id} (status={instance.status}); " + "the worker enrolls with the master on boot" + ) + + +@worker_plane_app.command("agent") +def worker_agent( + config: Path = typer.Option(Path("config/worker.example.yaml")), +): + """Run the miner-funded GPU worker agent loop. + + Registers with the master under a miner-signed binding, heartbeats to stay + active, pulls gpu work units, executes them on the worker's OWN broker, and + posts ExecutionProof-carrying results. Authenticates as the worker keypair and + never holds a provider API key. + """ + + settings = load_settings(config) + _configure_observability(settings) + agent = _build_worker_agent(settings, provider=settings.worker.deploy.provider) + typer.echo(f"Starting worker agent for pubkey {agent.worker_pubkey}") + asyncio.run(agent.run_forever()) + + +@worker_plane_app.command("deploy") +def worker_deploy( + provider: str = typer.Option( + ..., + "--provider", + help="Where to run the worker agent: lium | targon | local.", + ), + max_price: float | None = typer.Option( + None, + "--max-price", + help="Max price per GPU/hour bounding provider offer selection.", + ), + config: Path = typer.Option(Path("config/worker.example.yaml")), +): + """Deploy a worker agent locally or onto a rented provider instance. + + ``--provider local`` starts an agent against the local master. ``--provider + lium|targon`` requires the provider API key env (``LIUM_API_KEY`` / + ``TARGON_API_KEY``), selects an offer within ``--max-price`` (preferring an + exact GPU-count executor), and provisions the worker image. The provider key + is used only to authenticate provider calls and is NEVER sent to the master. + """ + + try: + normalized = normalize_provider(provider) + except UnsupportedProviderError as exc: + typer.echo(str(exc), err=True) + raise typer.Exit(code=2) from exc + + # Enforce the provider key BEFORE loading config or touching the network so a + # missing key is an actionable refusal with no side effects (VAL-AGENT-010). + api_key: str | None = None + if normalized != LOCAL_PROVIDER: + try: + api_key = require_provider_api_key(normalized) + except MissingProviderKeyError as exc: + typer.echo(str(exc), err=True) + raise typer.Exit(code=2) from exc + + settings = load_settings(config) + _configure_observability(settings) + + if normalized == LOCAL_PROVIDER: + _run_worker_local_deploy(settings, config) + return + + assert api_key is not None + try: + asyncio.run( + _run_worker_provider_deploy_async( + settings, provider=normalized, api_key=api_key, max_price=max_price + ) + ) + except NoOfferWithinBudgetError as exc: + typer.echo(str(exc), err=True) + raise typer.Exit(code=1) from exc + except WorkerDeployError as exc: + typer.echo(str(exc), err=True) + raise typer.Exit(code=1) from exc + + +@worker_plane_app.command("status") +def worker_status( + config: Path = typer.Option(Path("config/worker.example.yaml")), + hotkey: str | None = typer.Option( + None, + "--hotkey", + help="Show only the ACTIVE workers of this miner hotkey.", + ), +): + """Render the worker fleet from the master's ``GET /v1/workers``. + + Reads the fleet as the worker keypair (a registered worker is an eligible + coordination reader) and prints each worker's status, owner, provider, + last-seen timestamp, and attributed fault count (so the CLI fleet view agrees + with the ``GET /v1/workers`` JSON, VAL-CROSS-009). + """ + + settings = load_settings(config) + _configure_observability(settings) + signer = KeypairRequestSigner(create_worker_keypair(settings)) + client = WorkerCoordinationClient( + _worker_master_url(settings), + signer, + timeout_seconds=settings.worker.agent.request_timeout_seconds, + ) + workers = asyncio.run(client.list_workers(hotkey=hotkey)) + if not workers: + typer.echo("No workers registered.") + return + typer.echo( + f"{'WORKER_ID':<20} {'OWNER':<20} {'PROVIDER':<10} {'STATUS':<8} " + f"{'FAULTS':<6} LAST_SEEN" + ) + for worker in workers: + last_seen = ( + worker.last_heartbeat_at.isoformat() + if worker.last_heartbeat_at is not None + else "-" + ) + fault_count = len(getattr(worker, "faults", None) or []) + typer.echo( + f"{worker.worker_id:<20} {worker.miner_hotkey:<20} " + f"{worker.provider:<10} {worker.status:<8} {fault_count:<6} {last_seen}" + ) + + @challenge_app.command("create") def challenge_create( slug: str, diff --git a/src/base/compute/__init__.py b/src/base/compute/__init__.py new file mode 100644 index 00000000..a2c987b2 --- /dev/null +++ b/src/base/compute/__init__.py @@ -0,0 +1,64 @@ +"""Compute provider clients for the miner-funded GPU worker plane. + +Provider-agnostic contract (:mod:`base.compute.provider`) plus concrete clients +(Lium; Targon added alongside) used to provision, inspect, and tear down the GPU +instances that run worker agents, under mandatory cost guardrails. +""" + +from __future__ import annotations + +from base.compute.lium import LIUM_API_BASE_URL, LiumClient, LiumError +from base.compute.provider import ( + CostGuardrailError, + Instance, + InstanceSpec, + Offer, + ProviderClient, + ProviderError, +) +from base.compute.targon import ( + TARGON_API_BASE_URL, + BalanceUnavailableError, + InsufficientCreditsError, + TargonClient, + TargonError, +) +from base.compute.worker_deployment import ( + WORKER_IMAGE, + WORKER_IMAGE_DIGEST, + WORKER_IMAGE_TAG, + WORKER_STARTUP_COMMANDS, + build_lium_worker_template, + build_targon_worker_app, + is_loopback_url, + is_metachar_free, + is_pinned_digest, + pinned_image_reference, +) + +__all__ = [ + "LIUM_API_BASE_URL", + "TARGON_API_BASE_URL", + "WORKER_IMAGE", + "WORKER_IMAGE_DIGEST", + "WORKER_IMAGE_TAG", + "WORKER_STARTUP_COMMANDS", + "BalanceUnavailableError", + "CostGuardrailError", + "Instance", + "InstanceSpec", + "InsufficientCreditsError", + "LiumClient", + "LiumError", + "Offer", + "ProviderClient", + "ProviderError", + "TargonClient", + "TargonError", + "build_lium_worker_template", + "build_targon_worker_app", + "is_loopback_url", + "is_metachar_free", + "is_pinned_digest", + "pinned_image_reference", +] diff --git a/src/base/compute/lium.py b/src/base/compute/lium.py new file mode 100644 index 00000000..9dd7b0d0 --- /dev/null +++ b/src/base/compute/lium.py @@ -0,0 +1,486 @@ +"""Lium compute provider client (architecture.md sec 3.1). + +Thin async ``httpx`` wrapper over the Lium REST API (base ``https://lium.io/api``, +auth header ``X-API-Key``). It implements the :class:`~base.compute.provider` +contract with the cost guardrails baked in: + +* :meth:`LiumClient.provision` refuses an unbounded or over-priced spec BEFORE any + network call and always sends a bounded ``termination_hours``. +* Any failure after the pod is rented terminates + verifies the pod (try/finally), + so a failed provisioning never leaks a billable pod. +* :meth:`LiumClient.terminate` is idempotent (a ``404`` delete is success) and + :meth:`LiumClient.verify_terminated` reflects real pod absence via ``GET /pods``. + +The API key lives only in the request header; it is never logged, embedded in an +error message, or exposed via ``repr``. +""" + +from __future__ import annotations + +import logging +from collections.abc import AsyncIterator, Mapping, Sequence +from typing import Any + +import httpx + +from base.compute.provider import ( + CostGuardrailError, + Instance, + InstanceSpec, + Offer, + ProviderError, +) +from base.compute.worker_deployment import is_loopback_url + +logger = logging.getLogger(__name__) + +LIUM_API_BASE_URL = "https://lium.io/api" +_DEFAULT_SSH_KEY_NAME = "prism-mission-worker" + + +class LiumError(ProviderError): + """A Lium API request failed (non-2xx response or transport error).""" + + def __init__(self, message: str, *, status_code: int | None = None) -> None: + super().__init__(message) + self.status_code = status_code + + +class LiumClient: + """Async client for the Lium GPU rental API.""" + + def __init__( + self, + api_key: str, + *, + base_url: str = LIUM_API_BASE_URL, + transport: httpx.AsyncBaseTransport | None = None, + timeout_seconds: float = 30.0, + ) -> None: + self._api_key = api_key + self._base_url = base_url.rstrip("/") + self._transport = transport + self._timeout = timeout_seconds + + def __repr__(self) -> str: + return f"LiumClient(base_url={self._base_url!r})" + + # -- offers --------------------------------------------------------------- + + async def list_offers( + self, *, max_price_per_hour: float | None = None + ) -> list[Offer]: + response = await self._request("GET", "/executors") + offers: list[Offer] = [] + for item in _as_list(response.json(), "executors"): + offer = _parse_offer(item) + if offer is None: + continue + if ( + max_price_per_hour is not None + and offer.price_per_hour > max_price_per_hour + ): + continue + offers.append(offer) + return offers + + # -- lifecycle ------------------------------------------------------------ + + async def provision( + self, spec: InstanceSpec, *, offer: Offer | None = None + ) -> Instance: + lifetime = spec.max_lifetime_hours + if lifetime is None or lifetime <= 0: + raise CostGuardrailError( + "InstanceSpec.max_lifetime_hours must be a positive bound (hours)" + ) + if lifetime < 1: + raise CostGuardrailError( + "InstanceSpec.max_lifetime_hours must be at least 1 hour: Lium " + "termination_hours has 1-hour granularity, so a sub-hour bound " + "would truncate to termination_hours=0 and disable auto-termination" + ) + if spec.max_price_per_hour is None or spec.max_price_per_hour <= 0: + raise CostGuardrailError( + "InstanceSpec.max_price_per_hour must be a positive bound" + ) + if not spec.ssh_public_keys: + raise LiumError("Lium rent requires at least one SSH public key") + + selected = await self._resolve_offer(spec, offer) + + for public_key in spec.ssh_public_keys: + await self.ensure_ssh_key( + public_key=public_key, + name=spec.ssh_key_name or _DEFAULT_SSH_KEY_NAME, + ) + template_id = await self._resolve_template(spec) + + rent_body: dict[str, Any] = { + "pod_name": spec.name, + "user_public_key": list(spec.ssh_public_keys), + "termination_hours": int(lifetime), + "gpu_count": spec.gpu_count, + } + if template_id is not None: + rent_body["template_id"] = template_id + if spec.dockerfile_content is not None: + rent_body["dockerfile_content"] = spec.dockerfile_content + + rent = await self._request( + "POST", f"/executors/{selected.id}/rent", json_body=rent_body + ) + # The rent succeeded: a billable pod may now exist. Every subsequent + # failure -- an unparseable rent body, pod-id resolution, AND the status + # poll -- must best-effort terminate + verify before re-raising, so + # cleanup keys strictly off "rent HTTP call succeeded" rather than "pod + # id resolved". The pod-id extraction stays INSIDE the try so a 2xx rent + # with a non-JSON body cannot leak a just-rented pod. + pod_id: str | None = None + try: + pod_id = self._extract_pod_id(rent) + if pod_id is None: + pod_id = await self._find_pod_id(spec.name) + if pod_id is None: + raise LiumError( + "could not determine provisioned pod id from rent response" + ) + return await self.status(pod_id) + except BaseException: + await self._cleanup_after_rent(pod_id, spec.name) + raise + + async def status(self, instance_id: str) -> Instance: + response = await self._request("GET", f"/pods/{instance_id}") + return _parse_instance(response.json()) + + async def stream_logs(self, instance_id: str) -> AsyncIterator[str]: + async with self._client() as client: + async with client.stream("GET", f"/pods/{instance_id}/logs") as response: + if response.status_code >= 400: + await response.aread() + raise LiumError( + f"Lium GET /pods/{instance_id}/logs returned " + f"{response.status_code}", + status_code=response.status_code, + ) + async for line in response.aiter_lines(): + yield line + + async def terminate(self, instance_id: str) -> None: + response = await self._send("DELETE", f"/pods/{instance_id}") + if response.status_code == 404: + return + if response.status_code >= 400: + raise LiumError( + f"Lium DELETE /pods/{instance_id} returned {response.status_code}", + status_code=response.status_code, + ) + + async def verify_terminated(self, instance_id: str) -> bool: + for pod in await self.list_pods(): + if str(pod.get("id")) == str(instance_id): + return False + return True + + async def list_pods(self) -> list[dict[str, Any]]: + response = await self._request("GET", "/pods") + return _as_list(response.json(), "pods") + + # -- idempotent deploy helpers ------------------------------------------- + + async def ensure_ssh_key( + self, *, public_key: str, name: str | None = None + ) -> dict[str, Any]: + response = await self._request("GET", "/ssh-keys") + normalized = public_key.strip() + for key in _as_list(response.json(), "ssh_keys"): + if str(key.get("public_key", "")).strip() == normalized: + return key + body: dict[str, Any] = {"public_key": public_key} + if name is not None: + body["name"] = name + created = await self._request("POST", "/ssh-keys", json_body=body) + result = created.json() + return result if isinstance(result, dict) else {} + + async def ensure_template( + self, + *, + name: str, + docker_image: str, + docker_image_digest: str | None = None, + docker_image_tag: str | None = None, + internal_ports: Sequence[int] = (22,), + environment: Mapping[str, str] | None = None, + startup_commands: str | None = None, + is_private: bool = True, + container_start_immediately: bool = True, + ) -> str: + response = await self._request("GET", "/templates") + for template in _as_list(response.json(), "templates"): + if str(template.get("name")) == name and template.get("id"): + return str(template["id"]) + body: dict[str, Any] = { + "name": name, + "docker_image": docker_image, + "internal_ports": list(internal_ports), + "is_private": is_private, + "container_start_immediately": container_start_immediately, + } + if docker_image_digest: + body["docker_image_digest"] = docker_image_digest + if docker_image_tag: + body["docker_image_tag"] = docker_image_tag + if environment: + # Lium's edge WAF returns 403 "Request blocked" for ANY body carrying a + # loopback URL (http://127.0.0.1... / http://localhost...), so strip + # loopback-valued entries before POSTing the template. Such values are + # redundant anyway: the pod config defaults them to loopback and the + # agent resolves them at runtime (a reverse SSH tunnel reaches a local + # master). Non-loopback (real) values are preserved. + safe_environment = { + key: value + for key, value in environment.items() + if not is_loopback_url(value) + } + if safe_environment: + body["environment"] = safe_environment + # Lium rejects rents whose template startup_commands contain shell + # metacharacters ("Malicious startup command detected"), so a caller + # supplies a single metachar-free keep-alive here (e.g. "tail -f + # /dev/null"); omitted when None to preserve the image's own entrypoint. + if startup_commands is not None: + body["startup_commands"] = startup_commands + created = await self._request("POST", "/templates", json_body=body) + return str(created.json().get("id")) + + # -- account / proof inputs ---------------------------------------------- + + async def watchtower_digest(self) -> str: + response = await self._request("GET", "/watchtower/digest") + data = response.json() + if isinstance(data, Mapping): + digest = data.get("digest") + if digest: + return str(digest) + if isinstance(data, str) and data: + return data + raise LiumError("watchtower digest response missing 'digest'") + + async def balance(self) -> float: + response = await self._request("GET", "/users/me") + data = response.json() + balance = data.get("balance") if isinstance(data, Mapping) else None + if balance is None: + raise LiumError("users/me response missing 'balance'") + return float(balance) + + # -- internals ------------------------------------------------------------ + + async def _resolve_offer(self, spec: InstanceSpec, offer: Offer | None) -> Offer: + if offer is not None: + if ( + spec.max_price_per_hour is not None + and offer.price_per_hour > spec.max_price_per_hour + ): + raise CostGuardrailError( + f"offer {offer.id} at {offer.price_per_hour}/hr exceeds " + f"max_price_per_hour {spec.max_price_per_hour}" + ) + return offer + offers = await self.list_offers(max_price_per_hour=spec.max_price_per_hour) + if not offers: + raise CostGuardrailError( + "no Lium offer available within max_price_per_hour bound" + ) + return min(offers, key=lambda candidate: candidate.price_per_hour) + + async def _resolve_template(self, spec: InstanceSpec) -> str | None: + if spec.template_ref is not None: + return await self.ensure_template( + name=spec.template_ref, + docker_image=spec.image or "", + docker_image_digest=spec.image_digest, + internal_ports=spec.ports or (22,), + environment=spec.env, + startup_commands=spec.startup_commands, + ) + if spec.dockerfile_content is None: + raise LiumError("InstanceSpec requires template_ref or dockerfile_content") + return None + + async def _find_pod_id(self, pod_name: str) -> str | None: + for pod in await self.list_pods(): + if str(pod.get("pod_name")) == pod_name and pod.get("id"): + return str(pod["id"]) + return None + + async def _cleanup_after_rent(self, pod_id: str | None, pod_name: str) -> None: + if pod_id is None: + pod_id = await self._resolve_pod_id_quietly(pod_name) + if pod_id is None: + logger.warning( + "post-rent cleanup could not resolve a pod id for %r; the rented " + "pod (if any) will auto-terminate via termination_hours", + pod_name, + ) + return + await self._terminate_and_verify_quietly(pod_id) + + async def _resolve_pod_id_quietly(self, pod_name: str) -> str | None: + try: + return await self._find_pod_id(pod_name) + except Exception: # noqa: BLE001 - cleanup must not mask the original error + logger.warning("post-rent cleanup pod lookup failed for %r", pod_name) + return None + + async def _terminate_and_verify_quietly(self, instance_id: str) -> None: + try: + await self.terminate(instance_id) + except Exception: # noqa: BLE001 - cleanup must not mask the original error + logger.warning("cleanup terminate failed for pod %s", instance_id) + try: + await self.verify_terminated(instance_id) + except Exception: # noqa: BLE001 - cleanup must not mask the original error + logger.warning("cleanup verify_terminated failed for pod %s", instance_id) + + @staticmethod + def _extract_pod_id(response: httpx.Response) -> str | None: + if not response.content: + return None + try: + data = response.json() + except ValueError as exc: + raise LiumError( + "Lium rent returned a 2xx response with an unparseable body" + ) from exc + if not isinstance(data, Mapping): + return None + for key in ("id", "pod_id", "uuid"): + value = data.get(key) + if value: + return str(value) + pod = data.get("pod") + if isinstance(pod, Mapping) and pod.get("id"): + return str(pod["id"]) + return None + + def _headers(self) -> dict[str, str]: + return {"X-API-Key": self._api_key, "Accept": "application/json"} + + def _client(self) -> httpx.AsyncClient: + kwargs: dict[str, Any] = { + "base_url": self._base_url, + "timeout": self._timeout, + "headers": self._headers(), + } + if self._transport is not None: + kwargs["transport"] = self._transport + return httpx.AsyncClient(**kwargs) + + async def _send( + self, + method: str, + path: str, + *, + json_body: Any | None = None, + params: Mapping[str, Any] | None = None, + ) -> httpx.Response: + try: + async with self._client() as client: + return await client.request(method, path, json=json_body, params=params) + except httpx.HTTPError as exc: + raise LiumError(f"Lium request {method} {path} failed") from exc + + async def _request( + self, + method: str, + path: str, + *, + json_body: Any | None = None, + params: Mapping[str, Any] | None = None, + ) -> httpx.Response: + response = await self._send(method, path, json_body=json_body, params=params) + if response.status_code >= 400: + raise LiumError( + f"Lium {method} {path} returned {response.status_code}", + status_code=response.status_code, + ) + return response + + +def _as_list(data: Any, wrapper_key: str) -> list[dict[str, Any]]: + if isinstance(data, list): + return [item for item in data if isinstance(item, dict)] + if isinstance(data, Mapping): + inner = data.get(wrapper_key) + if isinstance(inner, list): + return [item for item in inner if isinstance(item, dict)] + return [] + + +def _parse_offer(item: Mapping[str, Any]) -> Offer | None: + offer_id = item.get("id") + price = _extract_price(item) + if not offer_id or price is None: + return None + return Offer( + id=str(offer_id), + gpu_type=_extract_gpu_type(item), + gpu_count=_extract_gpu_count(item), + price_per_hour=price, + provider="lium", + raw=item, + ) + + +def _extract_price(item: Mapping[str, Any]) -> float | None: + for key in ("price_per_hour", "price_per_gpu", "pending_price_per_hour"): + value = item.get(key) + if value is not None: + try: + return float(value) + except (TypeError, ValueError): + return None + return None + + +def _extract_gpu_type(item: Mapping[str, Any]) -> str: + name = item.get("machine_name") or item.get("gpu_type") or item.get("gpu_name") + if not name: + specs = item.get("specs") + if isinstance(specs, Mapping): + gpu = specs.get("gpu") + if isinstance(gpu, Mapping): + details = gpu.get("details") + if isinstance(details, list) and details: + first = details[0] + if isinstance(first, Mapping): + name = first.get("name") + return str(name) if name else "" + + +def _extract_gpu_count(item: Mapping[str, Any]) -> int: + count = item.get("gpu_count") + if count is None: + specs = item.get("specs") + if isinstance(specs, Mapping): + gpu = specs.get("gpu") + if isinstance(gpu, Mapping): + count = gpu.get("count") + try: + return int(count) if count is not None else 0 + except (TypeError, ValueError): + return 0 + + +def _parse_instance(data: Any) -> Instance: + if not isinstance(data, Mapping): + raise LiumError("unexpected pod response shape") + return Instance( + id=str(data.get("id", "")), + status=str(data.get("status", "")), + provider="lium", + raw=data, + ) diff --git a/src/base/compute/provider.py b/src/base/compute/provider.py new file mode 100644 index 00000000..d814fec9 --- /dev/null +++ b/src/base/compute/provider.py @@ -0,0 +1,115 @@ +"""Provider-agnostic compute contract for miner-funded GPU instances. + +The worker plane provisions GPU instances on external providers (Lium, Targon) +that a miner pays for. This module defines the shared contract every provider +client implements plus the cost guardrails that are part of that contract: + +* Every :class:`InstanceSpec` MUST carry a bounded ``max_lifetime_hours`` and a + ``max_price_per_hour``; :meth:`ProviderClient.provision` refuses a spec that is + unbounded (raising :class:`CostGuardrailError`) BEFORE any network call. +* Every ``provision`` code path, including exceptions, MUST attempt to + ``terminate`` + ``verify_terminated`` the instance it created so a failed + provisioning never leaks a billable pod (architecture.md sec 3.1). +""" + +from __future__ import annotations + +from collections.abc import AsyncIterator, Mapping +from dataclasses import dataclass, field +from typing import Any, Protocol, runtime_checkable + + +class ProviderError(RuntimeError): + """Base error for provider client operations.""" + + +class CostGuardrailError(ProviderError): + """A provision request violated a cost guardrail. + + Raised for an unbounded/non-positive lifetime, a missing price bound, or an + offer whose price exceeds the spec's ``max_price_per_hour``. It is a typed + error so callers can distinguish a guardrail refusal (never billed) from a + generic transport/API failure. + """ + + +@dataclass(frozen=True) +class InstanceSpec: + """A provider-agnostic request to provision one compute instance. + + ``max_lifetime_hours`` and ``max_price_per_hour`` are mandatory cost + guardrails: they are typed optional only so a caller can construct an invalid + spec that :meth:`ProviderClient.provision` will reject up front, but a valid + provisioning requires both to be positive. + """ + + name: str + template_ref: str | None = None + image: str | None = None + image_digest: str | None = None + dockerfile_content: str | None = None + env: Mapping[str, str] = field(default_factory=dict) + ports: tuple[int, ...] = (22,) + ssh_public_keys: tuple[str, ...] = () + ssh_key_name: str | None = None + startup_commands: str | None = None + max_lifetime_hours: float | None = None + max_price_per_hour: float | None = None + gpu_count: int = 1 + + +@dataclass(frozen=True) +class Offer: + """A rentable GPU offer exposed by a provider. + + ``price_per_hour`` is the per-GPU hourly price the cost guardrails filter on + (Lium prices per GPU; the mission budget is expressed per GPU/hour). + """ + + id: str + gpu_type: str + gpu_count: int + price_per_hour: float + provider: str = "" + raw: Mapping[str, Any] = field(default_factory=dict) + + +@dataclass(frozen=True) +class Instance: + """A provisioned compute instance (a Lium pod / Targon workload).""" + + id: str + status: str + provider: str = "" + raw: Mapping[str, Any] = field(default_factory=dict) + + +@runtime_checkable +class ProviderClient(Protocol): + """The contract every compute provider client implements.""" + + async def list_offers( + self, *, max_price_per_hour: float | None = None + ) -> list[Offer]: + """Return rentable offers, optionally filtered to ``<= price`` per hour.""" + ... + + async def provision(self, spec: InstanceSpec) -> Instance: + """Provision an instance for ``spec`` under the cost guardrails.""" + ... + + async def status(self, instance_id: str) -> Instance: + """Return the current state of a provisioned instance.""" + ... + + def stream_logs(self, instance_id: str) -> AsyncIterator[str]: + """Yield log lines for a provisioned instance.""" + ... + + async def terminate(self, instance_id: str) -> None: + """Terminate an instance; idempotent (a missing instance is success).""" + ... + + async def verify_terminated(self, instance_id: str) -> bool: + """Return ``True`` once the instance is absent from the provider.""" + ... diff --git a/src/base/compute/targon.py b/src/base/compute/targon.py new file mode 100644 index 00000000..4ed2b1c9 --- /dev/null +++ b/src/base/compute/targon.py @@ -0,0 +1,391 @@ +"""Targon compute provider client (architecture.md sec 3.1). + +Thin async ``httpx`` wrapper over the Targon REST API (base +``https://api.targon.com/tha/v2``, auth header ``Authorization: Bearer``). It +implements the :class:`~base.compute.provider` contract with the cost guardrails +baked in and two Targon-specific realities surfaced as typed errors: + +* Targon exposes NO balance/credits endpoint, so :meth:`TargonClient.balance` + raises :class:`BalanceUnavailableError` WITHOUT issuing any HTTP request rather + than returning a fake/silent value. +* A deploy that fails for insufficient credits is surfaced as a distinct typed + :class:`InsufficientCreditsError` and is NEVER retried. A deploy is a two-step + create-then-deploy flow (Targon has no single ``POST /workloads/deploy`` route): + ``POST /workloads`` (create -> ``uid``) then ``POST /workloads/{uid}/deploy``; + a credit failure at EITHER step raises and is not retried. + +The API key lives only in the request header; it is never logged, embedded in an +error message, or exposed via ``repr``. +""" + +from __future__ import annotations + +import logging +from collections.abc import AsyncIterator, Mapping +from typing import Any + +import httpx + +from base.compute.provider import ( + CostGuardrailError, + Instance, + InstanceSpec, + Offer, + ProviderError, +) + +logger = logging.getLogger(__name__) + +TARGON_API_BASE_URL = "https://api.targon.com/tha/v2" + +_CREDIT_TOKENS = ("insufficient", "credit", "payment required", "out of credit") + + +class TargonError(ProviderError): + """A Targon API request failed (non-2xx response or transport error).""" + + def __init__(self, message: str, *, status_code: int | None = None) -> None: + super().__init__(message) + self.status_code = status_code + + +class InsufficientCreditsError(TargonError): + """A Targon deploy was rejected for insufficient credits. + + Targon exposes no balance endpoint, so an insufficient-credit deploy failure + is only observable at deploy time. It is a distinct typed error (identifiable + via ``isinstance``) that callers must NOT retry. + """ + + +class BalanceUnavailableError(TargonError): + """Account balance is not queryable: Targon exposes no balance endpoint.""" + + +class TargonClient: + """Async client for the Targon GPU compute API.""" + + def __init__( + self, + api_key: str, + *, + base_url: str = TARGON_API_BASE_URL, + transport: httpx.AsyncBaseTransport | None = None, + timeout_seconds: float = 30.0, + ) -> None: + self._api_key = api_key + self._base_url = base_url.rstrip("/") + self._transport = transport + self._timeout = timeout_seconds + + def __repr__(self) -> str: + return f"TargonClient(base_url={self._base_url!r})" + + # -- offers --------------------------------------------------------------- + + async def list_offers( + self, *, max_price_per_hour: float | None = None + ) -> list[Offer]: + response = await self._request( + "GET", "/inventory", params={"type": "rental", "gpu": "true"} + ) + offers: list[Offer] = [] + for item in _as_list(response.json(), "items"): + offer = _parse_inventory_offer(item) + if offer is None: + continue + if ( + max_price_per_hour is not None + and offer.price_per_hour > max_price_per_hour + ): + continue + offers.append(offer) + return offers + + # -- apps / workloads listing -------------------------------------------- + + async def list_apps(self) -> list[dict[str, Any]]: + response = await self._request("GET", "/apps") + return _as_list(response.json(), "items") + + async def create_app( + self, name: str, *, project_id: str | None = None + ) -> dict[str, Any]: + body: dict[str, Any] = {"name": name} + if project_id is not None: + body["project_id"] = project_id + response = await self._request("POST", "/apps", json_body=body) + result = response.json() + return result if isinstance(result, dict) else {} + + async def list_workloads(self) -> list[dict[str, Any]]: + response = await self._request("GET", "/workloads") + return _as_list(response.json(), "items") + + async def workload_state(self, uid: str) -> dict[str, Any]: + response = await self._request("GET", f"/workloads/{uid}/state") + result = response.json() + return result if isinstance(result, dict) else {} + + async def workload_events(self, uid: str) -> dict[str, Any]: + response = await self._request("GET", f"/workloads/{uid}/events") + result = response.json() + return result if isinstance(result, dict) else {} + + # -- lifecycle ------------------------------------------------------------ + + async def provision(self, spec: InstanceSpec) -> Instance: + lifetime = spec.max_lifetime_hours + if lifetime is None or lifetime <= 0: + raise CostGuardrailError( + "InstanceSpec.max_lifetime_hours must be a positive bound (hours)" + ) + if lifetime < 1: + raise CostGuardrailError( + "InstanceSpec.max_lifetime_hours must be at least 1 hour: Targon " + "termination_hours has 1-hour granularity, so a sub-hour bound " + "would truncate to termination_hours=0 and disable auto-termination" + ) + if spec.max_price_per_hour is None or spec.max_price_per_hour <= 0: + raise CostGuardrailError( + "InstanceSpec.max_price_per_hour must be a positive bound" + ) + return await self.deploy(_build_workload_payload(spec, lifetime=lifetime)) + + async def deploy(self, workload: Mapping[str, Any]) -> Instance: + """Deploy a workload via Targon's two-step create-then-deploy flow. + + Targon has NO single ``POST /workloads/deploy`` route (confirmed against + the live API and the ``targon-sdk``): a workload is first CREATED + (``POST /workloads`` -> ``uid``) then DEPLOYED + (``POST /workloads/{uid}/deploy``). Both calls are part of ONE deploy + attempt -- an insufficient-credit failure at EITHER step is surfaced as a + typed :class:`InsufficientCreditsError` and is NEVER retried. + """ + create = await self._send("POST", "/workloads", json_body=dict(workload)) + self._raise_deploy_error(create, "POST /workloads") + uid = _extract_workload_uid(create.json()) + if not uid: + raise TargonError("Targon POST /workloads returned no workload uid") + deployed = await self._send("POST", f"/workloads/{uid}/deploy") + self._raise_deploy_error(deployed, f"POST /workloads/{uid}/deploy") + return _parse_workload_instance(deployed.json(), fallback_uid=uid) + + async def status(self, instance_id: str) -> Instance: + response = await self._request("GET", f"/workloads/{instance_id}") + return _parse_workload_instance(response.json()) + + async def stream_logs(self, instance_id: str) -> AsyncIterator[str]: + async with self._client() as client: + async with client.stream( + "GET", f"/workloads/{instance_id}/logs" + ) as response: + if response.status_code >= 400: + await response.aread() + raise TargonError( + f"Targon GET /workloads/{instance_id}/logs returned " + f"{response.status_code}", + status_code=response.status_code, + ) + async for line in response.aiter_lines(): + yield line + + async def terminate(self, instance_id: str) -> None: + response = await self._send("DELETE", f"/workloads/{instance_id}") + if response.status_code == 404: + return + if response.status_code >= 400: + raise TargonError( + f"Targon DELETE /workloads/{instance_id} returned " + f"{response.status_code}", + status_code=response.status_code, + ) + + async def verify_terminated(self, instance_id: str) -> bool: + response = await self._send("GET", f"/workloads/{instance_id}") + if response.status_code == 404: + return True + if response.status_code >= 400: + raise TargonError( + f"Targon GET /workloads/{instance_id} returned {response.status_code}", + status_code=response.status_code, + ) + instance = _parse_workload_instance(response.json()) + return instance.status.lower() in {"deleted", "terminated", "stopped"} + + async def balance(self) -> float: + raise BalanceUnavailableError( + "Targon exposes no balance endpoint; balance is only visible in the " + "web dashboard" + ) + + # -- internals ------------------------------------------------------------ + + def _raise_deploy_error(self, response: httpx.Response, route: str) -> None: + if response.status_code < 400: + return + if _is_insufficient_credits(response.status_code, response.text): + raise InsufficientCreditsError( + f"Targon {route} rejected for insufficient credits " + f"(status {response.status_code})", + status_code=response.status_code, + ) + raise TargonError( + f"Targon {route} returned {response.status_code}", + status_code=response.status_code, + ) + + def _headers(self) -> dict[str, str]: + return { + "Authorization": f"Bearer {self._api_key}", + "Accept": "application/json", + } + + def _client(self) -> httpx.AsyncClient: + kwargs: dict[str, Any] = { + "base_url": self._base_url, + "timeout": self._timeout, + "headers": self._headers(), + } + if self._transport is not None: + kwargs["transport"] = self._transport + return httpx.AsyncClient(**kwargs) + + async def _send( + self, + method: str, + path: str, + *, + json_body: Any | None = None, + params: Mapping[str, Any] | None = None, + ) -> httpx.Response: + try: + async with self._client() as client: + return await client.request(method, path, json=json_body, params=params) + except httpx.HTTPError as exc: + raise TargonError(f"Targon request {method} {path} failed") from exc + + async def _request( + self, + method: str, + path: str, + *, + json_body: Any | None = None, + params: Mapping[str, Any] | None = None, + ) -> httpx.Response: + response = await self._send(method, path, json_body=json_body, params=params) + if response.status_code >= 400: + raise TargonError( + f"Targon {method} {path} returned {response.status_code}", + status_code=response.status_code, + ) + return response + + +def _build_workload_payload(spec: InstanceSpec, *, lifetime: float) -> dict[str, Any]: + payload: dict[str, Any] = {"name": spec.name, "type": "RENTAL"} + if spec.template_ref: + payload["resource_name"] = spec.template_ref + if spec.image: + payload["image"] = spec.image + if spec.image_digest: + payload["image_digest"] = spec.image_digest + if spec.gpu_count: + payload["gpu_count"] = spec.gpu_count + payload["envs"] = [{"name": k, "value": v} for k, v in spec.env.items()] + if spec.ports: + payload["ports"] = [{"port": port} for port in spec.ports] + if spec.ssh_public_keys: + payload["ssh_public_keys"] = list(spec.ssh_public_keys) + payload["termination_hours"] = int(lifetime) + return payload + + +def _is_insufficient_credits(status_code: int, body_text: str) -> bool: + if status_code in (402, 403): + return True + lowered = (body_text or "").lower() + return any(token in lowered for token in _CREDIT_TOKENS) + + +def _as_list(data: Any, wrapper_key: str) -> list[dict[str, Any]]: + if isinstance(data, list): + return [item for item in data if isinstance(item, dict)] + if isinstance(data, Mapping): + inner = data.get(wrapper_key) + if isinstance(inner, list): + return [item for item in inner if isinstance(item, dict)] + return [] + + +def _parse_inventory_offer(item: Mapping[str, Any]) -> Offer | None: + name = item.get("name") or item.get("display_name") + cost = item.get("cost_per_hour") + if not name or cost is None: + return None + try: + cost_per_hour = float(cost) + except (TypeError, ValueError): + return None + if _coerce_int(item.get("available")) <= 0: + return None + gpu_count = _extract_gpu_count(item) + # Targon inventory quotes cost_per_hour for the WHOLE shape; Offer.price_per_hour + # is per-GPU (what the per-GPU max_price cap filters on). Fall back to the + # whole-shape cost when the count is unknown (avoids ZeroDivisionError). + price_per_gpu = cost_per_hour / gpu_count if gpu_count >= 1 else cost_per_hour + return Offer( + id=str(name), + gpu_type=_extract_gpu_type(item), + gpu_count=gpu_count, + price_per_hour=price_per_gpu, + provider="targon", + raw=item, + ) + + +def _extract_gpu_type(item: Mapping[str, Any]) -> str: + spec = item.get("spec") + if isinstance(spec, Mapping): + gpu_type = spec.get("gpu_type") + if gpu_type: + return str(gpu_type) + name = item.get("display_name") or item.get("name") + return str(name) if name else "" + + +def _extract_gpu_count(item: Mapping[str, Any]) -> int: + spec = item.get("spec") + if isinstance(spec, Mapping): + count = _coerce_int(spec.get("gpu_count")) + if count >= 1: + return count + return _coerce_int(item.get("gpu_count")) + + +def _coerce_int(value: Any) -> int: + try: + return int(value) if value is not None else 0 + except (TypeError, ValueError): + return 0 + + +def _parse_workload_instance(data: Any, *, fallback_uid: str | None = None) -> Instance: + if not isinstance(data, Mapping): + raise TargonError("unexpected workload response shape") + uid = data.get("uid") or data.get("id") or fallback_uid or "" + status = "" + state = data.get("state") + if isinstance(state, Mapping): + status = str(state.get("status", "")) + if not status: + status = str(data.get("status", "")) + return Instance(id=str(uid), status=status, provider="targon", raw=data) + + +def _extract_workload_uid(data: Any) -> str: + if isinstance(data, Mapping): + uid = data.get("uid") or data.get("id") + if uid: + return str(uid) + return "" diff --git a/src/base/compute/worker_deployment.py b/src/base/compute/worker_deployment.py new file mode 100644 index 00000000..561a7bd4 --- /dev/null +++ b/src/base/compute/worker_deployment.py @@ -0,0 +1,219 @@ +"""Declarative worker-image deploy definitions for both compute providers. + +The miner deploys the (future) worker image into a GPU instance they pay for. The +image is referenced the same way everywhere: BY DIGEST, so a rented pod runs the +exact bytes we published. This module builds the two declarative definitions that +carry that pin: + +* :func:`build_lium_worker_template` -- a Lium ``CustomTemplateRequest`` payload + for ``POST /templates`` (architecture.md sec 3.2 / lium-api.md), and +* :func:`build_targon_worker_app` -- a Targon app definition referencing the same + pinned image with an explicit GPU resource shape. + +For M1 the image defaults to the already-published ``prism-evaluator`` image +pinned by digest (a placeholder standing in for M2's ``docker/Dockerfile.worker`` +image). Both builders take ``image``/``image_digest`` as inputs so M2 swaps in the +worker image with a one-line change; nothing else about the definitions moves. + +The definitions embed NO credentials: environment plumbing is caller-provided and +carries only non-secret configuration (master URL, worker role, ...). +""" + +from __future__ import annotations + +import ipaddress +import re +from collections.abc import Mapping, Sequence +from typing import Any +from urllib.parse import urlsplit + +# Placeholder worker image, pinned by digest. IMPORTANT: this is a PRIVATE-namespace +# GHCR image and is NOT reliably pullable by a rented provider host -- on Lium a +# private-namespace image makes pod creation fail with CREATION_FAILED. It stands in +# for docker/Dockerfile.worker's PUBLICLY-published image; a real `base worker deploy` +# MUST override it with an explicit, publicly-pullable, digest-pinned image +# (worker.deploy.image + worker.deploy.image_digest / BASE_WORKER__DEPLOY__IMAGE*). +# See docs/miner/worker-plane.md ("Publishing the worker image") for the publish + +# digest-pin procedure. The two constants remain only as defaults for the M1 +# declarative-definition builders below and their tests. +WORKER_IMAGE = "ghcr.io/baseintelligence/prism-evaluator" +WORKER_IMAGE_TAG = "latest" +WORKER_IMAGE_DIGEST = ( + "sha256:713b39f13af69dbaf229e67fb682df8a2b7ac93dd02d9e60867ff021d4edb3c9" +) + +WORKER_TEMPLATE_NAME = "prism-worker" +WORKER_APP_NAME = "prism-worker" + +# SSH (22) is mandatory so the miner can reach the pod; the worker's local broker +# is bound to loopback inside the instance and needs no external mapping. +WORKER_INTERNAL_PORTS: tuple[int, ...] = (22,) + +# Metachar-free keep-alive for the Lium template startup_commands. Lium rejects a +# rent whose template startup_commands contain shell metacharacters ("Malicious +# startup command detected"; live-confirmed, library/lium-api.md), so the value +# MUST NOT contain '&&', ';', or '|'. The container's own entrypoint launches the +# agent; this keep-alive only guarantees the pod stays up (Lium's pod agent +# provides SSH independently of the container command). +WORKER_STARTUP_COMMANDS = "tail -f /dev/null" + +# Shell metacharacters Lium's rent guard rejects in a template startup command. +_SHELL_METACHARACTERS: tuple[str, ...] = ("&&", "||", ";", "|", "&", "`", "$(") + +# Default Targon GPU resource shape (inventory shape id + human GPU type). Real +# Targon inventory ids carry a size suffix (h100-small, b200-large per +# library/targon-api.md); a bare 'h100' would be rejected by a live deploy. +WORKER_GPU_SHAPE = "h100-small" +WORKER_GPU_TYPE = "H100" + +_DIGEST_RE = re.compile(r"^sha256:[0-9a-f]{64}$") + + +def is_pinned_digest(value: str) -> bool: + """Return ``True`` if ``value`` is a well-formed ``sha256:<64-hex>`` digest.""" + return bool(_DIGEST_RE.match(value)) + + +def is_loopback_url(value: str) -> bool: + """Return ``True`` if ``value`` is a URL whose host is loopback/unspecified. + + Lium's edge CDN/WAF returns ``403 "Request blocked"`` for ANY request body that + contains a loopback URL (``http://127.0.0.1...`` / ``http://localhost...``), and + ``base worker deploy`` bakes the pod env into the WAF-sensitive ``POST /templates`` + body. Such values must therefore be stripped before the template is created. A + non-URL string (a signature, a hotkey, a flag) is never a loopback URL. + """ + if not value: + return False + text = value.strip() + host = urlsplit(text).hostname + if host is None and "//" not in text: + host = urlsplit(f"//{text}").hostname + if not host: + return False + host = host.lower() + if host == "localhost" or host.endswith(".localhost"): + return True + try: + address = ipaddress.ip_address(host) + except ValueError: + return False + return address.is_loopback or address.is_unspecified + + +def is_metachar_free(command: str) -> bool: + """Return ``True`` if ``command`` carries no shell metacharacter Lium rejects. + + Lium's rent guard refuses a template startup command containing shell + metacharacters (``&&``, ``;``, ``|`` ...); a compliant keep-alive is a single + plain command such as ``tail -f /dev/null``. + """ + return not any(token in command for token in _SHELL_METACHARACTERS) + + +def _require_metachar_free(command: str) -> str: + if not is_metachar_free(command): + raise ValueError( + "startup_commands must be a single metachar-free command (no " + "'&&', ';', '|', ...): Lium rejects shell metacharacters" + ) + return command + + +def _require_digest(image_digest: str) -> str: + if not is_pinned_digest(image_digest): + raise ValueError( + "image_digest must be a pinned digest of the form 'sha256:<64 hex>'" + ) + return image_digest + + +def pinned_image_reference(image: str, digest: str, *, tag: str | None = None) -> str: + """Return a fully qualified, digest-pinned image reference. + + With a tag: ``:@``; without: ``@``. The + ``@`` suffix is what makes the reference immutable. + """ + _require_digest(digest) + if tag: + return f"{image}:{tag}@{digest}" + return f"{image}@{digest}" + + +def build_lium_worker_template( + *, + image: str = WORKER_IMAGE, + image_digest: str = WORKER_IMAGE_DIGEST, + image_tag: str | None = WORKER_IMAGE_TAG, + name: str = WORKER_TEMPLATE_NAME, + environment: Mapping[str, str] | None = None, + internal_ports: Sequence[int] = WORKER_INTERNAL_PORTS, + entrypoint: str = "", + startup_commands: str = WORKER_STARTUP_COMMANDS, + is_private: bool = True, + container_start_immediately: bool = True, +) -> dict[str, Any]: + """Build the Lium ``CustomTemplateRequest`` payload for the worker image. + + The returned dict is the exact JSON body for ``POST /templates``. ``image`` + is pinned by ``image_digest`` (validated to be ``sha256:<64 hex>``) and the + internal ports MUST include 22 (SSH). ``startup_commands`` is validated to be + metachar-free (Lium rejects shell metacharacters at rent time), defaulting to + a plain keep-alive so the pod stays up while the image entrypoint runs the + agent. + """ + _require_digest(image_digest) + _require_metachar_free(startup_commands) + ports = [int(port) for port in internal_ports] + if 22 not in ports: + raise ValueError("internal_ports must include 22 (SSH)") + return { + "name": name, + "docker_image": image, + "docker_image_tag": image_tag or "", + "docker_image_digest": image_digest, + "environment": dict(environment or {}), + "entrypoint": entrypoint, + "startup_commands": startup_commands, + "internal_ports": ports, + "is_private": is_private, + "container_start_immediately": container_start_immediately, + } + + +def build_targon_worker_app( + *, + image: str = WORKER_IMAGE, + image_digest: str = WORKER_IMAGE_DIGEST, + image_tag: str | None = WORKER_IMAGE_TAG, + name: str = WORKER_APP_NAME, + gpu_shape: str = WORKER_GPU_SHAPE, + gpu_type: str = WORKER_GPU_TYPE, + gpu_count: int = 1, + environment: Mapping[str, str] | None = None, + ports: Sequence[int] = WORKER_INTERNAL_PORTS, + min_replicas: int = 0, + max_replicas: int = 1, +) -> dict[str, Any]: + """Build the Targon app definition for the worker image. + + The app references a fully qualified, digest-pinned image and declares its GPU + resource shape (``resource`` inventory shape name + ``gpu_type``/``gpu_count``). + Environment is plumbed as Targon ``{"name", "value"}`` pairs and carries no + secrets. + """ + _require_digest(image_digest) + if gpu_count < 1: + raise ValueError("gpu_count must be >= 1") + return { + "name": name, + "image": pinned_image_reference(image, image_digest, tag=image_tag), + "image_digest": image_digest, + "resource": gpu_shape, + "gpu_type": gpu_type, + "gpu_count": gpu_count, + "min_replicas": min_replicas, + "max_replicas": max_replicas, + "envs": [{"name": k, "value": v} for k, v in dict(environment or {}).items()], + "ports": [int(port) for port in ports], + } diff --git a/src/base/config/settings.py b/src/base/config/settings.py index 04840083..5f992b70 100644 --- a/src/base/config/settings.py +++ b/src/base/config/settings.py @@ -224,6 +224,124 @@ class SecuritySettings(BaseModel): admin_token_file: str | None = None +class ComputeSettings(BaseModel): + """Miner-funded GPU worker plane (architecture.md sec 3.3). + + ALL worker-plane behavior is gated behind ``worker_plane_enabled`` (env + ``BASE_COMPUTE__WORKER_PLANE_ENABLED``); OFF (the default) preserves legacy + behavior byte-for-byte: the worker coordination surface is not mounted and + gpu units route to validators exactly as today. ``worker_heartbeat_ttl_seconds`` + is the freshness window: an ``active`` worker whose last heartbeat is older + than the TTL is reported ``stale`` and is not assignable. + """ + + worker_plane_enabled: bool = False + worker_heartbeat_ttl_seconds: int = 120 + worker_signature_ttl_seconds: int = 300 + worker_nonce_ttl_seconds: int = 86_400 + worker_health_interval_seconds: float = 60.0 + #: Number of DISTINCT-owner workers each gpu work unit is replicated across + #: when the worker plane is on. Degrades to 1 (with a recorded warning) when + #: fewer eligible distinct owners exist. + replication_factor: int = 2 + + +class WorkerAgentSettings(BaseModel): + """Miner-funded GPU worker agent runtime (architecture.md sec 3.2). + + The agent registers with the master under a miner-signed binding, heartbeats, + pulls gpu work units, executes them on its OWN local broker, and posts + ExecutionProof-carrying results. It authenticates as its worker keypair, never + as a metagraph validator permit, and never holds a provider API key. + ``master_url`` is required to reach the coordination plane; + ``heartbeat_interval_seconds`` left unset uses the interval the master + returns from ``register``. + """ + + master_url: str | None = None + gateway_url: str | None = None + capabilities: list[str] = Field(default_factory=lambda: ["gpu"]) + version: str | None = None + heartbeat_interval_seconds: int | None = None + poll_interval_seconds: float = 5.0 + request_timeout_seconds: float = 15.0 + #: Worker's OWN local Docker broker. Falls back to ``docker.broker_url``. + broker_url: str | None = None + broker_token: str | None = None + broker_token_file: str | None = "/run/secrets/base_broker_token" + #: Extra allowed eval-image prefixes (added to ``docker.broker_allowed_images``). + allowed_images: list[str] = Field(default_factory=list) + run_timeout_seconds: int = 3_600 + + +class WorkerDeploySettings(BaseModel): + """Provisioning inputs for ``base worker deploy`` (architecture.md sec 3.2). + + ``provider`` selects the deploy target (``local`` runs the agent on this host; + ``lium``/``targon`` provision a paid GPU instance running the worker image). + Image fields default to the published worker image when unset. Cost guardrails + (``max_price_per_hour``/``max_lifetime_hours``) bound provider provisioning. + ``startup_commands`` MUST be metachar-free (Lium rejects shell metacharacters + at rent time). + """ + + provider: Literal["local", "lium", "targon"] = "local" + provider_instance_ref: str | None = None + image: str | None = None + image_digest: str | None = None + #: Informational-only: a human-readable tag recorded alongside the pin. The + #: digest-pinned deploy path provisions BY DIGEST (``image`` + ``image_digest``) + #: and never consumes this value, so it never affects which image bytes run. + image_tag: str | None = None + template_name: str | None = None + gpu_count: int = 1 + max_price_per_hour: float | None = None + max_lifetime_hours: float = 1.0 + ssh_public_key: str | None = None + ssh_public_key_file: str | None = None + ssh_key_name: str | None = None + startup_commands: str = "tail -f /dev/null" + #: Seconds ``deploy --provider local`` polls the master for the worker to + #: reach ``active`` before reporting failure. + ready_timeout_seconds: float = 60.0 + + +class WorkerIdentitySettings(BaseModel): + """Worker keypair + miner binding material for the worker agent/CLI. + + The WORKER keypair signs coordination requests + ExecutionProofs; the MINER + keypair signs the enrollment binding (``worker-binding:{worker_pubkey}: + {miner_hotkey}:{nonce}``). Each key resolves from an sr25519 dev URI + (``//Worker``), a mnemonic, or a bittensor wallet, in that order, falling + back to ``network.wallet`` for the worker key. When the binding is signed + out-of-band (e.g. a Lium/Targon pod that never holds the miner key) supply the + pre-signed ``miner_hotkey`` + ``binding_signature`` + ``binding_nonce`` + instead of a miner key. + """ + + key_uri: str | None = None + key_mnemonic: str | None = None + wallet_name: str | None = None + wallet_hotkey: str | None = None + wallet_path: str | None = None + miner_hotkey: str | None = None + miner_key_uri: str | None = None + miner_key_mnemonic: str | None = None + miner_wallet_name: str | None = None + miner_wallet_hotkey: str | None = None + miner_wallet_path: str | None = None + binding_signature: str | None = None + binding_nonce: str | None = None + + +class WorkerSettings(BaseModel): + """Top-level ``base worker`` config: agent runtime, deploy, and identity.""" + + agent: WorkerAgentSettings = Field(default_factory=WorkerAgentSettings) + deploy: WorkerDeploySettings = Field(default_factory=WorkerDeploySettings) + identity: WorkerIdentitySettings = Field(default_factory=WorkerIdentitySettings) + + class ProviderEntry(BaseModel): """One configured LLM provider: its OpenAI-compatible base URL + key. @@ -453,6 +571,8 @@ class Settings(BaseModel): database: DatabaseSettings = Field(default_factory=DatabaseSettings) docker: DockerSettings = Field(default_factory=DockerSettings) security: SecuritySettings = Field(default_factory=SecuritySettings) + compute: ComputeSettings = Field(default_factory=ComputeSettings) + worker: WorkerSettings = Field(default_factory=WorkerSettings) gateway: GatewaySettings = Field(default_factory=GatewaySettings) observability: ObservabilitySettings = Field(default_factory=ObservabilitySettings) supervisor: SupervisorSettings = Field(default_factory=SupervisorSettings) diff --git a/src/base/coordination/__init__.py b/src/base/coordination/__init__.py new file mode 100644 index 00000000..9284548d --- /dev/null +++ b/src/base/coordination/__init__.py @@ -0,0 +1,17 @@ +"""Shared coordination-plane agent primitives (validator + worker planes).""" + +from base.coordination.agent_loop import ( + AgentCycleSummary, + BackoffPolicy, + backoff_sleep, + is_transient_error, + sleep_until, +) + +__all__ = [ + "AgentCycleSummary", + "BackoffPolicy", + "backoff_sleep", + "is_transient_error", + "sleep_until", +] diff --git a/src/base/coordination/agent_loop.py b/src/base/coordination/agent_loop.py new file mode 100644 index 00000000..efa2df91 --- /dev/null +++ b/src/base/coordination/agent_loop.py @@ -0,0 +1,84 @@ +"""Transport-agnostic loop primitives shared by coordination-plane agents. + +Both the validator agent (:mod:`base.validator.agent.runtime`) and the +miner-funded worker agent (:mod:`base.worker.runtime`) run the same +register/heartbeat/pull/execute/post shape. The pure pieces of that loop live +here so the two agents share one implementation instead of copy-pasting it: +bounded exponential backoff, the transient-vs-permanent error classification, +the per-pass summary, and the shutdown-aware sleep helpers. This is a verbatim +extraction from the validator agent, so validator behavior is unchanged. +""" + +from __future__ import annotations + +import asyncio +from dataclasses import dataclass + + +@dataclass(frozen=True) +class BackoffPolicy: + """Bounded exponential backoff for transient master/coordination failures. + + The delay grows geometrically with the number of consecutive failures and is + capped at ``max_seconds`` so an agent retries a briefly-unavailable master + without either giving up or busy-looping. + """ + + initial_seconds: float = 1.0 + max_seconds: float = 60.0 + multiplier: float = 2.0 + + def delay(self, consecutive_failures: int) -> float: + """Backoff delay after ``consecutive_failures`` failures (>=1).""" + + if consecutive_failures <= 0: + return 0.0 + raw = self.initial_seconds * (self.multiplier ** (consecutive_failures - 1)) + return min(self.max_seconds, max(0.0, raw)) + + +def is_transient_error(exc: BaseException) -> bool: + """Whether a coordination failure is worth retrying with backoff. + + Transport errors (no status code) and ``429``/``5xx`` master responses are + transient; a ``4xx`` (e.g. ``403`` ineligible, ``404`` not registered, + ``401`` auth) is a permanent client error that should fail fast. + """ + + status_code = getattr(exc, "status_code", None) + if status_code is None: + return True + return status_code == 429 or status_code >= 500 + + +@dataclass(frozen=True) +class AgentCycleSummary: + """Counts from one assignment-processing pass.""" + + pulled: int + completed: int + failed: int + + +async def sleep_until(shutdown_event: asyncio.Event, seconds: float) -> None: + """Sleep up to ``seconds``, waking early if ``shutdown_event`` fires.""" + + try: + await asyncio.wait_for(shutdown_event.wait(), timeout=seconds) + except TimeoutError: + return + + +async def backoff_sleep(shutdown_event: asyncio.Event | None, seconds: float) -> bool: + """Sleep ``seconds``; return ``False`` if shutdown fired during the wait.""" + + if shutdown_event is None: + await asyncio.sleep(seconds) + return True + if shutdown_event.is_set(): + return False + try: + await asyncio.wait_for(shutdown_event.wait(), timeout=seconds) + except TimeoutError: + return True + return False diff --git a/src/base/db/__init__.py b/src/base/db/__init__.py index 6c3e1a7a..c28e124f 100644 --- a/src/base/db/__init__.py +++ b/src/base/db/__init__.py @@ -22,6 +22,11 @@ ValidatorStatus, WorkAssignment, WorkAssignmentStatus, + WorkerAssignment, + WorkerFault, + WorkerRegistration, + WorkerRequestNonce, + WorkerStatus, WorkResult, ) from base.db.repositories import ChallengeRepository @@ -54,6 +59,11 @@ "ValidatorStatus", "WorkAssignment", "WorkAssignmentStatus", + "WorkerAssignment", + "WorkerFault", + "WorkerRegistration", + "WorkerRequestNonce", + "WorkerStatus", "WorkResult", "create_engine", "create_session_factory", diff --git a/src/base/db/models.py b/src/base/db/models.py index fbda4273..d4f55b0d 100644 --- a/src/base/db/models.py +++ b/src/base/db/models.py @@ -62,13 +62,35 @@ class ValidatorHealthEventType(StrEnum): class WorkAssignmentStatus(StrEnum): - """Lifecycle states for a work unit coordinated to a validator.""" + """Lifecycle states for a work unit coordinated to a validator. + + ``disputed`` is a terminal worker-plane outcome (architecture.md sec 3.3): + a gpu unit whose replica manifest hashes diverged is disputed and NEVER + forwarded to the challenge (before or after audit); it is only ever set by + worker-plane reconciliation, so the validator plane never observes it. + """ PENDING = "pending" ASSIGNED = "assigned" RUNNING = "running" COMPLETED = "completed" FAILED = "failed" + DISPUTED = "disputed" + + +class WorkerStatus(StrEnum): + """Lifecycle states for a miner-funded GPU worker (architecture.md sec 3.3). + + ``pending`` -> ``active`` -> ``stale`` -> ``retired``. ``active`` requires a + verified miner binding AND a heartbeat within the freshness window + (``compute.worker_heartbeat_ttl_seconds``); ``retired`` is terminal (a + retired worker is never assignable and a heartbeat never resurrects it). + """ + + PENDING = "pending" + ACTIVE = "active" + STALE = "stale" + RETIRED = "retired" class TimestampMixin: @@ -713,3 +735,195 @@ class MinerRequestNonce(Base): created_at: Mapped[datetime] = mapped_column( DateTime(timezone=True), nullable=False ) + + +class WorkerRegistration(Base, TimestampMixin): + """A miner-funded GPU worker enrolled with the master (architecture.md 3.3). + + A worker is bound to exactly one miner hotkey via a miner-signed sr25519 + binding (``worker-binding:{worker_pubkey}:{miner_hotkey}:{nonce}``). The + ``worker_pubkey`` is unique, so a pubkey has a single stable owner: a + re-registration under a DIFFERENT miner hotkey is rejected (no silent + rebind). ``worker_id`` is the public identifier used by the heartbeat route + and returned to the agent. + """ + + __tablename__ = "worker_registrations" + __table_args__ = ( + Index("ix_worker_registrations_status", "status"), + Index("ix_worker_registrations_miner_hotkey", "miner_hotkey"), + Index("ix_worker_registrations_last_heartbeat_at", "last_heartbeat_at"), + Index( + "ix_worker_registrations_status_miner_hotkey", + "status", + "miner_hotkey", + ), + ) + + id: Mapped[uuid.UUID] = mapped_column( + Uuid(as_uuid=True), primary_key=True, default=uuid.uuid4 + ) + worker_id: Mapped[str] = mapped_column(Text, nullable=False, unique=True) + worker_pubkey: Mapped[str] = mapped_column(Text, nullable=False, unique=True) + miner_hotkey: Mapped[str] = mapped_column(Text, nullable=False) + binding_signature: Mapped[str] = mapped_column(Text, nullable=False) + binding_nonce: Mapped[str] = mapped_column(Text, nullable=False) + provider: Mapped[str] = mapped_column(Text, nullable=False) + provider_instance_ref: Mapped[str | None] = mapped_column(Text) + capabilities: Mapped[list[str]] = mapped_column( + JSON, + nullable=False, + default=list, + server_default="[]", + ) + status: Mapped[WorkerStatus] = mapped_column( + Enum( + WorkerStatus, + name="worker_status", + values_callable=lambda obj: [e.value for e in obj], + native_enum=False, + ), + nullable=False, + default=WorkerStatus.PENDING, + server_default=WorkerStatus.PENDING.value, + ) + last_seen_meta: Mapped[dict[str, Any]] = mapped_column( + JSON, + nullable=False, + default=dict, + server_default="{}", + ) + last_heartbeat_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True)) + + +class WorkerFault(Base): + """A fault attributed to a worker replica during reconciliation/audit. + + Written when a worker's ``ExecutionProof.manifest_sha256`` diverges from the + authoritative validator replay (architecture.md sec 3.3). Recording a fault + does NOT change the worker's ``worker_registrations.status``; faults are + surfaced read-only in the fleet view. + """ + + __tablename__ = "worker_faults" + __table_args__ = ( + Index("ix_worker_faults_worker_id", "worker_id"), + Index("ix_worker_faults_work_unit_id", "work_unit_id"), + ) + + id: Mapped[uuid.UUID] = mapped_column( + Uuid(as_uuid=True), primary_key=True, default=uuid.uuid4 + ) + worker_id: Mapped[str] = mapped_column(Text, nullable=False) + work_unit_id: Mapped[str] = mapped_column(Text, nullable=False) + challenge_slug: Mapped[str | None] = mapped_column(Text) + detail: Mapped[str | None] = mapped_column(Text) + created_at: Mapped[datetime] = mapped_column( + DateTime(timezone=True), + nullable=False, + server_default=func.now(), + ) + + +class WorkerRequestNonce(Base): + """Replay protection for worker-plane nonces. + + Serves two purposes keyed on ``(hotkey, nonce)``: the miner binding nonce + (``hotkey`` = the miner hotkey) consumed at ``register`` and the request + envelope nonce (``hotkey`` = the worker pubkey) consumed on signed + heartbeat/fleet-read requests. Miner hotkeys and worker pubkeys never + collide, so one table safely dedups both. + """ + + __tablename__ = "worker_request_nonces" + __table_args__ = ( + UniqueConstraint( + "hotkey", "nonce", name="uq_worker_request_nonces_hotkey_nonce" + ), + Index("ix_worker_request_nonces_created_at", "created_at"), + Index("ix_worker_request_nonces_hotkey", "hotkey"), + ) + + id: Mapped[uuid.UUID] = mapped_column( + Uuid(as_uuid=True), primary_key=True, default=uuid.uuid4 + ) + hotkey: Mapped[str] = mapped_column(Text, nullable=False) + nonce: Mapped[str] = mapped_column(Text, nullable=False) + body_hash: Mapped[str] = mapped_column(Text, nullable=False) + created_at: Mapped[datetime] = mapped_column( + DateTime(timezone=True), nullable=False + ) + + +class WorkerAssignment(Base, TimestampMixin): + """A gpu work-unit replica coordinated to a miner-funded worker. + + Distinct from :class:`WorkAssignment` (the validator plane): a single gpu + work unit is replicated to multiple workers bound to DISTINCT owners + (architecture.md sec 3.3 R=2), so this table keys one row PER (work unit, + worker) rather than one row per unit. ``worker_pubkey`` is the authenticated + identity the worker pull/post routes gate on; ``miner_hotkey`` is the owner + used for anti-collusion accounting. The reported result and its + ``ExecutionProof`` (and the extracted ``manifest_sha256`` used for + reconciliation) are stored on the row. + """ + + __tablename__ = "worker_assignments" + __table_args__ = ( + UniqueConstraint( + "work_unit_id", + "worker_id", + name="uq_worker_assignments_work_unit_worker", + ), + Index("ix_worker_assignments_work_unit_id", "work_unit_id"), + Index("ix_worker_assignments_status", "status"), + Index("ix_worker_assignments_worker_pubkey", "worker_pubkey"), + Index( + "ix_worker_assignments_status_worker_pubkey", + "status", + "worker_pubkey", + ), + Index("ix_worker_assignments_status_deadline", "status", "deadline_at"), + ) + + id: Mapped[uuid.UUID] = mapped_column( + Uuid(as_uuid=True), primary_key=True, default=uuid.uuid4 + ) + challenge_slug: Mapped[str] = mapped_column(Text, nullable=False) + work_unit_id: Mapped[str] = mapped_column(Text, nullable=False) + submission_ref: Mapped[str] = mapped_column(Text, nullable=False) + worker_id: Mapped[str] = mapped_column(Text, nullable=False) + worker_pubkey: Mapped[str] = mapped_column(Text, nullable=False) + miner_hotkey: Mapped[str] = mapped_column(Text, nullable=False) + payload: Mapped[dict[str, Any]] = mapped_column( + JSON, + nullable=False, + default=dict, + server_default="{}", + ) + required_capability: Mapped[str] = mapped_column( + Text, nullable=False, default="gpu", server_default="gpu" + ) + status: Mapped[WorkAssignmentStatus] = mapped_column( + Enum( + WorkAssignmentStatus, + name="work_assignment_status", + values_callable=lambda obj: [e.value for e in obj], + native_enum=False, + ), + nullable=False, + default=WorkAssignmentStatus.PENDING, + server_default=WorkAssignmentStatus.PENDING.value, + ) + attempt_count: Mapped[int] = mapped_column( + Integer, nullable=False, default=0, server_default="0" + ) + max_attempts: Mapped[int] = mapped_column( + Integer, nullable=False, default=3, server_default="3" + ) + deadline_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True)) + last_progress_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True)) + checkpoint_ref: Mapped[str | None] = mapped_column(Text) + result_success: Mapped[bool | None] = mapped_column(Boolean) + result_payload: Mapped[dict[str, Any] | None] = mapped_column(JSON) + manifest_sha256: Mapped[str | None] = mapped_column(Text) diff --git a/src/base/master/app_proxy.py b/src/base/master/app_proxy.py index 22c19c91..ab33caa5 100644 --- a/src/base/master/app_proxy.py +++ b/src/base/master/app_proxy.py @@ -49,6 +49,19 @@ build_validator_coordination_router, build_validator_health_lifespan, ) +from base.master.worker_assignment import ( + WorkerAssignmentService, + build_worker_assignment_router, +) +from base.master.worker_coordination import ( + WorkerCoordinationService, + build_worker_coordination_router, + build_worker_health_lifespan, +) +from base.master.worker_unit_status import ( + WorkerUnitStatusService, + build_worker_unit_status_router, +) from base.schemas.challenge import ChallengeRecord, ChallengeStatus from base.security.miner_auth import ( MinerAuthError, @@ -60,6 +73,11 @@ ValidatorSignedRequestVerifier, build_validator_auth_dependency, ) +from base.security.worker_auth import ( + WorkerSignedRequestVerifier, + build_internal_bearer_auth, + build_worker_auth_dependency, +) from base.supervisor.challenge_image_updater import ( build_challenge_image_update_lifespan, ) @@ -126,6 +144,12 @@ ClientFactory = Callable[[], AbstractAsyncContextManager[httpx.AsyncClient]] ChallengeTokenProvider = Callable[[str], str] +#: Challenge slug whose prism<->master bridge shared token additionally +#: authenticates the admission fleet-read ``GET /v1/workers/active`` (in addition +#: to the signed-request path). prism reuses this same token as its +#: admission-query bearer (architecture.md sec 3.5). +WORKER_ADMISSION_BRIDGE_SLUG = "prism" + def is_blocked_proxy_path(path: str) -> bool: """Return whether a public proxy path targets a private challenge route.""" @@ -356,6 +380,12 @@ def create_proxy_app( validator_service: ValidatorCoordinationService | None = None, validator_verifier: ValidatorSignedRequestVerifier | None = None, validator_health_interval_seconds: float | None = None, + worker_service: WorkerCoordinationService | None = None, + worker_verifier: WorkerSignedRequestVerifier | None = None, + worker_health_interval_seconds: float | None = None, + worker_assignment_service: WorkerAssignmentService | None = None, + worker_assignment_verifier: WorkerSignedRequestVerifier | None = None, + worker_unit_status_service: WorkerUnitStatusService | None = None, assignment_coordination_service: AssignmentCoordinationService | None = None, llm_gateway_service: LLMGatewayService | None = None, orchestration_driver: MasterOrchestrationDriver | None = None, @@ -384,6 +414,9 @@ def create_proxy_app( build_validator_health_lifespan( validator_service, validator_health_interval_seconds ), + build_worker_health_lifespan( + worker_service, worker_health_interval_seconds + ), build_master_orchestration_lifespan( orchestration_driver, orchestration_interval_seconds ), @@ -677,6 +710,44 @@ async def proxy_path(slug: str, path: str, request: Request) -> Response: ) app.state.validator_coordination_service = validator_service + if worker_service is not None and worker_verifier is not None: + + def _worker_admission_tokens() -> list[str]: + try: + token = token_provider(WORKER_ADMISSION_BRIDGE_SLUG) + except Exception: # noqa: BLE001 - missing token => no internal bearer + return [] + return [token] if token else [] + + app.include_router( + build_worker_coordination_router( + service=worker_service, + auth_dependency=build_worker_auth_dependency(worker_verifier), + internal_auth=build_internal_bearer_auth(_worker_admission_tokens), + ) + ) + app.state.worker_coordination_service = worker_service + + if worker_assignment_service is not None and worker_assignment_verifier is not None: + app.include_router( + build_worker_assignment_router( + service=worker_assignment_service, + auth_dependency=build_worker_auth_dependency( + worker_assignment_verifier + ), + ) + ) + app.state.worker_assignment_service = worker_assignment_service + + if worker_unit_status_service is not None and worker_verifier is not None: + app.include_router( + build_worker_unit_status_router( + service=worker_unit_status_service, + auth_dependency=build_worker_auth_dependency(worker_verifier), + ) + ) + app.state.worker_unit_status_service = worker_unit_status_service + if assignment_coordination_service is not None and validator_verifier is not None: app.include_router( build_assignment_coordination_router( diff --git a/src/base/master/assignment.py b/src/base/master/assignment.py index 67a4522c..eeac1bd4 100644 --- a/src/base/master/assignment.py +++ b/src/base/master/assignment.py @@ -34,6 +34,15 @@ AGENT_CHALLENGE_SLUG = "agent-challenge" PRISM_SLUG = "prism" +#: Work-unit payload key naming the executor kind that must run a unit +#: (architecture.md sec 3.3). Absent/``worker`` gpu units route to the worker +#: replica plane when the flag is on; ``validator`` gpu units (e.g. dispute +#: AUDIT units) route to validators even with the worker plane on. Irrelevant to +#: cpu units, which always route to validators. +EXECUTOR_KIND_PAYLOAD_KEY = "executor_kind" +EXECUTOR_KIND_WORKER = "worker" +EXECUTOR_KIND_VALIDATOR = "validator" + DEFAULT_MAX_ATTEMPTS = 3 #: Payload marker set on a terminally-``failed`` work unit once it has been @@ -107,6 +116,19 @@ def capability_matches( return required in capabilities +def unit_executor_kind(payload: Mapping[str, Any] | None) -> str: + """Return the executor kind a unit's payload requests (defaults to worker). + + A gpu unit with no explicit kind (the common submission unit) is + worker-executed when the worker plane is on; only a unit explicitly marked + ``validator`` (a dispute AUDIT unit) is validator-executed, so it routes to a + validator even while the worker plane owns gpu. + """ + + kind = (payload or {}).get(EXECUTOR_KIND_PAYLOAD_KEY) + return kind if kind == EXECUTOR_KIND_VALIDATOR else EXECUTOR_KIND_WORKER + + class AssignmentService: """Create work units and assign pending ones across online validators.""" @@ -118,6 +140,7 @@ def __init__( gpu_serves_cpu: bool = True, capability_concurrency: Mapping[str, int] | None = None, default_max_attempts: int = DEFAULT_MAX_ATTEMPTS, + worker_plane_capabilities: frozenset[str] = frozenset(), ) -> None: self._session_factory = session_factory self._now_fn = now_fn @@ -128,6 +151,11 @@ def __init__( else capability_concurrency ) self._default_max_attempts = default_max_attempts + #: Capabilities owned by the worker plane: units requiring one of these + #: are NOT routed to validators here (the worker replica engine assigns + #: them). Empty (the default) preserves legacy validator routing exactly, + #: so a flag-OFF master behaves byte-for-byte as before. + self._worker_plane_capabilities = frozenset(worker_plane_capabilities) async def create_agent_challenge_work_units( self, @@ -331,6 +359,8 @@ async def _assign_pending_in_session( for unit in pending: capability = unit.required_capability + if self._worker_plane_owns(capability, unit.payload): + continue limit = self._capability_concurrency.get(capability) eligible = [ hotkey @@ -378,6 +408,13 @@ async def reclaim_stale_assignments( touched here; the subsequent :meth:`assign_pending` increments it as part of the reassignment. + Worker-plane units are skipped SYMMETRICALLY with :meth:`assign_pending`: + a unit whose capability is owned by the worker plane and that is not an + explicit ``validator`` executor-kind unit is maintained by the worker + replica engine (its primary is ``ASSIGNED`` with a NULL validator hotkey + by design), so it is neither assigned nor reclaimed here. With the worker + plane off this guard is inert and reclaim is byte-identical to legacy. + When ``session`` is provided the work runs inside the caller's transaction; otherwise a fresh transaction is opened and committed here. """ @@ -416,6 +453,8 @@ async def _reclaim_in_session(self, session: AsyncSession) -> ReclaimOutcome: ) for unit in inflight: + if self._worker_plane_owns(unit.required_capability, unit.payload): + continue hotkey = unit.assigned_validator_hotkey validator_offline = hotkey is None or hotkey not in online deadline = unit.deadline_at @@ -529,3 +568,24 @@ def _capability_matches(self, required: str, capabilities: set[str]) -> bool: return capability_matches( required, capabilities, gpu_serves_cpu=self._gpu_serves_cpu ) + + def _worker_plane_owns( + self, capability: str, payload: Mapping[str, Any] | None + ) -> bool: + """Whether a unit belongs to the worker replica plane, not this plane. + + A unit whose capability is owned by the worker plane and that is not an + explicit ``validator`` executor-kind unit (e.g. a dispute AUDIT unit) is + materialized and maintained by the worker replica engine, which sets the + primary ``ASSIGNED`` with a NULL ``assigned_validator_hotkey`` by design. + Such a unit must be skipped SYMMETRICALLY here: neither assigned to a + validator nor reclaimed as a stale-validator unit (a null hotkey is the + worker-owned marker, not "offline validator => reassignable"). Empty + ``_worker_plane_capabilities`` (flag OFF) makes this always ``False``, so + legacy validator assign/reclaim behavior is byte-identical. + """ + + return ( + capability in self._worker_plane_capabilities + and unit_executor_kind(payload) != EXECUTOR_KIND_VALIDATOR + ) diff --git a/src/base/master/challenge_work_source.py b/src/base/master/challenge_work_source.py index 767e808d..6e3ddbac 100644 --- a/src/base/master/challenge_work_source.py +++ b/src/base/master/challenge_work_source.py @@ -208,7 +208,73 @@ async def fold( ) +class HttpChallengeResultForwarder: + """Forward a reconciled worker result to the challenge over HTTP. + + Realizes :class:`base.master.worker_reconciliation.ChallengeResultForwarder`: + the master posts an accepted (reconciled) worker result to the challenge's + internal result route so the challenge folds exactly one outcome for the unit + (architecture.md sec 3.3). The challenge base URL + bearer token are resolved + from the registry exactly as the fold path does. + """ + + def __init__( + self, + registry: Any, + *, + timeout_seconds: float = 10.0, + retries: int = 3, + transport: httpx.AsyncBaseTransport | None = None, + ) -> None: + self._registry = registry + self._timeout_seconds = timeout_seconds + self._retries = retries + self._transport = transport + + async def forward_result( + self, + *, + challenge_slug: str, + work_unit_id: str, + submission_ref: str, + result_payload: Any, + ) -> None: + record = await _resolve(self._registry.get(challenge_slug)) + token = await _resolve(self._registry.get_token(challenge_slug)) + if not token: + raise RuntimeError( + f"challenge {challenge_slug!r} has no token for result forward" + ) + url = f"{record.internal_base_url.rstrip('/')}/internal/v1/work_units/result" + headers = { + "Authorization": f"Bearer {token}", + "X-Base-Challenge-Slug": challenge_slug, + "Accept": "application/json", + } + body = { + "work_unit_id": work_unit_id, + "submission_ref": submission_ref, + "result": dict(result_payload or {}), + } + last_error = "unknown error" + for _attempt in range(max(self._retries, 1)): + try: + async with httpx.AsyncClient( + timeout=self._timeout_seconds, transport=self._transport + ) as client: + response = await client.post(url, json=body, headers=headers) + response.raise_for_status() + return + except Exception as exc: # noqa: BLE001 - raised after retries exhausted + last_error = str(exc) + raise RuntimeError( + f"failed to forward result for work unit {work_unit_id} on " + f"{challenge_slug}: {last_error}" + ) + + __all__ = [ "HttpChallengeFoldTrigger", + "HttpChallengeResultForwarder", "HttpChallengeWorkSource", ] diff --git a/src/base/master/orchestration.py b/src/base/master/orchestration.py index 784eb5b0..17b27433 100644 --- a/src/base/master/orchestration.py +++ b/src/base/master/orchestration.py @@ -45,6 +45,15 @@ ) from base.master.reassignment import ReassignmentPassResult, run_reassignment_pass from base.master.validator_coordination import ValidatorCoordinationService +from base.master.worker_assignment_engine import ( + WorkerAssignmentEngine, + WorkerEnginePassResult, + run_worker_assignment_pass, +) +from base.master.worker_reconciliation import ( + ReconciliationPassResult, + WorkerReconciliationService, +) from base.schemas.challenge import ChallengeStatus logger = logging.getLogger(__name__) @@ -114,6 +123,12 @@ class OrchestrationPassResult: #: work-unit ids of agent-challenge units folded this pass (newly failed, or #: re-folded after a prior fold attempt failed). folded: list[str] + #: worker-plane engine outcome for this pass, or ``None`` when the worker + #: plane is disabled (flag OFF -> no engine constructed, legacy routing). + worker: WorkerEnginePassResult | None = None + #: worker-plane reconciliation outcome for this pass, or ``None`` when the + #: worker plane is disabled (flag OFF -> no reconciler constructed). + reconciliation: ReconciliationPassResult | None = None class MasterOrchestrationDriver: @@ -126,12 +141,16 @@ def __init__( validator_service: ValidatorCoordinationService, work_source: ChallengeWorkSource, fold_trigger: ChallengeFoldTrigger | None = None, + worker_assignment_engine: WorkerAssignmentEngine | None = None, + worker_reconciler: WorkerReconciliationService | None = None, seed: int | None = None, ) -> None: self._assignment_service = assignment_service self._validator_service = validator_service self._work_source = work_source self._fold_trigger = fold_trigger + self._worker_assignment_engine = worker_assignment_engine + self._worker_reconciler = worker_reconciler self._seed = seed async def bridge_pending_work(self) -> dict[str, list[str]]: @@ -174,7 +193,13 @@ async def bridge_pending_work(self) -> dict[str, list[str]]: return bridged async def run_once(self) -> OrchestrationPassResult: - """Bridge pending work, run the reassignment pass, then fold dead units.""" + """Bridge work, reassign, replicate + reconcile worker units, then fold. + + When the worker plane is on, the worker assignment/reassignment pass runs + (materializing gpu replicas) and reconciliation then folds reported + replicas into accept/dispute outcomes; both are ``None`` with the flag + off (legacy validator routing, byte-identical). + """ bridged = await self.bridge_pending_work() reassignment = await run_reassignment_pass( @@ -182,11 +207,22 @@ async def run_once(self) -> OrchestrationPassResult: assignment_service=self._assignment_service, seed=self._seed, ) + worker: WorkerEnginePassResult | None = None + if self._worker_assignment_engine is not None: + worker = await run_worker_assignment_pass( + engine=self._worker_assignment_engine, + seed=self._seed, + ) + reconciliation: ReconciliationPassResult | None = None + if self._worker_reconciler is not None: + reconciliation = await self._worker_reconciler.reconcile_once() folded = await self._fold_failed() return OrchestrationPassResult( bridged=bridged, reassignment=reassignment, folded=folded, + worker=worker, + reconciliation=reconciliation, ) async def _fold_failed(self) -> list[str]: diff --git a/src/base/master/worker_assignment.py b/src/base/master/worker_assignment.py new file mode 100644 index 00000000..50bd825e --- /dev/null +++ b/src/base/master/worker_assignment.py @@ -0,0 +1,406 @@ +"""Master worker assignment plane: worker-authenticated pull + result. + +The worker-facing counterpart to :mod:`base.master.assignment_coordination`. +Where the validator plane pulls from ``work_assignments`` (one row per unit), the +worker plane pulls from ``worker_assignments`` (one replica row per (unit, +worker)) so a gpu unit can be replicated across distinct-owner workers. Every +route authenticates as the WORKER keypair via the worker signed-request +dependency and gates on registration/liveness -- never on a metagraph validator +permit (VAL-AGENT-018): + +* ``POST /v1/workers/assignments/pull`` returns the ACTIVE worker's assigned/ + running gpu replicas (transitioning ``assigned`` -> ``running`` with a lease), + and only gpu-capability units (VAL-AGENT-007). A stale/retired worker receives + no new units. +* ``POST /v1/workers/assignments/{id}/result`` persists the reported result and + its ``ExecutionProof`` (VAL-AGENT-008), gated on replica ownership so a late or + foreign post cannot corrupt another worker's replica. + +The engine that CREATES worker replica rows (gpu routing, self-evaluation +exclusion, R=2, reassignment) is layered on top of :meth:`create_worker_assignment` +by the assignment-worker-plane feature. +""" + +from __future__ import annotations + +import uuid +from collections.abc import Callable, Mapping +from dataclasses import dataclass +from datetime import UTC, datetime, timedelta +from typing import Any + +from fastapi import APIRouter, Depends, HTTPException, status +from sqlalchemy import select +from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker + +from base.db.models import ( + WorkAssignmentStatus, + WorkerAssignment, + WorkerRegistration, + WorkerStatus, +) +from base.db.session import session_scope +from base.master.assignment import CAPABILITY_GPU +from base.master.worker_coordination import WorkerCoordinationService +from base.schemas.assignment import ( + AssignmentPullResponse, + AssignmentResultRequest, + AssignmentResultResponse, + AssignmentView, +) +from base.worker.proof import MANIFEST_SHA256_PAYLOAD_KEY, PROOF_PAYLOAD_KEY + +DEFAULT_LEASE_SECONDS = 900 + +_PULLABLE_STATUSES = (WorkAssignmentStatus.ASSIGNED, WorkAssignmentStatus.RUNNING) +_TERMINAL_STATUSES = (WorkAssignmentStatus.COMPLETED, WorkAssignmentStatus.FAILED) + + +class WorkerAssignmentNotFoundError(LookupError): + """No ``worker_assignments`` row with the given id (HTTP 404).""" + + +class WorkerAssignmentOwnershipError(PermissionError): + """The replica is not owned by the calling worker (HTTP 403).""" + + +@dataclass(frozen=True) +class WorkerResultOutcome: + """Outcome of a worker result post (idempotent when already terminal).""" + + status: str + result_ref: str | None + idempotent: bool + + +def _parse_uuid(value: str) -> uuid.UUID | None: + try: + return uuid.UUID(value) + except (ValueError, AttributeError, TypeError): + return None + + +def _extract_manifest_sha256(payload: Mapping[str, Any]) -> str | None: + """Read the manifest hash from a result payload's ExecutionProof envelope.""" + + proof = payload.get(PROOF_PAYLOAD_KEY) + if isinstance(proof, Mapping): + manifest = proof.get(MANIFEST_SHA256_PAYLOAD_KEY) + if isinstance(manifest, str) and manifest: + return manifest + manifest = payload.get(MANIFEST_SHA256_PAYLOAD_KEY) + return manifest if isinstance(manifest, str) and manifest else None + + +def worker_assignment_to_view(assignment: WorkerAssignment) -> AssignmentView: + """Convert a persisted worker replica row to the shared assignment view.""" + + return AssignmentView( + id=str(assignment.id), + challenge_slug=assignment.challenge_slug, + work_unit_id=assignment.work_unit_id, + submission_ref=assignment.submission_ref, + payload=dict(assignment.payload or {}), + required_capability=assignment.required_capability, + status=WorkAssignmentStatus(assignment.status).value, + attempt_count=assignment.attempt_count, + max_attempts=assignment.max_attempts, + deadline_at=assignment.deadline_at, + last_progress_at=assignment.last_progress_at, + checkpoint_ref=assignment.checkpoint_ref, + ) + + +class WorkerAssignmentService: + """Serve pull/result for gpu work-unit replicas coordinated to workers.""" + + def __init__( + self, + session_factory: async_sessionmaker[AsyncSession], + *, + worker_service: WorkerCoordinationService, + lease_seconds: int = DEFAULT_LEASE_SECONDS, + now_fn: Callable[[], datetime] = lambda: datetime.now(UTC), + ) -> None: + self._session_factory = session_factory + self._worker_service = worker_service + self._lease_seconds = lease_seconds + self._now_fn = now_fn + + async def create_worker_assignment( + self, + *, + work_unit_id: str, + challenge_slug: str, + submission_ref: str, + worker_id: str, + worker_pubkey: str, + miner_hotkey: str, + payload: Mapping[str, Any] | None = None, + required_capability: str = CAPABILITY_GPU, + status: WorkAssignmentStatus = WorkAssignmentStatus.ASSIGNED, + attempt_count: int = 1, + max_attempts: int = 3, + checkpoint_ref: str | None = None, + session: AsyncSession | None = None, + ) -> WorkerAssignment: + """Create one replica row binding ``work_unit_id`` to a worker. + + Low-level primitive shared by tests and the assignment engine (which adds + gpu routing, self-evaluation exclusion, and R=2 on top). The engine is + responsible for eligibility; this only persists the row. When ``session`` + is provided the row is added to the caller's transaction (the caller + commits) so a whole replica-creation pass stays atomic; otherwise a fresh + transaction is opened and committed here. + """ + + if session is not None: + return await self._create_in_session( + session, + work_unit_id=work_unit_id, + challenge_slug=challenge_slug, + submission_ref=submission_ref, + worker_id=worker_id, + worker_pubkey=worker_pubkey, + miner_hotkey=miner_hotkey, + payload=payload, + required_capability=required_capability, + status=status, + attempt_count=attempt_count, + max_attempts=max_attempts, + checkpoint_ref=checkpoint_ref, + ) + async with session_scope(self._session_factory) as own_session: + return await self._create_in_session( + own_session, + work_unit_id=work_unit_id, + challenge_slug=challenge_slug, + submission_ref=submission_ref, + worker_id=worker_id, + worker_pubkey=worker_pubkey, + miner_hotkey=miner_hotkey, + payload=payload, + required_capability=required_capability, + status=status, + attempt_count=attempt_count, + max_attempts=max_attempts, + checkpoint_ref=checkpoint_ref, + ) + + async def _create_in_session( + self, + session: AsyncSession, + *, + work_unit_id: str, + challenge_slug: str, + submission_ref: str, + worker_id: str, + worker_pubkey: str, + miner_hotkey: str, + payload: Mapping[str, Any] | None, + required_capability: str, + status: WorkAssignmentStatus, + attempt_count: int, + max_attempts: int, + checkpoint_ref: str | None, + ) -> WorkerAssignment: + now = self._now_fn() + row = WorkerAssignment( + challenge_slug=challenge_slug, + work_unit_id=work_unit_id, + submission_ref=submission_ref, + worker_id=worker_id, + worker_pubkey=worker_pubkey, + miner_hotkey=miner_hotkey, + payload=dict(payload or {}), + required_capability=required_capability, + status=status, + attempt_count=attempt_count, + max_attempts=max_attempts, + checkpoint_ref=checkpoint_ref, + created_at=now, + updated_at=now, + ) + session.add(row) + await session.flush() + return row + + async def pull(self, *, worker_pubkey: str) -> list[WorkerAssignment]: + """Return the ACTIVE worker's assigned/running gpu replicas. + + Transitions ``assigned`` replicas to ``running`` with a fresh lease. A + worker that is not currently ``active`` (pending/stale/retired) receives + no units, and only ``gpu``-capability replicas are ever returned. + """ + + now = self._now_fn() + lease = timedelta(seconds=self._lease_seconds) + async with session_scope(self._session_factory) as session: + worker = ( + await session.execute( + select(WorkerRegistration).where( + WorkerRegistration.worker_pubkey == worker_pubkey + ) + ) + ).scalar_one_or_none() + if worker is None: + return [] + if ( + self._worker_service.effective_status(worker, now) + != WorkerStatus.ACTIVE + ): + return [] + + rows = ( + ( + await session.execute( + select(WorkerAssignment) + .where( + WorkerAssignment.worker_pubkey == worker_pubkey, + WorkerAssignment.status.in_(_PULLABLE_STATUSES), + ) + .order_by( + WorkerAssignment.created_at, + WorkerAssignment.work_unit_id, + ) + ) + ) + .scalars() + .all() + ) + + pulled: list[WorkerAssignment] = [] + for unit in rows: + if unit.required_capability != CAPABILITY_GPU: + continue + if unit.status == WorkAssignmentStatus.ASSIGNED: + unit.status = WorkAssignmentStatus.RUNNING + unit.deadline_at = now + lease + unit.last_progress_at = now + pulled.append(unit) + return pulled + + async def post_result( + self, + *, + assignment_id: str, + worker_pubkey: str, + success: bool, + payload: Mapping[str, Any] | None = None, + checkpoint_ref: str | None = None, + ) -> WorkerResultOutcome: + """Persist a worker replica result and transition it to terminal state. + + Ownership-gated: a post for a replica assigned to a different worker (a + late or foreign post) is rejected without persisting anything. Idempotent + for an already-terminal replica. The reported ``ExecutionProof``'s + ``manifest_sha256`` is extracted onto the row for reconciliation. + """ + + now = self._now_fn() + parsed = _parse_uuid(assignment_id) + if parsed is None: + raise WorkerAssignmentNotFoundError(assignment_id) + async with session_scope(self._session_factory) as session: + unit = ( + await session.execute( + select(WorkerAssignment).where(WorkerAssignment.id == parsed) + ) + ).scalar_one_or_none() + if unit is None: + raise WorkerAssignmentNotFoundError(assignment_id) + if unit.worker_pubkey != worker_pubkey: + raise WorkerAssignmentOwnershipError(assignment_id) + if WorkAssignmentStatus(unit.status) in _TERMINAL_STATUSES: + return WorkerResultOutcome( + status=WorkAssignmentStatus(unit.status).value, + result_ref=str(unit.id), + idempotent=True, + ) + + result_payload = dict(payload or {}) + unit.result_success = success + unit.result_payload = result_payload + unit.manifest_sha256 = _extract_manifest_sha256(result_payload) + unit.status = ( + WorkAssignmentStatus.COMPLETED + if success + else WorkAssignmentStatus.FAILED + ) + unit.last_progress_at = now + if checkpoint_ref is not None: + unit.checkpoint_ref = checkpoint_ref + return WorkerResultOutcome( + status=WorkAssignmentStatus(unit.status).value, + result_ref=str(unit.id), + idempotent=False, + ) + + +def build_worker_assignment_router( + *, + service: WorkerAssignmentService, + auth_dependency: Callable[..., Any], +) -> APIRouter: + """Build the worker assignment router (pull + result). + + ``auth_dependency`` MUST authenticate the caller as a REGISTERED worker (not + validator): the worker plane never gates on a metagraph validator permit. + """ + + router = APIRouter() + + @router.post("/v1/workers/assignments/pull", response_model=AssignmentPullResponse) + async def pull_worker_assignments( + identity: Any = Depends(auth_dependency), + ) -> AssignmentPullResponse: + units = await service.pull(worker_pubkey=identity.hotkey) + return AssignmentPullResponse( + assignments=[worker_assignment_to_view(unit) for unit in units] + ) + + @router.post( + "/v1/workers/assignments/{assignment_id}/result", + response_model=AssignmentResultResponse, + ) + async def post_worker_result( + assignment_id: str, + payload: AssignmentResultRequest, + identity: Any = Depends(auth_dependency), + ) -> AssignmentResultResponse: + try: + outcome = await service.post_result( + assignment_id=assignment_id, + worker_pubkey=identity.hotkey, + success=payload.success, + payload=payload.payload, + checkpoint_ref=payload.checkpoint_ref, + ) + except WorkerAssignmentNotFoundError as exc: + raise HTTPException( + status_code=status.HTTP_404_NOT_FOUND, + detail="worker assignment not found", + ) from exc + except WorkerAssignmentOwnershipError as exc: + raise HTTPException( + status_code=status.HTTP_403_FORBIDDEN, + detail="worker assignment not owned by caller", + ) from exc + return AssignmentResultResponse( + status=outcome.status, + result_ref=outcome.result_ref, + idempotent=outcome.idempotent, + ) + + return router + + +__all__ = [ + "DEFAULT_LEASE_SECONDS", + "WorkerAssignmentNotFoundError", + "WorkerAssignmentOwnershipError", + "WorkerAssignmentService", + "WorkerResultOutcome", + "build_worker_assignment_router", + "worker_assignment_to_view", +] diff --git a/src/base/master/worker_assignment_engine.py b/src/base/master/worker_assignment_engine.py new file mode 100644 index 00000000..a50a13d2 --- /dev/null +++ b/src/base/master/worker_assignment_engine.py @@ -0,0 +1,529 @@ +"""Worker-plane replica-creation ENGINE (architecture.md sec 3.3). + +Layered on top of :class:`base.master.worker_assignment.WorkerAssignmentService` +(the low-level replica-row primitive ``create_worker_assignment``) and +:class:`base.master.worker_coordination.WorkerCoordinationService` (worker +eligibility via ``effective_status``). All of this is gated behind +``compute.worker_plane_enabled``: with the flag OFF the engine is never +constructed and gpu units route to validators exactly as legacy. + +The engine reads PENDING ``work_assignments`` rows whose ``required_capability`` +is gpu -- the validator :class:`base.master.assignment.AssignmentService` is +configured to SKIP these when the worker plane is on (``worker_plane_capabilities``) +-- and materializes ``worker_assignments`` replica rows for them: + +* **No self-evaluation**: a unit whose ``submission_ref`` (the submitting miner + hotkey) equals a worker's ``miner_hotkey`` is never assigned to that worker, + even when it is the only capacity (the unit simply waits, mirroring the legacy + no-eligible-executor behavior). +* **Replication R=** ``replication_factor`` (default 2) across DISTINCT owner + hotkeys -- two workers of the same owner never both serve one unit. +* **Graceful degradation** to a single replica (with a durably recorded warning + on the unit payload + a log record) when fewer than ``replication_factor`` + eligible distinct owners exist. +* **Per-worker gpu concurrency 1** (mirrors ``DEFAULT_CAPABILITY_CONCURRENCY``): + a worker already holding an in-flight gpu replica is not given a second. +* **Heartbeat TTL** drives eligibility: a stale/retired worker is never chosen. + +:meth:`reassign_stale_replicas` mirrors the legacy validator reassignment pass: +a replica whose lease deadline lapsed, or whose worker went stale/retired, is +moved to a different eligible worker (bounded by ``max_attempts`` per replica -- +exhausted replicas are terminally ``failed`` and never recreated); a unit whose +replicas are ALL exhausted is itself marked ``failed``. +""" + +from __future__ import annotations + +import logging +import random +from collections.abc import Callable, Sequence +from contextlib import AbstractAsyncContextManager +from dataclasses import dataclass, field +from datetime import UTC, datetime + +from sqlalchemy import func, select +from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker + +from base.db.models import ( + WorkAssignment, + WorkAssignmentStatus, + WorkerAssignment, + WorkerRegistration, + WorkerStatus, +) +from base.db.session import session_scope +from base.master.assignment import ( + CAPABILITY_GPU, + DEFAULT_CAPABILITY_CONCURRENCY, + EXECUTOR_KIND_VALIDATOR, + unit_executor_kind, +) +from base.master.worker_assignment import WorkerAssignmentService +from base.master.worker_coordination import WorkerCoordinationService + +logger = logging.getLogger(__name__) + +DEFAULT_REPLICATION_FACTOR = 2 + +#: Payload key stamped onto a primary ``work_assignments`` unit whose worker +#: replication was degraded below ``replication_factor`` (observable warning, +#: alongside the emitted log record). +DEGRADED_REPLICATION_PAYLOAD_KEY = "worker_replication_degraded" + +_ACTIVE_REPLICA_STATUSES = ( + WorkAssignmentStatus.ASSIGNED, + WorkAssignmentStatus.RUNNING, +) + + +@dataclass(frozen=True) +class WorkerAssignmentPassResult: + """Observable outcome of one worker-plane assignment pass. + + ``created`` maps each work-unit id to the worker ids it was replicated to + this pass; ``degraded`` lists the units assigned fewer than + ``replication_factor`` replicas (a warning was recorded for each). + """ + + created: dict[str, list[str]] = field(default_factory=dict) + degraded: list[str] = field(default_factory=list) + + +@dataclass(frozen=True) +class WorkerReassignmentPassResult: + """Observable outcome of one worker-plane reassignment pass. + + ``reassigned`` lists the work-unit ids that had a stale replica moved to a + new worker; ``failed_replicas`` lists ``{work_unit_id}:{worker_id}`` pairs + terminally failed (retries exhausted); ``failed_units`` lists the units whose + replicas are ALL exhausted (the unit itself is marked ``failed``). + """ + + reassigned: list[str] = field(default_factory=list) + failed_replicas: list[str] = field(default_factory=list) + failed_units: list[str] = field(default_factory=list) + + +def _as_utc(value: datetime) -> datetime: + return value if value.tzinfo is not None else value.replace(tzinfo=UTC) + + +class WorkerAssignmentEngine: + """Create + maintain gpu work-unit replicas across distinct-owner workers.""" + + def __init__( + self, + session_factory: async_sessionmaker[AsyncSession], + *, + assignment_service: WorkerAssignmentService, + worker_service: WorkerCoordinationService, + replication_factor: int = DEFAULT_REPLICATION_FACTOR, + required_capability: str = CAPABILITY_GPU, + gpu_concurrency: int | None = None, + now_fn: Callable[[], datetime] = lambda: datetime.now(UTC), + ) -> None: + self._session_factory = session_factory + self._assignment_service = assignment_service + self._worker_service = worker_service + self._replication_factor = max(1, replication_factor) + self._required_capability = required_capability + self._gpu_concurrency = ( + DEFAULT_CAPABILITY_CONCURRENCY.get(required_capability, 1) + if gpu_concurrency is None + else gpu_concurrency + ) + self._now_fn = now_fn + + def transaction(self) -> AbstractAsyncContextManager[AsyncSession]: + """Open a single committed transaction over the control-plane DB. + + Used to compose :meth:`reassign_stale_replicas` and + :meth:`assign_pending` into one atomic pass (mirrors + ``AssignmentService.transaction``). + """ + + return session_scope(self._session_factory) + + # -- assignment --------------------------------------------------------- + + async def assign_pending( + self, *, seed: int | None = None, session: AsyncSession | None = None + ) -> WorkerAssignmentPassResult: + """Replicate fresh gpu units across eligible distinct-owner workers. + + A unit that already has any replica row is left untouched (its existing + replicas are maintained by :meth:`reassign_stale_replicas`). A unit with + no eligible distinct-owner worker stays ``pending`` (never lost). + """ + + if session is not None: + return await self._assign_pending_in_session(session, seed=seed) + async with session_scope(self._session_factory) as own_session: + return await self._assign_pending_in_session(own_session, seed=seed) + + async def _assign_pending_in_session( + self, session: AsyncSession, *, seed: int | None + ) -> WorkerAssignmentPassResult: + rng = random.Random(seed) + created: dict[str, list[str]] = {} + degraded: list[str] = [] + + active_workers = await self._active_workers(session) + if not active_workers: + return WorkerAssignmentPassResult(created={}, degraded=[]) + + load = await self._inflight_load(session, active_workers) + + pending_units = ( + ( + await session.execute( + select(WorkAssignment) + .where( + WorkAssignment.required_capability == self._required_capability, + WorkAssignment.status == WorkAssignmentStatus.PENDING, + ) + .order_by(WorkAssignment.created_at, WorkAssignment.work_unit_id) + ) + ) + .scalars() + .all() + ) + + for unit in pending_units: + if unit_executor_kind(unit.payload) == EXECUTOR_KIND_VALIDATOR: + continue + existing = ( + await session.execute( + select(func.count()) + .select_from(WorkerAssignment) + .where( + WorkerAssignment.work_unit_id == unit.work_unit_id, + WorkerAssignment.challenge_slug == unit.challenge_slug, + ) + ) + ).scalar_one() + if existing: + continue + + chosen = self._select_owners( + active_workers, + load=load, + submitter_hotkey=unit.submission_ref, + excluded_owners=set(), + excluded_pubkeys=set(), + count=self._replication_factor, + rng=rng, + ) + if not chosen: + continue + + created_ids: list[str] = [] + for worker in chosen: + replica = await self._assignment_service.create_worker_assignment( + work_unit_id=unit.work_unit_id, + challenge_slug=unit.challenge_slug, + submission_ref=unit.submission_ref, + worker_id=worker.worker_id, + worker_pubkey=worker.worker_pubkey, + miner_hotkey=worker.miner_hotkey, + payload=dict(unit.payload or {}), + required_capability=unit.required_capability, + max_attempts=unit.max_attempts, + checkpoint_ref=unit.checkpoint_ref, + session=session, + ) + created_ids.append(replica.worker_id) + load[worker.worker_pubkey] = load.get(worker.worker_pubkey, 0) + 1 + + unit.status = WorkAssignmentStatus.ASSIGNED + unit.attempt_count = (unit.attempt_count or 0) + 1 + created[unit.work_unit_id] = created_ids + + if len(chosen) < self._replication_factor: + self._record_degradation(unit, replicas=len(chosen)) + degraded.append(unit.work_unit_id) + + await session.flush() + return WorkerAssignmentPassResult(created=created, degraded=degraded) + + # -- reassignment ------------------------------------------------------- + + async def reassign_stale_replicas( + self, *, seed: int | None = None, session: AsyncSession | None = None + ) -> WorkerReassignmentPassResult: + """Reassign lapsed/stale replicas; fail exhausted ones (bounded retries). + + Mirrors the legacy validator reassignment: a replica whose deadline + lapsed or whose worker is no longer ``active`` is moved to another + eligible worker (self-evaluation and distinct-owner rules still apply), + incrementing ``attempt_count``. A replica whose attempts are exhausted is + terminally ``failed`` and never recreated; a unit whose replicas are ALL + failed is itself marked ``failed``. + """ + + if session is not None: + return await self._reassign_in_session(session, seed=seed) + async with session_scope(self._session_factory) as own_session: + return await self._reassign_in_session(own_session, seed=seed) + + async def _reassign_in_session( + self, session: AsyncSession, *, seed: int | None + ) -> WorkerReassignmentPassResult: + rng = random.Random(seed) + now = self._now_fn() + reassigned: list[str] = [] + failed_replicas: list[str] = [] + failed_units: list[str] = [] + + active_workers = await self._active_workers(session) + active_pubkeys = {w.worker_pubkey for w in active_workers} + + replicas = (await session.execute(select(WorkerAssignment))).scalars().all() + load: dict[str, int] = {w.worker_pubkey: 0 for w in active_workers} + for row in replicas: + if ( + WorkAssignmentStatus(row.status) in _ACTIVE_REPLICA_STATUSES + and row.worker_pubkey in load + ): + load[row.worker_pubkey] += 1 + + # Owners currently covering each unit with a non-terminal replica. + covered: dict[str, set[str]] = {} + for row in replicas: + if WorkAssignmentStatus(row.status) in _ACTIVE_REPLICA_STATUSES: + covered.setdefault(row.work_unit_id, set()).add(row.miner_hotkey) + + for row in replicas: + if WorkAssignmentStatus(row.status) not in _ACTIVE_REPLICA_STATUSES: + continue + worker_active = row.worker_pubkey in active_pubkeys + deadline = row.deadline_at + deadline_passed = deadline is not None and _as_utc(deadline) < now + if worker_active and not deadline_passed: + continue + + unit_covered = covered.setdefault(row.work_unit_id, set()) + if (row.attempt_count or 0) >= row.max_attempts: + row.status = WorkAssignmentStatus.FAILED + row.last_progress_at = now + unit_covered.discard(row.miner_hotkey) + if row.worker_pubkey in load and load[row.worker_pubkey] > 0: + load[row.worker_pubkey] -= 1 + failed_replicas.append(f"{row.work_unit_id}:{row.worker_id}") + continue + + replacement = self._select_owners( + active_workers, + load=load, + submitter_hotkey=row.submission_ref, + excluded_owners=unit_covered - {row.miner_hotkey}, + excluded_pubkeys={row.worker_pubkey}, + count=1, + rng=rng, + ) + if not replacement: + continue + new_worker = replacement[0] + if row.worker_pubkey in load and load[row.worker_pubkey] > 0: + load[row.worker_pubkey] -= 1 + unit_covered.discard(row.miner_hotkey) + row.worker_id = new_worker.worker_id + row.worker_pubkey = new_worker.worker_pubkey + row.miner_hotkey = new_worker.miner_hotkey + row.status = WorkAssignmentStatus.ASSIGNED + row.attempt_count = (row.attempt_count or 0) + 1 + row.deadline_at = None + row.last_progress_at = now + load[new_worker.worker_pubkey] = load.get(new_worker.worker_pubkey, 0) + 1 + unit_covered.add(new_worker.miner_hotkey) + reassigned.append(row.work_unit_id) + + failed_units = await self._fail_exhausted_units(session, replicas) + await session.flush() + return WorkerReassignmentPassResult( + reassigned=reassigned, + failed_replicas=failed_replicas, + failed_units=failed_units, + ) + + async def _fail_exhausted_units( + self, session: AsyncSession, replicas: Sequence[WorkerAssignment] + ) -> list[str]: + replicas_by_unit: dict[tuple[str, str], list[WorkerAssignment]] = {} + for row in replicas: + replicas_by_unit.setdefault( + (row.challenge_slug, row.work_unit_id), [] + ).append(row) + + failed_units: list[str] = [] + primary_units = ( + ( + await session.execute( + select(WorkAssignment).where( + WorkAssignment.required_capability == self._required_capability, + WorkAssignment.status.in_(_ACTIVE_REPLICA_STATUSES), + ) + ) + ) + .scalars() + .all() + ) + for unit in primary_units: + unit_replicas = replicas_by_unit.get( + (unit.challenge_slug, unit.work_unit_id) + ) + if not unit_replicas: + continue + if all( + WorkAssignmentStatus(r.status) == WorkAssignmentStatus.FAILED + for r in unit_replicas + ): + unit.status = WorkAssignmentStatus.FAILED + failed_units.append(unit.work_unit_id) + return failed_units + + # -- helpers ------------------------------------------------------------ + + async def _active_workers(self, session: AsyncSession) -> list[WorkerRegistration]: + now = self._now_fn() + rows = ( + ( + await session.execute( + select(WorkerRegistration).order_by( + WorkerRegistration.created_at, + WorkerRegistration.worker_id, + ) + ) + ) + .scalars() + .all() + ) + return [ + row + for row in rows + if self._worker_service.effective_status(row, now) == WorkerStatus.ACTIVE + ] + + async def _inflight_load( + self, session: AsyncSession, active_workers: list[WorkerRegistration] + ) -> dict[str, int]: + load = {worker.worker_pubkey: 0 for worker in active_workers} + counts = ( + await session.execute( + select(WorkerAssignment.worker_pubkey, func.count()) + .where( + WorkerAssignment.required_capability == self._required_capability, + WorkerAssignment.status.in_(_ACTIVE_REPLICA_STATUSES), + ) + .group_by(WorkerAssignment.worker_pubkey) + ) + ).all() + for pubkey, count in counts: + if pubkey in load: + load[pubkey] = count + return load + + def _select_owners( + self, + active_workers: list[WorkerRegistration], + *, + load: dict[str, int], + submitter_hotkey: str, + excluded_owners: set[str], + excluded_pubkeys: set[str], + count: int, + rng: random.Random, + ) -> list[WorkerRegistration]: + """Pick up to ``count`` distinct-owner workers under the eligibility rules. + + Excludes the submitter's own workers (self-evaluation), owners already + covering the unit, the current replica's worker, and any worker at its + gpu concurrency cap. One worker per owner is considered (the lowest + ``worker_pubkey``), so two workers of the same owner never both count. + Selection is least-loaded-first with seeded-random tie-breaking, so a + fixed seed + identical inputs yield an identical choice. + """ + + by_owner: dict[str, WorkerRegistration] = {} + for worker in active_workers: + if worker.miner_hotkey == submitter_hotkey: + continue + if worker.miner_hotkey in excluded_owners: + continue + if worker.worker_pubkey in excluded_pubkeys: + continue + if load.get(worker.worker_pubkey, 0) >= self._gpu_concurrency: + continue + current = by_owner.get(worker.miner_hotkey) + if current is None or worker.worker_pubkey < current.worker_pubkey: + by_owner[worker.miner_hotkey] = worker + + chosen: list[WorkerRegistration] = [] + remaining = dict(by_owner) + for _ in range(count): + if not remaining: + break + min_load = min( + load.get(worker.worker_pubkey, 0) for worker in remaining.values() + ) + tied = sorted( + ( + owner + for owner, worker in remaining.items() + if load.get(worker.worker_pubkey, 0) == min_load + ) + ) + owner = rng.choice(tied) + chosen.append(remaining.pop(owner)) + return chosen + + def _record_degradation(self, unit: WorkAssignment, *, replicas: int) -> None: + payload = dict(unit.payload or {}) + payload[DEGRADED_REPLICATION_PAYLOAD_KEY] = replicas + unit.payload = payload + logger.warning( + "worker replication degraded to R=%d for unit %s (wanted R=%d): " + "only %d eligible distinct owner(s) available", + replicas, + unit.work_unit_id, + self._replication_factor, + replicas, + ) + + +@dataclass(frozen=True) +class WorkerEnginePassResult: + """Combined outcome of one full worker-plane engine pass.""" + + assignment: WorkerAssignmentPassResult + reassignment: WorkerReassignmentPassResult + + +async def run_worker_assignment_pass( + *, + engine: WorkerAssignmentEngine, + seed: int | None = None, +) -> WorkerEnginePassResult: + """Run reassignment (reclaim stale replicas) then fresh assignment. + + Both steps run in ONE atomic transaction so a partial failure rolls back + cleanly. Reassignment runs first so a replica freed this pass (worker gone + stale/retired, deadline lapsed) is immediately eligible for a fresh unit's + distinct-owner selection in the same pass. + """ + + async with engine.transaction() as session: + reassignment = await engine.reassign_stale_replicas(seed=seed, session=session) + assignment = await engine.assign_pending(seed=seed, session=session) + return WorkerEnginePassResult(assignment=assignment, reassignment=reassignment) + + +__all__ = [ + "DEFAULT_REPLICATION_FACTOR", + "DEGRADED_REPLICATION_PAYLOAD_KEY", + "WorkerAssignmentEngine", + "WorkerAssignmentPassResult", + "WorkerEnginePassResult", + "WorkerReassignmentPassResult", + "run_worker_assignment_pass", +] diff --git a/src/base/master/worker_coordination.py b/src/base/master/worker_coordination.py new file mode 100644 index 00000000..903799e3 --- /dev/null +++ b/src/base/master/worker_coordination.py @@ -0,0 +1,621 @@ +"""Master worker coordination plane (architecture.md sec 3.3). + +Persists the miner-funded GPU worker registry and its enrollment surface, +mirroring :mod:`base.master.validator_coordination`: + +* ``POST /v1/workers/register`` verifies the miner's sr25519 binding signature + against the (mock) metagraph, protects against binding-nonce replay, refuses a + silent rebind of a pubkey to a different miner hotkey, and creates a + ``pending`` worker (idempotent re-enroll for the same owner). +* ``POST /v1/workers/{worker_id}/heartbeat`` flips ``pending``/``stale`` -> + ``active`` (never resurrects ``retired``). +* ``GET /v1/workers`` is the fleet view (status/owner/provider/last-seen/faults). +* ``GET /v1/workers/active?hotkey=`` is the admission surface (exactly the ACTIVE + workers of a miner hotkey). + +Lifecycle: ``pending -> active -> stale -> retired``. ``active`` requires a +verified binding AND a heartbeat within ``worker_heartbeat_ttl_seconds``; +``retired`` is terminal. +""" + +from __future__ import annotations + +import asyncio +import contextlib +import logging +import uuid +from collections.abc import Callable, Mapping +from contextlib import AbstractAsyncContextManager, asynccontextmanager +from datetime import UTC, datetime, timedelta +from typing import Any, Protocol + +from fastapi import APIRouter, Depends, FastAPI, HTTPException, Query, Request, status +from sqlalchemy import select +from sqlalchemy.exc import IntegrityError +from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker + +from base.db.models import WorkerFault, WorkerRegistration, WorkerStatus +from base.db.session import session_scope +from base.schemas.worker import ( + WorkerFaultView, + WorkerHeartbeatRequest, + WorkerHeartbeatResponse, + WorkerListResponse, + WorkerRegisterRequest, + WorkerRegisterResponse, + WorkerView, +) +from base.security.miner_auth import SignatureVerifier, verify_substrate_signature +from base.security.validator_auth import body_sha256 +from base.security.worker_auth import ( + WorkerNonceStore, + WorkerReplayError, + worker_binding_message, +) + +logger = logging.getLogger(__name__) + +DEFAULT_WORKER_HEARTBEAT_TTL_SECONDS = 120 + + +class MinerNotInMetagraphError(ValueError): + """Binding miner hotkey is absent from the (mock) metagraph (HTTP 403).""" + + +class WorkerBindingSignatureError(ValueError): + """Binding signature does not verify against the miner hotkey (HTTP 401).""" + + +class WorkerRebindError(ValueError): + """A pubkey re-registration under a different miner hotkey (HTTP 409).""" + + +class WorkerNotRegisteredError(LookupError): + """Heartbeat for a ``worker_id`` without a row (HTTP 404).""" + + +class WorkerOwnershipError(ValueError): + """Heartbeat signer does not own the target ``worker_id`` (HTTP 403).""" + + +class MinerMembership(Protocol): + def is_registered(self, hotkey: str) -> bool: ... + + +def _as_utc(value: datetime) -> datetime: + return value if value.tzinfo is not None else value.replace(tzinfo=UTC) + + +class WorkerCoordinationService: + """Persist worker registration, heartbeat liveness, and lifecycle.""" + + def __init__( + self, + session_factory: async_sessionmaker[AsyncSession], + *, + miner_membership: MinerMembership, + binding_nonce_store: WorkerNonceStore, + signature_verifier: SignatureVerifier = verify_substrate_signature, + heartbeat_ttl_seconds: int = DEFAULT_WORKER_HEARTBEAT_TTL_SECONDS, + now_fn: Callable[[], datetime] = lambda: datetime.now(UTC), + worker_id_factory: Callable[[], str] = lambda: uuid.uuid4().hex, + ) -> None: + self._session_factory = session_factory + self._miner_membership = miner_membership + self._binding_nonce_store = binding_nonce_store + self._signature_verifier = signature_verifier + self.heartbeat_ttl_seconds = heartbeat_ttl_seconds + self._now_fn = now_fn + self._worker_id_factory = worker_id_factory + + def current_time(self) -> datetime: + return self._now_fn() + + def effective_status( + self, worker: WorkerRegistration, now: datetime + ) -> WorkerStatus: + """Derive the reported status from persisted status + heartbeat freshness. + + ``retired`` and ``pending`` are returned verbatim (``retired`` is + terminal; ``pending`` awaits a first heartbeat). An ``active``/``stale`` + worker is reported ``active`` only while its heartbeat is within the TTL, + so a query is correct regardless of whether the background staleness pass + has run yet. + """ + + current = WorkerStatus(worker.status) + if current in (WorkerStatus.RETIRED, WorkerStatus.PENDING): + return current + last = worker.last_heartbeat_at + if last is None: + return WorkerStatus.STALE + if now - _as_utc(last) > timedelta(seconds=self.heartbeat_ttl_seconds): + return WorkerStatus.STALE + return WorkerStatus.ACTIVE + + async def register( + self, + *, + worker_pubkey: str, + miner_hotkey: str, + binding_signature: str, + nonce: str, + provider: str, + provider_instance_ref: str | None, + capabilities: list[str], + last_seen_meta: Mapping[str, Any] | None = None, + ) -> WorkerRegistration: + """Enroll a worker after verifying the miner binding. + + Verifies (1) the miner hotkey is on the metagraph, (2) the sr25519 + binding signature, (3) the binding nonce has not been replayed, then + upserts the ``worker_registrations`` row as ``pending``. A pubkey already + bound to a DIFFERENT miner hotkey raises :class:`WorkerRebindError` (no + silent rebind); the same owner re-enrolling with a fresh nonce updates + the single existing row (idempotent, restart-safe). + """ + + now = self._now_fn() + if not self._miner_membership.is_registered(miner_hotkey): + raise MinerNotInMetagraphError(miner_hotkey) + + message = worker_binding_message( + worker_pubkey=worker_pubkey, miner_hotkey=miner_hotkey, nonce=nonce + ) + if not self._signature_verifier(miner_hotkey, message, binding_signature): + raise WorkerBindingSignatureError(worker_pubkey) + + await self._binding_nonce_store.reserve( + hotkey=miner_hotkey, + nonce=nonce, + body_hash=body_sha256(message), + created_at=now, + ) + + try: + async with session_scope(self._session_factory) as session: + return await self._register_in_session( + session, + now=now, + worker_pubkey=worker_pubkey, + miner_hotkey=miner_hotkey, + binding_signature=binding_signature, + nonce=nonce, + provider=provider, + provider_instance_ref=provider_instance_ref, + capabilities=capabilities, + last_seen_meta=last_seen_meta, + ) + except IntegrityError: + # Lost the first-register race on worker_pubkey: re-run as an update. + async with session_scope(self._session_factory) as session: + return await self._register_in_session( + session, + now=now, + worker_pubkey=worker_pubkey, + miner_hotkey=miner_hotkey, + binding_signature=binding_signature, + nonce=nonce, + provider=provider, + provider_instance_ref=provider_instance_ref, + capabilities=capabilities, + last_seen_meta=last_seen_meta, + ) + + async def _register_in_session( + self, + session: AsyncSession, + *, + now: datetime, + worker_pubkey: str, + miner_hotkey: str, + binding_signature: str, + nonce: str, + provider: str, + provider_instance_ref: str | None, + capabilities: list[str], + last_seen_meta: Mapping[str, Any] | None, + ) -> WorkerRegistration: + existing = ( + await session.execute( + select(WorkerRegistration).where( + WorkerRegistration.worker_pubkey == worker_pubkey + ) + ) + ).scalar_one_or_none() + + if existing is not None: + if existing.miner_hotkey != miner_hotkey: + raise WorkerRebindError(worker_pubkey) + existing.binding_signature = binding_signature + existing.binding_nonce = nonce + existing.provider = provider + existing.provider_instance_ref = provider_instance_ref + existing.capabilities = list(capabilities) + existing.status = WorkerStatus.PENDING + existing.last_heartbeat_at = None + if last_seen_meta is not None: + existing.last_seen_meta = dict(last_seen_meta) + return existing + + worker = WorkerRegistration( + worker_id=self._worker_id_factory(), + worker_pubkey=worker_pubkey, + miner_hotkey=miner_hotkey, + binding_signature=binding_signature, + binding_nonce=nonce, + provider=provider, + provider_instance_ref=provider_instance_ref, + capabilities=list(capabilities), + status=WorkerStatus.PENDING, + last_seen_meta=dict(last_seen_meta or {}), + last_heartbeat_at=None, + created_at=now, + ) + session.add(worker) + return worker + + async def heartbeat( + self, + *, + worker_id: str, + worker_pubkey: str, + last_seen_meta: Mapping[str, Any] | None = None, + ) -> tuple[WorkerRegistration, datetime]: + """Refresh liveness; flip ``pending``/``stale`` -> ``active``. + + ``retired`` is terminal: a heartbeat never resurrects it (recovery needs + a fresh registration). Raises :class:`WorkerNotRegisteredError` for an + unknown ``worker_id`` and :class:`WorkerOwnershipError` when the signer + does not own it. + """ + + now = self._now_fn() + async with session_scope(self._session_factory) as session: + worker = ( + await session.execute( + select(WorkerRegistration).where( + WorkerRegistration.worker_id == worker_id + ) + ) + ).scalar_one_or_none() + if worker is None: + raise WorkerNotRegisteredError(worker_id) + if worker.worker_pubkey != worker_pubkey: + raise WorkerOwnershipError(worker_id) + if WorkerStatus(worker.status) == WorkerStatus.RETIRED: + return worker, now + worker.last_heartbeat_at = now + worker.status = WorkerStatus.ACTIVE + if last_seen_meta is not None: + worker.last_seen_meta = dict(last_seen_meta) + return worker, now + + async def list_workers(self) -> list[WorkerRegistration]: + """Return all workers ordered for stable fleet observability.""" + + async with self._session_factory() as session: + rows = ( + ( + await session.execute( + select(WorkerRegistration).order_by( + WorkerRegistration.created_at, + WorkerRegistration.worker_id, + ) + ) + ) + .scalars() + .all() + ) + return list(rows) + + async def faults_by_worker(self) -> dict[str, list[WorkerFault]]: + """Return each worker's faults keyed by ``worker_id`` (time-ordered).""" + + async with self._session_factory() as session: + rows = ( + ( + await session.execute( + select(WorkerFault).order_by(WorkerFault.created_at) + ) + ) + .scalars() + .all() + ) + grouped: dict[str, list[WorkerFault]] = {} + for fault in rows: + grouped.setdefault(fault.worker_id, []).append(fault) + return grouped + + async def active_workers(self, miner_hotkey: str) -> list[WorkerRegistration]: + """Return exactly the ACTIVE workers bound to ``miner_hotkey``. + + Retired/stale/pending workers and other owners are excluded; staleness is + derived from the heartbeat freshness window at query time. + """ + + now = self._now_fn() + async with self._session_factory() as session: + rows = ( + ( + await session.execute( + select(WorkerRegistration) + .where(WorkerRegistration.miner_hotkey == miner_hotkey) + .order_by( + WorkerRegistration.created_at, + WorkerRegistration.worker_id, + ) + ) + ) + .scalars() + .all() + ) + return [ + row + for row in rows + if self.effective_status(row, now) == WorkerStatus.ACTIVE + ] + + async def detect_stale_workers( + self, *, session: AsyncSession | None = None + ) -> list[str]: + """Persist ``active`` -> ``stale`` for workers past the TTL. + + Edge-triggered: only ``active`` workers are considered. Returns the + ``worker_id``s that transitioned this pass. + """ + + now = self._now_fn() + if session is not None: + return await self._detect_stale_in_session(session, now) + async with session_scope(self._session_factory) as own_session: + return await self._detect_stale_in_session(own_session, now) + + async def _detect_stale_in_session( + self, session: AsyncSession, now: datetime + ) -> list[str]: + ttl = timedelta(seconds=self.heartbeat_ttl_seconds) + transitioned: list[str] = [] + rows = ( + ( + await session.execute( + select(WorkerRegistration).where( + WorkerRegistration.status == WorkerStatus.ACTIVE + ) + ) + ) + .scalars() + .all() + ) + for worker in rows: + last = worker.last_heartbeat_at + if last is None: + continue + if now - _as_utc(last) > ttl: + worker.status = WorkerStatus.STALE + transitioned.append(worker.worker_id) + return transitioned + + +def worker_fault_to_view(fault: WorkerFault) -> WorkerFaultView: + return WorkerFaultView( + work_unit_id=fault.work_unit_id, + challenge_slug=fault.challenge_slug, + detail=fault.detail, + created_at=fault.created_at, + ) + + +def worker_to_view( + worker: WorkerRegistration, + *, + service: WorkerCoordinationService, + now: datetime, + faults: list[WorkerFault] | None = None, +) -> WorkerView: + """Build the fleet view for a worker with its derived (effective) status.""" + + return WorkerView( + worker_id=worker.worker_id, + worker_pubkey=worker.worker_pubkey, + miner_hotkey=worker.miner_hotkey, + provider=worker.provider, + provider_instance_ref=worker.provider_instance_ref, + capabilities=list(worker.capabilities), + status=service.effective_status(worker, now).value, + last_heartbeat_at=worker.last_heartbeat_at, + created_at=worker.created_at, + faults=[worker_fault_to_view(fault) for fault in (faults or [])], + ) + + +def build_worker_coordination_router( + *, + service: WorkerCoordinationService, + auth_dependency: Callable[..., Any], + internal_auth: Callable[[str | None], bool] | None = None, +) -> APIRouter: + """Build the worker coordination router (register/heartbeat/fleet reads). + + ``auth_dependency`` authenticates heartbeat + fleet reads as a registered + worker (or, for reads, an eligible validator). Registration is authenticated + by the miner binding signature carried in the body, not this dependency. + + ``internal_auth`` (when provided) is an ADDITIONAL acceptor for the narrow + admission-support endpoint ``GET /v1/workers/active`` only: prism reuses the + prism<->master bridge shared bearer to read a hotkey's active workers + (architecture.md sec 3.5). It never weakens the signed-request path and is + NOT applied to the full-fleet ``GET /v1/workers`` view, which stays + signed-request only. + """ + + router = APIRouter() + + async def authorize_active_read(request: Request) -> Any: + """Dual-auth the admission fleet-read: internal bearer OR signed request. + + Accepts the prism<->master bridge shared bearer when configured; any + other caller (and any request without a matching bearer) falls through + to the unchanged signed-request path, so a bad/missing bearer with no + valid signature is still rejected 401/403. + """ + + if internal_auth is not None and internal_auth( + request.headers.get("authorization") + ): + return None + return await auth_dependency(request) + + @router.post("/v1/workers/register", response_model=WorkerRegisterResponse) + async def register_worker( + payload: WorkerRegisterRequest, + ) -> WorkerRegisterResponse: + try: + worker = await service.register( + worker_pubkey=payload.worker_pubkey, + miner_hotkey=payload.miner_hotkey, + binding_signature=payload.binding_signature, + nonce=payload.nonce, + provider=payload.provider, + provider_instance_ref=payload.provider_instance_ref, + capabilities=payload.capabilities, + last_seen_meta=payload.last_seen_meta, + ) + except MinerNotInMetagraphError as exc: + raise HTTPException( + status_code=status.HTTP_403_FORBIDDEN, + detail="miner hotkey is not in the metagraph", + ) from exc + except WorkerBindingSignatureError as exc: + raise HTTPException( + status_code=status.HTTP_401_UNAUTHORIZED, + detail="invalid worker binding signature", + ) from exc + except WorkerReplayError as exc: + raise HTTPException( + status_code=status.HTTP_409_CONFLICT, + detail="binding nonce already used", + ) from exc + except WorkerRebindError as exc: + raise HTTPException( + status_code=status.HTTP_409_CONFLICT, + detail="worker pubkey already bound to a different miner hotkey", + ) from exc + now = service.current_time() + return WorkerRegisterResponse( + worker=worker_to_view(worker, service=service, now=now), + heartbeat_ttl_seconds=service.heartbeat_ttl_seconds, + ) + + @router.post( + "/v1/workers/{worker_id}/heartbeat", + response_model=WorkerHeartbeatResponse, + ) + async def heartbeat_worker( + worker_id: str, + payload: WorkerHeartbeatRequest, + identity: Any = Depends(auth_dependency), + ) -> WorkerHeartbeatResponse: + try: + worker, now = await service.heartbeat( + worker_id=worker_id, + worker_pubkey=identity.hotkey, + last_seen_meta=payload.last_seen_meta, + ) + except WorkerNotRegisteredError as exc: + raise HTTPException( + status_code=status.HTTP_404_NOT_FOUND, + detail="worker not registered", + ) from exc + except WorkerOwnershipError as exc: + raise HTTPException( + status_code=status.HTTP_403_FORBIDDEN, + detail="worker id not owned by signer", + ) from exc + return WorkerHeartbeatResponse( + status=service.effective_status(worker, now).value, + now=now, + ) + + @router.get( + "/v1/workers", + response_model=WorkerListResponse, + dependencies=[Depends(auth_dependency)], + ) + async def list_workers() -> WorkerListResponse: + now = service.current_time() + rows = await service.list_workers() + faults = await service.faults_by_worker() + return WorkerListResponse( + workers=[ + worker_to_view( + row, service=service, now=now, faults=faults.get(row.worker_id) + ) + for row in rows + ] + ) + + @router.get( + "/v1/workers/active", + response_model=WorkerListResponse, + dependencies=[Depends(authorize_active_read)], + ) + async def list_active_workers( + hotkey: str = Query(...), + ) -> WorkerListResponse: + now = service.current_time() + rows = await service.active_workers(hotkey) + return WorkerListResponse( + workers=[worker_to_view(row, service=service, now=now) for row in rows] + ) + + return router + + +async def run_worker_health_loop( + service: WorkerCoordinationService, + *, + interval_seconds: float, + shutdown_event: asyncio.Event, +) -> None: + """Run the staleness pass every ``interval_seconds`` until shutdown.""" + + while not shutdown_event.is_set(): + try: + await service.detect_stale_workers() + except Exception: + logger.exception("worker staleness detection pass failed") + try: + await asyncio.wait_for(shutdown_event.wait(), timeout=interval_seconds) + except TimeoutError: + continue + + +def build_worker_health_lifespan( + service: WorkerCoordinationService | None, + interval_seconds: float | None, +) -> Callable[[FastAPI], AbstractAsyncContextManager[None]] | None: + """Build a FastAPI lifespan running the staleness loop (None when disabled).""" + + if service is None or interval_seconds is None or interval_seconds <= 0: + return None + + @asynccontextmanager + async def lifespan(_app: FastAPI) -> Any: + shutdown = asyncio.Event() + task = asyncio.create_task( + run_worker_health_loop( + service, + interval_seconds=interval_seconds, + shutdown_event=shutdown, + ) + ) + try: + yield + finally: + shutdown.set() + task.cancel() + with contextlib.suppress(asyncio.CancelledError): + await task + + return lifespan diff --git a/src/base/master/worker_reconciliation.py b/src/base/master/worker_reconciliation.py new file mode 100644 index 00000000..43a2f283 --- /dev/null +++ b/src/base/master/worker_reconciliation.py @@ -0,0 +1,411 @@ +"""Worker-plane reconciliation, disputes, and audit-fault attribution. + +Architecture.md sec 3.3: once the replicas of a gpu work unit report, the master +reconciles their ``ExecutionProof.manifest_sha256`` values. + +* **Matching hashes** -> the unit is accepted and EXACTLY ONE result is forwarded + to the challenge (single fold); the primary ``work_assignments`` row reaches + ``completed``. +* **Divergent hashes** -> the unit is marked ``disputed`` (a terminal state that + is NEVER forwarded to the challenge, before or after audit), and a validator- + executor AUDIT unit (a fresh ``work_assignments`` row, ``required_capability=gpu``, + ``executor_kind=validator``) is created. The validator replay is authoritative: + when the audit result lands, every worker whose manifest diverged from it gets a + ``worker_faults`` row (its ``worker_registrations.status`` is untouched), and the + disputed unit stays disputed (the submission is left unscored). +* **Single surviving proof** -> when the other replica can never report (its + worker went stale/retired and the replica's retries are exhausted with no other + eligible distinct owner, per the reassignment pass) the unit does NOT hang: it + is accepted from the one proof, forwarded once, and a degradation warning is + recorded. A replica that merely posts ``success=false`` (no proof) never + triggers a dispute -- dispute requires TWO differing manifest hashes. + +Late/foreign result posts cannot corrupt reconciliation: they are rejected by the +ownership gate in :class:`base.master.worker_assignment.WorkerAssignmentService` +before touching any replica state, so reconciliation only ever reads +legitimately-owned replica results. + +All of this is gated behind ``compute.worker_plane_enabled``: with the flag OFF +the reconciler is never constructed and no gpu unit is ever replicated, disputed, +or audited. +""" + +from __future__ import annotations + +import logging +from collections.abc import Callable, Mapping, Sequence +from contextlib import AbstractAsyncContextManager +from dataclasses import dataclass, field +from datetime import UTC, datetime +from typing import Any, Protocol + +from sqlalchemy import select +from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker + +from base.db.models import ( + WorkAssignment, + WorkAssignmentStatus, + WorkerAssignment, + WorkerFault, + WorkResult, +) +from base.db.session import session_scope +from base.master.assignment import ( + CAPABILITY_GPU, + EXECUTOR_KIND_PAYLOAD_KEY, + EXECUTOR_KIND_VALIDATOR, + unit_executor_kind, +) +from base.worker.proof import MANIFEST_SHA256_PAYLOAD_KEY, PROOF_PAYLOAD_KEY + +logger = logging.getLogger(__name__) + +#: Suffix appended to a disputed unit's id to form its validator AUDIT unit id. +AUDIT_WORK_UNIT_SUFFIX = ":audit" + +#: Audit-unit payload key naming the original (disputed) work-unit id. +AUDIT_OF_PAYLOAD_KEY = "audit_of_work_unit_id" +#: Audit-unit payload flag set once its outcome has been folded into faults, so a +#: resolved audit is not reprocessed on later passes. +AUDIT_RESOLVED_PAYLOAD_KEY = "audit_resolved" +#: Primary-unit payload flag recording that a unit was accepted from a single +#: surviving proof (degraded replication) -- the reconciliation warning channel +#: (mirrors the assignment-engine degradation marker for VAL-MASTER-007/017). +RECONCILE_DEGRADED_PAYLOAD_KEY = "worker_reconciliation_degraded" + +_TERMINAL_REPLICA_STATUSES = ( + WorkAssignmentStatus.COMPLETED, + WorkAssignmentStatus.FAILED, +) +_TERMINAL_UNIT_STATUSES = ( + WorkAssignmentStatus.COMPLETED, + WorkAssignmentStatus.FAILED, + WorkAssignmentStatus.DISPUTED, +) + + +def audit_work_unit_id(work_unit_id: str) -> str: + """Deterministic id of the validator AUDIT unit for a disputed unit.""" + + return f"{work_unit_id}{AUDIT_WORK_UNIT_SUFFIX}" + + +def _manifest_from_payload(payload: Mapping[str, Any] | None) -> str | None: + """Read the manifest hash from a result payload's ExecutionProof envelope.""" + + if not payload: + return None + proof = payload.get(PROOF_PAYLOAD_KEY) + if isinstance(proof, Mapping): + manifest = proof.get(MANIFEST_SHA256_PAYLOAD_KEY) + if isinstance(manifest, str) and manifest: + return manifest + manifest = payload.get(MANIFEST_SHA256_PAYLOAD_KEY) + return manifest if isinstance(manifest, str) and manifest else None + + +class ChallengeResultForwarder(Protocol): + """Forward exactly one accepted result to the challenge (fold seam). + + Mirrors :class:`base.master.orchestration.ChallengeFoldTrigger`: the + production HTTP implementation posts to the challenge's internal result route; + tests substitute a fake that counts forwards. + """ + + async def forward_result( + self, + *, + challenge_slug: str, + work_unit_id: str, + submission_ref: str, + result_payload: Mapping[str, Any], + ) -> None: ... + + +@dataclass(frozen=True) +class ReconciliationPassResult: + """Observable outcome of one reconciliation pass. + + ``accepted`` lists units accepted (result forwarded once); ``single_replica`` + is the subset accepted from a single surviving proof (degraded, warned); + ``disputed`` lists units marked disputed this pass; ``audit_units`` maps a + disputed unit id to its created audit unit id; ``faults`` lists + ``{work_unit_id}:{worker_id}`` pairs faulted from audit outcomes this pass. + """ + + accepted: list[str] = field(default_factory=list) + single_replica: list[str] = field(default_factory=list) + disputed: list[str] = field(default_factory=list) + audit_units: dict[str, str] = field(default_factory=dict) + faults: list[str] = field(default_factory=list) + + +class WorkerReconciliationService: + """Reconcile worker replica results; dispute divergence; fold audit faults.""" + + def __init__( + self, + session_factory: async_sessionmaker[AsyncSession], + *, + result_forwarder: ChallengeResultForwarder | None = None, + required_capability: str = CAPABILITY_GPU, + now_fn: Callable[[], datetime] = lambda: datetime.now(UTC), + ) -> None: + self._session_factory = session_factory + self._result_forwarder = result_forwarder + self._required_capability = required_capability + self._now_fn = now_fn + + def transaction(self) -> AbstractAsyncContextManager[AsyncSession]: + """Open a single committed transaction over the control-plane DB.""" + + return session_scope(self._session_factory) + + async def reconcile_once( + self, *, session: AsyncSession | None = None + ) -> ReconciliationPassResult: + """Reconcile all worker-replicated units and fold completed audits. + + Idempotent: an already-accepted unit (primary ``completed``) is never + re-forwarded, a disputed unit is never late-accepted, and a resolved + audit is never re-faulted. + """ + + if session is not None: + return await self._reconcile_in_session(session) + async with session_scope(self._session_factory) as own_session: + return await self._reconcile_in_session(own_session) + + async def _reconcile_in_session( + self, session: AsyncSession + ) -> ReconciliationPassResult: + now = self._now_fn() + accepted: list[str] = [] + single_replica: list[str] = [] + disputed: list[str] = [] + audit_units: dict[str, str] = {} + faults: list[str] = [] + + replicas = (await session.execute(select(WorkerAssignment))).scalars().all() + groups: dict[tuple[str, str], list[WorkerAssignment]] = {} + for replica in replicas: + groups.setdefault( + (replica.challenge_slug, replica.work_unit_id), [] + ).append(replica) + + for (slug, unit_id), rows in sorted(groups.items()): + primary = ( + await session.execute( + select(WorkAssignment).where( + WorkAssignment.challenge_slug == slug, + WorkAssignment.work_unit_id == unit_id, + ) + ) + ).scalar_one_or_none() + if primary is None: + continue + if WorkAssignmentStatus(primary.status) in _TERMINAL_UNIT_STATUSES: + continue + + reported = [ + row + for row in rows + if WorkAssignmentStatus(row.status) == WorkAssignmentStatus.COMPLETED + and row.result_success + and row.manifest_sha256 + ] + distinct = {row.manifest_sha256 for row in reported} + non_terminal = [ + row + for row in rows + if WorkAssignmentStatus(row.status) not in _TERMINAL_REPLICA_STATUSES + ] + + if len(distinct) >= 2: + primary.status = WorkAssignmentStatus.DISPUTED + primary.last_progress_at = now + audit_id = await self._ensure_audit_unit(session, primary, now) + disputed.append(unit_id) + audit_units[unit_id] = audit_id + continue + + # A replica that can still report must be awaited: dispute needs two + # DIFFERING hashes, not a missing one. + if non_terminal: + continue + # No proof at all: the assignment engine terminally fails the unit. + if not reported: + continue + + winner = reported[0] + if not await self._forward(primary, winner): + continue + primary.status = WorkAssignmentStatus.COMPLETED + primary.result_ref = str(winner.id) + primary.last_progress_at = now + accepted.append(unit_id) + if len(reported) < len(rows): + primary.payload = { + **(primary.payload or {}), + RECONCILE_DEGRADED_PAYLOAD_KEY: True, + } + single_replica.append(unit_id) + logger.warning( + "worker unit %s reconciled from a single surviving proof " + "(%d of %d replicas reported); accepted with degraded " + "replication", + unit_id, + len(reported), + len(rows), + ) + + faults = await self._resolve_audits(session, replicas, now) + + await session.flush() + return ReconciliationPassResult( + accepted=accepted, + single_replica=single_replica, + disputed=disputed, + audit_units=audit_units, + faults=faults, + ) + + async def _resolve_audits( + self, + session: AsyncSession, + replicas: Sequence[WorkerAssignment], + now: datetime, + ) -> list[str]: + faults: list[str] = [] + completed = ( + ( + await session.execute( + select(WorkAssignment).where( + WorkAssignment.status == WorkAssignmentStatus.COMPLETED + ) + ) + ) + .scalars() + .all() + ) + for audit in completed: + payload = audit.payload or {} + if unit_executor_kind(payload) != EXECUTOR_KIND_VALIDATOR: + continue + if payload.get(AUDIT_RESOLVED_PAYLOAD_KEY): + continue + original_id = payload.get(AUDIT_OF_PAYLOAD_KEY) + if not original_id: + continue + authoritative = await self._audit_manifest(session, audit) + if authoritative is None: + continue + for replica in replicas: + if ( + replica.challenge_slug != audit.challenge_slug + or replica.work_unit_id != original_id + or not replica.manifest_sha256 + ): + continue + if replica.manifest_sha256 != authoritative: + session.add( + WorkerFault( + worker_id=replica.worker_id, + work_unit_id=original_id, + challenge_slug=audit.challenge_slug, + detail=( + "manifest diverged from validator audit " + f"({replica.manifest_sha256} != {authoritative})" + ), + created_at=now, + ) + ) + faults.append(f"{original_id}:{replica.worker_id}") + audit.payload = {**payload, AUDIT_RESOLVED_PAYLOAD_KEY: True} + return faults + + async def _audit_manifest( + self, session: AsyncSession, audit: WorkAssignment + ) -> str | None: + result = ( + ( + await session.execute( + select(WorkResult) + .where(WorkResult.assignment_id == audit.id) + .order_by(WorkResult.created_at.desc()) + ) + ) + .scalars() + .first() + ) + if result is None: + return None + return _manifest_from_payload(result.payload) + + async def _ensure_audit_unit( + self, session: AsyncSession, primary: WorkAssignment, now: datetime + ) -> str: + audit_id = audit_work_unit_id(primary.work_unit_id) + existing = ( + await session.execute( + select(WorkAssignment).where( + WorkAssignment.challenge_slug == primary.challenge_slug, + WorkAssignment.work_unit_id == audit_id, + ) + ) + ).scalar_one_or_none() + if existing is not None: + return audit_id + payload = { + key: value + for key, value in (primary.payload or {}).items() + if key != RECONCILE_DEGRADED_PAYLOAD_KEY + } + payload[EXECUTOR_KIND_PAYLOAD_KEY] = EXECUTOR_KIND_VALIDATOR + payload[AUDIT_OF_PAYLOAD_KEY] = primary.work_unit_id + session.add( + WorkAssignment( + challenge_slug=primary.challenge_slug, + work_unit_id=audit_id, + submission_ref=primary.submission_ref, + payload=payload, + required_capability=self._required_capability, + status=WorkAssignmentStatus.PENDING, + attempt_count=0, + max_attempts=primary.max_attempts, + checkpoint_ref=primary.checkpoint_ref, + created_at=now, + updated_at=now, + ) + ) + return audit_id + + async def _forward(self, primary: WorkAssignment, winner: WorkerAssignment) -> bool: + if self._result_forwarder is None: + return True + try: + await self._result_forwarder.forward_result( + challenge_slug=primary.challenge_slug, + work_unit_id=primary.work_unit_id, + submission_ref=primary.submission_ref, + result_payload=dict(winner.result_payload or {}), + ) + except Exception: + logger.exception( + "failed to forward accepted result for unit %s", + primary.work_unit_id, + ) + return False + return True + + +__all__ = [ + "AUDIT_OF_PAYLOAD_KEY", + "AUDIT_RESOLVED_PAYLOAD_KEY", + "AUDIT_WORK_UNIT_SUFFIX", + "RECONCILE_DEGRADED_PAYLOAD_KEY", + "ChallengeResultForwarder", + "ReconciliationPassResult", + "WorkerReconciliationService", + "audit_work_unit_id", +] diff --git a/src/base/master/worker_unit_status.py b/src/base/master/worker_unit_status.py new file mode 100644 index 00000000..d95df01a --- /dev/null +++ b/src/base/master/worker_unit_status.py @@ -0,0 +1,215 @@ +"""Read-only worker-plane unit status surface (architecture.md sec 3.3). + +``GET /v1/workers/units`` makes the dispute -> audit -> invalidation -> fault +chain OPERATOR-DISCOVERABLE VIA API ALONE (VAL-CROSS-011): today the disputed +state and the audit unit's executor kind/outcome live only in ``worker_faults`` +detail strings, so the divergence story cannot be reconstructed without reading +the database. This surface exposes, per primary gpu unit: + +* the unit id + status (INCLUDING ``disputed``); +* its replicas (worker_id, owner miner hotkey, posted ``manifest_sha256``, and + whether a proof envelope was posted); +* for a disputed unit, the linked validator AUDIT unit's id, executor kind + (``validator``), and terminal outcome (``pending``/``passed``/ + ``mismatch-resolved``). + +It reads existing control-plane tables only (``work_assignments`` primaries + +their ``worker_assignments`` replicas + the validator audit ``work_assignments`` +row + ``worker_faults``); no schema change is required. + +Auth is the signed-request fleet-read (``CoordinationReadEligibility`` -- a +registered worker OR an eligible validator), IDENTICAL to ``GET /v1/workers`` and +distinct from the admission fleet-read's internal bridge bearer (never accepted +here). The router is mounted only when ``compute.worker_plane_enabled`` is on; +with the flag off it is unmounted (404) and legacy behavior is unchanged. +""" + +from __future__ import annotations + +from collections.abc import Callable, Mapping +from typing import Any + +from fastapi import APIRouter, Depends +from sqlalchemy import select +from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker + +from base.db.models import ( + WorkAssignment, + WorkAssignmentStatus, + WorkerAssignment, + WorkerFault, +) +from base.master.assignment import ( + CAPABILITY_GPU, + EXECUTOR_KIND_VALIDATOR, + unit_executor_kind, +) +from base.master.worker_reconciliation import ( + AUDIT_OF_PAYLOAD_KEY, + AUDIT_RESOLVED_PAYLOAD_KEY, +) +from base.schemas.worker import ( + WorkerAuditUnitView, + WorkerReplicaView, + WorkerUnitStatusListResponse, + WorkerUnitStatusView, +) +from base.worker.proof import PROOF_PAYLOAD_KEY + +#: Terminal audit outcomes exposed for a disputed unit's validator audit. +AUDIT_OUTCOME_PENDING = "pending" +AUDIT_OUTCOME_PASSED = "passed" +AUDIT_OUTCOME_MISMATCH_RESOLVED = "mismatch-resolved" + + +def _replica_has_proof(replica: WorkerAssignment) -> bool: + """Whether the replica posted an ExecutionProof envelope in its result.""" + + payload = replica.result_payload or {} + proof = payload.get(PROOF_PAYLOAD_KEY) + return isinstance(proof, Mapping) and bool(proof) + + +def _replica_to_view(replica: WorkerAssignment) -> WorkerReplicaView: + return WorkerReplicaView( + worker_id=replica.worker_id, + miner_hotkey=replica.miner_hotkey, + status=WorkAssignmentStatus(replica.status).value, + manifest_sha256=replica.manifest_sha256, + has_proof=_replica_has_proof(replica), + ) + + +def _audit_outcome(audit: WorkAssignment, *, faulted: bool) -> str: + """Derive the audit's terminal outcome from its resolution + fault state. + + ``pending`` until the validator replay has been folded into faults + (``AUDIT_RESOLVED_PAYLOAD_KEY``); once resolved, ``mismatch-resolved`` when a + worker fault was attributed for the disputed unit, else ``passed``. + """ + + payload = audit.payload or {} + if not payload.get(AUDIT_RESOLVED_PAYLOAD_KEY): + return AUDIT_OUTCOME_PENDING + return AUDIT_OUTCOME_MISMATCH_RESOLVED if faulted else AUDIT_OUTCOME_PASSED + + +class WorkerUnitStatusService: + """Read the worker-plane unit/replica/audit/fault chain for operators.""" + + def __init__( + self, + session_factory: async_sessionmaker[AsyncSession], + *, + required_capability: str = CAPABILITY_GPU, + ) -> None: + self._session_factory = session_factory + self._required_capability = required_capability + + async def list_units(self) -> list[WorkerUnitStatusView]: + """Return each primary gpu unit with replicas and (if disputed) its audit.""" + + async with self._session_factory() as session: + return await self._list_in_session(session) + + async def _list_in_session( + self, session: AsyncSession + ) -> list[WorkerUnitStatusView]: + gpu_units = ( + ( + await session.execute( + select(WorkAssignment).where( + WorkAssignment.required_capability == self._required_capability + ) + ) + ) + .scalars() + .all() + ) + primaries: list[WorkAssignment] = [] + audits_by_original: dict[tuple[str, str], WorkAssignment] = {} + for row in gpu_units: + payload = row.payload or {} + if unit_executor_kind(payload) == EXECUTOR_KIND_VALIDATOR: + original = payload.get(AUDIT_OF_PAYLOAD_KEY) + if isinstance(original, str) and original: + audits_by_original[(row.challenge_slug, original)] = row + else: + primaries.append(row) + + replicas = (await session.execute(select(WorkerAssignment))).scalars().all() + replicas_by_unit: dict[tuple[str, str], list[WorkerAssignment]] = {} + for replica in replicas: + replicas_by_unit.setdefault( + (replica.challenge_slug, replica.work_unit_id), [] + ).append(replica) + + faults = (await session.execute(select(WorkerFault))).scalars().all() + faulted_units = {(fault.challenge_slug, fault.work_unit_id) for fault in faults} + + views: list[WorkerUnitStatusView] = [] + for primary in sorted( + primaries, key=lambda row: (row.challenge_slug, row.work_unit_id) + ): + key = (primary.challenge_slug, primary.work_unit_id) + unit_replicas = sorted( + replicas_by_unit.get(key, []), + key=lambda row: (row.miner_hotkey, row.worker_id), + ) + audit_view: WorkerAuditUnitView | None = None + if WorkAssignmentStatus(primary.status) == WorkAssignmentStatus.DISPUTED: + audit = audits_by_original.get(key) + if audit is not None: + audit_view = WorkerAuditUnitView( + work_unit_id=audit.work_unit_id, + executor_kind=EXECUTOR_KIND_VALIDATOR, + outcome=_audit_outcome( + audit, + faulted=key in faulted_units, + ), + ) + views.append( + WorkerUnitStatusView( + work_unit_id=primary.work_unit_id, + challenge_slug=primary.challenge_slug, + submission_ref=primary.submission_ref, + status=WorkAssignmentStatus(primary.status).value, + replicas=[_replica_to_view(row) for row in unit_replicas], + audit=audit_view, + ) + ) + return views + + +def build_worker_unit_status_router( + *, + service: WorkerUnitStatusService, + auth_dependency: Callable[..., Any], +) -> APIRouter: + """Build the read-only worker unit-status router (``GET /v1/workers/units``). + + ``auth_dependency`` MUST be the signed-request fleet-read dependency + (``CoordinationReadEligibility``), identical to ``GET /v1/workers``; the + internal bridge bearer is never accepted here. + """ + + router = APIRouter() + + @router.get( + "/v1/workers/units", + response_model=WorkerUnitStatusListResponse, + dependencies=[Depends(auth_dependency)], + ) + async def list_worker_units() -> WorkerUnitStatusListResponse: + return WorkerUnitStatusListResponse(units=await service.list_units()) + + return router + + +__all__ = [ + "AUDIT_OUTCOME_MISMATCH_RESOLVED", + "AUDIT_OUTCOME_PASSED", + "AUDIT_OUTCOME_PENDING", + "WorkerUnitStatusService", + "build_worker_unit_status_router", +] diff --git a/src/base/schemas/worker.py b/src/base/schemas/worker.py new file mode 100644 index 00000000..689f7c15 --- /dev/null +++ b/src/base/schemas/worker.py @@ -0,0 +1,157 @@ +"""API schemas for the master worker coordination plane (architecture.md 3.3).""" + +from __future__ import annotations + +from datetime import datetime +from typing import Any + +from pydantic import BaseModel, Field + + +class WorkerFaultView(BaseModel): + """A fault attributed to a worker, surfaced read-only in the fleet view.""" + + work_unit_id: str + challenge_slug: str | None = None + detail: str | None = None + created_at: datetime + + +class WorkerView(BaseModel): + """Fleet view of a worker: status, owner, provider, last-seen, faults.""" + + worker_id: str + worker_pubkey: str + miner_hotkey: str + provider: str + provider_instance_ref: str | None = None + capabilities: list[str] = Field(default_factory=list) + status: str + last_heartbeat_at: datetime | None = None + created_at: datetime + faults: list[WorkerFaultView] = Field(default_factory=list) + + +class WorkerRegisterRequest(BaseModel): + """Body for ``POST /v1/workers/register``. + + ``binding_signature`` is the miner's sr25519 signature over + ``worker-binding:{worker_pubkey}:{miner_hotkey}:{nonce}``. + """ + + worker_pubkey: str + miner_hotkey: str + binding_signature: str + nonce: str + provider: str + provider_instance_ref: str | None = None + capabilities: list[str] = Field(default_factory=lambda: ["gpu"]) + last_seen_meta: dict[str, Any] | None = None + + +class WorkerRegisterResponse(BaseModel): + """Response for a successful worker registration.""" + + worker: WorkerView + heartbeat_ttl_seconds: int + + +class WorkerHeartbeatRequest(BaseModel): + """Body for ``POST /v1/workers/{worker_id}/heartbeat``.""" + + last_seen_meta: dict[str, Any] | None = None + + +class WorkerHeartbeatResponse(BaseModel): + """Response for a successful worker heartbeat.""" + + status: str + now: datetime + + +class WorkerListResponse(BaseModel): + """Response for ``GET /v1/workers`` and ``GET /v1/workers/active``.""" + + workers: list[WorkerView] = Field(default_factory=list) + + +class WorkerReplicaView(BaseModel): + """One worker replica of a gpu unit, surfaced in the unit-status view.""" + + worker_id: str + miner_hotkey: str + status: str + manifest_sha256: str | None = None + has_proof: bool = False + + +class WorkerAuditUnitView(BaseModel): + """The validator AUDIT unit linked to a disputed gpu unit. + + ``outcome`` is the audit's terminal state: ``pending`` (not yet resolved), + ``mismatch-resolved`` (resolved with a fault attributed to the divergent + worker(s)), or ``passed`` (resolved with no fault). + """ + + work_unit_id: str + executor_kind: str + outcome: str + + +class WorkerUnitStatusView(BaseModel): + """Per primary gpu unit status for operator dispute discovery (VAL-CROSS-011). + + Exposes the unit id + status (INCLUDING ``disputed``), its replicas + (worker/owner/manifest/proof-presence), and, when disputed, the linked + validator audit unit with its terminal outcome -- so the full dispute -> + audit -> invalidation -> fault chain is reconstructable via APIs alone. + """ + + work_unit_id: str + challenge_slug: str + submission_ref: str + status: str + replicas: list[WorkerReplicaView] = Field(default_factory=list) + audit: WorkerAuditUnitView | None = None + + +class WorkerUnitStatusListResponse(BaseModel): + """Response for ``GET /v1/workers/units``.""" + + units: list[WorkerUnitStatusView] = Field(default_factory=list) + + +class ProviderInfo(BaseModel): + """Provider/pod identity carried by an ExecutionProof (architecture 3.4).""" + + name: str + executor_id: str | None = None + pod_id: str | None = None + miner_hotkey: str | None = None + + +class WorkerSignature(BaseModel): + """The worker's sr25519 signature binding a manifest hash to a work unit.""" + + worker_pubkey: str + sig: str + + +class ExecutionProof(BaseModel): + """Proof envelope attached to every worker result (architecture 3.4). + + Tier 0 (mandatory, all backends) carries the deterministic + ``manifest_sha256`` plus the worker's sr25519 ``worker_signature`` over the + pinned message (``sha256(f"{manifest_sha256}:{unit_id}")``). Tier 1 adds + ``image_digest`` + a populated ``provider`` block; tier 2 adds a non-null + ``attestation``. The base worker plane emits tier 0; prism fills the higher + tiers. + """ + + version: int = 1 + tier: int = 0 + manifest_sha256: str + image_digest: str | None = None + provider: ProviderInfo | None = None + worker_signature: WorkerSignature + attestation: dict[str, Any] | None = None diff --git a/src/base/security/worker_auth.py b/src/base/security/worker_auth.py new file mode 100644 index 00000000..85de4ef3 --- /dev/null +++ b/src/base/security/worker_auth.py @@ -0,0 +1,361 @@ +"""Worker-plane auth: miner binding verification + worker signed requests. + +Two distinct authentications back the master worker surface (architecture.md +sec 3.3): + +* **Registration** is authenticated by the MINER's sr25519 signature over the + binding message ``worker-binding:{worker_pubkey}:{miner_hotkey}:{nonce}``, + verified against the (mock) metagraph. See :func:`worker_binding_message` and + :class:`MetagraphMinerMembership`. +* **Heartbeat + fleet reads** are authenticated by a hotkey-signed request in the + canonical scheme shared with the validator plane + (``METHOD\\nPATH\\nTIMESTAMP\\nNONCE\\nSHA256_HEX(body)``). The signer is the + WORKER keypair (heartbeat) or any registered worker / eligible validator (fleet + reads); the gate is registration status, NOT a metagraph validator permit. +""" + +from __future__ import annotations + +import hmac +import time +import uuid +from collections.abc import Callable, Iterable, Mapping +from dataclasses import dataclass +from datetime import UTC, datetime, timedelta +from typing import Protocol + +from fastapi import HTTPException, Request, status +from sqlalchemy import delete, select +from sqlalchemy.exc import IntegrityError +from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker + +from base.bittensor.metagraph_cache import MetagraphCache +from base.db.models import WorkerRegistration, WorkerRequestNonce +from base.db.session import session_scope +from base.security.challenge_auth import bearer_token +from base.security.miner_auth import SignatureVerifier, verify_substrate_signature +from base.security.validator_auth import ( + body_sha256, + canonical_validator_request, +) + +HOTKEY_HEADER = "X-Hotkey" +SIGNATURE_HEADER = "X-Signature" +NONCE_HEADER = "X-Nonce" +TIMESTAMP_HEADER = "X-Timestamp" + +WORKER_BINDING_MESSAGE_PREFIX = "worker-binding" + +INVALID_AUTH_DETAIL = "invalid signed request" +REPLAY_AUTH_DETAIL = "replayed request" +INELIGIBLE_DETAIL = "worker not eligible" + + +def worker_binding_message( + *, worker_pubkey: str, miner_hotkey: str, nonce: str +) -> bytes: + """Canonical miner binding message signed by the miner hotkey (sr25519).""" + + return ( + f"{WORKER_BINDING_MESSAGE_PREFIX}:{worker_pubkey}:{miner_hotkey}:{nonce}" + ).encode() + + +class WorkerAuthError(ValueError): + """Signed-request authentication failure (maps to HTTP 401).""" + + +class WorkerReplayError(WorkerAuthError): + """A previously-seen ``(hotkey, nonce)`` was reused (maps to HTTP 409).""" + + +class WorkerEligibilityError(ValueError): + """Signer is neither a registered worker nor an eligible validator (403).""" + + +@dataclass(frozen=True) +class WorkerIdentity: + """Verified identity of a worker-plane signed request.""" + + hotkey: str + nonce: str + timestamp: int + body_hash: str + canonical_request: str + + +class WorkerNonceStore(Protocol): + async def reserve( + self, + *, + hotkey: str, + nonce: str, + body_hash: str, + created_at: datetime, + ) -> None: ... + + +class WorkerAuthEligibility(Protocol): + async def is_eligible(self, hotkey: str) -> bool: ... + + +@dataclass +class MetagraphMinerMembership: + """Miner membership backed by the metagraph cache (on-graph, no permit). + + The binding is signed by a MINER, so membership only requires the hotkey to + be present in the metagraph snapshot (it need NOT hold a validator permit). + """ + + metagraph_cache: MetagraphCache + + def is_registered(self, hotkey: str) -> bool: + return hotkey in self.metagraph_cache.get() + + +class SqlAlchemyWorkerNonceStore: + """Persisted ``(hotkey, nonce)`` replay guard for worker-plane nonces.""" + + def __init__( + self, + session_factory: async_sessionmaker[AsyncSession], + *, + ttl_seconds: int = 86_400, + ) -> None: + self.session_factory = session_factory + self.ttl_seconds = ttl_seconds + + async def reserve( + self, + *, + hotkey: str, + nonce: str, + body_hash: str, + created_at: datetime, + ) -> None: + cutoff = created_at - timedelta(seconds=self.ttl_seconds) + try: + async with session_scope(self.session_factory) as session: + await session.execute( + delete(WorkerRequestNonce).where( + WorkerRequestNonce.created_at < cutoff + ) + ) + session.add( + WorkerRequestNonce( + id=uuid.uuid4(), + hotkey=hotkey, + nonce=nonce, + body_hash=body_hash, + created_at=created_at, + ) + ) + except IntegrityError as exc: + raise WorkerReplayError("nonce already used") from exc + + +@dataclass +class RegisteredWorkerEligibility: + """Eligible when the signer pubkey has a ``worker_registrations`` row. + + Any lifecycle status qualifies for authentication; status-sensitive gating + (e.g. a retired worker not being resurrected) is enforced by the service so + the caller still gets a meaningful, authenticated response. + """ + + session_factory: async_sessionmaker[AsyncSession] + + async def is_eligible(self, hotkey: str) -> bool: + async with self.session_factory() as session: + found = ( + await session.execute( + select(WorkerRegistration.id).where( + WorkerRegistration.worker_pubkey == hotkey + ) + ) + ).first() + return found is not None + + +@dataclass +class CoordinationReadEligibility: + """Eligible when the signer is a registered worker OR an eligible validator. + + Fleet reads (``GET /v1/workers`` / ``GET /v1/workers/active``) are + authenticated-but-not-admin: any registered coordination identity may read. + """ + + session_factory: async_sessionmaker[AsyncSession] + metagraph_cache: MetagraphCache + + async def is_eligible(self, hotkey: str) -> bool: + self.metagraph_cache.get() + if self.metagraph_cache.is_validator(hotkey): + return True + return await RegisteredWorkerEligibility(self.session_factory).is_eligible( + hotkey + ) + + +class WorkerSignedRequestVerifier: + """Verify a hotkey-signed worker-plane request (heartbeat + fleet reads).""" + + def __init__( + self, + *, + nonce_store: WorkerNonceStore, + eligibility: WorkerAuthEligibility, + signature_verifier: SignatureVerifier = verify_substrate_signature, + ttl_seconds: int = 300, + now_fn: Callable[[], float] = time.time, + ) -> None: + self.nonce_store = nonce_store + self.eligibility = eligibility + self.signature_verifier = signature_verifier + self.ttl_seconds = ttl_seconds + self.now_fn = now_fn + + async def verify( + self, + *, + method: str, + path: str, + query_string: str | bytes, + headers: Mapping[str, str], + body: bytes, + ) -> WorkerIdentity: + hotkey = _required_header(headers, HOTKEY_HEADER) + signature = _required_header(headers, SIGNATURE_HEADER) + nonce = _required_header(headers, NONCE_HEADER) + timestamp_raw = _required_header(headers, TIMESTAMP_HEADER) + + timestamp = _parse_epoch(timestamp_raw) + now = float(self.now_fn()) + if abs(now - timestamp) > self.ttl_seconds: + raise WorkerAuthError("stale signature") + + canonical_request = canonical_validator_request( + method=method, + path=path, + query_string=query_string, + timestamp=timestamp_raw, + nonce=nonce, + body=body, + ) + if not self.signature_verifier(hotkey, canonical_request.encode(), signature): + raise WorkerAuthError("invalid signature") + + if not await self.eligibility.is_eligible(hotkey): + raise WorkerEligibilityError("signer is not an eligible worker identity") + + body_hash = body_sha256(body) + await self.nonce_store.reserve( + hotkey=hotkey, + nonce=nonce, + body_hash=body_hash, + created_at=datetime.fromtimestamp(now, UTC), + ) + return WorkerIdentity( + hotkey=hotkey, + nonce=nonce, + timestamp=int(timestamp), + body_hash=body_hash, + canonical_request=canonical_request, + ) + + +def build_worker_auth_dependency( + verifier: WorkerSignedRequestVerifier, +) -> Callable[[Request], object]: + """Build a FastAPI dependency enforcing worker signed-request auth.""" + + async def authenticate(request: Request) -> WorkerIdentity: + body = await request.body() + try: + return await verifier.verify( + method=request.method, + path=request.url.path, + query_string=request.url.query, + headers=request.headers, + body=body, + ) + except WorkerReplayError as exc: + raise HTTPException( + status_code=status.HTTP_409_CONFLICT, detail=REPLAY_AUTH_DETAIL + ) from exc + except WorkerEligibilityError as exc: + raise HTTPException( + status_code=status.HTTP_403_FORBIDDEN, detail=INELIGIBLE_DETAIL + ) from exc + except WorkerAuthError as exc: + raise HTTPException( + status_code=status.HTTP_401_UNAUTHORIZED, detail=INVALID_AUTH_DETAIL + ) from exc + + return authenticate + + +def build_internal_bearer_auth( + accepted_tokens: Callable[[], Iterable[str]], +) -> Callable[[str | None], bool]: + """Authenticate the admission fleet-read via the internal bridge bearer. + + The prism admission rule queries ``GET /v1/workers/active?hotkey=`` reusing + the SAME shared token base holds for the prism<->master bridge + (architecture.md sec 3.5); prism is neither a registered worker nor an + eligible validator, so it cannot use the signed-request path. This returns a + predicate over the raw ``Authorization`` header that is true only when it + carries ``Bearer `` matching one of ``accepted_tokens`` (compared in + constant time). It is an ADDITIONAL acceptor layered on top of the + signed-request path and never weakens it; ``accepted_tokens`` is resolved per + call so token rotation is picked up without rebuilding the router. + """ + + def is_valid(authorization: str | None) -> bool: + presented = bearer_token(authorization) + if not presented: + return False + return any( + bool(expected) and hmac.compare_digest(presented, expected) + for expected in accepted_tokens() + ) + + return is_valid + + +def _parse_epoch(raw: str) -> float: + import math + + try: + value = float(raw) + except (TypeError, ValueError) as exc: + raise WorkerAuthError("invalid timestamp") from exc + if not math.isfinite(value): + raise WorkerAuthError("invalid timestamp") + return value + + +def _required_header(headers: Mapping[str, str], name: str) -> str: + value = headers.get(name) or headers.get(name.lower()) + if not value or not value.strip(): + raise WorkerAuthError(f"missing {name}") + return value.strip() + + +__all__ = [ + "WORKER_BINDING_MESSAGE_PREFIX", + "CoordinationReadEligibility", + "MetagraphMinerMembership", + "RegisteredWorkerEligibility", + "SqlAlchemyWorkerNonceStore", + "WorkerAuthEligibility", + "WorkerAuthError", + "WorkerEligibilityError", + "WorkerIdentity", + "WorkerNonceStore", + "WorkerReplayError", + "WorkerSignedRequestVerifier", + "build_internal_bearer_auth", + "build_worker_auth_dependency", + "worker_binding_message", +] diff --git a/src/base/validator/agent/runtime.py b/src/base/validator/agent/runtime.py index 45970157..123dbbd3 100644 --- a/src/base/validator/agent/runtime.py +++ b/src/base/validator/agent/runtime.py @@ -13,9 +13,15 @@ import asyncio import logging from collections.abc import Callable, Mapping -from dataclasses import dataclass from typing import Any +from base.coordination.agent_loop import ( + AgentCycleSummary, + BackoffPolicy, + backoff_sleep, + is_transient_error, + sleep_until, +) from base.validator.agent.coordination_client import CoordinationClient from base.validator.agent.executor import ( AssignmentContext, @@ -31,51 +37,6 @@ _DEFAULT_HEARTBEAT_INTERVAL = 60 -@dataclass(frozen=True) -class BackoffPolicy: - """Bounded exponential backoff for transient master/coordination failures. - - The delay grows geometrically with the number of consecutive failures and is - capped at ``max_seconds`` so the agent retries a briefly-unavailable master - without either giving up or busy-looping (architecture.md sec 2.2). - """ - - initial_seconds: float = 1.0 - max_seconds: float = 60.0 - multiplier: float = 2.0 - - def delay(self, consecutive_failures: int) -> float: - """Backoff delay after ``consecutive_failures`` failures (>=1).""" - - if consecutive_failures <= 0: - return 0.0 - raw = self.initial_seconds * (self.multiplier ** (consecutive_failures - 1)) - return min(self.max_seconds, max(0.0, raw)) - - -def _is_transient_error(exc: BaseException) -> bool: - """Whether a coordination failure is worth retrying with backoff. - - Transport errors (no status code) and ``429``/``5xx`` master responses are - transient; a ``4xx`` (e.g. ``403`` ineligible, ``404`` not registered, - ``401`` auth) is a permanent client error that should fail fast. - """ - - status_code = getattr(exc, "status_code", None) - if status_code is None: - return True - return status_code == 429 or status_code >= 500 - - -@dataclass(frozen=True) -class AgentCycleSummary: - """Counts from one assignment-processing pass.""" - - pulled: int - completed: int - failed: int - - class ValidatorAgent: """Coordinated executor: register, heartbeat, pull, execute, post results.""" @@ -134,7 +95,7 @@ async def register(self, shutdown_event: asyncio.Event | None = None) -> int: last_seen_meta=self._meta(), ) except Exception as exc: - if not _is_transient_error(exc): + if not is_transient_error(exc): raise failures += 1 delay = self._backoff.delay(failures) @@ -145,7 +106,7 @@ async def register(self, shutdown_event: asyncio.Event | None = None) -> int: exc, delay, ) - if not await self._backoff_sleep(shutdown_event, delay): + if not await backoff_sleep(shutdown_event, delay): raise continue self._registered_interval = response.heartbeat_interval_seconds @@ -183,7 +144,7 @@ async def run_heartbeat_loop(self, shutdown_event: asyncio.Event) -> None: delay = ( self._backoff.delay(failures) if failures else self.heartbeat_interval ) - await self._sleep_until(shutdown_event, delay) + await sleep_until(shutdown_event, delay) async def run_assignment_loop(self, shutdown_event: asyncio.Event) -> None: failures = 0 @@ -195,7 +156,7 @@ async def run_assignment_loop(self, shutdown_event: asyncio.Event) -> None: failures += 1 logger.exception("validator agent assignment pass failed") delay = self._backoff.delay(failures) if failures else self._poll_interval - await self._sleep_until(shutdown_event, delay) + await sleep_until(shutdown_event, delay) async def run_forever(self, shutdown_event: asyncio.Event | None = None) -> None: shutdown_event = shutdown_event or asyncio.Event() @@ -254,30 +215,6 @@ def _meta(self) -> dict[str, Any]: meta.update(dict(self._last_seen_meta_factory())) return meta - @staticmethod - async def _sleep_until(shutdown_event: asyncio.Event, seconds: float) -> None: - try: - await asyncio.wait_for(shutdown_event.wait(), timeout=seconds) - except TimeoutError: - return - - @staticmethod - async def _backoff_sleep( - shutdown_event: asyncio.Event | None, seconds: float - ) -> bool: - """Sleep ``seconds``; return ``False`` if shutdown fired during the wait.""" - - if shutdown_event is None: - await asyncio.sleep(seconds) - return True - if shutdown_event.is_set(): - return False - try: - await asyncio.wait_for(shutdown_event.wait(), timeout=seconds) - except TimeoutError: - return True - return False - __all__ = [ "AgentCycleSummary", diff --git a/src/base/worker/__init__.py b/src/base/worker/__init__.py new file mode 100644 index 00000000..4d73bedf --- /dev/null +++ b/src/base/worker/__init__.py @@ -0,0 +1,87 @@ +"""Miner-funded GPU worker agent runtime (architecture.md sec 3.2). + +A long-running client that enrolls with the master under a miner-signed binding, +heartbeats, pulls gpu work units, executes them on its own broker via the shared +:class:`AssignmentExecutor` seam, and posts results carrying an ``ExecutionProof`` +envelope. It authenticates as its worker keypair, never as a validator permit. +""" + +from __future__ import annotations + +from base.worker.coordination_client import ( + WorkerCoordinationClient, + WorkerCoordinationClientError, +) +from base.worker.deploy import ( + LOCAL_PROVIDER, + PROVIDER_KEY_ENV, + SUPPORTED_PROVIDERS, + MissingProviderKeyError, + NoOfferWithinBudgetError, + UnsupportedProviderError, + WorkerDeployError, + WorkerImageNotConfiguredError, + build_signed_binding, + build_worker_pod_env, + normalize_provider, + plan_provider_deployment, + rank_worker_offers, + require_provider_api_key, + require_worker_image, + select_worker_offer, +) +from base.worker.executor import ( + StubManifestExecutor, + WorkerProofError, + WorkerProofExecutor, + WorkerProvenance, +) +from base.worker.proof import ( + EXECUTION_PROOF_VERSION, + MANIFEST_SHA256_PAYLOAD_KEY, + PROOF_PAYLOAD_KEY, + build_execution_proof, + execution_proof_signing_payload, + verify_execution_proof, +) +from base.worker.runtime import ( + AgentCycleSummary, + BackoffPolicy, + WorkerAgent, + WorkerBinding, +) + +__all__ = [ + "EXECUTION_PROOF_VERSION", + "LOCAL_PROVIDER", + "MANIFEST_SHA256_PAYLOAD_KEY", + "PROOF_PAYLOAD_KEY", + "PROVIDER_KEY_ENV", + "SUPPORTED_PROVIDERS", + "AgentCycleSummary", + "BackoffPolicy", + "MissingProviderKeyError", + "NoOfferWithinBudgetError", + "StubManifestExecutor", + "UnsupportedProviderError", + "WorkerAgent", + "WorkerBinding", + "WorkerCoordinationClient", + "WorkerCoordinationClientError", + "WorkerDeployError", + "WorkerImageNotConfiguredError", + "WorkerProofError", + "WorkerProofExecutor", + "WorkerProvenance", + "build_execution_proof", + "build_signed_binding", + "build_worker_pod_env", + "execution_proof_signing_payload", + "normalize_provider", + "plan_provider_deployment", + "rank_worker_offers", + "require_provider_api_key", + "require_worker_image", + "select_worker_offer", + "verify_execution_proof", +] diff --git a/src/base/worker/coordination_client.py b/src/base/worker/coordination_client.py new file mode 100644 index 00000000..7cf6edb7 --- /dev/null +++ b/src/base/worker/coordination_client.py @@ -0,0 +1,207 @@ +"""HTTP client for the master worker coordination + assignment plane. + +The worker agent authenticates as its WORKER keypair (NOT a metagraph validator +permit, VAL-AGENT-018): registration carries the miner's binding signature in the +body, while heartbeat/pull/result are hotkey-signed requests in the canonical +scheme shared with the validator plane (the signer here is the worker keypair). +Pull/result reuse the assignment schemas so the executor seam is identical to the +validator agent's. +""" + +from __future__ import annotations + +import json +import time +from collections.abc import Callable, Mapping +from typing import Any +from urllib.parse import urlencode + +import httpx + +from base.schemas.assignment import ( + AssignmentPullResponse, + AssignmentResultResponse, + AssignmentView, +) +from base.schemas.worker import ( + WorkerHeartbeatResponse, + WorkerListResponse, + WorkerRegisterResponse, + WorkerView, +) +from base.validator.agent.signing import RequestSigner, build_signed_headers + + +class WorkerCoordinationClientError(RuntimeError): + """A worker coordination request failed (non-2xx or transport error).""" + + def __init__(self, message: str, *, status_code: int | None = None) -> None: + super().__init__(message) + self.status_code = status_code + + +class WorkerCoordinationClient: + """Async client for the worker-facing register/heartbeat/pull/result routes.""" + + def __init__( + self, + base_url: str, + signer: RequestSigner, + *, + transport: httpx.AsyncBaseTransport | None = None, + timeout_seconds: float = 15.0, + now_fn: Callable[[], float] = time.time, + ) -> None: + self._base_url = base_url.rstrip("/") + self._signer = signer + self._transport = transport + self._timeout = timeout_seconds + self._now_fn = now_fn + + @property + def worker_pubkey(self) -> str: + return self._signer.hotkey + + async def register( + self, + *, + miner_hotkey: str, + binding_signature: str, + nonce: str, + provider: str, + provider_instance_ref: str | None = None, + capabilities: list[str], + last_seen_meta: Mapping[str, Any] | None = None, + ) -> WorkerRegisterResponse: + payload: dict[str, Any] = { + "worker_pubkey": self._signer.hotkey, + "miner_hotkey": miner_hotkey, + "binding_signature": binding_signature, + "nonce": nonce, + "provider": provider, + "provider_instance_ref": provider_instance_ref, + "capabilities": list(capabilities), + } + if last_seen_meta is not None: + payload["last_seen_meta"] = dict(last_seen_meta) + data = await self._post("/v1/workers/register", payload, signed=False) + return WorkerRegisterResponse.model_validate(data) + + async def heartbeat( + self, + *, + worker_id: str, + last_seen_meta: Mapping[str, Any] | None = None, + ) -> WorkerHeartbeatResponse: + payload: dict[str, Any] = {} + if last_seen_meta is not None: + payload["last_seen_meta"] = dict(last_seen_meta) + data = await self._post( + f"/v1/workers/{worker_id}/heartbeat", payload, signed=True + ) + return WorkerHeartbeatResponse.model_validate(data) + + async def pull(self) -> list[AssignmentView]: + data = await self._post("/v1/workers/assignments/pull", {}, signed=True) + return AssignmentPullResponse.model_validate(data).assignments + + async def list_workers(self, *, hotkey: str | None = None) -> list[WorkerView]: + """Read the fleet view (``GET /v1/workers``) as this worker identity. + + With ``hotkey`` set, reads ``GET /v1/workers/active?hotkey=`` (exactly the + ACTIVE workers of that miner hotkey). The request is signed by the worker + keypair; a registered worker is an eligible coordination reader. + """ + + if hotkey is None: + data = await self._get("/v1/workers") + else: + data = await self._get("/v1/workers/active", query={"hotkey": hotkey}) + return WorkerListResponse.model_validate(data).workers + + async def post_result( + self, + assignment_id: str, + *, + success: bool, + payload: Mapping[str, Any] | None = None, + checkpoint_ref: str | None = None, + ) -> AssignmentResultResponse: + body: dict[str, Any] = {"success": success, "payload": dict(payload or {})} + if checkpoint_ref is not None: + body["checkpoint_ref"] = checkpoint_ref + data = await self._post( + f"/v1/workers/assignments/{assignment_id}/result", body, signed=True + ) + return AssignmentResultResponse.model_validate(data) + + async def _post( + self, path: str, payload: Mapping[str, Any], *, signed: bool + ) -> dict[str, Any]: + body = json.dumps(payload, separators=(",", ":")).encode() + headers = {"Content-Type": "application/json"} + if signed: + headers.update( + build_signed_headers( + self._signer, + method="POST", + path=path, + body=body, + now_fn=self._now_fn, + ) + ) + try: + async with self._build_client() as client: + response = await client.post(path, content=body, headers=headers) + except httpx.HTTPError as exc: + raise WorkerCoordinationClientError( + f"worker request to {path} failed: {exc}" + ) from exc + if response.status_code >= 400: + raise WorkerCoordinationClientError( + f"worker request to {path} returned {response.status_code}", + status_code=response.status_code, + ) + return response.json() + + async def _get( + self, path: str, *, query: Mapping[str, str] | None = None + ) -> dict[str, Any]: + query_string = urlencode(sorted(query.items())) if query else "" + request_path = f"{path}?{query_string}" if query_string else path + headers = build_signed_headers( + self._signer, + method="GET", + path=path, + body=b"", + query_string=query_string, + now_fn=self._now_fn, + ) + try: + async with self._build_client() as client: + response = await client.get(request_path, headers=headers) + except httpx.HTTPError as exc: + raise WorkerCoordinationClientError( + f"worker request to {path} failed: {exc}" + ) from exc + if response.status_code >= 400: + raise WorkerCoordinationClientError( + f"worker request to {path} returned {response.status_code}", + status_code=response.status_code, + ) + return response.json() + + def _build_client(self) -> httpx.AsyncClient: + kwargs: dict[str, Any] = { + "base_url": self._base_url, + "timeout": self._timeout, + } + if self._transport is not None: + kwargs["transport"] = self._transport + return httpx.AsyncClient(**kwargs) + + +__all__ = [ + "WorkerCoordinationClient", + "WorkerCoordinationClientError", +] diff --git a/src/base/worker/deploy.py b/src/base/worker/deploy.py new file mode 100644 index 00000000..dcd6a280 --- /dev/null +++ b/src/base/worker/deploy.py @@ -0,0 +1,308 @@ +"""Deploy planning + orchestration for ``base worker deploy`` (architecture 3.2). + +Pure, network-free helpers the ``base worker deploy`` CLI composes: + +* :func:`require_provider_api_key` enforces the provider key env BEFORE any + network call (a missing key is a typed, actionable refusal). +* :func:`select_worker_offer` / :func:`rank_worker_offers` pick a rentable offer + under the ``--max-price`` cap, preferring executors whose ``gpu_count`` matches + the request (a live-learned Lium constraint: renting a partial slice of a + multi-GPU executor 400s), with next-cheapest fallback ordering. +* :func:`build_signed_binding` produces the miner-signed enrollment binding. +* :func:`build_worker_pod_env` builds the env handed to a provisioned pod's agent + and NEVER includes a provider API key (architecture security invariant 1). + +The provider API key lives only in the CLI/agent environment: it authenticates +the provider client's own HTTP calls and is never placed in a pod env, a binding, +or any master-bound request. +""" + +from __future__ import annotations + +import logging +import os +import uuid +from collections.abc import Mapping +from typing import Any + +from base.compute.provider import Offer +from base.compute.worker_deployment import is_loopback_url, is_pinned_digest +from base.security.worker_auth import worker_binding_message +from base.validator.agent.signing import RequestSigner +from base.worker.runtime import WorkerBinding + +logger = logging.getLogger(__name__) + +LOCAL_PROVIDER = "local" +SUPPORTED_PROVIDERS: tuple[str, ...] = ("local", "lium", "targon") + +#: Provider -> the environment variable that MUST hold the miner's provider key. +PROVIDER_KEY_ENV: dict[str, str] = { + "lium": "LIUM_API_KEY", + "targon": "TARGON_API_KEY", +} + + +class WorkerDeployError(RuntimeError): + """A worker deploy could not be completed.""" + + +class UnsupportedProviderError(WorkerDeployError): + """The requested provider is not one of ``local``/``lium``/``targon``.""" + + +class MissingProviderKeyError(WorkerDeployError): + """A provider deploy was attempted without its API key env var set. + + Carries the missing ``env_var`` so the CLI can name it in an actionable + error; raised BEFORE any provider or master network call. + """ + + def __init__(self, provider: str, env_var: str) -> None: + self.provider = provider + self.env_var = env_var + super().__init__( + f"provider '{provider}' requires the {env_var} environment variable; " + "set it and retry. No provider or master call was made." + ) + + +class NoOfferWithinBudgetError(WorkerDeployError): + """No rentable offer is at or below the requested ``--max-price`` cap.""" + + +class WorkerImageNotConfiguredError(WorkerDeployError): + """A provider deploy lacks an explicit, digest-pinned worker image. + + A provider (``lium``/``targon``) deploy MUST pin a PUBLICLY-pullable worker + image; the baked-in placeholder is a PRIVATE-namespace GHCR image that fails + Lium pod creation (``CREATION_FAILED``), so we refuse to silently pin it and + surface an actionable error naming the config keys + env vars instead. + """ + + +def normalize_provider(provider: str) -> str: + """Normalize + validate a ``--provider`` value to a supported provider.""" + + normalized = provider.strip().lower() + if normalized not in SUPPORTED_PROVIDERS: + raise UnsupportedProviderError( + f"unsupported provider '{provider}'; expected one of " + f"{', '.join(SUPPORTED_PROVIDERS)}" + ) + return normalized + + +def require_provider_api_key( + provider: str, *, environ: Mapping[str, str] | None = None +) -> str: + """Return the provider's API key from the env, or raise a typed refusal. + + ``local`` needs no key. For ``lium``/``targon`` a missing/blank key raises + :class:`MissingProviderKeyError` naming the exact env var, BEFORE any network + call is made. + """ + + env = os.environ if environ is None else environ + env_var = PROVIDER_KEY_ENV.get(provider) + if env_var is None: + raise UnsupportedProviderError( + f"provider '{provider}' has no API key requirement" + ) + value = env.get(env_var) + if value is None or not value.strip(): + raise MissingProviderKeyError(provider, env_var) + return value + + +def require_worker_image( + *, image: str | None, image_digest: str | None, provider: str +) -> tuple[str, str]: + """Return the ``(image, digest)`` a provider deploy must pin, or refuse. + + A ``lium``/``targon`` deploy MUST reference a PUBLICLY-pullable, digest-pinned + worker image via ``worker.deploy.image`` + ``worker.deploy.image_digest`` (env + ``BASE_WORKER__DEPLOY__IMAGE`` / ``BASE_WORKER__DEPLOY__IMAGE_DIGEST``). We do + NOT fall back to a baked-in placeholder: a private-namespace GHCR image makes + Lium pod creation fail with ``CREATION_FAILED``, so silently pinning it would + provision a pod that never boots. Raises :class:`WorkerImageNotConfiguredError` + (naming the config keys) when unset or when the digest is not ``sha256:<64 hex>``. + See docs/miner/worker-plane.md ("Publishing the worker image"). + """ + + resolved_image = image.strip() if image else "" + resolved_digest = image_digest.strip() if image_digest else "" + if not resolved_image or not resolved_digest: + raise WorkerImageNotConfiguredError( + f"provider '{provider}' deploy requires an explicit worker image: set " + "worker.deploy.image + worker.deploy.image_digest (env " + "BASE_WORKER__DEPLOY__IMAGE / BASE_WORKER__DEPLOY__IMAGE_DIGEST) to a " + "PUBLICLY-pullable, digest-pinned image. The deploy refuses to pin a " + "baked-in placeholder because a private-namespace GHCR image fails Lium " + "pod creation (CREATION_FAILED). See docs/miner/worker-plane.md." + ) + if not is_pinned_digest(resolved_digest): + raise WorkerImageNotConfiguredError( + "worker.deploy.image_digest must be a pinned digest of the form " + f"'sha256:<64 hex>' (got {resolved_digest!r}); a mutable tag is not an " + "immutable pin. See docs/miner/worker-plane.md." + ) + return resolved_image, resolved_digest + + +def rank_worker_offers( + offers: list[Offer], *, gpu_count: int, max_price: float | None = None +) -> list[Offer]: + """Return offers within budget, best-first for worker deployment. + + Filters to ``price_per_hour <= max_price`` (when a cap is given), then orders + so an executor whose ``gpu_count`` equals the request comes first (renting a + partial slice of a multi-GPU executor 400s on Lium), breaking ties by ascending + price then id. The full ordered list supports next-cheapest fallback when a + rent loses an availability race. + """ + + eligible = [ + offer + for offer in offers + if max_price is None or offer.price_per_hour <= max_price + ] + return sorted( + eligible, + key=lambda offer: ( + offer.gpu_count != gpu_count, + offer.price_per_hour, + offer.id, + ), + ) + + +def select_worker_offer( + offers: list[Offer], *, gpu_count: int, max_price: float | None = None +) -> Offer: + """Select the best in-budget offer, or raise :class:`NoOfferWithinBudgetError`.""" + + ranked = rank_worker_offers(offers, gpu_count=gpu_count, max_price=max_price) + if not ranked: + cap = "unbounded" if max_price is None else f"{max_price}/GPU/hr" + raise NoOfferWithinBudgetError( + f"no rentable offer within budget ({cap}); nothing was provisioned" + ) + return ranked[0] + + +async def plan_provider_deployment( + client: Any, *, gpu_count: int, max_price: float | None +) -> Offer: + """List provider offers under the cap and select the best (no rent issued).""" + + offers = await client.list_offers(max_price_per_hour=max_price) + return select_worker_offer(offers, gpu_count=gpu_count, max_price=max_price) + + +def build_signed_binding( + *, + worker_pubkey: str, + miner_signer: RequestSigner, + nonce: str | None = None, +) -> WorkerBinding: + """Sign the enrollment binding with the miner keypair (fresh nonce default). + + The message is the pinned ``worker-binding:{worker_pubkey}:{miner_hotkey}: + {nonce}`` bytes; ``miner_signer`` is the MINER keypair (its hotkey must be on + the metagraph). A fresh ``nonce`` per call makes a restart's re-enroll + idempotent while a replay is rejected. + """ + + resolved_nonce = nonce or uuid.uuid4().hex + message = worker_binding_message( + worker_pubkey=worker_pubkey, + miner_hotkey=miner_signer.hotkey, + nonce=resolved_nonce, + ) + return WorkerBinding( + miner_hotkey=miner_signer.hotkey, + signature=miner_signer.sign(message), + nonce=resolved_nonce, + ) + + +def build_worker_pod_env( + *, + master_url: str, + provider: str, + binding: WorkerBinding, + worker_key_uri: str | None = None, + worker_key_mnemonic: str | None = None, + broker_url: str | None = None, + gateway_url: str | None = None, + extra: Mapping[str, str] | None = None, +) -> dict[str, str]: + """Build the ``BASE_``-prefixed env for a provisioned pod's worker agent. + + Carries only non-secret coordination config plus the miner-signed binding + (the pod never holds the miner key, so the binding is pre-signed here). It + NEVER contains a provider API key: the key authenticates the CLI's provider + client and must not leave the CLI/agent environment (architecture invariant 1). + + Loopback coordination URLs (master/broker/gateway) are OMITTED: Lium's edge WAF + 403s on any request body carrying a loopback URL and this env is baked into the + WAF-sensitive template body, while the pod config already defaults these to + loopback (the agent resolves them at runtime, reaching a local master via an SSH + tunnel). A real (public) URL is passed through as an override. + """ + + env: dict[str, str] = { + "BASE_COMPUTE__WORKER_PLANE_ENABLED": "true", + "BASE_WORKER__DEPLOY__PROVIDER": provider, + "BASE_WORKER__IDENTITY__MINER_HOTKEY": binding.miner_hotkey, + "BASE_WORKER__IDENTITY__BINDING_SIGNATURE": binding.signature, + "BASE_WORKER__IDENTITY__BINDING_NONCE": binding.nonce, + } + _set_non_loopback(env, "BASE_WORKER__AGENT__MASTER_URL", master_url) + if worker_key_uri: + env["BASE_WORKER__IDENTITY__KEY_URI"] = worker_key_uri + if worker_key_mnemonic: + env["BASE_WORKER__IDENTITY__KEY_MNEMONIC"] = worker_key_mnemonic + if broker_url: + _set_non_loopback(env, "BASE_WORKER__AGENT__BROKER_URL", broker_url) + if gateway_url: + _set_non_loopback(env, "BASE_WORKER__AGENT__GATEWAY_URL", gateway_url) + if extra: + for key, value in extra.items(): + if key not in PROVIDER_KEY_ENV.values(): + env[key] = value + return env + + +def _set_non_loopback(env: dict[str, str], key: str, value: str) -> None: + """Set ``env[key] = value`` unless ``value`` is a WAF-triggering loopback URL.""" + + if is_loopback_url(value): + logger.info( + "omitting loopback URL from worker pod env %s (Lium WAF 403s on loopback " + "URLs; the agent resolves it at runtime from config)", + key, + ) + return + env[key] = value + + +__all__ = [ + "LOCAL_PROVIDER", + "PROVIDER_KEY_ENV", + "SUPPORTED_PROVIDERS", + "MissingProviderKeyError", + "NoOfferWithinBudgetError", + "UnsupportedProviderError", + "WorkerDeployError", + "WorkerImageNotConfiguredError", + "build_signed_binding", + "build_worker_pod_env", + "normalize_provider", + "plan_provider_deployment", + "rank_worker_offers", + "require_provider_api_key", + "require_worker_image", + "select_worker_offer", +] diff --git a/src/base/worker/executor.py b/src/base/worker/executor.py new file mode 100644 index 00000000..82d028d8 --- /dev/null +++ b/src/base/worker/executor.py @@ -0,0 +1,135 @@ +"""Worker-plane execution: guarantee an ExecutionProof on every result. + +The worker agent reuses the validator :class:`AssignmentExecutor` seam to run +work on its OWN local broker. :class:`WorkerProofExecutor` wraps any such +executor so every SUCCESSFUL result carries an ``ExecutionProof`` envelope: it +passes an executor-emitted proof through unchanged, and otherwise builds and +signs a tier-0 proof from the run's ``manifest_sha256`` and the worker's static +provenance. Failed results (e.g. an unreachable broker) pass through untouched -- +a failed unit has no deterministic manifest to bind. + +:class:`StubManifestExecutor` is a CPU-only reference executor (no GPU, no +container) that deterministically "evaluates" a unit into a fixed manifest hash; +it backs the worker unit tests and ``base worker deploy --provider local``. +""" + +from __future__ import annotations + +import hashlib +from dataclasses import dataclass + +from base.schemas.worker import ProviderInfo +from base.validator.agent.executor import ( + AssignmentContext, + AssignmentExecutor, + ExecutionResult, + ProgressCallback, +) +from base.validator.agent.signing import RequestSigner +from base.worker.proof import ( + MANIFEST_SHA256_PAYLOAD_KEY, + PROOF_PAYLOAD_KEY, + build_execution_proof, +) + + +class WorkerProofError(RuntimeError): + """A successful result lacked the manifest needed to build a proof.""" + + +@dataclass(frozen=True) +class WorkerProvenance: + """Static provider/pod identity stamped into every proof a worker emits. + + ``image_digest`` + ``pod_id`` together promote a proof to tier 1 (architecture + sec 3.4); the base worker plane usually has neither, so proofs stay tier 0. + """ + + provider_name: str + miner_hotkey: str + executor_id: str | None = None + pod_id: str | None = None + image_digest: str | None = None + + def provider_info(self) -> ProviderInfo: + return ProviderInfo( + name=self.provider_name, + executor_id=self.executor_id, + pod_id=self.pod_id, + miner_hotkey=self.miner_hotkey, + ) + + def tier(self) -> int: + return 1 if self.image_digest and self.pod_id else 0 + + +class WorkerProofExecutor: + """Wrap an executor so every successful result carries an ExecutionProof.""" + + def __init__( + self, + inner: AssignmentExecutor, + *, + signer: RequestSigner, + provenance: WorkerProvenance, + ) -> None: + self._inner = inner + self._signer = signer + self._provenance = provenance + + async def execute( + self, context: AssignmentContext, *, progress: ProgressCallback + ) -> ExecutionResult: + result = await self._inner.execute(context, progress=progress) + if not result.success: + return result + payload = dict(result.payload) + if PROOF_PAYLOAD_KEY in payload: + return ExecutionResult( + success=True, + payload=payload, + checkpoint_ref=result.checkpoint_ref, + ) + manifest_sha256 = payload.get(MANIFEST_SHA256_PAYLOAD_KEY) + if not manifest_sha256: + raise WorkerProofError( + "successful result is missing manifest_sha256; cannot build " + "an ExecutionProof" + ) + proof = build_execution_proof( + signer=self._signer, + manifest_sha256=str(manifest_sha256), + unit_id=context.assignment.work_unit_id, + tier=self._provenance.tier(), + provider=self._provenance.provider_info(), + image_digest=self._provenance.image_digest, + ) + payload[PROOF_PAYLOAD_KEY] = proof.model_dump(mode="json") + return ExecutionResult( + success=True, payload=payload, checkpoint_ref=result.checkpoint_ref + ) + + +class StubManifestExecutor: + """Deterministic CPU stub: 'evaluate' a unit into a fixed manifest hash. + + The hash is a pure function of the work unit id, so two workers replaying the + same unit agree (reconciliation-friendly). Emits no proof itself -- the + :class:`WorkerProofExecutor` wrapping it produces the tier-0 envelope. + """ + + async def execute( + self, context: AssignmentContext, *, progress: ProgressCallback + ) -> ExecutionResult: + digest = hashlib.sha256(context.assignment.work_unit_id.encode()).hexdigest() + return ExecutionResult( + success=True, payload={MANIFEST_SHA256_PAYLOAD_KEY: digest} + ) + + +__all__ = [ + "StubManifestExecutor", + "WorkerProofError", + "WorkerProofExecutor", + "WorkerProvenance", +] diff --git a/src/base/worker/proof.py b/src/base/worker/proof.py new file mode 100644 index 00000000..3dc02d3e --- /dev/null +++ b/src/base/worker/proof.py @@ -0,0 +1,99 @@ +"""ExecutionProof construction and verification (architecture.md sec 3.4). + +Every worker result carries an ``ExecutionProof`` envelope. The tier-0 core is a +deterministic ``manifest_sha256`` plus the worker's sr25519 signature binding +that hash to the work unit. The signed message format is PINNED identically to +prism (VAL-AGENT-008 / VAL-PRISM-006): the signature is over +``sha256(manifest_sha256 + ":" + unit_id)`` -- the sha256 digest of the UTF-8 +bytes of the string ``{manifest_sha256}:{unit_id}`` -- so a proof produced by the +worker plane verifies with the same code as one produced by prism, and a proof +cannot be replayed across units. +""" + +from __future__ import annotations + +import hashlib +from typing import Any + +from base.schemas.worker import ExecutionProof, ProviderInfo, WorkerSignature +from base.security.miner_auth import SignatureVerifier, verify_substrate_signature +from base.validator.agent.signing import RequestSigner + +EXECUTION_PROOF_VERSION = 1 + +#: Result-payload key carrying the serialized :class:`ExecutionProof`. +PROOF_PAYLOAD_KEY = "execution_proof" +#: Result-payload key carrying the deterministic prism manifest hash. +MANIFEST_SHA256_PAYLOAD_KEY = "manifest_sha256" + + +def execution_proof_signing_payload(*, manifest_sha256: str, unit_id: str) -> bytes: + """The exact bytes an ExecutionProof signature covers (pinned format). + + ``sha256`` digest of the UTF-8 bytes of ``{manifest_sha256}:{unit_id}``. + """ + + message = f"{manifest_sha256}:{unit_id}".encode() + return hashlib.sha256(message).digest() + + +def build_execution_proof( + *, + signer: RequestSigner, + manifest_sha256: str, + unit_id: str, + tier: int = 0, + provider: ProviderInfo | None = None, + image_digest: str | None = None, + attestation: dict[str, Any] | None = None, +) -> ExecutionProof: + """Build and sign an ExecutionProof for ``unit_id`` under ``signer``. + + ``signer`` is the WORKER keypair; its public identity becomes + ``worker_signature.worker_pubkey`` and it signs the pinned message. + """ + + signature = signer.sign( + execution_proof_signing_payload( + manifest_sha256=manifest_sha256, unit_id=unit_id + ) + ) + return ExecutionProof( + version=EXECUTION_PROOF_VERSION, + tier=tier, + manifest_sha256=manifest_sha256, + image_digest=image_digest, + provider=provider, + worker_signature=WorkerSignature(worker_pubkey=signer.hotkey, sig=signature), + attestation=attestation, + ) + + +def verify_execution_proof( + proof: ExecutionProof, + *, + unit_id: str, + signature_verifier: SignatureVerifier = verify_substrate_signature, +) -> bool: + """Whether ``proof``'s worker signature verifies for ``unit_id`` (sr25519). + + Rejects a proof presented with a DIFFERENT ``unit_id`` than the one signed, + so a proof cannot be replayed across units. + """ + + payload = execution_proof_signing_payload( + manifest_sha256=proof.manifest_sha256, unit_id=unit_id + ) + return signature_verifier( + proof.worker_signature.worker_pubkey, payload, proof.worker_signature.sig + ) + + +__all__ = [ + "EXECUTION_PROOF_VERSION", + "MANIFEST_SHA256_PAYLOAD_KEY", + "PROOF_PAYLOAD_KEY", + "build_execution_proof", + "execution_proof_signing_payload", + "verify_execution_proof", +] diff --git a/src/base/worker/runtime.py b/src/base/worker/runtime.py new file mode 100644 index 00000000..65b891ff --- /dev/null +++ b/src/base/worker/runtime.py @@ -0,0 +1,259 @@ +"""Long-running miner-funded worker agent loop (architecture.md sec 3.2). + +Mirrors :class:`base.validator.agent.runtime.ValidatorAgent` (they share the pure +loop primitives in :mod:`base.coordination.agent_loop`) but enrolls as a WORKER: +it registers with the master under a miner-signed binding, heartbeats to stay +``active``, pulls gpu work units assigned to it, executes each via the +:class:`AssignmentExecutor` seam on its OWN local broker, and posts results that +always carry an ``ExecutionProof`` envelope. The agent authenticates every +pull/post as its worker keypair, never as a metagraph validator permit. + +Resilience: an execution failure (e.g. an unreachable local broker) is posted as +a ``success=false`` result and never crashes the loop; the agent keeps +heartbeating and pulling. Registration is an idempotent server-side upsert keyed +on the worker pubkey, so a restart re-registers into the same fleet entry. +""" + +from __future__ import annotations + +import asyncio +import logging +from collections.abc import Callable, Mapping +from dataclasses import dataclass +from typing import Any + +from base.coordination.agent_loop import ( + AgentCycleSummary, + BackoffPolicy, + backoff_sleep, + is_transient_error, + sleep_until, +) +from base.validator.agent.executor import ( + AssignmentContext, + AssignmentExecutor, + BrokerConfig, + gateway_env_for_assignment, +) +from base.worker.coordination_client import WorkerCoordinationClient + +logger = logging.getLogger(__name__) + +_ACTIVE_STATUSES = frozenset({"assigned", "running"}) +_DEFAULT_HEARTBEAT_INTERVAL = 30 + + +@dataclass(frozen=True) +class WorkerBinding: + """The miner-signed binding the agent presents at registration. + + The MINER signs ``worker-binding:{worker_pubkey}:{miner_hotkey}:{nonce}`` + (sr25519) out-of-band (by the deploy CLI); the agent only carries the signed + material. A fresh ``nonce`` is required per registration (a restart supplies + a new one), so a replayed binding is rejected while a same-owner re-enroll + with a fresh nonce is idempotent. + """ + + miner_hotkey: str + signature: str + nonce: str + + +class WorkerAgent: + """Coordinated GPU worker: register, heartbeat, pull, execute, post proofs.""" + + def __init__( + self, + *, + client: WorkerCoordinationClient, + executor: AssignmentExecutor, + broker: BrokerConfig, + binding: WorkerBinding, + provider: str, + provider_instance_ref: str | None = None, + capabilities: list[str] | None = None, + gateway_url: str = "", + heartbeat_interval_seconds: int | None = None, + poll_interval_seconds: float = 5.0, + last_seen_meta_factory: Callable[[], Mapping[str, Any]] | None = None, + backoff: BackoffPolicy | None = None, + ) -> None: + self._client = client + self._executor = executor + self._broker = broker + self._binding = binding + self._provider = provider + self._provider_instance_ref = provider_instance_ref + self._capabilities = list(capabilities or ["gpu"]) + self._gateway_url = gateway_url + self._configured_interval = heartbeat_interval_seconds + self._poll_interval = poll_interval_seconds + self._last_seen_meta_factory = last_seen_meta_factory + self._backoff = backoff or BackoffPolicy() + self._worker_id: str | None = None + self._heartbeat_ttl_seconds: int | None = None + + @property + def worker_pubkey(self) -> str: + return self._client.worker_pubkey + + @property + def worker_id(self) -> str | None: + return self._worker_id + + @property + def heartbeat_interval(self) -> int: + if self._configured_interval is not None: + return self._configured_interval + if self._heartbeat_ttl_seconds: + return max(1, self._heartbeat_ttl_seconds // 2) + return _DEFAULT_HEARTBEAT_INTERVAL + + async def register(self, shutdown_event: asyncio.Event | None = None) -> str: + """Register (idempotent upsert) and resolve the worker id + heartbeat TTL. + + Transient master failures (transport errors / ``429``/``5xx``) are retried + with bounded exponential backoff; a permanent error (``4xx``, e.g. a + forged binding or replayed nonce) fails fast so the caller surfaces it. A + set ``shutdown_event`` aborts the retry loop (re-raising the last error). + """ + + failures = 0 + while True: + try: + response = await self._client.register( + miner_hotkey=self._binding.miner_hotkey, + binding_signature=self._binding.signature, + nonce=self._binding.nonce, + provider=self._provider, + provider_instance_ref=self._provider_instance_ref, + capabilities=self._capabilities, + last_seen_meta=self._meta(), + ) + except Exception as exc: + if not is_transient_error(exc): + raise + failures += 1 + delay = self._backoff.delay(failures) + logger.warning( + "worker agent register attempt %d failed (%s); retrying in %.1fs", + failures, + exc, + delay, + ) + if not await backoff_sleep(shutdown_event, delay): + raise + continue + self._worker_id = response.worker.worker_id + self._heartbeat_ttl_seconds = response.heartbeat_ttl_seconds + return self._worker_id + + async def heartbeat_once(self) -> None: + if self._worker_id is None: + raise RuntimeError("worker agent must register before heartbeating") + await self._client.heartbeat( + worker_id=self._worker_id, last_seen_meta=self._meta() + ) + + async def process_pending_assignments(self) -> AgentCycleSummary: + """Pull, execute, and post results for all currently-assigned units.""" + + assignments = await self._client.pull() + completed = 0 + failed = 0 + for assignment in assignments: + if assignment.status not in _ACTIVE_STATUSES: + continue + if await self._execute_one(assignment): + completed += 1 + else: + failed += 1 + return AgentCycleSummary( + pulled=len(assignments), completed=completed, failed=failed + ) + + async def run_heartbeat_loop(self, shutdown_event: asyncio.Event) -> None: + failures = 0 + while not shutdown_event.is_set(): + try: + await self.heartbeat_once() + failures = 0 + except Exception: + failures += 1 + logger.exception("worker agent heartbeat failed") + delay = ( + self._backoff.delay(failures) if failures else self.heartbeat_interval + ) + await sleep_until(shutdown_event, delay) + + async def run_assignment_loop(self, shutdown_event: asyncio.Event) -> None: + failures = 0 + while not shutdown_event.is_set(): + try: + await self.process_pending_assignments() + failures = 0 + except Exception: + failures += 1 + logger.exception("worker agent assignment pass failed") + delay = self._backoff.delay(failures) if failures else self._poll_interval + await sleep_until(shutdown_event, delay) + + async def run_forever(self, shutdown_event: asyncio.Event | None = None) -> None: + shutdown_event = shutdown_event or asyncio.Event() + await self.register(shutdown_event) + await asyncio.gather( + self.run_heartbeat_loop(shutdown_event), + self.run_assignment_loop(shutdown_event), + ) + + async def _execute_one(self, assignment: Any) -> bool: + gateway_env = gateway_env_for_assignment( + assignment, gateway_url=self._gateway_url + ) + context = AssignmentContext( + assignment=assignment, gateway_env=gateway_env, broker=self._broker + ) + + try: + result = await self._executor.execute(context, progress=_noop_progress) + except Exception as exc: + logger.exception("worker agent execution failed for %s", assignment.id) + await self._client.post_result( + assignment.id, success=False, payload={"error": str(exc)} + ) + return False + + await self._client.post_result( + assignment.id, + success=result.success, + payload=dict(result.payload), + checkpoint_ref=result.checkpoint_ref, + ) + return result.success + + def _meta(self) -> dict[str, Any]: + meta: dict[str, Any] = { + "capabilities": list(self._capabilities), + "broker_url": self._broker.broker_url, + } + if self._last_seen_meta_factory is not None: + meta.update(dict(self._last_seen_meta_factory())) + return meta + + +async def _noop_progress( + *, + checkpoint_ref: str | None = None, + meta: Mapping[str, Any] | None = None, +) -> None: + """Progress heartbeats are not part of the worker plane yet (no-op).""" + + return None + + +__all__ = [ + "AgentCycleSummary", + "BackoffPolicy", + "WorkerAgent", + "WorkerBinding", +] diff --git a/tests/integration/conftest.py b/tests/integration/conftest.py index 3df7b2ac..942f7153 100644 --- a/tests/integration/conftest.py +++ b/tests/integration/conftest.py @@ -26,6 +26,10 @@ "validator_request_nonces", "validator_health_events", "validators", + "worker_faults", + "worker_registrations", + "worker_request_nonces", + "worker_assignments", "work_results", "work_assignments", "llm_usage_records", diff --git a/tests/integration/test_reclaim_guard_symmetry_postgres.py b/tests/integration/test_reclaim_guard_symmetry_postgres.py new file mode 100644 index 00000000..cf9a574e --- /dev/null +++ b/tests/integration/test_reclaim_guard_symmetry_postgres.py @@ -0,0 +1,274 @@ +"""Flag-ON ``run_once`` integration proof for reclaim/assign guard symmetry. + +A live :class:`MasterOrchestrationDriver.run_once` pass runs both the legacy +reassignment pass (detect -> reclaim -> assign) and the worker-plane engine. +Worker-owned prism PRIMARY units are ``ASSIGNED`` with a NULL +``assigned_validator_hotkey`` by design (owned by the ``worker_assignments`` +replica plane). Before this fix the legacy ``reclaim_stale_assignments`` read a +null hotkey as "offline validator => reassignable" and churned those primaries +back to PENDING every pass. This test proves that under +``BASE_COMPUTE__WORKER_PLANE_ENABLED`` semantics (validator +``AssignmentService`` configured with ``worker_plane_capabilities={"gpu"}``) a +worker-owned primary stays ASSIGNED across repeated passes with NO PENDING churn, +while a genuinely stale validator-owned unit is still reclaimed (reclaim is not +globally disabled). Runs against the mission test Postgres (127.0.0.1:15433). +""" + +from __future__ import annotations + +from collections.abc import Awaitable, Callable +from dataclasses import dataclass, field +from datetime import UTC, datetime +from typing import Any + +import pytest +from sqlalchemy import select + +from base.bittensor.metagraph_cache import MetagraphCache +from base.db import ( + Validator, + ValidatorStatus, + WorkAssignmentStatus, + WorkerAssignment, + WorkerRegistration, + WorkerStatus, + create_engine, + create_session_factory, + session_scope, +) +from base.db.models import WorkAssignment +from base.master.assignment import CAPABILITY_GPU, AssignmentService +from base.master.orchestration import ( + ChallengePendingWork, + MasterOrchestrationDriver, +) +from base.master.validator_coordination import ValidatorCoordinationService +from base.master.worker_assignment import WorkerAssignmentService +from base.master.worker_assignment_engine import WorkerAssignmentEngine +from base.master.worker_coordination import WorkerCoordinationService +from base.master.worker_reconciliation import WorkerReconciliationService +from base.security.worker_auth import ( + MetagraphMinerMembership, + SqlAlchemyWorkerNonceStore, +) + +pytestmark = pytest.mark.postgres + +NOW = datetime(2026, 6, 29, 12, 0, 0, tzinfo=UTC) +TTL = 120 + + +@dataclass +class _FakeWorkSource: + works: list[ChallengePendingWork] = field(default_factory=list) + + async def fetch_pending_work(self) -> list[ChallengePendingWork]: + return list(self.works) + + +@dataclass +class _FakeFoldTrigger: + calls: list[tuple[str, str, str, str]] = field(default_factory=list) + + async def fold( + self, *, challenge_slug: str, job_id: str, task_id: str, reason: str + ) -> None: + self.calls.append((challenge_slug, job_id, task_id, reason)) + + +class _FakeForwarder: + def __init__(self) -> None: + self.calls: list[str] = [] + + async def forward_result( + self, + *, + challenge_slug: str, + work_unit_id: str, + submission_ref: str, + result_payload: Any, + ) -> None: + self.calls.append(work_unit_id) + + +async def _add_worker(factory: Any, *, worker_pubkey: str, miner_hotkey: str) -> None: + async with session_scope(factory) as session: + session.add( + WorkerRegistration( + worker_id=f"wid-{worker_pubkey}", + worker_pubkey=worker_pubkey, + miner_hotkey=miner_hotkey, + binding_signature="sig", + binding_nonce=f"nonce-{worker_pubkey}", + provider="local", + provider_instance_ref="local-1", + capabilities=["gpu"], + status=WorkerStatus.ACTIVE, + last_heartbeat_at=NOW, + created_at=NOW, + updated_at=NOW, + ) + ) + + +async def _add_validator(factory: Any, hotkey: str, capabilities: list[str]) -> None: + async with session_scope(factory) as session: + session.add( + Validator( + hotkey=hotkey, + uid=None, + status=ValidatorStatus.ONLINE, + capabilities=list(capabilities), + version="1.0.0", + registered_at=NOW, + last_heartbeat_at=NOW, + ) + ) + + +async def _set_validator_status( + factory: Any, hotkey: str, status: ValidatorStatus +) -> None: + async with session_scope(factory) as session: + validator = ( + await session.execute(select(Validator).where(Validator.hotkey == hotkey)) + ).scalar_one() + validator.status = status + + +async def _rows(factory: Any) -> dict[str, WorkAssignment]: + async with factory() as session: + result = await session.execute(select(WorkAssignment)) + return {r.work_unit_id: r for r in result.scalars().all()} + + +async def _active_replicas(factory: Any, work_unit_id: str) -> list[WorkerAssignment]: + async with factory() as session: + rows = ( + ( + await session.execute( + select(WorkerAssignment).where( + WorkerAssignment.work_unit_id == work_unit_id, + WorkerAssignment.status.in_( + ( + WorkAssignmentStatus.ASSIGNED, + WorkAssignmentStatus.RUNNING, + ) + ), + ) + ) + ) + .scalars() + .all() + ) + return list(rows) + + +def _build_flag_on_driver(factory: Any) -> tuple[MasterOrchestrationDriver, Any, Any]: + cache = MetagraphCache(netuid=1, ttl_seconds=300) + worker_service = WorkerCoordinationService( + factory, + miner_membership=MetagraphMinerMembership(cache), + binding_nonce_store=SqlAlchemyWorkerNonceStore(factory), + heartbeat_ttl_seconds=TTL, + now_fn=lambda: NOW, + ) + worker_assignment_service = WorkerAssignmentService( + factory, worker_service=worker_service, now_fn=lambda: NOW + ) + worker_engine = WorkerAssignmentEngine( + factory, + assignment_service=worker_assignment_service, + worker_service=worker_service, + replication_factor=2, + now_fn=lambda: NOW, + ) + forwarder = _FakeForwarder() + reconciler = WorkerReconciliationService( + factory, result_forwarder=forwarder, now_fn=lambda: NOW + ) + # Flag-ON semantics: gpu units are owned by the worker plane, so the validator + # AssignmentService is configured to skip them in BOTH assign and reclaim. + assignment_service = AssignmentService( + factory, + now_fn=lambda: NOW, + worker_plane_capabilities=frozenset({CAPABILITY_GPU}), + ) + validator_service = ValidatorCoordinationService(factory, now_fn=lambda: NOW) + fold = _FakeFoldTrigger() + driver = MasterOrchestrationDriver( + assignment_service=assignment_service, + validator_service=validator_service, + work_source=_FakeWorkSource( + works=[ + ChallengePendingWork( + challenge_slug="prism", + submission_id="psub", + submission_ref="miner-H", + ), + ChallengePendingWork( + challenge_slug="agent-challenge", + submission_id="sub", + submission_ref="miner-C", + task_ids=("a",), + job_id="job", + ), + ] + ), + fold_trigger=fold, + worker_assignment_engine=worker_engine, + worker_reconciler=reconciler, + seed=1, + ) + return driver, forwarder, fold + + +async def test_flag_on_run_once_keeps_worker_primary_assigned_no_churn( + migrated_postgres_database: str, + cleanup_postgres_database: Callable[[], Awaitable[None]], +) -> None: + engine_db = create_engine(migrated_postgres_database) + factory = create_session_factory(engine_db) + driver, forwarder, _fold = _build_flag_on_driver(factory) + try: + # Two workers of distinct owners, neither owned by the submitter miner-H. + await _add_worker(factory, worker_pubkey="wp-a", miner_hotkey="miner-A") + await _add_worker(factory, worker_pubkey="wp-b", miner_hotkey="miner-B") + # A cpu validator holds the agent-challenge (validator-owned) unit. + await _add_validator(factory, "v1", ["cpu"]) + + # Pass 1: bridge + replicate the gpu primary across two workers and assign + # the cpu unit to the validator. + await driver.run_once() + rows = await _rows(factory) + assert rows["psub"].status == WorkAssignmentStatus.ASSIGNED + assert rows["psub"].assigned_validator_hotkey is None # worker-owned + assert len(await _active_replicas(factory, "psub")) == 2 + assert rows["sub:a"].status == WorkAssignmentStatus.ASSIGNED + assert rows["sub:a"].assigned_validator_hotkey == "v1" + + # The cpu validator crashes; its (validator-owned) unit must still be + # reclaimed by the legacy reclaim path on the next passes. + await _set_validator_status(factory, "v1", ValidatorStatus.OFFLINE) + + # Repeated passes must NOT churn the worker-owned primary to PENDING. + for _ in range(3): + await driver.run_once() + rows = await _rows(factory) + assert rows["psub"].status == WorkAssignmentStatus.ASSIGNED + assert rows["psub"].assigned_validator_hotkey is None + assert rows["psub"].attempt_count == 1 # never re-incremented + assert len(await _active_replicas(factory, "psub")) == 2 + + # Reconciliation never fired (no worker posted a result), so the primary + # was held ASSIGNED purely by the reclaim guard, not by acceptance. + assert forwarder.calls == [] + + # Reclaim is NOT globally disabled: the stale validator-owned cpu unit was + # reverted off the offline validator (reclaimed to pending, no re-assign + # target online). + rows = await _rows(factory) + assert rows["sub:a"].status == WorkAssignmentStatus.PENDING + assert rows["sub:a"].assigned_validator_hotkey is None + finally: + await engine_db.dispose() diff --git a/tests/integration/test_worker_assignment_migration_postgres.py b/tests/integration/test_worker_assignment_migration_postgres.py new file mode 100644 index 00000000..baf833d6 --- /dev/null +++ b/tests/integration/test_worker_assignment_migration_postgres.py @@ -0,0 +1,156 @@ +"""Alembic migration coverage for ``worker_assignments`` (revision 0010). + +Postgres counterpart to the sqlite drift test: ``alembic upgrade head`` creates +``worker_assignments`` with its replica columns and the ``(work_unit_id, +worker_id)`` uniqueness that keeps a restart from spawning orphan replicas +(VAL-AGENT-017); a downgrade to the prior worker-registry revision drops it while +leaving ``worker_registrations`` and the legacy control-plane tables intact. +""" + +from __future__ import annotations + +import asyncio +from pathlib import Path + +import pytest +from sqlalchemy import text + +from base.db import create_engine +from base.db.migrations import downgrade, upgrade + +pytestmark = pytest.mark.postgres + +ROOT = Path(__file__).resolve().parents[2] + +WORKER_ASSIGNMENT_COLUMNS = { + "id", + "challenge_slug", + "work_unit_id", + "submission_ref", + "worker_id", + "worker_pubkey", + "miner_hotkey", + "payload", + "required_capability", + "status", + "attempt_count", + "max_attempts", + "deadline_at", + "last_progress_at", + "checkpoint_ref", + "result_success", + "result_payload", + "manifest_sha256", + "created_at", + "updated_at", +} +PRESERVED_TABLES = ("worker_registrations", "work_assignments", "validators") + + +async def _columns(database_url: str, table: str) -> set[str]: + engine = create_engine(database_url) + try: + async with engine.connect() as connection: + rows = ( + await connection.execute( + text( + "SELECT column_name FROM information_schema.columns " + "WHERE table_schema = current_schema() AND table_name = :t" + ), + {"t": table}, + ) + ).scalars() + return set(rows.all()) + finally: + await engine.dispose() + + +async def _table_exists(database_url: str, table: str) -> bool: + engine = create_engine(database_url) + try: + async with engine.connect() as connection: + found = ( + await connection.execute( + text( + "SELECT 1 FROM information_schema.tables " + "WHERE table_schema = current_schema() AND table_name = :t" + ), + {"t": table}, + ) + ).first() + return found is not None + finally: + await engine.dispose() + + +async def _unique_column_sets(database_url: str, table: str) -> set[tuple[str, ...]]: + engine = create_engine(database_url) + try: + async with engine.connect() as connection: + rows = ( + await connection.execute( + text( + """ + SELECT tc.constraint_name, kcu.column_name + FROM information_schema.table_constraints tc + JOIN information_schema.key_column_usage kcu + ON tc.constraint_name = kcu.constraint_name + AND tc.table_schema = kcu.table_schema + WHERE tc.table_schema = current_schema() + AND tc.table_name = :t + AND tc.constraint_type = 'UNIQUE' + ORDER BY kcu.ordinal_position + """ + ), + {"t": table}, + ) + ).all() + finally: + await engine.dispose() + grouped: dict[str, list[str]] = {} + for constraint_name, column_name in rows: + grouped.setdefault(constraint_name, []).append(column_name) + return {tuple(sorted(cols)) for cols in grouped.values()} + + +async def test_worker_assignments_present_after_upgrade( + migrated_postgres_database: str, +) -> None: + columns = await _columns(migrated_postgres_database, "worker_assignments") + assert WORKER_ASSIGNMENT_COLUMNS <= columns + unique_sets = await _unique_column_sets( + migrated_postgres_database, "worker_assignments" + ) + assert ("work_unit_id", "worker_id") in unique_sets + + +async def test_worker_assignments_downgrade_upgrade_roundtrip( + migrated_postgres_database: str, +) -> None: + config_path = ROOT / "alembic.ini" + preserved_before = { + table: await _columns(migrated_postgres_database, table) + for table in PRESERVED_TABLES + } + + await asyncio.to_thread( + downgrade, + config_path, + database_url=migrated_postgres_database, + revision="0009_create_worker_registry", + ) + try: + assert not await _table_exists(migrated_postgres_database, "worker_assignments") + for table, columns in preserved_before.items(): + assert await _columns(migrated_postgres_database, table) == columns + finally: + await asyncio.to_thread( + upgrade, + config_path, + database_url=migrated_postgres_database, + revision="head", + ) + + assert await _table_exists(migrated_postgres_database, "worker_assignments") + for table, columns in preserved_before.items(): + assert await _columns(migrated_postgres_database, table) == columns diff --git a/tests/integration/test_worker_coordination_postgres.py b/tests/integration/test_worker_coordination_postgres.py new file mode 100644 index 00000000..a403dd85 --- /dev/null +++ b/tests/integration/test_worker_coordination_postgres.py @@ -0,0 +1,437 @@ +from __future__ import annotations + +import hashlib +from collections.abc import Awaitable, Callable +from datetime import UTC, datetime, timedelta +from typing import Any + +import pytest +from httpx import ASGITransport, AsyncClient +from sqlalchemy import func, select + +from base.bittensor.metagraph_cache import MetagraphCache +from base.db import ( + WorkerRegistration, + WorkerStatus, + create_engine, + create_session_factory, + session_scope, +) +from base.master.app_proxy import create_proxy_app +from base.master.worker_coordination import ( + WorkerCoordinationService, + WorkerRebindError, +) +from base.security.validator_auth import canonical_validator_request +from base.security.worker_auth import ( + CoordinationReadEligibility, + MetagraphMinerMembership, + SqlAlchemyWorkerNonceStore, + WorkerReplayError, + WorkerSignedRequestVerifier, + worker_binding_message, +) + +pytestmark = pytest.mark.postgres + +BASE_TS = datetime(2025, 6, 15, 12, 0, 0, tzinfo=UTC) +BASE_EPOCH = BASE_TS.timestamp() +TTL = 120 +MINER_H1 = "pg-miner-H1" +MINER_H2 = "pg-miner-H2" +VALIDATOR = "pg-val-permit" +INTERNAL_BRIDGE_TOKEN = "pg-prism-bridge-token" + + +class _Clock: + def __init__(self, value: datetime) -> None: + self.value = value + + def now(self) -> datetime: + return self.value + + +def _sign(hotkey: str, message: bytes) -> str: + return hashlib.sha256(hotkey.encode() + b":" + message).hexdigest() + + +def _fake_verifier(hotkey: str, message: bytes, signature: str) -> bool: + return signature == _sign(hotkey, message) + + +def _service(session_factory: object, clock: _Clock) -> WorkerCoordinationService: + cache = MetagraphCache(netuid=1, ttl_seconds=300) + cache.update_from_metagraph( + [MINER_H1, MINER_H2], validator_permits=[False, False], stakes=[0.0, 0.0] + ) + return WorkerCoordinationService( + session_factory, # type: ignore[arg-type] + miner_membership=MetagraphMinerMembership(cache), + binding_nonce_store=SqlAlchemyWorkerNonceStore(session_factory), # type: ignore[arg-type] + signature_verifier=_fake_verifier, + heartbeat_ttl_seconds=TTL, + now_fn=clock.now, + ) + + +async def _register( + service: WorkerCoordinationService, + *, + worker_pubkey: str, + miner_hotkey: str, + nonce: str, + provider: str = "local", +) -> WorkerRegistration: + message = worker_binding_message( + worker_pubkey=worker_pubkey, miner_hotkey=miner_hotkey, nonce=nonce + ) + return await service.register( + worker_pubkey=worker_pubkey, + miner_hotkey=miner_hotkey, + binding_signature=_sign(miner_hotkey, message), + nonce=nonce, + provider=provider, + provider_instance_ref="ref-1", + capabilities=["gpu"], + ) + + +# VAL-MASTER-001 +async def test_register_pending_then_heartbeat_active_on_postgres( + migrated_postgres_database: str, + cleanup_postgres_database: Callable[[], Awaitable[None]], +) -> None: + engine = create_engine(migrated_postgres_database) + session_factory = create_session_factory(engine) + clock = _Clock(BASE_TS) + service = _service(session_factory, clock) + try: + worker = await _register( + service, worker_pubkey="pg-wp-1", miner_hotkey=MINER_H1, nonce="n-1" + ) + assert WorkerStatus(worker.status) == WorkerStatus.PENDING + + rows = await service.list_workers() + assert [service.effective_status(r, clock.now()).value for r in rows] == [ + "pending" + ] + + clock.value = BASE_TS + timedelta(seconds=10) + updated, now = await service.heartbeat( + worker_id=worker.worker_id, worker_pubkey="pg-wp-1" + ) + assert now == clock.value + assert service.effective_status(updated, now) == WorkerStatus.ACTIVE + + active = await service.active_workers(MINER_H1) + assert [w.worker_pubkey for w in active] == ["pg-wp-1"] + finally: + await engine.dispose() + + +# VAL-MASTER-002 (c): replayed binding nonce rejected on postgres +async def test_replayed_binding_nonce_rejected_on_postgres( + migrated_postgres_database: str, + cleanup_postgres_database: Callable[[], Awaitable[None]], +) -> None: + engine = create_engine(migrated_postgres_database) + session_factory = create_session_factory(engine) + service = _service(session_factory, _Clock(BASE_TS)) + try: + await _register( + service, worker_pubkey="pg-wp-r", miner_hotkey=MINER_H1, nonce="dup" + ) + with pytest.raises(WorkerReplayError): + await _register( + service, worker_pubkey="pg-wp-r2", miner_hotkey=MINER_H1, nonce="dup" + ) + async with session_factory() as session: + count = await session.scalar( + select(func.count(WorkerRegistration.id)).where( + WorkerRegistration.worker_pubkey == "pg-wp-r2" + ) + ) + assert count == 0 + finally: + await engine.dispose() + + +# VAL-MASTER-016 +async def test_active_filtering_on_postgres( + migrated_postgres_database: str, + cleanup_postgres_database: Callable[[], Awaitable[None]], +) -> None: + engine = create_engine(migrated_postgres_database) + session_factory = create_session_factory(engine) + clock = _Clock(BASE_TS) + service = _service(session_factory, clock) + try: + active = await _register( + service, worker_pubkey="pg-h1-active", miner_hotkey=MINER_H1, nonce="a" + ) + stale = await _register( + service, worker_pubkey="pg-h1-stale", miner_hotkey=MINER_H1, nonce="b" + ) + other = await _register( + service, worker_pubkey="pg-h2-active", miner_hotkey=MINER_H2, nonce="c" + ) + + await service.heartbeat(worker_id=stale.worker_id, worker_pubkey="pg-h1-stale") + clock.value = BASE_TS + timedelta(seconds=TTL + 10) + await service.heartbeat( + worker_id=active.worker_id, worker_pubkey="pg-h1-active" + ) + await service.heartbeat(worker_id=other.worker_id, worker_pubkey="pg-h2-active") + + h1 = await service.active_workers(MINER_H1) + assert [w.worker_pubkey for w in h1] == ["pg-h1-active"] + h2 = await service.active_workers(MINER_H2) + assert [w.worker_pubkey for w in h2] == ["pg-h2-active"] + assert await service.active_workers("pg-unknown") == [] + finally: + await engine.dispose() + + +# VAL-MASTER-021 +async def test_cross_owner_rebind_rejected_on_postgres( + migrated_postgres_database: str, + cleanup_postgres_database: Callable[[], Awaitable[None]], +) -> None: + engine = create_engine(migrated_postgres_database) + session_factory = create_session_factory(engine) + clock = _Clock(BASE_TS) + service = _service(session_factory, clock) + try: + worker = await _register( + service, worker_pubkey="pg-wp-rb", miner_hotkey=MINER_H1, nonce="rb-1" + ) + await service.heartbeat(worker_id=worker.worker_id, worker_pubkey="pg-wp-rb") + + with pytest.raises(WorkerRebindError): + await _register( + service, worker_pubkey="pg-wp-rb", miner_hotkey=MINER_H2, nonce="rb-2" + ) + + async with session_factory() as session: + rows = ( + ( + await session.execute( + select(WorkerRegistration).where( + WorkerRegistration.worker_pubkey == "pg-wp-rb" + ) + ) + ) + .scalars() + .all() + ) + assert len(rows) == 1 + assert rows[0].miner_hotkey == MINER_H1 + finally: + await engine.dispose() + + +# VAL-MASTER-018 +async def test_retired_worker_terminal_on_postgres( + migrated_postgres_database: str, + cleanup_postgres_database: Callable[[], Awaitable[None]], +) -> None: + engine = create_engine(migrated_postgres_database) + session_factory = create_session_factory(engine) + clock = _Clock(BASE_TS) + service = _service(session_factory, clock) + try: + worker = await _register( + service, worker_pubkey="pg-wp-ret", miner_hotkey=MINER_H1, nonce="ret-1" + ) + await service.heartbeat(worker_id=worker.worker_id, worker_pubkey="pg-wp-ret") + + async with session_scope(session_factory) as session: + row = ( + await session.execute( + select(WorkerRegistration).where( + WorkerRegistration.worker_pubkey == "pg-wp-ret" + ) + ) + ).scalar_one() + row.status = WorkerStatus.RETIRED + + # sole worker of its hotkey, yet never listed active. + assert await service.active_workers(MINER_H1) == [] + + # heartbeat does not resurrect it. + clock.value = BASE_TS + timedelta(seconds=5) + updated, now = await service.heartbeat( + worker_id=worker.worker_id, worker_pubkey="pg-wp-ret" + ) + assert WorkerStatus(updated.status) == WorkerStatus.RETIRED + assert service.effective_status(updated, now) == WorkerStatus.RETIRED + assert await service.active_workers(MINER_H1) == [] + finally: + await engine.dispose() + + +# --------------------------------------------------------------------------- +# Admission fleet-read dual-auth over HTTP against real postgres (15433). +# master-admission-fleetread-auth: internal bridge bearer OR signed request on +# GET /v1/workers/active; the full fleet GET /v1/workers stays signed-only. +# --------------------------------------------------------------------------- + + +class _TokenRegistry: + def __init__(self, token: str) -> None: + self._token = token + + def get_token(self, slug: str) -> str: + if slug == "prism": + return self._token + raise RuntimeError(f"no token for {slug!r}") + + +class _HttpClock: + def __init__(self, epoch: float) -> None: + self.epoch = float(epoch) + + def time(self) -> float: + return self.epoch + + def now(self) -> datetime: + return datetime.fromtimestamp(self.epoch, UTC) + + +class FakeNonceStore: + async def reserve(self, **_kwargs: Any) -> None: + return None + + +class FakeCache: + def get(self) -> dict[str, int]: + return {} + + +def _signed_headers( + *, + method: str, + path: str, + query_string: str, + hotkey: str, + nonce: str, + epoch: float, +) -> dict[str, str]: + ts = str(int(epoch)) + canonical = canonical_validator_request( + method=method, + path=path, + query_string=query_string, + timestamp=ts, + nonce=nonce, + body=b"", + ) + return { + "X-Hotkey": hotkey, + "X-Signature": _sign(hotkey, canonical.encode()), + "X-Nonce": nonce, + "X-Timestamp": ts, + } + + +async def test_active_admission_dual_auth_over_http_on_postgres( + migrated_postgres_database: str, + cleanup_postgres_database: Callable[[], Awaitable[None]], +) -> None: + engine = create_engine(migrated_postgres_database) + session_factory = create_session_factory(engine) + clock = _HttpClock(BASE_EPOCH) + cache = MetagraphCache(netuid=1, ttl_seconds=300) + cache.update_from_metagraph( + [MINER_H1, MINER_H2, VALIDATOR], + validator_permits=[False, False, True], + stakes=[0.0, 0.0, 100.0], + ) + service = WorkerCoordinationService( + session_factory, + miner_membership=MetagraphMinerMembership(cache), + binding_nonce_store=SqlAlchemyWorkerNonceStore(session_factory), + signature_verifier=_fake_verifier, + heartbeat_ttl_seconds=TTL, + now_fn=clock.now, + ) + verifier = WorkerSignedRequestVerifier( + nonce_store=SqlAlchemyWorkerNonceStore(session_factory), + eligibility=CoordinationReadEligibility(session_factory, cache), + signature_verifier=_fake_verifier, + ttl_seconds=300, + now_fn=clock.time, + ) + app = create_proxy_app( + registry=_TokenRegistry(INTERNAL_BRIDGE_TOKEN), + nonce_store=FakeNonceStore(), + metagraph_cache=FakeCache(), # type: ignore[arg-type] + worker_service=service, + worker_verifier=verifier, + ) + transport = ASGITransport(app=app) + client = AsyncClient(transport=transport, base_url="http://testserver") + try: + # Two active workers on distinct owners (proves per-hotkey filtering). + w1 = await _register( + service, worker_pubkey="pg-http-h1", miner_hotkey=MINER_H1, nonce="h1" + ) + w2 = await _register( + service, worker_pubkey="pg-http-h2", miner_hotkey=MINER_H2, nonce="h2" + ) + await service.heartbeat(worker_id=w1.worker_id, worker_pubkey="pg-http-h1") + await service.heartbeat(worker_id=w2.worker_id, worker_pubkey="pg-http-h2") + + bearer = {"Authorization": f"Bearer {INTERNAL_BRIDGE_TOKEN}"} + + # (a) internal bearer returns exactly that hotkey's ACTIVE workers. + resp = await client.get( + "/v1/workers/active", params={"hotkey": MINER_H1}, headers=bearer + ) + assert resp.status_code == 200 + assert [w["worker_pubkey"] for w in resp.json()["workers"]] == ["pg-http-h1"] + + # (b) wrong bearer / missing bearer with no signature => 401/403. + wrong = await client.get( + "/v1/workers/active", + params={"hotkey": MINER_H1}, + headers={"Authorization": "Bearer nope"}, + ) + assert wrong.status_code in (401, 403) + missing = await client.get("/v1/workers/active", params={"hotkey": MINER_H1}) + assert missing.status_code in (401, 403) + + # (c) signed-request path still works unchanged. + signed = await client.get( + "/v1/workers/active", + params={"hotkey": MINER_H1}, + headers=_signed_headers( + method="GET", + path="/v1/workers/active", + query_string=f"hotkey={MINER_H1}", + hotkey=VALIDATOR, + nonce="sig-active-1", + epoch=clock.time(), + ), + ) + assert signed.status_code == 200 + assert [w["worker_pubkey"] for w in signed.json()["workers"]] == ["pg-http-h1"] + + # (d) full fleet view rejects the internal bearer (signed-only). + fleet_bearer = await client.get("/v1/workers", headers=bearer) + assert fleet_bearer.status_code in (401, 403) + fleet_signed = await client.get( + "/v1/workers", + headers=_signed_headers( + method="GET", + path="/v1/workers", + query_string="", + hotkey=VALIDATOR, + nonce="sig-fleet-1", + epoch=clock.time(), + ), + ) + assert fleet_signed.status_code == 200 + finally: + await client.aclose() + await engine.dispose() diff --git a/tests/integration/test_worker_migration_postgres.py b/tests/integration/test_worker_migration_postgres.py new file mode 100644 index 00000000..a1c85e8b --- /dev/null +++ b/tests/integration/test_worker_migration_postgres.py @@ -0,0 +1,128 @@ +from __future__ import annotations + +import asyncio +from pathlib import Path + +import pytest +from sqlalchemy import text + +from base.db import create_engine +from base.db.migrations import downgrade, upgrade + +pytestmark = pytest.mark.postgres + +ROOT = Path(__file__).resolve().parents[2] + +WORKER_REGISTRATION_COLUMNS = { + "worker_id", + "worker_pubkey", + "miner_hotkey", + "binding_signature", + "provider", + "provider_instance_ref", + "capabilities", + "status", + "last_heartbeat_at", + "created_at", +} +# Legacy tables the worker migration MUST NOT alter. +LEGACY_TABLES = ("validators", "work_assignments", "challenges") + + +async def _columns(database_url: str, table: str) -> set[str]: + engine = create_engine(database_url) + try: + async with engine.connect() as connection: + rows = ( + await connection.execute( + text( + "SELECT column_name FROM information_schema.columns " + "WHERE table_schema = current_schema() AND table_name = :t" + ), + {"t": table}, + ) + ).scalars() + return set(rows.all()) + finally: + await engine.dispose() + + +async def _table_exists(database_url: str, table: str) -> bool: + engine = create_engine(database_url) + try: + async with engine.connect() as connection: + found = ( + await connection.execute( + text( + "SELECT 1 FROM information_schema.tables " + "WHERE table_schema = current_schema() AND table_name = :t" + ), + {"t": table}, + ) + ).first() + return found is not None + finally: + await engine.dispose() + + +# VAL-MASTER-015 +async def test_worker_tables_present_after_upgrade( + migrated_postgres_database: str, +) -> None: + registration_columns = await _columns( + migrated_postgres_database, "worker_registrations" + ) + assert WORKER_REGISTRATION_COLUMNS <= registration_columns + assert await _table_exists(migrated_postgres_database, "worker_faults") + + legacy_before = { + table: await _columns(migrated_postgres_database, table) + for table in LEGACY_TABLES + } + for table, columns in legacy_before.items(): + assert columns, f"legacy table {table} unexpectedly empty/missing" + + +# VAL-MASTER-015: downgrade removes the worker tables, upgrade restores them. +async def test_worker_migration_downgrade_upgrade_roundtrip( + migrated_postgres_database: str, +) -> None: + config_path = ROOT / "alembic.ini" + + legacy_before = { + table: await _columns(migrated_postgres_database, table) + for table in LEGACY_TABLES + } + + # alembic env runs asyncio.run internally, so drive it from a worker thread + # to avoid nesting event loops inside this async test. + await asyncio.to_thread( + downgrade, + config_path, + database_url=migrated_postgres_database, + revision="0008_validator_subscriptions", + ) + try: + assert not await _table_exists( + migrated_postgres_database, "worker_registrations" + ) + assert not await _table_exists(migrated_postgres_database, "worker_faults") + assert not await _table_exists( + migrated_postgres_database, "worker_request_nonces" + ) + # legacy tables untouched by the downgrade. + for table, columns in legacy_before.items(): + assert await _columns(migrated_postgres_database, table) == columns + finally: + # restore head so the shared session-scoped test database is intact. + await asyncio.to_thread( + upgrade, + config_path, + database_url=migrated_postgres_database, + revision="head", + ) + + assert await _table_exists(migrated_postgres_database, "worker_registrations") + assert await _table_exists(migrated_postgres_database, "worker_faults") + for table, columns in legacy_before.items(): + assert await _columns(migrated_postgres_database, table) == columns diff --git a/tests/integration/test_worker_reconciliation_postgres.py b/tests/integration/test_worker_reconciliation_postgres.py new file mode 100644 index 00000000..1d1ec853 --- /dev/null +++ b/tests/integration/test_worker_reconciliation_postgres.py @@ -0,0 +1,313 @@ +"""Postgres parity for worker-plane reconciliation, disputes, and audit faults. + +Mirrors the SQLite unit behavior of :class:`WorkerReconciliationService` against +the throwaway test Postgres, confirming the ``disputed`` terminal status, the +validator-executor audit unit, and the ``worker_faults`` fault row persist and +reconcile identically across backends (VAL-MASTER-008/009/010). +""" + +from __future__ import annotations + +from collections.abc import Awaitable, Callable +from datetime import UTC, datetime +from typing import Any + +import pytest +from sqlalchemy import select + +from base.bittensor.metagraph_cache import MetagraphCache +from base.db import ( + Validator, + ValidatorStatus, + WorkAssignmentStatus, + WorkerAssignment, + WorkerFault, + WorkerRegistration, + WorkerStatus, + create_engine, + create_session_factory, + session_scope, +) +from base.db.models import WorkAssignment +from base.master.assignment import CAPABILITY_GPU, AssignmentService +from base.master.assignment_coordination import AssignmentCoordinationService +from base.master.worker_assignment import WorkerAssignmentService +from base.master.worker_assignment_engine import WorkerAssignmentEngine +from base.master.worker_coordination import WorkerCoordinationService +from base.master.worker_reconciliation import ( + AUDIT_OF_PAYLOAD_KEY, + WorkerReconciliationService, + audit_work_unit_id, +) +from base.security.worker_auth import ( + MetagraphMinerMembership, + SqlAlchemyWorkerNonceStore, +) +from base.worker.proof import MANIFEST_SHA256_PAYLOAD_KEY, PROOF_PAYLOAD_KEY + +pytestmark = pytest.mark.postgres + +NOW = datetime(2026, 6, 28, 12, 0, 0, tzinfo=UTC) +TTL = 120 +HASH_A = "a" * 64 +HASH_B = "b" * 64 +GPU_VALIDATOR = "gpu-val" + + +class _Forwarder: + def __init__(self) -> None: + self.calls: list[str] = [] + + async def forward_result( + self, + *, + challenge_slug: str, + work_unit_id: str, + submission_ref: str, + result_payload: Any, + ) -> None: + self.calls.append(work_unit_id) + + +def _proof_payload(manifest: str) -> dict[str, Any]: + return { + PROOF_PAYLOAD_KEY: { + "version": 1, + "tier": 0, + MANIFEST_SHA256_PAYLOAD_KEY: manifest, + }, + MANIFEST_SHA256_PAYLOAD_KEY: manifest, + } + + +async def _add_worker(factory: Any, *, worker_pubkey: str, miner_hotkey: str) -> None: + async with session_scope(factory) as session: + session.add( + WorkerRegistration( + worker_id=f"wid-{worker_pubkey}", + worker_pubkey=worker_pubkey, + miner_hotkey=miner_hotkey, + binding_signature="sig", + binding_nonce=f"nonce-{worker_pubkey}", + provider="local", + provider_instance_ref="local-1", + capabilities=["gpu"], + status=WorkerStatus.ACTIVE, + last_heartbeat_at=NOW, + created_at=NOW, + updated_at=NOW, + ) + ) + + +async def _add_gpu_unit(factory: Any, *, work_unit_id: str, submitter: str) -> None: + async with session_scope(factory) as session: + session.add( + WorkAssignment( + challenge_slug="prism", + work_unit_id=work_unit_id, + submission_ref=submitter, + payload={"run_spec": {"image": "img"}}, + required_capability="gpu", + status=WorkAssignmentStatus.PENDING, + attempt_count=0, + max_attempts=3, + created_at=NOW, + updated_at=NOW, + ) + ) + + +async def _replicas(factory: Any, work_unit_id: str) -> list[WorkerAssignment]: + async with factory() as session: + rows = ( + ( + await session.execute( + select(WorkerAssignment).where( + WorkerAssignment.work_unit_id == work_unit_id + ) + ) + ) + .scalars() + .all() + ) + return list(rows) + + +async def _unit(factory: Any, work_unit_id: str) -> WorkAssignment: + async with factory() as session: + return ( + await session.execute( + select(WorkAssignment).where( + WorkAssignment.work_unit_id == work_unit_id + ) + ) + ).scalar_one() + + +def _build(factory: Any) -> tuple[Any, ...]: + cache = MetagraphCache(netuid=1, ttl_seconds=300) + worker_service = WorkerCoordinationService( + factory, + miner_membership=MetagraphMinerMembership(cache), + binding_nonce_store=SqlAlchemyWorkerNonceStore(factory), + heartbeat_ttl_seconds=TTL, + now_fn=lambda: NOW, + ) + assignment_service = WorkerAssignmentService( + factory, worker_service=worker_service, now_fn=lambda: NOW + ) + engine = WorkerAssignmentEngine( + factory, + assignment_service=assignment_service, + worker_service=worker_service, + replication_factor=2, + now_fn=lambda: NOW, + ) + forwarder = _Forwarder() + reconciler = WorkerReconciliationService( + factory, result_forwarder=forwarder, now_fn=lambda: NOW + ) + validator_plane = AssignmentService( + factory, + now_fn=lambda: NOW, + worker_plane_capabilities=frozenset({CAPABILITY_GPU}), + ) + coordination = AssignmentCoordinationService(factory, now_fn=lambda: NOW) + return ( + worker_service, + assignment_service, + engine, + forwarder, + reconciler, + validator_plane, + coordination, + ) + + +async def _post( + service: WorkerAssignmentService, factory: Any, unit: str, owner: str, h: str +) -> None: + replica = next(r for r in await _replicas(factory, unit) if r.miner_hotkey == owner) + await service.post_result( + assignment_id=str(replica.id), + worker_pubkey=replica.worker_pubkey, + success=True, + payload=_proof_payload(h), + ) + + +async def test_dispute_and_audit_faults_on_postgres( + migrated_postgres_database: str, + cleanup_postgres_database: Callable[[], Awaitable[None]], +) -> None: + engine_db = create_engine(migrated_postgres_database) + factory = create_session_factory(engine_db) + ( + worker_service, + assignment_service, + engine, + forwarder, + reconciler, + validator_plane, + coordination, + ) = _build(factory) + try: + await _add_worker(factory, worker_pubkey="wp-a", miner_hotkey="miner-A") + await _add_worker(factory, worker_pubkey="wp-b", miner_hotkey="miner-B") + await _add_gpu_unit(factory, work_unit_id="U", submitter="miner-H") + await engine.assign_pending(seed=1) + + await _post(assignment_service, factory, "U", "miner-A", HASH_A) + await _post(assignment_service, factory, "U", "miner-B", HASH_B) + + result = await reconciler.reconcile_once() + assert result.disputed == ["U"] + assert forwarder.calls == [] + assert (await _unit(factory, "U")).status == WorkAssignmentStatus.DISPUTED + + audit = await _unit(factory, audit_work_unit_id("U")) + assert audit.required_capability == CAPABILITY_GPU + assert audit.payload[AUDIT_OF_PAYLOAD_KEY] == "U" + + # A validator replays; its authoritative manifest agrees with A. + async with session_scope(factory) as session: + session.add( + Validator( + hotkey=GPU_VALIDATOR, + uid=None, + status=ValidatorStatus.ONLINE, + capabilities=["gpu"], + version="1.0.0", + registered_at=NOW, + last_heartbeat_at=NOW, + ) + ) + await validator_plane.assign_pending(seed=1) + audit = await _unit(factory, audit_work_unit_id("U")) + assert audit.assigned_validator_hotkey == GPU_VALIDATOR + await coordination.post_result( + assignment_id=str(audit.id), + hotkey=GPU_VALIDATOR, + success=True, + payload=_proof_payload(HASH_A), + ) + + result = await reconciler.reconcile_once() + assert {p.split(":")[-1] for p in result.faults} == {"wid-wp-b"} + + async with factory() as session: + faults = (await session.execute(select(WorkerFault))).scalars().all() + assert len(faults) == 1 + assert faults[0].worker_id == "wid-wp-b" + assert faults[0].work_unit_id == "U" + + # Fault recording leaves the worker's status unchanged. + async with factory() as session: + worker_b = ( + await session.execute( + select(WorkerRegistration).where( + WorkerRegistration.worker_id == "wid-wp-b" + ) + ) + ).scalar_one() + assert WorkerStatus(worker_b.status) == WorkerStatus.ACTIVE + # The disputed unit is never forwarded, even after audit resolution. + assert forwarder.calls == [] + assert (await _unit(factory, "U")).status == WorkAssignmentStatus.DISPUTED + finally: + await engine_db.dispose() + + +async def test_matching_hashes_accept_on_postgres( + migrated_postgres_database: str, + cleanup_postgres_database: Callable[[], Awaitable[None]], +) -> None: + engine_db = create_engine(migrated_postgres_database) + factory = create_session_factory(engine_db) + ( + _worker_service, + assignment_service, + engine, + forwarder, + reconciler, + _validator_plane, + _coordination, + ) = _build(factory) + try: + await _add_worker(factory, worker_pubkey="wp-a", miner_hotkey="miner-A") + await _add_worker(factory, worker_pubkey="wp-b", miner_hotkey="miner-B") + await _add_gpu_unit(factory, work_unit_id="U", submitter="miner-H") + await engine.assign_pending(seed=1) + + await _post(assignment_service, factory, "U", "miner-A", HASH_A) + await _post(assignment_service, factory, "U", "miner-B", HASH_A) + + result = await reconciler.reconcile_once() + await reconciler.reconcile_once() + + assert result.accepted == ["U"] + assert forwarder.calls == ["U"] + assert (await _unit(factory, "U")).status == WorkAssignmentStatus.COMPLETED + finally: + await engine_db.dispose() diff --git a/tests/integration/test_worker_unit_status_postgres.py b/tests/integration/test_worker_unit_status_postgres.py new file mode 100644 index 00000000..b8f4b2b2 --- /dev/null +++ b/tests/integration/test_worker_unit_status_postgres.py @@ -0,0 +1,437 @@ +"""Worker-plane unit-status read endpoint over real postgres (15433). + +master-unit-status-read-endpoint / VAL-CROSS-011: the full dispute -> audit -> +invalidation -> fault chain must be operator-discoverable via ``GET +/v1/workers/units`` alone. These tests drive the REAL reconciliation machinery +(``WorkerAssignmentEngine`` + ``WorkerReconciliationService`` + the validator +plane) against the mission test postgres, then read the chain back through the +service and the signed HTTP surface. +""" + +from __future__ import annotations + +import hashlib +from collections.abc import Awaitable, Callable +from datetime import UTC, datetime +from typing import Any + +import pytest +from httpx import ASGITransport, AsyncClient +from sqlalchemy import select + +from base.bittensor.metagraph_cache import MetagraphCache +from base.db import ( + Validator, + ValidatorStatus, + WorkAssignment, + WorkAssignmentStatus, + WorkerRegistration, + WorkerStatus, + create_engine, + create_session_factory, + session_scope, +) +from base.master.app_proxy import create_proxy_app +from base.master.assignment import CAPABILITY_GPU, AssignmentService +from base.master.assignment_coordination import AssignmentCoordinationService +from base.master.worker_assignment import WorkerAssignmentService +from base.master.worker_assignment_engine import WorkerAssignmentEngine +from base.master.worker_coordination import WorkerCoordinationService +from base.master.worker_reconciliation import ( + WorkerReconciliationService, + audit_work_unit_id, +) +from base.master.worker_unit_status import ( + AUDIT_OUTCOME_MISMATCH_RESOLVED, + AUDIT_OUTCOME_PENDING, + WorkerUnitStatusService, +) +from base.security.validator_auth import canonical_validator_request +from base.security.worker_auth import ( + CoordinationReadEligibility, + MetagraphMinerMembership, + SqlAlchemyWorkerNonceStore, + WorkerSignedRequestVerifier, +) +from base.worker.proof import MANIFEST_SHA256_PAYLOAD_KEY, PROOF_PAYLOAD_KEY + +pytestmark = pytest.mark.postgres + +NOW = datetime(2026, 7, 7, 12, 0, 0, tzinfo=UTC) +NOW_EPOCH = NOW.timestamp() +TTL = 120 + +MINER_A = "us-miner-A" +MINER_B = "us-miner-B" +MINER_H = "us-miner-H" +VALIDATOR = "us-val-permit" +STRANGER = "us-stranger" +GPU_VALIDATOR = "us-gpu-val" + +HASH_A = "a" * 64 +HASH_B = "b" * 64 +INTERNAL_BRIDGE_TOKEN = "us-prism-bridge-token" + + +class _Clock: + def __init__(self, moment: datetime) -> None: + self.moment = moment + + def now(self) -> datetime: + return self.moment + + +def _sign(hotkey: str, message: bytes) -> str: + return hashlib.sha256(hotkey.encode() + b":" + message).hexdigest() + + +def _fake_verifier(hotkey: str, message: bytes, signature: str) -> bool: + return signature == _sign(hotkey, message) + + +def _proof_payload(manifest: str) -> dict[str, Any]: + return { + PROOF_PAYLOAD_KEY: { + "version": 1, + "tier": 0, + MANIFEST_SHA256_PAYLOAD_KEY: manifest, + }, + MANIFEST_SHA256_PAYLOAD_KEY: manifest, + } + + +class _TokenRegistry: + def __init__(self, token: str) -> None: + self._token = token + + def get_token(self, slug: str) -> str: + if slug == "prism": + return self._token + raise RuntimeError(f"no token for {slug!r}") + + +class FakeNonceStore: + async def reserve(self, **_kwargs: Any) -> None: + return None + + +class FakeCache: + def get(self) -> dict[str, int]: + return {} + + +async def _add_worker(factory: Any, *, worker_pubkey: str, miner_hotkey: str) -> str: + worker_id = f"wid-{worker_pubkey}" + async with session_scope(factory) as session: + session.add( + WorkerRegistration( + worker_id=worker_id, + worker_pubkey=worker_pubkey, + miner_hotkey=miner_hotkey, + binding_signature="sig", + binding_nonce=f"nonce-{worker_pubkey}", + provider="local", + provider_instance_ref="local-1", + capabilities=["gpu"], + status=WorkerStatus.ACTIVE, + last_heartbeat_at=NOW, + created_at=NOW, + updated_at=NOW, + ) + ) + return worker_id + + +async def _add_gpu_unit(factory: Any, *, work_unit_id: str, submitter: str) -> None: + async with session_scope(factory) as session: + session.add( + WorkAssignment( + challenge_slug="prism", + work_unit_id=work_unit_id, + submission_ref=submitter, + payload={"run_spec": {"image": "img"}}, + required_capability="gpu", + status=WorkAssignmentStatus.PENDING, + attempt_count=0, + max_attempts=3, + created_at=NOW, + updated_at=NOW, + ) + ) + + +async def _add_gpu_validator(factory: Any) -> None: + async with session_scope(factory) as session: + session.add( + Validator( + hotkey=GPU_VALIDATOR, + uid=None, + status=ValidatorStatus.ONLINE, + capabilities=["gpu"], + version="1.0.0", + registered_at=NOW, + last_heartbeat_at=NOW, + ) + ) + + +async def _replicas(factory: Any, work_unit_id: str) -> list[Any]: + from base.db import WorkerAssignment + + async with factory() as session: + rows = ( + ( + await session.execute( + select(WorkerAssignment).where( + WorkerAssignment.work_unit_id == work_unit_id + ) + ) + ) + .scalars() + .all() + ) + return list(rows) + + +async def _post_replica( + svc: WorkerAssignmentService, + factory: Any, + *, + work_unit_id: str, + miner_hotkey: str, + manifest: str, +) -> None: + replica = next( + r + for r in await _replicas(factory, work_unit_id) + if r.miner_hotkey == miner_hotkey + ) + await svc.post_result( + assignment_id=str(replica.id), + worker_pubkey=replica.worker_pubkey, + success=True, + payload=_proof_payload(manifest), + ) + + +def _build_services(factory: Any, clock: _Clock) -> dict[str, Any]: + cache = MetagraphCache(netuid=1, ttl_seconds=300) + worker_service = WorkerCoordinationService( + factory, + miner_membership=MetagraphMinerMembership(cache), + binding_nonce_store=SqlAlchemyWorkerNonceStore(factory), + heartbeat_ttl_seconds=TTL, + now_fn=clock.now, + ) + worker_assignment_service = WorkerAssignmentService( + factory, worker_service=worker_service, now_fn=clock.now + ) + engine = WorkerAssignmentEngine( + factory, + assignment_service=worker_assignment_service, + worker_service=worker_service, + replication_factor=2, + now_fn=clock.now, + ) + reconciler = WorkerReconciliationService(factory, now_fn=clock.now) + validator_plane = AssignmentService( + factory, + now_fn=clock.now, + worker_plane_capabilities=frozenset({CAPABILITY_GPU}), + ) + validator_coordination = AssignmentCoordinationService(factory, now_fn=clock.now) + status_service = WorkerUnitStatusService(factory) + return { + "worker_service": worker_service, + "worker_assignment_service": worker_assignment_service, + "engine": engine, + "reconciler": reconciler, + "validator_plane": validator_plane, + "validator_coordination": validator_coordination, + "status_service": status_service, + } + + +async def test_disputed_unit_discoverable_on_postgres( + migrated_postgres_database: str, + cleanup_postgres_database: Callable[[], Awaitable[None]], +) -> None: + engine = create_engine(migrated_postgres_database) + factory = create_session_factory(engine) + clock = _Clock(NOW) + svc = _build_services(factory, clock) + try: + await _add_worker(factory, worker_pubkey="wp-a", miner_hotkey=MINER_A) + await _add_worker(factory, worker_pubkey="wp-b", miner_hotkey=MINER_B) + await _add_gpu_unit(factory, work_unit_id="U", submitter=MINER_H) + await svc["engine"].assign_pending(seed=1) + + await _post_replica( + svc["worker_assignment_service"], + factory, + work_unit_id="U", + miner_hotkey=MINER_A, + manifest=HASH_A, + ) + await _post_replica( + svc["worker_assignment_service"], + factory, + work_unit_id="U", + miner_hotkey=MINER_B, + manifest=HASH_B, + ) + await svc["reconciler"].reconcile_once() + + # Before audit resolution: disputed + linked validator audit unit pending. + units = await svc["status_service"].list_units() + assert [u.work_unit_id for u in units] == ["U"] + unit = units[0] + assert unit.status == WorkAssignmentStatus.DISPUTED.value + assert {r.manifest_sha256 for r in unit.replicas} == {HASH_A, HASH_B} + assert all(r.has_proof for r in unit.replicas) + assert unit.audit is not None + assert unit.audit.work_unit_id == audit_work_unit_id("U") + assert unit.audit.executor_kind == "validator" + assert unit.audit.outcome == AUDIT_OUTCOME_PENDING + + # Validator replays the audit (authoritative HASH_A); reconcile resolves it. + await _add_gpu_validator(factory) + await svc["validator_plane"].assign_pending(seed=1) + async with factory() as session: + audit = ( + await session.execute( + select(WorkAssignment).where( + WorkAssignment.work_unit_id == audit_work_unit_id("U") + ) + ) + ).scalar_one() + await svc["validator_coordination"].post_result( + assignment_id=str(audit.id), + hotkey=GPU_VALIDATOR, + success=True, + payload=_proof_payload(HASH_A), + ) + await svc["reconciler"].reconcile_once() + + units = await svc["status_service"].list_units() + unit = units[0] + assert unit.status == WorkAssignmentStatus.DISPUTED.value + assert unit.audit is not None + assert unit.audit.executor_kind == "validator" + assert unit.audit.outcome == AUDIT_OUTCOME_MISMATCH_RESOLVED + finally: + await engine.dispose() + + +async def test_matched_unit_accepted_with_both_hashes_on_postgres( + migrated_postgres_database: str, + cleanup_postgres_database: Callable[[], Awaitable[None]], +) -> None: + engine = create_engine(migrated_postgres_database) + factory = create_session_factory(engine) + clock = _Clock(NOW) + svc = _build_services(factory, clock) + try: + await _add_worker(factory, worker_pubkey="wp-a", miner_hotkey=MINER_A) + await _add_worker(factory, worker_pubkey="wp-b", miner_hotkey=MINER_B) + await _add_gpu_unit(factory, work_unit_id="U", submitter=MINER_H) + await svc["engine"].assign_pending(seed=1) + + await _post_replica( + svc["worker_assignment_service"], + factory, + work_unit_id="U", + miner_hotkey=MINER_A, + manifest=HASH_A, + ) + await _post_replica( + svc["worker_assignment_service"], + factory, + work_unit_id="U", + miner_hotkey=MINER_B, + manifest=HASH_A, + ) + await svc["reconciler"].reconcile_once() + + units = await svc["status_service"].list_units() + assert [u.work_unit_id for u in units] == ["U"] + unit = units[0] + assert unit.status == WorkAssignmentStatus.COMPLETED.value + assert unit.audit is None + assert [r.manifest_sha256 for r in unit.replicas] == [HASH_A, HASH_A] + assert all(r.has_proof for r in unit.replicas) + assert {r.miner_hotkey for r in unit.replicas} == {MINER_A, MINER_B} + finally: + await engine.dispose() + + +def _signed_headers(*, hotkey: str, nonce: str) -> dict[str, str]: + ts = str(int(NOW_EPOCH)) + canonical = canonical_validator_request( + method="GET", + path="/v1/workers/units", + query_string="", + timestamp=ts, + nonce=nonce, + body=b"", + ) + return { + "X-Hotkey": hotkey, + "X-Signature": _sign(hotkey, canonical.encode()), + "X-Nonce": nonce, + "X-Timestamp": ts, + } + + +async def test_unit_status_http_signed_auth_on_postgres( + migrated_postgres_database: str, + cleanup_postgres_database: Callable[[], Awaitable[None]], +) -> None: + engine = create_engine(migrated_postgres_database) + factory = create_session_factory(engine) + cache = MetagraphCache(netuid=1, ttl_seconds=300) + cache.update_from_metagraph( + [MINER_A, VALIDATOR], + validator_permits=[False, True], + stakes=[0.0, 100.0], + ) + verifier = WorkerSignedRequestVerifier( + nonce_store=SqlAlchemyWorkerNonceStore(factory), + eligibility=CoordinationReadEligibility(factory, cache), + signature_verifier=_fake_verifier, + ttl_seconds=300, + now_fn=lambda: NOW_EPOCH, + ) + app = create_proxy_app( + registry=_TokenRegistry(INTERNAL_BRIDGE_TOKEN), + nonce_store=FakeNonceStore(), + metagraph_cache=FakeCache(), # type: ignore[arg-type] + worker_verifier=verifier, + worker_unit_status_service=WorkerUnitStatusService(factory), + ) + transport = ASGITransport(app=app) + client = AsyncClient(transport=transport, base_url="http://testserver") + try: + await _add_gpu_unit(factory, work_unit_id="U", submitter=MINER_H) + + # (a) valid signed request from an eligible validator succeeds. + ok = await client.get( + "/v1/workers/units", + headers=_signed_headers(hotkey=VALIDATOR, nonce="u-ok"), + ) + assert ok.status_code == 200 + assert [u["work_unit_id"] for u in ok.json()["units"]] == ["U"] + + # (b) missing signature => 401/403. + missing = await client.get("/v1/workers/units") + assert missing.status_code in (401, 403) + + # (c) internal bridge bearer is NOT accepted on this surface. + bearer = await client.get( + "/v1/workers/units", + headers={"Authorization": f"Bearer {INTERNAL_BRIDGE_TOKEN}"}, + ) + assert bearer.status_code in (401, 403) + finally: + await client.aclose() + await engine.dispose() diff --git a/tests/unit/test_compute_lium_client.py b/tests/unit/test_compute_lium_client.py new file mode 100644 index 00000000..4aa7c25a --- /dev/null +++ b/tests/unit/test_compute_lium_client.py @@ -0,0 +1,843 @@ +"""Offline respx tests for :class:`base.compute.lium.LiumClient`. + +Every test mocks Lium HTTP via respx; no credentials and no real network are +required. These pin the provider contract assertions VAL-PROV-001/003/004/005/ +011/017/018 plus secret hygiene for the Lium client. +""" + +from __future__ import annotations + +import json +import logging + +import httpx +import pytest +import respx + +from base.compute import ( + CostGuardrailError, + Instance, + InstanceSpec, + LiumClient, + LiumError, + Offer, +) +from base.compute.lium import ( + _as_list, + _extract_gpu_count, + _extract_gpu_type, + _extract_price, + _parse_instance, + _parse_offer, +) + +BASE = "https://lium.io/api" + +# GET /executors payload shaped like the real API (price_per_gpu + machine_name), +# mixing offers above and below a $1.00/GPU/hr bound. +EXECUTORS = [ + {"id": "a", "machine_name": "RTX 5090", "gpu_count": 8, "price_per_gpu": 0.95}, + {"id": "b", "machine_name": "H100", "gpu_count": 1, "price_per_gpu": 2.0}, + {"id": "c", "machine_name": "RTX 4090", "gpu_count": 2, "price_per_gpu": 0.5}, +] + + +def _spec(**overrides: object) -> InstanceSpec: + base: dict[str, object] = { + "name": "mission-pod", + "template_ref": "prism-worker", + "image": "ghcr.io/base/worker", + "ssh_public_keys": ("ssh-ed25519 AAAA",), + "max_lifetime_hours": 1, + "max_price_per_hour": 1.5, + } + base.update(overrides) + return InstanceSpec(**base) # type: ignore[arg-type] + + +def _offer(price: float = 0.95, offer_id: str = "exec-1") -> Offer: + return Offer(id=offer_id, gpu_type="RTX 5090", gpu_count=1, price_per_hour=price) + + +# -- VAL-PROV-001 ------------------------------------------------------------- + + +@respx.mock +async def test_list_offers_filters_by_max_price() -> None: + respx.get(f"{BASE}/executors").mock( + return_value=httpx.Response(200, json=EXECUTORS) + ) + offers = await LiumClient("k").list_offers(max_price_per_hour=1.0) + assert {o.id for o in offers} == {"a", "c"} + for offer in offers: + assert offer.price_per_hour <= 1.0 + assert offer.gpu_type + assert offer.gpu_count >= 1 + + +@respx.mock +async def test_list_offers_without_bound_returns_all_priced() -> None: + respx.get(f"{BASE}/executors").mock( + return_value=httpx.Response(200, json=EXECUTORS) + ) + offers = await LiumClient("k").list_offers() + assert {o.id for o in offers} == {"a", "b", "c"} + + +@respx.mock +async def test_list_offers_accepts_price_per_hour_field() -> None: + payload = [{"id": "x", "gpu_type": "H100", "gpu_count": 1, "price_per_hour": 0.8}] + respx.get(f"{BASE}/executors").mock(return_value=httpx.Response(200, json=payload)) + offers = await LiumClient("k").list_offers(max_price_per_hour=1.0) + assert offers[0].id == "x" + assert offers[0].price_per_hour == 0.8 + + +@respx.mock +async def test_list_offers_skips_offers_without_price() -> None: + payload = [{"id": "x", "machine_name": "H100", "gpu_count": 1}] + respx.get(f"{BASE}/executors").mock(return_value=httpx.Response(200, json=payload)) + assert await LiumClient("k").list_offers() == [] + + +# -- VAL-PROV-003 ------------------------------------------------------------- + + +@respx.mock +@pytest.mark.parametrize( + "overrides", + [ + {"max_lifetime_hours": None}, + {"max_lifetime_hours": 0}, + {"max_lifetime_hours": -1}, + {"max_price_per_hour": None}, + ], +) +async def test_provision_refuses_unbounded_spec_without_network( + overrides: dict[str, object], +) -> None: + client = LiumClient("k") + with pytest.raises(CostGuardrailError): + await client.provision(_spec(**overrides)) + assert respx.calls.call_count == 0 + + +@respx.mock +@pytest.mark.parametrize("lifetime", [0.5, 0.99, 0.1]) +async def test_provision_rejects_sub_hour_lifetime_without_network( + lifetime: float, +) -> None: + # A sub-1-hour bound would truncate to termination_hours=0 (auto-termination + # disabled -> real-money guardrail bypass); refuse it before any network call. + client = LiumClient("k") + with pytest.raises(CostGuardrailError): + await client.provision(_spec(max_lifetime_hours=lifetime)) + assert respx.calls.call_count == 0 + + +@respx.mock +async def test_provision_never_sends_zero_termination_hours() -> None: + routes = _mock_happy_path() + # A fractional lifetime >= 1 truncates DOWN (never below 1, never 0). + await LiumClient("k").provision(_spec(max_lifetime_hours=1.9), offer=_offer()) + body = json.loads(routes["rent"].calls.last.request.content) + assert body["termination_hours"] == 1 + assert body["termination_hours"] >= 1 + + +# -- VAL-PROV-004 ------------------------------------------------------------- + + +def _mock_happy_path(*, template: list | None = None, keys: list | None = None) -> dict: + routes = { + "ssh_get": respx.get(f"{BASE}/ssh-keys").mock( + return_value=httpx.Response( + 200, + json=keys + if keys is not None + else [{"id": "k1", "public_key": "ssh-ed25519 AAAA"}], + ) + ), + "ssh_post": respx.post(f"{BASE}/ssh-keys").mock( + return_value=httpx.Response(200, json={"id": "k-new"}) + ), + "tpl_get": respx.get(f"{BASE}/templates").mock( + return_value=httpx.Response( + 200, + json=template + if template is not None + else [{"id": "tpl-1", "name": "prism-worker"}], + ) + ), + "tpl_post": respx.post(f"{BASE}/templates").mock( + return_value=httpx.Response(200, json={"id": "tpl-new"}) + ), + "rent": respx.post(f"{BASE}/executors/exec-1/rent").mock( + return_value=httpx.Response(200, json={"id": "pod-1", "status": "PENDING"}) + ), + "status": respx.get(f"{BASE}/pods/pod-1").mock( + return_value=httpx.Response(200, json={"id": "pod-1", "status": "RUNNING"}) + ), + } + return routes + + +@respx.mock +async def test_provision_sends_termination_hours_and_ssh_key() -> None: + routes = _mock_happy_path() + instance = await LiumClient("k").provision( + _spec(max_lifetime_hours=2), offer=_offer() + ) + assert isinstance(instance, Instance) + assert instance.id == "pod-1" + assert instance.status == "RUNNING" + body = json.loads(routes["rent"].calls.last.request.content) + assert body["termination_hours"] == 2 + assert body["pod_name"] == "mission-pod" + assert body["user_public_key"] == ["ssh-ed25519 AAAA"] + assert body["template_id"] == "tpl-1" + + +@respx.mock +async def test_provision_rejects_overpriced_offer_without_rent() -> None: + rent = respx.post(f"{BASE}/executors/exec-1/rent") + client = LiumClient("k") + with pytest.raises(CostGuardrailError): + await client.provision(_spec(max_price_per_hour=1.0), offer=_offer(price=2.0)) + assert rent.call_count == 0 + assert respx.calls.call_count == 0 + + +@respx.mock +async def test_provision_selects_cheapest_within_budget_when_no_offer() -> None: + respx.get(f"{BASE}/executors").mock( + return_value=httpx.Response(200, json=EXECUTORS) + ) + _mock_happy_path() + # cheapest under the bound is "c" at 0.5 -> rent goes to /executors/c/rent + rent_c = respx.post(f"{BASE}/executors/c/rent").mock( + return_value=httpx.Response(200, json={"id": "pod-1", "status": "PENDING"}) + ) + await LiumClient("k").provision(_spec(max_price_per_hour=1.0)) + assert rent_c.call_count == 1 + + +@respx.mock +async def test_provision_raises_guardrail_when_no_offer_within_budget() -> None: + respx.get(f"{BASE}/executors").mock( + return_value=httpx.Response(200, json=EXECUTORS) + ) + rent = respx.post(f"{BASE}/executors/b/rent") + with pytest.raises(CostGuardrailError): + await LiumClient("k").provision(_spec(max_price_per_hour=0.1)) + assert rent.call_count == 0 + + +# -- VAL-PROV-005 ------------------------------------------------------------- + + +@respx.mock +async def test_terminate_is_idempotent() -> None: + route = respx.delete(f"{BASE}/pods/pod-1").mock( + side_effect=[httpx.Response(200), httpx.Response(404)] + ) + client = LiumClient("k") + await client.terminate("pod-1") + await client.terminate("pod-1") + assert route.call_count == 2 + + +@respx.mock +async def test_terminate_raises_on_non_404_error() -> None: + respx.delete(f"{BASE}/pods/pod-1").mock(return_value=httpx.Response(500)) + with pytest.raises(LiumError): + await LiumClient("k").terminate("pod-1") + + +@respx.mock +async def test_verify_terminated_reflects_pod_presence() -> None: + respx.get(f"{BASE}/pods").mock( + side_effect=[ + httpx.Response(200, json=[{"id": "pod-1", "status": "RUNNING"}]), + httpx.Response(200, json=[]), + ] + ) + client = LiumClient("k") + assert await client.verify_terminated("pod-1") is False + assert await client.verify_terminated("pod-1") is True + + +# -- VAL-PROV-011 ------------------------------------------------------------- + + +@respx.mock +async def test_watchtower_digest_returned_verbatim() -> None: + digest = "sha256:" + "b" * 64 + respx.get(f"{BASE}/watchtower/digest").mock( + return_value=httpx.Response( + 200, json={"digest": digest, "signature": "0xabc", "timestamp": 1} + ) + ) + assert await LiumClient("k").watchtower_digest() == digest + + +@respx.mock +async def test_watchtower_digest_missing_field_raises() -> None: + respx.get(f"{BASE}/watchtower/digest").mock( + return_value=httpx.Response(200, json={"signature": "0xabc"}) + ) + with pytest.raises(LiumError): + await LiumClient("k").watchtower_digest() + + +# -- VAL-PROV-017 ------------------------------------------------------------- + + +@respx.mock +async def test_provision_failure_path_terminates_and_verifies() -> None: + respx.get(f"{BASE}/ssh-keys").mock( + return_value=httpx.Response(200, json=[{"public_key": "ssh-ed25519 AAAA"}]) + ) + respx.get(f"{BASE}/templates").mock( + return_value=httpx.Response(200, json=[{"id": "tpl-1", "name": "prism-worker"}]) + ) + respx.post(f"{BASE}/executors/exec-1/rent").mock( + return_value=httpx.Response(200, json={"id": "pod-1"}) + ) + # The post-rent status poll fails mid-provision. + respx.get(f"{BASE}/pods/pod-1").mock(return_value=httpx.Response(500)) + delete = respx.delete(f"{BASE}/pods/pod-1").mock(return_value=httpx.Response(200)) + pods = respx.get(f"{BASE}/pods").mock(return_value=httpx.Response(200, json=[])) + + with pytest.raises(LiumError): + await LiumClient("k").provision(_spec(), offer=_offer()) + + assert delete.call_count == 1 + assert pods.called # verify_terminated polled GET /pods after the DELETE + + +# -- VAL-PROV-018 ------------------------------------------------------------- + + +@respx.mock +async def test_ensure_helpers_idempotent_when_present() -> None: + routes = _mock_happy_path() + client = LiumClient("k") + # Repeated planning runs never re-create an existing template/key. + await client.provision(_spec(), offer=_offer()) + await client.provision(_spec(), offer=_offer()) + assert routes["ssh_post"].call_count == 0 + assert routes["tpl_post"].call_count == 0 + body = json.loads(routes["rent"].calls.last.request.content) + assert body["template_id"] == "tpl-1" + + +@respx.mock +async def test_ensure_helpers_create_once_when_absent() -> None: + routes = _mock_happy_path(template=[], keys=[]) + await LiumClient("k").provision(_spec(), offer=_offer()) + assert routes["ssh_post"].call_count == 1 + assert routes["tpl_post"].call_count == 1 + body = json.loads(routes["rent"].calls.last.request.content) + assert body["template_id"] == "tpl-new" + + +@respx.mock +async def test_ensure_template_returns_existing_id() -> None: + respx.get(f"{BASE}/templates").mock( + return_value=httpx.Response(200, json=[{"id": "tpl-9", "name": "prism-worker"}]) + ) + post = respx.post(f"{BASE}/templates") + result = await LiumClient("k").ensure_template( + name="prism-worker", docker_image="img" + ) + assert result == "tpl-9" + assert post.call_count == 0 + + +@respx.mock +async def test_ensure_ssh_key_creates_when_absent() -> None: + respx.get(f"{BASE}/ssh-keys").mock(return_value=httpx.Response(200, json=[])) + post = respx.post(f"{BASE}/ssh-keys").mock( + return_value=httpx.Response(200, json={"id": "k-new", "public_key": "ssh x"}) + ) + result = await LiumClient("k").ensure_ssh_key(public_key="ssh x", name="deploy") + assert result["id"] == "k-new" + assert post.call_count == 1 + assert json.loads(post.calls.last.request.content)["public_key"] == "ssh x" + + +@respx.mock +async def test_ensure_template_body_pins_digest_tag_env_and_ports() -> None: + respx.get(f"{BASE}/templates").mock(return_value=httpx.Response(200, json=[])) + post = respx.post(f"{BASE}/templates").mock( + return_value=httpx.Response(200, json={"id": "tpl-new"}) + ) + digest = "sha256:" + "d" * 64 + template_id = await LiumClient("k").ensure_template( + name="prism-worker", + docker_image="ghcr.io/base/worker", + docker_image_digest=digest, + docker_image_tag="v1", + internal_ports=(22, 8080), + environment={"ROLE": "worker"}, + ) + assert template_id == "tpl-new" + body = json.loads(post.calls.last.request.content) + assert body["docker_image_digest"] == digest + assert body["docker_image_tag"] == "v1" + assert body["environment"] == {"ROLE": "worker"} + assert body["internal_ports"] == [22, 8080] + assert body["is_private"] is True + + +@respx.mock +async def test_ensure_template_includes_startup_commands_when_given() -> None: + respx.get(f"{BASE}/templates").mock(return_value=httpx.Response(200, json=[])) + post = respx.post(f"{BASE}/templates").mock( + return_value=httpx.Response(200, json={"id": "tpl-new"}) + ) + await LiumClient("k").ensure_template( + name="prism-worker", + docker_image="ghcr.io/base/worker", + startup_commands="tail -f /dev/null", + ) + body = json.loads(post.calls.last.request.content) + assert body["startup_commands"] == "tail -f /dev/null" + + +@respx.mock +async def test_ensure_template_omits_startup_commands_when_none() -> None: + respx.get(f"{BASE}/templates").mock(return_value=httpx.Response(200, json=[])) + post = respx.post(f"{BASE}/templates").mock( + return_value=httpx.Response(200, json={"id": "tpl-new"}) + ) + await LiumClient("k").ensure_template( + name="prism-worker", docker_image="ghcr.io/base/worker" + ) + body = json.loads(post.calls.last.request.content) + assert "startup_commands" not in body + + +@respx.mock +async def test_provision_forwards_spec_startup_commands_to_template() -> None: + routes = _mock_happy_path(template=[]) + spec = _spec(startup_commands="tail -f /dev/null") + await LiumClient("k").provision(spec, offer=_offer()) + body = json.loads(routes["tpl_post"].calls.last.request.content) + assert body["startup_commands"] == "tail -f /dev/null" + + +# -- loopback env hygiene: Lium edge-WAF 403s on loopback URLs in the body ----- + + +@respx.mock +async def test_ensure_template_strips_loopback_urls_from_environment() -> None: + respx.get(f"{BASE}/templates").mock(return_value=httpx.Response(200, json=[])) + post = respx.post(f"{BASE}/templates").mock( + return_value=httpx.Response(200, json={"id": "tpl-new"}) + ) + await LiumClient("k").ensure_template( + name="prism-worker", + docker_image="ghcr.io/base/worker", + environment={ + "BASE_WORKER__AGENT__MASTER_URL": "http://127.0.0.1:8081", + "BASE_WORKER__AGENT__BROKER_URL": "http://localhost:8082", + "ROLE": "worker", + }, + ) + body = json.loads(post.calls.last.request.content) + blob = json.dumps(body) + assert "127.0.0.1" not in blob + assert "localhost" not in blob + # Non-loopback env still travels. + assert body["environment"] == {"ROLE": "worker"} + + +@respx.mock +async def test_ensure_template_omits_environment_when_all_loopback() -> None: + respx.get(f"{BASE}/templates").mock(return_value=httpx.Response(200, json=[])) + post = respx.post(f"{BASE}/templates").mock( + return_value=httpx.Response(200, json={"id": "tpl-new"}) + ) + await LiumClient("k").ensure_template( + name="prism-worker", + docker_image="ghcr.io/base/worker", + environment={"BASE_WORKER__AGENT__MASTER_URL": "http://127.0.0.1:8081"}, + ) + body = json.loads(post.calls.last.request.content) + assert "environment" not in body + + +@respx.mock +async def test_provision_template_post_carries_no_loopback_url() -> None: + routes = _mock_happy_path(template=[]) + spec = _spec( + env={ + "BASE_WORKER__AGENT__MASTER_URL": "http://127.0.0.1:8081", + "ROLE": "worker", + } + ) + await LiumClient("k").provision(spec, offer=_offer()) + body = json.loads(routes["tpl_post"].calls.last.request.content) + blob = json.dumps(body) + assert "127.0.0.1" not in blob + assert "localhost" not in blob + assert body["environment"] == {"ROLE": "worker"} + + +# -- status / logs / balance ------------------------------------------------- + + +@respx.mock +async def test_status_parses_pod_detail() -> None: + respx.get(f"{BASE}/pods/pod-1").mock( + return_value=httpx.Response(200, json={"id": "pod-1", "status": "RUNNING"}) + ) + instance = await LiumClient("k").status("pod-1") + assert instance.id == "pod-1" + assert instance.status == "RUNNING" + assert instance.provider == "lium" + + +@respx.mock +async def test_stream_logs_yields_lines() -> None: + respx.get(f"{BASE}/pods/pod-1/logs").mock( + return_value=httpx.Response(200, text="line-1\nline-2\n") + ) + lines = [line async for line in LiumClient("k").stream_logs("pod-1")] + assert lines == ["line-1", "line-2"] + + +@respx.mock +async def test_stream_logs_raises_on_error() -> None: + respx.get(f"{BASE}/pods/pod-1/logs").mock(return_value=httpx.Response(404)) + with pytest.raises(LiumError): + async for _ in LiumClient("k").stream_logs("pod-1"): + pass + + +@respx.mock +async def test_balance_returns_float() -> None: + respx.get(f"{BASE}/users/me").mock( + return_value=httpx.Response(200, json={"balance": 9.99}) + ) + assert await LiumClient("k").balance() == pytest.approx(9.99) + + +@respx.mock +async def test_request_error_raises_lium_error() -> None: + respx.get(f"{BASE}/users/me").mock(return_value=httpx.Response(500)) + with pytest.raises(LiumError) as exc_info: + await LiumClient("k").balance() + assert exc_info.value.status_code == 500 + + +# -- secret hygiene ----------------------------------------------------------- + + +@respx.mock +async def test_api_key_sent_only_in_header() -> None: + route = respx.get(f"{BASE}/users/me").mock( + return_value=httpx.Response(200, json={"balance": 1.0}) + ) + await LiumClient("MY-SECRET").balance() + assert route.calls.last.request.headers["X-API-Key"] == "MY-SECRET" + + +@respx.mock +async def test_api_key_never_in_repr_str_logs_or_errors( + caplog: pytest.LogCaptureFixture, +) -> None: + sentinel = "SENTINEL-LIUM-KEY-XYZ" + client = LiumClient(sentinel) + assert sentinel not in repr(client) + assert sentinel not in str(client) + respx.get(f"{BASE}/users/me").mock(return_value=httpx.Response(500)) + with caplog.at_level(logging.DEBUG): + with pytest.raises(LiumError) as exc_info: + await client.balance() + assert sentinel not in str(exc_info.value) + assert sentinel not in caplog.text + + +# -- provision fallback / cleanup edge paths --------------------------------- + + +@respx.mock +async def test_provision_falls_back_to_pod_lookup_by_name() -> None: + respx.get(f"{BASE}/ssh-keys").mock( + return_value=httpx.Response(200, json=[{"public_key": "ssh-ed25519 AAAA"}]) + ) + respx.get(f"{BASE}/templates").mock( + return_value=httpx.Response(200, json=[{"id": "tpl-1", "name": "prism-worker"}]) + ) + # Rent response carries no id -> the client resolves it via GET /pods by name. + respx.post(f"{BASE}/executors/exec-1/rent").mock( + return_value=httpx.Response(200, json={}) + ) + respx.get(f"{BASE}/pods").mock( + return_value=httpx.Response( + 200, json=[{"id": "pod-7", "pod_name": "mission-pod"}] + ) + ) + respx.get(f"{BASE}/pods/pod-7").mock( + return_value=httpx.Response(200, json={"id": "pod-7", "status": "RUNNING"}) + ) + instance = await LiumClient("k").provision(_spec(), offer=_offer()) + assert instance.id == "pod-7" + + +@respx.mock +async def test_provision_raises_when_pod_id_undeterminable() -> None: + respx.get(f"{BASE}/ssh-keys").mock( + return_value=httpx.Response(200, json=[{"public_key": "ssh-ed25519 AAAA"}]) + ) + respx.get(f"{BASE}/templates").mock( + return_value=httpx.Response(200, json=[{"id": "tpl-1", "name": "prism-worker"}]) + ) + respx.post(f"{BASE}/executors/exec-1/rent").mock( + return_value=httpx.Response(200, json={}) + ) + respx.get(f"{BASE}/pods").mock(return_value=httpx.Response(200, json=[])) + with pytest.raises(LiumError): + await LiumClient("k").provision(_spec(), offer=_offer()) + + +@respx.mock +async def test_provision_cleanup_swallows_cleanup_errors_and_reraises() -> None: + respx.get(f"{BASE}/ssh-keys").mock( + return_value=httpx.Response(200, json=[{"public_key": "ssh-ed25519 AAAA"}]) + ) + respx.get(f"{BASE}/templates").mock( + return_value=httpx.Response(200, json=[{"id": "tpl-1", "name": "prism-worker"}]) + ) + respx.post(f"{BASE}/executors/exec-1/rent").mock( + return_value=httpx.Response(200, json={"id": "pod-1"}) + ) + respx.get(f"{BASE}/pods/pod-1").mock(return_value=httpx.Response(500)) + # Cleanup itself fails, but the original provisioning error must still surface. + delete = respx.delete(f"{BASE}/pods/pod-1").mock(return_value=httpx.Response(500)) + respx.get(f"{BASE}/pods").mock(return_value=httpx.Response(500)) + with pytest.raises(LiumError): + await LiumClient("k").provision(_spec(), offer=_offer()) + assert delete.call_count == 1 + + +@respx.mock +async def test_provision_cleanup_uses_rent_pod_id_when_status_fails() -> None: + respx.get(f"{BASE}/ssh-keys").mock( + return_value=httpx.Response(200, json=[{"public_key": "ssh-ed25519 AAAA"}]) + ) + respx.get(f"{BASE}/templates").mock( + return_value=httpx.Response(200, json=[{"id": "tpl-1", "name": "prism-worker"}]) + ) + # Live rent shape {"success": true, "pod_id": "..."}; the status poll then + # fails and cleanup must terminate using the id from the rent response. + respx.post(f"{BASE}/executors/exec-1/rent").mock( + return_value=httpx.Response(200, json={"success": True, "pod_id": "pod-1"}) + ) + respx.get(f"{BASE}/pods/pod-1").mock(return_value=httpx.Response(500)) + delete = respx.delete(f"{BASE}/pods/pod-1").mock(return_value=httpx.Response(200)) + pods = respx.get(f"{BASE}/pods").mock(return_value=httpx.Response(200, json=[])) + + with pytest.raises(LiumError): + await LiumClient("k").provision(_spec(), offer=_offer()) + + assert delete.call_count == 1 + assert pods.called + + +@respx.mock +async def test_provision_cleanup_survives_transient_pods_failure_in_resolution() -> ( + None +): + respx.get(f"{BASE}/ssh-keys").mock( + return_value=httpx.Response(200, json=[{"public_key": "ssh-ed25519 AAAA"}]) + ) + respx.get(f"{BASE}/templates").mock( + return_value=httpx.Response(200, json=[{"id": "tpl-1", "name": "prism-worker"}]) + ) + # Rent SUCCEEDS but returns no id, so pod-id resolution falls back to + # GET /pods, which fails transiently on the first call. Because cleanup keys + # off "rent succeeded" (not "pod id resolved"), the just-rented pod is still + # found by name on the cleanup retry and deleted. + respx.post(f"{BASE}/executors/exec-1/rent").mock( + return_value=httpx.Response(200, json={"success": True}) + ) + respx.get(f"{BASE}/pods").mock( + side_effect=[ + httpx.Response(500), + httpx.Response(200, json=[{"id": "pod-9", "pod_name": "mission-pod"}]), + httpx.Response(200, json=[]), + ] + ) + delete = respx.delete(f"{BASE}/pods/pod-9").mock(return_value=httpx.Response(200)) + + with pytest.raises(LiumError): + await LiumClient("k").provision(_spec(), offer=_offer()) + + assert delete.call_count == 1 + + +@respx.mock +async def test_provision_cleanup_after_unparseable_rent_body() -> None: + respx.get(f"{BASE}/ssh-keys").mock( + return_value=httpx.Response(200, json=[{"public_key": "ssh-ed25519 AAAA"}]) + ) + respx.get(f"{BASE}/templates").mock( + return_value=httpx.Response(200, json=[{"id": "tpl-1", "name": "prism-worker"}]) + ) + # A 2xx rent with a NON-JSON/garbage body: pod-id extraction raises while + # parsing. Because the rent HTTP call SUCCEEDED, a billable pod may now exist, + # so provision must still terminate + verify the just-rented pod (resolved by + # name) before re-raising -- no leaked pod. + respx.post(f"{BASE}/executors/exec-1/rent").mock( + return_value=httpx.Response(200, content=b"not json") + ) + pods = respx.get(f"{BASE}/pods").mock( + side_effect=[ + httpx.Response(200, json=[{"id": "pod-3", "pod_name": "mission-pod"}]), + httpx.Response(200, json=[]), + ] + ) + delete = respx.delete(f"{BASE}/pods/pod-3").mock(return_value=httpx.Response(200)) + + with pytest.raises(LiumError): + await LiumClient("k").provision(_spec(), offer=_offer()) + + assert delete.call_count == 1 # the just-rented pod was terminated + assert pods.call_count == 2 # resolved by name, then verify_terminated polled + + +@respx.mock +async def test_provision_with_dockerfile_content_omits_template() -> None: + respx.get(f"{BASE}/ssh-keys").mock( + return_value=httpx.Response(200, json=[{"public_key": "ssh-ed25519 AAAA"}]) + ) + rent = respx.post(f"{BASE}/executors/exec-1/rent").mock( + return_value=httpx.Response(200, json={"id": "pod-1", "status": "PENDING"}) + ) + respx.get(f"{BASE}/pods/pod-1").mock( + return_value=httpx.Response(200, json={"id": "pod-1", "status": "RUNNING"}) + ) + templates = respx.get(f"{BASE}/templates") + spec = InstanceSpec( + name="mission-pod", + template_ref=None, + dockerfile_content="FROM ubuntu:22.04", + ssh_public_keys=("ssh-ed25519 AAAA",), + max_lifetime_hours=1, + max_price_per_hour=1.5, + ) + await LiumClient("k").provision(spec, offer=_offer()) + body = json.loads(rent.calls.last.request.content) + assert body["dockerfile_content"] == "FROM ubuntu:22.04" + assert "template_id" not in body + assert templates.call_count == 0 # no template ensure for the dockerfile path + + +@respx.mock +async def test_provision_requires_template_or_dockerfile() -> None: + respx.get(f"{BASE}/ssh-keys").mock( + return_value=httpx.Response(200, json=[{"public_key": "ssh-ed25519 AAAA"}]) + ) + rent = respx.post(f"{BASE}/executors/exec-1/rent") + spec = InstanceSpec( + name="mission-pod", + template_ref=None, + dockerfile_content=None, + ssh_public_keys=("ssh-ed25519 AAAA",), + max_lifetime_hours=1, + max_price_per_hour=1.5, + ) + with pytest.raises(LiumError): + await LiumClient("k").provision(spec, offer=_offer()) + assert rent.call_count == 0 + + +@respx.mock +async def test_provision_requires_ssh_key() -> None: + spec = InstanceSpec( + name="mission-pod", + template_ref="prism-worker", + ssh_public_keys=(), + max_lifetime_hours=1, + max_price_per_hour=1.5, + ) + with pytest.raises(LiumError): + await LiumClient("k").provision(spec, offer=_offer()) + assert respx.calls.call_count == 0 + + +@respx.mock +async def test_transport_error_wrapped_as_lium_error() -> None: + respx.get(f"{BASE}/users/me").mock(side_effect=httpx.ConnectError("down")) + with pytest.raises(LiumError): + await LiumClient("k").balance() + + +@respx.mock +async def test_balance_missing_field_raises() -> None: + respx.get(f"{BASE}/users/me").mock(return_value=httpx.Response(200, json={})) + with pytest.raises(LiumError): + await LiumClient("k").balance() + + +@respx.mock +async def test_watchtower_digest_accepts_plain_string() -> None: + respx.get(f"{BASE}/watchtower/digest").mock( + return_value=httpx.Response(200, json="sha256:" + "c" * 64) + ) + assert await LiumClient("k").watchtower_digest() == "sha256:" + "c" * 64 + + +@respx.mock +async def test_list_offers_accepts_wrapped_payload() -> None: + respx.get(f"{BASE}/executors").mock( + return_value=httpx.Response(200, json={"executors": EXECUTORS}) + ) + offers = await LiumClient("k").list_offers(max_price_per_hour=1.0) + assert {o.id for o in offers} == {"a", "c"} + + +# -- parsing helpers ---------------------------------------------------------- + + +def test_extract_price_prefers_explicit_then_per_gpu() -> None: + assert _extract_price({"price_per_hour": 1.5}) == 1.5 + assert _extract_price({"price_per_gpu": 0.9}) == 0.9 + assert _extract_price({"pending_price_per_hour": 0.7}) == 0.7 + assert _extract_price({}) is None + assert _extract_price({"price_per_gpu": "not-a-number"}) is None + + +def test_extract_gpu_type_falls_back_to_specs_details() -> None: + item = {"specs": {"gpu": {"details": [{"name": "NVIDIA H100"}]}}} + assert _extract_gpu_type(item) == "NVIDIA H100" + assert _extract_gpu_type({}) == "" + + +def test_extract_gpu_count_falls_back_to_specs_count() -> None: + assert _extract_gpu_count({"gpu_count": 4}) == 4 + assert _extract_gpu_count({"specs": {"gpu": {"count": 2}}}) == 2 + assert _extract_gpu_count({}) == 0 + assert _extract_gpu_count({"gpu_count": "x"}) == 0 + + +def test_parse_offer_skips_items_without_id_or_price() -> None: + assert _parse_offer({"machine_name": "H100", "gpu_count": 1}) is None + parsed = _parse_offer({"id": "z", "price_per_gpu": 1.0}) + assert parsed is not None + assert parsed.id == "z" + + +def test_parse_instance_rejects_non_mapping() -> None: + with pytest.raises(LiumError): + _parse_instance(["not", "a", "mapping"]) + + +def test_as_list_handles_unexpected_shapes() -> None: + assert _as_list("nope", "executors") == [] + assert _as_list({"executors": "nope"}, "executors") == [] + assert _as_list([{"a": 1}, "skip"], "executors") == [{"a": 1}] diff --git a/tests/unit/test_compute_provider.py b/tests/unit/test_compute_provider.py new file mode 100644 index 00000000..b8f7a766 --- /dev/null +++ b/tests/unit/test_compute_provider.py @@ -0,0 +1,42 @@ +"""Tests for the provider-agnostic compute contract (:mod:`base.compute.provider`).""" + +from __future__ import annotations + +from base.compute import ( + CostGuardrailError, + Instance, + InstanceSpec, + LiumClient, + Offer, + ProviderClient, + ProviderError, +) + + +def test_cost_guardrail_error_is_provider_error() -> None: + assert issubclass(CostGuardrailError, ProviderError) + + +def test_instance_spec_defaults_are_conservative() -> None: + spec = InstanceSpec(name="pod") + # Guardrail fields default to unset so provision() can reject an unbounded spec. + assert spec.max_lifetime_hours is None + assert spec.max_price_per_hour is None + assert spec.ports == (22,) + assert spec.gpu_count == 1 + assert spec.ssh_public_keys == () + + +def test_offer_and_instance_are_frozen_dataclasses() -> None: + offer = Offer(id="o", gpu_type="H100", gpu_count=1, price_per_hour=2.0) + instance = Instance(id="p", status="RUNNING") + assert offer.price_per_hour == 2.0 + assert instance.status == "RUNNING" + + +def test_lium_client_satisfies_provider_protocol() -> None: + client = LiumClient("k") + # runtime_checkable protocol: verifies the full method surface is present. + assert isinstance(client, ProviderClient) + typed: ProviderClient = client + assert typed is client diff --git a/tests/unit/test_compute_secret_hygiene.py b/tests/unit/test_compute_secret_hygiene.py new file mode 100644 index 00000000..2e17cf86 --- /dev/null +++ b/tests/unit/test_compute_secret_hygiene.py @@ -0,0 +1,85 @@ +"""Secret-hygiene tests across BOTH compute provider clients (VAL-PROV-006). + +Constructs each client with a sentinel API key, exercises a successful mocked +call and a failing (HTTP 500) mocked call, then asserts the sentinel never leaks +into ``repr``, ``str``, any log record (root logger, DEBUG level), or the +stringified raised exception. +""" + +from __future__ import annotations + +import logging + +import httpx +import pytest +import respx + +from base.compute import LiumClient, LiumError, TargonClient, TargonError + +LIUM_BASE = "https://lium.io/api" +TARGON_BASE = "https://api.targon.com/tha/v2" + +LIUM_SENTINEL = "SENTINEL-LIUM-KEY-XYZ" +TARGON_SENTINEL = "SENTINEL-TARGON-KEY-XYZ" + + +@respx.mock +async def test_lium_client_never_leaks_key( + caplog: pytest.LogCaptureFixture, +) -> None: + client = LiumClient(LIUM_SENTINEL) + assert LIUM_SENTINEL not in repr(client) + assert LIUM_SENTINEL not in str(client) + + with caplog.at_level(logging.DEBUG, logger=""): + # Successful call. + respx.get(f"{LIUM_BASE}/users/me").mock( + return_value=httpx.Response(200, json={"balance": 1.0}) + ) + assert await client.balance() == pytest.approx(1.0) + + # Failing call (HTTP 500). + respx.get(f"{LIUM_BASE}/pods").mock(return_value=httpx.Response(500)) + with pytest.raises(LiumError) as exc_info: + await client.list_pods() + + assert LIUM_SENTINEL not in str(exc_info.value) + assert LIUM_SENTINEL not in repr(exc_info.value) + assert LIUM_SENTINEL not in caplog.text + + +@respx.mock +async def test_targon_client_never_leaks_key( + caplog: pytest.LogCaptureFixture, +) -> None: + client = TargonClient(TARGON_SENTINEL) + assert TARGON_SENTINEL not in repr(client) + assert TARGON_SENTINEL not in str(client) + + with caplog.at_level(logging.DEBUG, logger=""): + # Successful call. + respx.get(f"{TARGON_BASE}/workloads").mock( + return_value=httpx.Response(200, json={"items": []}) + ) + assert await client.list_workloads() == [] + + # Failing call (HTTP 500). + respx.get(f"{TARGON_BASE}/apps").mock(return_value=httpx.Response(500)) + with pytest.raises(TargonError) as exc_info: + await client.list_apps() + + assert TARGON_SENTINEL not in str(exc_info.value) + assert TARGON_SENTINEL not in repr(exc_info.value) + assert TARGON_SENTINEL not in caplog.text + + +@respx.mock +async def test_targon_deploy_failure_never_leaks_key() -> None: + client = TargonClient(TARGON_SENTINEL) + respx.post(f"{TARGON_BASE}/workloads").mock( + return_value=httpx.Response(402, json={"error": "payment required"}) + ) + with pytest.raises(TargonError) as exc_info: + await client.deploy({"name": "x"}) + assert TARGON_SENTINEL not in str(exc_info.value) + assert TARGON_SENTINEL not in repr(exc_info.value) diff --git a/tests/unit/test_compute_targon_client.py b/tests/unit/test_compute_targon_client.py new file mode 100644 index 00000000..c779d376 --- /dev/null +++ b/tests/unit/test_compute_targon_client.py @@ -0,0 +1,632 @@ +"""Offline respx tests for :class:`base.compute.targon.TargonClient`. + +Every test mocks Targon HTTP via respx; no credentials and no real network are +required. These pin the provider contract assertions VAL-PROV-002/003/006/007/008 +for the Targon client. +""" + +from __future__ import annotations + +import json + +import httpx +import pytest +import respx + +from base.compute import ( + BalanceUnavailableError, + CostGuardrailError, + Instance, + InstanceSpec, + InsufficientCreditsError, + ProviderError, + TargonClient, + TargonError, +) +from base.compute.targon import ( + _as_list, + _is_insufficient_credits, + _parse_inventory_offer, + _parse_workload_instance, +) + +BASE = "https://api.targon.com/tha/v2" + +# GET /inventory?type=rental&gpu=true payload shaped like the real API: a bare +# list of fixed GPU shapes carrying cost_per_hour + availability, mixing zero and +# non-zero availability and prices above/below a $3.00/GPU/hr bound. +INVENTORY = [ + { + "name": "h100", + "display_name": "H100 x1", + "type": "rental", + "gpu": True, + "spec": {"gpu_type": "H100", "gpu_count": 1}, + "cost_per_hour": 2.5, + "available": 4, + }, + { + "name": "h200", + "display_name": "H200 x1", + "type": "rental", + "gpu": True, + "spec": {"gpu_type": "H200", "gpu_count": 1}, + "cost_per_hour": 3.29, + "available": 2, + }, + { + "name": "h100x8", + "display_name": "H100 x8", + "type": "rental", + "gpu": True, + "spec": {"gpu_type": "H100", "gpu_count": 8}, + "cost_per_hour": 20.0, + "available": 0, + }, +] + + +def _spec(**overrides: object) -> InstanceSpec: + base: dict[str, object] = { + "name": "mission-workload", + "template_ref": "h100", + "image": "ghcr.io/base/worker", + "ssh_public_keys": ("ssh-ed25519 AAAA",), + "max_lifetime_hours": 1, + "max_price_per_hour": 3.0, + } + base.update(overrides) + return InstanceSpec(**base) # type: ignore[arg-type] + + +# -- VAL-PROV-002 ------------------------------------------------------------- + + +@respx.mock +async def test_list_offers_excludes_zero_availability_and_over_price() -> None: + respx.get(f"{BASE}/inventory").mock( + return_value=httpx.Response(200, json=INVENTORY) + ) + offers = await TargonClient("k").list_offers(max_price_per_hour=3.0) + assert {o.id for o in offers} == {"h100"} + only = offers[0] + assert only.price_per_hour == pytest.approx(2.5) + assert isinstance(only.price_per_hour, float) + assert only.gpu_type == "H100" + assert only.gpu_count == 1 + assert only.provider == "targon" + assert only.raw["cost_per_hour"] == pytest.approx(2.5) + + +@respx.mock +async def test_list_offers_without_bound_excludes_only_zero_availability() -> None: + respx.get(f"{BASE}/inventory").mock( + return_value=httpx.Response(200, json=INVENTORY) + ) + offers = await TargonClient("k").list_offers() + assert {o.id for o in offers} == {"h100", "h200"} + for offer in offers: + assert offer.price_per_hour > 0 + + +@respx.mock +async def test_list_offers_sends_rental_gpu_query_params() -> None: + route = respx.get(f"{BASE}/inventory").mock( + return_value=httpx.Response(200, json=INVENTORY) + ) + await TargonClient("k").list_offers() + request = route.calls.last.request + assert request.url.params["type"] == "rental" + assert request.url.params["gpu"] == "true" + + +@respx.mock +async def test_list_offers_accepts_wrapped_payload() -> None: + respx.get(f"{BASE}/inventory").mock( + return_value=httpx.Response(200, json={"items": INVENTORY}) + ) + offers = await TargonClient("k").list_offers() + assert {o.id for o in offers} == {"h100", "h200"} + + +@respx.mock +async def test_list_offers_normalizes_price_to_per_gpu() -> None: + inventory = [ + { + "name": "h100x4", + "display_name": "H100 x4", + "type": "rental", + "gpu": True, + "spec": {"gpu_type": "H100", "gpu_count": 4}, + "cost_per_hour": 10.0, + "available": 2, + } + ] + respx.get(f"{BASE}/inventory").mock( + return_value=httpx.Response(200, json=inventory) + ) + offers = await TargonClient("k").list_offers() + assert len(offers) == 1 + offer = offers[0] + # Whole-shape cost is 10.0 for 4 GPUs -> per-GPU price is 2.5. + assert offer.price_per_hour == pytest.approx(2.5) + assert offer.gpu_count == 4 + assert offer.raw["cost_per_hour"] == pytest.approx(10.0) + + +@respx.mock +async def test_multi_gpu_offer_filtered_by_per_gpu_max_price() -> None: + inventory = [ + { + "name": "h100x4", + "spec": {"gpu_type": "H100", "gpu_count": 4}, + "cost_per_hour": 10.0, + "available": 2, + }, + { + "name": "h200x8", + "spec": {"gpu_type": "H200", "gpu_count": 8}, + "cost_per_hour": 40.0, + "available": 2, + }, + ] + respx.get(f"{BASE}/inventory").mock( + return_value=httpx.Response(200, json=inventory) + ) + # Per-GPU: h100x4 -> 2.5/gpu (kept), h200x8 -> 5.0/gpu (dropped by 3.0 cap). + # A whole-shape comparison (10 and 40) would have wrongly dropped BOTH. + offers = await TargonClient("k").list_offers(max_price_per_hour=3.0) + assert {o.id for o in offers} == {"h100x4"} + assert offers[0].price_per_hour == pytest.approx(2.5) + + +@respx.mock +async def test_multi_gpu_offer_count_from_numeric_field_filters_correctly() -> None: + inventory = [ + { + "name": "h100x8", + # ``spec`` is present but carries no usable gpu_count; the count lives + # ONLY in the top-level numeric field. Per-GPU price must still be + # cost/count so the per-GPU cap filters the shape correctly. + "spec": {"gpu_type": "H100", "gpu_count": 0}, + "gpu_count": 8, + "cost_per_hour": 16.0, + "available": 2, + } + ] + respx.get(f"{BASE}/inventory").mock( + return_value=httpx.Response(200, json=inventory) + ) + # Per-GPU: 16.0 / 8 = 2.0 (within the 3.0/GPU cap). A whole-shape fallback + # (16.0 > 3.0) would have WRONGLY dropped this in-budget multi-GPU shape. + offers = await TargonClient("k").list_offers(max_price_per_hour=3.0) + assert len(offers) == 1 + assert offers[0].gpu_count == 8 + assert offers[0].price_per_hour == pytest.approx(2.0) + + +# -- VAL-PROV-003 (provision guardrails, no network) -------------------------- + + +@respx.mock +@pytest.mark.parametrize( + "overrides", + [ + {"max_lifetime_hours": None}, + {"max_lifetime_hours": 0}, + {"max_lifetime_hours": -1}, + {"max_lifetime_hours": 0.5}, + {"max_price_per_hour": None}, + {"max_price_per_hour": 0}, + ], +) +async def test_provision_refuses_unbounded_spec_without_network( + overrides: dict[str, object], +) -> None: + create = respx.post(f"{BASE}/workloads") + with pytest.raises(CostGuardrailError): + await TargonClient("k").provision(_spec(**overrides)) + assert create.call_count == 0 + assert respx.calls.call_count == 0 + + +async def test_provision_sub_hour_lifetime_message_mentions_truncation() -> None: + with pytest.raises(CostGuardrailError) as exc_info: + await TargonClient("k").provision(_spec(max_lifetime_hours=0.5)) + message = str(exc_info.value) + assert "at least 1 hour" in message + assert "termination_hours" in message + + +# -- deploy call shape (two-step create-then-deploy) -------------------------- + + +@respx.mock +async def test_provision_creates_then_deploys_with_workload_body() -> None: + create = respx.post(f"{BASE}/workloads").mock( + return_value=httpx.Response(200, json={"uid": "wl-1"}) + ) + deploy = respx.post(f"{BASE}/workloads/wl-1/deploy").mock( + return_value=httpx.Response( + 200, json={"uid": "wl-1", "state": {"status": "PENDING"}} + ) + ) + instance = await TargonClient("k").provision(_spec(max_lifetime_hours=2)) + assert isinstance(instance, Instance) + assert instance.id == "wl-1" + assert instance.status == "PENDING" + assert instance.provider == "targon" + assert create.call_count == 1 + assert deploy.call_count == 1 + body = json.loads(create.calls.last.request.content) + assert body["name"] == "mission-workload" + assert body["type"] == "RENTAL" + assert body["resource_name"] == "h100" + assert body["image"] == "ghcr.io/base/worker" + assert body["termination_hours"] == 2 + assert body["envs"] == [] or isinstance(body["envs"], list) + # The deploy step carries no body (matches the SDK / live route). + assert not deploy.calls.last.request.content + + +@respx.mock +async def test_deploy_uses_create_uid_when_deploy_response_omits_it() -> None: + respx.post(f"{BASE}/workloads").mock( + return_value=httpx.Response(200, json={"uid": "wl-42"}) + ) + respx.post(f"{BASE}/workloads/wl-42/deploy").mock( + return_value=httpx.Response(200, json={"state": {"status": "DEPLOYING"}}) + ) + instance = await TargonClient("k").deploy({"name": "custom"}) + assert instance.id == "wl-42" + assert instance.status == "DEPLOYING" + + +@respx.mock +async def test_deploy_passes_explicit_payload_verbatim_to_create() -> None: + create = respx.post(f"{BASE}/workloads").mock( + return_value=httpx.Response(200, json={"uid": "wl-9"}) + ) + respx.post(f"{BASE}/workloads/wl-9/deploy").mock( + return_value=httpx.Response(200, json={"uid": "wl-9", "status": "RUNNING"}) + ) + payload = {"name": "custom", "resource_name": "h200", "image": "img"} + instance = await TargonClient("k").deploy(payload) + assert instance.id == "wl-9" + assert instance.status == "RUNNING" + assert json.loads(create.calls.last.request.content) == payload + + +@respx.mock +async def test_deploy_raises_when_create_returns_no_uid() -> None: + create = respx.post(f"{BASE}/workloads").mock( + return_value=httpx.Response(200, json={"state": {"status": "PENDING"}}) + ) + deploy = respx.post(f"{BASE}/workloads//deploy") + with pytest.raises(TargonError): + await TargonClient("k").deploy({"name": "x"}) + assert create.call_count == 1 + assert deploy.call_count == 0 + + +@respx.mock +async def test_provision_includes_env_and_ports() -> None: + create = respx.post(f"{BASE}/workloads").mock( + return_value=httpx.Response(200, json={"uid": "wl-1"}) + ) + respx.post(f"{BASE}/workloads/wl-1/deploy").mock( + return_value=httpx.Response(200, json={"uid": "wl-1"}) + ) + spec = _spec(env={"ROLE": "worker"}, ports=(22, 8080)) + await TargonClient("k").provision(spec) + body = json.loads(create.calls.last.request.content) + assert {"name": "ROLE", "value": "worker"} in body["envs"] + assert body["ports"] == [{"port": 22}, {"port": 8080}] + assert body["ssh_public_keys"] == ["ssh-ed25519 AAAA"] + + +# -- VAL-PROV-007 (insufficient credits, typed, no retry) --------------------- + + +@respx.mock +async def test_deploy_402_on_create_raises_insufficient_credits_no_retry() -> None: + create = respx.post(f"{BASE}/workloads").mock( + return_value=httpx.Response(402, json={"error": "payment required"}) + ) + deploy = respx.post(url__regex=rf"{BASE}/workloads/.+/deploy") + with pytest.raises(InsufficientCreditsError) as exc_info: + await TargonClient("k").provision(_spec()) + assert isinstance(exc_info.value, ProviderError) + assert isinstance(exc_info.value, TargonError) + assert create.call_count == 1 + # A credit failure on create never proceeds to (or retries) the deploy step. + assert deploy.call_count == 0 + + +@respx.mock +async def test_deploy_402_on_deploy_step_raises_insufficient_credits_no_retry() -> None: + create = respx.post(f"{BASE}/workloads").mock( + return_value=httpx.Response(200, json={"uid": "wl-1"}) + ) + deploy = respx.post(f"{BASE}/workloads/wl-1/deploy").mock( + return_value=httpx.Response(402, json={"error": "insufficient credits"}) + ) + with pytest.raises(InsufficientCreditsError): + await TargonClient("k").provision(_spec()) + assert create.call_count == 1 + assert deploy.call_count == 1 + + +@respx.mock +async def test_deploy_403_raises_insufficient_credits_no_retry() -> None: + create = respx.post(f"{BASE}/workloads").mock( + return_value=httpx.Response(403, text="forbidden") + ) + with pytest.raises(InsufficientCreditsError): + await TargonClient("k").deploy({"name": "x"}) + assert create.call_count == 1 + + +@respx.mock +async def test_deploy_credit_body_raises_insufficient_credits() -> None: + create = respx.post(f"{BASE}/workloads").mock( + return_value=httpx.Response(400, json={"error": "insufficient credits"}) + ) + with pytest.raises(InsufficientCreditsError): + await TargonClient("k").deploy({"name": "x"}) + assert create.call_count == 1 + + +@respx.mock +async def test_deploy_generic_error_raises_targon_error_not_credits() -> None: + create = respx.post(f"{BASE}/workloads").mock( + return_value=httpx.Response(500, text="boom") + ) + with pytest.raises(TargonError) as exc_info: + await TargonClient("k").deploy({"name": "x"}) + assert not isinstance(exc_info.value, InsufficientCreditsError) + assert create.call_count == 1 + + +# -- VAL-PROV-008 (balance unavailable, typed, zero HTTP) --------------------- + + +@respx.mock +async def test_balance_raises_typed_error_with_no_http() -> None: + with pytest.raises(BalanceUnavailableError) as exc_info: + await TargonClient("k").balance() + assert isinstance(exc_info.value, ProviderError) + assert isinstance(exc_info.value, TargonError) + assert respx.calls.call_count == 0 + + +# -- apps / workloads listing ------------------------------------------------- + + +@respx.mock +async def test_list_apps_returns_items() -> None: + respx.get(f"{BASE}/apps").mock( + return_value=httpx.Response(200, json={"items": [{"uid": "app-1"}]}) + ) + apps = await TargonClient("k").list_apps() + assert apps == [{"uid": "app-1"}] + + +@respx.mock +async def test_create_app_posts_name() -> None: + route = respx.post(f"{BASE}/apps").mock( + return_value=httpx.Response(200, json={"uid": "app-2", "name": "worker"}) + ) + result = await TargonClient("k").create_app("worker") + assert result["uid"] == "app-2" + assert json.loads(route.calls.last.request.content)["name"] == "worker" + + +@respx.mock +async def test_list_workloads_returns_items() -> None: + respx.get(f"{BASE}/workloads").mock( + return_value=httpx.Response(200, json={"items": [{"uid": "wl-1"}]}) + ) + workloads = await TargonClient("k").list_workloads() + assert workloads == [{"uid": "wl-1"}] + + +@respx.mock +async def test_status_parses_workload_detail() -> None: + respx.get(f"{BASE}/workloads/wl-1").mock( + return_value=httpx.Response( + 200, json={"uid": "wl-1", "state": {"status": "RUNNING"}} + ) + ) + instance = await TargonClient("k").status("wl-1") + assert instance.id == "wl-1" + assert instance.status == "RUNNING" + + +@respx.mock +async def test_workload_state_and_events() -> None: + respx.get(f"{BASE}/workloads/wl-1/state").mock( + return_value=httpx.Response(200, json={"uid": "wl-1", "status": "RUNNING"}) + ) + respx.get(f"{BASE}/workloads/wl-1/events").mock( + return_value=httpx.Response(200, json={"items": [{"event_type": "created"}]}) + ) + client = TargonClient("k") + state = await client.workload_state("wl-1") + events = await client.workload_events("wl-1") + assert state["status"] == "RUNNING" + assert events["items"][0]["event_type"] == "created" + + +# -- logs --------------------------------------------------------------------- + + +@respx.mock +async def test_stream_logs_yields_lines() -> None: + respx.get(f"{BASE}/workloads/wl-1/logs").mock( + return_value=httpx.Response(200, text="line-1\nline-2\n") + ) + lines = [line async for line in TargonClient("k").stream_logs("wl-1")] + assert lines == ["line-1", "line-2"] + + +@respx.mock +async def test_stream_logs_raises_on_error() -> None: + respx.get(f"{BASE}/workloads/wl-1/logs").mock(return_value=httpx.Response(404)) + with pytest.raises(TargonError): + async for _ in TargonClient("k").stream_logs("wl-1"): + pass + + +# -- terminate / verify_terminated ------------------------------------------- + + +@respx.mock +async def test_terminate_is_idempotent() -> None: + route = respx.delete(f"{BASE}/workloads/wl-1").mock( + side_effect=[httpx.Response(200), httpx.Response(404)] + ) + client = TargonClient("k") + await client.terminate("wl-1") + await client.terminate("wl-1") + assert route.call_count == 2 + + +@respx.mock +async def test_terminate_raises_on_non_404_error() -> None: + respx.delete(f"{BASE}/workloads/wl-1").mock(return_value=httpx.Response(500)) + with pytest.raises(TargonError): + await TargonClient("k").terminate("wl-1") + + +@respx.mock +async def test_verify_terminated_reflects_absence() -> None: + respx.get(f"{BASE}/workloads/wl-1").mock( + side_effect=[ + httpx.Response(200, json={"uid": "wl-1", "state": {"status": "RUNNING"}}), + httpx.Response(404), + ] + ) + client = TargonClient("k") + assert await client.verify_terminated("wl-1") is False + assert await client.verify_terminated("wl-1") is True + + +@respx.mock +async def test_verify_terminated_true_for_deleted_status() -> None: + respx.get(f"{BASE}/workloads/wl-1").mock( + return_value=httpx.Response( + 200, json={"uid": "wl-1", "state": {"status": "deleted"}} + ) + ) + assert await TargonClient("k").verify_terminated("wl-1") is True + + +@respx.mock +async def test_verify_terminated_raises_on_non_404_error() -> None: + respx.get(f"{BASE}/workloads/wl-1").mock(return_value=httpx.Response(500)) + with pytest.raises(TargonError): + await TargonClient("k").verify_terminated("wl-1") + + +@respx.mock +async def test_create_app_sends_project_id() -> None: + route = respx.post(f"{BASE}/apps").mock( + return_value=httpx.Response(200, json={"uid": "app-3"}) + ) + await TargonClient("k").create_app("worker", project_id="proj-1") + assert json.loads(route.calls.last.request.content)["project_id"] == "proj-1" + + +@respx.mock +async def test_workload_state_and_create_app_tolerate_non_dict() -> None: + respx.get(f"{BASE}/workloads/wl-1/state").mock( + return_value=httpx.Response(200, json=["unexpected"]) + ) + respx.post(f"{BASE}/apps").mock( + return_value=httpx.Response(200, json=["unexpected"]) + ) + client = TargonClient("k") + assert await client.workload_state("wl-1") == {} + assert await client.create_app("worker") == {} + + +# -- transport / parsing edge paths ------------------------------------------ + + +@respx.mock +async def test_request_error_raises_targon_error() -> None: + respx.get(f"{BASE}/workloads").mock(return_value=httpx.Response(500)) + with pytest.raises(TargonError) as exc_info: + await TargonClient("k").list_workloads() + assert exc_info.value.status_code == 500 + + +@respx.mock +async def test_transport_error_wrapped_as_targon_error() -> None: + respx.get(f"{BASE}/workloads").mock(side_effect=httpx.ConnectError("down")) + with pytest.raises(TargonError): + await TargonClient("k").list_workloads() + + +@respx.mock +async def test_api_key_sent_only_in_authorization_header() -> None: + route = respx.get(f"{BASE}/workloads").mock( + return_value=httpx.Response(200, json={"items": []}) + ) + await TargonClient("MY-SECRET").list_workloads() + assert route.calls.last.request.headers["Authorization"] == "Bearer MY-SECRET" + + +def test_is_insufficient_credits_classifies_codes_and_bodies() -> None: + assert _is_insufficient_credits(402, "") is True + assert _is_insufficient_credits(403, "") is True + assert _is_insufficient_credits(400, "insufficient credits") is True + assert _is_insufficient_credits(400, "no credit left") is True + assert _is_insufficient_credits(500, "boom") is False + assert _is_insufficient_credits(401, "unauthorized") is False + + +def test_parse_inventory_offer_skips_zero_and_missing() -> None: + zero = {"name": "x", "available": 0, "cost_per_hour": 1} + assert _parse_inventory_offer(zero) is None + assert _parse_inventory_offer({"available": 1, "cost_per_hour": 1}) is None + assert _parse_inventory_offer({"name": "x", "available": 1}) is None + offer = _parse_inventory_offer( + {"name": "x", "available": 2, "cost_per_hour": "1.5"} + ) + assert offer is not None + assert offer.price_per_hour == pytest.approx(1.5) + + +def test_parse_inventory_offer_normalizes_price_per_gpu() -> None: + multi = _parse_inventory_offer( + { + "name": "h100x8", + "spec": {"gpu_count": 8}, + "cost_per_hour": 20.0, + "available": 3, + } + ) + assert multi is not None + assert multi.price_per_hour == pytest.approx(2.5) + assert multi.gpu_count == 8 + # Missing gpu_count falls back to the whole-shape cost (no ZeroDivisionError). + no_count = _parse_inventory_offer( + {"name": "x", "cost_per_hour": 1.5, "available": 1} + ) + assert no_count is not None + assert no_count.price_per_hour == pytest.approx(1.5) + assert no_count.gpu_count == 0 + + +def test_parse_workload_instance_rejects_non_mapping() -> None: + with pytest.raises(TargonError): + _parse_workload_instance(["not", "a", "mapping"]) + + +def test_as_list_handles_shapes() -> None: + assert _as_list("nope", "items") == [] + assert _as_list({"items": "nope"}, "items") == [] + assert _as_list([{"a": 1}, "skip"], "items") == [{"a": 1}] + assert _as_list({"items": [{"a": 1}]}, "items") == [{"a": 1}] diff --git a/tests/unit/test_compute_worker_deployment.py b/tests/unit/test_compute_worker_deployment.py new file mode 100644 index 00000000..533d47bb --- /dev/null +++ b/tests/unit/test_compute_worker_deployment.py @@ -0,0 +1,326 @@ +"""Offline tests for the worker deployment definitions (VAL-PROV-009/010/016). + +These pin the well-formedness of the declarative deploy definitions the miner's +worker image ships as: + +* the Lium ``CustomTemplateRequest`` payload (VAL-PROV-009), and +* the Targon app definition (VAL-PROV-010), + +each pinning the docker image BY DIGEST, plus an explicit assertion that the whole +compute package makes NO real network calls under respx strict mode and needs no +provider credentials (VAL-PROV-016). +""" + +from __future__ import annotations + +import json +import re + +import httpx +import pytest +import respx + +from base.compute import ( + LiumClient, + TargonClient, + build_lium_worker_template, + build_targon_worker_app, + is_metachar_free, + pinned_image_reference, +) +from base.compute.worker_deployment import ( + WORKER_IMAGE, + WORKER_IMAGE_DIGEST, + WORKER_INTERNAL_PORTS, + WORKER_STARTUP_COMMANDS, + is_loopback_url, + is_pinned_digest, +) + +_DIGEST_RE = re.compile(r"^sha256:[0-9a-f]{64}$") + +# Substrings that would indicate a leaked credential or secret in a definition. +_SECRET_MARKERS = ( + "X-API-Key", + "Authorization", + "Bearer", + "LIUM_API_KEY", + "TARGON_API_KEY", + "SENTINEL", + "password", + "secret", +) + + +def _assert_no_secrets(definition: object) -> None: + blob = json.dumps(definition) + for marker in _SECRET_MARKERS: + assert marker not in blob + + +# -- shared placeholder digest ------------------------------------------------ + + +def test_placeholder_image_pinned_by_digest() -> None: + assert WORKER_IMAGE == "ghcr.io/baseintelligence/prism-evaluator" + assert _DIGEST_RE.match(WORKER_IMAGE_DIGEST) + assert is_pinned_digest(WORKER_IMAGE_DIGEST) is True + assert is_pinned_digest("sha256:short") is False + assert is_pinned_digest("latest") is False + + +def test_pinned_image_reference_is_fully_qualified_and_immutable() -> None: + ref = pinned_image_reference(WORKER_IMAGE, WORKER_IMAGE_DIGEST, tag="latest") + assert ref == f"{WORKER_IMAGE}:latest@{WORKER_IMAGE_DIGEST}" + assert f"@{WORKER_IMAGE_DIGEST}" in ref + # Tag is optional; the digest pin is what makes the reference immutable. + ref_no_tag = pinned_image_reference(WORKER_IMAGE, WORKER_IMAGE_DIGEST) + assert ref_no_tag == f"{WORKER_IMAGE}@{WORKER_IMAGE_DIGEST}" + + +def test_pinned_image_reference_rejects_unpinned_digest() -> None: + with pytest.raises(ValueError): + pinned_image_reference(WORKER_IMAGE, "latest") + + +# -- VAL-PROV-009 : Lium worker template -------------------------------------- + + +def test_lium_template_is_well_formed_and_pins_digest() -> None: + template = build_lium_worker_template() + assert template["name"] + assert template["docker_image"] == WORKER_IMAGE + assert _DIGEST_RE.match(template["docker_image_digest"]) + assert 22 in template["internal_ports"] + assert template["is_private"] is True + _assert_no_secrets(template) + + +def test_lium_template_digest_is_sha256_shaped() -> None: + template = build_lium_worker_template() + assert _DIGEST_RE.match(template["docker_image_digest"]) is not None + + +def test_lium_template_plumbs_environment() -> None: + template = build_lium_worker_template( + environment={"BASE_MASTER_URL": "http://master:8000", "WORKER_ROLE": "gpu"} + ) + assert template["environment"]["BASE_MASTER_URL"] == "http://master:8000" + assert template["environment"]["WORKER_ROLE"] == "gpu" + + +def test_lium_template_default_environment_is_empty_dict() -> None: + template = build_lium_worker_template() + assert template["environment"] == {} + + +def test_lium_template_accepts_swapped_image_and_digest() -> None: + digest = "sha256:" + "a" * 64 + template = build_lium_worker_template( + image="ghcr.io/baseintelligence/base-worker", + image_digest=digest, + image_tag="v2", + ) + assert template["docker_image"] == "ghcr.io/baseintelligence/base-worker" + assert template["docker_image_digest"] == digest + assert template["docker_image_tag"] == "v2" + + +def test_lium_template_custom_ports_still_include_ssh() -> None: + template = build_lium_worker_template(internal_ports=(22, 8082)) + assert template["internal_ports"] == [22, 8082] + + +def test_lium_template_rejects_ports_without_ssh() -> None: + with pytest.raises(ValueError): + build_lium_worker_template(internal_ports=(8082,)) + + +def test_lium_template_rejects_malformed_digest() -> None: + with pytest.raises(ValueError): + build_lium_worker_template(image_digest="not-a-digest") + + +# -- startup_commands metachar-free constraint (live-learned Lium guard) ------- + + +def test_lium_template_default_startup_commands_is_metachar_free() -> None: + template = build_lium_worker_template() + assert template["startup_commands"] == WORKER_STARTUP_COMMANDS + assert template["startup_commands"] == "tail -f /dev/null" + assert is_metachar_free(template["startup_commands"]) is True + + +def test_lium_template_accepts_custom_metachar_free_startup_commands() -> None: + template = build_lium_worker_template(startup_commands="sleep infinity") + assert template["startup_commands"] == "sleep infinity" + + +@pytest.mark.parametrize( + "command", + [ + "service ssh start && tail -f /dev/null", + "a; b", + "a | b", + "a || b", + "echo `id`", + "echo $(whoami)", + ], +) +def test_lium_template_rejects_startup_commands_with_metacharacters( + command: str, +) -> None: + assert is_metachar_free(command) is False + with pytest.raises(ValueError): + build_lium_worker_template(startup_commands=command) + + +# -- VAL-PROV-010 : Targon app definition ------------------------------------- + + +def test_targon_app_is_well_formed_and_references_pinned_image() -> None: + app = build_targon_worker_app() + assert app["name"] + # Fully qualified image reference with an immutable digest pin. + assert app["image"].startswith(f"{WORKER_IMAGE}:") + assert f"@{WORKER_IMAGE_DIGEST}" in app["image"] + assert _DIGEST_RE.match(app["image_digest"]) is not None + _assert_no_secrets(app) + + +def test_targon_app_declares_gpu_resource_shape() -> None: + app = build_targon_worker_app() + assert app["resource"] + assert app["gpu_type"] + assert isinstance(app["gpu_count"], int) + assert app["gpu_count"] >= 1 + + +def test_targon_app_default_gpu_shape_is_live_valid_suffixed_id() -> None: + # Real Targon inventory ids carry a size suffix (h100-small, b200-large per + # library/targon-api.md); a bare 'h100' would be rejected by a live deploy. + from base.compute.worker_deployment import WORKER_GPU_SHAPE + + assert WORKER_GPU_SHAPE == "h100-small" + assert build_targon_worker_app()["resource"] == "h100-small" + + +def test_targon_app_plumbs_environment_as_name_value_pairs() -> None: + app = build_targon_worker_app(environment={"BASE_MASTER_URL": "http://master:8000"}) + assert {"name": "BASE_MASTER_URL", "value": "http://master:8000"} in app["envs"] + + +def test_targon_app_default_environment_is_empty_list() -> None: + app = build_targon_worker_app() + assert app["envs"] == [] + + +def test_targon_app_ports_include_ssh() -> None: + app = build_targon_worker_app() + assert 22 in app["ports"] + + +def test_targon_app_accepts_swapped_image_and_digest() -> None: + digest = "sha256:" + "b" * 64 + app = build_targon_worker_app( + image="ghcr.io/baseintelligence/base-worker", + image_digest=digest, + image_tag="v2", + gpu_shape="h200", + gpu_type="H200", + gpu_count=2, + ) + assert app["image"] == f"ghcr.io/baseintelligence/base-worker:v2@{digest}" + assert app["image_digest"] == digest + assert app["resource"] == "h200" + assert app["gpu_type"] == "H200" + assert app["gpu_count"] == 2 + + +def test_targon_app_rejects_malformed_digest() -> None: + with pytest.raises(ValueError): + build_targon_worker_app(image_digest="nope") + + +def test_targon_app_rejects_non_positive_gpu_count() -> None: + with pytest.raises(ValueError): + build_targon_worker_app(gpu_count=0) + + +# -- VAL-PROV-016 : offline, no credentials, respx strict mode ---------------- + + +@respx.mock(assert_all_mocked=True) +async def test_respx_strict_mode_blocks_unmocked_request() -> None: + # No route is registered: a real call would egress to lium.io. Under respx + # strict mode the request is refused BEFORE leaving the process, proving the + # compute suite performs zero real network I/O. + with pytest.raises(AssertionError): + await LiumClient("k").balance() + assert respx.calls.call_count == 0 + + +@respx.mock(assert_all_mocked=True) +async def test_targon_respx_strict_mode_blocks_unmocked_request() -> None: + with pytest.raises(AssertionError): + await TargonClient("k").list_workloads() + assert respx.calls.call_count == 0 + + +@respx.mock +async def test_clients_need_no_provider_credentials_env( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.delenv("LIUM_API_KEY", raising=False) + monkeypatch.delenv("TARGON_API_KEY", raising=False) + monkeypatch.delenv("BASE_LIVE_PROVIDER_TESTS", raising=False) + respx.get("https://lium.io/api/users/me").mock( + return_value=httpx.Response(200, json={"balance": 1.0}) + ) + respx.get("https://api.targon.com/tha/v2/workloads").mock( + return_value=httpx.Response(200, json={"items": []}) + ) + assert await LiumClient("dummy").balance() == pytest.approx(1.0) + assert await TargonClient("dummy").list_workloads() == [] + + +def test_default_internal_ports_include_ssh() -> None: + assert 22 in WORKER_INTERNAL_PORTS + + +# -- loopback detection (Lium edge-WAF 403 on loopback URLs) ------------------- + + +@pytest.mark.parametrize( + "value", + [ + "http://127.0.0.1:8081", + "http://localhost:3100", + "https://localhost", + "http://[::1]:8082", + "127.0.0.1:8081", + "localhost", + "http://0.0.0.0:9000", + "http://127.0.0.53", + ], +) +def test_is_loopback_url_detects_loopback(value: str) -> None: + assert is_loopback_url(value) is True + + +@pytest.mark.parametrize( + "value", + [ + "https://master.example.com", + "http://master:3100", + "http://10.0.0.5:8081", + "http://broker.internal:8082", + "worker-binding:pubkey:hotkey:nonce", + "0xdeadbeef", + "gpu", + "true", + "", + ], +) +def test_is_loopback_url_ignores_non_loopback(value: str) -> None: + assert is_loopback_url(value) is False diff --git a/tests/unit/test_image_updater.py b/tests/unit/test_image_updater.py index 19bb4121..14f667c3 100644 --- a/tests/unit/test_image_updater.py +++ b/tests/unit/test_image_updater.py @@ -10,6 +10,7 @@ import json import logging +import threading from collections.abc import Sequence from pathlib import Path @@ -33,6 +34,45 @@ DIGEST_B = "sha256:" + "b" * 64 DIGEST_C = "sha256:" + "c" * 64 IMAGE = "ghcr.io/baseintelligence/base-master:latest" +IMAGE_UPDATER_LOGGER = "base.supervisor.image_updater" + + +class _AttachedHandler: + """Capture records on a named logger, immune to root-logging churn. + + Mirrors the pattern in test_supervisor_weights.py: other tests reconfigure + root logging / raise logger levels (importing ``bittensor`` bumps every + already-created logger to CRITICAL) or disable loggers (alembic env.py), + which breaks ``caplog`` in a full-suite run, so we attach directly to the + target logger and reset its level/disabled flag. + """ + + def __init__(self, logger_name: str) -> None: + self._logger = logging.getLogger(logger_name) + self.messages: list[str] = [] + self._lock = threading.Lock() + self._was_disabled = False + + class _H(logging.Handler): + def __init__(inner) -> None: + super().__init__(level=logging.DEBUG) + + def emit(inner, record: logging.LogRecord) -> None: + with self._lock: + self.messages.append(record.getMessage()) + + self._handler = _H() + + def __enter__(self) -> _AttachedHandler: + self._was_disabled = self._logger.disabled + self._logger.disabled = False + self._logger.setLevel(logging.DEBUG) + self._logger.addHandler(self._handler) + return self + + def __exit__(self, *exc: object) -> None: + self._logger.removeHandler(self._handler) + self._logger.disabled = self._was_disabled class FakeClock: @@ -170,19 +210,17 @@ def test_unpinned_current_image_is_updated() -> None: # --------------------------------------------------------------------------- -def test_resolver_failure_logs_and_skips_update( - caplog: pytest.LogCaptureFixture, -) -> None: +def test_resolver_failure_logs_and_skips_update() -> None: runner = FakeRunner({"base-admin": f"{IMAGE}@{DIGEST_A}"}) def resolver(reference: ImageReference) -> str: raise RuntimeError("registry unreachable") updater = make_updater(runner, resolver) - with caplog.at_level(logging.WARNING): + with _AttachedHandler(IMAGE_UPDATER_LOGGER) as handler: updater.run_once() assert runner.update_calls == [] - assert any("digest resolution failed" in rec.message for rec in caplog.records) + assert any("digest resolution failed" in msg for msg in handler.messages) def test_resolver_failure_does_not_block_other_targets() -> None: @@ -208,9 +246,7 @@ def resolver(reference: ImageReference) -> str: assert runner.update_calls[0][-1] == "base-other" -def test_untagged_image_rejected_without_any_docker_calls( - caplog: pytest.LogCaptureFixture, -) -> None: +def test_untagged_image_rejected_without_any_docker_calls() -> None: runner = FakeRunner() targets = ( ImageUpdateTarget( @@ -218,30 +254,26 @@ def test_untagged_image_rejected_without_any_docker_calls( image="ghcr.io/baseintelligence/base-master", ), ) - with caplog.at_level(logging.ERROR): + with _AttachedHandler(IMAGE_UPDATER_LOGGER) as handler: make_updater(runner, make_resolver(DIGEST_B), targets).run_once() assert runner.calls == [] - assert any("untagged image" in rec.message for rec in caplog.records) + assert any("untagged image" in msg for msg in handler.messages) -def test_non_sha256_resolver_result_rejected( - caplog: pytest.LogCaptureFixture, -) -> None: +def test_non_sha256_resolver_result_rejected() -> None: runner = FakeRunner({"base-admin": f"{IMAGE}@{DIGEST_A}"}) - with caplog.at_level(logging.ERROR): + with _AttachedHandler(IMAGE_UPDATER_LOGGER) as handler: make_updater(runner, make_resolver("md5:deadbeef")).run_once() assert runner.update_calls == [] - assert any("refusing un-pinned update" in rec.message for rec in caplog.records) + assert any("refusing un-pinned update" in msg for msg in handler.messages) -def test_missing_service_skipped_without_update( - caplog: pytest.LogCaptureFixture, -) -> None: +def test_missing_service_skipped_without_update() -> None: runner = FakeRunner({}) - with caplog.at_level(logging.WARNING): + with _AttachedHandler(IMAGE_UPDATER_LOGGER) as handler: make_updater(runner, make_resolver(DIGEST_B)).run_once() assert runner.update_calls == [] - assert any("cannot inspect service" in rec.message for rec in caplog.records) + assert any("cannot inspect service" in msg for msg in handler.messages) # --------------------------------------------------------------------------- diff --git a/tests/unit/test_no_external_egress.py b/tests/unit/test_no_external_egress.py new file mode 100644 index 00000000..e6092a0e --- /dev/null +++ b/tests/unit/test_no_external_egress.py @@ -0,0 +1,85 @@ +"""Unit tests for the offline network-egress guard (VAL-CROSS-006 part a). + +The guard (``scripts/mission/no_external_egress.py``) is the reusable plugin the +flags-OFF legacy regression uses to prove BOTH repos' default suites run with ZERO real +egress to lium.io / api.targon.com: it blocks any DNS/connect to a non-loopback host +while leaving loopback (the local test PostgreSQL, in-process stubs) and AF_UNIX sockets +working. +""" + +from __future__ import annotations + +import socket +import sys +from pathlib import Path + +import pytest + +sys.path.insert(0, str(Path(__file__).resolve().parents[2] / "scripts" / "mission")) + +import no_external_egress as guard # noqa: E402 + + +@pytest.fixture(autouse=True) +def _restore_socket(): + guard.uninstall() + yield + guard.uninstall() + + +def test_error_is_oserror_subclass_so_probes_skip_not_error() -> None: + # A "is the network up?" probe catches OSError and skips; the guard must not surface + # as a RuntimeError that would error test collection instead. + assert issubclass(guard.ExternalEgressBlocked, OSError) + + +def test_blocks_external_dns_resolution() -> None: + guard.install() + for host in ("lium.io", "api.targon.com", "huggingface.co", "example.com"): + with pytest.raises(guard.ExternalEgressBlocked): + socket.getaddrinfo(host, 443) + + +def test_allows_loopback_resolution() -> None: + guard.install() + assert socket.getaddrinfo("127.0.0.1", 15433) + assert socket.getaddrinfo("localhost", 80) + assert socket.getaddrinfo("::1", 80) + + +def test_blocks_external_socket_connect() -> None: + guard.install() + sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + try: + with pytest.raises(guard.ExternalEgressBlocked): + sock.connect(("1.1.1.1", 80)) + finally: + sock.close() + + +def test_afunix_and_loopback_addresses_are_not_blocked() -> None: + guard.install() + # AF_UNIX addresses are filesystem paths (multiprocessing's /tmp/pymp-* sockets), + # never external, so they must pass the address guard untouched. + guard._guard_address("/tmp/pymp-abc123/listener-0") + guard._guard_address(("127.0.0.1", 15433)) + guard._guard_address(("::1", 8080, 0, 0)) + + +def test_httpx_call_to_provider_host_is_blocked() -> None: + httpx = pytest.importorskip("httpx") + guard.install() + # httpx maps the underlying OSError from the transport to its own ConnectError, but + # the guard's message is preserved -- the real egress never leaves the host. + with pytest.raises(httpx.ConnectError) as excinfo: + httpx.get("https://lium.io/api/users/me", timeout=5) + assert "blocked real network egress" in str(excinfo.value) + + +def test_install_is_idempotent_and_uninstall_restores() -> None: + original = socket.getaddrinfo + guard.install() + guard.install() + assert socket.getaddrinfo is not original + guard.uninstall() + assert socket.getaddrinfo is original diff --git a/tests/unit/test_orchestration.py b/tests/unit/test_orchestration.py index c02c58b3..f0816973 100644 --- a/tests/unit/test_orchestration.py +++ b/tests/unit/test_orchestration.py @@ -221,6 +221,10 @@ async def test_run_once_bridges_and_assigns_to_online_validators() -> None: # Two tasks split across two online validators (balanced). assert {r.assigned_validator_hotkey for r in rows} == {"v1", "v2"} assert result.folded == [] + # Flag OFF (no worker engine/reconciler wired): the worker plane is inert + # and the pass folds/routes exactly as it did before the worker plane. + assert result.worker is None + assert result.reconciliation is None finally: await engine.dispose() diff --git a/tests/unit/test_reassignment.py b/tests/unit/test_reassignment.py index f1474414..51198ecc 100644 --- a/tests/unit/test_reassignment.py +++ b/tests/unit/test_reassignment.py @@ -28,6 +28,9 @@ ) from base.db.models import WorkAssignment, WorkAssignmentStatus from base.master.assignment import ( + CAPABILITY_GPU, + EXECUTOR_KIND_PAYLOAD_KEY, + EXECUTOR_KIND_VALIDATOR, RESUME_CHECKPOINT_PAYLOAD_KEY, AssignmentService, ) @@ -91,6 +94,41 @@ async def _one(factory) -> WorkAssignment: return rows[0] +async def _add_gpu_primary( + factory, + work_unit_id: str, + *, + submission_ref: str, + assigned_validator_hotkey: str | None = None, + payload: dict | None = None, + status: WorkAssignmentStatus = WorkAssignmentStatus.ASSIGNED, + attempt_count: int = 1, + max_attempts: int = 3, +) -> None: + """Insert a gpu primary ``work_assignments`` row. + + A worker-owned prism primary is ``ASSIGNED`` with a NULL validator hotkey by + design (the worker replica plane owns it via ``worker_assignments`` rows). + """ + + async with session_scope(factory) as session: + session.add( + WorkAssignment( + challenge_slug="prism", + work_unit_id=work_unit_id, + submission_ref=submission_ref, + payload=payload or {}, + required_capability=CAPABILITY_GPU, + assigned_validator_hotkey=assigned_validator_hotkey, + status=status, + attempt_count=attempt_count, + max_attempts=max_attempts, + created_at=NOW, + updated_at=NOW, + ) + ) + + # VAL-ASSIGN-024 async def test_offline_validator_work_reverts_and_reassigns() -> None: engine, factory = await _setup() @@ -419,3 +457,79 @@ async def _boom(**_: object) -> dict[str, str]: assert row.assigned_validator_hotkey == "v1" finally: await engine.dispose() + + +# Reclaim/assign guard symmetry: with the worker plane ON a worker-owned prism +# primary (gpu, NULL validator hotkey by design) must NOT be reclaimed as a +# stale-validator unit -- the null hotkey means "owned by the worker replica +# plane", not "offline validator => reassignable". +async def test_reclaim_skips_worker_owned_primary_when_worker_plane_on() -> None: + engine, factory = await _setup() + try: + service = AssignmentService( + factory, + now_fn=lambda: NOW, + worker_plane_capabilities=frozenset({CAPABILITY_GPU}), + ) + await _add_gpu_primary(factory, "psub", submission_ref="miner-H") + + outcome = await service.reclaim_stale_assignments() + + assert outcome.reverted == [] + assert outcome.failed == [] + row = await _one(factory) + assert row.status == WorkAssignmentStatus.ASSIGNED + assert row.assigned_validator_hotkey is None + assert row.attempt_count == 1 # untouched -> no churn + finally: + await engine.dispose() + + +# Flag OFF: reclaim is byte-identical to legacy -- a null-hotkey gpu primary IS +# reverted to pending (legacy reads a null hotkey as "offline => reassignable"). +async def test_reclaim_reverts_null_hotkey_gpu_primary_when_flag_off() -> None: + engine, factory = await _setup() + try: + service = AssignmentService(factory, now_fn=lambda: NOW) + await _add_gpu_primary(factory, "psub", submission_ref="miner-H") + + outcome = await service.reclaim_stale_assignments() + + assert outcome.reverted == ["psub"] + assert outcome.failed == [] + row = await _one(factory) + assert row.status == WorkAssignmentStatus.PENDING + assert row.assigned_validator_hotkey is None + finally: + await engine.dispose() + + +# Flag ON does not globally disable reclaim: a validator-owned gpu AUDIT unit +# (executor_kind=validator) whose validator is offline IS still reclaimed, since +# the guard only spares worker-owned units (executor_kind != validator). +async def test_reclaim_still_reclaims_stale_validator_gpu_audit_unit_when_flag_on() -> ( + None +): + engine, factory = await _setup() + try: + service = AssignmentService( + factory, + now_fn=lambda: NOW, + worker_plane_capabilities=frozenset({CAPABILITY_GPU}), + ) + await _add_gpu_primary( + factory, + "U:audit", + submission_ref="miner-H", + assigned_validator_hotkey="gone-val", + payload={EXECUTOR_KIND_PAYLOAD_KEY: EXECUTOR_KIND_VALIDATOR}, + ) + + outcome = await service.reclaim_stale_assignments() + + assert outcome.reverted == ["U:audit"] + row = await _one(factory) + assert row.status == WorkAssignmentStatus.PENDING + assert row.assigned_validator_hotkey is None + finally: + await engine.dispose() diff --git a/tests/unit/test_worker_agent.py b/tests/unit/test_worker_agent.py new file mode 100644 index 00000000..74d432b4 --- /dev/null +++ b/tests/unit/test_worker_agent.py @@ -0,0 +1,631 @@ +"""WorkerAgent runtime behavior against the real master proxy app. + +Drives the agent's :class:`WorkerCoordinationClient` + :class:`WorkerAgent` +through the in-process master (httpx ASGITransport) wired with the live worker +coordination + worker assignment services, using REAL sr25519 keypairs and the +production signature verifier. Covers the feature's expected behaviors: + +* register under a miner-signed binding -> visible via the master (VAL-AGENT-002) +* forged binding rejected with a clear error (VAL-AGENT-003) +* replayed nonce rejected; fresh nonce re-enrolls (VAL-AGENT-004) +* heartbeats keep active + advance last-seen (VAL-AGENT-005); missing heartbeats + go stale after the TTL (VAL-AGENT-006) +* pull only gpu units (VAL-AGENT-007); posted results carry the ExecutionProof + (VAL-AGENT-008) +* unreachable broker -> clean failed unit, agent stays alive (VAL-AGENT-016) +* restart re-registers idempotently -> one fleet entry (VAL-AGENT-017) +""" + +from __future__ import annotations + +from collections.abc import AsyncIterator +from datetime import UTC, datetime +from typing import Any + +import pytest +from httpx import ASGITransport +from sqlalchemy import func, select + +from base.bittensor.metagraph_cache import MetagraphCache +from base.db import ( + Base, + WorkAssignmentStatus, + WorkerAssignment, + WorkerRegistration, + create_engine, + create_session_factory, +) +from base.master.app_proxy import create_proxy_app +from base.master.worker_assignment import WorkerAssignmentService +from base.master.worker_coordination import WorkerCoordinationService +from base.security.worker_auth import ( + CoordinationReadEligibility, + MetagraphMinerMembership, + RegisteredWorkerEligibility, + SqlAlchemyWorkerNonceStore, + WorkerSignedRequestVerifier, + worker_binding_message, +) +from base.validator.agent.executor import ( + AssignmentContext, + BrokerConfig, + ExecutionResult, + ProgressCallback, +) +from base.validator.agent.signing import KeypairRequestSigner +from base.worker import ( + StubManifestExecutor, + WorkerAgent, + WorkerBinding, + WorkerCoordinationClient, + WorkerCoordinationClientError, + WorkerProofExecutor, + WorkerProvenance, + verify_execution_proof, +) +from base.worker.runtime import BackoffPolicy + +NOW_EPOCH = 1_750_000_000.0 +TTL_SECONDS = 120 +FAST_BACKOFF = BackoffPolicy(initial_seconds=0.0, max_seconds=0.0, multiplier=2.0) + + +class _LazyKeypair: + """Defer sr25519 keypair creation until first real use. + + ``bittensor``'s import reconfigures the stdlib logging levels of every + already-created logger, which breaks unrelated ``caplog`` suites. Deferring + the import to test-execution time avoids that; the dunder guard is essential + because pytest's collection probes module globals via ``getattr(obj, + "__test__")`` and would otherwise trigger the import at collection time. + """ + + def __init__(self, uri: str) -> None: + self._uri = uri + self._kp: Any = None + + def _resolve(self) -> Any: + if self._kp is None: + import bittensor as bt + + self._kp = bt.Keypair.create_from_uri(self._uri) + return self._kp + + def __getattr__(self, name: str) -> Any: + if name.startswith("__") and name.endswith("__"): + raise AttributeError(name) + return getattr(self._resolve(), name) + + +MINER = _LazyKeypair("//AgentMiner1") +MINER2 = _LazyKeypair("//AgentMiner2") +WORKER = _LazyKeypair("//AgentWorker1") +BROKER = BrokerConfig(broker_url="http://127.0.0.1:65533", broker_token="t") + + +class FakeClock: + def __init__(self, epoch: float) -> None: + self.epoch = float(epoch) + + def time(self) -> float: + return self.epoch + + def now(self) -> datetime: + return datetime.fromtimestamp(self.epoch, UTC) + + +class FakeNonceStore: + async def reserve(self, **_: Any) -> None: + return None + + +class FakeCache: + def get(self) -> dict[str, int]: + return {} + + +class ConnectionRefusedExecutor: + """Executor that fails as an unreachable local broker would.""" + + async def execute( + self, context: AssignmentContext, *, progress: ProgressCallback + ) -> ExecutionResult: + raise ConnectionError("broker connection refused") + + +def _binding(miner: Any, *, worker_pubkey: str, nonce: str) -> WorkerBinding: + message = worker_binding_message( + worker_pubkey=worker_pubkey, miner_hotkey=miner.ss58_address, nonce=nonce + ) + return WorkerBinding( + miner_hotkey=miner.ss58_address, + signature="0x" + bytes(miner.sign(message)).hex(), + nonce=nonce, + ) + + +class Harness: + def __init__( + self, + transport: ASGITransport, + session_factory: Any, + clock: FakeClock, + service: WorkerCoordinationService, + assignment_service: WorkerAssignmentService, + ) -> None: + self.transport = transport + self.session_factory = session_factory + self.clock = clock + self.service = service + self.assignment_service = assignment_service + + def make_agent( + self, + *, + keypair: Any = WORKER, + binding: WorkerBinding, + executor: Any | None = None, + ) -> WorkerAgent: + signer = KeypairRequestSigner(keypair) + client = WorkerCoordinationClient( + "http://testserver", + signer, + transport=self.transport, + now_fn=self.clock.time, + ) + proof_executor = executor or WorkerProofExecutor( + StubManifestExecutor(), + signer=signer, + provenance=WorkerProvenance( + provider_name="local", miner_hotkey=binding.miner_hotkey + ), + ) + return WorkerAgent( + client=client, + executor=proof_executor, + broker=BROKER, + binding=binding, + provider="local", + provider_instance_ref="local-1", + capabilities=["gpu"], + backoff=FAST_BACKOFF, + ) + + async def fleet(self) -> list[dict[str, Any]]: + async with self.session_factory() as session: + rows = (await session.execute(select(WorkerRegistration))).scalars().all() + now = self.clock.now() + return [ + { + "worker_pubkey": r.worker_pubkey, + "miner_hotkey": r.miner_hotkey, + "status": self.service.effective_status(r, now).value, + "last_heartbeat_at": r.last_heartbeat_at, + } + for r in rows + ] + + async def count_pubkey(self, worker_pubkey: str) -> int: + async with self.session_factory() as session: + return await session.scalar( + select(func.count(WorkerRegistration.id)).where( + WorkerRegistration.worker_pubkey == worker_pubkey + ) + ) + + async def assignment_row(self, work_unit_id: str) -> WorkerAssignment | None: + async with self.session_factory() as session: + return ( + await session.execute( + select(WorkerAssignment).where( + WorkerAssignment.work_unit_id == work_unit_id + ) + ) + ).scalar_one_or_none() + + +async def _build_harness() -> tuple[Harness, Any]: + engine = create_engine("sqlite+aiosqlite:///:memory:") + async with engine.begin() as connection: + await connection.run_sync(Base.metadata.create_all) + session_factory = create_session_factory(engine) + + cache = MetagraphCache(netuid=1, ttl_seconds=300) + cache.update_from_metagraph( + [MINER.ss58_address, MINER2.ss58_address], + validator_permits=[False, False], + stakes=[0.0, 0.0], + ) + clock = FakeClock(NOW_EPOCH) + service = WorkerCoordinationService( + session_factory, + miner_membership=MetagraphMinerMembership(cache), + binding_nonce_store=SqlAlchemyWorkerNonceStore(session_factory), + heartbeat_ttl_seconds=TTL_SECONDS, + now_fn=clock.now, + ) + verifier = WorkerSignedRequestVerifier( + nonce_store=SqlAlchemyWorkerNonceStore(session_factory), + eligibility=CoordinationReadEligibility(session_factory, cache), + ttl_seconds=300, + now_fn=clock.time, + ) + assignment_service = WorkerAssignmentService( + session_factory, worker_service=service, now_fn=clock.now + ) + assignment_verifier = WorkerSignedRequestVerifier( + nonce_store=SqlAlchemyWorkerNonceStore(session_factory), + eligibility=RegisteredWorkerEligibility(session_factory), + ttl_seconds=300, + now_fn=clock.time, + ) + app = create_proxy_app( + registry=object(), + nonce_store=FakeNonceStore(), + metagraph_cache=FakeCache(), # type: ignore[arg-type] + worker_service=service, + worker_verifier=verifier, + worker_assignment_service=assignment_service, + worker_assignment_verifier=assignment_verifier, + ) + transport = ASGITransport(app=app) + return ( + Harness(transport, session_factory, clock, service, assignment_service), + engine, + ) + + +@pytest.fixture +async def harness() -> AsyncIterator[Harness]: + h, engine = await _build_harness() + try: + yield h + finally: + await engine.dispose() + + +# VAL-AGENT-002 +async def test_register_makes_worker_visible(harness: Harness) -> None: + agent = harness.make_agent( + binding=_binding(MINER, worker_pubkey=WORKER.ss58_address, nonce="n1") + ) + worker_id = await agent.register() + assert worker_id + fleet = await harness.fleet() + assert len(fleet) == 1 + assert fleet[0]["worker_pubkey"] == WORKER.ss58_address + assert fleet[0]["miner_hotkey"] == MINER.ss58_address + + +# VAL-AGENT-003 +async def test_forged_binding_rejected(harness: Harness) -> None: + # Sign the binding with the WRONG key (worker signs, claiming to be miner). + forged = _binding(WORKER, worker_pubkey=WORKER.ss58_address, nonce="n1") + forged = WorkerBinding( + miner_hotkey=MINER.ss58_address, + signature=forged.signature, + nonce="n1", + ) + agent = harness.make_agent(binding=forged) + with pytest.raises(WorkerCoordinationClientError) as exc: + await agent.register() + assert exc.value.status_code == 401 + assert await harness.count_pubkey(WORKER.ss58_address) == 0 + + +# VAL-AGENT-004 +async def test_replayed_nonce_rejected_fresh_nonce_reenrolls(harness: Harness) -> None: + agent = harness.make_agent( + binding=_binding(MINER, worker_pubkey=WORKER.ss58_address, nonce="dup") + ) + await agent.register() + + replay = harness.make_agent( + binding=_binding(MINER, worker_pubkey=WORKER.ss58_address, nonce="dup") + ) + with pytest.raises(WorkerCoordinationClientError) as exc: + await replay.register() + assert exc.value.status_code == 409 + + fresh = harness.make_agent( + binding=_binding(MINER, worker_pubkey=WORKER.ss58_address, nonce="fresh") + ) + await fresh.register() + assert await harness.count_pubkey(WORKER.ss58_address) == 1 + + +# VAL-AGENT-005 +async def test_heartbeats_keep_active_and_advance_last_seen(harness: Harness) -> None: + agent = harness.make_agent( + binding=_binding(MINER, worker_pubkey=WORKER.ss58_address, nonce="n1") + ) + await agent.register() + + harness.clock.epoch = NOW_EPOCH + 10 + await agent.heartbeat_once() + first = (await harness.fleet())[0] + assert first["status"] == "active" + + harness.clock.epoch = NOW_EPOCH + 20 + await agent.heartbeat_once() + second = (await harness.fleet())[0] + assert second["status"] == "active" + assert second["last_heartbeat_at"] > first["last_heartbeat_at"] + + +# VAL-AGENT-006 +async def test_missing_heartbeats_go_stale_after_ttl(harness: Harness) -> None: + agent = harness.make_agent( + binding=_binding(MINER, worker_pubkey=WORKER.ss58_address, nonce="n1") + ) + await agent.register() + await agent.heartbeat_once() + assert (await harness.fleet())[0]["status"] == "active" + + harness.clock.epoch = NOW_EPOCH + TTL_SECONDS + 5 + assert (await harness.fleet())[0]["status"] == "stale" + + +# VAL-AGENT-007 + VAL-AGENT-008 +async def test_pull_gpu_only_and_post_carries_proof(harness: Harness) -> None: + agent = harness.make_agent( + binding=_binding(MINER, worker_pubkey=WORKER.ss58_address, nonce="n1") + ) + worker_id = await agent.register() + await agent.heartbeat_once() + + await harness.assignment_service.create_worker_assignment( + work_unit_id="gpu-unit", + challenge_slug="prism", + submission_ref="ref-gpu", + worker_id=worker_id, + worker_pubkey=WORKER.ss58_address, + miner_hotkey=MINER.ss58_address, + required_capability="gpu", + ) + await harness.assignment_service.create_worker_assignment( + work_unit_id="cpu-unit", + challenge_slug="prism", + submission_ref="ref-cpu", + worker_id=worker_id, + worker_pubkey=WORKER.ss58_address, + miner_hotkey=MINER.ss58_address, + required_capability="cpu", + ) + + summary = await agent.process_pending_assignments() + assert summary.pulled == 1 + assert summary.completed == 1 + + gpu_row = await harness.assignment_row("gpu-unit") + assert gpu_row is not None + assert WorkAssignmentStatus(gpu_row.status) == WorkAssignmentStatus.COMPLETED + assert gpu_row.result_success is True + from base.schemas.worker import ExecutionProof + + assert gpu_row.result_payload is not None + proof = ExecutionProof.model_validate(gpu_row.result_payload["execution_proof"]) + assert proof.version == 1 + assert proof.provider is not None + assert proof.provider.name == "local" + assert proof.provider.miner_hotkey == MINER.ss58_address + assert proof.worker_signature.worker_pubkey == WORKER.ss58_address + assert verify_execution_proof(proof, unit_id="gpu-unit") is True + + cpu_row = await harness.assignment_row("cpu-unit") + assert cpu_row is not None + assert WorkAssignmentStatus(cpu_row.status) == WorkAssignmentStatus.ASSIGNED + assert cpu_row.result_success is None + + +# VAL-AGENT-016 +async def test_unreachable_broker_fails_cleanly_agent_survives( + harness: Harness, +) -> None: + agent = harness.make_agent( + binding=_binding(MINER, worker_pubkey=WORKER.ss58_address, nonce="n1"), + executor=ConnectionRefusedExecutor(), + ) + worker_id = await agent.register() + await agent.heartbeat_once() + await harness.assignment_service.create_worker_assignment( + work_unit_id="gpu-unit", + challenge_slug="prism", + submission_ref="ref-gpu", + worker_id=worker_id, + worker_pubkey=WORKER.ss58_address, + miner_hotkey=MINER.ss58_address, + ) + + summary = await agent.process_pending_assignments() + assert summary.pulled == 1 + assert summary.failed == 1 + + row = await harness.assignment_row("gpu-unit") + assert row is not None + assert WorkAssignmentStatus(row.status) == WorkAssignmentStatus.FAILED + assert row.result_success is False + assert "broker" in (row.result_payload or {}).get("error", "") + + # The agent did not crash: it keeps heartbeating and pulling. + harness.clock.epoch = NOW_EPOCH + 10 + await agent.heartbeat_once() + assert (await harness.fleet())[0]["status"] == "active" + followup = await agent.process_pending_assignments() + assert followup.pulled == 0 # the failed unit is terminal, not re-pulled + + +# VAL-AGENT-017 +async def test_restart_reregisters_idempotently_single_entry( + harness: Harness, +) -> None: + agent = harness.make_agent( + binding=_binding(MINER, worker_pubkey=WORKER.ss58_address, nonce="n1") + ) + worker_id = await agent.register() + await agent.heartbeat_once() + await harness.assignment_service.create_worker_assignment( + work_unit_id="gpu-unit", + challenge_slug="prism", + submission_ref="ref-gpu", + worker_id=worker_id, + worker_pubkey=WORKER.ss58_address, + miner_hotkey=MINER.ss58_address, + ) + before = await harness.assignment_row("gpu-unit") + assert before is not None + before_status = WorkAssignmentStatus(before.status) + before_attempts = before.attempt_count + + # Restart: same worker keypair, FRESH binding nonce. + restarted = harness.make_agent( + binding=_binding(MINER, worker_pubkey=WORKER.ss58_address, nonce="n2") + ) + restart_worker_id = await restarted.register() + await restarted.heartbeat_once() + + assert await harness.count_pubkey(WORKER.ss58_address) == 1 + assert restart_worker_id == worker_id + fleet = await harness.fleet() + assert len(fleet) == 1 + assert fleet[0]["miner_hotkey"] == MINER.ss58_address + assert fleet[0]["status"] == "active" + + after = await harness.assignment_row("gpu-unit") + assert after is not None + assert WorkAssignmentStatus(after.status) == before_status + assert after.attempt_count == before_attempts + + +# VAL-AGENT-017 (cross-owner rebind never silently rebinds) +async def test_cross_owner_rebind_rejected(harness: Harness) -> None: + agent = harness.make_agent( + binding=_binding(MINER, worker_pubkey=WORKER.ss58_address, nonce="n1") + ) + await agent.register() + await agent.heartbeat_once() + + rebind = harness.make_agent( + binding=_binding(MINER2, worker_pubkey=WORKER.ss58_address, nonce="n2") + ) + with pytest.raises(WorkerCoordinationClientError) as exc: + await rebind.register() + assert exc.value.status_code == 409 + assert await harness.count_pubkey(WORKER.ss58_address) == 1 + assert (await harness.fleet())[0]["miner_hotkey"] == MINER.ss58_address + + +class _StubWorker: + def __init__(self, worker_id: str) -> None: + self.worker_id = worker_id + + +class _StubRegisterResponse: + def __init__(self, worker_id: str, ttl: int) -> None: + self.worker = _StubWorker(worker_id) + self.heartbeat_ttl_seconds = ttl + + +class _FlakyClient: + """In-memory client to exercise the agent's retry/loop resilience paths.""" + + def __init__(self, *, fail_registers: int = 0) -> None: + self.worker_pubkey = "wp-stub" + self._fail_registers = fail_registers + self.register_calls = 0 + self.heartbeat_calls = 0 + self.on_heartbeat: Any = None + self.raise_heartbeats: set[int] = set() + + async def register(self, **_: Any) -> _StubRegisterResponse: + self.register_calls += 1 + if self.register_calls <= self._fail_registers: + raise WorkerCoordinationClientError("register down", status_code=503) + return _StubRegisterResponse("w-1", 60) + + async def heartbeat(self, **_: Any) -> None: + self.heartbeat_calls += 1 + if self.heartbeat_calls in self.raise_heartbeats: + raise WorkerCoordinationClientError("hb down", status_code=503) + if self.on_heartbeat is not None: + self.on_heartbeat() + + async def pull(self) -> list[Any]: + return [] + + async def post_result(self, *_: Any, **__: Any) -> None: + return None + + +def _stub_agent(client: _FlakyClient) -> WorkerAgent: + return WorkerAgent( + client=client, # type: ignore[arg-type] + executor=StubManifestExecutor(), + broker=BROKER, + binding=WorkerBinding(miner_hotkey="m", signature="0xsig", nonce="n"), + provider="local", + heartbeat_interval_seconds=0, + poll_interval_seconds=0.0, + backoff=FAST_BACKOFF, + ) + + +# VAL-AGENT-016: transient master failures are retried, not fatal. +async def test_register_retries_transient_then_succeeds() -> None: + client = _FlakyClient(fail_registers=2) + agent = _stub_agent(client) + worker_id = await agent.register() + assert worker_id == "w-1" + assert client.register_calls == 3 + + +async def test_register_fails_fast_on_permanent_error() -> None: + class _Permanent(_FlakyClient): + async def register(self, **_: Any): + self.register_calls += 1 + raise WorkerCoordinationClientError("forged", status_code=401) + + client = _Permanent() + agent = _stub_agent(client) + with pytest.raises(WorkerCoordinationClientError) as exc: + await agent.register() + assert exc.value.status_code == 401 + assert client.register_calls == 1 + + +# VAL-AGENT-016: a heartbeat failure never crashes the loop. +async def test_run_heartbeat_loop_survives_failures() -> None: + import asyncio + + client = _FlakyClient() + client.raise_heartbeats = {1} + agent = _stub_agent(client) + await agent.register() + + shutdown = asyncio.Event() + + def _stop_after_two() -> None: + if client.heartbeat_calls >= 2: + shutdown.set() + + client.on_heartbeat = _stop_after_two + await asyncio.wait_for(agent.run_heartbeat_loop(shutdown), timeout=5.0) + assert client.heartbeat_calls >= 2 # kept beating past the first failure + + +async def test_run_forever_registers_then_loops_until_shutdown() -> None: + import asyncio + + client = _FlakyClient() + agent = _stub_agent(client) + shutdown = asyncio.Event() + + def _stop_after_two() -> None: + if client.heartbeat_calls >= 2: + shutdown.set() + + client.on_heartbeat = _stop_after_two + await asyncio.wait_for(agent.run_forever(shutdown), timeout=5.0) + assert client.register_calls == 1 + assert client.heartbeat_calls >= 2 + assert agent.worker_id == "w-1" diff --git a/tests/unit/test_worker_assignment.py b/tests/unit/test_worker_assignment.py new file mode 100644 index 00000000..cff7ac75 --- /dev/null +++ b/tests/unit/test_worker_assignment.py @@ -0,0 +1,553 @@ +"""Worker assignment surface: gpu-only pull, proof-carrying result, worker auth. + +Covers the worker-plane pull/result routes against the real proxy app (mock +metagraph, ASGI transport): + +* VAL-AGENT-007: an active worker pulls only gpu-capability replicas. +* VAL-AGENT-008: a posted result carries and persists the ExecutionProof. +* VAL-AGENT-018: pull/post authenticate as the WORKER identity and gate on + registration/liveness, never on a metagraph validator permit. +""" + +from __future__ import annotations + +import hashlib +from collections.abc import AsyncIterator +from datetime import UTC, datetime +from typing import Any + +import pytest +from httpx import ASGITransport, AsyncClient +from sqlalchemy import select + +from base.bittensor.metagraph_cache import MetagraphCache +from base.db import ( + Base, + WorkAssignmentStatus, + WorkerAssignment, + WorkerStatus, + create_engine, + create_session_factory, + session_scope, +) +from base.master.app_proxy import create_proxy_app +from base.master.worker_assignment import WorkerAssignmentService +from base.master.worker_coordination import WorkerCoordinationService +from base.security.validator_auth import canonical_validator_request +from base.security.worker_auth import ( + CoordinationReadEligibility, + MetagraphMinerMembership, + RegisteredWorkerEligibility, + SqlAlchemyWorkerNonceStore, + WorkerSignedRequestVerifier, + worker_binding_message, +) + +NOW_EPOCH = 1_750_000_000.0 +TTL_SECONDS = 120 + +MINER_H1 = "miner-H1" +MINER_H2 = "miner-H2" +VALIDATOR = "val-permit" +STRANGER = "stranger-key" + + +class FakeClock: + def __init__(self, epoch: float) -> None: + self.epoch = float(epoch) + + def time(self) -> float: + return self.epoch + + def now(self) -> datetime: + return datetime.fromtimestamp(self.epoch, UTC) + + +class FakeNonceStore: + async def reserve(self, **_: Any) -> None: + return None + + +class FakeCache: + def get(self) -> dict[str, int]: + return {} + + +def _sign(hotkey: str, message: bytes) -> str: + return hashlib.sha256(hotkey.encode() + b":" + message).hexdigest() + + +def _fake_verifier(hotkey: str, message: bytes, signature: str) -> bool: + return signature == _sign(hotkey, message) + + +def _signed_headers( + *, + method: str, + path: str, + body: bytes, + hotkey: str, + nonce: str, + timestamp: float, +) -> dict[str, str]: + ts = str(int(timestamp)) + canonical = canonical_validator_request( + method=method, + path=path, + query_string="", + timestamp=ts, + nonce=nonce, + body=body, + ) + return { + "X-Hotkey": hotkey, + "X-Signature": _sign(hotkey, canonical.encode()), + "X-Nonce": nonce, + "X-Timestamp": ts, + "Content-Type": "application/json", + } + + +class Harness: + def __init__( + self, + client: AsyncClient, + session_factory: Any, + clock: FakeClock, + service: WorkerCoordinationService, + assignment_service: WorkerAssignmentService, + ) -> None: + self.client = client + self.session_factory = session_factory + self.clock = clock + self.service = service + self.assignment_service = assignment_service + self._nonce = 0 + + def _next_nonce(self, prefix: str) -> str: + self._nonce += 1 + return f"{prefix}-{self._nonce}" + + async def enroll_active(self, *, worker_pubkey: str, miner_hotkey: str) -> str: + """Register + heartbeat a worker so it is ACTIVE; return its worker_id.""" + + nonce = self._next_nonce("bind") + message = worker_binding_message( + worker_pubkey=worker_pubkey, miner_hotkey=miner_hotkey, nonce=nonce + ) + resp = await self.client.post( + "/v1/workers/register", + json={ + "worker_pubkey": worker_pubkey, + "miner_hotkey": miner_hotkey, + "binding_signature": _sign(miner_hotkey, message), + "nonce": nonce, + "provider": "local", + "provider_instance_ref": "local-1", + }, + ) + assert resp.status_code == 200, resp.text + worker_id = resp.json()["worker"]["worker_id"] + hb = await self.heartbeat(worker_id=worker_id, signer_pubkey=worker_pubkey) + assert hb.status_code == 200, hb.text + return worker_id + + async def register_pending(self, *, worker_pubkey: str, miner_hotkey: str) -> str: + nonce = self._next_nonce("bind") + message = worker_binding_message( + worker_pubkey=worker_pubkey, miner_hotkey=miner_hotkey, nonce=nonce + ) + resp = await self.client.post( + "/v1/workers/register", + json={ + "worker_pubkey": worker_pubkey, + "miner_hotkey": miner_hotkey, + "binding_signature": _sign(miner_hotkey, message), + "nonce": nonce, + "provider": "local", + "provider_instance_ref": "local-1", + }, + ) + assert resp.status_code == 200, resp.text + return resp.json()["worker"]["worker_id"] + + async def heartbeat(self, *, worker_id: str, signer_pubkey: str) -> Any: + body = b"{}" + path = f"/v1/workers/{worker_id}/heartbeat" + headers = _signed_headers( + method="POST", + path=path, + body=body, + hotkey=signer_pubkey, + nonce=self._next_nonce("hb"), + timestamp=self.clock.time(), + ) + return await self.client.post(path, content=body, headers=headers) + + async def pull(self, *, signer_pubkey: str) -> Any: + body = b"{}" + path = "/v1/workers/assignments/pull" + headers = _signed_headers( + method="POST", + path=path, + body=body, + hotkey=signer_pubkey, + nonce=self._next_nonce("pull"), + timestamp=self.clock.time(), + ) + return await self.client.post(path, content=body, headers=headers) + + async def post_result( + self, + *, + assignment_id: str, + signer_pubkey: str, + success: bool, + payload: dict[str, Any], + ) -> Any: + import json + + body = json.dumps( + {"success": success, "payload": payload}, separators=(",", ":") + ).encode() + path = f"/v1/workers/assignments/{assignment_id}/result" + headers = _signed_headers( + method="POST", + path=path, + body=body, + hotkey=signer_pubkey, + nonce=self._next_nonce("res"), + timestamp=self.clock.time(), + ) + return await self.client.post(path, content=body, headers=headers) + + async def seed_assignment( + self, + *, + work_unit_id: str, + worker_id: str, + worker_pubkey: str, + miner_hotkey: str, + required_capability: str = "gpu", + ) -> WorkerAssignment: + return await self.assignment_service.create_worker_assignment( + work_unit_id=work_unit_id, + challenge_slug="prism", + submission_ref=f"ref-{work_unit_id}", + worker_id=worker_id, + worker_pubkey=worker_pubkey, + miner_hotkey=miner_hotkey, + payload={"run_spec": {"image": "img", "command": ["run"]}}, + required_capability=required_capability, + ) + + async def row(self, assignment_id: str) -> WorkerAssignment | None: + import uuid + + async with self.session_factory() as session: + return ( + await session.execute( + select(WorkerAssignment).where( + WorkerAssignment.id == uuid.UUID(assignment_id) + ) + ) + ).scalar_one_or_none() + + async def set_status(self, worker_pubkey: str, status: WorkerStatus) -> None: + from base.db import WorkerRegistration + + async with session_scope(self.session_factory) as session: + row = ( + await session.execute( + select(WorkerRegistration).where( + WorkerRegistration.worker_pubkey == worker_pubkey + ) + ) + ).scalar_one() + row.status = status + + +async def _build_harness() -> tuple[Harness, Any]: + engine = create_engine("sqlite+aiosqlite:///:memory:") + async with engine.begin() as connection: + await connection.run_sync(Base.metadata.create_all) + session_factory = create_session_factory(engine) + + cache = MetagraphCache(netuid=1, ttl_seconds=300) + cache.update_from_metagraph( + [MINER_H1, MINER_H2, VALIDATOR], + validator_permits=[False, False, True], + stakes=[0.0, 0.0, 100.0], + ) + clock = FakeClock(NOW_EPOCH) + service = WorkerCoordinationService( + session_factory, + miner_membership=MetagraphMinerMembership(cache), + binding_nonce_store=SqlAlchemyWorkerNonceStore(session_factory), + signature_verifier=_fake_verifier, + heartbeat_ttl_seconds=TTL_SECONDS, + now_fn=clock.now, + ) + verifier = WorkerSignedRequestVerifier( + nonce_store=SqlAlchemyWorkerNonceStore(session_factory), + eligibility=CoordinationReadEligibility(session_factory, cache), + signature_verifier=_fake_verifier, + ttl_seconds=300, + now_fn=clock.time, + ) + assignment_service = WorkerAssignmentService( + session_factory, worker_service=service, now_fn=clock.now + ) + # The assignment surface uses a WORKER-ONLY verifier (no validator permit). + assignment_verifier = WorkerSignedRequestVerifier( + nonce_store=SqlAlchemyWorkerNonceStore(session_factory), + eligibility=RegisteredWorkerEligibility(session_factory), + signature_verifier=_fake_verifier, + ttl_seconds=300, + now_fn=clock.time, + ) + + app = create_proxy_app( + registry=object(), + nonce_store=FakeNonceStore(), + metagraph_cache=FakeCache(), # type: ignore[arg-type] + worker_service=service, + worker_verifier=verifier, + worker_assignment_service=assignment_service, + worker_assignment_verifier=assignment_verifier, + ) + transport = ASGITransport(app=app) + client = AsyncClient(transport=transport, base_url="http://testserver") + return Harness(client, session_factory, clock, service, assignment_service), engine + + +@pytest.fixture +async def harness() -> AsyncIterator[Harness]: + h, engine = await _build_harness() + try: + yield h + finally: + await h.client.aclose() + await engine.dispose() + + +# VAL-AGENT-007 +async def test_pull_returns_only_gpu_units(harness: Harness) -> None: + worker_id = await harness.enroll_active(worker_pubkey="wp-1", miner_hotkey=MINER_H1) + gpu = await harness.seed_assignment( + work_unit_id="gpu-unit", + worker_id=worker_id, + worker_pubkey="wp-1", + miner_hotkey=MINER_H1, + required_capability="gpu", + ) + await harness.seed_assignment( + work_unit_id="cpu-unit", + worker_id=worker_id, + worker_pubkey="wp-1", + miner_hotkey=MINER_H1, + required_capability="cpu", + ) + + resp = await harness.pull(signer_pubkey="wp-1") + assert resp.status_code == 200, resp.text + assignments = resp.json()["assignments"] + assert [a["work_unit_id"] for a in assignments] == ["gpu-unit"] + assert assignments[0]["required_capability"] == "gpu" + assert assignments[0]["status"] == "running" + + # The gpu replica transitioned assigned -> running; the cpu one is untouched. + gpu_row = await harness.row(str(gpu.id)) + assert gpu_row is not None + assert WorkAssignmentStatus(gpu_row.status) == WorkAssignmentStatus.RUNNING + + +# VAL-AGENT-008 +async def test_posted_result_carries_and_persists_execution_proof( + harness: Harness, +) -> None: + import bittensor as bt + + from base.validator.agent.signing import KeypairRequestSigner + from base.worker.proof import build_execution_proof, verify_execution_proof + + worker_id = await harness.enroll_active(worker_pubkey="wp-1", miner_hotkey=MINER_H1) + unit = await harness.seed_assignment( + work_unit_id="unit-proof", + worker_id=worker_id, + worker_pubkey="wp-1", + miner_hotkey=MINER_H1, + ) + await harness.pull(signer_pubkey="wp-1") + + proof_signer = KeypairRequestSigner(bt.Keypair.create_from_uri("//Worker1")) + proof = build_execution_proof( + signer=proof_signer, + manifest_sha256="d" * 64, + unit_id="unit-proof", + ) + resp = await harness.post_result( + assignment_id=str(unit.id), + signer_pubkey="wp-1", + success=True, + payload={"execution_proof": proof.model_dump(mode="json")}, + ) + assert resp.status_code == 200, resp.text + assert resp.json()["status"] == "completed" + + row = await harness.row(str(unit.id)) + assert row is not None + assert row.result_success is True + assert row.result_payload is not None + stored_proof = row.result_payload["execution_proof"] + assert stored_proof["version"] == 1 + assert row.manifest_sha256 == "d" * 64 + # The stored worker signature verifies against the pinned message format. + from base.schemas.worker import ExecutionProof + + parsed = ExecutionProof.model_validate(stored_proof) + assert verify_execution_proof(parsed, unit_id="unit-proof") is True + + +# VAL-AGENT-018 (a): unregistered key cannot pull or post +async def test_unregistered_key_pull_rejected(harness: Harness) -> None: + resp = await harness.pull(signer_pubkey=STRANGER) + assert resp.status_code == 403 + + +async def test_validator_permit_only_key_cannot_pull(harness: Harness) -> None: + # VAL-AGENT-018 (d): a metagraph validator permit is NOT a worker credential. + resp = await harness.pull(signer_pubkey=VALIDATOR) + assert resp.status_code == 403 + + +# VAL-AGENT-018 (b): stale/retired worker receives no NEW units +async def test_stale_worker_receives_no_units(harness: Harness) -> None: + worker_id = await harness.enroll_active(worker_pubkey="wp-1", miner_hotkey=MINER_H1) + await harness.seed_assignment( + work_unit_id="gpu-unit", + worker_id=worker_id, + worker_pubkey="wp-1", + miner_hotkey=MINER_H1, + ) + # advance clock beyond TTL so the worker is effectively stale + harness.clock.epoch = NOW_EPOCH + TTL_SECONDS + 10 + resp = await harness.pull(signer_pubkey="wp-1") + assert resp.status_code == 200 + assert resp.json()["assignments"] == [] + + +async def test_retired_worker_receives_no_units(harness: Harness) -> None: + worker_id = await harness.enroll_active(worker_pubkey="wp-1", miner_hotkey=MINER_H1) + await harness.seed_assignment( + work_unit_id="gpu-unit", + worker_id=worker_id, + worker_pubkey="wp-1", + miner_hotkey=MINER_H1, + ) + await harness.set_status("wp-1", WorkerStatus.RETIRED) + resp = await harness.pull(signer_pubkey="wp-1") + assert resp.status_code == 200 + assert resp.json()["assignments"] == [] + + +async def test_pending_worker_receives_no_units(harness: Harness) -> None: + worker_id = await harness.register_pending( + worker_pubkey="wp-1", miner_hotkey=MINER_H1 + ) + await harness.seed_assignment( + work_unit_id="gpu-unit", + worker_id=worker_id, + worker_pubkey="wp-1", + miner_hotkey=MINER_H1, + ) + resp = await harness.pull(signer_pubkey="wp-1") + assert resp.status_code == 200 + assert resp.json()["assignments"] == [] + + +# VAL-AGENT-018 (c): active registered worker pull succeeds without a permit +async def test_active_worker_pull_succeeds_without_permit(harness: Harness) -> None: + worker_id = await harness.enroll_active(worker_pubkey="wp-1", miner_hotkey=MINER_H1) + await harness.seed_assignment( + work_unit_id="gpu-unit", + worker_id=worker_id, + worker_pubkey="wp-1", + miner_hotkey=MINER_H1, + ) + resp = await harness.pull(signer_pubkey="wp-1") + assert resp.status_code == 200 + assert [a["work_unit_id"] for a in resp.json()["assignments"]] == ["gpu-unit"] + + +# Foreign/late post rejected (VAL-MASTER-019 baseline) +async def test_foreign_worker_post_rejected(harness: Harness) -> None: + w1 = await harness.enroll_active(worker_pubkey="wp-1", miner_hotkey=MINER_H1) + await harness.enroll_active(worker_pubkey="wp-2", miner_hotkey=MINER_H2) + unit = await harness.seed_assignment( + work_unit_id="gpu-unit", + worker_id=w1, + worker_pubkey="wp-1", + miner_hotkey=MINER_H1, + ) + resp = await harness.post_result( + assignment_id=str(unit.id), + signer_pubkey="wp-2", + success=True, + payload={"manifest_sha256": "e" * 64}, + ) + assert resp.status_code == 403 + row = await harness.row(str(unit.id)) + assert row is not None + assert row.result_success is None + + +async def test_post_result_is_idempotent(harness: Harness) -> None: + worker_id = await harness.enroll_active(worker_pubkey="wp-1", miner_hotkey=MINER_H1) + unit = await harness.seed_assignment( + work_unit_id="gpu-unit", + worker_id=worker_id, + worker_pubkey="wp-1", + miner_hotkey=MINER_H1, + ) + first = await harness.post_result( + assignment_id=str(unit.id), + signer_pubkey="wp-1", + success=True, + payload={"manifest_sha256": "e" * 64}, + ) + assert first.status_code == 200 + assert first.json()["idempotent"] is False + second = await harness.post_result( + assignment_id=str(unit.id), + signer_pubkey="wp-1", + success=True, + payload={"manifest_sha256": "f" * 64}, + ) + assert second.status_code == 200 + assert second.json()["idempotent"] is True + row = await harness.row(str(unit.id)) + assert row is not None + assert row.manifest_sha256 == "e" * 64 + + +async def test_post_result_failure_records_success_false(harness: Harness) -> None: + worker_id = await harness.enroll_active(worker_pubkey="wp-1", miner_hotkey=MINER_H1) + unit = await harness.seed_assignment( + work_unit_id="gpu-unit", + worker_id=worker_id, + worker_pubkey="wp-1", + miner_hotkey=MINER_H1, + ) + resp = await harness.post_result( + assignment_id=str(unit.id), + signer_pubkey="wp-1", + success=False, + payload={"error": "broker down"}, + ) + assert resp.status_code == 200 + assert resp.json()["status"] == "failed" + row = await harness.row(str(unit.id)) + assert row is not None + assert row.result_success is False + assert row.manifest_sha256 is None diff --git a/tests/unit/test_worker_assignment_engine.py b/tests/unit/test_worker_assignment_engine.py new file mode 100644 index 00000000..1077271a --- /dev/null +++ b/tests/unit/test_worker_assignment_engine.py @@ -0,0 +1,579 @@ +"""Worker-plane replica-creation ENGINE behavior (architecture.md sec 3.3). + +Covers the master-side gpu routing the engine adds on top of the +worker-agent-runtime seams (``WorkerAssignmentService.create_worker_assignment``, +``WorkerCoordinationService.effective_status``): + +* VAL-MASTER-003: gpu units replicate only to ACTIVE workers. +* VAL-MASTER-004 / VAL-MASTER-005: self-evaluation exclusion, holding even when + the submitter's own worker is the sole capacity (the unit waits). +* VAL-MASTER-006: R=2 across DISTINCT owner hotkeys; a same-owner pair never + satisfies R=2. +* VAL-MASTER-007: graceful degradation to R=1 with a recorded warning. +* VAL-MASTER-011: per-replica deadline/reassignment; a unit with ALL replicas + attempts-exhausted ends ``failed``. +* VAL-MASTER-012: heartbeat TTL drives active->stale eligibility (and recovery). +* VAL-MASTER-013: flag OFF routes gpu units to validators; flag ON routes them + away from validators to worker replicas. +* VAL-MASTER-014: flag OFF leaves the worker assignment surface inert (404). +* VAL-MASTER-020: per-worker gpu concurrency is 1. +""" + +from __future__ import annotations + +import logging +from collections.abc import AsyncIterator +from dataclasses import dataclass +from datetime import UTC, datetime, timedelta +from typing import Any + +import pytest +from httpx import ASGITransport, AsyncClient +from sqlalchemy import select + +from base.bittensor.metagraph_cache import MetagraphCache +from base.db import ( + Base, + Validator, + ValidatorStatus, + WorkAssignmentStatus, + WorkerAssignment, + WorkerRegistration, + WorkerStatus, + create_engine, + create_session_factory, + session_scope, +) +from base.db.models import WorkAssignment +from base.master.app_proxy import create_proxy_app +from base.master.assignment import AssignmentService +from base.master.worker_assignment import WorkerAssignmentService +from base.master.worker_assignment_engine import ( + DEGRADED_REPLICATION_PAYLOAD_KEY, + WorkerAssignmentEngine, + run_worker_assignment_pass, +) +from base.master.worker_coordination import WorkerCoordinationService +from base.security.worker_auth import ( + MetagraphMinerMembership, + SqlAlchemyWorkerNonceStore, +) + +NOW = datetime(2026, 6, 27, 12, 0, 0, tzinfo=UTC) +TTL = 120 + +# Submitter H and the distinct owner hotkeys A/B/C (none is the submitter). +MINER_H = "miner-H" +MINER_A = "miner-A" +MINER_B = "miner-B" +MINER_C = "miner-C" + + +class FakeClock: + def __init__(self, moment: datetime) -> None: + self.moment = moment + + def now(self) -> datetime: + return self.moment + + +class _FakeNonceStore: + async def reserve(self, **_: Any) -> None: + return None + + +class _FakeCache: + def get(self) -> dict[str, int]: + return {} + + +@dataclass +class Env: + db_engine: Any + factory: Any + clock: FakeClock + worker_service: WorkerCoordinationService + worker_assignment_service: WorkerAssignmentService + engine: WorkerAssignmentEngine + + +async def _build_env(*, replication_factor: int = 2) -> Env: + db_engine = create_engine("sqlite+aiosqlite:///:memory:") + async with db_engine.begin() as connection: + await connection.run_sync(Base.metadata.create_all) + factory = create_session_factory(db_engine) + clock = FakeClock(NOW) + cache = MetagraphCache(netuid=1, ttl_seconds=300) + worker_service = WorkerCoordinationService( + factory, + miner_membership=MetagraphMinerMembership(cache), + binding_nonce_store=SqlAlchemyWorkerNonceStore(factory), + heartbeat_ttl_seconds=TTL, + now_fn=clock.now, + ) + worker_assignment_service = WorkerAssignmentService( + factory, worker_service=worker_service, now_fn=clock.now + ) + engine = WorkerAssignmentEngine( + factory, + assignment_service=worker_assignment_service, + worker_service=worker_service, + replication_factor=replication_factor, + now_fn=clock.now, + ) + return Env( + db_engine=db_engine, + factory=factory, + clock=clock, + worker_service=worker_service, + worker_assignment_service=worker_assignment_service, + engine=engine, + ) + + +async def _add_worker( + factory: Any, + *, + worker_pubkey: str, + miner_hotkey: str, + status: WorkerStatus = WorkerStatus.ACTIVE, + last_heartbeat_at: datetime | None = NOW, + created_at: datetime = NOW, +) -> str: + worker_id = f"wid-{worker_pubkey}" + async with session_scope(factory) as session: + session.add( + WorkerRegistration( + worker_id=worker_id, + worker_pubkey=worker_pubkey, + miner_hotkey=miner_hotkey, + binding_signature="sig", + binding_nonce=f"nonce-{worker_pubkey}", + provider="local", + provider_instance_ref="local-1", + capabilities=["gpu"], + status=status, + last_heartbeat_at=last_heartbeat_at, + created_at=created_at, + updated_at=created_at, + ) + ) + return worker_id + + +async def _add_gpu_unit( + factory: Any, + *, + work_unit_id: str, + submitter: str, + max_attempts: int = 3, + created_at: datetime = NOW, +) -> None: + async with session_scope(factory) as session: + session.add( + WorkAssignment( + challenge_slug="prism", + work_unit_id=work_unit_id, + submission_ref=submitter, + payload={"run_spec": {"image": "img"}}, + required_capability="gpu", + status=WorkAssignmentStatus.PENDING, + attempt_count=0, + max_attempts=max_attempts, + created_at=created_at, + updated_at=created_at, + ) + ) + + +async def _replicas(factory: Any, work_unit_id: str) -> list[WorkerAssignment]: + async with factory() as session: + rows = ( + ( + await session.execute( + select(WorkerAssignment) + .where(WorkerAssignment.work_unit_id == work_unit_id) + .order_by(WorkerAssignment.worker_id) + ) + ) + .scalars() + .all() + ) + return list(rows) + + +async def _unit(factory: Any, work_unit_id: str) -> WorkAssignment: + async with factory() as session: + return ( + await session.execute( + select(WorkAssignment).where( + WorkAssignment.work_unit_id == work_unit_id + ) + ) + ).scalar_one() + + +@pytest.fixture +async def env() -> AsyncIterator[Env]: + e = await _build_env() + try: + yield e + finally: + await e.db_engine.dispose() + + +# VAL-MASTER-003 +async def test_gpu_units_replicate_only_to_active_workers(env: Env) -> None: + await _add_worker(env.factory, worker_pubkey="wp-a", miner_hotkey=MINER_A) + await _add_worker(env.factory, worker_pubkey="wp-b", miner_hotkey=MINER_B) + await _add_worker( + env.factory, + worker_pubkey="wp-pending", + miner_hotkey=MINER_C, + status=WorkerStatus.PENDING, + last_heartbeat_at=None, + ) + await _add_worker( + env.factory, + worker_pubkey="wp-stale", + miner_hotkey="miner-D", + status=WorkerStatus.STALE, + last_heartbeat_at=NOW - timedelta(seconds=TTL + 5), + ) + await _add_worker( + env.factory, + worker_pubkey="wp-retired", + miner_hotkey="miner-E", + status=WorkerStatus.RETIRED, + ) + await _add_gpu_unit(env.factory, work_unit_id="U", submitter=MINER_H) + + await env.engine.assign_pending(seed=1) + + replicas = await _replicas(env.factory, "U") + assert len(replicas) == 2 + assert {r.worker_pubkey for r in replicas} == {"wp-a", "wp-b"} + inactive = {"wp-pending", "wp-stale", "wp-retired"} + assert all(r.worker_pubkey not in inactive for r in replicas) + + +# VAL-MASTER-004 +async def test_self_evaluation_exclusion(env: Env) -> None: + # H owns an active, capable worker, but it must never serve H's own unit. + await _add_worker(env.factory, worker_pubkey="wp-h", miner_hotkey=MINER_H) + await _add_worker(env.factory, worker_pubkey="wp-a", miner_hotkey=MINER_A) + await _add_worker(env.factory, worker_pubkey="wp-b", miner_hotkey=MINER_B) + await _add_gpu_unit(env.factory, work_unit_id="U", submitter=MINER_H) + + # Repeated passes never place the unit on H's worker. + await env.engine.assign_pending(seed=1) + await env.engine.assign_pending(seed=1) + + replicas = await _replicas(env.factory, "U") + assert len(replicas) == 2 + owners = {r.miner_hotkey for r in replicas} + assert MINER_H not in owners + assert owners == {MINER_A, MINER_B} + + +# VAL-MASTER-005 +async def test_self_evaluation_holds_when_only_own_worker_available(env: Env) -> None: + await _add_worker(env.factory, worker_pubkey="wp-h", miner_hotkey=MINER_H) + await _add_gpu_unit(env.factory, work_unit_id="U", submitter=MINER_H) + + await env.engine.assign_pending(seed=1) + + # The unit is never assigned to H's worker; it simply waits. + assert await _replicas(env.factory, "U") == [] + assert (await _unit(env.factory, "U")).status == WorkAssignmentStatus.PENDING + + # A distinct-owner worker activates -> the unit is assigned promptly. + await _add_worker(env.factory, worker_pubkey="wp-a", miner_hotkey=MINER_A) + await env.engine.assign_pending(seed=1) + + replicas = await _replicas(env.factory, "U") + assert [r.miner_hotkey for r in replicas] == [MINER_A] + assert (await _unit(env.factory, "U")).status == WorkAssignmentStatus.ASSIGNED + + +# VAL-MASTER-006 +async def test_replication_two_across_distinct_owners(env: Env) -> None: + await _add_worker(env.factory, worker_pubkey="wp-a", miner_hotkey=MINER_A) + await _add_worker(env.factory, worker_pubkey="wp-b", miner_hotkey=MINER_B) + await _add_gpu_unit(env.factory, work_unit_id="U", submitter=MINER_H) + + await env.engine.assign_pending(seed=1) + + replicas = await _replicas(env.factory, "U") + assert len(replicas) == 2 + assert len({r.miner_hotkey for r in replicas}) == 2 + unit = await _unit(env.factory, "U") + assert DEGRADED_REPLICATION_PAYLOAD_KEY not in unit.payload + + +# VAL-MASTER-006 (negative): two workers of the SAME owner never satisfy R=2. +async def test_same_owner_pair_never_satisfies_r2(env: Env) -> None: + await _add_worker(env.factory, worker_pubkey="wp-a1", miner_hotkey=MINER_A) + await _add_worker(env.factory, worker_pubkey="wp-a2", miner_hotkey=MINER_A) + await _add_gpu_unit(env.factory, work_unit_id="U", submitter=MINER_H) + + await env.engine.assign_pending(seed=1) + + replicas = await _replicas(env.factory, "U") + assert len(replicas) == 1 + assert replicas[0].miner_hotkey == MINER_A + + +# VAL-MASTER-007 +async def test_degradation_to_r1_records_warning( + env: Env, caplog: pytest.LogCaptureFixture +) -> None: + # Only one distinct eligible owner: H's worker is excluded by self-eval. + await _add_worker(env.factory, worker_pubkey="wp-h", miner_hotkey=MINER_H) + await _add_worker(env.factory, worker_pubkey="wp-a", miner_hotkey=MINER_A) + await _add_gpu_unit(env.factory, work_unit_id="U", submitter=MINER_H) + + with caplog.at_level( + logging.WARNING, logger="base.master.worker_assignment_engine" + ): + result = await env.engine.assign_pending(seed=1) + + replicas = await _replicas(env.factory, "U") + assert [r.miner_hotkey for r in replicas] == [MINER_A] + assert result.degraded == ["U"] + # The warning is durably recorded on the unit AND emitted as a log record. + unit = await _unit(env.factory, "U") + assert unit.payload[DEGRADED_REPLICATION_PAYLOAD_KEY] == 1 + assert any( + "degraded" in record.message and "U" in record.getMessage() + for record in caplog.records + ) + + +# VAL-MASTER-011: a lapsed replica reassigns to another eligible worker. +async def test_lapsed_replica_reassigned_to_other_worker(env: Env) -> None: + await _add_worker(env.factory, worker_pubkey="wp-a", miner_hotkey=MINER_A) + await _add_worker(env.factory, worker_pubkey="wp-b", miner_hotkey=MINER_B) + await _add_gpu_unit(env.factory, work_unit_id="U", submitter=MINER_H) + await env.engine.assign_pending(seed=1) + + # A's replica lease lapses; a fresh distinct-owner worker C is available. + async with session_scope(env.factory) as session: + row = ( + await session.execute( + select(WorkerAssignment).where( + WorkerAssignment.work_unit_id == "U", + WorkerAssignment.miner_hotkey == MINER_A, + ) + ) + ).scalar_one() + row.deadline_at = NOW - timedelta(seconds=10) + await _add_worker(env.factory, worker_pubkey="wp-c", miner_hotkey=MINER_C) + + result = await env.engine.reassign_stale_replicas(seed=1) + + assert result.reassigned == ["U"] + replicas = await _replicas(env.factory, "U") + owners = {r.miner_hotkey for r in replicas} + # B's replica untouched; A's replica moved to the new distinct owner C. + assert owners == {MINER_B, MINER_C} + moved = next(r for r in replicas if r.miner_hotkey == MINER_C) + assert moved.attempt_count == 2 + assert moved.deadline_at is None + assert WorkAssignmentStatus(moved.status) == WorkAssignmentStatus.ASSIGNED + untouched = next(r for r in replicas if r.miner_hotkey == MINER_B) + assert untouched.attempt_count == 1 + + +# VAL-MASTER-011: all replicas exhausted -> the unit itself fails. +async def test_all_replicas_exhausted_fails_unit(env: Env) -> None: + await _add_worker(env.factory, worker_pubkey="wp-a", miner_hotkey=MINER_A) + await _add_worker(env.factory, worker_pubkey="wp-b", miner_hotkey=MINER_B) + await _add_gpu_unit( + env.factory, work_unit_id="U", submitter=MINER_H, max_attempts=1 + ) + await env.engine.assign_pending(seed=1) + + # Both replicas are at their attempt cap and lapse with no other capacity. + async with session_scope(env.factory) as session: + rows = ( + ( + await session.execute( + select(WorkerAssignment).where(WorkerAssignment.work_unit_id == "U") + ) + ) + .scalars() + .all() + ) + for row in rows: + row.deadline_at = NOW - timedelta(seconds=10) + + result = await env.engine.reassign_stale_replicas(seed=1) + + assert result.failed_units == ["U"] + assert len(result.failed_replicas) == 2 + replicas = await _replicas(env.factory, "U") + assert all( + WorkAssignmentStatus(r.status) == WorkAssignmentStatus.FAILED for r in replicas + ) + assert (await _unit(env.factory, "U")).status == WorkAssignmentStatus.FAILED + + +# VAL-MASTER-012 +async def test_heartbeat_ttl_drives_eligibility(env: Env) -> None: + worker_pubkey = "wp-a" + await _add_worker(env.factory, worker_pubkey=worker_pubkey, miner_hotkey=MINER_A) + await _add_gpu_unit(env.factory, work_unit_id="U", submitter=MINER_H) + + # Clock advances past the TTL -> the worker is stale and not assignable. + env.clock.moment = NOW + timedelta(seconds=TTL + 10) + async with env.factory() as session: + worker = ( + await session.execute( + select(WorkerRegistration).where( + WorkerRegistration.worker_pubkey == worker_pubkey + ) + ) + ).scalar_one() + assert ( + env.worker_service.effective_status(worker, env.clock.now()) + == WorkerStatus.STALE + ) + + await env.engine.assign_pending(seed=1) + assert await _replicas(env.factory, "U") == [] + assert (await _unit(env.factory, "U")).status == WorkAssignmentStatus.PENDING + + # The worker heartbeats again -> active -> regains eligibility. + async with session_scope(env.factory) as session: + worker = ( + await session.execute( + select(WorkerRegistration).where( + WorkerRegistration.worker_pubkey == worker_pubkey + ) + ) + ).scalar_one() + worker.last_heartbeat_at = env.clock.now() + assert ( + env.worker_service.effective_status(worker, env.clock.now()) + == WorkerStatus.ACTIVE + ) + + await env.engine.assign_pending(seed=1) + replicas = await _replicas(env.factory, "U") + assert [r.worker_pubkey for r in replicas] == [worker_pubkey] + + +# VAL-MASTER-020 +async def test_per_worker_gpu_concurrency_is_one(env: Env) -> None: + await _add_worker(env.factory, worker_pubkey="wp-a", miner_hotkey=MINER_A) + await _add_gpu_unit(env.factory, work_unit_id="unit-a", submitter=MINER_H) + await _add_gpu_unit(env.factory, work_unit_id="unit-b", submitter=MINER_H) + + await env.engine.assign_pending(seed=1) + + # The single worker holds one in-flight replica; the second unit waits. + assert len(await _replicas(env.factory, "unit-a")) == 1 + assert await _replicas(env.factory, "unit-b") == [] + assert (await _unit(env.factory, "unit-b")).status == WorkAssignmentStatus.PENDING + + # The first replica finishes -> the worker is eligible for the queued unit. + async with session_scope(env.factory) as session: + row = ( + await session.execute( + select(WorkerAssignment).where( + WorkerAssignment.work_unit_id == "unit-a" + ) + ) + ).scalar_one() + row.status = WorkAssignmentStatus.COMPLETED + + await env.engine.assign_pending(seed=1) + replicas = await _replicas(env.factory, "unit-b") + assert [r.worker_pubkey for r in replicas] == ["wp-a"] + + +# VAL-MASTER-013: flag OFF routes gpu units to validators (never to workers). +async def test_flag_off_routes_gpu_to_validator(env: Env) -> None: + validator_plane = AssignmentService(env.factory, now_fn=env.clock.now) + async with session_scope(env.factory) as session: + session.add( + Validator( + hotkey="gpu-val", + uid=None, + status=ValidatorStatus.ONLINE, + capabilities=["gpu"], + version="1.0.0", + registered_at=NOW, + last_heartbeat_at=NOW, + ) + ) + # An active worker exists but the flag is OFF, so it is ignored. + await _add_worker(env.factory, worker_pubkey="wp-a", miner_hotkey=MINER_A) + await _add_gpu_unit(env.factory, work_unit_id="U", submitter=MINER_H) + + await validator_plane.assign_pending(seed=1) + + unit = await _unit(env.factory, "U") + assert unit.status == WorkAssignmentStatus.ASSIGNED + assert unit.assigned_validator_hotkey == "gpu-val" + assert await _replicas(env.factory, "U") == [] + + +# VAL-MASTER-013: flag ON routes gpu units AWAY from validators to workers. +async def test_flag_on_routes_gpu_away_from_validators(env: Env) -> None: + validator_plane = AssignmentService( + env.factory, now_fn=env.clock.now, worker_plane_capabilities=frozenset({"gpu"}) + ) + async with session_scope(env.factory) as session: + session.add( + Validator( + hotkey="gpu-val", + uid=None, + status=ValidatorStatus.ONLINE, + capabilities=["gpu"], + version="1.0.0", + registered_at=NOW, + last_heartbeat_at=NOW, + ) + ) + await _add_worker(env.factory, worker_pubkey="wp-a", miner_hotkey=MINER_A) + await _add_worker(env.factory, worker_pubkey="wp-b", miner_hotkey=MINER_B) + await _add_gpu_unit(env.factory, work_unit_id="U", submitter=MINER_H) + + # The validator plane skips the gpu unit; the engine materializes replicas. + await validator_plane.assign_pending(seed=1) + assert (await _unit(env.factory, "U")).assigned_validator_hotkey is None + assert (await _unit(env.factory, "U")).status == WorkAssignmentStatus.PENDING + + await run_worker_assignment_pass(engine=env.engine, seed=1) + replicas = await _replicas(env.factory, "U") + assert len(replicas) == 2 + assert {r.miner_hotkey for r in replicas} == {MINER_A, MINER_B} + + +# VAL-MASTER-014: with the plane OFF the worker assignment surface is inert (404). +async def test_flag_off_worker_assignment_surface_inert() -> None: + db_engine = create_engine("sqlite+aiosqlite:///:memory:") + async with db_engine.begin() as connection: + await connection.run_sync(Base.metadata.create_all) + + app = create_proxy_app( + registry=object(), + nonce_store=_FakeNonceStore(), + metagraph_cache=_FakeCache(), # type: ignore[arg-type] + worker_assignment_service=None, + worker_assignment_verifier=None, + ) + transport = ASGITransport(app=app) + client = AsyncClient(transport=transport, base_url="http://testserver") + try: + pull = await client.post("/v1/workers/assignments/pull", json={}) + assert pull.status_code == 404 + result = await client.post("/v1/workers/assignments/abc/result", json={}) + assert result.status_code == 404 + health = await client.get("/health") + assert health.status_code == 200 + finally: + await client.aclose() + await db_engine.dispose() diff --git a/tests/unit/test_worker_auth.py b/tests/unit/test_worker_auth.py new file mode 100644 index 00000000..a0368799 --- /dev/null +++ b/tests/unit/test_worker_auth.py @@ -0,0 +1,199 @@ +from __future__ import annotations + +import hashlib +from typing import Any + +import pytest + +from base.bittensor.metagraph_cache import MetagraphCache +from base.db import ( + Base, + WorkerRegistration, + WorkerStatus, + create_engine, + create_session_factory, + session_scope, +) +from base.security.validator_auth import canonical_validator_request +from base.security.worker_auth import ( + CoordinationReadEligibility, + MetagraphMinerMembership, + RegisteredWorkerEligibility, + SqlAlchemyWorkerNonceStore, + WorkerAuthError, + WorkerEligibilityError, + WorkerReplayError, + WorkerSignedRequestVerifier, + worker_binding_message, +) + +MINER = "miner-1" +VALIDATOR = "validator-1" +WORKER = "worker-pubkey-1" + + +def _sign(hotkey: str, message: bytes) -> str: + return hashlib.sha256(hotkey.encode() + b":" + message).hexdigest() + + +def _fake_verifier(hotkey: str, message: bytes, signature: str) -> bool: + return signature == _sign(hotkey, message) + + +def test_worker_binding_message_matches_pinned_format() -> None: + assert ( + worker_binding_message(worker_pubkey="wp", miner_hotkey="mh", nonce="n1") + == b"worker-binding:wp:mh:n1" + ) + + +def test_metagraph_miner_membership_requires_only_graph_presence() -> None: + cache = MetagraphCache(netuid=1, ttl_seconds=300) + # miner is present WITHOUT a validator permit; validator holds a permit. + cache.update_from_metagraph( + [MINER, VALIDATOR], validator_permits=[False, True], stakes=[0.0, 1.0] + ) + membership = MetagraphMinerMembership(cache) + assert membership.is_registered(MINER) is True + assert membership.is_registered("off-graph") is False + + +async def _session_factory_with_worker(status: WorkerStatus) -> Any: + engine = create_engine("sqlite+aiosqlite:///:memory:") + async with engine.begin() as connection: + await connection.run_sync(Base.metadata.create_all) + session_factory = create_session_factory(engine) + async with session_scope(session_factory) as session: + session.add( + WorkerRegistration( + worker_id="w-1", + worker_pubkey=WORKER, + miner_hotkey=MINER, + binding_signature="sig", + binding_nonce="n1", + provider="local", + provider_instance_ref=None, + capabilities=["gpu"], + status=status, + ) + ) + return session_factory, engine + + +async def test_registered_worker_eligibility_matches_any_status() -> None: + session_factory, engine = await _session_factory_with_worker(WorkerStatus.RETIRED) + try: + eligibility = RegisteredWorkerEligibility(session_factory) + assert await eligibility.is_eligible(WORKER) is True + assert await eligibility.is_eligible("unknown") is False + finally: + await engine.dispose() + + +async def test_coordination_read_eligibility_allows_worker_or_validator() -> None: + session_factory, engine = await _session_factory_with_worker(WorkerStatus.ACTIVE) + try: + cache = MetagraphCache(netuid=1, ttl_seconds=300) + cache.update_from_metagraph( + [MINER, VALIDATOR], validator_permits=[False, True], stakes=[0.0, 1.0] + ) + eligibility = CoordinationReadEligibility(session_factory, cache) + assert await eligibility.is_eligible(WORKER) is True # registered worker + assert await eligibility.is_eligible(VALIDATOR) is True # metagraph validator + assert await eligibility.is_eligible(MINER) is False # miner, not a worker + assert await eligibility.is_eligible("stranger") is False + finally: + await engine.dispose() + + +async def _verify( + verifier: WorkerSignedRequestVerifier, + *, + hotkey: str, + nonce: str, + signature: str | None = None, + method: str = "GET", + path: str = "/v1/workers", + timestamp: str = "1000", +) -> Any: + canonical = canonical_validator_request( + method=method, + path=path, + query_string="", + timestamp=timestamp, + nonce=nonce, + body=b"", + ) + sig = signature if signature is not None else _sign(hotkey, canonical.encode()) + return await verifier.verify( + method=method, + path=path, + query_string="", + headers={ + "X-Hotkey": hotkey, + "X-Signature": sig, + "X-Nonce": nonce, + "X-Timestamp": timestamp, + }, + body=b"", + ) + + +async def test_signed_request_verifier_happy_path_and_replay() -> None: + session_factory, engine = await _session_factory_with_worker(WorkerStatus.ACTIVE) + try: + cache = MetagraphCache(netuid=1, ttl_seconds=300) + cache.update_from_metagraph([VALIDATOR], validator_permits=[True], stakes=[1.0]) + verifier = WorkerSignedRequestVerifier( + nonce_store=SqlAlchemyWorkerNonceStore(session_factory), + eligibility=CoordinationReadEligibility(session_factory, cache), + signature_verifier=_fake_verifier, + ttl_seconds=300, + now_fn=lambda: 1000.0, + ) + identity = await _verify(verifier, hotkey=WORKER, nonce="rq-1") + assert identity.hotkey == WORKER + + with pytest.raises(WorkerReplayError): + await _verify(verifier, hotkey=WORKER, nonce="rq-1") + finally: + await engine.dispose() + + +async def test_signed_request_verifier_rejects_bad_signature_and_ineligible() -> None: + session_factory, engine = await _session_factory_with_worker(WorkerStatus.ACTIVE) + try: + cache = MetagraphCache(netuid=1, ttl_seconds=300) + cache.update_from_metagraph([VALIDATOR], validator_permits=[True], stakes=[1.0]) + verifier = WorkerSignedRequestVerifier( + nonce_store=SqlAlchemyWorkerNonceStore(session_factory), + eligibility=CoordinationReadEligibility(session_factory, cache), + signature_verifier=_fake_verifier, + ttl_seconds=300, + now_fn=lambda: 1000.0, + ) + with pytest.raises(WorkerAuthError): + await _verify(verifier, hotkey=WORKER, nonce="rq-2", signature="bad") + + with pytest.raises(WorkerEligibilityError): + await _verify(verifier, hotkey="stranger", nonce="rq-3") + finally: + await engine.dispose() + + +async def test_signed_request_verifier_rejects_stale_timestamp() -> None: + session_factory, engine = await _session_factory_with_worker(WorkerStatus.ACTIVE) + try: + cache = MetagraphCache(netuid=1, ttl_seconds=300) + cache.update_from_metagraph([VALIDATOR], validator_permits=[True], stakes=[1.0]) + verifier = WorkerSignedRequestVerifier( + nonce_store=SqlAlchemyWorkerNonceStore(session_factory), + eligibility=CoordinationReadEligibility(session_factory, cache), + signature_verifier=_fake_verifier, + ttl_seconds=10, + now_fn=lambda: 1_000_000.0, + ) + with pytest.raises(WorkerAuthError): + await _verify(verifier, hotkey=WORKER, nonce="rq-4", timestamp="1000") + finally: + await engine.dispose() diff --git a/tests/unit/test_worker_cli.py b/tests/unit/test_worker_cli.py new file mode 100644 index 00000000..80b24d45 --- /dev/null +++ b/tests/unit/test_worker_cli.py @@ -0,0 +1,404 @@ +"""CLI tests for the top-level ``base worker`` app (VAL-AGENT-001/010/011/012). + +Exercises the Typer surface with :class:`CliRunner`: the ``base worker`` group is +DISTINCT from the legacy ``base master worker`` Swarm-node group; ``deploy`` +enforces the provider key before any network/config work; provider planning bounds +selection by ``--max-price`` (preferring an exact GPU-count executor) and NEVER +places a provider key in the provisioned pod env; ``local`` reports the active +worker; and ``status`` renders the fleet. +""" + +from __future__ import annotations + +from datetime import UTC, datetime +from types import SimpleNamespace + +import pytest +from typer.testing import CliRunner + +import base.cli_app.main as cli_main +from base.compute.provider import Instance, InstanceSpec, Offer + +runner = CliRunner() + +_SENTINEL_KEY = "SENTINEL-PROVIDER-KEY-DO-NOT-LEAK" + + +class _FakeKeypair: + def __init__(self, address: str) -> None: + self.ss58_address = address + + def sign(self, message: bytes) -> bytes: + return b"sig:" + message + + +@pytest.fixture +def _fake_keys(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr( + cli_main, "create_worker_keypair", lambda _s: _FakeKeypair("worker-ss58") + ) + monkeypatch.setattr( + cli_main, + "create_worker_miner_keypair", + lambda _s: _FakeKeypair("miner-ss58"), + ) + monkeypatch.setattr(cli_main, "_configure_observability", lambda _s: None) + + +# -- VAL-AGENT-001: command surface ------------------------------------------- + + +def test_worker_help_lists_agent_deploy_status() -> None: + result = runner.invoke(cli_main.app, ["worker", "--help"]) + assert result.exit_code == 0 + assert "agent" in result.stdout + assert "deploy" in result.stdout + assert "status" in result.stdout + + +def test_legacy_master_worker_group_intact() -> None: + result = runner.invoke(cli_main.app, ["master", "worker", "--help"]) + assert result.exit_code == 0 + for legacy in ("token", "list", "label", "drain", "rm", "inspect"): + assert legacy in result.stdout + + +# -- VAL-AGENT-010: provider-key refusal -------------------------------------- + + +@pytest.mark.parametrize( + ("provider", "env_var"), + [("lium", "LIUM_API_KEY"), ("targon", "TARGON_API_KEY")], +) +def test_deploy_without_provider_key_refuses_before_any_work( + provider: str, env_var: str, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.delenv(env_var, raising=False) + + def _fail_load(*_a: object, **_k: object) -> object: + raise AssertionError("load_settings must not run before the key check") + + monkeypatch.setattr(cli_main, "load_settings", _fail_load) + result = runner.invoke(cli_main.app, ["worker", "deploy", "--provider", provider]) + assert result.exit_code == 2 + assert env_var in result.stderr + assert "No provider or master call was made" in result.stderr + + +def test_deploy_rejects_unknown_provider() -> None: + result = runner.invoke(cli_main.app, ["worker", "deploy", "--provider", "aws"]) + assert result.exit_code == 2 + assert "unsupported provider" in result.stderr + + +# -- VAL-AGENT-012 + 011: planning bounds + provider-key hygiene -------------- + + +def test_deploy_provider_selects_in_budget_offer_without_leaking_key( + monkeypatch: pytest.MonkeyPatch, tmp_path: object, _fake_keys: None +) -> None: + captured: dict[str, object] = {} + + class _FakeLiumClient: + def __init__(self, api_key: str) -> None: + captured["api_key"] = api_key + + async def list_offers( + self, *, max_price_per_hour: float | None = None + ) -> list[Offer]: + captured["max_price"] = max_price_per_hour + return [ + Offer( + id="multi-cheap", + gpu_type="H100", + gpu_count=8, + price_per_hour=0.4, + ), + Offer( + id="single-fit", + gpu_type="H100", + gpu_count=1, + price_per_hour=0.7, + ), + Offer( + id="over-cap", + gpu_type="H100", + gpu_count=1, + price_per_hour=5.0, + ), + ] + + async def provision(self, spec: object, *, offer: Offer) -> Instance: + captured["spec"] = spec + captured["offer"] = offer + return Instance(id="pod-xyz", status="PENDING") + + monkeypatch.setenv("LIUM_API_KEY", _SENTINEL_KEY) + monkeypatch.setenv("BASE_WORKER__DEPLOY__IMAGE", "ghcr.io/public/base-worker") + monkeypatch.setenv("BASE_WORKER__DEPLOY__IMAGE_DIGEST", "sha256:" + "a" * 64) + monkeypatch.setattr(cli_main, "LiumClient", _FakeLiumClient) + + result = runner.invoke( + cli_main.app, + [ + "worker", + "deploy", + "--provider", + "lium", + "--max-price", + "1.0", + "--config", + "config/worker.example.yaml", + ], + ) + assert result.exit_code == 0, result.stdout + # In-budget + exact-gpu-count offer preferred over the cheaper 8-GPU node. + assert captured["max_price"] == 1.0 + offer = captured["offer"] + assert isinstance(offer, Offer) + assert offer.id == "single-fit" + assert "single-fit" in result.stdout + assert "pod-xyz" in result.stdout + # VAL-AGENT-011: the provider key authenticated the client but is absent from + # the provisioned pod env handed to the worker agent. + assert captured["api_key"] == _SENTINEL_KEY + spec = captured["spec"] + assert isinstance(spec, InstanceSpec) + env = dict(spec.env) + assert _SENTINEL_KEY not in repr(env) + assert "LIUM_API_KEY" not in env + assert env["BASE_WORKER__IDENTITY__MINER_HOTKEY"] == "miner-ss58" + # The explicitly-configured worker image is pinned (no silent placeholder). + assert spec.image == "ghcr.io/public/base-worker" + assert spec.image_digest == "sha256:" + "a" * 64 + # The loopback master_url from the example config never reaches the pod env + # (Lium's edge WAF 403s on loopback URLs baked into the template body). + assert "127.0.0.1" not in repr(env) + + +def test_deploy_provider_all_over_cap_provisions_nothing( + monkeypatch: pytest.MonkeyPatch, _fake_keys: None +) -> None: + provisioned = {"count": 0} + + class _FakeLiumClient: + def __init__(self, api_key: str) -> None: + pass + + async def list_offers( + self, *, max_price_per_hour: float | None = None + ) -> list[Offer]: + return [ + Offer(id="a", gpu_type="H100", gpu_count=1, price_per_hour=9.0), + ] + + async def provision(self, *a: object, **k: object) -> Instance: + provisioned["count"] += 1 + return Instance(id="never", status="PENDING") + + monkeypatch.setenv("LIUM_API_KEY", _SENTINEL_KEY) + monkeypatch.setenv("BASE_WORKER__DEPLOY__IMAGE", "ghcr.io/public/base-worker") + monkeypatch.setenv("BASE_WORKER__DEPLOY__IMAGE_DIGEST", "sha256:" + "a" * 64) + monkeypatch.setattr(cli_main, "LiumClient", _FakeLiumClient) + + result = runner.invoke( + cli_main.app, + [ + "worker", + "deploy", + "--provider", + "lium", + "--max-price", + "1.0", + "--config", + "config/worker.example.yaml", + ], + ) + assert result.exit_code == 1 + assert provisioned["count"] == 0 + assert "no rentable offer within budget" in result.stderr + + +def test_deploy_provider_without_worker_image_refuses( + monkeypatch: pytest.MonkeyPatch, _fake_keys: None +) -> None: + """A provider deploy refuses when no explicit worker image + digest is set. + + The M1 placeholder image is a PRIVATE-namespace GHCR image that fails Lium pod + creation, so the deploy must not silently pin it: an unset image is a clear, + actionable refusal that provisions nothing. + """ + + listed = {"count": 0} + + class _FakeLiumClient: + def __init__(self, api_key: str) -> None: + pass + + async def list_offers( + self, *, max_price_per_hour: float | None = None + ) -> list[Offer]: + listed["count"] += 1 + return [Offer(id="a", gpu_type="H100", gpu_count=1, price_per_hour=0.5)] + + async def provision(self, *a: object, **k: object) -> Instance: + raise AssertionError("provision must not run without a configured image") + + monkeypatch.setenv("LIUM_API_KEY", _SENTINEL_KEY) + monkeypatch.delenv("BASE_WORKER__DEPLOY__IMAGE", raising=False) + monkeypatch.delenv("BASE_WORKER__DEPLOY__IMAGE_DIGEST", raising=False) + monkeypatch.setattr(cli_main, "LiumClient", _FakeLiumClient) + + result = runner.invoke( + cli_main.app, + [ + "worker", + "deploy", + "--provider", + "lium", + "--max-price", + "1.0", + "--config", + "config/worker.example.yaml", + ], + ) + assert result.exit_code == 1 + # Refused before any provider network call (image check is fail-fast). + assert listed["count"] == 0 + assert "worker.deploy.image" in result.stderr + assert "BASE_WORKER__DEPLOY__IMAGE" in result.stderr + + +# -- VAL-AGENT-009 (unit-level): local deploy reports the active worker -------- + + +def test_deploy_local_reports_active_worker( + monkeypatch: pytest.MonkeyPatch, _fake_keys: None +) -> None: + monkeypatch.setattr( + cli_main, + "_spawn_worker_agent_process", + lambda _config: SimpleNamespace(pid=4321, terminate=lambda: None), + ) + + async def _fake_wait(*_a: object, **_k: object) -> object: + return SimpleNamespace( + worker_id="wrk-1", + worker_pubkey="worker-ss58", + miner_hotkey="miner-ss58", + provider="local", + status="active", + ) + + monkeypatch.setattr(cli_main, "_wait_worker_active", _fake_wait) + + result = runner.invoke( + cli_main.app, + [ + "worker", + "deploy", + "--provider", + "local", + "--config", + "config/worker.example.yaml", + ], + ) + assert result.exit_code == 0, result.stdout + assert "pid=4321" in result.stdout + assert "wrk-1" in result.stdout + assert "active" in result.stdout + + +# -- status renders the fleet ------------------------------------------------- + + +def test_status_renders_fleet_from_list_workers( + monkeypatch: pytest.MonkeyPatch, _fake_keys: None +) -> None: + class _FakeClient: + def __init__(self, *a: object, **k: object) -> None: + pass + + async def list_workers(self, *, hotkey: str | None = None) -> list[object]: + return [ + SimpleNamespace( + worker_id="wrk-1", + miner_hotkey="miner-ss58", + provider="lium", + status="active", + last_heartbeat_at=datetime(2026, 1, 1, tzinfo=UTC), + ) + ] + + monkeypatch.setattr(cli_main, "WorkerCoordinationClient", _FakeClient) + result = runner.invoke( + cli_main.app, ["worker", "status", "--config", "config/worker.example.yaml"] + ) + assert result.exit_code == 0, result.stdout + assert "wrk-1" in result.stdout + assert "miner-ss58" in result.stdout + assert "lium" in result.stdout + assert "active" in result.stdout + + +def test_status_surfaces_fault_count_agreeing_with_fleet_view( + monkeypatch: pytest.MonkeyPatch, _fake_keys: None +) -> None: + """The CLI fleet view shows each worker's attributed fault count (VAL-CROSS-009).""" + + class _FakeClient: + def __init__(self, *a: object, **k: object) -> None: + pass + + async def list_workers(self, *, hotkey: str | None = None) -> list[object]: + return [ + SimpleNamespace( + worker_id="clean-1", + miner_hotkey="owner-a", + provider="local", + status="active", + last_heartbeat_at=datetime(2026, 1, 1, tzinfo=UTC), + faults=[], + ), + SimpleNamespace( + worker_id="faulted-1", + miner_hotkey="owner-b", + provider="local", + status="active", + last_heartbeat_at=datetime(2026, 1, 1, tzinfo=UTC), + faults=[ + SimpleNamespace(work_unit_id="u1"), + SimpleNamespace(work_unit_id="u2"), + ], + ), + ] + + monkeypatch.setattr(cli_main, "WorkerCoordinationClient", _FakeClient) + result = runner.invoke( + cli_main.app, ["worker", "status", "--config", "config/worker.example.yaml"] + ) + assert result.exit_code == 0, result.stdout + assert "FAULTS" in result.stdout + lines = { + line.split()[0]: line for line in result.stdout.splitlines() if line.split() + } + assert " 0 " in f" {lines['clean-1'].split()[4]} " + assert lines["faulted-1"].split()[4] == "2" + + +def test_status_reports_empty_fleet( + monkeypatch: pytest.MonkeyPatch, _fake_keys: None +) -> None: + class _FakeClient: + def __init__(self, *a: object, **k: object) -> None: + pass + + async def list_workers(self, *, hotkey: str | None = None) -> list[object]: + return [] + + monkeypatch.setattr(cli_main, "WorkerCoordinationClient", _FakeClient) + result = runner.invoke( + cli_main.app, ["worker", "status", "--config", "config/worker.example.yaml"] + ) + assert result.exit_code == 0 + assert "No workers registered." in result.stdout diff --git a/tests/unit/test_worker_coordination.py b/tests/unit/test_worker_coordination.py new file mode 100644 index 00000000..7cc51984 --- /dev/null +++ b/tests/unit/test_worker_coordination.py @@ -0,0 +1,761 @@ +from __future__ import annotations + +import asyncio +import hashlib +import json +from collections.abc import AsyncIterator +from datetime import UTC, datetime, timedelta +from typing import Any + +import pytest +from httpx import ASGITransport, AsyncClient +from sqlalchemy import func, select + +from base.bittensor.metagraph_cache import MetagraphCache +from base.db import ( + Base, + WorkerFault, + WorkerRegistration, + WorkerStatus, + create_engine, + create_session_factory, + session_scope, +) +from base.master.app_proxy import create_proxy_app +from base.master.worker_coordination import ( + WorkerCoordinationService, + run_worker_health_loop, +) +from base.security.validator_auth import canonical_validator_request +from base.security.worker_auth import ( + CoordinationReadEligibility, + MetagraphMinerMembership, + SqlAlchemyWorkerNonceStore, + WorkerSignedRequestVerifier, + worker_binding_message, +) + +NOW_EPOCH = 1_750_000_000.0 +TTL_SECONDS = 120 + +# Metagraph hotkeys: H1/H2/H3 are miners (present, no permit); VAL holds a +# validator permit; STRANGER is off-graph and unregistered. +MINER_H1 = "miner-H1" +MINER_H2 = "miner-H2" +MINER_H3 = "miner-H3" +VALIDATOR = "val-permit" +STRANGER = "stranger-key" + +# The prism<->master bridge shared token the admission fleet-read also accepts. +INTERNAL_BRIDGE_TOKEN = "prism-bridge-shared-token" + + +class _TokenRegistry: + """Minimal registry exposing ``get_token`` like the real challenge registry. + + Mirrors production semantics: an unknown/unavailable slug raises + ``RuntimeError`` (exercises the token-resolver's defensive guard). + """ + + def __init__(self, tokens: dict[str, str]) -> None: + self._tokens = tokens + + def get_token(self, slug: str) -> str: + token = self._tokens.get(slug) + if not token: + raise RuntimeError(f"no token for {slug!r}") + return token + + +class FakeClock: + def __init__(self, epoch: float) -> None: + self.epoch = float(epoch) + + def time(self) -> float: + return self.epoch + + def now(self) -> datetime: + return datetime.fromtimestamp(self.epoch, UTC) + + +class FakeNonceStore: + async def reserve(self, **kwargs: Any) -> None: + return None + + +class FakeCache: + def get(self) -> dict[str, int]: + return {} + + +def _sign(hotkey: str, message: bytes) -> str: + return hashlib.sha256(hotkey.encode() + b":" + message).hexdigest() + + +def _fake_verifier(hotkey: str, message: bytes, signature: str) -> bool: + return signature == _sign(hotkey, message) + + +def _signed_headers( + *, + method: str, + path: str, + query_string: str = "", + body: bytes, + hotkey: str, + nonce: str, + timestamp: float, +) -> dict[str, str]: + ts = str(int(timestamp)) + canonical = canonical_validator_request( + method=method, + path=path, + query_string=query_string, + timestamp=ts, + nonce=nonce, + body=body, + ) + return { + "X-Hotkey": hotkey, + "X-Signature": _sign(hotkey, canonical.encode()), + "X-Nonce": nonce, + "X-Timestamp": ts, + "Content-Type": "application/json", + } + + +class Harness: + def __init__( + self, + client: AsyncClient, + session_factory: Any, + clock: FakeClock, + service: WorkerCoordinationService, + ) -> None: + self.client = client + self.session_factory = session_factory + self.clock = clock + self.service = service + self._nonce = 0 + + def _next_nonce(self, prefix: str) -> str: + self._nonce += 1 + return f"{prefix}-{self._nonce}" + + async def register( + self, + *, + worker_pubkey: str, + miner_hotkey: str, + provider: str = "local", + provider_instance_ref: str | None = "local-1", + capabilities: list[str] | None = None, + nonce: str | None = None, + signature: str | None = None, + ) -> Any: + nonce = nonce or self._next_nonce("bind") + message = worker_binding_message( + worker_pubkey=worker_pubkey, miner_hotkey=miner_hotkey, nonce=nonce + ) + signature = signature if signature is not None else _sign(miner_hotkey, message) + payload: dict[str, Any] = { + "worker_pubkey": worker_pubkey, + "miner_hotkey": miner_hotkey, + "binding_signature": signature, + "nonce": nonce, + "provider": provider, + "provider_instance_ref": provider_instance_ref, + } + if capabilities is not None: + payload["capabilities"] = capabilities + return await self.client.post("/v1/workers/register", json=payload) + + async def heartbeat(self, *, worker_id: str, signer_pubkey: str) -> Any: + body = b"{}" + path = f"/v1/workers/{worker_id}/heartbeat" + headers = _signed_headers( + method="POST", + path=path, + body=body, + hotkey=signer_pubkey, + nonce=self._next_nonce("hb"), + timestamp=self.clock.time(), + ) + return await self.client.post(path, content=body, headers=headers) + + async def list_workers(self, *, signer: str = VALIDATOR) -> Any: + path = "/v1/workers" + headers = _signed_headers( + method="GET", + path=path, + body=b"", + hotkey=signer, + nonce=self._next_nonce("list"), + timestamp=self.clock.time(), + ) + return await self.client.get(path, headers=headers) + + async def active(self, hotkey: str, *, signer: str = VALIDATOR) -> Any: + path = "/v1/workers/active" + query = f"hotkey={hotkey}" + headers = _signed_headers( + method="GET", + path=path, + query_string=query, + body=b"", + hotkey=signer, + nonce=self._next_nonce("active"), + timestamp=self.clock.time(), + ) + return await self.client.get(path, params={"hotkey": hotkey}, headers=headers) + + async def active_bearer(self, hotkey: str, *, token: str | None) -> Any: + headers = {} if token is None else {"Authorization": f"Bearer {token}"} + return await self.client.get( + "/v1/workers/active", params={"hotkey": hotkey}, headers=headers + ) + + async def list_bearer(self, *, token: str) -> Any: + return await self.client.get( + "/v1/workers", headers={"Authorization": f"Bearer {token}"} + ) + + async def worker_row(self, worker_pubkey: str) -> WorkerRegistration | None: + async with self.session_factory() as session: + return ( + await session.execute( + select(WorkerRegistration).where( + WorkerRegistration.worker_pubkey == worker_pubkey + ) + ) + ).scalar_one_or_none() + + async def count(self, worker_pubkey: str) -> int: + async with self.session_factory() as session: + return await session.scalar( + select(func.count(WorkerRegistration.id)).where( + WorkerRegistration.worker_pubkey == worker_pubkey + ) + ) + + async def set_status(self, worker_pubkey: str, status: WorkerStatus) -> None: + async with session_scope(self.session_factory) as session: + row = ( + await session.execute( + select(WorkerRegistration).where( + WorkerRegistration.worker_pubkey == worker_pubkey + ) + ) + ).scalar_one() + row.status = status + + async def add_fault(self, worker_id: str, work_unit_id: str) -> None: + async with session_scope(self.session_factory) as session: + session.add( + WorkerFault( + worker_id=worker_id, + work_unit_id=work_unit_id, + challenge_slug="prism", + detail="manifest hash divergence", + created_at=self.clock.now(), + ) + ) + + +async def _build_harness( + *, mount_worker_plane: bool = True, internal_token: str | None = None +) -> tuple[Harness, Any]: + engine = create_engine("sqlite+aiosqlite:///:memory:") + async with engine.begin() as connection: + await connection.run_sync(Base.metadata.create_all) + session_factory = create_session_factory(engine) + + cache = MetagraphCache(netuid=1, ttl_seconds=300) + cache.update_from_metagraph( + [MINER_H1, MINER_H2, MINER_H3, VALIDATOR], + validator_permits=[False, False, False, True], + stakes=[0.0, 0.0, 0.0, 100.0], + ) + clock = FakeClock(NOW_EPOCH) + service = WorkerCoordinationService( + session_factory, + miner_membership=MetagraphMinerMembership(cache), + binding_nonce_store=SqlAlchemyWorkerNonceStore(session_factory), + signature_verifier=_fake_verifier, + heartbeat_ttl_seconds=TTL_SECONDS, + now_fn=clock.now, + ) + verifier = WorkerSignedRequestVerifier( + nonce_store=SqlAlchemyWorkerNonceStore(session_factory), + eligibility=CoordinationReadEligibility(session_factory, cache), + signature_verifier=_fake_verifier, + ttl_seconds=300, + now_fn=clock.time, + ) + + registry: Any = object() + if internal_token is not None: + registry = _TokenRegistry({"prism": internal_token}) + + app = create_proxy_app( + registry=registry, + nonce_store=FakeNonceStore(), + metagraph_cache=FakeCache(), # type: ignore[arg-type] + worker_service=service if mount_worker_plane else None, + worker_verifier=verifier if mount_worker_plane else None, + ) + transport = ASGITransport(app=app) + client = AsyncClient(transport=transport, base_url="http://testserver") + return Harness(client, session_factory, clock, service), engine + + +@pytest.fixture +async def harness() -> AsyncIterator[Harness]: + h, engine = await _build_harness() + try: + yield h + finally: + await h.client.aclose() + await engine.dispose() + + +# VAL-MASTER-001 +async def test_register_pending_then_heartbeat_active(harness: Harness) -> None: + response = await harness.register(worker_pubkey="wp-1", miner_hotkey=MINER_H1) + assert response.status_code == 200 + body = response.json() + assert body["worker"]["status"] == "pending" + assert body["worker"]["miner_hotkey"] == MINER_H1 + assert body["worker"]["provider"] == "local" + assert body["heartbeat_ttl_seconds"] == TTL_SECONDS + worker_id = body["worker"]["worker_id"] + + listed = await harness.list_workers() + assert listed.status_code == 200 + entry = listed.json()["workers"][0] + assert entry["status"] == "pending" + assert entry["worker_pubkey"] == "wp-1" + assert entry["last_heartbeat_at"] is None + + harness.clock.epoch = NOW_EPOCH + 10 + hb = await harness.heartbeat(worker_id=worker_id, signer_pubkey="wp-1") + assert hb.status_code == 200 + assert hb.json()["status"] == "active" + + listed = await harness.list_workers() + entry = listed.json()["workers"][0] + assert entry["status"] == "active" + assert entry["last_heartbeat_at"] is not None + + +# VAL-AGENT-015 +async def test_fleet_view_exposes_status_owner_provider_last_seen_faults( + harness: Harness, +) -> None: + reg1 = await harness.register( + worker_pubkey="wp-a", miner_hotkey=MINER_H1, provider="lium" + ) + reg2 = await harness.register( + worker_pubkey="wp-b", miner_hotkey=MINER_H2, provider="targon" + ) + id_a = reg1.json()["worker"]["worker_id"] + id_b = reg2.json()["worker"]["worker_id"] + + await harness.heartbeat(worker_id=id_a, signer_pubkey="wp-a") + await harness.heartbeat(worker_id=id_b, signer_pubkey="wp-b") + await harness.add_fault(id_b, "unit-42") + + # wp-a heartbeats again just past nothing; wp-b goes stale past the TTL. + harness.clock.epoch = NOW_EPOCH + TTL_SECONDS + 5 + await harness.heartbeat(worker_id=id_a, signer_pubkey="wp-a") + + listed = await harness.list_workers() + assert listed.status_code == 200 + workers = {w["worker_pubkey"]: w for w in listed.json()["workers"]} + for key in ("status", "miner_hotkey", "provider", "last_heartbeat_at", "faults"): + assert key in workers["wp-a"] + assert key in workers["wp-b"] + assert workers["wp-a"]["status"] == "active" + assert workers["wp-a"]["miner_hotkey"] == MINER_H1 + assert workers["wp-a"]["provider"] == "lium" + assert workers["wp-b"]["status"] == "stale" + assert workers["wp-b"]["provider"] == "targon" + assert len(workers["wp-b"]["faults"]) == 1 + assert workers["wp-b"]["faults"][0]["work_unit_id"] == "unit-42" + assert workers["wp-a"]["faults"] == [] + + +# VAL-MASTER-002 (a): forged/invalid signature +async def test_register_invalid_signature_rejected(harness: Harness) -> None: + response = await harness.register( + worker_pubkey="wp-x", miner_hotkey=MINER_H1, signature="deadbeef" + ) + assert response.status_code == 401 + assert await harness.worker_row("wp-x") is None + + +# VAL-MASTER-002 (b): miner not in metagraph +async def test_register_miner_not_in_metagraph_rejected(harness: Harness) -> None: + response = await harness.register( + worker_pubkey="wp-y", miner_hotkey="off-graph-miner" + ) + assert response.status_code == 403 + assert await harness.worker_row("wp-y") is None + + +# VAL-MASTER-002 (c) + VAL-AGENT-004: replay rejected, fresh nonce re-enroll ok +async def test_replayed_nonce_rejected_fresh_nonce_reenrolls(harness: Harness) -> None: + first = await harness.register( + worker_pubkey="wp-z", miner_hotkey=MINER_H1, nonce="fixed-nonce" + ) + assert first.status_code == 200 + + replay = await harness.register( + worker_pubkey="wp-z2", miner_hotkey=MINER_H1, nonce="fixed-nonce" + ) + assert replay.status_code == 409 + assert await harness.worker_row("wp-z2") is None + + fresh = await harness.register( + worker_pubkey="wp-z", miner_hotkey=MINER_H1, nonce="fresh-nonce" + ) + assert fresh.status_code == 200 + assert await harness.count("wp-z") == 1 + + +# VAL-MASTER-021: cross-owner rebind never silently rebinds +async def test_rebind_to_different_hotkey_rejected(harness: Harness) -> None: + reg = await harness.register(worker_pubkey="wp-rb", miner_hotkey=MINER_H1) + worker_id = reg.json()["worker"]["worker_id"] + await harness.heartbeat(worker_id=worker_id, signer_pubkey="wp-rb") + + rebind = await harness.register(worker_pubkey="wp-rb", miner_hotkey=MINER_H2) + assert rebind.status_code == 409 + + row = await harness.worker_row("wp-rb") + assert row is not None + assert row.miner_hotkey == MINER_H1 + assert await harness.count("wp-rb") == 1 + + listed = await harness.list_workers() + entries = [w for w in listed.json()["workers"] if w["worker_pubkey"] == "wp-rb"] + assert len(entries) == 1 + assert entries[0]["miner_hotkey"] == MINER_H1 + assert entries[0]["status"] == "active" + + +# VAL-MASTER-016: active endpoint filters by hotkey + status +async def test_active_endpoint_filters_by_hotkey_and_status(harness: Harness) -> None: + # H1: one active + one stale; H2: one active; H3: only pending. + a1 = await harness.register(worker_pubkey="h1-active", miner_hotkey=MINER_H1) + a1_id = a1.json()["worker"]["worker_id"] + s1 = await harness.register(worker_pubkey="h1-stale", miner_hotkey=MINER_H1) + s1_id = s1.json()["worker"]["worker_id"] + b1 = await harness.register(worker_pubkey="h2-active", miner_hotkey=MINER_H2) + b1_id = b1.json()["worker"]["worker_id"] + await harness.register(worker_pubkey="h3-pending", miner_hotkey=MINER_H3) + + await harness.heartbeat(worker_id=s1_id, signer_pubkey="h1-stale") + # advance so h1-stale's heartbeat is now beyond the TTL, then heartbeat the + # others fresh at the later time. + harness.clock.epoch = NOW_EPOCH + TTL_SECONDS + 10 + await harness.heartbeat(worker_id=a1_id, signer_pubkey="h1-active") + await harness.heartbeat(worker_id=b1_id, signer_pubkey="h2-active") + + h1 = await harness.active(MINER_H1) + assert h1.status_code == 200 + h1_workers = [w["worker_pubkey"] for w in h1.json()["workers"]] + assert h1_workers == ["h1-active"] + + h2 = await harness.active(MINER_H2) + assert [w["worker_pubkey"] for w in h2.json()["workers"]] == ["h2-active"] + + h3 = await harness.active(MINER_H3) + assert h3.status_code == 200 + assert h3.json()["workers"] == [] + + unknown = await harness.active("nobody") + assert unknown.status_code == 200 + assert unknown.json()["workers"] == [] + + # After H1's active worker also passes the TTL, the query is empty. + harness.clock.epoch = NOW_EPOCH + 2 * (TTL_SECONDS + 10) + 10 + h1_later = await harness.active(MINER_H1) + assert h1_later.json()["workers"] == [] + + +# VAL-MASTER-018: retired is terminal +async def test_retired_worker_is_terminal(harness: Harness) -> None: + reg = await harness.register(worker_pubkey="wp-ret", miner_hotkey=MINER_H1) + worker_id = reg.json()["worker"]["worker_id"] + await harness.heartbeat(worker_id=worker_id, signer_pubkey="wp-ret") + await harness.set_status("wp-ret", WorkerStatus.RETIRED) + + # (a)/(b) excluded from active listing even as sole worker of its hotkey. + active = await harness.active(MINER_H1) + assert active.json()["workers"] == [] + + # still listed with status retired in the fleet view. + listed = await harness.list_workers() + entry = [w for w in listed.json()["workers"] if w["worker_pubkey"] == "wp-ret"][0] + assert entry["status"] == "retired" + + # (c) a heartbeat never resurrects a retired worker. + harness.clock.epoch = NOW_EPOCH + 5 + hb = await harness.heartbeat(worker_id=worker_id, signer_pubkey="wp-ret") + assert hb.status_code == 200 + assert hb.json()["status"] == "retired" + row = await harness.worker_row("wp-ret") + assert row is not None + assert WorkerStatus(row.status) == WorkerStatus.RETIRED + assert (await harness.active(MINER_H1)).json()["workers"] == [] + + +async def test_heartbeat_unknown_worker_returns_404(harness: Harness) -> None: + # A registered worker signs, but targets an unknown worker_id. + await harness.register(worker_pubkey="wp-known", miner_hotkey=MINER_H1) + hb = await harness.heartbeat(worker_id="does-not-exist", signer_pubkey="wp-known") + assert hb.status_code == 404 + + +async def test_heartbeat_foreign_owner_returns_403(harness: Harness) -> None: + reg = await harness.register(worker_pubkey="wp-owner", miner_hotkey=MINER_H1) + other = await harness.register(worker_pubkey="wp-other", miner_hotkey=MINER_H2) + owner_id = reg.json()["worker"]["worker_id"] + # wp-other is a registered (eligible) signer but does not own owner_id. + hb = await harness.heartbeat(worker_id=owner_id, signer_pubkey="wp-other") + assert hb.status_code == 403 + assert other.status_code == 200 + + +# Fleet reads are authenticated-but-not-admin: worker OR validator identities. +async def test_fleet_read_auth_accepts_worker_and_validator_rejects_stranger( + harness: Harness, +) -> None: + await harness.register(worker_pubkey="wp-auth", miner_hotkey=MINER_H1) + + by_validator = await harness.list_workers(signer=VALIDATOR) + assert by_validator.status_code == 200 + + by_worker = await harness.list_workers(signer="wp-auth") + assert by_worker.status_code == 200 + + by_stranger = await harness.list_workers(signer=STRANGER) + assert by_stranger.status_code == 403 + + +async def test_fleet_read_rejects_unsigned_request(harness: Harness) -> None: + unsigned = await harness.client.get("/v1/workers") + assert unsigned.status_code == 401 + + +# --------------------------------------------------------------------------- +# Admission fleet-read dual-auth: internal bridge bearer OR signed request. +# (master-admission-fleetread-auth; supports VAL-CROSS-004 / VAL-MASTER-016) +# --------------------------------------------------------------------------- + + +async def _active_worker(h: Harness, *, worker_pubkey: str, miner_hotkey: str) -> None: + reg = await h.register(worker_pubkey=worker_pubkey, miner_hotkey=miner_hotkey) + worker_id = reg.json()["worker"]["worker_id"] + h.clock.epoch = NOW_EPOCH + 10 + await h.heartbeat(worker_id=worker_id, signer_pubkey=worker_pubkey) + + +# (a) internal-bearer request returns that hotkey's ACTIVE workers. +async def test_active_endpoint_accepts_internal_bridge_bearer() -> None: + h, engine = await _build_harness(internal_token=INTERNAL_BRIDGE_TOKEN) + try: + await _active_worker(h, worker_pubkey="wp-a", miner_hotkey=MINER_H1) + await _active_worker(h, worker_pubkey="wp-other", miner_hotkey=MINER_H2) + + resp = await h.active_bearer(MINER_H1, token=INTERNAL_BRIDGE_TOKEN) + assert resp.status_code == 200 + assert [w["worker_pubkey"] for w in resp.json()["workers"]] == ["wp-a"] + + # An unknown / no-worker hotkey is a 2xx empty list, not an error. + empty = await h.active_bearer("nobody", token=INTERNAL_BRIDGE_TOKEN) + assert empty.status_code == 200 + assert empty.json()["workers"] == [] + finally: + await h.client.aclose() + await engine.dispose() + + +# (b) missing/wrong bearer AND no valid signature => 401/403. +async def test_active_endpoint_rejects_bad_bearer_without_signature() -> None: + h, engine = await _build_harness(internal_token=INTERNAL_BRIDGE_TOKEN) + try: + await _active_worker(h, worker_pubkey="wp-b", miner_hotkey=MINER_H1) + + wrong = await h.active_bearer(MINER_H1, token="not-the-token") + assert wrong.status_code in (401, 403) + + missing = await h.active_bearer(MINER_H1, token=None) + assert missing.status_code in (401, 403) + finally: + await h.client.aclose() + await engine.dispose() + + +# (c) signed-request path still works unchanged when internal auth is configured. +async def test_active_endpoint_signed_request_unchanged_with_internal_auth() -> None: + h, engine = await _build_harness(internal_token=INTERNAL_BRIDGE_TOKEN) + try: + await _active_worker(h, worker_pubkey="wp-c", miner_hotkey=MINER_H1) + + by_validator = await h.active(MINER_H1, signer=VALIDATOR) + assert by_validator.status_code == 200 + assert [w["worker_pubkey"] for w in by_validator.json()["workers"]] == ["wp-c"] + + by_worker = await h.active(MINER_H1, signer="wp-c") + assert by_worker.status_code == 200 + + by_stranger = await h.active(MINER_H1, signer=STRANGER) + assert by_stranger.status_code == 403 + finally: + await h.client.aclose() + await engine.dispose() + + +# (d) full fleet view still rejects the internal bearer (signed-request only). +async def test_full_fleet_endpoint_rejects_internal_bearer() -> None: + h, engine = await _build_harness(internal_token=INTERNAL_BRIDGE_TOKEN) + try: + await _active_worker(h, worker_pubkey="wp-d", miner_hotkey=MINER_H1) + + by_bearer = await h.list_bearer(token=INTERNAL_BRIDGE_TOKEN) + assert by_bearer.status_code in (401, 403) + + # sanity: the signed-request path to the full fleet still works. + signed = await h.list_workers(signer=VALIDATOR) + assert signed.status_code == 200 + finally: + await h.client.aclose() + await engine.dispose() + + +# With no bridge token configured the internal bearer is inert (signed-only). +async def test_active_endpoint_bearer_inert_without_configured_token() -> None: + h, engine = await _build_harness(internal_token=None) + try: + await _active_worker(h, worker_pubkey="wp-e", miner_hotkey=MINER_H1) + + bearer = await h.active_bearer(MINER_H1, token=INTERNAL_BRIDGE_TOKEN) + assert bearer.status_code in (401, 403) + + signed = await h.active(MINER_H1, signer=VALIDATOR) + assert signed.status_code == 200 + finally: + await h.client.aclose() + await engine.dispose() + + +# Flag OFF (worker plane not mounted) => surface is inert (404), legacy safe. +async def test_worker_surface_absent_when_plane_disabled() -> None: + h, engine = await _build_harness(mount_worker_plane=False) + try: + listed = await h.client.get("/v1/workers") + assert listed.status_code == 404 + health = await h.client.get("/health") + assert health.status_code == 200 + finally: + await h.client.aclose() + await engine.dispose() + + +# --------------------------------------------------------------------------- +# Service-level staleness detection +# --------------------------------------------------------------------------- + + +async def _service_only() -> tuple[WorkerCoordinationService, Any, FakeClock, Any]: + engine = create_engine("sqlite+aiosqlite:///:memory:") + async with engine.begin() as connection: + await connection.run_sync(Base.metadata.create_all) + session_factory = create_session_factory(engine) + cache = MetagraphCache(netuid=1, ttl_seconds=300) + cache.update_from_metagraph([MINER_H1], validator_permits=[False], stakes=[0.0]) + clock = FakeClock(NOW_EPOCH) + service = WorkerCoordinationService( + session_factory, + miner_membership=MetagraphMinerMembership(cache), + binding_nonce_store=SqlAlchemyWorkerNonceStore(session_factory), + signature_verifier=_fake_verifier, + heartbeat_ttl_seconds=TTL_SECONDS, + now_fn=clock.now, + ) + return service, session_factory, clock, engine + + +async def test_detect_stale_workers_transitions_active_to_stale() -> None: + service, session_factory, clock, engine = await _service_only() + try: + worker = await service.register( + worker_pubkey="wp-s", + miner_hotkey=MINER_H1, + binding_signature=_sign( + MINER_H1, + worker_binding_message( + worker_pubkey="wp-s", miner_hotkey=MINER_H1, nonce="n1" + ), + ), + nonce="n1", + provider="local", + provider_instance_ref=None, + capabilities=["gpu"], + ) + await service.heartbeat(worker_id=worker.worker_id, worker_pubkey="wp-s") + + # within TTL: no transition + clock.epoch = NOW_EPOCH + TTL_SECONDS + assert await service.detect_stale_workers() == [] + + # beyond TTL: transitions to stale (edge-triggered, once) + clock.epoch = NOW_EPOCH + TTL_SECONDS + 1 + assert await service.detect_stale_workers() == [worker.worker_id] + clock.epoch = NOW_EPOCH + TTL_SECONDS + 500 + assert await service.detect_stale_workers() == [] + + async with session_factory() as session: + row = ( + await session.execute( + select(WorkerRegistration).where( + WorkerRegistration.worker_pubkey == "wp-s" + ) + ) + ).scalar_one() + assert WorkerStatus(row.status) == WorkerStatus.STALE + finally: + await engine.dispose() + + +async def test_run_worker_health_loop_repeats_until_shutdown() -> None: + shutdown = asyncio.Event() + + class _Recording: + def __init__(self) -> None: + self.calls = 0 + + async def detect_stale_workers(self) -> list[str]: + self.calls += 1 + if self.calls >= 3: + shutdown.set() + return [] + + service = _Recording() + await asyncio.wait_for( + run_worker_health_loop( + service, # type: ignore[arg-type] + interval_seconds=0.001, + shutdown_event=shutdown, + ), + timeout=2, + ) + assert service.calls >= 3 + + +def test_json_body_roundtrip_is_deterministic() -> None: + # guard: the register body encoding used by the harness is stable JSON. + assert json.loads(json.dumps({"a": 1})) == {"a": 1} + assert timedelta(seconds=TTL_SECONDS).total_seconds() == TTL_SECONDS diff --git a/tests/unit/test_worker_deploy.py b/tests/unit/test_worker_deploy.py new file mode 100644 index 00000000..7b1d9edc --- /dev/null +++ b/tests/unit/test_worker_deploy.py @@ -0,0 +1,314 @@ +"""Offline unit tests for worker deploy planning (VAL-AGENT-010/011/012). + +Pure, network-free coverage of :mod:`base.worker.deploy`: provider-key +enforcement, in-budget offer selection with the GPU-count preference + fallback, +the miner-signed binding, and the pod env's provider-key hygiene. +""" + +from __future__ import annotations + +import ast +from dataclasses import dataclass +from pathlib import Path + +import pytest + +from base.compute.provider import Offer +from base.compute.worker_deployment import is_metachar_free +from base.worker.deploy import ( + PROVIDER_KEY_ENV, + MissingProviderKeyError, + NoOfferWithinBudgetError, + UnsupportedProviderError, + WorkerImageNotConfiguredError, + build_signed_binding, + build_worker_pod_env, + normalize_provider, + plan_provider_deployment, + rank_worker_offers, + require_provider_api_key, + require_worker_image, + select_worker_offer, +) + +_REPO_ROOT = Path(__file__).resolve().parents[2] +_DOCKERFILE = _REPO_ROOT / "docker" / "Dockerfile.worker" + + +def _offer(offer_id: str, *, price: float, gpu_count: int = 1) -> Offer: + return Offer( + id=offer_id, + gpu_type="H100", + gpu_count=gpu_count, + price_per_hour=price, + provider="lium", + ) + + +@dataclass +class _FakeSigner: + hotkey: str + + def sign(self, message: bytes) -> str: + return "0x" + message.hex() + + +class _FakeProviderClient: + def __init__(self, offers: list[Offer]) -> None: + self._offers = offers + self.list_calls: list[float | None] = [] + self.provision_calls = 0 + + async def list_offers( + self, *, max_price_per_hour: float | None = None + ) -> list[Offer]: + self.list_calls.append(max_price_per_hour) + return [ + offer + for offer in self._offers + if max_price_per_hour is None or offer.price_per_hour <= max_price_per_hour + ] + + async def provision(self, *_: object, **__: object) -> None: + self.provision_calls += 1 + + +# -- provider normalization + key enforcement (VAL-AGENT-010) ----------------- + + +def test_normalize_provider_accepts_supported() -> None: + assert normalize_provider("Lium") == "lium" + assert normalize_provider(" TARGON ") == "targon" + assert normalize_provider("local") == "local" + + +def test_normalize_provider_rejects_unknown() -> None: + with pytest.raises(UnsupportedProviderError): + normalize_provider("aws") + + +@pytest.mark.parametrize("provider", ["lium", "targon"]) +def test_require_provider_api_key_raises_when_missing(provider: str) -> None: + with pytest.raises(MissingProviderKeyError) as exc: + require_provider_api_key(provider, environ={}) + assert exc.value.env_var == PROVIDER_KEY_ENV[provider] + assert exc.value.provider == provider + assert PROVIDER_KEY_ENV[provider] in str(exc.value) + + +@pytest.mark.parametrize("provider", ["lium", "targon"]) +def test_require_provider_api_key_raises_when_blank(provider: str) -> None: + with pytest.raises(MissingProviderKeyError): + require_provider_api_key(provider, environ={PROVIDER_KEY_ENV[provider]: " "}) + + +def test_require_provider_api_key_returns_value() -> None: + assert require_provider_api_key("lium", environ={"LIUM_API_KEY": "k"}) == "k" + + +# -- offer selection (VAL-AGENT-012 + live gpu-count constraint) -------------- + + +def test_rank_worker_offers_filters_by_max_price() -> None: + offers = [_offer("a", price=0.5), _offer("b", price=2.0), _offer("c", price=1.0)] + ranked = rank_worker_offers(offers, gpu_count=1, max_price=1.0) + assert [o.id for o in ranked] == ["a", "c"] + assert all(o.price_per_hour <= 1.0 for o in ranked) + + +def test_rank_worker_offers_prefers_matching_gpu_count_over_cheaper() -> None: + # A cheaper 8-GPU node would 400 on a gpu_count=1 rent, so the exact + # single-GPU node is preferred even though it is pricier. + offers = [ + _offer("multi", price=0.5, gpu_count=8), + _offer("single", price=1.0, gpu_count=1), + ] + ranked = rank_worker_offers(offers, gpu_count=1, max_price=1.5) + assert ranked[0].id == "single" + assert ranked[1].id == "multi" # next-cheapest fallback retained in order + + +def test_select_worker_offer_returns_best_in_budget() -> None: + offers = [_offer("a", price=0.9), _offer("b", price=0.3)] + assert select_worker_offer(offers, gpu_count=1, max_price=1.0).id == "b" + + +def test_select_worker_offer_raises_when_all_over_cap() -> None: + offers = [_offer("a", price=2.0), _offer("b", price=3.0)] + with pytest.raises(NoOfferWithinBudgetError): + select_worker_offer(offers, gpu_count=1, max_price=1.0) + + +def test_select_worker_offer_raises_when_empty() -> None: + with pytest.raises(NoOfferWithinBudgetError): + select_worker_offer([], gpu_count=1, max_price=1.0) + + +async def test_plan_provider_deployment_selects_in_budget_offer() -> None: + client = _FakeProviderClient( + [_offer("cheap", price=0.5), _offer("pricey", price=5.0)] + ) + offer = await plan_provider_deployment(client, gpu_count=1, max_price=1.0) + assert offer.id == "cheap" + assert client.provision_calls == 0 # planning never rents + + +async def test_plan_provider_deployment_all_over_cap_provisions_nothing() -> None: + client = _FakeProviderClient([_offer("a", price=2.0), _offer("b", price=3.0)]) + with pytest.raises(NoOfferWithinBudgetError): + await plan_provider_deployment(client, gpu_count=1, max_price=1.0) + assert client.provision_calls == 0 + + +# -- binding + pod env hygiene (VAL-AGENT-011) -------------------------------- + + +def test_build_signed_binding_signs_pinned_message() -> None: + binding = build_signed_binding( + worker_pubkey="worker-pk", + miner_signer=_FakeSigner("miner-hk"), + nonce="n1", + ) + assert binding.miner_hotkey == "miner-hk" + assert binding.nonce == "n1" + expected = b"worker-binding:worker-pk:miner-hk:n1".hex() + assert binding.signature == "0x" + expected + + +def test_build_signed_binding_generates_fresh_nonce() -> None: + first = build_signed_binding(worker_pubkey="wp", miner_signer=_FakeSigner("mh")) + second = build_signed_binding(worker_pubkey="wp", miner_signer=_FakeSigner("mh")) + assert first.nonce != second.nonce + + +def test_build_worker_pod_env_never_carries_provider_key() -> None: + binding = build_signed_binding( + worker_pubkey="wp", miner_signer=_FakeSigner("mh"), nonce="n1" + ) + env = build_worker_pod_env( + master_url="http://master:3100", + provider="lium", + binding=binding, + worker_key_uri="//Worker", + broker_url="http://127.0.0.1:8082", + extra={ + "LIUM_API_KEY": "SENTINEL-KEY-12345", + "TARGON_API_KEY": "SENTINEL-TARGON", + "SAFE_VALUE": "ok", + }, + ) + blob = repr(env) + assert "SENTINEL-KEY-12345" not in blob + assert "SENTINEL-TARGON" not in blob + assert "LIUM_API_KEY" not in env + assert "TARGON_API_KEY" not in env + assert env["BASE_WORKER__AGENT__MASTER_URL"] == "http://master:3100" + assert env["BASE_WORKER__IDENTITY__BINDING_SIGNATURE"] == binding.signature + assert env["BASE_WORKER__IDENTITY__MINER_HOTKEY"] == "mh" + assert env["SAFE_VALUE"] == "ok" + + +# -- loopback master_url hygiene (Lium edge-WAF 403; VAL follow-up) ------------ + + +def test_build_worker_pod_env_omits_loopback_urls() -> None: + # Lium's edge WAF 403s on any request body carrying a loopback URL, and this + # env is baked into the WAF-sensitive POST /templates body. A loopback + # coordination URL is also redundant (the pod config defaults to loopback and + # the agent resolves it at runtime), so it must not travel in the env. + binding = build_signed_binding( + worker_pubkey="wp", miner_signer=_FakeSigner("mh"), nonce="n1" + ) + env = build_worker_pod_env( + master_url="http://127.0.0.1:8081", + provider="lium", + binding=binding, + broker_url="http://127.0.0.1:8082", + gateway_url="http://localhost:8081", + ) + blob = repr(env) + assert "127.0.0.1" not in blob + assert "localhost" not in blob + assert "BASE_WORKER__AGENT__MASTER_URL" not in env + assert "BASE_WORKER__AGENT__BROKER_URL" not in env + assert "BASE_WORKER__AGENT__GATEWAY_URL" not in env + # The binding + provider still travel (they are not loopback URLs). + assert env["BASE_WORKER__IDENTITY__MINER_HOTKEY"] == "mh" + assert env["BASE_WORKER__DEPLOY__PROVIDER"] == "lium" + + +def test_build_worker_pod_env_keeps_public_urls() -> None: + binding = build_signed_binding( + worker_pubkey="wp", miner_signer=_FakeSigner("mh"), nonce="n1" + ) + env = build_worker_pod_env( + master_url="https://master.example.com", + provider="lium", + binding=binding, + broker_url="http://broker.internal:8082", + gateway_url="https://gateway.example.com", + ) + assert env["BASE_WORKER__AGENT__MASTER_URL"] == "https://master.example.com" + assert env["BASE_WORKER__AGENT__BROKER_URL"] == "http://broker.internal:8082" + assert env["BASE_WORKER__AGENT__GATEWAY_URL"] == "https://gateway.example.com" + + +# -- worker image config (no silent private-image pin; VAL follow-up) ---------- + + +def test_require_worker_image_returns_configured_image_and_digest() -> None: + digest = "sha256:" + "a" * 64 + assert require_worker_image( + image="ghcr.io/public/base-worker", image_digest=digest, provider="lium" + ) == ("ghcr.io/public/base-worker", digest) + + +@pytest.mark.parametrize("provider", ["lium", "targon"]) +def test_require_worker_image_raises_when_unset(provider: str) -> None: + with pytest.raises(WorkerImageNotConfiguredError) as exc: + require_worker_image(image=None, image_digest=None, provider=provider) + message = str(exc.value) + assert "worker.deploy.image" in message + assert "BASE_WORKER__DEPLOY__IMAGE" in message + + +def test_require_worker_image_requires_both_image_and_digest() -> None: + digest = "sha256:" + "a" * 64 + with pytest.raises(WorkerImageNotConfiguredError): + require_worker_image( + image="ghcr.io/public/base-worker", image_digest=None, provider="lium" + ) + with pytest.raises(WorkerImageNotConfiguredError): + require_worker_image(image=None, image_digest=digest, provider="lium") + + +def test_require_worker_image_rejects_malformed_digest() -> None: + with pytest.raises(WorkerImageNotConfiguredError): + require_worker_image( + image="ghcr.io/public/base-worker", + image_digest="latest", + provider="lium", + ) + + +# -- worker image entrypoint (VAL-AGENT-013/014, live metachar constraint) ----- + + +def _dockerfile_cmd_tokens() -> list[str]: + for line in _DOCKERFILE.read_text(encoding="utf-8").splitlines(): + stripped = line.strip() + if stripped.startswith("CMD ["): + tokens = ast.literal_eval(stripped[len("CMD ") :]) + assert isinstance(tokens, list) + return [str(token) for token in tokens] + raise AssertionError("Dockerfile.worker has no exec-form CMD") + + +def test_worker_dockerfile_entrypoint_starts_the_agent() -> None: + tokens = _dockerfile_cmd_tokens() + assert tokens[:3] == ["base", "worker", "agent"] + + +def test_worker_dockerfile_entrypoint_is_metachar_free() -> None: + for token in _dockerfile_cmd_tokens(): + assert is_metachar_free(token), token diff --git a/tests/unit/test_worker_proof.py b/tests/unit/test_worker_proof.py new file mode 100644 index 00000000..eef16598 --- /dev/null +++ b/tests/unit/test_worker_proof.py @@ -0,0 +1,106 @@ +"""ExecutionProof construction/verification tests (VAL-AGENT-008). + +Uses REAL sr25519 keypairs (bittensor) so the pinned signature format is +exercised end-to-end: the signature is over ``sha256(f"{manifest}:{unit_id}")`` +and verifies against the worker pubkey, but not when replayed onto another unit. +""" + +from __future__ import annotations + +import hashlib + +from base.security.miner_auth import verify_substrate_signature +from base.validator.agent.signing import KeypairRequestSigner +from base.worker.proof import ( + EXECUTION_PROOF_VERSION, + build_execution_proof, + execution_proof_signing_payload, + verify_execution_proof, +) + +MANIFEST = "a" * 64 +UNIT_ID = "submission-123" + + +def _signer(uri: str = "//WorkerAlice") -> KeypairRequestSigner: + import bittensor as bt + + return KeypairRequestSigner(bt.Keypair.create_from_uri(uri)) + + +def test_signing_payload_is_sha256_of_pinned_string() -> None: + payload = execution_proof_signing_payload(manifest_sha256=MANIFEST, unit_id=UNIT_ID) + assert payload == hashlib.sha256(f"{MANIFEST}:{UNIT_ID}".encode()).digest() + + +def test_build_execution_proof_is_tier0_with_worker_signature() -> None: + signer = _signer() + proof = build_execution_proof( + signer=signer, manifest_sha256=MANIFEST, unit_id=UNIT_ID + ) + assert proof.version == EXECUTION_PROOF_VERSION == 1 + assert proof.tier == 0 + assert proof.manifest_sha256 == MANIFEST + assert proof.attestation is None + assert proof.worker_signature.worker_pubkey == signer.hotkey + + +def test_worker_signature_verifies_over_pinned_message() -> None: + signer = _signer() + proof = build_execution_proof( + signer=signer, manifest_sha256=MANIFEST, unit_id=UNIT_ID + ) + assert verify_execution_proof(proof, unit_id=UNIT_ID) is True + + # Independent verification against the pubkey using the shared sr25519 helper. + payload = execution_proof_signing_payload(manifest_sha256=MANIFEST, unit_id=UNIT_ID) + assert verify_substrate_signature( + proof.worker_signature.worker_pubkey, payload, proof.worker_signature.sig + ) + + +def test_proof_cannot_be_replayed_across_units() -> None: + signer = _signer() + proof = build_execution_proof( + signer=signer, manifest_sha256=MANIFEST, unit_id=UNIT_ID + ) + assert verify_execution_proof(proof, unit_id="different-unit") is False + + +def test_tampered_manifest_hash_fails_verification() -> None: + signer = _signer() + proof = build_execution_proof( + signer=signer, manifest_sha256=MANIFEST, unit_id=UNIT_ID + ) + proof.manifest_sha256 = "b" * 64 + assert verify_execution_proof(proof, unit_id=UNIT_ID) is False + + +def test_corrupted_signature_fails_verification() -> None: + signer = _signer() + proof = build_execution_proof( + signer=signer, manifest_sha256=MANIFEST, unit_id=UNIT_ID + ) + proof.worker_signature.sig = "0x" + "00" * 64 + assert verify_execution_proof(proof, unit_id=UNIT_ID) is False + + +def test_tier1_when_image_digest_and_pod_present() -> None: + from base.schemas.worker import ProviderInfo + + signer = _signer() + proof = build_execution_proof( + signer=signer, + manifest_sha256=MANIFEST, + unit_id=UNIT_ID, + tier=1, + image_digest="sha256:" + "c" * 64, + provider=ProviderInfo( + name="lium", executor_id="ex-1", pod_id="pod-1", miner_hotkey="miner-H1" + ), + ) + assert proof.tier == 1 + assert proof.image_digest == "sha256:" + "c" * 64 + assert proof.provider is not None + assert proof.provider.pod_id == "pod-1" + assert verify_execution_proof(proof, unit_id=UNIT_ID) is True diff --git a/tests/unit/test_worker_reconciliation.py b/tests/unit/test_worker_reconciliation.py new file mode 100644 index 00000000..585408bb --- /dev/null +++ b/tests/unit/test_worker_reconciliation.py @@ -0,0 +1,659 @@ +"""Worker-plane reconciliation, disputes, and audit-fault attribution. + +Covers the master-side reconciliation the reconciliation-and-disputes feature +adds on top of the worker replica plane (``WorkerAssignmentEngine`` + +``WorkerAssignmentService.post_result``): + +* VAL-MASTER-008: matching manifest hashes => exactly one result forwarded to the + challenge and the unit reaches ``completed``. +* VAL-MASTER-009: divergent manifest hashes => the unit is ``disputed``, NOTHING + is forwarded (before or after audit), and a validator-executor audit unit is + created (assignable to validators, never to workers). +* VAL-MASTER-010: the validator audit outcome writes ``worker_faults`` for the + divergent worker(s) only, visible in the fleet view, without mutating status. +* VAL-MASTER-017: a single surviving proof (the other replica exhausted) resolves + deterministically (accepted, forwarded once, degradation warning), and a mere + ``success=false`` replica never triggers a dispute. +* VAL-MASTER-019: late/foreign result posts are rejected and reconciliation uses + only legitimately-owned replica results. +""" + +from __future__ import annotations + +import logging +from collections.abc import AsyncIterator, Mapping +from dataclasses import dataclass +from datetime import UTC, datetime, timedelta +from typing import Any + +import pytest +from sqlalchemy import select + +from base.bittensor.metagraph_cache import MetagraphCache +from base.db import ( + Base, + Validator, + ValidatorStatus, + WorkAssignmentStatus, + WorkerAssignment, + WorkerFault, + WorkerRegistration, + WorkerStatus, + create_engine, + create_session_factory, + session_scope, +) +from base.db.models import WorkAssignment +from base.master.assignment import ( + CAPABILITY_GPU, + EXECUTOR_KIND_PAYLOAD_KEY, + EXECUTOR_KIND_VALIDATOR, + AssignmentService, +) +from base.master.assignment_coordination import AssignmentCoordinationService +from base.master.worker_assignment import ( + WorkerAssignmentOwnershipError, + WorkerAssignmentService, +) +from base.master.worker_assignment_engine import WorkerAssignmentEngine +from base.master.worker_coordination import WorkerCoordinationService +from base.master.worker_reconciliation import ( + AUDIT_OF_PAYLOAD_KEY, + RECONCILE_DEGRADED_PAYLOAD_KEY, + WorkerReconciliationService, + audit_work_unit_id, +) +from base.security.worker_auth import ( + MetagraphMinerMembership, + SqlAlchemyWorkerNonceStore, +) +from base.worker.proof import MANIFEST_SHA256_PAYLOAD_KEY, PROOF_PAYLOAD_KEY + +NOW = datetime(2026, 6, 28, 12, 0, 0, tzinfo=UTC) +TTL = 120 + +MINER_H = "miner-H" +MINER_A = "miner-A" +MINER_B = "miner-B" +MINER_C = "miner-C" + +HASH_A = "a" * 64 +HASH_B = "b" * 64 + +GPU_VALIDATOR = "gpu-val" + + +class FakeClock: + def __init__(self, moment: datetime) -> None: + self.moment = moment + + def now(self) -> datetime: + return self.moment + + +class _FakeForwarder: + """Records every challenge result forward so tests can count them.""" + + def __init__(self) -> None: + self.calls: list[dict[str, Any]] = [] + + async def forward_result( + self, + *, + challenge_slug: str, + work_unit_id: str, + submission_ref: str, + result_payload: Mapping[str, Any], + ) -> None: + self.calls.append( + { + "challenge_slug": challenge_slug, + "work_unit_id": work_unit_id, + "submission_ref": submission_ref, + "result_payload": result_payload, + } + ) + + +@dataclass +class Env: + db_engine: Any + factory: Any + clock: FakeClock + worker_service: WorkerCoordinationService + worker_assignment_service: WorkerAssignmentService + engine: WorkerAssignmentEngine + forwarder: _FakeForwarder + reconciler: WorkerReconciliationService + validator_plane: AssignmentService + validator_coordination: AssignmentCoordinationService + + +async def _build_env() -> Env: + db_engine = create_engine("sqlite+aiosqlite:///:memory:") + async with db_engine.begin() as connection: + await connection.run_sync(Base.metadata.create_all) + factory = create_session_factory(db_engine) + clock = FakeClock(NOW) + cache = MetagraphCache(netuid=1, ttl_seconds=300) + worker_service = WorkerCoordinationService( + factory, + miner_membership=MetagraphMinerMembership(cache), + binding_nonce_store=SqlAlchemyWorkerNonceStore(factory), + heartbeat_ttl_seconds=TTL, + now_fn=clock.now, + ) + worker_assignment_service = WorkerAssignmentService( + factory, worker_service=worker_service, now_fn=clock.now + ) + engine = WorkerAssignmentEngine( + factory, + assignment_service=worker_assignment_service, + worker_service=worker_service, + replication_factor=2, + now_fn=clock.now, + ) + forwarder = _FakeForwarder() + reconciler = WorkerReconciliationService( + factory, result_forwarder=forwarder, now_fn=clock.now + ) + validator_plane = AssignmentService( + factory, + now_fn=clock.now, + worker_plane_capabilities=frozenset({CAPABILITY_GPU}), + ) + validator_coordination = AssignmentCoordinationService(factory, now_fn=clock.now) + return Env( + db_engine=db_engine, + factory=factory, + clock=clock, + worker_service=worker_service, + worker_assignment_service=worker_assignment_service, + engine=engine, + forwarder=forwarder, + reconciler=reconciler, + validator_plane=validator_plane, + validator_coordination=validator_coordination, + ) + + +async def _add_worker( + factory: Any, + *, + worker_pubkey: str, + miner_hotkey: str, + status: WorkerStatus = WorkerStatus.ACTIVE, + last_heartbeat_at: datetime | None = NOW, +) -> str: + worker_id = f"wid-{worker_pubkey}" + async with session_scope(factory) as session: + session.add( + WorkerRegistration( + worker_id=worker_id, + worker_pubkey=worker_pubkey, + miner_hotkey=miner_hotkey, + binding_signature="sig", + binding_nonce=f"nonce-{worker_pubkey}", + provider="local", + provider_instance_ref="local-1", + capabilities=["gpu"], + status=status, + last_heartbeat_at=last_heartbeat_at, + created_at=NOW, + updated_at=NOW, + ) + ) + return worker_id + + +async def _add_gpu_unit( + factory: Any, + *, + work_unit_id: str, + submitter: str, + max_attempts: int = 3, +) -> None: + async with session_scope(factory) as session: + session.add( + WorkAssignment( + challenge_slug="prism", + work_unit_id=work_unit_id, + submission_ref=submitter, + payload={"run_spec": {"image": "img"}}, + required_capability="gpu", + status=WorkAssignmentStatus.PENDING, + attempt_count=0, + max_attempts=max_attempts, + created_at=NOW, + updated_at=NOW, + ) + ) + + +async def _add_gpu_validator(factory: Any, *, hotkey: str = GPU_VALIDATOR) -> None: + async with session_scope(factory) as session: + session.add( + Validator( + hotkey=hotkey, + uid=None, + status=ValidatorStatus.ONLINE, + capabilities=["gpu"], + version="1.0.0", + registered_at=NOW, + last_heartbeat_at=NOW, + ) + ) + + +def _proof_payload(manifest: str) -> dict[str, Any]: + return { + PROOF_PAYLOAD_KEY: { + "version": 1, + "tier": 0, + MANIFEST_SHA256_PAYLOAD_KEY: manifest, + }, + MANIFEST_SHA256_PAYLOAD_KEY: manifest, + } + + +async def _replicas(factory: Any, work_unit_id: str) -> list[WorkerAssignment]: + async with factory() as session: + rows = ( + ( + await session.execute( + select(WorkerAssignment) + .where(WorkerAssignment.work_unit_id == work_unit_id) + .order_by(WorkerAssignment.miner_hotkey) + ) + ) + .scalars() + .all() + ) + return list(rows) + + +async def _unit(factory: Any, work_unit_id: str) -> WorkAssignment: + async with factory() as session: + return ( + await session.execute( + select(WorkAssignment).where( + WorkAssignment.work_unit_id == work_unit_id + ) + ) + ).scalar_one() + + +async def _maybe_unit(factory: Any, work_unit_id: str) -> WorkAssignment | None: + async with factory() as session: + return ( + await session.execute( + select(WorkAssignment).where( + WorkAssignment.work_unit_id == work_unit_id + ) + ) + ).scalar_one_or_none() + + +async def _post_replica( + env: Env, + *, + work_unit_id: str, + miner_hotkey: str, + manifest: str | None = None, + success: bool = True, +) -> None: + replica = next( + r + for r in await _replicas(env.factory, work_unit_id) + if r.miner_hotkey == miner_hotkey + ) + payload = _proof_payload(manifest) if manifest is not None else {"error": "boom"} + await env.worker_assignment_service.post_result( + assignment_id=str(replica.id), + worker_pubkey=replica.worker_pubkey, + success=success, + payload=payload, + ) + + +@pytest.fixture +async def env() -> AsyncIterator[Env]: + e = await _build_env() + try: + yield e + finally: + await e.db_engine.dispose() + + +# VAL-MASTER-008 +async def test_matching_hashes_forward_exactly_one(env: Env) -> None: + await _add_worker(env.factory, worker_pubkey="wp-a", miner_hotkey=MINER_A) + await _add_worker(env.factory, worker_pubkey="wp-b", miner_hotkey=MINER_B) + await _add_gpu_unit(env.factory, work_unit_id="U", submitter=MINER_H) + await env.engine.assign_pending(seed=1) + + await _post_replica(env, work_unit_id="U", miner_hotkey=MINER_A, manifest=HASH_A) + await _post_replica(env, work_unit_id="U", miner_hotkey=MINER_B, manifest=HASH_A) + + result = await env.reconciler.reconcile_once() + # Reconciling again must not forward a second time. + await env.reconciler.reconcile_once() + + assert result.accepted == ["U"] + assert result.disputed == [] + assert len(env.forwarder.calls) == 1 + assert env.forwarder.calls[0]["work_unit_id"] == "U" + unit = await _unit(env.factory, "U") + assert unit.status == WorkAssignmentStatus.COMPLETED + assert await _maybe_unit(env.factory, audit_work_unit_id("U")) is None + + +# VAL-MASTER-009 +async def test_divergent_hashes_dispute_and_audit_unit(env: Env) -> None: + await _add_worker(env.factory, worker_pubkey="wp-a", miner_hotkey=MINER_A) + await _add_worker(env.factory, worker_pubkey="wp-b", miner_hotkey=MINER_B) + await _add_gpu_unit(env.factory, work_unit_id="U", submitter=MINER_H) + await env.engine.assign_pending(seed=1) + + await _post_replica(env, work_unit_id="U", miner_hotkey=MINER_A, manifest=HASH_A) + await _post_replica(env, work_unit_id="U", miner_hotkey=MINER_B, manifest=HASH_B) + + result = await env.reconciler.reconcile_once() + + assert result.disputed == ["U"] + assert result.accepted == [] + assert env.forwarder.calls == [] + unit = await _unit(env.factory, "U") + assert unit.status == WorkAssignmentStatus.DISPUTED + + audit = await _unit(env.factory, audit_work_unit_id("U")) + assert audit.required_capability == CAPABILITY_GPU + assert audit.payload[EXECUTOR_KIND_PAYLOAD_KEY] == EXECUTOR_KIND_VALIDATOR + assert audit.payload[AUDIT_OF_PAYLOAD_KEY] == "U" + assert audit.status == WorkAssignmentStatus.PENDING + + +# VAL-MASTER-009: the audit unit routes to a validator and never to a worker. +async def test_audit_unit_routes_to_validator_not_worker(env: Env) -> None: + await _add_worker(env.factory, worker_pubkey="wp-a", miner_hotkey=MINER_A) + await _add_worker(env.factory, worker_pubkey="wp-b", miner_hotkey=MINER_B) + await _add_gpu_unit(env.factory, work_unit_id="U", submitter=MINER_H) + await env.engine.assign_pending(seed=1) + await _post_replica(env, work_unit_id="U", miner_hotkey=MINER_A, manifest=HASH_A) + await _post_replica(env, work_unit_id="U", miner_hotkey=MINER_B, manifest=HASH_B) + await env.reconciler.reconcile_once() + + audit_id = audit_work_unit_id("U") + + # The worker engine never picks up the validator-executor audit unit. + await env.engine.assign_pending(seed=1) + assert await _replicas(env.factory, audit_id) == [] + + # The validator plane assigns it to an online gpu validator. + await _add_gpu_validator(env.factory) + await env.validator_plane.assign_pending(seed=1) + audit = await _unit(env.factory, audit_id) + assert audit.assigned_validator_hotkey == GPU_VALIDATOR + assert audit.status == WorkAssignmentStatus.ASSIGNED + + +# VAL-MASTER-010 +async def test_audit_outcome_faults_divergent_worker_only(env: Env) -> None: + await _add_worker(env.factory, worker_pubkey="wp-a", miner_hotkey=MINER_A) + await _add_worker(env.factory, worker_pubkey="wp-b", miner_hotkey=MINER_B) + await _add_gpu_unit(env.factory, work_unit_id="U", submitter=MINER_H) + await env.engine.assign_pending(seed=1) + # A agrees with the validator (HASH_A); B diverges (HASH_B). + await _post_replica(env, work_unit_id="U", miner_hotkey=MINER_A, manifest=HASH_A) + await _post_replica(env, work_unit_id="U", miner_hotkey=MINER_B, manifest=HASH_B) + await env.reconciler.reconcile_once() + + # A validator replays the audit unit; its authoritative manifest is HASH_A. + await _add_gpu_validator(env.factory) + await env.validator_plane.assign_pending(seed=1) + audit = await _unit(env.factory, audit_work_unit_id("U")) + await env.validator_coordination.post_result( + assignment_id=str(audit.id), + hotkey=GPU_VALIDATOR, + success=True, + payload=_proof_payload(HASH_A), + ) + + result = await env.reconciler.reconcile_once() + + # Exactly the divergent worker (B) is faulted; A is not. + faulted_workers = {pair.split(":")[-1] for pair in result.faults} + assert faulted_workers == {"wid-wp-b"} + + async with env.factory() as session: + faults = (await session.execute(select(WorkerFault))).scalars().all() + assert len(faults) == 1 + assert faults[0].worker_id == "wid-wp-b" + assert faults[0].work_unit_id == "U" + + faults_by_worker = await env.worker_service.faults_by_worker() + assert "wid-wp-b" in faults_by_worker + assert "wid-wp-a" not in faults_by_worker + + # The divergent worker's status is unchanged by fault recording. + async with env.factory() as session: + worker_b = ( + await session.execute( + select(WorkerRegistration).where( + WorkerRegistration.worker_id == "wid-wp-b" + ) + ) + ).scalar_one() + assert WorkerStatus(worker_b.status) == WorkerStatus.ACTIVE + + # The disputed unit is never forwarded, before or after audit resolution. + assert env.forwarder.calls == [] + assert (await _unit(env.factory, "U")).status == WorkAssignmentStatus.DISPUTED + + +# VAL-MASTER-010: faults are recorded once, not re-written each pass. +async def test_audit_faults_are_idempotent(env: Env) -> None: + await _add_worker(env.factory, worker_pubkey="wp-a", miner_hotkey=MINER_A) + await _add_worker(env.factory, worker_pubkey="wp-b", miner_hotkey=MINER_B) + await _add_gpu_unit(env.factory, work_unit_id="U", submitter=MINER_H) + await env.engine.assign_pending(seed=1) + await _post_replica(env, work_unit_id="U", miner_hotkey=MINER_A, manifest=HASH_A) + await _post_replica(env, work_unit_id="U", miner_hotkey=MINER_B, manifest=HASH_B) + await env.reconciler.reconcile_once() + await _add_gpu_validator(env.factory) + await env.validator_plane.assign_pending(seed=1) + audit = await _unit(env.factory, audit_work_unit_id("U")) + await env.validator_coordination.post_result( + assignment_id=str(audit.id), + hotkey=GPU_VALIDATOR, + success=True, + payload=_proof_payload(HASH_A), + ) + + await env.reconciler.reconcile_once() + await env.reconciler.reconcile_once() + + async with env.factory() as session: + faults = (await session.execute(select(WorkerFault))).scalars().all() + assert len(faults) == 1 + + +# VAL-MASTER-017 +async def test_single_surviving_proof_accepts_with_warning( + env: Env, caplog: pytest.LogCaptureFixture +) -> None: + await _add_worker(env.factory, worker_pubkey="wp-a", miner_hotkey=MINER_A) + await _add_worker(env.factory, worker_pubkey="wp-b", miner_hotkey=MINER_B) + await _add_gpu_unit( + env.factory, work_unit_id="U", submitter=MINER_H, max_attempts=1 + ) + await env.engine.assign_pending(seed=1) + + # A posts a valid proof; B's worker goes stale and its retries are exhausted. + await _post_replica(env, work_unit_id="U", miner_hotkey=MINER_A, manifest=HASH_A) + async with session_scope(env.factory) as session: + worker_b = ( + await session.execute( + select(WorkerRegistration).where( + WorkerRegistration.miner_hotkey == MINER_B + ) + ) + ).scalar_one() + worker_b.status = WorkerStatus.STALE + worker_b.last_heartbeat_at = NOW - timedelta(seconds=TTL + 60) + # B's replica lapses with no other eligible distinct owner -> FAILED. + await env.engine.reassign_stale_replicas(seed=1) + + with caplog.at_level(logging.WARNING, logger="base.master.worker_reconciliation"): + result = await env.reconciler.reconcile_once() + + assert result.accepted == ["U"] + assert result.disputed == [] + assert result.single_replica == ["U"] + assert len(env.forwarder.calls) == 1 + unit = await _unit(env.factory, "U") + assert unit.status == WorkAssignmentStatus.COMPLETED + assert unit.payload[RECONCILE_DEGRADED_PAYLOAD_KEY] is True + assert await _maybe_unit(env.factory, audit_work_unit_id("U")) is None + assert any("U" in record.getMessage() for record in caplog.records) + + +# VAL-MASTER-017: a bare success=false replica never triggers a dispute. +async def test_failed_replica_alone_never_disputes(env: Env) -> None: + await _add_worker(env.factory, worker_pubkey="wp-a", miner_hotkey=MINER_A) + await _add_worker(env.factory, worker_pubkey="wp-b", miner_hotkey=MINER_B) + await _add_gpu_unit(env.factory, work_unit_id="U", submitter=MINER_H) + await env.engine.assign_pending(seed=1) + + await _post_replica(env, work_unit_id="U", miner_hotkey=MINER_A, manifest=HASH_A) + # B reports an execution failure with no proof. + await _post_replica( + env, work_unit_id="U", miner_hotkey=MINER_B, manifest=None, success=False + ) + + result = await env.reconciler.reconcile_once() + + assert result.disputed == [] + assert result.accepted == ["U"] + assert len(env.forwarder.calls) == 1 + assert (await _unit(env.factory, "U")).status == WorkAssignmentStatus.COMPLETED + assert await _maybe_unit(env.factory, audit_work_unit_id("U")) is None + + +# VAL-MASTER-017: reconciliation waits while a replica may still report. +async def test_reconciliation_waits_for_in_flight_replica(env: Env) -> None: + await _add_worker(env.factory, worker_pubkey="wp-a", miner_hotkey=MINER_A) + await _add_worker(env.factory, worker_pubkey="wp-b", miner_hotkey=MINER_B) + await _add_gpu_unit(env.factory, work_unit_id="U", submitter=MINER_H) + await env.engine.assign_pending(seed=1) + + await _post_replica(env, work_unit_id="U", miner_hotkey=MINER_A, manifest=HASH_A) + # B has not reported and is still in flight. + + result = await env.reconciler.reconcile_once() + + assert result.accepted == [] + assert result.disputed == [] + assert env.forwarder.calls == [] + assert (await _unit(env.factory, "U")).status == WorkAssignmentStatus.ASSIGNED + + +# VAL-MASTER-019 +async def test_late_and_foreign_posts_rejected(env: Env) -> None: + await _add_worker(env.factory, worker_pubkey="wp-a", miner_hotkey=MINER_A) + await _add_worker(env.factory, worker_pubkey="wp-b", miner_hotkey=MINER_B) + await _add_worker(env.factory, worker_pubkey="wp-c", miner_hotkey=MINER_C) + await _add_gpu_unit(env.factory, work_unit_id="U", submitter=MINER_H) + await env.engine.assign_pending(seed=1) + + # Capture W1 = A's original replica id, then lapse it and reassign to C. + a_replica = next( + r for r in await _replicas(env.factory, "U") if r.miner_hotkey == MINER_A + ) + a_replica_id = str(a_replica.id) + async with session_scope(env.factory) as session: + row = ( + await session.execute( + select(WorkerAssignment).where(WorkerAssignment.id == a_replica.id) + ) + ).scalar_one() + row.deadline_at = NOW - timedelta(seconds=10) + # Make A's worker stale so the replica is reclaimed to a fresh owner. + worker_a = ( + await session.execute( + select(WorkerRegistration).where( + WorkerRegistration.miner_hotkey == MINER_A + ) + ) + ).scalar_one() + worker_a.status = WorkerStatus.STALE + worker_a.last_heartbeat_at = NOW - timedelta(seconds=TTL + 60) + await env.engine.reassign_stale_replicas(seed=1) + + # The replica is now owned by C (W2); A's old post is foreign now. + reassigned = next( + r for r in await _replicas(env.factory, "U") if str(r.id) == a_replica_id + ) + assert reassigned.miner_hotkey == MINER_C + + # W1 (A) posting for its OLD replica is rejected on ownership. + with pytest.raises(WorkerAssignmentOwnershipError): + await env.worker_assignment_service.post_result( + assignment_id=a_replica_id, + worker_pubkey="wp-a", + success=True, + payload=_proof_payload(HASH_B), + ) + # A foreign worker (B) posting for C's replica is rejected too. + with pytest.raises(WorkerAssignmentOwnershipError): + await env.worker_assignment_service.post_result( + assignment_id=a_replica_id, + worker_pubkey="wp-b", + success=True, + payload=_proof_payload(HASH_B), + ) + + # The replica state is intact: still owned by C, still running, no manifest. + reassigned = next( + r for r in await _replicas(env.factory, "U") if str(r.id) == a_replica_id + ) + assert reassigned.miner_hotkey == MINER_C + assert reassigned.manifest_sha256 is None + assert WorkAssignmentStatus(reassigned.status) == WorkAssignmentStatus.ASSIGNED + + # Both legitimately-owned replicas now post equal hashes -> clean accept. + await _post_replica(env, work_unit_id="U", miner_hotkey=MINER_B, manifest=HASH_A) + await _post_replica(env, work_unit_id="U", miner_hotkey=MINER_C, manifest=HASH_A) + result = await env.reconciler.reconcile_once() + + assert result.accepted == ["U"] + assert len(env.forwarder.calls) == 1 + assert (await _unit(env.factory, "U")).status == WorkAssignmentStatus.COMPLETED + + +# VAL-MASTER-019: a late post to an already-reconciled replica is a no-op. +async def test_late_post_to_reconciled_unit_is_noop(env: Env) -> None: + await _add_worker(env.factory, worker_pubkey="wp-a", miner_hotkey=MINER_A) + await _add_worker(env.factory, worker_pubkey="wp-b", miner_hotkey=MINER_B) + await _add_gpu_unit(env.factory, work_unit_id="U", submitter=MINER_H) + await env.engine.assign_pending(seed=1) + await _post_replica(env, work_unit_id="U", miner_hotkey=MINER_A, manifest=HASH_A) + await _post_replica(env, work_unit_id="U", miner_hotkey=MINER_B, manifest=HASH_A) + await env.reconciler.reconcile_once() + + a_replica = next( + r for r in await _replicas(env.factory, "U") if r.miner_hotkey == MINER_A + ) + outcome = await env.worker_assignment_service.post_result( + assignment_id=str(a_replica.id), + worker_pubkey="wp-a", + success=True, + payload=_proof_payload(HASH_B), + ) + assert outcome.idempotent is True + # The stored manifest is unchanged and no extra forward happened. + a_replica = next( + r for r in await _replicas(env.factory, "U") if r.miner_hotkey == MINER_A + ) + assert a_replica.manifest_sha256 == HASH_A + assert len(env.forwarder.calls) == 1 diff --git a/tests/unit/test_worker_unit_status.py b/tests/unit/test_worker_unit_status.py new file mode 100644 index 00000000..6331af80 --- /dev/null +++ b/tests/unit/test_worker_unit_status.py @@ -0,0 +1,634 @@ +"""Worker-plane unit-status read endpoint (master-unit-status-read-endpoint). + +Makes the dispute -> audit -> invalidation -> fault chain OPERATOR-DISCOVERABLE +VIA API ALONE (VAL-CROSS-011). ``GET /v1/workers/units`` is a signed read-only +master surface (auth: ``CoordinationReadEligibility`` -- same as +``GET /v1/workers``, NOT the internal bridge bearer) that exposes, per primary +gpu unit: + +* the unit id + status (INCLUDING ``disputed``); +* its replicas (worker_id, owner miner hotkey, posted manifest_sha256, proof + presence); +* for a disputed unit, the linked validator AUDIT unit's id, executor kind + (``validator``), and terminal outcome (pending/passed/mismatch-resolved). + +Flag OFF => the router is unmounted (404). Signed-request auth is enforced +(bad/missing signature => 401/403) and the internal bridge bearer is NOT accepted +here (that acceptor is only on the narrow admission fleet-read). +""" + +from __future__ import annotations + +import hashlib +from collections.abc import AsyncIterator +from datetime import UTC, datetime +from typing import Any + +import pytest +from httpx import ASGITransport, AsyncClient + +from base.bittensor.metagraph_cache import MetagraphCache +from base.db import ( + Base, + WorkAssignment, + WorkAssignmentStatus, + WorkerAssignment, + WorkerFault, + create_engine, + create_session_factory, + session_scope, +) +from base.master.app_proxy import create_proxy_app +from base.master.assignment import ( + EXECUTOR_KIND_PAYLOAD_KEY, + EXECUTOR_KIND_VALIDATOR, +) +from base.master.worker_reconciliation import ( + AUDIT_OF_PAYLOAD_KEY, + AUDIT_RESOLVED_PAYLOAD_KEY, + audit_work_unit_id, +) +from base.master.worker_unit_status import ( + AUDIT_OUTCOME_MISMATCH_RESOLVED, + AUDIT_OUTCOME_PASSED, + AUDIT_OUTCOME_PENDING, + WorkerUnitStatusService, +) +from base.security.validator_auth import canonical_validator_request +from base.security.worker_auth import ( + CoordinationReadEligibility, + SqlAlchemyWorkerNonceStore, + WorkerSignedRequestVerifier, +) +from base.worker.proof import MANIFEST_SHA256_PAYLOAD_KEY, PROOF_PAYLOAD_KEY + +NOW = datetime(2026, 7, 7, 12, 0, 0, tzinfo=UTC) +NOW_EPOCH = NOW.timestamp() +TTL = 120 + +MINER_A = "miner-A" +MINER_B = "miner-B" +MINER_H = "miner-H" +VALIDATOR = "val-permit" +STRANGER = "stranger-key" + +HASH_A = "a" * 64 +HASH_B = "b" * 64 +INTERNAL_BRIDGE_TOKEN = "prism-bridge-shared-token" + + +def _sign(hotkey: str, message: bytes) -> str: + return hashlib.sha256(hotkey.encode() + b":" + message).hexdigest() + + +def _fake_verifier(hotkey: str, message: bytes, signature: str) -> bool: + return signature == _sign(hotkey, message) + + +class _TokenRegistry: + def __init__(self, token: str) -> None: + self._token = token + + def get_token(self, slug: str) -> str: + if slug == "prism": + return self._token + raise RuntimeError(f"no token for {slug!r}") + + +class FakeNonceStore: + async def reserve(self, **_kwargs: Any) -> None: + return None + + +class FakeCache: + def get(self) -> dict[str, int]: + return {} + + +def _proof_payload(manifest: str) -> dict[str, Any]: + return { + PROOF_PAYLOAD_KEY: { + "version": 1, + "tier": 0, + MANIFEST_SHA256_PAYLOAD_KEY: manifest, + }, + MANIFEST_SHA256_PAYLOAD_KEY: manifest, + } + + +async def _add_primary( + factory: Any, + *, + work_unit_id: str, + submitter: str, + status: WorkAssignmentStatus, + challenge_slug: str = "prism", +) -> None: + async with session_scope(factory) as session: + session.add( + WorkAssignment( + challenge_slug=challenge_slug, + work_unit_id=work_unit_id, + submission_ref=submitter, + payload={"run_spec": {"image": "img"}}, + required_capability="gpu", + status=status, + attempt_count=1, + max_attempts=3, + created_at=NOW, + updated_at=NOW, + ) + ) + + +async def _add_audit_unit( + factory: Any, + *, + original_id: str, + status: WorkAssignmentStatus, + resolved: bool, + challenge_slug: str = "prism", +) -> None: + payload: dict[str, Any] = { + EXECUTOR_KIND_PAYLOAD_KEY: EXECUTOR_KIND_VALIDATOR, + AUDIT_OF_PAYLOAD_KEY: original_id, + } + if resolved: + payload[AUDIT_RESOLVED_PAYLOAD_KEY] = True + async with session_scope(factory) as session: + session.add( + WorkAssignment( + challenge_slug=challenge_slug, + work_unit_id=audit_work_unit_id(original_id), + submission_ref=MINER_H, + payload=payload, + required_capability="gpu", + status=status, + attempt_count=0, + max_attempts=3, + created_at=NOW, + updated_at=NOW, + ) + ) + + +async def _add_replica( + factory: Any, + *, + work_unit_id: str, + worker_id: str, + miner_hotkey: str, + manifest: str | None, + status: WorkAssignmentStatus = WorkAssignmentStatus.COMPLETED, + challenge_slug: str = "prism", +) -> None: + payload = _proof_payload(manifest) if manifest is not None else None + async with session_scope(factory) as session: + session.add( + WorkerAssignment( + challenge_slug=challenge_slug, + work_unit_id=work_unit_id, + submission_ref=MINER_H, + worker_id=worker_id, + worker_pubkey=f"wp-{worker_id}", + miner_hotkey=miner_hotkey, + payload={}, + required_capability="gpu", + status=status, + attempt_count=1, + max_attempts=3, + result_success=manifest is not None, + result_payload=payload, + manifest_sha256=manifest, + created_at=NOW, + updated_at=NOW, + ) + ) + + +async def _add_fault( + factory: Any, + *, + worker_id: str, + work_unit_id: str, + challenge_slug: str = "prism", +) -> None: + async with session_scope(factory) as session: + session.add( + WorkerFault( + worker_id=worker_id, + work_unit_id=work_unit_id, + challenge_slug=challenge_slug, + detail="manifest diverged from validator audit", + created_at=NOW, + ) + ) + + +async def _build_app( + *, mount: bool = True, internal_token: str | None = None +) -> tuple[AsyncClient, Any, WorkerUnitStatusService, Any]: + engine = create_engine("sqlite+aiosqlite:///:memory:") + async with engine.begin() as connection: + await connection.run_sync(Base.metadata.create_all) + factory = create_session_factory(engine) + + cache = MetagraphCache(netuid=1, ttl_seconds=300) + cache.update_from_metagraph( + [MINER_A, MINER_B, MINER_H, VALIDATOR], + validator_permits=[False, False, False, True], + stakes=[0.0, 0.0, 0.0, 100.0], + ) + verifier = WorkerSignedRequestVerifier( + nonce_store=SqlAlchemyWorkerNonceStore(factory), + eligibility=CoordinationReadEligibility(factory, cache), + signature_verifier=_fake_verifier, + ttl_seconds=300, + now_fn=lambda: NOW_EPOCH, + ) + service = WorkerUnitStatusService(factory) + + registry: Any = object() + if internal_token is not None: + registry = _TokenRegistry(internal_token) + + app = create_proxy_app( + registry=registry, + nonce_store=FakeNonceStore(), + metagraph_cache=FakeCache(), # type: ignore[arg-type] + worker_verifier=verifier if mount else None, + worker_unit_status_service=service if mount else None, + ) + transport = ASGITransport(app=app) + client = AsyncClient(transport=transport, base_url="http://testserver") + return client, factory, service, engine + + +def _signed_headers( + *, hotkey: str, nonce: str, sign_hotkey: str | None = None +) -> dict[str, str]: + ts = str(int(NOW_EPOCH)) + canonical = canonical_validator_request( + method="GET", + path="/v1/workers/units", + query_string="", + timestamp=ts, + nonce=nonce, + body=b"", + ) + signer = sign_hotkey if sign_hotkey is not None else hotkey + return { + "X-Hotkey": hotkey, + "X-Signature": _sign(signer, canonical.encode()), + "X-Nonce": nonce, + "X-Timestamp": ts, + } + + +@pytest.fixture +async def env() -> AsyncIterator[tuple[AsyncClient, Any, WorkerUnitStatusService]]: + client, factory, service, engine = await _build_app() + try: + yield client, factory, service + finally: + await client.aclose() + await engine.dispose() + + +# --- service-level assertions -------------------------------------------------- + + +async def test_matched_unit_shows_completed_with_both_replica_hashes( + env: tuple[AsyncClient, Any, WorkerUnitStatusService], +) -> None: + _, factory, service = env + await _add_primary( + factory, + work_unit_id="U-ok", + submitter=MINER_H, + status=WorkAssignmentStatus.COMPLETED, + ) + await _add_replica( + factory, + work_unit_id="U-ok", + worker_id="wa", + miner_hotkey=MINER_A, + manifest=HASH_A, + ) + await _add_replica( + factory, + work_unit_id="U-ok", + worker_id="wb", + miner_hotkey=MINER_B, + manifest=HASH_A, + ) + + units = await service.list_units() + assert len(units) == 1 + unit = units[0] + assert unit.work_unit_id == "U-ok" + assert unit.status == WorkAssignmentStatus.COMPLETED.value + assert unit.audit is None + assert {r.miner_hotkey for r in unit.replicas} == {MINER_A, MINER_B} + assert [r.manifest_sha256 for r in unit.replicas] == [HASH_A, HASH_A] + assert all(r.has_proof for r in unit.replicas) + assert {r.worker_id for r in unit.replicas} == {"wa", "wb"} + + +async def test_disputed_unit_shows_audit_pending_before_resolution( + env: tuple[AsyncClient, Any, WorkerUnitStatusService], +) -> None: + _, factory, service = env + await _add_primary( + factory, + work_unit_id="U-dis", + submitter=MINER_H, + status=WorkAssignmentStatus.DISPUTED, + ) + await _add_replica( + factory, + work_unit_id="U-dis", + worker_id="wa", + miner_hotkey=MINER_A, + manifest=HASH_A, + ) + await _add_replica( + factory, + work_unit_id="U-dis", + worker_id="wb", + miner_hotkey=MINER_B, + manifest=HASH_B, + ) + await _add_audit_unit( + factory, + original_id="U-dis", + status=WorkAssignmentStatus.PENDING, + resolved=False, + ) + + units = await service.list_units() + # The validator audit unit is nested, never a top-level primary entry. + assert [u.work_unit_id for u in units] == ["U-dis"] + unit = units[0] + assert unit.status == WorkAssignmentStatus.DISPUTED.value + assert unit.audit is not None + assert unit.audit.work_unit_id == audit_work_unit_id("U-dis") + assert unit.audit.executor_kind == EXECUTOR_KIND_VALIDATOR + assert unit.audit.outcome == AUDIT_OUTCOME_PENDING + assert {r.manifest_sha256 for r in unit.replicas} == {HASH_A, HASH_B} + + +async def test_disputed_unit_audit_mismatch_resolved_when_fault_recorded( + env: tuple[AsyncClient, Any, WorkerUnitStatusService], +) -> None: + _, factory, service = env + await _add_primary( + factory, + work_unit_id="U-dis", + submitter=MINER_H, + status=WorkAssignmentStatus.DISPUTED, + ) + await _add_replica( + factory, + work_unit_id="U-dis", + worker_id="wa", + miner_hotkey=MINER_A, + manifest=HASH_A, + ) + await _add_replica( + factory, + work_unit_id="U-dis", + worker_id="wb", + miner_hotkey=MINER_B, + manifest=HASH_B, + ) + await _add_audit_unit( + factory, + original_id="U-dis", + status=WorkAssignmentStatus.COMPLETED, + resolved=True, + ) + await _add_fault(factory, worker_id="wb", work_unit_id="U-dis") + + units = await service.list_units() + unit = units[0] + assert unit.audit is not None + assert unit.audit.outcome == AUDIT_OUTCOME_MISMATCH_RESOLVED + + +async def test_disputed_unit_audit_passed_when_resolved_without_fault( + env: tuple[AsyncClient, Any, WorkerUnitStatusService], +) -> None: + _, factory, service = env + await _add_primary( + factory, + work_unit_id="U-dis", + submitter=MINER_H, + status=WorkAssignmentStatus.DISPUTED, + ) + await _add_replica( + factory, + work_unit_id="U-dis", + worker_id="wa", + miner_hotkey=MINER_A, + manifest=HASH_A, + ) + await _add_replica( + factory, + work_unit_id="U-dis", + worker_id="wb", + miner_hotkey=MINER_B, + manifest=HASH_B, + ) + await _add_audit_unit( + factory, + original_id="U-dis", + status=WorkAssignmentStatus.COMPLETED, + resolved=True, + ) + + units = await service.list_units() + unit = units[0] + assert unit.audit is not None + assert unit.audit.outcome == AUDIT_OUTCOME_PASSED + + +async def test_fault_keyed_by_challenge_and_unit_does_not_collide_across_slugs( + env: tuple[AsyncClient, Any, WorkerUnitStatusService], +) -> None: + """Two disputed units sharing a work_unit_id under DIFFERENT challenge slugs + must not collide: a fault recorded for one slug's unit must not flip the + other slug's (fault-free) audit to ``mismatch-resolved``.""" + + _, factory, service = env + shared_unit_id = "U-shared" + for slug in ("prism", "other-challenge"): + await _add_primary( + factory, + work_unit_id=shared_unit_id, + submitter=MINER_H, + status=WorkAssignmentStatus.DISPUTED, + challenge_slug=slug, + ) + await _add_replica( + factory, + work_unit_id=shared_unit_id, + worker_id=f"wa-{slug}", + miner_hotkey=MINER_A, + manifest=HASH_A, + challenge_slug=slug, + ) + await _add_replica( + factory, + work_unit_id=shared_unit_id, + worker_id=f"wb-{slug}", + miner_hotkey=MINER_B, + manifest=HASH_B, + challenge_slug=slug, + ) + await _add_audit_unit( + factory, + original_id=shared_unit_id, + status=WorkAssignmentStatus.COMPLETED, + resolved=True, + challenge_slug=slug, + ) + # Only the prism unit's diverging worker is faulted. + await _add_fault( + factory, + worker_id="wb-prism", + work_unit_id=shared_unit_id, + challenge_slug="prism", + ) + + units = await service.list_units() + by_slug = {u.challenge_slug: u for u in units} + assert by_slug["prism"].audit is not None + assert by_slug["prism"].audit.outcome == AUDIT_OUTCOME_MISMATCH_RESOLVED + assert by_slug["other-challenge"].audit is not None + assert by_slug["other-challenge"].audit.outcome == AUDIT_OUTCOME_PASSED + + +async def test_replica_without_proof_reports_absent_proof( + env: tuple[AsyncClient, Any, WorkerUnitStatusService], +) -> None: + _, factory, service = env + await _add_primary( + factory, + work_unit_id="U-mix", + submitter=MINER_H, + status=WorkAssignmentStatus.ASSIGNED, + ) + await _add_replica( + factory, + work_unit_id="U-mix", + worker_id="wa", + miner_hotkey=MINER_A, + manifest=HASH_A, + ) + await _add_replica( + factory, + work_unit_id="U-mix", + worker_id="wb", + miner_hotkey=MINER_B, + manifest=None, + status=WorkAssignmentStatus.RUNNING, + ) + + unit = (await service.list_units())[0] + by_worker = {r.worker_id: r for r in unit.replicas} + assert by_worker["wa"].has_proof is True + assert by_worker["wa"].manifest_sha256 == HASH_A + assert by_worker["wb"].has_proof is False + assert by_worker["wb"].manifest_sha256 is None + + +# --- HTTP surface + auth ------------------------------------------------------- + + +async def test_signed_request_returns_units( + env: tuple[AsyncClient, Any, WorkerUnitStatusService], +) -> None: + client, factory, _ = env + await _add_primary( + factory, + work_unit_id="U-ok", + submitter=MINER_H, + status=WorkAssignmentStatus.COMPLETED, + ) + await _add_replica( + factory, + work_unit_id="U-ok", + worker_id="wa", + miner_hotkey=MINER_A, + manifest=HASH_A, + ) + + resp = await client.get( + "/v1/workers/units", + headers=_signed_headers(hotkey=VALIDATOR, nonce="units-1"), + ) + assert resp.status_code == 200 + body = resp.json() + assert [u["work_unit_id"] for u in body["units"]] == ["U-ok"] + assert body["units"][0]["status"] == "completed" + + +async def test_missing_signature_rejected( + env: tuple[AsyncClient, Any, WorkerUnitStatusService], +) -> None: + client, _, _ = env + resp = await client.get("/v1/workers/units") + assert resp.status_code in (401, 403) + + +async def test_bad_signature_rejected( + env: tuple[AsyncClient, Any, WorkerUnitStatusService], +) -> None: + client, _, _ = env + resp = await client.get( + "/v1/workers/units", + headers=_signed_headers( + hotkey=VALIDATOR, nonce="units-bad", sign_hotkey=STRANGER + ), + ) + assert resp.status_code in (401, 403) + + +async def test_unregistered_identity_rejected( + env: tuple[AsyncClient, Any, WorkerUnitStatusService], +) -> None: + client, _, _ = env + resp = await client.get( + "/v1/workers/units", + headers=_signed_headers(hotkey=STRANGER, nonce="units-str"), + ) + assert resp.status_code in (401, 403) + + +async def test_internal_bridge_bearer_not_accepted() -> None: + client, factory, _service, engine = await _build_app( + internal_token=INTERNAL_BRIDGE_TOKEN + ) + try: + resp = await client.get( + "/v1/workers/units", + headers={"Authorization": f"Bearer {INTERNAL_BRIDGE_TOKEN}"}, + ) + assert resp.status_code in (401, 403) + finally: + await client.aclose() + await engine.dispose() + + +async def test_flag_off_router_unmounted_404() -> None: + client, _factory, _service, engine = await _build_app(mount=False) + try: + resp = await client.get( + "/v1/workers/units", + headers=_signed_headers(hotkey=VALIDATOR, nonce="off-1"), + ) + assert resp.status_code == 404 + finally: + await client.aclose() + await engine.dispose()