app/backend/tests/flow/test_runs.py

"""What a run is made of: isolation, the per-node record, and what it reports.

The service itself needs a database, so what is checked here is the part that
decides whether a run is correct — that two runs of one flow cannot see each
other's messages, that every node executed is reported once, and that the
parameters a caller sends are refused before anything runs if they are wrong.
"""

import pytest

from fluksio.flow.events import EventBus
from fluksio.flow.messages import DType, MessageSpec
from fluksio.flow.nodes import Node
from fluksio.flow.pipeline import CacheHit, NodeOutcome, Pipeline, run_cache_key
from fluksio.flow.runs import (
    MetricSink,
    RunRejected,
    batch_issues,
    collect_result,
    digest_of,
    required_labels,
    required_resources,
    seed_values,
)
from fluksio.flow.schemas import FlowDef, FlowInput, NodeDef, Resources
from fluksio.flow.state import MemoryState


def spec(name: str, dtype: DType = DType.FLOAT, **kwargs) -> MessageSpec:
    return MessageSpec(name=name, dtype=dtype, **kwargs)


def make_node(node_id: str, flow: str, f, requires=(), provides=()) -> Node:
    node = Node(f=f, requires=list(requires), provides=list(provides), name=node_id)
    node.assign_flow(flow, node_id)
    return node


def double_flow() -> FlowDef:
    """A flow with one input, one node and one declared output."""
    return FlowDef(
        name="study",
        mode="batch",
        inputs=[FlowInput(spec=spec("lr"), initial=0.1)],
        outputs=["loss"],
        nodes=[
            NodeDef(
                id="train",
                requires=[spec("lr")],
                provides=[spec("loss")],
            )
        ],
    )


def build(flow: FlowDef, state: MemoryState, observer=None) -> Pipeline:
    """The pipeline a run drives, without the controller that normally builds it."""
    node = make_node(
        "train",
        flow.name,
        lambda lr, params: {"loss": lr * 2},
        requires=[spec("lr")],
        provides=[spec("loss")],
    )
    return Pipeline(nodes=[node], state=state, observer=observer)


# -----------------------------------------------------------------------------
# Isolation — the reason a run has a state backend of its own
# -----------------------------------------------------------------------------


def test_two_runs_of_one_flow_do_not_see_each_other():
    flow = double_flow()
    first, second = MemoryState(), MemoryState()

    build(flow, first).run(seed_values(flow, {"lr": 0.5}))
    build(flow, second).run(seed_values(flow, {"lr": 4.0}))

    assert collect_result(flow, first) == {"loss": 1.0}
    assert collect_result(flow, second) == {"loss": 8.0}


def test_result_falls_back_to_everything_the_flow_holds():
    flow = double_flow()
    flow.outputs = []
    state = MemoryState()

    build(flow, state).run(seed_values(flow, {"lr": 1.0}))

    # The input is part of what the flow ended up holding; the engine's own
    # bookkeeping keys are not.
    assert collect_result(flow, state) == {"lr": 1.0, "loss": 2.0}


# -----------------------------------------------------------------------------
# The record — every node a run executed, reported once
# -----------------------------------------------------------------------------


def test_observer_sees_every_node_once():
    flow = double_flow()
    seen = []

    build(flow, MemoryState(), observer=seen.append).run(seed_values(flow, {"lr": 1.0}))

    assert [(o.node, o.ok) for o in seen] == [("study.train", True)]
    assert seen[0].outputs == 1


def test_observer_reports_a_failing_node_with_its_error():
    seen = []

    def boom(params):
        raise ValueError("no convergence")

    node = make_node("train", "study", boom, provides=[spec("loss")])
    Pipeline(nodes=[node], state=MemoryState(), observer=seen.append).run()

    assert len(seen) == 1
    assert not seen[0].ok
    assert "no convergence" in seen[0].error
    # The traceback rides along in the logs, so a failure can be read back off
    # the run rather than reproduced under `run --local`.
    assert "Traceback" in seen[0].logs
    assert "ValueError: no convergence" in seen[0].logs


