app/worker/fluksio_worker/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 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
#: How much of an artifact is held in memory at a time, matching the engine's
#: own store. Repeated rather than imported: nothing of the engine is here.
CHUNK = 1024 * 1024

#: 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)):
            return _store_bytes(bytes(source), name or "artifact.bin", media_type)
        path = str(source)
        return _store_file(path, name or os.path.basename(path), 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 _store_file(path: str, name: str, media_type: str) -> dict[str, Any]:
    """The same, for a file already on disk — read a chunk at a time.

    A video segment or a checkpoint is as likely to be handed over as a path as
    as bytes, and reading one into memory to hash it defeats the point of it
    being a file.
    """
    digester = hashlib.sha256()
    size = 0
    with open(path, "rb") as handle:
        for chunk in iter(lambda: handle.read(CHUNK), b""):
            digester.update(chunk)
            size += len(chunk)
    digest = "sha256:" + digester.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)
            handle_fd, temporary = tempfile.mkstemp(dir=os.path.dirname(target))
            with os.fdopen(handle_fd, "wb") as out, open(path, "rb") as source:
                for chunk in iter(lambda: source.read(CHUNK), b""):
                    out.write(chunk)
            os.replace(temporary, target)
    else:
        with open(path, "rb") as handle:
            _put_over_http(handle, name, media_type, size)
    return {"digest": digest, "size": size, "media_type": media_type, "name": name}


def _put_over_http(
    data: Any, name: str, media_type: str, length: int | None = None
) -> 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"
    )
    if length is not None:
        # urllib streams a file object only when it does not have to measure it.
        request.add_header("Content-Length", str(length))
    _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(CHUNK):
                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()