Oleks
201 lines
8.3 KiB
Python
201 lines
8.3 KiB
Python
"""Unit tests for the pure scheduling analysis on toy DAGs."""
|
||
|
||
from nix_estimator import schedule
|
||
|
||
|
||
def test_linear_chain_has_no_parallelism():
|
||
# a -> b -> c, 1 min each: span == work, peak == 1
|
||
dur = {"a": 1.0, "b": 1.0, "c": 1.0}
|
||
preds = {"a": [], "b": ["a"], "c": ["b"]}
|
||
span, chain = schedule.critical_path(dur, preds)
|
||
assert span == 3.0
|
||
assert chain == ["a", "b", "c"]
|
||
assert schedule.work(dur) == 3.0
|
||
assert schedule.peak_concurrency(dur, preds) == 1
|
||
# more machines cannot beat the chain
|
||
assert schedule.makespan(dur, preds, 1) == 3.0
|
||
assert schedule.makespan(dur, preds, 8) == 3.0
|
||
|
||
|
||
def test_wide_fanout_parallelizes():
|
||
# root -> {l1..l4}, then sink depends on all leaves
|
||
dur = {"root": 1.0, "l1": 2.0, "l2": 2.0, "l3": 2.0, "l4": 2.0, "sink": 1.0}
|
||
preds = {"root": [], "l1": ["root"], "l2": ["root"], "l3": ["root"],
|
||
"l4": ["root"], "sink": ["l1", "l2", "l3", "l4"]}
|
||
assert schedule.work(dur) == 10.0 # 1 + 4×2 + 1
|
||
span, _ = schedule.critical_path(dur, preds)
|
||
assert span == 4.0 # root(1) + one leaf(2) + sink(1)
|
||
assert schedule.peak_concurrency(dur, preds) == 4 # 4 leaves at once
|
||
# 1 machine == total work
|
||
assert schedule.makespan(dur, preds, 1) == 10.0
|
||
# 4 machines: root, then 4 leaves in parallel (2), then sink -> 4
|
||
assert schedule.makespan(dur, preds, 4) == 4.0
|
||
# a 5th machine cannot help beyond the width
|
||
assert schedule.makespan(dur, preds, 5) == schedule.makespan(dur, preds, 4)
|
||
|
||
|
||
def test_long_pole_dominates_span():
|
||
# one 40-min derivation gates a pile of tiny ones
|
||
dur = {"onnx": 40.0, **{f"lib{i}": 1.0 for i in range(20)}, "img": 1.0}
|
||
preds = {"onnx": [], **{f"lib{i}": [] for i in range(20)},
|
||
"img": ["onnx"] + [f"lib{i}" for i in range(20)]}
|
||
span, _ = schedule.critical_path(dur, preds)
|
||
assert span == 41.0 # onnx -> img
|
||
# even with many machines, makespan is pinned by the long pole
|
||
assert schedule.makespan(dur, preds, 16) == 41.0
|
||
|
||
|
||
def test_critical_path_keeps_zero_duration_predecessor():
|
||
# Issue #9: a 0-minute node (e.g. a cached/history entry) must not be
|
||
# dropped from the reconstructed chain. z(0) -> a(1) -> b(1).
|
||
dur = {"z": 0.0, "a": 1.0, "b": 1.0}
|
||
preds = {"z": [], "a": ["z"], "b": ["a"]}
|
||
span, chain = schedule.critical_path(dur, preds)
|
||
assert span == 2.0 # 0 + 1 + 1
|
||
assert chain == ["z", "a", "b"] # full chain, z not truncated
|
||
|
||
|
||
def test_makespan_return_schedule_is_consistent():
|
||
# Issue #16: return_schedule exposes per-machine (task, start, finish) rows.
|
||
dur = {"root": 1.0, "l1": 2.0, "l2": 2.0, "l3": 2.0, "l4": 2.0, "sink": 1.0}
|
||
preds = {"root": [], "l1": ["root"], "l2": ["root"], "l3": ["root"],
|
||
"l4": ["root"], "sink": ["l1", "l2", "l3", "l4"]}
|
||
ms, sched = schedule.makespan(dur, preds, 4, return_schedule=True)
|
||
# scalar makespan matches the back-compat return
|
||
assert ms == schedule.makespan(dur, preds, 4)
|
||
assert set(sched) == {0, 1, 2, 3}
|
||
|
||
seen = []
|
||
for machine, rows in sched.items():
|
||
prev_finish = -1.0
|
||
for task, start, finish in rows:
|
||
# finish == start + dur, and nothing runs past the makespan
|
||
assert finish == start + dur[task]
|
||
assert start >= 0.0 and finish <= ms + 1e-9
|
||
# no overlap on a single machine: rows are ordered, non-overlapping
|
||
assert start >= prev_finish - 1e-9
|
||
prev_finish = finish
|
||
seen.append(task)
|
||
# every task placed exactly once
|
||
assert sorted(seen) == sorted(dur)
|
||
|
||
|
||
def test_makespan_schedule_respects_dependencies():
|
||
# A task cannot start before every predecessor has finished.
|
||
dur = {"a": 1.0, "b": 2.0, "c": 1.0}
|
||
preds = {"a": [], "b": ["a"], "c": ["b"]}
|
||
_, sched = schedule.makespan(dur, preds, 2, return_schedule=True)
|
||
starts = {task: start for rows in sched.values() for task, start, _ in rows}
|
||
finish = {task: fin for rows in sched.values() for task, _, fin in rows}
|
||
for node, ps in preds.items():
|
||
for p in ps:
|
||
assert starts[node] >= finish[p] - 1e-9
|
||
|
||
|
||
def test_max_jobs_speeds_wide_graph_of_small_drvs():
|
||
# Issue #13: many independent small builds on a single node. With one job
|
||
# slot they serialize; max_jobs>1 runs several at once -> shorter makespan.
|
||
dur = {f"n{i}": 1.0 for i in range(8)}
|
||
preds = {f"n{i}": [] for i in range(8)}
|
||
one = schedule.makespan(dur, preds, machines=1, max_jobs=1)
|
||
four = schedule.makespan(dur, preds, machines=1, max_jobs=4)
|
||
assert one == 8.0 # 8 builds serialize on one slot
|
||
assert four == 2.0 # 4 lanes -> two rounds of four
|
||
assert four < one
|
||
|
||
|
||
def test_max_jobs_lanes_exposed_in_schedule():
|
||
# machines*max_jobs distinct lanes appear in the returned schedule.
|
||
dur = {f"n{i}": 1.0 for i in range(6)}
|
||
preds = {f"n{i}": [] for i in range(6)}
|
||
_, sched = schedule.makespan(
|
||
dur, preds, machines=2, max_jobs=3, return_schedule=True
|
||
)
|
||
assert set(sched) == set(range(6)) # 2 nodes * 3 jobs = 6 lanes
|
||
|
||
|
||
def test_ram_budget_caps_concurrency():
|
||
# Issue #15: 4 job slots on one node, but only enough RAM for 2 of these
|
||
# 3-GB builds at a time -> they run two-at-a-time, doubling the makespan
|
||
# versus the RAM-unconstrained case.
|
||
dur = {f"n{i}": 1.0 for i in range(4)}
|
||
preds = {f"n{i}": [] for i in range(4)}
|
||
gb = {f"n{i}": 3.0 for i in range(4)}
|
||
free = schedule.makespan(dur, preds, 1, max_jobs=4) # no RAM cap
|
||
capped = schedule.makespan(
|
||
dur, preds, 1, max_jobs=4, gb_per_job=gb, node_ram_gb=6.0
|
||
)
|
||
assert free == 1.0 # all 4 at once
|
||
assert capped == 2.0 # 2 + 2 -> two rounds
|
||
assert capped > free
|
||
|
||
|
||
def test_ram_oversized_job_still_runs_alone():
|
||
# A single build needing more than the whole node budget must not deadlock:
|
||
# it runs alone (need clamped to the budget).
|
||
dur = {"big": 1.0, "small": 1.0}
|
||
preds = {"big": [], "small": []}
|
||
gb = {"big": 100.0, "small": 1.0}
|
||
ms = schedule.makespan(
|
||
dur, preds, 1, max_jobs=2, gb_per_job=gb, node_ram_gb=4.0
|
||
)
|
||
# big monopolises RAM -> small waits -> serialized: 2.0
|
||
assert ms == 2.0
|
||
|
||
|
||
def test_ram_schedule_no_overlap_within_budget():
|
||
# With the RAM cap the returned schedule stays consistent and never exceeds
|
||
# the node budget at any instant.
|
||
dur = {f"n{i}": 1.0 + 0.1 * i for i in range(6)}
|
||
preds = {f"n{i}": [] for i in range(6)}
|
||
gb = {f"n{i}": 2.0 for i in range(6)}
|
||
ms, sched = schedule.makespan(
|
||
dur, preds, 2, max_jobs=3, gb_per_job=gb, node_ram_gb=4.0,
|
||
return_schedule=True,
|
||
)
|
||
# reconstruct per-node concurrent RAM over time from the intervals
|
||
intervals_by_node: dict[int, list[tuple[float, float, float]]] = {}
|
||
for lane, rows in sched.items():
|
||
node = lane // 3
|
||
for task, start, finish in rows:
|
||
intervals_by_node.setdefault(node, []).append((start, finish, gb[task]))
|
||
edges = sorted({t for ivs in intervals_by_node.values()
|
||
for s, f, _ in ivs for t in (s, f)})
|
||
for node, ivs in intervals_by_node.items():
|
||
for probe in edges:
|
||
used = sum(g for s, f, g in ivs if s <= probe < f)
|
||
assert used <= 4.0 + 1e-9
|
||
|
||
|
||
def test_max_jobs_zero_rejected():
|
||
dur = {"a": 1.0}
|
||
preds = {"a": []}
|
||
try:
|
||
schedule.makespan(dur, preds, 1, max_jobs=0)
|
||
except ValueError:
|
||
pass
|
||
else:
|
||
raise AssertionError("max_jobs=0 must raise")
|
||
|
||
|
||
def test_makespan_perf_smoke_10k_nodes():
|
||
# Issue #10: heap-based dispatch must handle a large DAG quickly.
|
||
import random
|
||
import time
|
||
|
||
rng = random.Random(1234) # seeded: deterministic across runs
|
||
n = 10_000
|
||
dur = {f"n{i}": float(rng.randint(1, 5)) for i in range(n)}
|
||
preds: dict[str, list[str]] = {}
|
||
for i in range(n):
|
||
# each node depends on a few earlier nodes -> acyclic by construction
|
||
k = rng.randint(0, 3) if i else 0
|
||
preds[f"n{i}"] = [f"n{rng.randint(0, i - 1)}" for _ in range(k)] if i else []
|
||
preds[f"n{i}"] = list(dict.fromkeys(preds[f"n{i}"])) # dedupe
|
||
|
||
start = time.perf_counter()
|
||
ms = schedule.makespan(dur, preds, machines=32)
|
||
elapsed = time.perf_counter() - start
|
||
assert ms > 0.0
|
||
assert elapsed < 5.0, f"makespan too slow: {elapsed:.2f}s"
|