Split the worker into a distribution of its own
A cluster or GPU host installs `pip install fluksio-worker` and gets the agent and the runner, not psycopg, numpy and the MCP SDK. The engine depends on it as a workspace member, so the file it launches node code with is the same file a remote worker runs — which is what keeps a node unable to tell the difference. Copying the two files by hand still works. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
This commit is contained in:
Claude Opus 5
@@ -33,6 +33,11 @@ COPY ./backend/pyproject.toml ./backend/alembic.ini /app/backend/
|
||||
|
||||
COPY ./backend/fluksio /app/backend/fluksio
|
||||
|
||||
# The worker is a workspace member of its own, so the engine's sync needs it
|
||||
# present to resolve the dependency on it. It is also what the engine runs
|
||||
# node code with.
|
||||
COPY ./worker /app/worker
|
||||
|
||||
# Sync the project
|
||||
# Ref: https://docs.astral.sh/uv/guides/integration/docker/#intermediate-layers
|
||||
RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
|
||||
@@ -15,12 +15,12 @@ from typing import Any
|
||||
|
||||
from fastapi import APIRouter, Depends, HTTPException, Request, WebSocket
|
||||
from fastapi.responses import PlainTextResponse
|
||||
from fluksio_worker import worker_main
|
||||
from jwt.exceptions import InvalidTokenError
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
from fluksio.api.deps import get_current_active_superuser, get_current_user
|
||||
from fluksio.core import security
|
||||
from fluksio.flow import worker_main
|
||||
from fluksio.flow.remote import PROTOCOL, RemoteWorker, RemoteWorkerHub
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
@@ -103,7 +103,7 @@ 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.
|
||||
standard library only, and with nothing of the engine importable in it.
|
||||
"""
|
||||
return worker_main.__file__ and open(worker_main.__file__).read()
|
||||
|
||||
|
||||
@@ -20,6 +20,7 @@ from typing import Any, cast
|
||||
|
||||
from fastapi import FastAPI
|
||||
from fastapi.concurrency import run_in_threadpool
|
||||
from fluksio_worker.worker_main import load_function
|
||||
|
||||
from fluksio.core.config import settings
|
||||
from fluksio.flow.alerts import AlertManager
|
||||
@@ -60,7 +61,6 @@ from fluksio.flow.secrets import SecretNotFound, resolve_params
|
||||
from fluksio.flow.state import MemoryState, StateBackend
|
||||
from fluksio.flow.store import LIB_DIR, FlowNotFound, FlowStore, LibNotFound
|
||||
from fluksio.flow.supervision import Supervisor
|
||||
from fluksio.flow.worker_main import load_function
|
||||
from fluksio.flow.workers import PythonWorkerPool
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -1,385 +0,0 @@
|
||||
"""One user-code process: compiles a node's source once, then runs it on demand.
|
||||
|
||||
This runs under the *user* venv's interpreter, so a node's imports resolve
|
||||
against what the Modules page installed rather than against the engine's own
|
||||
packages. Nothing of the app is importable here — the file is handed to the
|
||||
interpreter by path and is deliberately standard library only.
|
||||
|
||||
The protocol is one JSON object per line: a request arrives on stdin, the reply
|
||||
goes out on a private duplicate of fd 1 taken before anything else can write to
|
||||
it. fd 1 itself is pointed at stderr, so a stray ``write(1, ...)`` — from a
|
||||
native library, or a node printing during an import — lands in the server log
|
||||
instead of corrupting the reply stream.
|
||||
|
||||
A node may also send lines back *while* it is still running: anything carrying
|
||||
an ``event`` key is a report rather than the answer, and the engine keeps
|
||||
reading. That is what makes a training curve visible during the hours it takes
|
||||
to draw, and what tells the engine a long node is alive rather than hung —
|
||||
each event resets its deadline, so the timeout measures silence rather than
|
||||
duration.
|
||||
|
||||
What travels that way is not a log. A node that produces values over time is a
|
||||
generator: every ``yield`` is a dict keyed by output port, published the moment
|
||||
it happens, and whatever the generator returns at the end is the node's result.
|
||||
Nothing leaves a node except through a port it declared, which is the whole
|
||||
point — a number worth keeping is an output, not a side effect. Where a yield
|
||||
cannot reach, because the value comes from inside somebody else's callback,
|
||||
``fluksio.emit(loss=0.3)`` writes the same ports the same way.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import sys
|
||||
|
||||
# Handed to the interpreter by path, so CPython puts this file's own directory
|
||||
# — the engine's ``app/flow`` — on ``sys.path[0]``. A node doing ``import
|
||||
# queue`` would then get the engine's queue module rather than the standard
|
||||
# library's. Drop it before anything else can import. (``-P`` and
|
||||
# ``PYTHONSAFEPATH`` do this at startup, but both are 3.11+.)
|
||||
_HERE = os.path.dirname(os.path.abspath(__file__))
|
||||
sys.path[:] = [p for p in sys.path if os.path.abspath(p or ".") != _HERE]
|
||||
|
||||
import contextlib
|
||||
import hashlib
|
||||
import inspect
|
||||
import io
|
||||
import json
|
||||
import tempfile
|
||||
import time
|
||||
import traceback
|
||||
import urllib.parse
|
||||
import urllib.request
|
||||
from collections.abc import Callable
|
||||
from types import ModuleType
|
||||
from typing import Any, cast
|
||||
|
||||
#: Enough of a print to debug with. A node printing in a loop must not fill the
|
||||
#: pipe or the reply.
|
||||
MAX_LOG = 16 * 1024
|
||||
|
||||
#: Where the artifact store is, from this worker's point of view: a directory
|
||||
#: when it shares the engine's filesystem, a URL when it does not.
|
||||
ARTIFACT_DIR_ENV = "FLUKSIO_ARTIFACT_DIR"
|
||||
ARTIFACT_URL_ENV = "FLUKSIO_ARTIFACT_URL"
|
||||
ARTIFACT_TOKEN_ENV = "FLUKSIO_ARTIFACT_TOKEN"
|
||||
ARTIFACT_CACHE_ENV = "FLUKSIO_ARTIFACT_CACHE"
|
||||
ARTIFACT_TIMEOUT_S = 300
|
||||
|
||||
#: The reply channel, opened by ``main``. Also what an event line goes down.
|
||||
_RPC: Any = None
|
||||
#: The call being served, so an event can say which one it belongs to.
|
||||
_CALL_ID = ""
|
||||
|
||||
|
||||
def _emit(event: dict[str, Any]) -> None:
|
||||
"""Send one line back without ending the call."""
|
||||
if _RPC is None:
|
||||
return
|
||||
event["call_id"] = _CALL_ID
|
||||
event["ts"] = time.time()
|
||||
_RPC.write(json.dumps(event) + "\n")
|
||||
_RPC.flush()
|
||||
|
||||
|
||||
class _Reporter(ModuleType):
|
||||
"""``import fluksio`` — the parts of a node's job that need the engine."""
|
||||
|
||||
def emit(self, **ports: Any) -> None:
|
||||
"""Publish on this node's output ports without returning yet.
|
||||
|
||||
Yielding is the better way to say this and should be preferred; use
|
||||
this where a yield cannot reach — inside a training framework's
|
||||
callback, say, which calls you rather than the other way round. It is
|
||||
the same publication either way: the values go to the ports the node
|
||||
declared, and are checked against them.
|
||||
"""
|
||||
_emit({"event": "emit", "outputs": dict(ports)})
|
||||
|
||||
def save_artifact(
|
||||
self,
|
||||
source: Any,
|
||||
name: str = "",
|
||||
media_type: str = "application/octet-stream",
|
||||
) -> dict[str, Any]:
|
||||
"""Store bytes or a file and return the reference to return onward.
|
||||
|
||||
Return the reference from an ``artifact`` port: a checkpoint is far too
|
||||
big to be a message, and the reference is what the next node opens.
|
||||
"""
|
||||
if isinstance(source, (bytes, bytearray)):
|
||||
data = bytes(source)
|
||||
name = name or "artifact.bin"
|
||||
else:
|
||||
path = str(source)
|
||||
with open(path, "rb") as handle:
|
||||
data = handle.read()
|
||||
name = name or os.path.basename(path)
|
||||
return _store_bytes(data, name, media_type)
|
||||
|
||||
def load_artifact(self, ref: dict[str, Any]) -> str:
|
||||
"""Fetch an artifact and hand back a local path to read it from."""
|
||||
digest = str((ref or {}).get("digest") or "")
|
||||
if not digest.startswith("sha256:"):
|
||||
raise ValueError("not an artifact reference")
|
||||
return _fetch(digest)
|
||||
|
||||
|
||||
def _store_bytes(data: bytes, name: str, media_type: str) -> dict[str, Any]:
|
||||
"""Write to the artifact store, whichever end of it this worker can see.
|
||||
|
||||
A worker in the engine's own container writes the file; one on another host
|
||||
puts it over HTTP. Node code cannot tell the difference, which is the point
|
||||
— the same flow runs either place. The hashing is repeated from
|
||||
``fluksio.flow.artifacts`` rather than imported, because nothing of the app is
|
||||
importable here.
|
||||
"""
|
||||
digest = "sha256:" + hashlib.sha256(data).hexdigest()
|
||||
directory = os.environ.get(ARTIFACT_DIR_ENV)
|
||||
if directory:
|
||||
target = os.path.join(directory, digest[7:9], digest[7:])
|
||||
if not os.path.exists(target):
|
||||
os.makedirs(os.path.dirname(target), exist_ok=True)
|
||||
# Written beside the target and moved, so a reader never opens a
|
||||
# half-written artifact.
|
||||
handle, temporary = tempfile.mkstemp(dir=os.path.dirname(target))
|
||||
with os.fdopen(handle, "wb") as out:
|
||||
out.write(data)
|
||||
os.replace(temporary, target)
|
||||
else:
|
||||
_put_over_http(data, name, media_type)
|
||||
return {
|
||||
"digest": digest,
|
||||
"size": len(data),
|
||||
"media_type": media_type,
|
||||
"name": name,
|
||||
}
|
||||
|
||||
|
||||
def _put_over_http(data: bytes, name: str, media_type: str) -> None:
|
||||
base = os.environ.get(ARTIFACT_URL_ENV)
|
||||
if not base:
|
||||
raise RuntimeError(
|
||||
"this worker has no artifact store — neither "
|
||||
f"{ARTIFACT_DIR_ENV} nor {ARTIFACT_URL_ENV} is set"
|
||||
)
|
||||
query = urllib.parse.urlencode({"name": name, "media_type": media_type})
|
||||
request = urllib.request.Request(
|
||||
f"{base.rstrip('/')}?{query}", data=data, method="PUT"
|
||||
)
|
||||
_authorize(request)
|
||||
with urllib.request.urlopen(request, timeout=ARTIFACT_TIMEOUT_S):
|
||||
pass
|
||||
|
||||
|
||||
def _fetch(digest: str) -> str:
|
||||
"""A local path holding this artifact's bytes, downloading it if needed."""
|
||||
directory = os.environ.get(ARTIFACT_DIR_ENV)
|
||||
if directory:
|
||||
target = os.path.join(directory, digest[7:9], digest[7:])
|
||||
if not os.path.exists(target):
|
||||
raise FileNotFoundError(digest)
|
||||
return target
|
||||
|
||||
# Cached by digest: content-addressing means a file once fetched is never
|
||||
# stale, so a sweep pulls a shared input across the network once.
|
||||
cache = os.environ.get(ARTIFACT_CACHE_ENV) or os.path.join(
|
||||
tempfile.gettempdir(), "fluksio-artifacts"
|
||||
)
|
||||
target = os.path.join(cache, digest[7:])
|
||||
if os.path.exists(target):
|
||||
return target
|
||||
base = os.environ.get(ARTIFACT_URL_ENV)
|
||||
if not base:
|
||||
raise RuntimeError(f"this worker cannot reach the artifact store: {digest}")
|
||||
os.makedirs(cache, exist_ok=True)
|
||||
request = urllib.request.Request(f"{base.rstrip('/')}/{digest}")
|
||||
_authorize(request)
|
||||
handle, temporary = tempfile.mkstemp(dir=cache)
|
||||
with urllib.request.urlopen(request, timeout=ARTIFACT_TIMEOUT_S) as response:
|
||||
with os.fdopen(handle, "wb") as out:
|
||||
while chunk := response.read(1024 * 1024):
|
||||
out.write(chunk)
|
||||
os.replace(temporary, target)
|
||||
return target
|
||||
|
||||
|
||||
def _authorize(request: Any) -> None:
|
||||
token = os.environ.get(ARTIFACT_TOKEN_ENV)
|
||||
if token:
|
||||
request.add_header("Authorization", f"Bearer {token}")
|
||||
|
||||
|
||||
def _install_reporter() -> None:
|
||||
"""Put ``fluksio`` on the import path of every node this worker runs."""
|
||||
module = _Reporter("fluksio")
|
||||
module.__doc__ = "Report metrics and progress from inside a node."
|
||||
sys.modules["fluksio"] = module
|
||||
|
||||
|
||||
def load_function(flow: str, node_id: str, code: str) -> Callable[..., Any]:
|
||||
"""Compile a node's source and return the function to run.
|
||||
|
||||
A node file defines ``process(...)``; if it defines exactly one public
|
||||
function under another name, that one is used.
|
||||
"""
|
||||
digest = hashlib.md5(code.encode()).hexdigest()[:8]
|
||||
module_name = f"_fluksio_node_{flow}_{node_id}_{digest}"
|
||||
|
||||
module = ModuleType(module_name)
|
||||
module.__dict__["__name__"] = module_name
|
||||
sys.modules[module_name] = module
|
||||
exec(compile(code, f"<node {flow}.{node_id}>", "exec"), module.__dict__)
|
||||
|
||||
if callable(getattr(module, "process", None)):
|
||||
return cast(Callable[..., Any], module.process)
|
||||
|
||||
functions = [
|
||||
value
|
||||
for name, value in vars(module).items()
|
||||
if callable(value)
|
||||
and not name.startswith("_")
|
||||
and getattr(value, "__module__", None) == module_name
|
||||
]
|
||||
if len(functions) == 1:
|
||||
return cast(Callable[..., Any], functions[0])
|
||||
raise ValueError(
|
||||
"Define a function named 'process' — this file has "
|
||||
f"{len(functions)} functions to choose from."
|
||||
)
|
||||
|
||||
|
||||
class _Capped(io.StringIO):
|
||||
"""Keeps the first ``MAX_LOG`` characters and quietly drops the rest."""
|
||||
|
||||
def write(self, text: str) -> int:
|
||||
room = MAX_LOG - self.tell()
|
||||
if room > 0:
|
||||
super().write(text[:room])
|
||||
return len(text)
|
||||
|
||||
|
||||
def _short_error(exc: BaseException) -> str:
|
||||
"""One line a node author can act on — the engine's rule, applied here."""
|
||||
if isinstance(exc, SyntaxError):
|
||||
# Its own message already names the compiled file, which is noise here.
|
||||
return f"{type(exc).__name__}: {exc.msg} (line {exc.lineno})"
|
||||
|
||||
frames = [
|
||||
frame
|
||||
for frame in traceback.extract_tb(exc.__traceback__)
|
||||
if frame.filename.startswith("<node ")
|
||||
]
|
||||
where = f" (line {frames[-1].lineno})" if frames else ""
|
||||
return f"{type(exc).__name__}: {exc}{where}"
|
||||
|
||||
|
||||
def _node_traceback(exc: BaseException) -> str:
|
||||
"""The traceback from the node's own code onward; the rest is this loop."""
|
||||
frames = traceback.extract_tb(exc.__traceback__)
|
||||
start = next(
|
||||
(i for i, frame in enumerate(frames) if frame.filename.startswith("<node ")),
|
||||
0,
|
||||
)
|
||||
return "".join(
|
||||
["Traceback (most recent call last):\n"]
|
||||
+ traceback.format_list(frames[start:])
|
||||
+ traceback.format_exception_only(type(exc), exc)
|
||||
)
|
||||
|
||||
|
||||
def _handle(request: dict[str, Any], cache: dict[tuple[str, str], Any]) -> Any:
|
||||
"""Compile if needed, then run — the part that may raise the user's error."""
|
||||
flow, node = request["flow"], request["node"]
|
||||
source = request.get("source") or ""
|
||||
digest = hashlib.md5(source.encode()).hexdigest()
|
||||
|
||||
cached = cache.get((flow, node))
|
||||
if cached is None or cached[0] != digest:
|
||||
cache[(flow, node)] = (digest, load_function(flow, node, source))
|
||||
function = cache[(flow, node)][1]
|
||||
|
||||
if request["op"] == "compile":
|
||||
return None
|
||||
|
||||
result = function(**(request.get("kwargs") or {}))
|
||||
if inspect.isgenerator(result):
|
||||
result = _drain(result)
|
||||
try:
|
||||
json.dumps(result)
|
||||
except (TypeError, ValueError):
|
||||
# The typed-message contract only carries JSON, and a node running out
|
||||
# of process is where that stops being a formality.
|
||||
raise ValueError(
|
||||
f"returned {type(result).__name__}, which cannot be sent back as "
|
||||
"JSON — return numbers, strings, booleans, lists or dicts."
|
||||
) from None
|
||||
return result
|
||||
|
||||
|
||||
def _drain(generator: Any) -> Any:
|
||||
"""Run a generator node, publishing each yield as it happens.
|
||||
|
||||
Two shapes work, and they mean the same thing. Yield throughout and
|
||||
``return`` the result at the end, which is the explicit one; or just yield,
|
||||
and the last one is the result. Either way what the node *produces over
|
||||
time* leaves through its ports while it is still running, and what it
|
||||
*ends up with* is its return value.
|
||||
"""
|
||||
pending: Any = None
|
||||
have_pending = False
|
||||
try:
|
||||
while True:
|
||||
value = next(generator)
|
||||
# Held one behind: until the next yield arrives this might be the
|
||||
# last one, and the last one is the result rather than an emission.
|
||||
if have_pending:
|
||||
_emit({"event": "emit", "outputs": pending})
|
||||
pending, have_pending = value, True
|
||||
except StopIteration as stop:
|
||||
if stop.value is not None:
|
||||
# It returned something, so every yield was an emission.
|
||||
if have_pending:
|
||||
_emit({"event": "emit", "outputs": pending})
|
||||
return stop.value
|
||||
return pending if have_pending else None
|
||||
|
||||
|
||||
def main() -> None:
|
||||
global _RPC, _CALL_ID
|
||||
|
||||
_RPC = os.fdopen(os.dup(1), "w")
|
||||
# Everything the node writes to the real stdout now goes to the server log.
|
||||
os.dup2(2, 1)
|
||||
_install_reporter()
|
||||
|
||||
cache: dict[tuple[str, str], Any] = {}
|
||||
for line in sys.stdin:
|
||||
if not line.strip():
|
||||
continue
|
||||
request = json.loads(line)
|
||||
_CALL_ID = str(request.get("call_id") or "")
|
||||
captured = _Capped()
|
||||
response: dict[str, Any] = {"call_id": _CALL_ID}
|
||||
try:
|
||||
with (
|
||||
contextlib.redirect_stdout(captured),
|
||||
contextlib.redirect_stderr(captured),
|
||||
):
|
||||
response["result"] = _handle(request, cache)
|
||||
response["ok"] = True
|
||||
except Exception as exc:
|
||||
response["ok"] = False
|
||||
response["error"] = {
|
||||
"type": type(exc).__name__,
|
||||
"message": str(exc),
|
||||
"short": _short_error(exc),
|
||||
"traceback": _node_traceback(exc),
|
||||
}
|
||||
response["logs"] = captured.getvalue()
|
||||
_RPC.write(json.dumps(response) + "\n")
|
||||
_RPC.flush()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -26,11 +26,15 @@ from collections.abc import Callable
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from fluksio_worker import worker_main as _worker_main
|
||||
|
||||
from fluksio.flow.events import EventBus
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
WORKER_MAIN = Path(__file__).with_name("worker_main.py")
|
||||
#: The runner, in the worker distribution — the same file a remote worker runs,
|
||||
#: so a node cannot tell which kind of worker it is on.
|
||||
WORKER_MAIN = Path(_worker_main.__file__)
|
||||
|
||||
#: Importing what a node needs can be slow the first time; compiling is not
|
||||
#: something a person is watching a spinner for.
|
||||
|
||||
@@ -8,6 +8,7 @@ from fastapi import FastAPI
|
||||
from fastapi.concurrency import run_in_threadpool
|
||||
from fastapi.responses import JSONResponse
|
||||
from fastapi.routing import APIRoute
|
||||
from fluksio_worker.worker_main import ARTIFACT_DIR_ENV
|
||||
from starlette.middleware.cors import CORSMiddleware
|
||||
|
||||
from fluksio.api.main import api_router
|
||||
@@ -33,7 +34,6 @@ from fluksio.flow.secrets import init_secrets
|
||||
from fluksio.flow.state import MemoryState, RedisState, StateBackend
|
||||
from fluksio.flow.store import FlowStore
|
||||
from fluksio.flow.watchdog import LoopWatchdog
|
||||
from fluksio.flow.worker_main import ARTIFACT_DIR_ENV
|
||||
from fluksio.flow.workers import PythonWorkerPool
|
||||
|
||||
|
||||
|
||||
@@ -1,308 +0,0 @@
|
||||
"""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
|
||||
|
||||
|
||||
def _worker_main() -> Path:
|
||||
"""The user-code runner: beside this file once deployed, or in a checkout.
|
||||
|
||||
On a worker host the two files sit together, which is what the install
|
||||
instructions say. Run straight from a clone and the runner is one directory
|
||||
over, in the engine's own package — worth finding, so trying this out does
|
||||
not start with copying files around.
|
||||
"""
|
||||
here = Path(__file__).resolve().parent
|
||||
for candidate in (here / "worker_main.py", here.parent / "flow" / "worker_main.py"):
|
||||
if candidate.exists():
|
||||
return candidate
|
||||
return here / "worker_main.py"
|
||||
|
||||
|
||||
WORKER_MAIN = _worker_main()
|
||||
|
||||
|
||||
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())
|
||||
@@ -28,8 +28,12 @@ dependencies = [
|
||||
"influxdb-client[async]>=1.40.0",
|
||||
"croniter>=1.3.0",
|
||||
"mcp>=1.29,<2",
|
||||
"fluksio-worker>=0.1,<0.2",
|
||||
]
|
||||
|
||||
[tool.uv.sources]
|
||||
fluksio-worker = { workspace = true }
|
||||
|
||||
[dependency-groups]
|
||||
dev = [
|
||||
"pytest<8.0.0,>=7.4.3",
|
||||
@@ -42,6 +46,11 @@ dev = [
|
||||
[tool.hatch.build.targets.wheel]
|
||||
packages = ["fluksio"]
|
||||
|
||||
# Named rather than excluded: this directory also holds a working tree's
|
||||
# runtime data — `flow-data/` with the user venv in it — which is not source.
|
||||
[tool.hatch.build.targets.sdist]
|
||||
include = ["fluksio", "tests", "scripts", "alembic.ini", "pyproject.toml", "README.md"]
|
||||
|
||||
[build-system]
|
||||
requires = ["hatchling"]
|
||||
build-backend = "hatchling.build"
|
||||
@@ -92,12 +101,6 @@ ignore = [
|
||||
# Printing is what this one is about: node code is user code, and `print` is
|
||||
# how it says things.
|
||||
"tests/flow/test_logs.py" = ["ARG001", "T201"]
|
||||
# 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.
|
||||
"fluksio/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.
|
||||
"fluksio/worker/fluksio_worker.py" = ["T201"]
|
||||
|
||||
[tool.ruff.lint.pyupgrade]
|
||||
# Preserve types, even if a file imports `from __future__ import annotations`.
|
||||
|
||||
@@ -6,10 +6,10 @@ import time
|
||||
from collections.abc import Iterator
|
||||
|
||||
import pytest
|
||||
from fluksio_worker.worker_main import ARTIFACT_DIR_ENV
|
||||
|
||||
from fluksio.flow.artifacts import ArtifactStore
|
||||
from fluksio.flow.messages import DType, MessageSpec
|
||||
from fluksio.flow.worker_main import ARTIFACT_DIR_ENV
|
||||
from fluksio.flow.workers import NodeCancelled, NodeTimeout, PythonWorkerPool
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user