diff --git a/nix_estimator/costmodel.py b/nix_estimator/costmodel.py index 35d3754..d750bc8 100644 --- a/nix_estimator/costmodel.py +++ b/nix_estimator/costmodel.py @@ -55,6 +55,31 @@ HEAVY: dict[str, tuple[float, float]] = { "gnutls": (5.0, 0.50), } +# substring-of-store-name -> peak RAM (GB) a single build of it needs. Coarse: +# LTO / whole-program C++ and big compilers are the memory hogs; most library +# builds sit near 1 GB. Used to cap how many jobs a node can run concurrently +# (issue #15) — sum of co-resident jobs' RAM must fit the node budget. +RAM_HEAVY: dict[str, float] = { + "qtwebengine": 10.0, + "chromium": 10.0, # LTO link phase is the peak + "webkitgtk": 8.0, + "tensorflow": 8.0, + "pytorch": 8.0, + "libtorch": 8.0, + "ghc": 6.0, + "llvm": 4.0, + "clang": 4.0, + "rustc": 4.0, + "gcc": 3.0, + "qtbase": 3.0, + "onnxruntime": 3.0, + "opencv": 2.0, + "boost": 2.0, + "nodejs": 2.0, + "grpc": 2.0, +} +DEFAULT_RAM_GB: float = 1.0 # most library derivations + DEFAULT: tuple[float, float] = (1.0, 0.30) # most library derivations FIXED_OUTPUT: tuple[float, float] = (0.3, 0.0) # source fetches: network-bound TRIVIAL: tuple[float, float] = (0.05, 0.0) # NixOS assembly glue: ~instant @@ -112,6 +137,30 @@ _HEAVY_RE: dict[re.Pattern[str], tuple[float, float]] = { } +_RAM_RE: dict[re.Pattern[str], float] = { + re.compile(r"^" + re.escape(key.rstrip("-")) + r"(?:-|$)"): val + for key, val in RAM_HEAVY.items() +} + + +def ram_gb(drv_path: str, drv: dict | None = None) -> float: + """Peak RAM (GB) a single build of ``drv_path`` needs (coarse, issue #15). + + Fixed-output fetches and NixOS glue barely use memory; the heavy compilers + and LTO links in ``RAM_HEAVY`` dominate. Everything else defaults to ~1 GB. + """ + name = store_name(drv_path) + if drv is not None and is_fixed_output(drv): + return DEFAULT_RAM_GB + if is_trivial(name) or is_shim(name): + return DEFAULT_RAM_GB + low = name.lower() + for pattern, val in _RAM_RE.items(): + if pattern.search(low): + return val + return DEFAULT_RAM_GB + + def is_shim(name: str) -> bool: """True for wrapper/doc/dev-style names that must not be costed as HEAVY.""" return bool(_SHIM_RE.search(name)) diff --git a/nix_estimator/estimate.py b/nix_estimator/estimate.py index 10ee3bd..1238726 100644 --- a/nix_estimator/estimate.py +++ b/nix_estimator/estimate.py @@ -17,15 +17,26 @@ class Estimate: avg_parallelism: float # work / span longest_chain: list[tuple[str, float]] # (store-name, min@8c) on crit path grid: dict[tuple[int, int], float] # (nodes, cores) -> makespan minutes + max_jobs: int = 1 # per-node concurrent builds the grid was computed at + node_ram_gb: float | None = None # per-node RAM budget the grid used (None=∞) recommendation: dict = field(default_factory=dict) -def _durations(preds_nodes, closure, history, cores): +def _durations(preds_nodes, closure, history, cores, max_jobs=1): + """Per-derivation minutes at ``cores`` cores per node. + + With ``max_jobs > 1`` the node's cores are shared across concurrent builds, + so each job sees ~``cores/max_jobs`` effective cores (steady-state + approximation — issue #13). Heavily parallel builds pay for the split; small, + weakly-parallel derivations barely notice, so wider job counts win on graphs + of many small drvs. + """ + eff_cores = max(1.0, cores / max_jobs) dur, scaling = {}, {} for d in preds_nodes: m8, s = costmodel.cost(d, closure.get(d), history) scaling[d] = s - dur[d] = costmodel.scale_to_cores(m8, s, cores) + dur[d] = costmodel.scale_to_cores(m8, s, eff_cores) return dur @@ -38,6 +49,8 @@ def estimate( node_grid=(1, 2, 3, 4, 6, 8, 12, 16), core_grid=(8, 16, 32), baseline_cores: int = 8, + max_jobs: int = 1, + node_ram_gb: float | None = None, knee_threshold: float = 0.12, ) -> Estimate: """Compute parallelism metrics and the (nodes, cores) makespan grid.""" @@ -59,11 +72,21 @@ def estimate( max(base_dur[d] for d in chain) / span if chain and span else 0.0 ) + # Per-derivation peak RAM, used only when a node budget is set (issue #15). + gb_per_job = ( + {d: costmodel.ram_gb(d, closure.get(d)) for d in nodes} + if node_ram_gb is not None + else None + ) + grid: dict[tuple[int, int], float] = {} for cores in core_grid: - dur_c = _durations(nodes, closure, history, cores) + dur_c = _durations(nodes, closure, history, cores, max_jobs) for p in node_grid: - grid[(p, cores)] = schedule.makespan(dur_c, preds, p) + grid[(p, cores)] = schedule.makespan( + dur_c, preds, p, max_jobs=max_jobs, + gb_per_job=gb_per_job, node_ram_gb=node_ram_gb, + ) rec = _recommend( grid, peak, core_grid, node_grid, knee_threshold, span_dominator_frac @@ -77,6 +100,8 @@ def estimate( avg_parallelism=avg, longest_chain=[(costmodel.store_name(d), base_dur[d]) for d in chain], grid=grid, + max_jobs=max_jobs, + node_ram_gb=node_ram_gb, recommendation=rec, ) @@ -89,10 +114,13 @@ def _recommend(grid, peak, core_grid, node_grid, knee, span_dominator_frac=0.0): Nodes: the fewest nodes within ``knee`` of the best makespan at that core count (adding nodes past the graph width or past the span floor is waste). """ + # "Single node" reference is the smallest node count actually in the grid — + # callers may pass a node_grid that omits 1 (issue #18), so never assume it. + single = node_grid[0] best_cores = core_grid[0] for c in core_grid[1:]: - prev = grid[(1, best_cores)] - cur = grid[(1, c)] + prev = grid[(single, best_cores)] + cur = grid[(single, c)] if prev and (prev - cur) / prev >= knee: best_cores = c at = {p: grid[(p, best_cores)] for p in node_grid} @@ -102,7 +130,13 @@ def _recommend(grid, peak, core_grid, node_grid, knee, span_dominator_frac=0.0): if best_makespan and at[p] <= best_makespan * (1 + knee): chosen_nodes = p break - chosen_nodes = min(chosen_nodes, max(1, peak)) + # Clamp to the graph width, then snap back onto node_grid: the clamp target + # (peak) need not be a grid value, and if it undershoots every grid entry a + # bare ``at[chosen_nodes]`` would KeyError (issue #18). Pick the largest grid + # value <= the target, or the smallest grid value if the target is below all. + clamp_target = min(chosen_nodes, max(1, peak)) + below = [p for p in node_grid if p <= clamp_target] + chosen_nodes = max(below) if below else min(node_grid) if span_dominator_frac >= 0.5: pole = ( "The long pole is one heavy derivation " @@ -120,7 +154,7 @@ def _recommend(grid, peak, core_grid, node_grid, knee, span_dominator_frac=0.0): "nodes": chosen_nodes, "cores_per_node": best_cores, "est_makespan_min": round(at[chosen_nodes], 1), - "one_big_node_min": round(grid[(1, core_grid[-1])], 1), + "one_big_node_min": round(grid[(single, core_grid[-1])], 1), "note": ( f"~{chosen_nodes}×{best_cores}-core. Node ceiling (graph width) = " f"{peak}; beyond it nodes idle. {pole}" diff --git a/nix_estimator/schedule.py b/nix_estimator/schedule.py index 4e647b7..8a1afff 100644 --- a/nix_estimator/schedule.py +++ b/nix_estimator/schedule.py @@ -87,15 +87,49 @@ def peak_concurrency(dur: dict[str, float], preds: dict[str, list[str]]) -> int: return peak -def makespan(dur: dict[str, float], preds: dict[str, list[str]], machines: int, - priority: dict[str, float] | None = None) -> float: +Assignment = tuple[str, float, float] # (task, start, finish) + + +def makespan( + dur: dict[str, float], + preds: dict[str, list[str]], + machines: int, + priority: dict[str, float] | None = None, + *, + max_jobs: int = 1, + gb_per_job: dict[str, float] | None = None, + node_ram_gb: float | None = None, + return_schedule: bool = False, +) -> float | tuple[float, dict[int, list[Assignment]]]: """Estimated wall-clock with ``machines`` builders, greedy list scheduling. Ready tasks are dispatched highest-priority first (default: longest path to a sink — the classic critical-path heuristic, near-optimal in practice). + + ``max_jobs`` is Nix's per-node concurrent-build setting: each of the + ``machines`` nodes offers ``max_jobs`` job slots, so up to + ``machines * max_jobs`` tasks run at once (issue #13). Core-sharing between + co-resident jobs (each gets ~``cores/max_jobs`` effective cores) is modelled + upstream in ``estimate()`` by pre-scaling ``dur`` — this function stays pure + and simply schedules onto ``machines * max_jobs`` lanes. + + With ``return_schedule=True`` returns ``(makespan, sched)`` where ``sched`` + maps each lane id ``0..machines*max_jobs-1`` to its ``(task, start, finish)`` + assignments in dispatch order (for timeline/Gantt rendering, issue #16). The + default scalar return is unchanged for back-compat. + + ``node_ram_gb`` (with per-task ``gb_per_job``, default 1 GB each) is a second, + orthogonal constraint (issue #15): a node's co-resident jobs must fit its RAM + budget, so a memory-hungry mix caps concurrency below ``max_jobs`` and pushes + the makespan up. Left ``None`` the RAM check is skipped (fast path). """ if machines < 1: raise ValueError("machines must be >= 1") + if max_jobs < 1: + raise ValueError("max_jobs must be >= 1") + n_nodes = machines + lanes = machines * max_jobs + machines = lanes # remaining code schedules onto flat lanes prio = priority or _path_to_sink(dur, preds) succ = _succs(preds) indeg = {n: len(preds.get(n, ())) for n in dur} @@ -105,26 +139,71 @@ def makespan(dur: dict[str, float], preds: dict[str, list[str]], machines: int, ready: list[tuple[float, str]] = [(-prio[n], n) for n, d in indeg.items() if d == 0] heapq.heapify(ready) - running: list[tuple[float, str]] = [] # (finish_time, node) min-heap + # Concrete machine ids drawn from a min-heap free pool so assignments land on + # a stable, lowest-available builder — needed to attribute each task to a + # machine without overlap (issue #16). + free_slots: list[int] = list(range(machines)) + heapq.heapify(free_slots) + # (finish, lane, node, ram) min-heap. node = lane // max_jobs groups the + # ``max_jobs`` lanes that share one physical builder's RAM budget. + running: list[tuple[float, int, int, float]] = [] + sched: dict[int, list[Assignment]] = {m: [] for m in range(machines)} + ram_capped = node_ram_gb is not None + node_ram_used: dict[int, float] = {node: 0.0 for node in range(n_nodes)} t = 0.0 - free = machines done = 0 total = len(dur) while done < total: - while free > 0 and ready: - _, n = heapq.heappop(ready) - heapq.heappush(running, (t + dur[n], n)) - free -= 1 + if not ram_capped: + while free_slots and ready: + _, n = heapq.heappop(ready) + m = heapq.heappop(free_slots) + finish = t + dur[n] + heapq.heappush(running, (finish, m, m // max_jobs, 0.0)) + sched[m].append((n, t, finish)) + else: + # RAM-aware placement: dispatch the highest-priority ready task onto a + # free lane whose node still has headroom. A single job always fits + # alone (its need is clamped to the node budget), so progress is + # guaranteed while anything is running. + while free_slots and ready: + _, n = ready[0] + need = min((gb_per_job or {}).get(n, 1.0), node_ram_gb) + stash: list[int] = [] + placed = False + while free_slots: + m = heapq.heappop(free_slots) + node = m // max_jobs + if node_ram_used[node] + need <= node_ram_gb + 1e-9: + heapq.heappop(ready) + node_ram_used[node] += need + finish = t + dur[n] + heapq.heappush(running, (finish, m, node, need)) + sched[m].append((n, t, finish)) + placed = True + break + stash.append(m) + for s in stash: + heapq.heappush(free_slots, s) + if not placed: + break # no node can host it now — wait for a job to finish if not running: raise ValueError("deadlock — cycle in graph") - ft, n = heapq.heappop(running) + ft, m, node, need = heapq.heappop(running) + # the finishing task's node id is recoverable from the lane, but we look + # its store name up from the assignment we recorded on that lane. + n = sched[m][-1][0] t = ft - free += 1 + heapq.heappush(free_slots, m) + if ram_capped: + node_ram_used[node] -= need done += 1 for c in succ.get(n, ()): indeg[c] -= 1 if indeg[c] == 0: heapq.heappush(ready, (-prio[c], c)) + if return_schedule: + return t, sched return t diff --git a/tests/test_estimate.py b/tests/test_estimate.py index 3f5b827..7644100 100644 --- a/tests/test_estimate.py +++ b/tests/test_estimate.py @@ -48,6 +48,69 @@ def test_unsorted_grids_duplicates_deduped(): 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. diff --git a/tests/test_schedule.py b/tests/test_schedule.py index 8de1e1e..d11c946 100644 --- a/tests/test_schedule.py +++ b/tests/test_schedule.py @@ -55,6 +55,129 @@ def test_critical_path_keeps_zero_duration_predecessor(): 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