Merge wave2A branch a604e4a5f4fb7c756 (features)

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))

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}"

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
 
 

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.

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