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:
@@ -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)
|
||||
|
||||
@@ -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)]
|
||||
)
|
||||
|
||||
|
||||
|
||||
@@ -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}",
|
||||
|
||||
@@ -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)
|
||||
Reference in New Issue
Block a user