def test_a_failing_observer_does_not_take_the_node_down():
    def refuse(_outcome):
        raise RuntimeError("the database is gone")

    node = make_node(
        "train", "study", lambda params: {"loss": 1.0}, provides=[spec("loss")]
    )
    pipeline = Pipeline(nodes=[node], state=MemoryState(), observer=refuse)
    pipeline.run()

    assert pipeline.state["study.loss"] == 1.0


# -----------------------------------------------------------------------------
# Producing values before returning
#
# A number worth keeping is an output, not a log. A node that produces over
# time is a generator, and each yield is published on the port it names — so
# the run's metrics are its streaming outputs rather than something recorded
# beside them.
# -----------------------------------------------------------------------------


def training_flow() -> FlowDef:
    return FlowDef(
        name="study",
        mode="batch",
        inputs=[FlowInput(spec=spec("steps", DType.INT), initial=3)],
        outputs=["final_loss"],
        nodes=[
            NodeDef(
                id="train",
                requires=[spec("steps", DType.INT)],
                provides=[
                    spec("loss", stream=True),
                    spec("final_loss"),
                ],
            )
        ],
    )


def training_node(flow: str = "study") -> Node:
    def train(steps, params):
        loss = 1.0
        for _ in range(steps):
            loss = loss / 2
            yield {"loss": loss}
        return {"final_loss": loss}

    return make_node(
        "train",
        flow,
        train,
        requires=[spec("steps", DType.INT)],
        provides=[spec("loss", stream=True), spec("final_loss")],
    )


def test_each_yield_is_published_and_the_return_is_the_result():
    flow = training_flow()
    state = MemoryState()
    seen: list[tuple[str, dict]] = []
    pipeline = Pipeline(
        nodes=[training_node()],
        state=state,
        emission_observer=lambda node, outputs: seen.append((node, outputs)),
    )

    pipeline.run(seed_values(flow, {"steps": 3}))

    # It returned something, so every yield was a value produced on the way —
    # each published on `loss` the moment it happened.
    assert [outputs["study.loss"] for _node, outputs in seen] == [0.5, 0.25, 0.125]
    assert all(node == "study.train" for node, _ in seen)
    # The latest of them is what the message holds, as for any producer.
    assert state["study.loss"] == 0.125
    # And what it returned is the node's output, and so the run's result.
    assert collect_result(flow, state) == {"final_loss": 0.125}


def test_without_a_return_the_last_yield_is_the_result():
    def train(params):
        yield {"loss": 1.0}
        yield {"loss": 0.5}

    node = make_node("train", "study", train, provides=[spec("loss", stream=True)])
    seen: list[tuple[str, dict]] = []
    pipeline = Pipeline(
        nodes=[node],
        state=MemoryState(),
        emission_observer=lambda n, o: seen.append((n, o)),
    )
    pipeline.run()

    # The last yield is the node's output rather than an emission, so it is
    # not counted twice.
    assert [outputs["study.loss"] for _node, outputs in seen] == [1.0]
    assert pipeline.state["study.loss"] == 0.5


def test_emissions_are_checked_against_the_port_they_name():
    def wrong(params):
        yield {"loss": "not a number"}
        return {"final_loss": 1.0}

    node = make_node(
        "train",
        "study",
        wrong,
        provides=[spec("loss", stream=True), spec("final_loss")],
    )
    seen: list[NodeOutcome] = []
    Pipeline(nodes=[node], state=MemoryState(), observer=seen.append).run()

    # A wrong type is a failed node, exactly as it is for a return value —
    # which is the point of emissions going out through declared ports.
    assert not seen[0].ok
    assert "loss" in seen[0].error


