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)

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)

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

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.