app/backend/tests/flow/test_placement.py

"""Which machine a node runs on, and what happens when none of them is free.

The two questions the engine used to answer separately — where does a node run,
and how much of that machine does it get — are one question here. What is
checked is that the answer covers every machine attached rather than only the
host the engine is on, that a claim is given back however the call ends, and
that a node with nowhere to go waits and says so instead of failing.
"""

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

import pytest

from fluksio.flow.placement import Placer
from fluksio.flow.remote import RemoteWorker, RemoteWorkerHub, WorkerInventory
from fluksio.flow.resources import ResourceAccountant
from fluksio.flow.schemas import NodeDef, Resources


@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:
    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",
    labels: tuple[str, ...] = ("gpu",),
    cpus: int = 8,
    gpus: int = 1,
    ram_mb: int | None = None,
):
    socket = FakeSocket()
    worker = RemoteWorker(
        name=name,
        labels=list(labels),
        send=socket.send_json,
        loop=loop,
        max_parallel=4,
        inventory=WorkerInventory(cpus=cpus, gpus=gpus, ram_mb=ram_mb),
    )
    hub.attach(worker)
    return worker, socket


def placer_over(hub: RemoteWorkerHub | None = None, cpus: int = 4, gpus: int = 0):
    return Placer(local=ResourceAccountant(cpus=cpus, gpus=gpus), hub=hub)


# -----------------------------------------------------------------------------
# Where a node goes
# -----------------------------------------------------------------------------


def test_a_node_asking_for_a_card_finds_the_machine_that_has_one(loop):
    """No label anywhere — the size of the ask is what places it.

    This is the whole point of a worker reporting its inventory: the engine has
    no GPU, so a node that needs one used to be clamped down to none and run
    here anyway.
    """
    hub = RemoteWorkerHub()
    placer = placer_over(hub)
    worker, _ = attach(hub, loop, labels=())

    with placer.claim(Resources(cpus=2, gpus=1)) as (target, allocation):
        assert target.worker is worker
        assert allocation.gpus == (0,)
        assert worker.accountant.snapshot()["gpus"]["free"] == 0

    assert worker.accountant.snapshot()["gpus"]["free"] == 1


def test_a_node_that_fits_here_stays_here(loop):
    """No network, and a worker pool that is already warm."""
    hub = RemoteWorkerHub()
    placer = placer_over(hub)
    attach(hub, loop)

    with placer.claim(Resources(cpus=2)) as (target, _):
        assert target.worker is None
        assert target.name == "engine"


def test_a_required_label_waits_for_its_machine_rather_than_running_here(loop):
    hub = RemoteWorkerHub()
    placer = placer_over(hub)
    worker, _ = attach(hub, loop, labels=("gpu",))

    with placer.claim(Resources(cpus=2), device="gpu", policy="require") as (
        target,
        _,
    ):
        assert target.worker is worker


def test_a_preferred_label_falls_back_here_and_is_still_accounted(loop):
    """The fallback used to skip the books entirely, so it held nothing."""
    hub = RemoteWorkerHub()
    placer = placer_over(hub)

    with placer.claim(Resources(cpus=3), device="gpu", policy="prefer") as (
        target,
        allocation,
    ):
        assert target.worker is None
        assert allocation.cpus == 3
        assert placer.local.snapshot()["cpus"]["free"] == 1


def test_asking_for_more_than_anything_has_gets_what_there_is(loop):
    """A flow written on a cluster still has to run on a laptop."""
    placer = placer_over(cpus=2)

    with placer.claim(Resources(cpus=64, gpus=4)) as (target, allocation):
        assert target.name == "engine"
        assert allocation.cpus == 2
        assert allocation.gpus == ()


def test_what_is_clamped_to_is_a_machine_that_exists(loop):
    """Each dimension taken separately can describe a machine nobody has.

    A big CPU box and a small GPU box are not one machine with both. Clamping
    to the largest of each would ask for something that never comes free, which
    is the hung run the clamping is there to avoid.
    """
    hub = RemoteWorkerHub()
    placer = placer_over(hub, cpus=2)
    attach(hub, loop, name="big", labels=(), cpus=64, gpus=0)
    attach(hub, loop, name="small-gpu", labels=(), cpus=4, gpus=1)

    with placer.claim(Resources(cpus=64, gpus=1)) as (target, allocation):
        assert target.name == "small-gpu"
        assert (allocation.cpus, allocation.gpus) == (4, (0,))