def test_an_emission_that_nothing_declares_names_what_was_emitted():
    """A mistyped metric name is how a training curve goes missing."""
    events = []

    def stray(params):
        yield {"undeclared": 1.0}
        return {"final_loss": 2.0}

    bus = EventBus()
    bus.publish = events.append  # type: ignore[method-assign]
    node = make_node("train", "study", stray, provides=[spec("final_loss")])
    state = MemoryState()
    Pipeline(nodes=[node], state=state, events=bus).run()

    (error,) = [e for e in events if e["type"] == "node_error"]
    assert "undeclared" in error["error"]
    assert "final_loss" in error["error"]
    assert "study.undeclared" not in state


def test_a_mistyped_port_fails_at_the_yield_that_produced_it():
    """Not at the one after it: a training loop's second pass can be an hour."""
    passes = []

    def stray(params):
        passes.append(1)
        yield {"undeclared": 1.0}
        passes.append(2)
        yield {"loss": 0.5}

    node = make_node("train", "study", stray, provides=[spec("loss", stream=True)])
    seen: list[NodeOutcome] = []
    Pipeline(nodes=[node], state=MemoryState(), observer=seen.append).run()

    assert not seen[0].ok
    assert "undeclared" in seen[0].error
    assert passes == [1]


def test_emissions_reach_the_run_as_a_series_with_a_step_each():
    sink = MetricSink("run-1", batch=1)
    sink.handle("study.train", {"study.loss": 1.0, "study.tag": "ignored"})
    sink.handle("study.train", {"study.loss": 0.5})

    # Numbers become the run's series; anything else is on the run some other
    # way — as its result, or as an artifact.
    assert sink._steps == {"study.loss": 1}


def test_a_written_batch_says_it_is_there():
    """What a screen watching a live run waits on.

    The rows are still the record — this is a nudge carrying names, one per
    batch rather than one per point, so nothing lands on the hot path.
    """
    said = []
    sink = MetricSink("run-1", batch=1, on_flush=said.append)

    sink.handle("study.train", {"study.loss": 1.0, "study.tag": "ignored"})
    assert said == [["study.loss"]]

    # Held back until the batch is due, then announced once for all of it.
    quiet = MetricSink("run-2", batch=10, interval=3600, on_flush=said.append)
    quiet.handle("study.train", {"study.loss": 1.0, "study.acc": 0.5})
    assert said == [["study.loss"]]
    quiet.flush()
    assert said[-1] == ["study.acc", "study.loss"]

    # Nothing to write is nothing to say.
    quiet.flush()
    assert len(said) == 2


# -----------------------------------------------------------------------------
# What a caller may ask for
# -----------------------------------------------------------------------------


def test_parameters_must_be_declared_inputs():
    flow = double_flow()
    with pytest.raises(RunRejected, match="not an input"):
        seed_values(flow, {"learning_rate": 0.1})


def test_parameters_are_type_checked_before_anything_runs():
    flow = double_flow()
    with pytest.raises(RunRejected, match="lr"):
        seed_values(flow, {"lr": "fast"})


def test_a_rate_limited_port_cannot_be_run_as_a_batch():
    flow = double_flow()
    flow.nodes[0].provides = [spec("loss", interval=30)]
    # Without a queue there is no timer to release what an interval holds, so
    # the value would be dropped rather than delayed.
    assert batch_issues(flow)
    assert not batch_issues(double_flow())


def test_a_node_with_no_stored_code_cannot_be_run():
    """The store answers a template for one, and a template publishes nothing."""

    class Store:
        def __init__(self, *, written: bool):
            self.written = written

        def has_node_source(self, flow, node_id, draft=False):
            return self.written

    flow = double_flow()
    (issue,) = batch_issues(flow, Store(written=False))
    assert "no stored code" in issue
    assert not batch_issues(flow, Store(written=True))
    # A shared node runs the library's copy, so it is not this node's to have.
    flow.nodes[0].source_ref = "shared"
    assert not batch_issues(flow, Store(written=False))


