app/backend/app/flow/worker_main.py

"""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
    ``app.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 {}), params=request.get("params") 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()