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.messages import DType, MessageSpec
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 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_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):
    seen = []
    result = 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=seen.append,
    )

    assert result == {"out": 3}
    assert [event["outputs"] for event in seen] == [{"out": 1}, {"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()