def test_a_streaming_port_may_thin_itself_out():
    flow = double_flow()
    flow.nodes[0].provides = [spec("loss", interval=0.5, stream=True)]
    # On a curve, an interval is asking for the canvas not to be flooded —
    # the run's history still keeps every value.
    assert not batch_issues(flow)


def test_the_digest_identifies_the_inputs_not_their_order():
    assert digest_of({"a": 1, "b": 2}, 3) == digest_of({"b": 2, "a": 1}, 3)
    assert digest_of({"a": 1}, 3) != digest_of({"a": 1}, 4)


def test_labels_come_from_the_nodes_that_ask_for_a_device():
    flow = double_flow()
    assert required_labels(flow) == []
    flow.nodes[0].device = "gpu"
    assert required_labels(flow) == ["gpu"]


def test_a_preferred_device_does_not_hold_a_run_back():
    flow = double_flow()
    flow.nodes[0].device = "gpu"
    flow.nodes[0].device_policy = "prefer"
    # It runs on the engine when no such worker is attached, so waiting for
    # one would be waiting for something the run does not need.
    assert required_labels(flow) == []


def test_a_device_on_a_node_no_worker_could_run_is_not_waited_for():
    """A connector is an entry point loaded here, so a worker cannot take it.

    Holding the run for one is a run that never starts — and it looked exactly
    like a GPU box somebody had forgotten to turn on.
    """
    flow = double_flow()
    flow.nodes[0].type = "mqtt"
    flow.nodes[0].device = "gpu"
    assert required_labels(flow) == []


def test_what_a_run_needs_is_the_largest_single_machine_it_asks_for():
    flow = double_flow()
    assert required_resources(flow) is None

    flow.nodes[0].resources = Resources(cpus=2, gpus=1)
    flow.nodes.append(NodeDef(id="fit", resources=Resources(cpus=8)))
    assert required_resources(flow) == {
        "cpus": 8,
        "gpus": 1,
        "ram_mb": 0,
        "device": "",
    }


def test_a_runs_seed_fills_an_input_of_that_name():
    flow = double_flow()
    flow.inputs.append(FlowInput(spec=spec("seed", DType.INT), initial=0))

    # Otherwise the field that tells two runs of one configuration apart would
    # only look like the number the flow draws from.
    assert seed_values(flow, {}, seed=7)["study.seed"] == 7
    # An explicit parameter still wins.
    assert seed_values(flow, {"seed": 3}, seed=7)["study.seed"] == 3


def test_a_flow_without_a_seed_input_ignores_the_runs_seed():
    flow = double_flow()
    assert seed_values(flow, {"lr": 1.0}, seed=7) == {"study.lr": 1.0}


# -----------------------------------------------------------------------------
# The stage cache — a node whose inputs have not changed is not run again
# -----------------------------------------------------------------------------


class FakeCache:
    """A stage cache with no database behind it, and a record of what it was asked."""

    def __init__(
        self, entries: dict[str, dict | None] | None = None, flow: str = "study"
    ) -> None:
        self.entries = entries or {}
        self.flow = flow
        self.asked: list[str] = []

    def lookup(self, key: str):
        self.asked.append(key)
        if key in self.entries:
            return CacheHit(
                flow=self.flow, outputs=self.entries[key], metrics_run="earlier"
            )
        return None


def counting_node(flow: str = "study") -> tuple[Node, list[int]]:
    """A node that says how many times it actually ran."""
    calls: list[int] = []

    def train(lr, params):
        calls.append(1)
        return {"loss": lr * 2}

    node = make_node(
        "train", flow, train, requires=[spec("lr")], provides=[spec("loss")]
    )
    node.fingerprint = "fp-train"
    return node, calls


