From a0852ff479c9a769c6ccc29fdceed2dfbd0376b3 Mon Sep 17 00:00:00 2001 From: Oleks Date: Tue, 7 Jul 2026 19:45:25 +0300 Subject: [PATCH] =?UTF-8?q?nix-estimator=20v0.1:=20DAG=20=E2=86=92=20node?= =?UTF-8?q?=C3=97core=20builder-config=20estimator?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Reads the to-build derivation DAG from Nix (derivation show -r + --dry-run), applies a heuristic/history build-cost model with Amdahl core-scaling, and computes critical-path / peak-concurrency / list-scheduled makespan across a node×core grid to recommend a remote-builder shape. - schedule.py: pure critical-path, peak-concurrency, p-machine list scheduler - costmodel.py: heuristic table + --history override + core re-weighting - graph.py: DAG extraction via nix - estimate.py/cli.py: sweep + knee recommendation + report - tests: scheduler validated on toy DAGs (all pass) Prior-art gap documented (Hubble, nixbuild.net) in README. --- .gitignore | 10 +++ README.md | 113 ++++++++++++++++++++++++++++++ flake.nix | 35 ++++++++++ nix_estimator/__init__.py | 14 ++++ nix_estimator/cli.py | 111 +++++++++++++++++++++++++++++ nix_estimator/costmodel.py | 104 +++++++++++++++++++++++++++ nix_estimator/estimate.py | 97 ++++++++++++++++++++++++++ nix_estimator/graph.py | 78 +++++++++++++++++++++ nix_estimator/schedule.py | 139 +++++++++++++++++++++++++++++++++++++ pyproject.toml | 26 +++++++ tests/test_schedule.py | 45 ++++++++++++ 11 files changed, 772 insertions(+) create mode 100644 .gitignore create mode 100644 README.md create mode 100644 flake.nix create mode 100644 nix_estimator/__init__.py create mode 100644 nix_estimator/cli.py create mode 100644 nix_estimator/costmodel.py create mode 100644 nix_estimator/estimate.py create mode 100644 nix_estimator/graph.py create mode 100644 nix_estimator/schedule.py create mode 100644 pyproject.toml create mode 100644 tests/test_schedule.py diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..e043c10 --- /dev/null +++ b/.gitignore @@ -0,0 +1,10 @@ +__pycache__/ +*.pyc +.pytest_cache/ +*.egg-info/ +dist/ +build/ +result +result-* +.direnv/ +.venv/ diff --git a/README.md b/README.md new file mode 100644 index 0000000..f301956 --- /dev/null +++ b/README.md @@ -0,0 +1,113 @@ +# nix-estimator + +Estimate the **best remote-builder configuration — how many builder nodes × +how many cores per node — for a given Nix build**, before you provision anything. + +It reads the *to-build* derivation DAG from Nix, applies a build-time cost model, +and computes the parallel-computing quantities that bound each dimension, then +sweeps `(nodes × cores)` through a list-scheduler to recommend a shape. + +``` +$ nix-estimate .#packages.aarch64-linux.mempalace-image-arm64 --system aarch64-linux +``` + +## Why + +More builder *nodes* only help while the dependency graph is **wide** (many +independent derivations ready at once). Once a build narrows to a long pole — +one big derivation like `onnxruntime`, `llvm`, or `ghc` — no number of nodes +helps, because Nix builds a single derivation on a single machine. That long +pole is moved only by **cores on one node**. nix-estimator makes this concrete +for *your* build instead of leaving it to intuition. + +## The model + +The build is a DAG of derivations. Two classic quantities govern it: + +| quantity | meaning | bounds | +|----------|---------|--------| +| **work** `T₁` | Σ of every build time | wall-clock on 1 core / 1 node | +| **span** `T∞` | longest *dependent* chain | floor no parallelism beats | +| **width** | peak simultaneously-ready builds | **node-count ceiling** | + +Brent's bound `T_p ≈ max(T₁/p, T∞)` says it all: node-parallelism is capped by +`width` and by `T₁/T∞` (average available parallelism); core-parallelism is set +by the heaviest single derivation's internal `-j` scaling. nix-estimator reports +all of these plus a `(nodes × cores) → makespan` grid from a critical-path +list-scheduler, and picks the diminishing-returns knee. + +### The cost model + +Nix does not predict per-derivation build time, so `costmodel.py` approximates +it: a coarse heuristic table (heavy: onnxruntime/llvm/ghc/rustc/…; default: ~1 +min; fixed-output fetches: network-bound) plus an Amdahl `core_scaling` per +derivation. **Pass `--history `** (`{name: minutes}` mined from real `nom` +/ nix build logs) for accuracy — history always wins over the heuristic. + +## Usage + +```sh +# human report (default) +nix-estimate .#foo --system aarch64-linux + +# machine-readable +nix-estimate .#foo --json + +# sweep custom grids, use measured build times +nix-estimate .#foo --nodes 1,2,4,8 --cores 8,16,32 --history builds.json +``` + +Output: closure size, to-build count, `work`/`span`/`avg`/`peak`, the critical +path (the long pole), a makespan grid, and a `RECOMMENDATION` (nodes × cores ≈ +minutes, vs. one-big-node). + +Requires `nix` on `PATH`. Pure stdlib otherwise. `nix develop` for a dev shell. + +```sh +nix develop # python + pytest + nix +pytest # unit tests for the scheduler (toy DAGs, no Nix needed) +``` + +## Layout + +``` +nix_estimator/ + graph.py DAG + to-build set via `nix derivation show -r` + `--dry-run` + costmodel.py heuristic/history build-time + Amdahl core-scaling + schedule.py critical path, peak concurrency, p-machine list scheduler (pure) + estimate.py orchestration + node×core sweep + knee recommendation + cli.py `nix-estimate` entrypoint + report +tests/ scheduler unit tests on toy graphs +``` + +## Prior art (and where this sits) + +The node×core **provisioning recommendation for Nix is an unfilled gap.** The +ingredients exist separately; nix-estimator assembles them: + +- **[Hubble](https://gitlab.inria.fr/lcourtes/hubble)** (INRIA, Ludovic Courtès) — + the closest analytical core: a SimGrid simulator of scheduling strategies on + the Nix/Hydra DAG with a `critical-path` tool and speedup plots. But it + evaluates *scheduling algorithms* as a (dormant, ~2010-era) research artifact + — it does not output a fleet-size recommendation. +- **[nixbuild.net](https://docs.nixbuild.net/remote-builds/)** — solves sizing + *operationally*, auto-assigning CPU/RAM per derivation from history, but as a + closed hosted autoscaler with no DAG-level analysis you can run yourself. +- **DAG extractors** — [nom](https://github.com/maralorn/nix-output-monitor), + [nix-tree](https://github.com/utdemir/nix-tree), + [nix-visualize](https://github.com/craigmbooth/nix-visualize) — topology only, + no timing/scheduling. +- **Build-time cost models** — Uber's + [CI at Scale](https://arxiv.org/pdf/2501.03440) (NGBoost per-target Bazel + build-time prediction) — methodology to borrow for a learned cost model. +- **Theory** — list scheduling, Brent's theorem, HEFT. nix-estimator's novelty + is *applying* it to the Nix derivation DAG for provisioning, not the theory. + +## Status + +v0.1 — heuristic cost model, working scheduler + recommendation. Roadmap: +mine real per-derivation timings from `nom`/nix logs into a history file +(biggest accuracy win), RAM-per-core as a second constraint, and a +`--provision` mode that emits an `EphemeralBuilder` spec for the chosen shape. + +MIT. diff --git a/flake.nix b/flake.nix new file mode 100644 index 0000000..fba617d --- /dev/null +++ b/flake.nix @@ -0,0 +1,35 @@ +{ + description = "nix-estimator — best builder configuration (nodes × cores) for a Nix build"; + + inputs.nixpkgs.url = "github:NixOS/nixpkgs/nixos-unstable"; + + outputs = { self, nixpkgs }: + let + systems = [ "x86_64-linux" "aarch64-linux" "aarch64-darwin" ]; + forAll = f: nixpkgs.lib.genAttrs systems (s: f nixpkgs.legacyPackages.${s}); + in + { + packages = forAll (pkgs: { + default = pkgs.python3Packages.buildPythonApplication { + pname = "nix-estimator"; + version = "0.1.0"; + pyproject = true; + src = ./.; + build-system = [ pkgs.python3Packages.hatchling ]; + # runtime needs `nix` on PATH to introspect derivations + makeWrapperArgs = [ "--prefix" "PATH" ":" "${pkgs.nix}/bin" ]; + nativeCheckInputs = [ pkgs.python3Packages.pytest ]; + checkPhase = "pytest"; + }; + }); + + devShells = forAll (pkgs: { + default = pkgs.mkShell { + packages = [ + (pkgs.python3.withPackages (ps: [ ps.pytest ])) + pkgs.nix + ]; + }; + }); + }; +} diff --git a/nix_estimator/__init__.py b/nix_estimator/__init__.py new file mode 100644 index 0000000..adf6f3f --- /dev/null +++ b/nix_estimator/__init__.py @@ -0,0 +1,14 @@ +"""nix-estimator — estimate the best builder configuration (node × core parallelism) +for a Nix flake attribute or derivation. + +The tool extracts the *to-build* derivation DAG from Nix, applies a build-cost +model, and computes the parallel-computing quantities that bound each dimension: + + - work (T1) = sum of all build times → wall-clock on one core + - span (T∞) = longest chain of dependent builds → floor no parallelism beats + - width = peak simultaneous-ready builds → node-count ceiling + +then sweeps (nodes × cores) through a list-scheduler to recommend a shape. +""" + +__version__ = "0.1.0" diff --git a/nix_estimator/cli.py b/nix_estimator/cli.py new file mode 100644 index 0000000..877c86c --- /dev/null +++ b/nix_estimator/cli.py @@ -0,0 +1,111 @@ +"""Command-line entrypoint: ``nix-estimate ``.""" + +from __future__ import annotations + +import argparse +import json +import sys + +from . import __version__, graph +from .estimate import estimate + + +def _load_history(path: str | None) -> dict[str, float] | None: + if not path: + return None + with open(path) as fh: + return json.load(fh) + + +def _report(attr: str, est) -> str: + L = [] + L.append(f"nix-estimator {__version__} — {attr}") + L.append("=" * 64) + L.append(f"derivations in closure : {est.nodes_evaluated}") + L.append(f"to build (cache-miss) : {est.to_build}") + if est.to_build == 0: + L.append("\nNothing to build — the whole closure substitutes from cache.") + return "\n".join(L) + L.append("") + L.append(f"work T1 (@8 cores) : {est.work_min:7.1f} min (1 core, 1 node)") + L.append(f"span T∞ (@8 cores) : {est.span_min:7.1f} min (critical-path floor)") + L.append(f"avg parallelism T1/T∞: {est.avg_parallelism:7.1f} (sustained node use)") + L.append(f"peak parallelism : {est.peak_parallelism:7d} (node ceiling; more idle)") + L.append("") + L.append("critical path (long pole, top 6):") + for name in est.longest_chain[-6:]: + L.append(f" → {name}") + L.append("") + cores = sorted({c for _, c in est.grid}) + nodes = sorted({p for p, _ in est.grid}) + L.append("estimated makespan (minutes) rows=nodes cols=cores/node") + L.append(" nodes │ " + " ".join(f"{c:>7}c" for c in cores)) + L.append(" ──────┼" + "─" * (10 * len(cores))) + for p in nodes: + row = " ".join(f"{est.grid[(p, c)]:7.0f} " for c in cores) + L.append(f" {p:5d} │ {row}") + L.append("") + r = est.recommendation + L.append("RECOMMENDATION") + L.append(f" {r['nodes']} node(s) × {r['cores_per_node']} cores" + f" ≈ {r['est_makespan_min']} min") + L.append(f" (one big {cores[-1]}-core node alone ≈ {r['one_big_node_min']} min)") + L.append(f" {r['note']}") + return "\n".join(L) + + +def main(argv: list[str] | None = None) -> int: + ap = argparse.ArgumentParser( + prog="nix-estimate", + description="Estimate the best builder configuration (nodes × cores) " + "for a Nix flake attr or derivation.") + ap.add_argument("attr", help="flake attr, e.g. .#packages.aarch64-linux.foo") + ap.add_argument("--system", help="e.g. aarch64-linux") + ap.add_argument("--history", help="JSON {name: minutes} from real build logs") + ap.add_argument("--nix-arg", action="append", default=[], + help="extra arg passed through to nix (repeatable)") + ap.add_argument("--nodes", default="1,2,3,4,6,8,12,16", + help="comma list of node counts to sweep") + ap.add_argument("--cores", default="8,16,32", + help="comma list of cores-per-node to sweep") + ap.add_argument("--json", action="store_true", help="emit JSON not a report") + ap.add_argument("--version", action="version", version=__version__) + args = ap.parse_args(argv) + + node_grid = tuple(int(x) for x in args.nodes.split(",")) + core_grid = tuple(int(x) for x in args.cores.split(",")) + + try: + closure = graph.derivation_closure(args.attr, args.system, args.nix_arg) + except Exception as e: # noqa: BLE001 + print(f"error: `nix derivation show` failed: {e}", file=sys.stderr) + return 2 + to_build = graph.to_build_set(args.attr, args.system, args.nix_arg) + if to_build is None: + print("warning: could not parse `nix build --dry-run`; assuming a COLD " + "cache (whole closure builds). Numbers are an upper bound.", + file=sys.stderr) + preds, nodes = graph.build_dag(closure, to_build) + + est = estimate(closure, preds, nodes, + history=_load_history(args.history), + node_grid=node_grid, core_grid=core_grid) + + if args.json: + print(json.dumps({ + "attr": args.attr, + "to_build": est.to_build, + "work_min": est.work_min, + "span_min": est.span_min, + "peak_parallelism": est.peak_parallelism, + "avg_parallelism": est.avg_parallelism, + "grid": {f"{p}x{c}": v for (p, c), v in est.grid.items()}, + "recommendation": est.recommendation, + }, indent=2)) + else: + print(_report(args.attr, est)) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/nix_estimator/costmodel.py b/nix_estimator/costmodel.py new file mode 100644 index 0000000..13a4291 --- /dev/null +++ b/nix_estimator/costmodel.py @@ -0,0 +1,104 @@ +"""Build-cost model for Nix derivations. + +Nix does not predict per-derivation build time, so we approximate it. Each +derivation is assigned: + + - ``minutes`` : rough wall-clock build time on a baseline **8-core** node. + - ``core_scaling`` : Amdahl parallel fraction in [0, 1] — how much of the build + speeds up with more cores (1.0 = perfectly parallel, + 0.0 = single-threaded). Used to re-weight when estimating a + different cores-per-node. + +The weights are deliberately coarse. For accuracy, pass a ``history`` mapping +(``{pname_or_storename: minutes}``) mined from real ``nom`` / nix build logs — +history always wins over the heuristic table. +""" + +from __future__ import annotations + +# substring-of-store-name -> (minutes @ 8 cores, core_scaling) +# Ordered roughly by how distinctive the key is; first substring match wins. +HEAVY: dict[str, tuple[float, float]] = { + "qtwebengine": (140.0, 0.90), + "chromium": (150.0, 0.90), + "webkitgtk": (90.0, 0.90), + "tensorflow": (90.0, 0.85), + "pytorch": (80.0, 0.85), + "libtorch": (80.0, 0.85), + "onnxruntime": (40.0, 0.85), + "llvm": (45.0, 0.90), + "ghc": (55.0, 0.60), + "rustc": (35.0, 0.80), + "gcc": (40.0, 0.85), + "qtbase": (40.0, 0.85), + "clang": (30.0, 0.90), + "opencv": (20.0, 0.80), + "grpc": (20.0, 0.80), + "compiler-rt": (18.0, 0.90), + "nodejs": (18.0, 0.85), + "boost": (14.0, 0.70), + "linux-": (14.0, 0.90), + "openblas": (12.0, 0.90), + "tokenizers": (10.0, 0.70), # rust + "glibc": (8.0, 0.70), + "hf-xet": (8.0, 0.70), # rust + "protobuf": (8.0, 0.70), + "glslang": (8.0, 0.70), + "cpython": (6.0, 0.50), + "python3": (6.0, 0.50), + "cryptography": (6.0, 0.60), + "maturin": (6.0, 0.60), + "cmake": (5.0, 0.60), + "openssl": (5.0, 0.50), + "gnutls": (5.0, 0.50), +} + +DEFAULT: tuple[float, float] = (1.0, 0.30) # most library derivations +FIXED_OUTPUT: tuple[float, float] = (0.3, 0.0) # source fetches: network-bound + + +def is_fixed_output(drv: dict) -> bool: + """A fixed-output derivation (fetchurl/fetchgit/...) — network, not CPU.""" + for out in (drv.get("outputs") or {}).values(): + if out.get("hash") or out.get("hashAlgo"): + return True + return False + + +def store_name(drv_path: str) -> str: + """`/nix/store/-.drv` -> `` (hash + .drv stripped).""" + base = drv_path.rsplit("/", 1)[-1] + if base.endswith(".drv"): + base = base[:-4] + return base.split("-", 1)[1] if "-" in base else base + + +def cost(drv_path: str, drv: dict | None = None, + history: dict[str, float] | None = None) -> tuple[float, float]: + """Return ``(minutes @ 8 cores, core_scaling)`` for a derivation.""" + name = store_name(drv_path) + if history: + if name in history: + return (float(history[name]), 0.7) + # allow history keyed by pname prefix + for key, mins in history.items(): + if name.startswith(key): + return (float(mins), 0.7) + if drv is not None and is_fixed_output(drv): + return FIXED_OUTPUT + low = name.lower() + for key, val in HEAVY.items(): + if key in low: + return val + return DEFAULT + + +def scale_to_cores(minutes8: float, core_scaling: float, cores: int) -> float: + """Amdahl re-weight of an 8-core baseline time to ``cores`` cores. + + speedup(8 -> c) = 1 / ((1 - s) + s * 8 / c); duration divides by it. + For c < 8 the serial part dominates and heavy builds get *slower*. + """ + s = max(0.0, min(1.0, core_scaling)) + factor = (1.0 - s) + s * (8.0 / max(1, cores)) + return minutes8 * factor diff --git a/nix_estimator/estimate.py b/nix_estimator/estimate.py new file mode 100644 index 0000000..e11f5c2 --- /dev/null +++ b/nix_estimator/estimate.py @@ -0,0 +1,97 @@ +"""Orchestration: DAG + cost model -> node×core sweep -> recommendation.""" + +from __future__ import annotations + +from dataclasses import dataclass, field + +from . import costmodel, schedule + + +@dataclass +class Estimate: + nodes_evaluated: int + to_build: int + work_min: float # T1 @ 8 cores + span_min: float # T∞ @ 8 cores + peak_parallelism: int # node ceiling + avg_parallelism: float # work / span + longest_chain: list[str] # store-names on the critical path + grid: dict[tuple[int, int], float] # (nodes, cores) -> makespan minutes + recommendation: dict = field(default_factory=dict) + + +def _durations(preds_nodes, closure, history, cores): + 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) + return dur + + +def estimate(closure: dict[str, dict], preds: dict[str, list[str]], + nodes: set[str], *, history: dict[str, float] | None = None, + node_grid=(1, 2, 3, 4, 6, 8, 12, 16), + core_grid=(8, 16, 32), + baseline_cores: int = 8, + knee_threshold: float = 0.12) -> Estimate: + """Compute parallelism metrics and the (nodes, cores) makespan grid.""" + base_dur = _durations(nodes, closure, history, baseline_cores) + work = schedule.work(base_dur) + span, chain = schedule.critical_path(base_dur, preds) + peak = schedule.peak_concurrency(base_dur, preds) + avg = work / span if span else 0.0 + + grid: dict[tuple[int, int], float] = {} + for cores in core_grid: + dur_c = _durations(nodes, closure, history, cores) + for p in node_grid: + grid[(p, cores)] = schedule.makespan(dur_c, preds, p) + + rec = _recommend(grid, peak, core_grid, node_grid, knee_threshold) + return Estimate( + nodes_evaluated=len(closure), + to_build=len(nodes), + work_min=work, + span_min=span, + peak_parallelism=peak, + avg_parallelism=avg, + longest_chain=[costmodel.store_name(d) for d in chain], + grid=grid, + recommendation=rec, + ) + + +def _recommend(grid, peak, core_grid, node_grid, knee): + """Pick a (nodes, cores) at the diminishing-returns knee. + + Cores: the largest core count that still meaningfully cuts makespan at a + single node (the long pole is one derivation — cores, not nodes, move it). + 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). + """ + best_cores = core_grid[0] + for c in core_grid[1:]: + prev = grid[(1, best_cores)] + cur = grid[(1, c)] + if prev and (prev - cur) / prev >= knee: + best_cores = c + at = {p: grid[(p, best_cores)] for p in node_grid} + best_makespan = min(at.values()) + chosen_nodes = node_grid[-1] + for p in node_grid: + if best_makespan and at[p] <= best_makespan * (1 + knee): + chosen_nodes = p + break + chosen_nodes = min(chosen_nodes, max(1, peak)) + return { + "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), + "note": ( + f"~{chosen_nodes}×{best_cores}-core. Node ceiling (graph width) = " + f"{peak}; beyond it nodes idle. The long pole is one derivation, so " + f"cores/node — not node count — sets the floor." + ), + } diff --git a/nix_estimator/graph.py b/nix_estimator/graph.py new file mode 100644 index 0000000..3d53a8d --- /dev/null +++ b/nix_estimator/graph.py @@ -0,0 +1,78 @@ +"""Extract the to-build derivation DAG for a flake attr, using Nix itself. + +``nix derivation show -r `` emits JSON for the *whole* closure: +``{ drvPath: { name, outputs, inputDrvs, env, ... } }``. + +``nix build --dry-run `` reports which of those will actually be **built** +(vs substituted from a binary cache). We intersect the two so the estimate only +counts real compiles — a warm cache collapses the graph, and that must be +reflected. ``inputDrvs`` gives the edges (a derivation depends on its inputDrvs). +""" + +from __future__ import annotations + +import json +import re +import subprocess + + +def _run(cmd: list[str], check: bool = True) -> subprocess.CompletedProcess: + return subprocess.run(cmd, capture_output=True, text=True, check=check) + + +def derivation_closure(attr: str, system: str | None = None, + extra_args: list[str] | None = None) -> dict[str, dict]: + """Full derivation closure as ``{drvPath: derivation-json}``.""" + cmd = ["nix", "derivation", "show", "-r", attr] + if system: + cmd += ["--system", system] + if extra_args: + cmd += extra_args + return json.loads(_run(cmd).stdout) + + +def to_build_set(attr: str, system: str | None = None, + extra_args: list[str] | None = None) -> set[str] | None: + """Set of ``.drv`` paths Nix says *will be built*, or ``None`` if unparseable. + + ``None`` signals the caller to fall back to "the whole closure" (a cold-cache + over-estimate) rather than silently reporting zero work. + """ + cmd = ["nix", "build", "--dry-run", attr] + if system: + cmd += ["--system", system] + if extra_args: + cmd += extra_args + proc = _run(cmd, check=False) + text = proc.stderr + "\n" + proc.stdout + built: set[str] = set() + grabbing = False + for line in text.splitlines(): + s = line.strip() + if re.search(r"will be built", s): + grabbing = True + continue + if re.search(r"will be fetched|will be copied", s): + grabbing = False + continue + if grabbing: + m = re.match(r"(/nix/store/\S+\.drv)\b", s) + if m: + built.add(m.group(1)) + elif s and not s.startswith("/nix/store"): + grabbing = False + return built or None + + +def build_dag(closure: dict[str, dict], + to_build: set[str] | None) -> tuple[dict[str, list[str]], set[str]]: + """Return ``(preds, nodes)`` where ``preds[d]`` are the in-graph derivations + ``d`` depends on. Restricted to ``to_build`` when given (edges to already-cached + inputs are dropped — they contribute no build time). + """ + nodes = set(closure) if to_build is None else (set(to_build) & set(closure)) + preds: dict[str, list[str]] = {} + for d in nodes: + ins = closure[d].get("inputDrvs") or {} + preds[d] = [p for p in ins if p in nodes] + return preds, nodes diff --git a/nix_estimator/schedule.py b/nix_estimator/schedule.py new file mode 100644 index 0000000..6f9536b --- /dev/null +++ b/nix_estimator/schedule.py @@ -0,0 +1,139 @@ +"""DAG scheduling analysis: critical path, peak concurrency, and p-machine makespan. + +All functions take: + - ``dur`` : ``{node: minutes}`` + - ``preds`` : ``{node: [dependency nodes]}`` (edges point dependency -> dependent) + +and are pure (no Nix, no I/O) so they unit-test on toy graphs. +""" + +from __future__ import annotations + +import heapq +from collections import defaultdict + + +def _succs(preds: dict[str, list[str]]) -> dict[str, list[str]]: + succ: dict[str, list[str]] = defaultdict(list) + for node, ps in preds.items(): + for p in ps: + succ[p].append(node) + return succ + + +def _topo(dur: dict[str, float], preds: dict[str, list[str]]) -> list[str]: + indeg = {n: len(preds.get(n, ())) for n in dur} + succ = _succs(preds) + q = [n for n, d in indeg.items() if d == 0] + order: list[str] = [] + while q: + n = q.pop() + order.append(n) + for c in succ.get(n, ()): + indeg[c] -= 1 + if indeg[c] == 0: + q.append(c) + if len(order) != len(dur): + raise ValueError("dependency graph has a cycle") + return order + + +def work(dur: dict[str, float]) -> float: + """Total build time on a single core (T1).""" + return float(sum(dur.values())) + + +def critical_path(dur: dict[str, float], + preds: dict[str, list[str]]) -> tuple[float, list[str]]: + """Span (T∞) and the longest dependent chain that realises it.""" + order = _topo(dur, preds) + finish: dict[str, float] = {} + back: dict[str, str | None] = {} + for n in order: + best_p, best_f = None, 0.0 + for p in preds.get(n, ()): + if finish[p] > best_f: + best_f, best_p = finish[p], p + finish[n] = best_f + dur[n] + back[n] = best_p + if not finish: + return 0.0, [] + end = max(finish, key=lambda k: finish[k]) + chain = [] + cur: str | None = end + while cur is not None: + chain.append(cur) + cur = back[cur] + chain.reverse() + return finish[end], chain + + +def peak_concurrency(dur: dict[str, float], preds: dict[str, list[str]]) -> int: + """Max simultaneously-runnable builds in the ASAP (unlimited-machine) + schedule — the hard ceiling on useful node count.""" + order = _topo(dur, preds) + start: dict[str, float] = {} + events: list[tuple[float, int]] = [] + for n in order: + est = max((start[p] + dur[p] for p in preds.get(n, ())), default=0.0) + start[n] = est + events.append((est, +1)) + events.append((est + dur[n], -1)) + events.sort(key=lambda e: (e[0], e[1])) # ends (-1) before starts at a tie + cur = peak = 0 + for _, delta in events: + cur += delta + peak = max(peak, cur) + return peak + + +def makespan(dur: dict[str, float], preds: dict[str, list[str]], machines: int, + priority: dict[str, float] | None = None) -> float: + """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). + """ + if machines < 1: + raise ValueError("machines must be >= 1") + prio = priority or _path_to_sink(dur, preds) + succ = _succs(preds) + indeg = {n: len(preds.get(n, ())) for n in dur} + ready = [n for n, d in indeg.items() if d == 0] + ready.sort(key=lambda n: prio[n]) # pop() takes the last = highest priority + running: list[tuple[float, str]] = [] # (finish_time, node) min-heap + t = 0.0 + free = machines + done = 0 + total = len(dur) + while done < total: + while free > 0 and ready: + n = ready.pop() + heapq.heappush(running, (t + dur[n], n)) + free -= 1 + if not running: + raise ValueError("deadlock — cycle in graph") + ft, n = heapq.heappop(running) + t = ft + free += 1 + done += 1 + newly = [] + for c in succ.get(n, ()): + indeg[c] -= 1 + if indeg[c] == 0: + newly.append(c) + if newly: + ready.extend(newly) + ready.sort(key=lambda n: prio[n]) + return t + + +def _path_to_sink(dur: dict[str, float], + preds: dict[str, list[str]]) -> dict[str, float]: + """Longest weighted path from each node to a sink (critical-path priority).""" + order = _topo(dur, preds) + succ = _succs(preds) + d: dict[str, float] = {} + for n in reversed(order): + d[n] = dur[n] + max((d[c] for c in succ.get(n, ())), default=0.0) + return d diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..4a26281 --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,26 @@ +[build-system] +requires = ["hatchling"] +build-backend = "hatchling.build" + +[project] +name = "nix-estimator" +version = "0.1.0" +description = "Estimate the best remote-builder configuration (nodes × cores) for a Nix build" +readme = "README.md" +requires-python = ">=3.10" +license = "MIT" +authors = [{ name = "Oleks" }] +keywords = ["nix", "build", "scheduling", "critical-path", "remote-builders"] +dependencies = [] # stdlib only — shells out to `nix` + +[project.scripts] +nix-estimate = "nix_estimator.cli:main" + +[project.optional-dependencies] +dev = ["pytest>=7"] + +[tool.hatch.build.targets.wheel] +packages = ["nix_estimator"] + +[tool.pytest.ini_options] +testpaths = ["tests"] diff --git a/tests/test_schedule.py b/tests/test_schedule.py new file mode 100644 index 0000000..e635d9a --- /dev/null +++ b/tests/test_schedule.py @@ -0,0 +1,45 @@ +"""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