app/backend/tests/flow/test_workers.py

"""Python nodes run in a worker process, and stay there when things go wrong."""

import sys
import threading
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.controller import EmitSink
from fluksio.flow.messages import DType, MessageSpec
from fluksio.flow.nodes import Node
from fluksio.flow.nodes.base import NodeOutputError
from fluksio.flow.workers import NodeCancelled, NodeTimeout, PythonWorkerPool


@pytest.fixture
def pool() -> Iterator[PythonWorkerPool]:
    # One worker: a respawn is then provably the same slot coming back.
    worker_pool = PythonWorkerPool(python=sys.executable, size=1)
    worker_pool.start()
    yield worker_pool
    worker_pool.stop()


def _sink_node(node_id: str, port: str, dtype: DType) -> tuple[Node, list[dict]]:
    """A node declaring one streaming port, and the list its emissions land in."""
    node = Node(
        f=lambda: None,
        requires=[],
        provides=[MessageSpec(name=port, dtype=dtype, stream=True)],
        name=node_id,
    )
    node.assign_flow("demo", node_id)
    published: list[dict] = []
    node._pipeline = type(
        "Recorder",
        (),
        {
            "publish_emission": staticmethod(
                lambda _node, outputs: published.append(outputs)
            )
        },
    )()
    return node, published


def run(pool: PythonWorkerPool, code: str, node: str = "demo", **kwargs):
    return pool.run("demo", node, code, kwargs, f"demo.{node}", timeout=5)


def test_a_node_returns_its_value_and_what_it_printed(pool, capsys):
    result = run(
        pool,
        "def process(value, factor):\n"
        "    print('seen', value)\n"
        "    return {'out': value * factor}\n",
        value=21,
        factor=2,
    )
    assert result == {"out": 42}
    # The proxy writes them to stdout, which is where the engine's tee is.
    assert "seen 21" in capsys.readouterr().out


def test_a_failure_keeps_its_class_and_points_at_the_node(pool):
    with pytest.raises(Exception) as caught:
        run(pool, "def process():\n    raise ValueError('bad input')\n")

    # The engine renders a node error as "<class>: <message>", so both have to
    # survive the trip.
    assert type(caught.value).__name__ == "ValueError"
    assert str(caught.value) == "bad input"
    assert "<node demo." in caught.value.remote_traceback
    assert "ValueError: bad input" in caught.value.remote_traceback


def test_a_node_that_kills_its_worker_is_an_ordinary_error(pool):
    with pytest.raises(Exception, match="worker died"):
        run(pool, "import os\n\n\ndef process():\n    os._exit(1)\n")

    assert run(pool, "def process():\n    return {'out': 1}\n") == {"out": 1}


def test_a_node_that_runs_too_long_is_killed_and_the_pool_recovers(pool):
    started = time.monotonic()
    with pytest.raises(NodeTimeout):
        pool.run(
            "demo",
            "slow",
            "import time\n\n\ndef process():\n    time.sleep(30)\n",
            {},
            "demo.slow",
            timeout=1,
        )
    assert time.monotonic() - started < 10

    # The killed worker's slot is refilled on the next call.
    assert run(pool, "def process():\n    return {'out': 2}\n") == {"out": 2}


def test_a_slow_import_is_not_charged_to_the_nodes_timeout(pool):
    # Module-level work is what a node's imports are, and it happens once, on a
    # cold worker. Paying for it out of the per-call budget makes a node with
    # heavy imports impossible to run at all: the timeout kills the worker, so
    # the next attempt is cold again and starts over.
    code = "import time\n\ntime.sleep(1.5)\n\n\ndef process():\n    return {'out': 1}\n"
    assert pool.run("demo", "heavy", code, {}, "demo.heavy", timeout=0.5) == {"out": 1}

    # And the same worker does not pay for them a second time.
    started = time.monotonic()
    assert pool.run("demo", "heavy", code, {}, "demo.heavy", timeout=0.5) == {"out": 1}
    assert time.monotonic() - started < 1.0


def test_a_running_node_can_be_cancelled(pool):
    def stop_it() -> None:
        for _ in range(100):
            if pool.cancel("demo.slow"):
                return
            time.sleep(0.05)

    stopper = threading.Thread(target=stop_it)
    stopper.start()
    try:
        with pytest.raises(NodeCancelled):
            pool.run(
                "demo",
                "slow",
                "import time\n\n\ndef process():\n    time.sleep(30)\n",
                {},
                "demo.slow",
                timeout=30,
            )
    finally:
        stopper.join()


