app/backend/tests/flow/test_remote.py

"""A worker on another host, and the thread-to-loop bridge that reaches it.

The socket lives on an event loop; a node executes on a worker thread. These
run a real loop in a thread of its own, because that split is the whole
difficulty: what is checked is that a call handed across it comes back — with
its answer, with the reports it made on the way, or with a failure that says
what happened rather than a wait that never ends.
"""

import asyncio
import gc
import threading
import time
import warnings
from collections.abc import Iterator

import pytest

from fluksio.flow import remote
from fluksio.flow.remote import (
    NoWorker,
    RemoteWorker,
    RemoteWorkerHub,
    WorkerInventory,
)
from fluksio.flow.workers import (
    MAX_REMOTE_NAME,
    NodeTimeout,
    RemoteError,
    _remote_class,
)


@pytest.fixture
def loop() -> Iterator[asyncio.AbstractEventLoop]:
    """An event loop running in a thread, as the server's does."""
    running = asyncio.new_event_loop()
    thread = threading.Thread(target=running.run_forever, daemon=True)
    thread.start()
    yield running
    running.call_soon_threadsafe(running.stop)
    thread.join(timeout=5)
    running.close()


class FakeSocket:
    """Stands in for the websocket: records frames, and says when one lands."""

    def __init__(self) -> None:
        self.sent: list[dict] = []
        self.arrived = threading.Event()

    async def send_json(self, payload: dict) -> None:
        self.sent.append(payload)
        self.arrived.set()


def attach(hub: RemoteWorkerHub, loop: asyncio.AbstractEventLoop, name: str = "gpu1"):
    socket = FakeSocket()
    worker = RemoteWorker(
        name=name, labels=["gpu"], send=socket.send_json, loop=loop, max_parallel=2
    )
    hub.attach(worker)
    return worker, socket


def call_in_thread(target) -> threading.Thread:
    thread = threading.Thread(target=target, daemon=True)
    thread.start()
    return thread


def test_a_call_crosses_to_the_thread_and_the_answer_comes_back(loop):
    hub = RemoteWorkerHub()
    worker, socket = attach(hub, loop)
    result: dict = {}

    thread = call_in_thread(
        lambda: result.update(
            value=hub.run(
                "gpu", "flow", "node", "src", {"x": 1}, "flow.node", timeout=5
            )
        )
    )
    assert socket.arrived.wait(5)
    assert socket.sent[0]["source"] == "src"
    assert socket.sent[0]["kwargs"] == {"x": 1}

    worker.deliver(
        {"call_id": socket.sent[0]["call_id"], "ok": True, "result": {"out": 2}}
    )
    thread.join(timeout=5)
    assert result["value"] == {"out": 2}


def test_reports_arrive_before_the_answer_and_a_heartbeat_is_not_one(loop):
    hub = RemoteWorkerHub()
    worker, socket = attach(hub, loop)
    seen: list[dict] = []
    result: dict = {}

    thread = call_in_thread(
        lambda: result.update(
            value=hub.run(
                "gpu",
                "flow",
                "node",
                "src",
                {},
                "flow.node",
                timeout=5,
                run_id="r1",
                on_event=seen.append,
            )
        )
    )
    assert socket.arrived.wait(5)
    call_id = socket.sent[0]["call_id"]
    # The call names its run, which is how a metric finds the run that made it.
    assert call_id == "r1:flow.node"

    worker.deliver(
        {"call_id": call_id, "event": "metric", "name": "loss", "value": 1.0}
    )
    worker.deliver({"call_id": call_id, "event": "heartbeat"})
    worker.deliver({"call_id": call_id, "ok": True, "result": {"done": True}})

    thread.join(timeout=5)
    assert result["value"] == {"done": True}
    # Liveness is not a measurement; only the metric is kept.
    assert [event["event"] for event in seen] == ["metric"]


def test_a_failure_keeps_its_class_across_the_socket(loop):
    hub = RemoteWorkerHub()
    worker, socket = attach(hub, loop)
    caught: list[Exception] = []

    def call() -> None:
        try:
            hub.run("gpu", "flow", "node", "src", {}, "flow.node", timeout=5)
        except Exception as exc:
            caught.append(exc)

    thread = call_in_thread(call)
    assert socket.arrived.wait(5)
    worker.deliver(
        {
            "call_id": socket.sent[0]["call_id"],
            "ok": False,
            "error": {"type": "ValueError", "message": "bad input", "traceback": "tb"},
        }
    )
    thread.join(timeout=5)
    assert type(caught[0]).__name__ == "ValueError"
    assert str(caught[0]) == "bad input"