def test_a_cache_hit_restores_the_outputs_without_running_the_node():
    flow = double_flow()
    node, calls = counting_node()
    state = MemoryState()
    seen: list[NodeOutcome] = []
    key = run_cache_key("fp-train", {"study.lr": 0.5}, "study")
    cache = FakeCache({key: {"study.loss": 99.0}})

    Pipeline(nodes=[node], state=state, observer=seen.append, run_cache=cache).run(
        seed_values(flow, {"lr": 0.5})
    )

    assert calls == []
    # Restored into this run's own state, which is where everything
    # downstream of it looks — its namespace holds nothing otherwise.
    assert collect_result(flow, state) == {"loss": 99.0}
    assert seen[0].cached and seen[0].cache_key == key
    # Where its series is, since the hit replayed none of it.
    assert seen[0].cached_from == "earlier"


def test_a_hit_from_another_flow_restores_under_this_flow_s_names():
    """The same node reached through two flows publishes under two names."""
    flow = double_flow()
    node, calls = counting_node()
    state = MemoryState()
    seen: list[NodeOutcome] = []
    key = run_cache_key("fp-train", {"study.lr": 0.5}, "study")
    # Recorded by a run of "other", which is what a node shared between two
    # flows gets a hit from — the values are the same, the namespace is not.
    cache = FakeCache({key: {"other.loss": 99.0}}, flow="other")

    Pipeline(nodes=[node], state=state, observer=seen.append, run_cache=cache).run(
        seed_values(flow, {"lr": 0.5})
    )

    assert calls == []
    assert collect_result(flow, state) == {"loss": 99.0}
    # And stored that way too, so the next run to reuse this one finds names
    # it can requalify from a flow that really did publish them.
    assert seen[0].output_values == {"study.loss": 99.0}


def test_a_miss_runs_the_node_and_carries_what_would_be_stored():
    flow = double_flow()
    node, calls = counting_node()
    seen: list[NodeOutcome] = []
    cache = FakeCache()

    Pipeline(
        nodes=[node], state=MemoryState(), observer=seen.append, run_cache=cache
    ).run(seed_values(flow, {"lr": 0.5}))

    assert calls == [1]
    assert cache.asked == [run_cache_key("fp-train", {"study.lr": 0.5}, "study")]
    assert not seen[0].cached
    assert seen[0].cache_key and seen[0].output_values == {"study.loss": 1.0}


def test_the_key_follows_the_inputs():
    first = run_cache_key("fp", {"lr": 0.5})
    assert first != run_cache_key("fp", {"lr": 0.6})
    assert first != run_cache_key("other", {"lr": 0.5})
    assert first == run_cache_key("fp", {"lr": 0.5})


def test_the_key_does_not_follow_the_flow_an_input_hangs_in():
    # The same node reading the same value through two flows did the same work,
    # so a run of one is a hit for the other.
    assert run_cache_key("fp", {"study.lr": 0.5}, "study") == run_cache_key(
        "fp", {"quick.lr": 0.5}, "quick"
    )
    # A name belonging to neither keeps its prefix: reading another flow's
    # message is part of what this execution is.
    assert run_cache_key("fp", {"other.lr": 0.5}, "study") != run_cache_key(
        "fp", {"other.lr": 0.5}, "other"
    )


def test_an_artifact_input_counts_as_its_digest():
    digest = "sha256:" + "0" * 64
    # The same bytes under another name, of a size recorded differently, are
    # the same input — the reference is a handle, the digest is the content.
    assert run_cache_key(
        "fp", {"data": {"digest": digest, "name": "a.csv", "size": 3}}
    ) == run_cache_key("fp", {"data": {"digest": digest, "name": "b.csv", "size": 3}})


def test_a_node_with_no_fingerprint_is_never_looked_up():
    """Built-in and connector nodes, and anything declared `cache=False`."""
    flow = double_flow()
    node, calls = counting_node()
    node.fingerprint = ""
    seen: list[NodeOutcome] = []
    cache = FakeCache()

    Pipeline(
        nodes=[node], state=MemoryState(), observer=seen.append, run_cache=cache
    ).run(seed_values(flow, {"lr": 0.5}))

    assert calls == [1]
    assert cache.asked == []
    assert seen[0].cache_key == ""