def test_a_result_that_is_not_json_is_refused(pool):
    with pytest.raises(Exception, match="cannot be sent back as JSON"):
        run(pool, "def process():\n    return {'out': {1, 2}}\n")


def test_compiling_reports_where_the_source_is_wrong(pool):
    assert pool.compile("demo", "broken", "def process()\n    return {}\n")
    assert pool.compile("demo", "fine", "def process():\n    return {}\n") is None


def test_a_node_imports_the_standard_library_not_the_engines_own_modules(pool):
    # The worker script lives in app/flow, which holds queue.py, secrets.py and
    # more; the interpreter would put that directory first on sys.path.
    result = run(
        pool,
        "import queue\nimport secrets\n\n\n"
        "def process():\n"
        "    return {'out': [queue.Queue().qsize(), len(secrets.token_hex(4))]}\n",
    )
    assert result == {"out": [0, 8]}


def test_the_engines_secrets_are_not_in_a_workers_environment(pool, monkeypatch):
    monkeypatch.setenv("SECRET_KEY", "not-for-nodes")
    monkeypatch.setenv("DATABASE_URL", "not-for-nodes")
    monkeypatch.setenv("FLUKSIO_HARMLESS", "fine")
    # A fresh process, so it is built from the environment set just now.
    pool.respawn_all()

    result = run(
        pool,
        "import os\n\n\n"
        "def process():\n"
        "    return {'out': [k for k in ('SECRET_KEY', 'DATABASE_URL',\n"
        "                                'FLUKSIO_HARMLESS') if k in os.environ]}\n",
    )
    assert result == {"out": ["FLUKSIO_HARMLESS"]}


def test_a_pool_can_stop_while_a_node_is_running(pool):
    # One slot, taken by a node that will not finish on its own, and a second
    # call queued behind it. The engine's node threads are not daemons, so a
    # wait here is a shutdown that never completes.
    outcomes: list[str] = []

    def call(node: str) -> None:
        try:
            pool.run(
                "demo",
                node,
                "import time\n\n\ndef process():\n    time.sleep(60)\n",
                {},
                f"demo.{node}",
                timeout=60,
            )
            outcomes.append("returned")
        except Exception as exc:
            outcomes.append(type(exc).__name__)

    busy = threading.Thread(target=call, args=("busy",))
    busy.start()
    # Let the first one take the slot, so the second is blocked acquiring it.
    time.sleep(1)
    waiting = threading.Thread(target=call, args=("waiting",))
    waiting.start()
    time.sleep(0.2)

    pool.stop()
    for thread in (busy, waiting):
        thread.join(timeout=10)
        assert not thread.is_alive()
    assert len(outcomes) == 2

    with pytest.raises(Exception, match="shutting down"):
        run(pool, "def process():\n    return {'out': 1}\n")


# -----------------------------------------------------------------------------
# Producing values before returning: a generator node
# -----------------------------------------------------------------------------


def test_a_generator_node_publishes_each_yield_and_returns_the_end(pool):
    seen = []
    result = pool.run(
        "demo",
        "train",
        "def process():\n"
        "    for step in range(3):\n"
        "        yield {'loss': 1.0 / (step + 1)}\n"
        "    return {'weights': 'w', 'final_loss': 0.25}\n",
        {},
        "demo.train",
        timeout=5,
        run_id="r1",
        on_event=seen.append,
    )

    # What it returned is the node's output; what it yielded went out as it
    # happened, on the same ports.
    assert result == {"weights": "w", "final_loss": 0.25}
    assert [event["outputs"] for event in seen] == [
        {"loss": 1.0},
        {"loss": 0.5},
        {"loss": 1 / 3},
    ]
    # Every frame says which call it belongs to, so a sweep can tell them apart.
    assert {event["call_id"] for event in seen} == {"r1:demo.train"}


def test_without_a_return_the_last_yield_is_the_result(pool):
    """The worker sends every yield; the engine is what holds the last one back."""
    node, published = _sink_node("count", "out", DType.INT)
    sink = EmitSink()
    sink.node = node

    result = sink.wrap(
        lambda: pool.run(
            "demo",
            "count",
            "def process():\n"
            "    yield {'out': 1}\n"
            "    yield {'out': 2}\n"
            "    yield {'out': 3}\n",
            {},
            "demo.count",
            timeout=5,
            on_event=sink.handle,
        )
    )()

    # The last yield is the node's output rather than an emission, so it is not
    # counted twice.
    assert result == {"out": 3}
    assert published == [{"demo.out": 1}, {"demo.out": 2}]