def test_a_class_name_off_the_wire_is_bounded_and_identifier_safe():
    cls = _remote_class("Bad Name\n" + "x" * 10_000)
    assert cls.__name__.isidentifier()
    assert len(cls.__name__) == MAX_REMOTE_NAME
    assert issubclass(cls, RemoteError)


def test_a_worker_that_goes_away_fails_the_call_rather_than_hanging(loop):
    hub = RemoteWorkerHub()
    worker, socket = attach(hub, loop)
    caught: list[Exception] = []

    def call() -> None:
        try:
            hub.run("gpu", "flow", "node", "src", {}, "flow.node", timeout=30)
        except Exception as exc:
            caught.append(exc)

    thread = call_in_thread(call)
    assert socket.arrived.wait(5)
    # Pulling the cable mid-training: the node fails, and does not wait out its
    # thirty-second deadline to do it.
    worker.detach()
    thread.join(timeout=5)

    assert not thread.is_alive()
    assert isinstance(caught[0], RemoteError)
    assert "went away" in str(caught[0])


def test_silence_past_the_deadline_is_a_timeout(loop):
    hub = RemoteWorkerHub()
    attach(hub, loop)
    caught: list[Exception] = []

    def call() -> None:
        try:
            hub.run("gpu", "flow", "node", "src", {}, "flow.node", timeout=0.3)
        except Exception as exc:
            caught.append(exc)

    call_in_thread(call).join(timeout=5)
    assert isinstance(caught[0], NodeTimeout)


def test_a_label_nothing_carries_is_named_rather_than_waited_on(loop):
    hub = RemoteWorkerHub()
    attach(hub, loop)

    with pytest.raises(NoWorker, match="tpu"):
        hub.run("tpu", "flow", "node", "src", {}, "flow.node", timeout=5)
    # Compiling against a machine that is not attached is not a broken node —
    # a node importing torch is correct there and missing here.
    assert hub.compile("tpu", "flow", "node", "src") is None


def test_reattaching_replaces_the_old_socket(loop):
    hub = RemoteWorkerHub()
    first, _ = attach(hub, loop)
    second, _ = attach(hub, loop)

    assert first.gone
    assert hub.pick("gpu") is second
    assert [worker.name for worker in hub.workers()] == ["gpu1"]


def test_cancelling_a_run_reaches_only_that_run(loop):
    hub = RemoteWorkerHub()
    worker, socket = attach(hub, loop)

    def call(run_id: str) -> None:
        try:
            hub.run("gpu", "flow", "node", "src", {}, "flow.node", 30, run_id=run_id)
        except Exception:
            pass

    call_in_thread(lambda: call("run-a"))
    assert socket.arrived.wait(5)
    socket.arrived.clear()
    call_in_thread(lambda: call("run-b"))
    assert socket.arrived.wait(5)

    assert hub.cancel_run("run-a") == 1
    cancels = [frame for frame in socket.sent if frame.get("op") == "cancel"]
    assert [frame["call_id"] for frame in cancels] == ["run-a:flow.node"]


def test_a_cancel_after_the_loop_closed_leaves_no_coroutine_behind():
    closed = asyncio.new_event_loop()
    closed.close()
    worker = RemoteWorker(
        name="gpu1", labels=[], send=FakeSocket().send_json, loop=closed
    )

    with warnings.catch_warnings(record=True) as caught:
        warnings.simplefilter("always")
        worker.cancel("run-a:flow.node")
        gc.collect()

    assert not [w for w in caught if "never awaited" in str(w.message)]


# -----------------------------------------------------------------------------
# Liveness, and what a heartbeat is evidence of
#
# The agent beats every ten seconds while it is executing. That says the agent
# is alive; it says nothing about the node, which is why it bounds the silence
# deadline and not the node's own timeout.
# -----------------------------------------------------------------------------


