Files
neuron-tai/docs/adr/0014-relay-outbound-client.md
2026-07-10 01:30:07 +03:00

124 lines
5.1 KiB
Markdown

# ADR-0014: Relay outbound client for NAT/internet pipeline hops
## Status: Accepted, amended 2026-07-10
## Context
ADR-0010 describes the relay server: a public WebSocket hub where nodes behind NAT
connect outbound and register as reachable peers. That ADR focused on the *inbound*
side: how the tracker reaches a behind-NAT node for the initial chat request.
The *pipeline hop* problem is different: when node A has the head shard and node B
(behind NAT) has the tail shard, node A must forward binary activations to node B
for *every generated token*. Direct HTTP from A to B is blocked. The relay must
carry this per-hop activation traffic.
### Why this is harder than tracker → node
The tracker-to-node relay (ADR-0010) proxies a single JSON request. The activation
hop carries raw bfloat16 tensors — binary data that must survive round-tripping
through the relay's JSON message envelope without precision loss.
Also, the relay `/rpc/{peer_id}` endpoint (one WebSocket connection per request)
must be opened and closed for every token in the autoregressive loop. Latency
of connection setup matters.
## Options considered
**A. Relay hop (persistent per Route Session, chosen)**
Node A opens a WebSocket to `wss://relay/rpc/{peer_id_B}`, sends activation requests
sequentially for the Route Session, then closes it when generation ends. The relay's
`_handle_rpc` forwards each request to B's persistent connection via the existing
`relay-http-request` envelope mechanism.
Pros: reuses the existing relay server unchanged. Each hop is independent; failures don't
affect other requests.
The original implementation opened and closed this socket per token. It was amended
to retain one requester socket per downstream relay address for the generation, so
connection setup is amortized across all tokens.
**B. Multiplexed persistent tunnel**
Node A sends multiple concurrent Route Sessions over a shared WebSocket and demultiplexes
responses by request id.
Pros: amortises connection setup across tokens.
Cons: requires session-level state on the relay; complicates relay shutdown/failover;
the current relay is stateless by design. Deferred for a future optimization.
**C. Tracker-proxied activations**
Route all activation traffic through the tracker's HTTP proxy.
Cons: the tracker is the control plane, not the data plane. High-volume binary tensor
traffic through the tracker would saturate it. Rejected.
## Decision
Option A — one sequential WebSocket per relayed Activation Seam and Route Session.
Each activation still has a unique request id for response correlation, while
`X-Meshnet-Session` remains stable for KV state. Option B remains a possible
connection-count optimization for high-concurrency workloads.
## Protocol
```
Node A opens WS once → wss://relay/rpc/{peer_id_B}
Node A sends repeatedly:
{
"request_id": "<hex>",
"method": "POST",
"path": "/forward",
"headers": { "X-Meshnet-Shape": "...", "X-Meshnet-Start-Layer": "12", ... },
"body_base64": "<base64(bfloat16 tensor)>"
}
Relay forwards to Node B as relay-http-request envelope.
Node B's RelayHttpBridge decodes body_base64, calls POST /forward locally.
Response:
{
"request_id": "<hex>",
"status": 200,
"headers": { "x-meshnet-shape": "...", "content-type": "application/octet-stream" },
"body_base64": "<base64(output tensor)>" ← for binary responses
# OR
"body": "<json string>" ← for text (last-hop decode)
}
Relay sends each response back to Node A without closing the requester socket.
Node A continues the pipeline and closes the socket when generation ends.
```
### Binary data through JSON: base64
Raw bfloat16 bytes cannot safely transit JSON (no UTF-8 guarantee, lossy decode).
`body_base64` carries the tensor as base64; the bridge decodes it before calling
the local HTTP endpoint, and re-encodes the response. No precision loss.
Text responses (final hop, `application/json`) use `body` (plain string) for efficiency.
### Fallback
If `_relay_hop` raises (relay unreachable, peer disconnected), `_run_downstream_pipeline`
logs a warning and retries via direct HTTP. If both fail, the hop returns a pipeline error
string and the token is skipped.
### Tracker injection
The tracker's `_handle_proxy_chat` includes `relay_addr` in each downstream hop dict
when the node has one registered:
```json
{"endpoint": "http://172.29.x.x:7002", "start_layer": 12, "relay_addr": "wss://relay/rpc/abc123"}
```
The head node reads `relay_addr` from the injected `X-Meshnet-Route` header and calls
`_relay_hop` instead of direct HTTP.
## Consequences
- Nodes behind NAT (WSL2, 5G, home routers) can now participate in distributed pipeline inference without opening firewall ports
- `relay_addr` is a stable registration field; nodes without a relay omit it and receive direct HTTP hops
- WebSocket/TCP/TLS setup occurs once per relayed Activation Seam per Route Session, not once per generated token
- Base64 encoding increases payload size by ~33%; acceptable for prototype
- The relay server remains stateless and horizontally scalable; only the persistent per-peer `/ws` connections are stateful