def test_emit_reaches_the_same_ports_from_inside_a_callback(pool):
    # A value produced somewhere a yield cannot reach — inside a framework's
    # callback — is still an output rather than a log.
    seen = []
    result = pool.run(
        "demo",
        "fit",
        "import fluksio\n"
        "def process():\n"
        "    def on_epoch(n):\n"
        "        fluksio.emit(loss=1.0 / (n + 1))\n"
        "    for epoch in range(2):\n"
        "        on_epoch(epoch)\n"
        "    return {'done': True}\n",
        {},
        "demo.fit",
        timeout=5,
        on_event=seen.append,
    )

    assert result == {"done": True}
    assert [event["outputs"] for event in seen] == [{"loss": 1.0}, {"loss": 0.5}]


def test_a_plain_function_still_just_returns(pool):
    seen = []
    assert run(pool, "def process():\n    return {'out': 7}\n") == {"out": 7}
    assert seen == []


def test_events_hold_off_the_timeout_but_silence_does_not(pool):
    # The deadline measures silence: a node yielding every 0.05s stays alive
    # well past a 0.3s timeout, which is what a two-hour training needs.
    result = pool.run(
        "demo",
        "slow",
        "import time\n"
        "def process():\n"
        "    for step in range(12):\n"
        "        time.sleep(0.05)\n"
        "        yield {'beat': step}\n"
        "    return {'done': True}\n",
        {},
        "demo.slow",
        timeout=0.3,
        run_id="r2",
        on_event=lambda _event: None,
    )
    assert result == {"done": True}

    with pytest.raises(NodeTimeout):
        pool.run(
            "demo",
            "quiet",
            "import time\ndef process():\n    time.sleep(2)\n    return {}\n",
            {},
            "demo.quiet",
            timeout=0.3,
        )


def test_cancelling_one_run_leaves_the_same_node_in_another_alone(pool):
    # Keyed by (run, node): cancelling a config of a sweep must not kill the
    # rest of it. With one slot the second run is not executing, so the check
    # is that the pool refuses to find it rather than killing the wrong worker.
    started = threading.Event()

    def hold():
        try:
            pool.run(
                "demo",
                "hold",
                "import time\ndef process():\n    time.sleep(5)\n    return {}\n",
                {},
                "demo.hold",
                timeout=10,
                run_id="run-a",
            )
        except Exception:
            pass
        finally:
            started.set()

    thread = threading.Thread(target=hold, daemon=True)
    thread.start()
    time.sleep(0.5)

    assert pool.cancel("demo.hold", run_id="run-b") is False
    assert pool.cancel("demo.hold", run_id="run-a") is True
    started.wait(timeout=5)
    thread.join(timeout=5)


def test_a_node_saves_and_loads_an_artifact(tmp_path):
    # Bytes never travel as a message: the node stores them and returns a
    # reference, which the next node opens.
    store = ArtifactStore(tmp_path / "artifacts")
    pool = PythonWorkerPool(
        python=sys.executable, size=1, env={ARTIFACT_DIR_ENV: str(store.root)}
    )
    pool.start()
    try:
        ref = pool.run(
            "demo",
            "save",
            "import fluksio\n"
            "def process():\n"
            "    return {'weights': fluksio.save_artifact(b'x' * 2048, 'w.npz')}\n",
            {},
            "demo.save",
            timeout=10,
        )["weights"]

        assert ref["size"] == 2048
        assert MessageSpec(name="weights", dtype=DType.ARTIFACT).check(ref) is None
        assert store.path(ref["digest"]) is not None

        loaded = pool.run(
            "demo",
            "load",
            "import fluksio\n"
            "def process(weights):\n"
            "    with open(fluksio.load_artifact(weights), 'rb') as f:\n"
            "        return {'size': len(f.read())}\n",
            {"weights": ref},
            "demo.load",
            timeout=10,
        )
        assert loaded == {"size": 2048}
    finally:
        pool.stop()


# -----------------------------------------------------------------------------
# No timeout at all
#
# The default: silence is a node working, not a node stuck. What still ends a
# call is the worker dying, which arrives as its pipe closing rather than as a
# deadline.
# -----------------------------------------------------------------------------


def test_a_node_with_no_timeout_runs_past_what_the_default_would_have_killed(pool):
    assert pool.run(
        "demo",
        "patient",
        "import time\n\n\ndef process():\n    time.sleep(2)\n    return {'out': 1}\n",
        {},
        "demo.patient",
        timeout=0,
    ) == {"out": 1}


