nix-estimator/tests/test_estimate.py

"""Tests that estimate() normalizes arbitrary-order node/core grids (issue #11)."""

from nix_estimator.estimate import estimate, peak_ram_gb_per_node


def _toy():
    """A root fanning out to heavy leaves, joined at a sink.

    Wide enough that node count matters, heavy enough that cores matter.
    Durations are pinned via ``history`` so the grid is deterministic.
    """

    def p(name):
        return f"/nix/store/{'0' * 32}-{name}.drv"

    leaves = [p(f"heavylib{i}") for i in range(4)]
    root, sink = p("root"), p("sink")
    closure = {d: {} for d in [root, *leaves, sink]}
    preds = {root: [], sink: leaves, **{leaf: [root] for leaf in leaves}}
    nodes = set(closure)
    history = {"root": 1.0, "sink": 1.0, **{f"heavylib{i}": 20.0 for i in range(4)}}
    return closure, preds, nodes, history


def test_unsorted_core_grid_matches_sorted():
    closure, preds, nodes, history = _toy()
    ref = estimate(closure, preds, nodes, history=history, core_grid=(8, 16, 32))
    shuf = estimate(closure, preds, nodes, history=history, core_grid=(32, 8, 16))
    assert shuf.recommendation == ref.recommendation
    # normalization must not silently swallow the top core count
    assert (
        shuf.recommendation["one_big_node_min"]
        == ref.recommendation["one_big_node_min"]
    )


def test_unsorted_node_grid_matches_sorted():
    closure, preds, nodes, history = _toy()
    ref = estimate(closure, preds, nodes, history=history, node_grid=(1, 2, 4, 8))
    shuf = estimate(closure, preds, nodes, history=history, node_grid=(8, 1, 4, 2))
    assert shuf.recommendation == ref.recommendation


def test_unsorted_grids_duplicates_deduped():
    closure, preds, nodes, history = _toy()
    ref = estimate(closure, preds, nodes, history=history, core_grid=(8, 16))
    dup = estimate(closure, preds, nodes, history=history, core_grid=(16, 8, 8, 16))
    assert dup.recommendation == ref.recommendation


def test_recommend_node_grid_min_above_peak():
    # Issue #18: if every node_grid value exceeds the graph's peak concurrency,
    # the clamp target (peak) is below all grid keys. _recommend must snap back
    # onto an existing grid value instead of KeyError-ing on at[peak].
    def p(name):
        return f"/nix/store/{'0' * 32}-{name}.drv"

    # linear chain: peak concurrency == 1, well below any node_grid entry.
    a, b, c = p("a"), p("b"), p("c")
    closure = {a: {}, b: {}, c: {}}
    preds = {a: [], b: [a], c: [b]}
    nodes = set(closure)
    history = {"a": 1.0, "b": 1.0, "c": 1.0}
    est = estimate(closure, preds, nodes, history=history, node_grid=(4, 8, 16))
    assert est.peak_parallelism == 1
    # snapped down to the smallest available grid value
    assert est.recommendation["nodes"] == 4
    assert "est_makespan_min" in est.recommendation


def test_estimate_max_jobs_helps_wide_small_graph():
    # Issue #13: a fan of many small, weakly-parallel derivations. max_jobs>1
    # packs several per node; the per-job core split barely hurts small drvs, so
    # the single-node makespan drops.
    def p(name):
        return f"/nix/store/{'0' * 32}-{name}.drv"

    leaves = [p(f"small{i}") for i in range(12)]
    root, sink = p("root"), p("sink")
    closure = {d: {} for d in [root, *leaves, sink]}
    preds = {root: [], sink: leaves, **{leaf: [root] for leaf in leaves}}
    nodes = set(closure)
    history = {"root": 0.5, "sink": 0.5, **{f"small{i}": 2.0 for i in range(12)}}
    base = estimate(closure, preds, nodes, history=history, max_jobs=1)
    packed = estimate(closure, preds, nodes, history=history, max_jobs=4)
    assert packed.max_jobs == 4
    # single node, 8 cores: 12 small builds serialize at mj=1 but pack at mj=4
    assert packed.grid[(1, 8)] < base.grid[(1, 8)]