# -----------------------------------------------------------------------------
# Waiting
# -----------------------------------------------------------------------------


def test_a_node_waits_for_a_busy_machine_and_says_so(loop):
    hub = RemoteWorkerHub()
    events: list[dict] = []
    placer = Placer(
        local=ResourceAccountant(cpus=2),
        hub=hub,
        events=type("Bus", (), {"publish": lambda self, e: events.append(e)})(),
    )
    started = threading.Event()

    def second() -> None:
        with placer.claim(Resources(cpus=2), node="study.b"):
            started.set()

    with placer.claim(Resources(cpus=2), node="study.a"):
        thread = threading.Thread(target=second, daemon=True)
        thread.start()
        # Long enough to have run if nothing was holding it back.
        assert not started.wait(0.3)
        assert placer.snapshot()["waiting"][0]["node"] == "study.b"
        assert events[0]["type"] == "node_queued"
        assert "2 cpu(s)" in events[0]["detail"]

    thread.join(timeout=5)
    assert started.is_set()
    assert placer.snapshot()["waiting"] == []


def test_a_worker_attaching_wakes_what_was_waiting_for_one(loop):
    """A node waiting inside one machine's books could not hear this."""
    hub = RemoteWorkerHub()
    placer = placer_over(hub)
    hub.on_change = placer.wake
    ran = threading.Event()

    def needs_a_card() -> None:
        with placer.claim(Resources(cpus=1, gpus=1), node="study.fit") as (target, _):
            if target.worker is not None:
                ran.set()

    thread = threading.Thread(target=needs_a_card, daemon=True)
    thread.start()
    # Nothing has a GPU yet, so it is clamped to none and runs here — the wait
    # only happens while a machine that *could* take it is busy.
    thread.join(timeout=5)
    assert not thread.is_alive()

    worker, _ = attach(hub, loop, labels=())
    held = worker.accountant.try_take(cpus=1, gpus=1)
    thread = threading.Thread(target=needs_a_card, daemon=True)
    thread.start()
    assert not ran.wait(0.3)

    worker.accountant.give_back(held)
    thread.join(timeout=5)
    assert ran.is_set()


def test_everything_comes_back_when_a_node_fails(loop):
    placer = placer_over(cpus=4, gpus=1)
    with pytest.raises(ValueError):
        with placer.claim(Resources(cpus=4, gpus=1)):
            raise ValueError("the node raised")

    assert placer.local.snapshot()["cpus"]["free"] == 4
    assert placer.local.snapshot()["gpus"]["free"] == 1


def test_a_worker_going_away_mid_call_gives_its_share_back(loop):
    hub = RemoteWorkerHub()
    placer = placer_over(hub)
    worker, _ = attach(hub, loop, labels=())

    with pytest.raises(RuntimeError):
        with placer.claim(Resources(cpus=2, gpus=1)) as (target, _):
            assert target.worker is worker
            hub.detach(worker.name)
            raise RuntimeError("worker went away mid-call")

    # The books are the worker's own, so they go with it — what matters is that
    # nothing is left holding them and the next claim is not blocked.
    assert worker.accountant.snapshot()["gpus"]["free"] == 1


def test_many_nodes_at_once_all_finish(loop):
    """The placer must not deadlock under contention; it is on every call."""
    placer = placer_over(cpus=4)
    done = []

    def work() -> None:
        with placer.claim(Resources(cpus=2), node="study.n"):
            time.sleep(0.01)
            done.append(1)

    threads = [threading.Thread(target=work) for _ in range(8)]
    for thread in threads:
        thread.start()
    for thread in threads:
        thread.join(timeout=10)

    assert len(done) == 8
    assert placer.local.snapshot()["cpus"]["free"] == 4