def test_a_worker_that_dies_still_fails_promptly_with_no_timeout(pool):
    started = time.monotonic()
    with pytest.raises(Exception, match="worker died"):
        pool.run(
            "demo",
            "doomed",
            "import os\n\n\ndef process():\n    os._exit(1)\n",
            {},
            "demo.doomed",
            timeout=0,
        )
    # Not waiting out a poll interval: the pipe closing is what wakes the read.
    assert time.monotonic() - started < 5

    assert run(pool, "def process():\n    return {'out': 3}\n") == {"out": 3}


def test_a_node_with_no_timeout_can_still_be_cancelled(pool):
    def stop_it() -> None:
        for _ in range(100):
            if pool.cancel("demo.slow"):
                return
            time.sleep(0.05)

    stopper = threading.Thread(target=stop_it)
    stopper.start()
    try:
        with pytest.raises(NodeCancelled):
            pool.run(
                "demo",
                "slow",
                "import time\n\n\ndef process():\n    time.sleep(30)\n",
                {},
                "demo.slow",
                timeout=0,
            )
    finally:
        stopper.join()


def test_an_emission_on_an_undeclared_port_fails_the_call(pool):
    """The engine's sink raises, and that has to reach the node's author.

    At the *first* yield, which is where the mistake is: the worker sends each
    one as it happens, so a mistyped port is checked before the node goes back
    to work rather than when the yield after it arrives.
    """
    node, _ = _sink_node("gen", "loss", DType.FLOAT)
    sink = EmitSink()
    sink.node = node

    started = time.monotonic()
    with pytest.raises(NodeOutputError, match="lss"):
        sink.wrap(
            lambda: pool.run(
                "demo",
                "gen",
                "import time\n\n\ndef process():\n"
                "    yield {'lss': 1.0}\n"
                "    time.sleep(30)\n"
                "    yield {'loss': 0.5}\n",
                {},
                "demo.gen",
                timeout=0,
                on_event=sink.handle,
            )
        )()
    # It did not wait out the node: the failure stopped the call.
    assert time.monotonic() - started < 10

    # The worker was retired rather than left mid-generator, so the slot works.
    assert run(pool, "def process():\n    return {'out': 4}\n") == {"out": 4}


# -----------------------------------------------------------------------------
# Pools per environment
#
# A thread cap is read when the process imports the library, so a warm worker
# cannot be told a different one. A node that declared resources therefore gets
# a pool started with its own environment, and nodes deriving the same one
# share it rather than paying for a cold start each.
# -----------------------------------------------------------------------------


def test_a_declared_environment_gets_a_pool_of_its_own(pool):
    child = pool.for_env({"OMP_NUM_THREADS": "2"})

    assert child is not pool
    assert child.env["OMP_NUM_THREADS"] == "2"
    # The same environment is the same pool: warm workers are the point.
    assert pool.for_env({"OMP_NUM_THREADS": "2"}) is child
    assert pool.for_env({"OMP_NUM_THREADS": "3"}) is not child
    # Nothing to derive means the shared pool, with no second process anywhere.
    assert pool.for_env({}) is pool


def test_a_child_worker_is_started_with_what_it_was_given(pool):
    child = pool.for_env({"OMP_NUM_THREADS": "2"})

    seen = child.run(
        "demo",
        "env",
        "import os\n\n\ndef process():\n"
        "    return {'threads': os.environ.get('OMP_NUM_THREADS', '')}\n",
        {},
        "demo.env",
        timeout=10,
    )
    assert seen == {"threads": "2"}


def test_retiring_workers_reaches_the_children(pool):
    child = pool.for_env({"OMP_NUM_THREADS": "2"})
    before = child._generation

    pool.respawn_all()

    assert child._generation > before


def test_cancelling_reaches_a_node_running_in_a_child(pool):
    child = pool.for_env({"OMP_NUM_THREADS": "2"})
    started = threading.Event()
    failed: list[Exception] = []

    def call() -> None:
        try:
            child.run(
                "demo",
                "slow",
                "import time\n\n\ndef process():\n"
                "    print('up', flush=True)\n"
                "    time.sleep(30)\n"
                "    return {'out': 1}\n",
                {},
                "demo.slow",
                timeout=0,
                on_event=lambda _event: started.set(),
            )
        except Exception as exc:  # noqa: BLE001 — the point is that it stopped
            failed.append(exc)

    thread = threading.Thread(target=call)
    thread.start()
    # The node has no events, so wait for the process rather than for a frame.
    deadline = time.monotonic() + 10
    while time.monotonic() < deadline and not child._running:
        time.sleep(0.01)

    assert pool.cancel("demo.slow"), "the parent must find a child's node"
    thread.join(timeout=10)
    assert failed