# gpu-broker-mcp

A stateless MCP server that brokers GPU compute access for AI agents. Agents
discover nodes, reserve capacity, dispatch inference, and poll results through
four MCP tools — without managing SSH keys, node IPs, or provider APIs
directly.

**SDK**: `mcp==2.0.0` (Python SDK v2, `mcp.server.MCPServer`)
**Target spec**: MCP specification revision 2026-07-28
**Transport**: Streamable HTTP, stateless mode (`stateless_http=True`,
`json_response=True`). No sessions, no `Mcp-Session-Id`, no sticky routing.

## Architecture

```
┌─────────────────────────────────────────────────────────────┐
│  Agent (MCP client)                                         │
│  Calls: list_nodes → reserve_node → dispatch_inference      │
│         → get_result (poll)                                 │
└────────────────────────┬────────────────────────────────────┘
                         │ JSON-RPC over Streamable HTTP
                         │ (stateless, any replica)
┌────────────────────────▼────────────────────────────────────┐
│  gpu-broker-mcp server                                      │
│                                                             │
│  ┌──────────────────┐  ┌──────────────────┐                 │
│  │ HMAC-SHA256       │  │ NodePool ABC     │                 │
│  │ Handle signing    │  │  ├ FakeNodePool   │                │
│  │ & validation      │  │  └ VastNodePool   │                │
│  └──────────────────┘  └──────────────────┘                 │
│                                                             │
│  ┌──────────────────┐  ┌──────────────────┐                 │
│  │ Error taxonomy    │  │ JobStore ABC     │                 │
│  │ (single enum,     │  │  └ InMemoryStore  │                │
│  │  structured JSON) │  │    (per-replica)  │                │
│  └──────────────────┘  └──────────────────┘                 │
└────────────────────────┬────────────────────────────────────┘
                         │ SSH (VastNodePool only)
┌────────────────────────▼────────────────────────────────────┐
│  GPU node (e.g. Vast.ai RTX 3090)                           │
│  Runs inference workload, returns stdout                    │
└─────────────────────────────────────────────────────────────┘
```

The broker runs locally. It is a client of GPU nodes, not resident on them —
it does CPU-bound HMAC signing and JSON serialization, nothing that benefits
from a GPU.

## Why signed handles instead of sessions

Reservation state lives inside the handle itself: a base64-encoded JSON
payload (node ID, expiry, scope) concatenated with its HMAC-SHA256 signature.
The secret comes from `GPU_BROKER_SECRET` and the server refuses to boot if
it's unset.

This means any replica sharing the secret can validate a handle it never
issued. There is no session table, no `Mcp-Session-Id` header, and no sticky
routing requirement. A load balancer can route any request to any replica.
Handles are scoped (`reserve` vs `task`) so a reservation handle cannot be
replayed as a task ID or vice versa — misuse returns `HANDLE_SCOPE_INVALID`.

What the in-memory `JobStore` does lose across replicas is job status lookup:
replica B cannot tell you the status of a job dispatched to replica A. This is
a shared-backend requirement (Redis, Postgres) rather than a flaw in the
stateless design. The signature validation — the security-critical part — is
fully portable.

## Tools

| Tool | Params | Returns |
|---|---|---|
| `list_nodes` | — | JSON array of available nodes (id, model, vram, price, load) |
| `reserve_node` | `node_id`, `ttl_seconds` | Signed reservation handle |
| `dispatch_inference` | `handle`, `payload` | `{"task_id": "...", "status": "pending"}` |
| `get_result` | `task_id` | `{"status": "pending\|completed\|failed", "output": ..., "error": ...}` |

Tool signatures are stable across backends — swapping `FakeNodePool` for
`VastNodePool` changes no client-visible interface.

### Caching note

`list_nodes` returns `meta.ttlMs` and `meta.cacheScope` on its tool result.
This is a **local convention** — SEP-2549 governs `tools/list` and
`resources/list` responses, not individual `tools/call` results. Clients that
recognize it can cache; those that don't will simply re-call.

### Routing headers

The server emits `Mcp-Method` and `Mcp-Name` headers for gateway routing but
does not enforce them server-side. The enforcement point is the edge (API
gateway, reverse proxy), not the broker itself.

