nix-estimator v0.1: DAG → node×core builder-config estimator
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.
This commit is contained in:
Oleks
+10
@@ -0,0 +1,10 @@
|
||||
__pycache__/
|
||||
*.pyc
|
||||
.pytest_cache/
|
||||
*.egg-info/
|
||||
dist/
|
||||
build/
|
||||
result
|
||||
result-*
|
||||
.direnv/
|
||||
.venv/
|
||||
@@ -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 <json>`** (`{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.
|
||||
@@ -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
|
||||
];
|
||||
};
|
||||
});
|
||||
};
|
||||
}
|
||||
@@ -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"
|
||||
@@ -0,0 +1,111 @@
|
||||
"""Command-line entrypoint: ``nix-estimate <flake-attr>``."""
|
||||
|
||||
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())
|
||||
@@ -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/<hash>-<name>.drv` -> `<name>` (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
|
||||
@@ -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."
|
||||
),
|
||||
}
|
||||
@@ -0,0 +1,78 @@
|
||||
"""Extract the to-build derivation DAG for a flake attr, using Nix itself.
|
||||
|
||||
``nix derivation show -r <attr>`` emits JSON for the *whole* closure:
|
||||
``{ drvPath: { name, outputs, inputDrvs, env, ... } }``.
|
||||
|
||||
``nix build --dry-run <attr>`` 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
|
||||
@@ -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
|
||||
@@ -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"]
|
||||
@@ -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
|
||||
Reference in New Issue
Block a user