def test_estimate_ram_budget_lengthens_makespan():
    # Issue #15: many RAM-heavy builds (ghc ~6 GB each) on nodes with a tight
    # RAM budget run fewer-at-a-time than max_jobs allows -> longer makespan than
    # the unconstrained run at the same (nodes, cores, max_jobs).
    def p(name):
        return f"/nix/store/{'0' * 32}-{name}.drv"

    leaves = [p(f"ghc-{i}") for i in range(6)]
    root = p("root")
    closure = {d: {} for d in [root, *leaves]}
    preds = {root: [], **{leaf: [root] for leaf in leaves}}
    nodes = set(closure)
    history = {"root": 0.5}  # leaves fall through to the ghc heuristic (RAM ~6GB)
    unc = estimate(closure, preds, nodes, history=history, max_jobs=4)
    capped = estimate(
        closure, preds, nodes, history=history, max_jobs=4, node_ram_gb=8.0
    )
    assert capped.node_ram_gb == 8.0
    # 8 GB fits only one 6-GB ghc build per node at a time, so a single node is
    # slower under the RAM cap than when 4 could pack.
    assert capped.grid[(1, 8)] > unc.grid[(1, 8)]


def test_knee_picks_biggest_core_that_helps():
    # one heavy single-threaded-ish pole: recommendation stays sane and the
    # reported one-big-node uses the largest core count regardless of input order.
    closure, preds, nodes, history = _toy()
    est = estimate(closure, preds, nodes, history=history, core_grid=(32, 16, 8))
    assert est.recommendation["cores_per_node"] in (8, 16, 32)
    # node choice never exceeds the graph width (4 heavy leaves)
    assert est.recommendation["nodes"] <= 4


def _p(name):
    return f"/nix/store/{'0' * 32}-{name}.drv"


def test_peak_ram_reflects_heavy_derivation():
    """A chromium (10 GB) build forces peak RAM >= 10, not the 1 GB default (#19)."""
    root = _p("root")
    heavy = _p("chromium")
    closure = {root: {}, heavy: {}}
    preds = {root: [], heavy: [root]}
    nodes = {root, heavy}
    history = {"root": 1.0, "chromium": 30.0}
    peak = peak_ram_gb_per_node(
        closure, preds, nodes, cores=16, machines=1, history=history, max_jobs=1
    )
    assert peak >= 10.0


def test_peak_ram_sums_co_resident_jobs():
    """With max_jobs>1, two heavy leaves co-resident on one node sum their RAM."""
    root = _p("root")
    # gcc-<v> matches the anchored ^gcc(-|$) RAM_HEAVY key -> 3 GB each
    leaves = [_p(f"gcc-{v}") for v in (13, 14)]
    closure = {d: {} for d in [root, *leaves]}
    preds = {root: [], **{leaf: [root] for leaf in leaves}}
    nodes = set(closure)
    history = {"root": 0.1, "gcc-13": 20.0, "gcc-14": 20.0}
    # one node, two job slots -> both gcc builds run together
    peak = peak_ram_gb_per_node(
        closure, preds, nodes, cores=16, machines=1, history=history, max_jobs=2
    )
    assert peak >= 6.0  # 3 + 3


def test_peak_ram_empty_closure_is_zero():
    assert peak_ram_gb_per_node({}, {}, set(), cores=8, machines=1) == 0.0


def _fetch_drv():
    """A fixed-output derivation: an output with a hash makes is_fixed_output true."""
    return {"outputs": {"out": {"hash": "abc123", "hashAlgo": "sha256"}}}


def test_parallelism_excludes_fixed_output_fetches():
    # Issue #21: root fans out to 3 CPU builds and 5 fetches. peak_parallelism
    # must count only the CPU-bound drvs (3), and the fetch stats must be exact.
    root = _p("root")
    cpu = [_p(f"lib{i}") for i in range(3)]
    fetches = [_p(f"src{i}") for i in range(5)]
    closure = {root: {}, **{d: {} for d in cpu}, **{d: _fetch_drv() for d in fetches}}
    preds = {root: [], **{d: [root] for d in cpu}, **{d: [root] for d in fetches}}
    nodes = set(closure)
    history = {"root": 1.0, **{f"lib{i}": 5.0 for i in range(3)}}
    est = estimate(closure, preds, nodes, history=history)

    # 3 CPU leaves run concurrently; the 5 fetches don't occupy a build slot.
    assert est.peak_parallelism == 3
    assert est.cpu_bound_count == 4  # root + 3 libs
    assert est.fetch_count == 5
    # each fixed-output fetch costs FIXED_OUTPUT = 0.3 min at baseline cores
    assert est.fetch_min == 5 * 0.3
    assert est.avg_parallelism > 0.0


def test_peak_parallelism_ignores_pure_fetch_fanout():
    # A root with only fetch children: without the exclusion peak would be huge;
    # with it the CPU ceiling is just the single non-fetch root.
    root = _p("root")
    fetches = [_p(f"src{i}") for i in range(20)]
    closure = {root: {}, **{d: _fetch_drv() for d in fetches}}
    preds = {root: [], **{d: [root] for d in fetches}}
    nodes = set(closure)
    est = estimate(closure, preds, nodes, history={"root": 1.0})
    assert est.fetch_count == 20
    assert est.cpu_bound_count == 1
    assert est.peak_parallelism == 1