def test_heartbeats_do_not_stave_off_a_nodes_own_timeout(loop):
    hub = RemoteWorkerHub()
    worker, socket = attach(hub, loop)
    caught: list[Exception] = []

    def call() -> None:
        try:
            hub.run("gpu", "flow", "node", "src", {}, "flow.node", timeout=0.5)
        except Exception as exc:
            caught.append(exc)

    thread = call_in_thread(call)
    assert socket.arrived.wait(5)
    call_id = socket.sent[0]["call_id"]
    # Beating faster than the deadline. A node that reports its progress is
    # held off; one whose agent is merely alive is not.
    for _ in range(10):
        worker.deliver({"call_id": call_id, "event": "heartbeat"})
        time.sleep(0.1)

    thread.join(timeout=5)
    assert isinstance(caught[0], NodeTimeout)


def test_with_no_timeout_total_silence_is_still_bounded(loop, monkeypatch):
    monkeypatch.setattr(remote, "SILENCE_S", 0.3)
    hub = RemoteWorkerHub()
    attach(hub, loop)
    caught: list[Exception] = []

    def call() -> None:
        try:
            hub.run("gpu", "flow", "node", "src", {}, "flow.node", timeout=0)
        except Exception as exc:
            caught.append(exc)

    call_in_thread(call).join(timeout=5)
    assert isinstance(caught[0], RemoteError)
    assert not isinstance(caught[0], NodeTimeout)
    assert "presumed gone" in str(caught[0])


def test_with_no_timeout_a_beating_worker_is_left_to_finish(loop, monkeypatch):
    monkeypatch.setattr(remote, "SILENCE_S", 0.3)
    hub = RemoteWorkerHub()
    worker, socket = attach(hub, loop)
    result: dict = {}

    thread = call_in_thread(
        lambda: result.update(
            value=hub.run("gpu", "flow", "node", "src", {}, "flow.node", timeout=0)
        )
    )
    assert socket.arrived.wait(5)
    call_id = socket.sent[0]["call_id"]
    for _ in range(6):
        worker.deliver({"call_id": call_id, "event": "heartbeat"})
        time.sleep(0.1)
    worker.deliver({"call_id": call_id, "ok": True, "result": {"done": True}})

    thread.join(timeout=5)
    assert result["value"] == {"done": True}


def test_an_allocation_rides_along_with_the_call(loop):
    """Thread caps and devices reach the worker on the frame that needs them.

    The worker starts a process per call, so this is the only moment it can
    apply them: a library reads them when it is imported and never again.
    """
    hub = RemoteWorkerHub()
    worker, socket = attach(hub, loop)

    thread = call_in_thread(
        lambda: hub.run(
            "gpu",
            "flow",
            "node",
            "src",
            {},
            "flow.node",
            timeout=5,
            env={"OMP_NUM_THREADS": "4", "CUDA_VISIBLE_DEVICES": "0"},
        )
    )
    assert socket.arrived.wait(5)
    assert socket.sent[0]["env"] == {
        "OMP_NUM_THREADS": "4",
        "CUDA_VISIBLE_DEVICES": "0",
    }

    worker.deliver({"call_id": socket.sent[0]["call_id"], "ok": True, "result": None})
    thread.join(timeout=5)


def test_a_call_with_nothing_allocated_carries_no_env(loop):
    hub = RemoteWorkerHub()
    worker, socket = attach(hub, loop)

    thread = call_in_thread(
        lambda: hub.run("gpu", "flow", "node", "src", {}, "flow.node", timeout=5)
    )
    assert socket.arrived.wait(5)
    assert "env" not in socket.sent[0]

    worker.deliver({"call_id": socket.sent[0]["call_id"], "ok": True, "result": None})
    thread.join(timeout=5)


def test_a_worker_that_reports_nothing_is_the_machine_it_used_to_be():
    """Inventory is read leniently: absent is a default, junk is a default.

    A worker saying nothing has to keep being scheduled by its label, and one
    that learns to report something new must not need this engine taught about
    it first.
    """
    assert WorkerInventory.from_hello({}) == WorkerInventory(cpus=1, gpus=0)
    assert WorkerInventory.from_hello(None) == WorkerInventory()

    read = WorkerInventory.from_hello(
        {"cpus": "8", "gpus": 2, "ram_mb": 64000, "tpus": 4}
    )
    assert (read.cpus, read.gpus, read.ram_mb) == (8, 2, 64000)

    junk = WorkerInventory.from_hello({"cpus": "many", "gpus": None, "ram_mb": 0})
    assert (junk.cpus, junk.gpus, junk.ram_mb) == (1, 0, None)