nix-estimator/tests/test_estimate.py

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

from nix_estimator.estimate import estimate


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