stroblmeandClaude Opus 5
647644ebbd
The worker held each yield one behind, because the last one is the node's result when the generator returns nothing of its own. Only the engine knows what ports a node declared, so the check happened when the *next* yield arrived — a pass late, which for a training loop is however long one epoch takes. The worker now sends every yield as it happens and returns whatever its generator returned; EmitSink holds the last one back and decides at the end of the call what it was. Old "emit" frames are still handled, so a remote agent that has not been restarted keeps working. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
459 lines
18 KiB
Python
459 lines
18 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)
|
|
|
|
# ---------------------------------------------------------------------
|
|
# The authoring API, inert.
|
|
#
|
|
# A node generated by `fluksio sync` imports the caller's own module, and
|
|
# that module says `from fluksio import Port, node, Flow` at the top —
|
|
# which, in here, is this. The declarations were read at sync time and are
|
|
# already in the flow document, so what they have to do now is import
|
|
# without doing anything: the decorators hand the function back, and `flow`
|
|
# builds nothing.
|
|
# ---------------------------------------------------------------------
|
|
|
|
def Port(self, *args: Any, **kwargs: Any) -> Any: # noqa: N802
|
|
"""A port declaration, already read by `fluksio sync`."""
|
|
return _Declared()
|
|
|
|
def node(self, *args: Any, **kwargs: Any) -> Any:
|
|
"""The decorator, which here gives the function straight back."""
|
|
return lambda fn: fn
|
|
|
|
def use(self, *args: Any, **kwargs: Any) -> Any:
|
|
"""One use of a node in a flow, already read by `fluksio sync`."""
|
|
return _Declared()
|
|
|
|
def Flow(self, *args: Any, **kwargs: Any) -> Any: # noqa: N802
|
|
"""A flow declaration, already read by `fluksio sync`."""
|
|
return _Declared()
|
|
|
|
|
|
class _Declared:
|
|
"""Stands in for a declaration whose work was done before the run.
|
|
|
|
Tolerant on purpose: a module may keep one at module level and touch it in
|
|
ways a node never exercises, and none of that should fail an import.
|
|
"""
|
|
|
|
def __getattr__(self, name: str) -> Any:
|
|
if name.startswith("__"):
|
|
raise AttributeError(name)
|
|
return _Declared()
|
|
|
|
def __call__(self, *args: Any, **kwargs: Any) -> Any:
|
|
return _Declared()
|
|
|
|
def __repr__(self) -> str:
|
|
return "<fluksio declaration>"
|
|
|
|
|
|
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()
|
|
|
|
if request["op"] == "compile" and not request.get("keep", True):
|
|
# A question about a draft, not a load: does this source compile? The
|
|
# answer is all the editor wants, and keeping it would throw away the
|
|
# published source this worker is still serving calls from.
|
|
load_function(flow, node, source)
|
|
return None
|
|
|
|
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)
|
|
return result
|
|
|
|
|
|
def _drain(generator: Any) -> Any:
|
|
"""Run a generator node, sending each yield the moment 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. Which of the two a node is cannot be known
|
|
until it ends, so the last yield has to be held back somewhere — and that
|
|
somewhere is the engine, which is the only side that knows what ports the
|
|
node declared. Holding it here instead cost a pass: a yield naming a port
|
|
that does not exist was only checked once the *next* one arrived.
|
|
"""
|
|
try:
|
|
while True:
|
|
_emit({"event": "yield", "outputs": next(generator)})
|
|
except StopIteration as stop:
|
|
return stop.value
|
|
|
|
|
|
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()
|
|
# ``id`` back untouched: the engine matches it against what it sent, so
|
|
# a pipe that has slipped a call is caught rather than handing one node
|
|
# another node's answer.
|
|
response: dict[str, Any] = {"id": request.get("id"), "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()
|
|
try:
|
|
# allow_nan=False: a bare NaN is what json.dumps would write, and
|
|
# nothing downstream can read it back — the ports refuse one too,
|
|
# but this is the crossing, so refuse it where it is still the
|
|
# node's own reply rather than a row somebody queries later.
|
|
reply = json.dumps(response, allow_nan=False)
|
|
except (TypeError, ValueError, RecursionError):
|
|
# Encoding the reply is also the check that the node returned
|
|
# something the typed-message contract can carry — the result is
|
|
# the only part of this dict a node controls, so it is the only
|
|
# part that can fail. Doing it once is why there is no separate
|
|
# dumps of the result above.
|
|
result = response.pop("result", None)
|
|
try:
|
|
json.dumps(result)
|
|
except (TypeError, ValueError, RecursionError):
|
|
trouble = (
|
|
f"returned {type(result).__name__}, which cannot be sent "
|
|
"back as JSON — return numbers, strings, booleans, lists "
|
|
"or dicts."
|
|
)
|
|
else:
|
|
# It encodes with NaN allowed, so that is what is wrong with it.
|
|
trouble = (
|
|
"returned a NaN or an infinity, which JSON cannot carry — "
|
|
"publish None, or a number that says the measurement had "
|
|
"nothing in it."
|
|
)
|
|
response["ok"] = False
|
|
response["error"] = {
|
|
"type": "ValueError",
|
|
"message": trouble,
|
|
"short": f"ValueError: {trouble}",
|
|
"traceback": "",
|
|
}
|
|
reply = json.dumps(response)
|
|
_RPC.write(reply + "\n")
|
|
_RPC.flush()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|