# -----------------------------------------------------------------------------
# What a run is told before it starts
# -----------------------------------------------------------------------------


def test_a_run_is_only_held_for_a_wait_that_will_end(loop):
    """An ask nothing can grant is cut down and runs; holding it would hang."""
    placer = placer_over(cpus=4)

    assert placer.satisfiable(None) is None
    assert placer.satisfiable({"cpus": 2}) is None
    # No machine here has a GPU and nothing can be started with one.
    assert placer.satisfiable({"cpus": 1, "gpus": 1}) is None


def test_a_run_waits_for_a_machine_something_can_start(loop):
    class Cluster:
        def covers(self, cpus, gpus, ram_mb, label=None):
            return gpus <= 2

        def provision(self, cpus, gpus, ram_mb, label=None):
            pass

        def reconcile(self, attached):
            pass

        def shapes(self, label=None):
            return [(16, 2, 65536)]

        def status(self):
            return {"name": "hpc"}

    placer = placer_over(cpus=4)
    placer.provisioners = [Cluster()]

    assert placer.satisfiable({"cpus": 1, "gpus": 1}) == "1 gpu(s) and 1 cpu(s)"
    assert placer.satisfiable({"cpus": 1, "gpus": 0}) is None


# -----------------------------------------------------------------------------
# The whole path, through a real worker
# -----------------------------------------------------------------------------


def test_a_declared_node_is_held_to_its_share(tmp_path):
    """Four nodes wanting two cores each, on a machine with four.

    What is checked is the pair: never more than the machine has in flight at
    once, and every one of them told what it was given — the two halves that
    together are the oversubscription this is for. Also that it does not
    deadlock, since the claim is taken before a worker slot and both block.
    """
    import sys

    from fluksio.flow.controller import FlowController, RunContext
    from fluksio.flow.messages import DType, MessageSpec
    from fluksio.flow.schemas import FlowDef
    from fluksio.flow.state import MemoryState
    from fluksio.flow.store import FlowStore
    from fluksio.flow.workers import PythonWorkerPool

    store = FlowStore(tmp_path / "flows")
    store.write_flow(
        FlowDef(
            name="study",
            mode="batch",
            nodes=[
                NodeDef(
                    id="fit",
                    provides=[MessageSpec(name="threads", dtype=DType.INT)],
                    resources=Resources(cpus=2, env={"XLA_FLAGS": "--x=false"}),
                )
            ],
        )
    )
    store.write_node_source(
        "study",
        "fit",
        "import os, time\n\n\ndef process():\n"
        "    time.sleep(0.2)\n"
        "    return {'threads': int(os.environ['OMP_NUM_THREADS'])}\n",
    )

    pool = PythonWorkerPool(python=sys.executable, size=4)
    pool.start()
    accountant = ResourceAccountant(cpus=4)
    placer = Placer(local=accountant)
    controller = FlowController(
        store, workers=pool, resources=accountant, placer=placer
    )
    answers: list[object] = []
    in_flight: list[int] = []

    def once(index: int) -> None:
        pipeline = controller.build_run_pipeline(
            store.read_flow("study"),
            state=MemoryState(),
            run=RunContext(run_id=f"r{index}"),
        )
        pipeline.run({})
        answers.append(pipeline.state.get("study.threads"))

    def watch(until: threading.Event) -> None:
        while not until.is_set():
            in_flight.append(4 - int(accountant.snapshot()["cpus"]["free"]))
            time.sleep(0.01)

    finished = threading.Event()
    watcher = threading.Thread(target=watch, args=(finished,))
    watcher.start()
    threads = [threading.Thread(target=once, args=(i,)) for i in range(4)]
    try:
        for thread in threads:
            thread.start()
        for thread in threads:
            thread.join(timeout=60)
            assert not thread.is_alive(), "a claim and a worker slot deadlocked"
    finally:
        finished.set()
        watcher.join(timeout=5)
        pool.stop()

    assert answers == [2, 2, 2, 2]
    assert max(in_flight) <= 4
    assert accountant.snapshot()["cpus"]["free"] == 4
    # One environment, so one extra pool however many nodes derived it.
    assert len(pool._children) == 1