diff --git a/backend/app/api/main.py b/backend/app/api/main.py index 79ff173..42c6820 100644 --- a/backend/app/api/main.py +++ b/backend/app/api/main.py @@ -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) diff --git a/backend/app/api/routes/artifacts.py b/backend/app/api/routes/artifacts.py index 44d393e..64ea46f 100644 --- a/backend/app/api/routes/artifacts.py +++ b/backend/app/api/routes/artifacts.py @@ -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)] ) diff --git a/backend/app/api/routes/flows.py b/backend/app/api/routes/flows.py index 4aa7c70..5335c8d 100644 --- a/backend/app/api/routes/flows.py +++ b/backend/app/api/routes/flows.py @@ -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}", diff --git a/backend/app/api/routes/workers.py b/backend/app/api/routes/workers.py new file mode 100644 index 0000000..7939710 --- /dev/null +++ b/backend/app/api/routes/workers.py @@ -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) diff --git a/backend/app/core/security.py b/backend/app/core/security.py index 8ae60f6..74d094a 100644 --- a/backend/app/core/security.py +++ b/backend/app/core/security.py @@ -29,6 +29,10 @@ ALGORITHM = "HS256" OAUTH_ALGORITHM = "RS256" #: The one scope an MCP token carries. MCP_SCOPE = "mcp" +#: What a remote worker's credential says it is for. Its own audience, so an +#: agent's token cannot attach a worker and a worker's cannot call the API. +WORKER_SCOPE = "worker" +WORKER_AUDIENCE = "fluksio-worker" def create_access_token(subject: str | Any, expires_delta: timedelta) -> str: @@ -124,6 +128,47 @@ def create_oauth_access_token( ) +def create_worker_token(name: str, expires_delta: timedelta) -> str: + """A credential a remote worker presents when it dials in. + + Signed with the same keypair the agent tokens use, so the whole set can be + revoked by rotating one key, and told apart from them by its audience: a + worker's token grants no API access, and an agent's opens no worker + connection. + """ + now = datetime.now(timezone.utc) + payload = { + "sub": name, + "iss": settings.oauth_issuer, + "aud": WORKER_AUDIENCE, + "iat": now, + "exp": now + expires_delta, + "scope": WORKER_SCOPE, + } + return jwt.encode( + payload, + oauth_key().private_bytes( + encoding=serialization.Encoding.PEM, + format=serialization.PrivateFormat.PKCS8, + encryption_algorithm=serialization.NoEncryption(), + ), + algorithm=OAUTH_ALGORITHM, + headers={"kid": "fluksio-oauth"}, + ) + + +def decode_worker_token(token: str) -> dict[str, Any]: + """Validate a worker credential. Raises ``InvalidTokenError`` if it does not.""" + payload: dict[str, Any] = jwt.decode( + token, + oauth_key().public_key(), + algorithms=[OAUTH_ALGORITHM], + audience=WORKER_AUDIENCE, + issuer=settings.oauth_issuer, + ) + return payload + + def decode_oauth_token(token: str) -> dict[str, Any]: """Validate an MCP token. Raises ``InvalidTokenError`` if it does not hold.""" payload: dict[str, Any] = jwt.decode( diff --git a/backend/app/flow/controller.py b/backend/app/flow/controller.py index 6157db4..fb6d42d 100644 --- a/backend/app/flow/controller.py +++ b/backend/app/flow/controller.py @@ -42,6 +42,7 @@ from app.flow.nodes import ( TriggerNode, ) from app.flow.pipeline import NodeOutcome, Pipeline, ValidationIssue, ValueSource +from app.flow.remote import RemoteWorkerHub from app.flow.schemas import ( BrainEdge, BrainGraph, @@ -257,11 +258,14 @@ class FlowController: execution: ExecutionService | None = None, alerts: AlertManager | None = None, workers: PythonWorkerPool | None = None, + remote: RemoteWorkerHub | None = None, ) -> None: self.store = store # Without a pool, python nodes are compiled and run in this process — # which is what the tests do, and what a bare `Pipeline` has always done. self.workers = workers + # Workers on other hosts. A node without a device never touches it. + self.remote = remote self.state = state if state is not None else MemoryState() self.events = events self.max_workers = max_workers @@ -469,20 +473,53 @@ class FlowController: else: # The code never runs here: it is loaded in a worker, and # the node calls that worker instead of a local function. - problem = self.workers.compile(owner, local, code) + # A node bound to a device is loaded on *that* machine — + # one importing torch is correct on the GPU box and a + # missing module here, so checking it here would fail a + # node that is fine. + remote_only = ( + node_def.device + and node_def.device_policy == "require" + and self.remote is not None + ) + if remote_only and self.remote is not None: + problem = self.remote.compile( + node_def.device or "", owner, local, code + ) + else: + problem = self.workers.compile(owner, local, code) if problem: entry.status = NodeStatus.ERROR entry.error = problem return entry + timeout = node_def.timeout or settings.FLOW_NODE_TIMEOUT function = self.workers.proxy( owner, local, code, node_id=node_id, - timeout=node_def.timeout or settings.FLOW_NODE_TIMEOUT, + timeout=timeout, run_id=run.run_id if run else "", on_event=run.on_event if run else None, ) + if node_def.device and self.remote is not None: + # A node with a device runs on a worker carrying that + # label. Which worker is decided per call, so one that + # attaches after this flow was built is used without + # anything being rebuilt. + function = self.remote.proxy( + node_def.device, + owner, + local, + code, + node_id=node_id, + timeout=timeout, + run_id=run.run_id if run else "", + on_event=run.on_event if run else None, + fallback=( + function if node_def.device_policy == "prefer" else None + ), + ) node = Node( f=function, requires=_bound(node_def.requires), @@ -587,8 +624,17 @@ class FlowController: issues=[issue for issue in issues if not issue.flow or issue.flow == name], ) - def compile_check(self, flow: str, node_id: str, code: str) -> str | None: - """Does this source load? Returns what to show the author, or None.""" + def compile_check( + self, flow: str, node_id: str, code: str, device: str | None = None + ) -> str | None: + """Does this source load? Returns what to show the author, or None. + + Loaded on the machine that will run it: a node bound to a device is + checked against that worker's packages, because a missing import here + says nothing about whether it is missing there. + """ + if device and self.remote is not None: + return self.remote.compile(device, flow, node_id, code) if self.workers is not None: return self.workers.compile(flow, node_id, code) try: diff --git a/backend/app/flow/remote.py b/backend/app/flow/remote.py new file mode 100644 index 0000000..4a7a362 --- /dev/null +++ b/backend/app/flow/remote.py @@ -0,0 +1,356 @@ +"""Workers on other hosts, reached over a socket they opened themselves. + +The engine runs where the automations are; a GPU sits somewhere else. Those +two are usually not on the same network, and the one that can be dialled is +the engine — so a worker connects *out* to it and the engine answers over the +connection it was given. That also means nothing has to expose Redis, which is +the thing a remote worker must never be handed. + +What travels is the protocol the local worker pool already speaks: one JSON +object per line becomes one JSON frame, the node's source rides along with +every call so no code has to be distributed, and the reports a node makes +while it runs arrive the same way they do over a pipe. A node cannot tell +which kind of worker it is running on, which is the point — the same flow runs +in both places. + +The awkward part is that the socket lives on the event loop and a node +executes on a worker thread. A call therefore hands its frame to the loop and +blocks on a queue of its own until the loop puts the answer there. +""" + +from __future__ import annotations + +import asyncio +import logging +import queue +import sys +import threading +import time +from collections.abc import Callable +from typing import Any + +from app.flow.workers import NodeTimeout, RemoteError, _remote_class + +logger = logging.getLogger(__name__) + +#: How long the loop is given to accept a frame we are handing it. +SEND_TIMEOUT_S = 30.0 +#: A worker that has said nothing for this long is treated as gone. It sends a +#: heartbeat while it is executing, so this only ever catches a dead socket. +SILENCE_S = 90.0 +#: Protocol version this engine speaks. A worker announcing anything else is +#: refused rather than half-understood. +PROTOCOL = 1 + + +class NoWorker(RemoteError): + """Nothing is attached that carries the label this node asked for.""" + + +class RemoteWorker: + """One attached worker, and the calls it has in flight.""" + + def __init__( + self, + name: str, + labels: list[str], + send: Callable[[dict[str, Any]], Any], + loop: asyncio.AbstractEventLoop, + max_parallel: int = 1, + info: dict[str, Any] | None = None, + ) -> None: + self.name = name + self.labels = set(labels) + self.info = info or {} + self.attached_at = time.time() + self.last_seen = time.time() + self._send = send + self._loop = loop + self._slots = threading.Semaphore(max_parallel) + self.max_parallel = max_parallel + self._pending: dict[str, queue.Queue[dict[str, Any] | None]] = {} + self._lock = threading.Lock() + self._gone = False + + # ------------------------------------------------------------------------- + # From the socket's side, on the event loop + # ------------------------------------------------------------------------- + + def deliver(self, message: dict[str, Any]) -> None: + """Hand a frame to whichever call is waiting for it.""" + self.last_seen = time.time() + call_id = str(message.get("call_id") or "") + with self._lock: + inbox = self._pending.get(call_id) + if inbox is not None: + inbox.put(message) + + def detach(self) -> None: + """The socket closed: wake everything still waiting on it.""" + self._gone = True + with self._lock: + inboxes = list(self._pending.values()) + for inbox in inboxes: + # None is "no more answers are coming", which the caller turns into + # a failed node rather than a wait that never ends. + inbox.put(None) + + # ------------------------------------------------------------------------- + # From a node's side, on a worker thread + # ------------------------------------------------------------------------- + + def request( + self, + payload: dict[str, Any], + timeout: float, + on_event: Callable[[dict[str, Any]], None] | None = None, + ) -> dict[str, Any]: + if self._gone: + raise RemoteError(f"worker '{self.name}' is no longer attached") + call_id = str(payload["call_id"]) + inbox: queue.Queue[dict[str, Any] | None] = queue.Queue() + # Blocking here is the backpressure, exactly as taking a slot is in the + # local pool. + self._slots.acquire() + with self._lock: + self._pending[call_id] = inbox + try: + future = asyncio.run_coroutine_threadsafe(self._send(payload), self._loop) + future.result(timeout=SEND_TIMEOUT_S) + + while True: + try: + # Reset per frame: the deadline measures silence, so a node + # reporting its progress is never mistaken for a hung one. + message = inbox.get(timeout=timeout) + except queue.Empty: + self.cancel(call_id) + raise NodeTimeout( + f"'{self.name}' was silent for {timeout}s" + ) from None + if message is None: + raise RemoteError(f"worker '{self.name}' went away mid-call") + kind = message.get("event") + if kind == "heartbeat": + continue + if kind: + if on_event is not None: + try: + on_event(message) + except Exception: + logger.exception("Could not record a worker event") + continue + return message + except Exception as exc: + if isinstance(exc, (NodeTimeout, RemoteError)): + raise + raise RemoteError(f"worker '{self.name}': {exc}") from exc + finally: + with self._lock: + self._pending.pop(call_id, None) + self._slots.release() + + @property + def gone(self) -> bool: + return self._gone + + @property + def in_flight(self) -> int: + """Calls this worker has not answered yet.""" + with self._lock: + return len(self._pending) + + def calls_of(self, run_id: str) -> list[str]: + with self._lock: + return [call for call in self._pending if call.startswith(f"{run_id}:")] + + def cancel(self, call_id: str) -> None: + """Ask the worker to kill what it is running for this call.""" + try: + asyncio.run_coroutine_threadsafe( + self._send({"op": "cancel", "call_id": call_id}), self._loop + ).result(timeout=SEND_TIMEOUT_S) + except Exception: + logger.warning("Could not cancel '%s' on '%s'", call_id, self.name) + + +class RemoteWorkerHub: + """Every attached worker, and which of them a node may run on.""" + + def __init__(self) -> None: + self._workers: dict[str, RemoteWorker] = {} + self._lock = threading.Lock() + + # ------------------------------------------------------------------------- + # Attachment + # ------------------------------------------------------------------------- + + def attach(self, worker: RemoteWorker) -> None: + with self._lock: + existing = self._workers.get(worker.name) + if existing is not None: + # A worker that reconnects after a network drop: the old socket + # is dead whether or not it has noticed yet. + existing.detach() + self._workers[worker.name] = worker + logger.info( + "Worker '%s' attached with labels %s", worker.name, sorted(worker.labels) + ) + + def detach(self, name: str) -> None: + with self._lock: + worker = self._workers.pop(name, None) + if worker is not None: + worker.detach() + logger.info("Worker '%s' detached", name) + + def workers(self) -> list[RemoteWorker]: + with self._lock: + return list(self._workers.values()) + + def labels(self) -> set[str]: + """Every label something attached right now carries.""" + with self._lock: + return { + label for worker in self._workers.values() for label in worker.labels + } + + def pick(self, label: str) -> RemoteWorker | None: + """A worker carrying this label, least busy first. + + Resolved per call rather than when the flow was built, so a worker that + attaches after a run was submitted picks the work up without anything + being rebuilt. + """ + with self._lock: + candidates = [ + worker + for worker in self._workers.values() + if not worker.gone and (label == worker.name or label in worker.labels) + ] + if not candidates: + return None + return min(candidates, key=lambda worker: worker.in_flight) + + # ------------------------------------------------------------------------- + # Running a node on one + # ------------------------------------------------------------------------- + + def run( + self, + label: str, + flow: str, + node: str, + source: str, + kwargs: dict[str, Any], + params: dict[str, Any] | None, + node_id: str, + timeout: float, + run_id: str = "", + on_event: Callable[[dict[str, Any]], None] | None = None, + ) -> Any: + worker = self.pick(label) + if worker is None: + raise NoWorker(f"no worker labelled '{label}' is attached") + response = worker.request( + { + "op": "run", + "call_id": f"{run_id}:{node_id}" if run_id else node_id, + "flow": flow, + "node": node, + "source": source, + "kwargs": kwargs, + "params": params or {}, + "run": {"id": run_id} if run_id else None, + "timeout": timeout, + }, + timeout=timeout, + on_event=on_event, + ) + logs = response.get("logs") + if logs: + sys.stdout.write(logs) + if response.get("ok"): + return response.get("result") + error = response.get("error") or {} + raise _remote_class(str(error.get("type") or "RemoteError"))( + str(error.get("message") or "the node failed"), + str(error.get("traceback") or ""), + ) + + def compile( + self, label: str, flow: str, node: str, source: str, timeout: float = 60.0 + ) -> str | None: + """Load this source on the worker that will run it. + + Which machine compiles matters: a node importing torch is fine on the + GPU box and a ``ModuleNotFoundError`` on the engine, so checking it + here would fail a node that is perfectly correct. When nothing is + attached there is nothing to check against, and ``None`` says so — the + node is not broken, it is waiting. + """ + worker = self.pick(label) + if worker is None: + return None + try: + response = worker.request( + { + "op": "compile", + "call_id": f"compile:{flow}.{node}", + "flow": flow, + "node": node, + "source": source, + }, + timeout=timeout, + ) + except RemoteError as exc: + return f"{type(exc).__name__}: {exc}" + if response.get("ok"): + return None + error = response.get("error") or {} + return str(error.get("short") or "The node could not be loaded.") + + def proxy( + self, + label: str, + flow: str, + node: str, + source: str, + node_id: str, + timeout: float, + run_id: str = "", + on_event: Callable[[dict[str, Any]], None] | None = None, + fallback: Callable[..., Any] | None = None, + ) -> Callable[..., Any]: + """What a node with a device runs instead of its own function. + + ``fallback`` is the local pool's proxy, used when the node only prefers + the label rather than requiring it. + """ + + def call(params: dict[str, Any] | None = None, **kwargs: Any) -> Any: + if fallback is not None and self.pick(label) is None: + return fallback(params=params, **kwargs) + return self.run( + label, + flow, + node, + source, + kwargs, + params, + node_id, + timeout, + run_id=run_id, + on_event=on_event, + ) + + return call + + def cancel_run(self, run_id: str) -> int: + """Kill whatever this run has executing on any attached worker.""" + stopped = 0 + for worker in self.workers(): + for call in worker.calls_of(run_id): + worker.cancel(call) + stopped += 1 + return stopped diff --git a/backend/app/flow/runs.py b/backend/app/flow/runs.py index 9b33040..d6d90eb 100644 --- a/backend/app/flow/runs.py +++ b/backend/app/flow/runs.py @@ -75,6 +75,8 @@ METRIC_BATCH = 500 METRIC_FLUSH_S = 2.0 #: Progress is for whoever is watching, so it is throttled hard. PROGRESS_INTERVAL_S = 1.0 +#: How often a run waiting for a worker looks again. +WAIT_RETRY_S = 15.0 #: Where a run's state lives, so it can never collide with the engine's own. RUN_NAMESPACE = "run" @@ -392,6 +394,8 @@ class RunService: # Keyed by run, so a sweep cancelling one config leaves the others # training. workers.cancel_run(run_id) + if self.controller.remote is not None: + self.controller.remote.cancel_run(run_id) return True def _flow_of(self, run_id: str) -> str | None: @@ -421,6 +425,9 @@ class RunService: failures = 0 while not self._stop.is_set(): try: + # Runs put back to wait for a worker come due here. The claim + # below blocks for a second, so this is about once a second. + self.queue.move_due(time.time()) items = self.queue.claim(CLAIM_COUNT, CLAIM_BLOCK_MS) failures = 0 except Exception as exc: @@ -429,16 +436,59 @@ class RunService: self._stop.wait(min(30.0, 2.0**failures)) continue for item in items: + if not item.run_id: + self.queue.ack(item) + continue + missing = self._missing_labels(item.run_id) + if missing: + # Left in the queue rather than failed: submitting a run + # before turning the GPU box on is a normal way to work, and + # the run says what it is waiting for while it waits. + self._waiting(item.run_id, missing) + self._defer(item) + continue # Acknowledged before it runs: from here on the row is the # record, and a lease that stops moving is what says otherwise. self.queue.ack(item) - if not item.run_id: - continue try: self._pool.submit(self._drive, item.run_id) except RuntimeError: logger.warning("Run %s not started: shutting down", item.run_id) + def _missing_labels(self, run_id: str) -> list[str]: + """Worker labels this run needs that nothing attached carries.""" + with Session(db_engine) as session: + run = session.get(Run, run_id) + needed = list(run.labels) if run else [] + if not needed: + return [] + hub = self.controller.remote + available = hub.labels() | {w.name for w in hub.workers()} if hub else set() + # A node that only prefers its label runs locally instead, so it is not + # a reason to hold the run back; that is decided per node at call time. + return sorted(set(needed) - available) + + def _waiting(self, run_id: str, missing: list[str]) -> None: + reason = f"Waiting for a worker labelled {', '.join(missing)}" + try: + with Session(db_engine) as session: + session.exec( + update(Run) + .where(col(Run.id) == run_id, col(Run.status) == "queued") + .values(status_reason=reason) + ) + session.commit() + except Exception: + logger.exception("Could not record what run %s is waiting for", run_id) + + def _defer(self, item: WorkItem) -> None: + """Put an item back for later, and let go of this delivery.""" + try: + self.queue.add_delayed(item, time.time() + WAIT_RETRY_S) + self.queue.ack(item) + except Exception: + logger.exception("Could not defer run %s", item.run_id) + def _keep_leases(self) -> None: """Say the local runs are alive, and clean up after engines that died.""" last_sweep = 0.0 diff --git a/backend/app/main.py b/backend/app/main.py index 6733f23..51cb3b8 100644 --- a/backend/app/main.py +++ b/backend/app/main.py @@ -26,6 +26,7 @@ from app.flow.nodes.http import close_shared_client from app.flow.pipeline import ValueSource from app.flow.plugins import load_plugins from app.flow.queue import MemoryWorkQueue, RedisWorkQueue, WorkQueue +from app.flow.remote import RemoteWorkerHub from app.flow.runs import RUN_STATE_TTL, RunService from app.flow.secrets import init_secrets from app.flow.state import MemoryState, RedisState, StateBackend @@ -116,6 +117,8 @@ async def lifespan(app: FastAPI) -> AsyncIterator[None]: ) pool.start() app.state.worker_pool = pool + worker_hub = RemoteWorkerHub() + app.state.worker_hub = worker_hub controller = FlowController( store=store, state=_state_backend(), @@ -125,6 +128,7 @@ async def lifespan(app: FastAPI) -> AsyncIterator[None]: execution=execution, alerts=alerts, workers=pool, + remote=worker_hub, ) app.state.flow_controller = controller # A "dashboard" alert channel puts its alert into the graph. Bound here diff --git a/backend/app/worker/__init__.py b/backend/app/worker/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/backend/app/worker/fluksio_worker.py b/backend/app/worker/fluksio_worker.py new file mode 100644 index 0000000..458f652 --- /dev/null +++ b/backend/app/worker/fluksio_worker.py @@ -0,0 +1,291 @@ +"""The agent that runs Fluksio nodes on a machine the engine cannot reach. + +Copy this file and ``worker_main.py`` onto the box with the GPU, point it at +the engine, and it dials in:: + + pip install websockets + python fluksio_worker.py --url wss://api.example.com/api/v1/workers/attach \\ + --token "$FLUKSIO_WORKER_TOKEN" --labels gpu --python /opt/venv/bin/python + +It connects *out*, so the engine needs no route back and nothing has to expose +Redis. What it then does is what the engine's own worker pool does: hold a few +subprocesses running ``worker_main.py``, hand each call to one, and pass back +everything that comes out — including the metrics a training loop reports +while it is still running. + +Deliberately one file with one dependency. Nothing of the app is imported +here; a worker host installs Python, ``websockets``, and whatever the nodes +themselves need. +""" + +from __future__ import annotations + +import argparse +import asyncio +import contextlib +import json +import logging +import os +import signal +import subprocess +import sys +from pathlib import Path +from typing import Any + +try: + import websockets +except ImportError: # pragma: no cover - the one dependency, named plainly + print( + "This needs the 'websockets' package: pip install websockets", file=sys.stderr + ) + raise SystemExit(1) from None + +logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s") +log = logging.getLogger("fluksio-worker") + +PROTOCOL = 1 +#: Sent while a call is running, so the engine can tell working from wedged. +HEARTBEAT_S = 10.0 +#: Reconnection backs off to this and no further. +MAX_BACKOFF_S = 30.0 +WORKER_MAIN = Path(__file__).with_name("worker_main.py") + + +class Subprocess: + """One user-code process and the framing of one call over its pipes.""" + + def __init__(self, python: str, env: dict[str, str]) -> None: + self.proc = subprocess.Popen( + [python, str(WORKER_MAIN)], + stdin=subprocess.PIPE, + stdout=subprocess.PIPE, + close_fds=True, + env=env, + ) + self._buffer = bytearray() + + def alive(self) -> bool: + return self.proc.poll() is None + + def send(self, request: dict[str, Any]) -> None: + assert self.proc.stdin is not None + self.proc.stdin.write((json.dumps(request) + "\n").encode()) + self.proc.stdin.flush() + + def read_line(self) -> str: + """One line, blocking. Empty when the process is gone. + + Read at the file-descriptor level so a report arriving mid-call is + passed on the moment it is written rather than when a buffer fills. + """ + assert self.proc.stdout is not None + fd = self.proc.stdout.fileno() + while True: + end = self._buffer.find(b"\n") + if end >= 0: + line = bytes(self._buffer[: end + 1]) + del self._buffer[: end + 1] + return line.decode(errors="replace") + try: + chunk = os.read(fd, 65536) + except OSError: + chunk = b"" + if not chunk: + self._buffer.clear() + return "" + self._buffer += chunk + + def kill(self) -> None: + with contextlib.suppress(OSError): + self.proc.send_signal(signal.SIGKILL) + with contextlib.suppress(subprocess.TimeoutExpired): + self.proc.wait(timeout=5) + + +class Agent: + """Holds the connection, and one subprocess per call in flight.""" + + def __init__(self, args: argparse.Namespace) -> None: + self.args = args + self.env = dict(os.environ) + self.env["FLUKSIO_ARTIFACT_URL"] = args.artifact_url or _artifacts_from( + args.url + ) + self.env["FLUKSIO_ARTIFACT_TOKEN"] = args.token + self.running: dict[str, Subprocess] = {} + + async def serve_forever(self) -> None: + backoff = 1.0 + while True: + try: + await self._session() + backoff = 1.0 + except Exception as exc: + log.warning("disconnected: %s — retrying in %.0fs", exc, backoff) + await asyncio.sleep(backoff) + backoff = min(MAX_BACKOFF_S, backoff * 2) + + async def _session(self) -> None: + url = f"{self.args.url}?token={self.args.token}" + async with websockets.connect(url, max_size=None, ping_interval=20) as socket: + await socket.send( + json.dumps( + { + "op": "hello", + "protocol": PROTOCOL, + "name": self.args.name, + "labels": self.args.labels, + "python": self.args.python, + "max_parallel": self.args.parallel, + "venv_digest": _venv_digest(self.args.python), + } + ) + ) + welcome = json.loads(await socket.recv()) + if welcome.get("op") != "welcome": + raise RuntimeError(str(welcome.get("reason") or "refused")) + log.info( + "attached to %s as '%s' with labels %s", + self.args.url, + welcome.get("name"), + self.args.labels, + ) + + async for raw in socket: + message = json.loads(raw) + op = message.get("op") + if op in ("run", "compile"): + # Compiling is loading the source, which is the same trip + # through a subprocess a call is — and has to happen here + # rather than on the engine, because "does this import" + # is a question about *this* machine's packages. + log.info("%s %s", op, message.get("call_id")) + task = asyncio.create_task(self._run(socket, message)) + # Without this a failure in the task is only noticed when + # it is garbage collected, which reads as a call that + # vanished. + task.add_done_callback(_report_failure) + elif op == "cancel": + self._cancel(str(message.get("call_id") or "")) + + async def _run(self, socket: Any, request: dict[str, Any]) -> None: + """Execute one call in a subprocess, streaming what it says back.""" + call_id = str(request.get("call_id") or "") + loop = asyncio.get_running_loop() + worker = Subprocess(self.args.python, self.env) + self.running[call_id] = worker + + async def beat() -> None: + while True: + await asyncio.sleep(HEARTBEAT_S) + with contextlib.suppress(Exception): + await socket.send( + json.dumps({"call_id": call_id, "event": "heartbeat"}) + ) + + heartbeat = asyncio.create_task(beat()) + try: + await loop.run_in_executor(None, worker.send, request) + while True: + line = await loop.run_in_executor(None, worker.read_line) + if not line: + await socket.send( + json.dumps( + { + "call_id": call_id, + "ok": False, + "error": { + "type": "NodeCancelled" + if worker.proc.returncode + else "RemoteError", + "message": "the node process stopped", + "short": "the node process stopped", + "traceback": "", + }, + } + ) + ) + return + await socket.send(line.strip()) + # Anything without an `event` is the answer; the call is over. + if not json.loads(line).get("event"): + return + finally: + heartbeat.cancel() + self.running.pop(call_id, None) + worker.kill() + + def _cancel(self, call_id: str) -> None: + worker = self.running.get(call_id) + if worker is not None: + log.info("cancelling %s", call_id) + worker.kill() + + +def _report_failure(task: asyncio.Task[Any]) -> None: + if not task.cancelled() and task.exception() is not None: + log.exception("call failed", exc_info=task.exception()) + + +def _artifacts_from(url: str) -> str: + """The artifact endpoint beside the socket, so one URL configures both.""" + base = url.replace("wss://", "https://").replace("ws://", "http://") + return base.rsplit("/workers/attach", 1)[0] + "/artifacts" + + +def _venv_digest(python: str) -> str: + """What is installed here, so the engine can say when it has drifted.""" + try: + listing = subprocess.run( + [python, "-m", "pip", "freeze"], + capture_output=True, + text=True, + timeout=60, + check=False, + ).stdout + except (OSError, subprocess.SubprocessError): + return "" + import hashlib + + return hashlib.sha256(listing.encode()).hexdigest()[:16] + + +def main() -> int: + parser = argparse.ArgumentParser(description="Run Fluksio nodes on this machine.") + parser.add_argument("--url", required=True, help="wss://…/api/v1/workers/attach") + parser.add_argument( + "--token", + default=os.environ.get("FLUKSIO_WORKER_TOKEN", ""), + help="issued by POST /api/v1/workers/tokens", + ) + parser.add_argument("--name", default=os.uname().nodename) + parser.add_argument( + "--labels", + default="", + help="comma-separated, e.g. gpu,cuda12 — what a node's device matches", + ) + parser.add_argument( + "--python", + default=sys.executable, + help="the interpreter node code runs on; point it at the venv with torch", + ) + parser.add_argument("--parallel", type=int, default=1) + parser.add_argument("--artifact-url", default="") + args = parser.parse_args() + + if not args.token: + parser.error("a token is required (--token or FLUKSIO_WORKER_TOKEN)") + args.labels = [part.strip() for part in args.labels.split(",") if part.strip()] + if not WORKER_MAIN.exists(): + parser.error(f"{WORKER_MAIN} is missing — copy it beside this file") + + agent = Agent(args) + try: + asyncio.run(agent.serve_forever()) + except KeyboardInterrupt: + log.info("stopping") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/backend/pyproject.toml b/backend/pyproject.toml index 57052b1..eb72a9d 100644 --- a/backend/pyproject.toml +++ b/backend/pyproject.toml @@ -92,6 +92,9 @@ ignore = [ # This one takes its own directory off sys.path before the rest of its imports # run, which is the whole point of doing it there. "app/flow/worker_main.py" = ["E402"] +# A standalone script copied onto another machine: it has no logger configured +# before it tells you the one dependency it is missing. +"app/worker/fluksio_worker.py" = ["T201"] [tool.ruff.lint.pyupgrade] # Preserve types, even if a file imports `from __future__ import annotations`. diff --git a/backend/tests/flow/test_remote.py b/backend/tests/flow/test_remote.py new file mode 100644 index 0000000..e0e4922 --- /dev/null +++ b/backend/tests/flow/test_remote.py @@ -0,0 +1,225 @@ +"""A worker on another host, and the thread-to-loop bridge that reaches it. + +The socket lives on an event loop; a node executes on a worker thread. These +run a real loop in a thread of its own, because that split is the whole +difficulty: what is checked is that a call handed across it comes back — with +its answer, with the reports it made on the way, or with a failure that says +what happened rather than a wait that never ends. +""" + +import asyncio +import threading +from collections.abc import Iterator + +import pytest + +from app.flow.remote import NoWorker, RemoteWorker, RemoteWorkerHub +from app.flow.workers import NodeTimeout, RemoteError + + +@pytest.fixture +def loop() -> Iterator[asyncio.AbstractEventLoop]: + """An event loop running in a thread, as the server's does.""" + running = asyncio.new_event_loop() + thread = threading.Thread(target=running.run_forever, daemon=True) + thread.start() + yield running + running.call_soon_threadsafe(running.stop) + thread.join(timeout=5) + running.close() + + +class FakeSocket: + """Stands in for the websocket: records frames, and says when one lands.""" + + def __init__(self) -> None: + self.sent: list[dict] = [] + self.arrived = threading.Event() + + async def send_json(self, payload: dict) -> None: + self.sent.append(payload) + self.arrived.set() + + +def attach(hub: RemoteWorkerHub, loop: asyncio.AbstractEventLoop, name: str = "gpu1"): + socket = FakeSocket() + worker = RemoteWorker( + name=name, labels=["gpu"], send=socket.send_json, loop=loop, max_parallel=2 + ) + hub.attach(worker) + return worker, socket + + +def call_in_thread(target) -> threading.Thread: + thread = threading.Thread(target=target, daemon=True) + thread.start() + return thread + + +def test_a_call_crosses_to_the_thread_and_the_answer_comes_back(loop): + hub = RemoteWorkerHub() + worker, socket = attach(hub, loop) + result: dict = {} + + thread = call_in_thread( + lambda: result.update( + value=hub.run( + "gpu", "flow", "node", "src", {"x": 1}, {}, "flow.node", timeout=5 + ) + ) + ) + assert socket.arrived.wait(5) + assert socket.sent[0]["source"] == "src" + assert socket.sent[0]["kwargs"] == {"x": 1} + + worker.deliver( + {"call_id": socket.sent[0]["call_id"], "ok": True, "result": {"out": 2}} + ) + thread.join(timeout=5) + assert result["value"] == {"out": 2} + + +def test_reports_arrive_before_the_answer_and_a_heartbeat_is_not_one(loop): + hub = RemoteWorkerHub() + worker, socket = attach(hub, loop) + seen: list[dict] = [] + result: dict = {} + + thread = call_in_thread( + lambda: result.update( + value=hub.run( + "gpu", + "flow", + "node", + "src", + {}, + {}, + "flow.node", + timeout=5, + run_id="r1", + on_event=seen.append, + ) + ) + ) + assert socket.arrived.wait(5) + call_id = socket.sent[0]["call_id"] + # The call names its run, which is how a metric finds the run that made it. + assert call_id == "r1:flow.node" + + worker.deliver( + {"call_id": call_id, "event": "metric", "name": "loss", "value": 1.0} + ) + worker.deliver({"call_id": call_id, "event": "heartbeat"}) + worker.deliver({"call_id": call_id, "ok": True, "result": {"done": True}}) + + thread.join(timeout=5) + assert result["value"] == {"done": True} + # Liveness is not a measurement; only the metric is kept. + assert [event["event"] for event in seen] == ["metric"] + + +def test_a_failure_keeps_its_class_across_the_socket(loop): + hub = RemoteWorkerHub() + worker, socket = attach(hub, loop) + caught: list[Exception] = [] + + def call() -> None: + try: + hub.run("gpu", "flow", "node", "src", {}, {}, "flow.node", timeout=5) + except Exception as exc: + caught.append(exc) + + thread = call_in_thread(call) + assert socket.arrived.wait(5) + worker.deliver( + { + "call_id": socket.sent[0]["call_id"], + "ok": False, + "error": {"type": "ValueError", "message": "bad input", "traceback": "tb"}, + } + ) + thread.join(timeout=5) + assert type(caught[0]).__name__ == "ValueError" + assert str(caught[0]) == "bad input" + + +def test_a_worker_that_goes_away_fails_the_call_rather_than_hanging(loop): + hub = RemoteWorkerHub() + worker, socket = attach(hub, loop) + caught: list[Exception] = [] + + def call() -> None: + try: + hub.run("gpu", "flow", "node", "src", {}, {}, "flow.node", timeout=30) + except Exception as exc: + caught.append(exc) + + thread = call_in_thread(call) + assert socket.arrived.wait(5) + # Pulling the cable mid-training: the node fails, and does not wait out its + # thirty-second deadline to do it. + worker.detach() + thread.join(timeout=5) + + assert not thread.is_alive() + assert isinstance(caught[0], RemoteError) + assert "went away" in str(caught[0]) + + +def test_silence_past_the_deadline_is_a_timeout(loop): + hub = RemoteWorkerHub() + attach(hub, loop) + caught: list[Exception] = [] + + def call() -> None: + try: + hub.run("gpu", "flow", "node", "src", {}, {}, "flow.node", timeout=0.3) + except Exception as exc: + caught.append(exc) + + call_in_thread(call).join(timeout=5) + assert isinstance(caught[0], NodeTimeout) + + +def test_a_label_nothing_carries_is_named_rather_than_waited_on(loop): + hub = RemoteWorkerHub() + attach(hub, loop) + + with pytest.raises(NoWorker, match="tpu"): + hub.run("tpu", "flow", "node", "src", {}, {}, "flow.node", timeout=5) + # Compiling against a machine that is not attached is not a broken node — + # a node importing torch is correct there and missing here. + assert hub.compile("tpu", "flow", "node", "src") is None + + +def test_reattaching_replaces_the_old_socket(loop): + hub = RemoteWorkerHub() + first, _ = attach(hub, loop) + second, _ = attach(hub, loop) + + assert first.gone + assert hub.pick("gpu") is second + assert [worker.name for worker in hub.workers()] == ["gpu1"] + + +def test_cancelling_a_run_reaches_only_that_run(loop): + hub = RemoteWorkerHub() + worker, socket = attach(hub, loop) + + def call(run_id: str) -> None: + try: + hub.run( + "gpu", "flow", "node", "src", {}, {}, "flow.node", 30, run_id=run_id + ) + except Exception: + pass + + call_in_thread(lambda: call("run-a")) + assert socket.arrived.wait(5) + socket.arrived.clear() + call_in_thread(lambda: call("run-b")) + assert socket.arrived.wait(5) + + assert hub.cancel_run("run-a") == 1 + cancels = [frame for frame in socket.sent if frame.get("op") == "cancel"] + assert [frame["call_id"] for frame in cancels] == ["run-a:flow.node"] diff --git a/frontend/src/client/schemas.gen.ts b/frontend/src/client/schemas.gen.ts index 4ea7a3c..4138bfc 100644 --- a/frontend/src/client/schemas.gen.ts +++ b/frontend/src/client/schemas.gen.ts @@ -2458,6 +2458,39 @@ export const TokenSchema = { title: 'Token' } as const; +export const TokenIssuedSchema = { + properties: { + name: { + type: 'string', + title: 'Name' + }, + token: { + type: 'string', + title: 'Token' + }, + expires_days: { + type: 'integer', + title: 'Expires Days', + default: 365 + } + }, + type: 'object', + required: ['name', 'token'], + title: 'TokenIssued' +} as const; + +export const TokenRequestSchema = { + properties: { + name: { + type: 'string', + title: 'Name' + } + }, + type: 'object', + required: ['name'], + title: 'TokenRequest' +} as const; + export const TriggerRequestSchema = { properties: { values: { @@ -2879,6 +2912,55 @@ refresh_s, range_s}\`\`: it publishes \`\`{range_s, interval_s}\`\` to a flow answers with on \`\`message\`\`.` } as const; +export const WorkerInfoSchema = { + properties: { + name: { + type: 'string', + title: 'Name' + }, + labels: { + items: { + type: 'string' + }, + type: 'array', + title: 'Labels' + }, + max_parallel: { + type: 'integer', + title: 'Max Parallel', + default: 1 + }, + in_flight: { + type: 'integer', + title: 'In Flight', + default: 0 + }, + attached_at: { + type: 'number', + title: 'Attached At', + default: 0 + }, + last_seen: { + type: 'number', + title: 'Last Seen', + default: 0 + }, + python: { + type: 'string', + title: 'Python', + default: '' + }, + venv_digest: { + type: 'string', + title: 'Venv Digest', + default: '' + } + }, + type: 'object', + required: ['name'], + title: 'WorkerInfo' +} as const; + export const app__api__routes__dashboards__PublishRequestSchema = { properties: { version: { diff --git a/frontend/src/client/sdk.gen.ts b/frontend/src/client/sdk.gen.ts index 27f62ab..56b7e92 100644 --- a/frontend/src/client/sdk.gen.ts +++ b/frontend/src/client/sdk.gen.ts @@ -3,7 +3,7 @@ import type { CancelablePromise } from './core/CancelablePromise'; import { OpenAPI } from './core/OpenAPI'; import { request as __request } from './core/request'; -import type { AlertsReadAlertsConfigResponse, AlertsSaveAlertsConfigData, AlertsSaveAlertsConfigResponse, AlertsTestChannelData, AlertsTestChannelResponse, ArtifactsPutArtifactData, ArtifactsPutArtifactResponse, ArtifactsGetArtifactData, ArtifactsGetArtifactResponse, DashboardsReadDashboardsResponse, DashboardsReadDashboardData, DashboardsReadDashboardResponse, DashboardsCreateDashboardData, DashboardsCreateDashboardResponse, DashboardsSaveDashboardData, DashboardsSaveDashboardResponse, DashboardsDeleteDashboardData, DashboardsDeleteDashboardResponse, DashboardsPublishDashboardData, DashboardsPublishDashboardResponse, DashboardsDiscardDashboardDraftData, DashboardsDiscardDashboardDraftResponse, DashboardsRenameDashboardData, DashboardsRenameDashboardResponse, FlowsReadFlowsResponse, FlowsReadNodeTypesResponse, FlowsReadGraphResponse, FlowsReadLibraryResponse, FlowsDeleteSharedNodeData, FlowsDeleteSharedNodeResponse, FlowsReadFlowData, FlowsReadFlowResponse, FlowsSaveFlowData, FlowsSaveFlowResponse, FlowsDeleteFlowData, FlowsDeleteFlowResponse, FlowsPublishFlowData, FlowsPublishFlowResponse, FlowsDiscardDraftData, FlowsDiscardDraftResponse, FlowsRenameFlowData, FlowsRenameFlowResponse, FlowsReadNodeSourceData, FlowsReadNodeSourceResponse, FlowsSaveNodeSourceData, FlowsSaveNodeSourceResponse, FlowsShareNodeData, FlowsShareNodeResponse, FlowsUnshareNodeData, FlowsUnshareNodeResponse, FlowsStartFlowData, FlowsStartFlowResponse, FlowsStopFlowData, FlowsStopFlowResponse, FlowsPauseFlowData, FlowsPauseFlowResponse, FlowsResumeFlowData, FlowsResumeFlowResponse, FlowsStepFlowData, FlowsStepFlowResponse, FlowsValidateFlowData, FlowsValidateFlowResponse, FlowsRunFlowData, FlowsRunFlowResponse, FlowsTriggerNodeData, FlowsTriggerNodeResponse, FlowsCancelNodeData, FlowsCancelNodeResponse, FlowsReadFlowStateData, FlowsReadFlowStateResponse, FlowsReadMessageHistoryData, FlowsReadMessageHistoryResponse, LoginLoginAccessTokenData, LoginLoginAccessTokenResponse, LoginTestTokenResponse, LoginRecoverPasswordData, LoginRecoverPasswordResponse, LoginResetPasswordData, LoginResetPasswordResponse, LoginRecoverPasswordHtmlContentData, LoginRecoverPasswordHtmlContentResponse, MessagesReadMessagesResponse, MessagesPublishMessageData, MessagesPublishMessageResponse, MessagesReadMessageHistoryData, MessagesReadMessageHistoryResponse, ModulesReadModulesResponse, ModulesApplyModulesData, ModulesApplyModulesResponse, OauthRegisterClientData, OauthRegisterClientResponse, OauthAuthorizeValidateData, OauthAuthorizeValidateResponse, OauthAuthorizeData, OauthAuthorizeResponse, OauthTokenData, OauthTokenResponse, OauthReadClientsResponse, OauthRevokeClientData, OauthRevokeClientResponse, ObservabilityReadSummaryResponse, ObservabilityReadTimeseriesData, ObservabilityReadTimeseriesResponse, ObservabilityReadFlowRollupsData, ObservabilityReadFlowRollupsResponse, ObservabilityReadRunsData, ObservabilityReadRunsResponse, ObservabilityReadEventsData, ObservabilityReadEventsResponse, ObservabilityReadDeadLettersData, ObservabilityReadDeadLettersResponse, PrivateCreateUserData, PrivateCreateUserResponse, RunsCreateRunData, RunsCreateRunResponse, RunsCreateSweepData, RunsCreateSweepResponse, RunsReadRunsData, RunsReadRunsResponse, RunsReadRunData, RunsReadRunResponse, RunsCancelRunData, RunsCancelRunResponse, RunsReadMetricsData, RunsReadMetricsResponse, RunsCompareMetricData, RunsCompareMetricResponse, SecretsReadSecretsResponse, SecretsSaveSecretData, SecretsSaveSecretResponse, SecretsDeleteSecretData, SecretsDeleteSecretResponse, UsersReadUsersData, UsersReadUsersResponse, UsersCreateUserData, UsersCreateUserResponse, UsersReadUserMeResponse, UsersDeleteUserMeResponse, UsersUpdateUserMeData, UsersUpdateUserMeResponse, UsersUpdatePasswordMeData, UsersUpdatePasswordMeResponse, UsersRegisterUserData, UsersRegisterUserResponse, UsersReadUserByIdData, UsersReadUserByIdResponse, UsersUpdateUserData, UsersUpdateUserResponse, UsersDeleteUserData, UsersDeleteUserResponse, UtilsTestEmailData, UtilsTestEmailResponse, UtilsHealthCheckResponse, UtilsHealthResponse } from './types.gen'; +import type { AlertsReadAlertsConfigResponse, AlertsSaveAlertsConfigData, AlertsSaveAlertsConfigResponse, AlertsTestChannelData, AlertsTestChannelResponse, ArtifactsPutArtifactData, ArtifactsPutArtifactResponse, ArtifactsGetArtifactData, ArtifactsGetArtifactResponse, DashboardsReadDashboardsResponse, DashboardsReadDashboardData, DashboardsReadDashboardResponse, DashboardsCreateDashboardData, DashboardsCreateDashboardResponse, DashboardsSaveDashboardData, DashboardsSaveDashboardResponse, DashboardsDeleteDashboardData, DashboardsDeleteDashboardResponse, DashboardsPublishDashboardData, DashboardsPublishDashboardResponse, DashboardsDiscardDashboardDraftData, DashboardsDiscardDashboardDraftResponse, DashboardsRenameDashboardData, DashboardsRenameDashboardResponse, FlowsReadFlowsResponse, FlowsReadNodeTypesResponse, FlowsReadGraphResponse, FlowsReadLibraryResponse, FlowsDeleteSharedNodeData, FlowsDeleteSharedNodeResponse, FlowsReadFlowData, FlowsReadFlowResponse, FlowsSaveFlowData, FlowsSaveFlowResponse, FlowsDeleteFlowData, FlowsDeleteFlowResponse, FlowsPublishFlowData, FlowsPublishFlowResponse, FlowsDiscardDraftData, FlowsDiscardDraftResponse, FlowsRenameFlowData, FlowsRenameFlowResponse, FlowsReadNodeSourceData, FlowsReadNodeSourceResponse, FlowsSaveNodeSourceData, FlowsSaveNodeSourceResponse, FlowsShareNodeData, FlowsShareNodeResponse, FlowsUnshareNodeData, FlowsUnshareNodeResponse, FlowsStartFlowData, FlowsStartFlowResponse, FlowsStopFlowData, FlowsStopFlowResponse, FlowsPauseFlowData, FlowsPauseFlowResponse, FlowsResumeFlowData, FlowsResumeFlowResponse, FlowsStepFlowData, FlowsStepFlowResponse, FlowsValidateFlowData, FlowsValidateFlowResponse, FlowsRunFlowData, FlowsRunFlowResponse, FlowsTriggerNodeData, FlowsTriggerNodeResponse, FlowsCancelNodeData, FlowsCancelNodeResponse, FlowsReadFlowStateData, FlowsReadFlowStateResponse, FlowsReadMessageHistoryData, FlowsReadMessageHistoryResponse, LoginLoginAccessTokenData, LoginLoginAccessTokenResponse, LoginTestTokenResponse, LoginRecoverPasswordData, LoginRecoverPasswordResponse, LoginResetPasswordData, LoginResetPasswordResponse, LoginRecoverPasswordHtmlContentData, LoginRecoverPasswordHtmlContentResponse, MessagesReadMessagesResponse, MessagesPublishMessageData, MessagesPublishMessageResponse, MessagesReadMessageHistoryData, MessagesReadMessageHistoryResponse, ModulesReadModulesResponse, ModulesApplyModulesData, ModulesApplyModulesResponse, OauthRegisterClientData, OauthRegisterClientResponse, OauthAuthorizeValidateData, OauthAuthorizeValidateResponse, OauthAuthorizeData, OauthAuthorizeResponse, OauthTokenData, OauthTokenResponse, OauthReadClientsResponse, OauthRevokeClientData, OauthRevokeClientResponse, ObservabilityReadSummaryResponse, ObservabilityReadTimeseriesData, ObservabilityReadTimeseriesResponse, ObservabilityReadFlowRollupsData, ObservabilityReadFlowRollupsResponse, ObservabilityReadRunsData, ObservabilityReadRunsResponse, ObservabilityReadEventsData, ObservabilityReadEventsResponse, ObservabilityReadDeadLettersData, ObservabilityReadDeadLettersResponse, PrivateCreateUserData, PrivateCreateUserResponse, RunsCreateRunData, RunsCreateRunResponse, RunsCreateSweepData, RunsCreateSweepResponse, RunsReadRunsData, RunsReadRunsResponse, RunsReadRunData, RunsReadRunResponse, RunsCancelRunData, RunsCancelRunResponse, RunsReadMetricsData, RunsReadMetricsResponse, RunsCompareMetricData, RunsCompareMetricResponse, SecretsReadSecretsResponse, SecretsSaveSecretData, SecretsSaveSecretResponse, SecretsDeleteSecretData, SecretsDeleteSecretResponse, UsersReadUsersData, UsersReadUsersResponse, UsersCreateUserData, UsersCreateUserResponse, UsersReadUserMeResponse, UsersDeleteUserMeResponse, UsersUpdateUserMeData, UsersUpdateUserMeResponse, UsersUpdatePasswordMeData, UsersUpdatePasswordMeResponse, UsersRegisterUserData, UsersRegisterUserResponse, UsersReadUserByIdData, UsersReadUserByIdResponse, UsersUpdateUserData, UsersUpdateUserResponse, UsersDeleteUserData, UsersDeleteUserResponse, UtilsTestEmailData, UtilsTestEmailResponse, UtilsHealthCheckResponse, UtilsHealthResponse, WorkersReadWorkersResponse, WorkersIssueTokenData, WorkersIssueTokenResponse, WorkersReadRuntimeResponse } from './types.gen'; export class AlertsService { /** @@ -1841,4 +1841,58 @@ export class UtilsService { url: '/api/v1/utils/health/' }); } +} + +export class WorkersService { + /** + * Read Workers + * What is attached, and how busy it is. + * @returns WorkerInfo Successful Response + * @throws ApiError + */ + public static readWorkers(): CancelablePromise { + return __request(OpenAPI, { + method: 'GET', + url: '/api/v1/workers' + }); + } + + /** + * Issue Token + * 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. + * @param data The data for the request. + * @param data.requestBody + * @returns TokenIssued Successful Response + * @throws ApiError + */ + public static issueToken(data: WorkersIssueTokenData): CancelablePromise { + return __request(OpenAPI, { + method: 'POST', + url: '/api/v1/workers/tokens', + body: data.requestBody, + mediaType: 'application/json', + errors: { + 422: 'Validation Error' + } + }); + } + + /** + * Read Runtime + * 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. + * @returns string Successful Response + * @throws ApiError + */ + public static readRuntime(): CancelablePromise { + return __request(OpenAPI, { + method: 'GET', + url: '/api/v1/workers/runtime' + }); + } } \ No newline at end of file diff --git a/frontend/src/client/types.gen.ts b/frontend/src/client/types.gen.ts index d3457dc..2abb8c1 100644 --- a/frontend/src/client/types.gen.ts +++ b/frontend/src/client/types.gen.ts @@ -850,6 +850,16 @@ export type Token = { token_type?: string; }; +export type TokenIssued = { + name: string; + token: string; + expires_days?: number; +}; + +export type TokenRequest = { + name: string; +}; + export type TriggerRequest = { values?: { [key: string]: unknown; @@ -955,6 +965,17 @@ export type WidgetDef = { export type type = 'stat' | 'gauge' | 'chart' | 'markdown' | 'agenda' | 'notification' | 'button' | 'switch' | 'slider' | 'input' | 'dropdown'; +export type WorkerInfo = { + name: string; + labels?: Array<(string)>; + max_parallel?: number; + in_flight?: number; + attached_at?: number; + last_seen?: number; + python?: string; + venv_digest?: string; +}; + export type AlertsReadAlertsConfigResponse = (AlertsConfig); export type AlertsSaveAlertsConfigData = { @@ -1445,4 +1466,14 @@ export type UtilsHealthCheckResponse = (boolean); export type UtilsHealthResponse = ({ [key: string]: unknown; -}); \ No newline at end of file +}); + +export type WorkersReadWorkersResponse = (Array); + +export type WorkersIssueTokenData = { + requestBody: TokenRequest; +}; + +export type WorkersIssueTokenResponse = (TokenIssued); + +export type WorkersReadRuntimeResponse = (string); \ No newline at end of file