stroblmeandClaude Opus 5
2c369ac75f
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>
386 lines
14 KiB
Python
386 lines
14 KiB
Python
"""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 engine 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
|
|
# — ``fluksio_worker`` — 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 engine
|
|
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()
|