## Error taxonomy

Every tool error returns structured JSON with `code`, `message`, `retryable`,
and optional `retry_after_seconds`. Agents should branch on `code`, never on
`message` — messages are human-readable diagnostics and may change.

| Code | Retryable | When it fires |
|---|---|---|
| `NVML_VERSION_MISMATCH` | No | NVIDIA management library version does not match the driver on the GPU host |
| `DRIVER_LIBRARY_MISMATCH` | No | CUDA driver/library version conflict on the GPU host |
| `DPKG_LOCK_CONTENTION` | Yes | Package manager lock held by another process on the GPU host (e.g. unattended-upgrades) |
| `DOCKER_SOCKET_PERMISSION_DENIED` | No | Container runtime socket inaccessible on the GPU host |
| `INSUFFICIENT_VRAM` | No | Not enough GPU memory for the requested workload |
| `NODE_UNREACHABLE` | Yes | Cannot connect to the GPU node (SSH timeout, connection refused, DNS failure) |
| `RESERVATION_EXPIRED` | No | The signed handle's TTL has elapsed |
| `HANDLE_SIGNATURE_INVALID` | No | HMAC signature does not match — tampered, wrong secret, or malformed handle |
| `HANDLE_SCOPE_INVALID` | No | Handle scope mismatch (e.g. passing a task handle where a reservation handle is expected) |
| `TLS_PROXY_FAILURE` | Yes | TLS termination or proxy-layer failure between broker and node |
| `JOB_NOT_FOUND` | No | Signature valid but job absent from this replica's store (expected with in-memory store across replicas) |

The host-level errors (`NVML_VERSION_MISMATCH` through
`DOCKER_SOCKET_PERMISSION_DENIED`) are mapped from SSH stderr strings in
`vast.py:_raise_from_stderr`. The patterns are based on known failure modes
from Vast.ai GPU hosts but have not yet been validated against captured
production strings. Task 3 will capture verbatim error output and refine the
matching patterns.

## Quickstart

### Fake mode (no GPU, no API key)

```bash
export GPU_BROKER_SECRET="any-secret-string"
python src/gpu_broker/server.py
# Server at http://127.0.0.1:8000/mcp
```

### Vast.ai mode (real GPU)

```bash
export GPU_BROKER_SECRET="any-secret-string"
export VASTAI_API_KEY="your-vast-api-key"

# Find and rent a node
python vast_manage.py search --gpu "RTX 3090" --max-price 0.30
python vast_manage.py rent <offer_id>
python vast_manage.py wait <instance_id>

# Start the broker (auto-detects VASTAI_API_KEY)
python src/gpu_broker/server.py

# When done
python vast_manage.py destroy <instance_id>
```

## Running tests

```bash
uv run pytest tests/ -v
```

Tests include:

- Handle round-trip (reserve → dispatch → get_result)
- Tampered signature rejection
- Expired handle rejection
- Scope mismatch rejection
- JOB_NOT_FOUND for cross-replica lookup
- **Subprocess statelessness test**: boots three real HTTP servers (A and B
  share a secret, C has a different one), issues a task from A, confirms A
  returns `pending`, B returns `JOB_NOT_FOUND`, and C returns
  `HANDLE_SIGNATURE_INVALID`
- Secret-unset boot refusal
- Serialization round-trip for every error variant

## Current scope and limitations

This is a working prototype, not a production system.

- **FakeNodePool** returns a static list of three nodes and does not dispatch
  real inference. Useful for testing tool interactions and handle mechanics.
- **VastNodePool** queries the Vast.ai API for running instances and
  dispatches inference via SSH. It does real work but has no connection
  pooling, retry logic, or SSH key management beyond the system default.
- **InMemoryJobStore** loses all state on restart and cannot share job status
  across replicas. A production deployment needs a shared backend (Redis,
  Postgres).
- The error taxonomy patterns for host-level failures are educated guesses
  based on known failure modes. They need validation against real captured
  stderr from GPU hosts.
- No authentication on the MCP endpoint itself — any client that can reach
  the HTTP port can call tools. Production needs an auth layer in front.
- No rate limiting, no request size limits, no audit logging.