Remote workers: a GPU box dials in and runs the nodes bound to it

The engine runs where the automations are and the GPU is somewhere else,
usually behind a different network — so the worker connects out and the engine
answers over the socket it was given. Nothing has to expose Redis, and the
same connection works through the tunnel the hosted access will use.

What travels is the protocol the local pool already speaks, so a node cannot
tell which kind of worker it is on. A node declares device: gpu and
device_policy, the label is resolved per call (a worker attaching later needs
no rebuild), and a run whose labels nothing carries waits in the queue saying
what it waits for rather than failing — submit from the couch, the GPU box
picks it up when it is switched on.

Two things had to move with it. Compiling now happens on the machine that will
run the node: a node importing torch is correct on the GPU box and a missing
module on the engine, so checking it here failed nodes that were fine. And the
artifact endpoint accepts a worker's own credential, because storing a
checkpoint is exactly what that credential is for — and only that.

Verified against the real split: the training ran on this host (its checkpoint
names the machine and a numpy the engine does not have), streamed 40 metric
points back mid-run, and the evaluate node read the checkpoint on the engine.
Cancel kills the remote training; pulling the worker fails the run in six
seconds instead of waiting out its ten-minute timeout.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01AD8SfVhzXBG2nAfFcVh3iD
This commit is contained in:
2026-08-18 17:52:42 +02:00
co-authored by Claude Fable 5
parent b1cb8b41bc
commit e302ba1a43
16 changed files with 1401 additions and 11 deletions
+2
View File
@@ -15,6 +15,7 @@ from app.api.routes import (
secrets,
users,
utils,
workers,
)
api_router = APIRouter()
@@ -31,6 +32,7 @@ api_router.include_router(modules.router)
api_router.include_router(observability.router)
api_router.include_router(runs.router)
api_router.include_router(artifacts.router)
api_router.include_router(workers.router)
# Always mounted so the generated SDK stays the same shape; the endpoints
# themselves refuse to work unless MCP is switched on.
api_router.include_router(oauth.router)
+32 -2
View File
@@ -9,13 +9,43 @@ from typing import Any
from fastapi import APIRouter, Depends, HTTPException, Query, Request
from fastapi.responses import StreamingResponse
from jwt.exceptions import InvalidTokenError
from pydantic import BaseModel
from sqlmodel import Session
from app.api.deps import get_current_user
from app.api.deps import user_from_token
from app.core import security
from app.core.db import engine
from app.flow.artifacts import ArtifactStore
def artifact_caller(request: Request) -> str:
"""Who may move artifacts: a signed-in person, or an attached worker.
A worker's node stores its checkpoints through this endpoint, so its own
credential has to open it — and only it. The token is no use anywhere else
in the API, which is why this check is here rather than in the shared
dependency every other route uses.
"""
header = request.headers.get("Authorization", "")
token = header[7:] if header.lower().startswith("bearer ") else ""
if not token:
raise HTTPException(status_code=401, detail="Not authenticated")
try:
claims = security.decode_worker_token(token)
except InvalidTokenError:
pass
else:
return f"worker:{claims.get('sub')}"
with Session(engine) as session:
user = user_from_token(session, token)
if user is None:
raise HTTPException(status_code=401, detail="Not authenticated")
return user.email
router = APIRouter(
prefix="/artifacts", tags=["artifacts"], dependencies=[Depends(get_current_user)]
prefix="/artifacts", tags=["artifacts"], dependencies=[Depends(artifact_caller)]
)
+7 -1
View File
@@ -491,8 +491,14 @@ async def save_node_source(
await run_in_threadpool(
controller.store.write_node_source, name, node_id, source.code, True
)
node_def = next((n for n in definition.nodes if n.id == node_id), None)
device = (
node_def.device
if node_def is not None and node_def.device_policy == "require"
else None
)
error = await run_in_threadpool(
controller.compile_check, name, node_id, source.code
controller.compile_check, name, node_id, source.code, device
)
return NodeStatusPublic(
id=f"{name}.{node_id}",
+165
View File
@@ -0,0 +1,165 @@
"""Remote workers: how one attaches, and what is attached right now.
A worker dials in rather than being dialled: the GPU box and the engine are
usually on different networks, and only one of them can be reached. It presents
a token minted here, says what it can do, and then answers calls on the socket
it opened.
"""
from __future__ import annotations
import asyncio
import logging
from datetime import timedelta
from typing import Any
from fastapi import APIRouter, Depends, HTTPException, Request, WebSocket
from fastapi.responses import PlainTextResponse
from jwt.exceptions import InvalidTokenError
from pydantic import BaseModel, Field
from app.api.deps import get_current_active_superuser, get_current_user
from app.core import security
from app.flow import worker_main
from app.flow.remote import PROTOCOL, RemoteWorker, RemoteWorkerHub
logger = logging.getLogger(__name__)
router = APIRouter(prefix="/workers", tags=["workers"])
#: Long, because a worker is a machine somebody set up once and left running.
TOKEN_DAYS = 365
class WorkerInfo(BaseModel):
name: str
labels: list[str] = Field(default_factory=list)
max_parallel: int = 1
in_flight: int = 0
attached_at: float = 0.0
last_seen: float = 0.0
python: str = ""
venv_digest: str = ""
class TokenRequest(BaseModel):
name: str
class TokenIssued(BaseModel):
name: str
token: str
expires_days: int = TOKEN_DAYS
def _hub(app: Any) -> RemoteWorkerHub:
hub: RemoteWorkerHub | None = getattr(app.state, "worker_hub", None)
if hub is None:
raise HTTPException(status_code=503, detail="Remote workers are not available")
return hub
@router.get(
"", response_model=list[WorkerInfo], dependencies=[Depends(get_current_user)]
)
def read_workers(request: Request) -> Any:
"""What is attached, and how busy it is."""
return [
WorkerInfo(
name=worker.name,
labels=sorted(worker.labels),
max_parallel=worker.max_parallel,
in_flight=worker.in_flight,
attached_at=worker.attached_at,
last_seen=worker.last_seen,
python=str(worker.info.get("python") or ""),
venv_digest=str(worker.info.get("venv_digest") or ""),
)
for worker in _hub(request.app).workers()
]
@router.post(
"/tokens",
response_model=TokenIssued,
dependencies=[Depends(get_current_active_superuser)],
)
def issue_token(body: TokenRequest) -> Any:
"""Mint the credential a worker presents when it dials in.
Shown once. It is signed with the same keypair the agent tokens use, so
rotating that key revokes every worker along with them.
"""
token = security.create_worker_token(body.name, timedelta(days=TOKEN_DAYS))
return TokenIssued(name=body.name, token=token)
@router.get(
"/runtime",
response_class=PlainTextResponse,
dependencies=[Depends(get_current_user)],
)
def read_runtime() -> str:
"""The worker's own code, so a fresh host installs by fetching one file.
It is the same module the engine's local workers run — deliberately
standard library only, and with nothing of the app importable in it.
"""
return worker_main.__file__ and open(worker_main.__file__).read()
@router.websocket("/attach")
async def attach(websocket: WebSocket, token: str = "") -> None:
"""A worker's connection, for as long as it holds.
The token goes in the query string for the same reason the dashboard's
does: a websocket handshake carries no headers of its own.
"""
try:
claims = security.decode_worker_token(token)
except InvalidTokenError:
await websocket.close(code=1008)
return
await websocket.accept()
try:
hello = await asyncio.wait_for(websocket.receive_json(), timeout=30)
except (TimeoutError, asyncio.TimeoutError, ValueError):
await websocket.close(code=1002)
return
if hello.get("op") != "hello" or int(hello.get("protocol", 0)) != PROTOCOL:
await websocket.send_json(
{"op": "refused", "reason": f"this engine speaks protocol {PROTOCOL}"}
)
await websocket.close(code=1002)
return
# The token names the worker; what it calls itself is a suggestion, so two
# hosts cannot fight over one identity by claiming the same name.
name = str(claims.get("sub") or hello.get("name") or "worker")
hub = _hub(websocket.app)
worker = RemoteWorker(
name=name,
labels=[str(label) for label in (hello.get("labels") or [])],
send=websocket.send_json,
loop=asyncio.get_running_loop(),
max_parallel=max(1, int(hello.get("max_parallel") or 1)),
info={
"python": hello.get("python"),
"venv_digest": hello.get("venv_digest"),
},
)
hub.attach(worker)
await websocket.send_json({"op": "welcome", "protocol": PROTOCOL, "name": name})
try:
while True:
message = await websocket.receive_json()
worker.deliver(message)
except Exception:
# Any way this ends is the same thing: the socket is gone, and whatever
# was waiting on it has to be told rather than left hanging.
logger.info("Worker '%s' disconnected", name)
finally:
hub.detach(name)