oleks/nix-estimator
1
0
Fork 0
Code Issues 6 Pull Requests Actions Packages Projects Releases Wiki Activity

Merge wave2B branch (CLI integration #12-#16)

Browse Source
This commit is contained in:
Oleks
2026-07-08 17:43:17 +03:00
parent dbf72e45bc 5f1a75fb12
commit cc2eabe321
5 changed files with 353 additions and 33 deletions
Download Patch File Download Diff File Expand all files Collapse all files
README.md
+39 -5
View File
@@ -58,8 +58,37 @@ nix-estimate .#foo --nodes 1,2,4,8 --cores 8,16,32 --history builds.json
# ignore the cache — estimate a from-scratch build of the whole closure
nix-estimate .#foo --cold
# model Nix per-node concurrency + a RAM budget (issues #13/#15)
nix-estimate .#foo --max-jobs 4 --node-ram-gb 16
# print a per-node Gantt of the recommended shape's schedule (issue #16)
nix-estimate .#foo --timeline
# emit an EphemeralBuilder CR for the recommendation (issue #14) —
# YAML goes to stdout, the human report to stderr, so this pipes cleanly
nix-estimate .#foo --provision | kubectl apply -f -
nix-estimate .#foo --provision runner --provision-name my-builder
```
### Mining a history file (issue #12)
`--history` is only as good as the timings you feed it. The `mine` subcommand
runs real builds and extracts per-derivation minutes from Nix's
`--log-format internal-json` stream:
```sh
# build once, print {name: minutes} JSON to stdout
nix-estimate mine .#foo
# build 3× and merge (median, to shrug off noise) into a file
nix-estimate mine .#foo .#bar --repeat 3 -o builds.json
nix-estimate .#foo --history builds.json
```
The plain `nix-estimate <attr> [flags]` invocation is unchanged — `mine` is the
only subcommand and is selected by the leading literal token.
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).
@@ -79,7 +108,9 @@ nix_estimator/
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
miner.py mine {name: minutes} from internal-json build logs (pure parser)
provision.py render the recommendation as an EphemeralBuilder CR (YAML)
cli.py `nix-estimate` entrypoint + report + `mine` subcommand
tests/ scheduler unit tests on toy graphs
```
@@ -108,9 +139,12 @@ ingredients exist separately; nix-estimator assembles them:
## 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.
v0.1 — heuristic cost model, working scheduler + recommendation. Now also:
history mining from internal-json build logs (`mine` subcommand, issue #12),
`--max-jobs` per-node concurrency and `--node-ram-gb` as a second constraint
(issues #13/#15), a `--timeline` per-node Gantt (issue #16), and a
`--provision` mode that emits an `EphemeralBuilder` spec for the chosen shape
(issue #14). Roadmap: a learned cost model, and calibrating the estimates
against mined makespans.
MIT.
nix_estimator/cli.py
+187 -27
View File
@@ -1,4 +1,4 @@
"""Command-line entrypoint: ``nix-estimate <flake-attr>``."""
"""Command-line entrypoint: ``nix-estimate <flake-attr>`` (+ the ``mine`` subcommand)."""
from __future__ import annotations
@@ -6,7 +6,7 @@ import argparse
import json
import sys
from . import __version__, graph
from . import __version__, costmodel, estimate as estimate_mod, graph, miner, provision
from .estimate import estimate
@@ -27,6 +27,11 @@ def _report(attr: str, est, cold: bool = False) -> str:
if est.to_build == 0:
L.append("\nNothing to build — the whole closure substitutes from cache.")
return "\n".join(L)
if est.max_jobs != 1 or est.node_ram_gb is not None:
knobs = [f"max-jobs/node = {est.max_jobs}"]
if est.node_ram_gb is not None:
knobs.append(f"node RAM budget = {est.node_ram_gb:g} GB")
L.append("per-node config : " + ", ".join(knobs))
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)")
@@ -61,12 +66,31 @@ def _report(attr: str, est, cold: bool = False) -> str:
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.",
)
def _gantt(sched: dict[int, list], makespan_min: float, max_jobs: int,
width: int = 48) -> str:
"""Compact per-lane text Gantt for a computed schedule (issue #16).
One row per lane; lanes are grouped ``max_jobs`` to a physical node. A busy
slice is ``█``, idle is ``·``; the tasks on each lane are listed in dispatch
order after the bar.
"""
L = ["", f"timeline — recommended shape (makespan {makespan_min:.1f} min):"]
scale = (width / makespan_min) if makespan_min > 0 else 0.0
for lane in sorted(sched):
bar = ["·"] * width
for _task, start, finish in sched[lane]:
a = min(width - 1, int(start * scale))
b = min(width, max(a + 1, int(round(finish * scale))))
for i in range(a, b):
bar[i] = "█"
names = ", ".join(costmodel.store_name(t) for t, _, _ in sched[lane]) or "idle"
node, job = lane // max_jobs, lane % max_jobs
tag = f"node{node}.j{job}" if max_jobs > 1 else f"node{node}"
L.append(f" {tag:>10} │{''.join(bar)}│ {names}")
return "\n".join(L)
def _add_estimate_args(ap: argparse.ArgumentParser) -> None:
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")
@@ -92,10 +116,43 @@ def main(argv: list[str] | None = None) -> int:
ap.add_argument(
"--cores", default="8,16,32", help="comma list of cores-per-node to sweep"
)
ap.add_argument(
"--max-jobs",
type=int,
default=1,
help="Nix per-node concurrent builds (nix.settings.max-jobs); "
"co-resident jobs share the node's cores (issue #13)",
)
ap.add_argument(
"--node-ram-gb",
type=float,
default=None,
help="per-node RAM budget in GB; caps co-resident jobs to what fits "
"(issue #15). Unset = unconstrained.",
)
ap.add_argument(
"--timeline",
action="store_true",
help="render a per-node Gantt of the recommended shape's schedule",
)
ap.add_argument(
"--provision",
nargs="?",
const="nix-builder",
default=None,
metavar="CLASS",
help="emit an EphemeralBuilder YAML for the recommendation to stdout "
"(builder class, default 'nix-builder'); the report goes to stderr",
)
ap.add_argument(
"--provision-name",
default=None,
help="metadata.name for the EphemeralBuilder (default derived from shape)",
)
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)
def _run_estimate(args, ap: argparse.ArgumentParser) -> int:
def _parse_grid(raw: str, label: str) -> tuple[int, ...]:
# Sort + dedupe ascending: _recommend assumes ordered grids, so
# `--cores 32,8,16` must mean the same as `--cores 8,16,32`.
@@ -111,6 +168,11 @@ def main(argv: list[str] | None = None) -> int:
vals.add(v)
return tuple(sorted(vals))
if args.max_jobs < 1:
ap.error("--max-jobs must be >= 1")
if args.node_ram_gb is not None and args.node_ram_gb <= 0:
ap.error("--node-ram-gb must be > 0")
node_grid = _parse_grid(args.nodes, "nodes")
core_grid = _parse_grid(args.cores, "cores")
@@ -132,35 +194,133 @@ def main(argv: list[str] | None = None) -> int:
)
preds, nodes = graph.build_dag(closure, to_build)
history = _load_history(args.history)
est = estimate(
closure,
preds,
nodes,
history=_load_history(args.history),
history=history,
node_grid=node_grid,
core_grid=core_grid,
max_jobs=args.max_jobs,
node_ram_gb=args.node_ram_gb,
)
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,
)
# Recompute the recommended shape's concrete schedule for --timeline / --json.
rec = est.recommendation
sched = None
makespan_min = 0.0
if est.to_build and (args.timeline or args.json):
makespan_min, sched = estimate_mod.schedule_for(
closure, preds, nodes,
cores=rec["cores_per_node"], machines=rec["nodes"],
history=history, max_jobs=args.max_jobs, node_ram_gb=args.node_ram_gb,
)
if args.json:
payload = {
"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,
"max_jobs": est.max_jobs,
"node_ram_gb": est.node_ram_gb,
"grid": {f"{p}x{c}": v for (p, c), v in est.grid.items()},
"recommendation": est.recommendation,
}
if args.timeline and sched is not None:
payload["timeline"] = {
"makespan_min": makespan_min,
"assignments": {
str(lane): [
[costmodel.store_name(t), s, f] for t, s, f in items
]
for lane, items in sched.items()
},
}
primary = json.dumps(payload, indent=2)
else:
print(_report(args.attr, est, cold=args.cold))
primary = _report(args.attr, est, cold=args.cold)
if args.timeline and sched is not None:
primary += "\n" + _gantt(sched, makespan_min, args.max_jobs)
# --provision: YAML is the machine artifact on stdout; the human/JSON report
# is context, pushed to stderr so `nix-estimate ... --provision | kubectl
# apply -f -` works cleanly.
if args.provision is not None:
yaml = provision.render_ephemeral_builder(
nodes=rec["nodes"],
cores=rec["cores_per_node"],
system=args.system,
est_makespan_min=rec.get("est_makespan_min"),
builder_class=args.provision,
name=args.provision_name,
)
sys.stdout.write(yaml)
print(primary, file=sys.stderr)
else:
print(primary)
return 0
def _run_mine(args) -> int:
histories = [
miner.mine_history(args.installables, extra_args=args.nix_arg)
for _ in range(args.repeat)
]
merged = miner.merge_histories(histories)
text = json.dumps(merged, indent=2, sort_keys=True)
if args.output:
with open(args.output, "w") as fh:
fh.write(text + "\n")
print(f"wrote {len(merged)} timings to {args.output}", file=sys.stderr)
else:
print(text)
return 0
def main(argv: list[str] | None = None) -> int:
argv = list(sys.argv[1:] if argv is None else argv)
# `mine` is the only subcommand; anything else is the default estimate path,
# so `nix-estimate <attr> [flags]` keeps working unchanged. We dispatch on a
# leading literal "mine" token (a flake attr never equals "mine").
if argv and argv[0] == "mine":
ap = argparse.ArgumentParser(
prog="nix-estimate mine",
description="Mine {name: minutes} build timings from real nix builds.",
)
ap.add_argument("installables", nargs="+", help="flake attrs / drvs to build")
ap.add_argument(
"--repeat", type=int, default=1,
help="build N times and merge (median) to smooth noise",
)
ap.add_argument(
"--nix-arg", action="append", default=[],
help="extra arg passed through to nix (repeatable)",
)
ap.add_argument(
"-o", "--output", help="write JSON here (default: stdout)"
)
args = ap.parse_args(argv[1:])
if args.repeat < 1:
ap.error("--repeat must be >= 1")
return _run_mine(args)
ap = argparse.ArgumentParser(
prog="nix-estimate",
description="Estimate the best builder configuration (nodes × cores) "
"for a Nix flake attr or derivation.",
epilog="subcommand: `nix-estimate mine <installables...>` mines a "
"{name: minutes} history file from real builds for --history.",
)
_add_estimate_args(ap)
ap.add_argument("--version", action="version", version=__version__)
args = ap.parse_args(argv)
return _run_estimate(args, ap)
if __name__ == "__main__":
raise SystemExit(main())
nix_estimator/costmodel.py
+1 -1
View File
@@ -205,7 +205,7 @@ def cost(
return DEFAULT
def scale_to_cores(minutes8: float, core_scaling: float, cores: int) -> float:
def scale_to_cores(minutes8: float, core_scaling: float, cores: float) -> 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.
nix_estimator/estimate.py
+29
View File
@@ -106,6 +106,35 @@ def estimate(
)
def schedule_for(
closure: dict[str, dict],
preds: dict[str, list[str]],
nodes: set[str],
*,
cores: int,
machines: int,
history: dict[str, float] | None = None,
max_jobs: int = 1,
node_ram_gb: float | None = None,
):
"""Concrete (makespan, per-lane assignments) for one (machines × cores) shape.
Reuses the same per-derivation durations the grid was built from, so the
makespan here matches ``grid[(machines, cores)]``. Used by ``--timeline`` to
render a Gantt of the recommended shape (issue #16).
"""
dur = _durations(nodes, closure, history, cores, max_jobs)
gb_per_job = (
{d: costmodel.ram_gb(d, closure.get(d)) for d in nodes}
if node_ram_gb is not None
else None
)
return schedule.makespan(
dur, preds, machines, max_jobs=max_jobs,
gb_per_job=gb_per_job, node_ram_gb=node_ram_gb, return_schedule=True,
)
def _recommend(grid, peak, core_grid, node_grid, knee, span_dominator_frac=0.0):
"""Pick a (nodes, cores) at the diminishing-returns knee.
tests/test_cli.py
+97
View File
@@ -77,3 +77,100 @@ def test_non_integer_cores_rejected(stub_graph, capsys):
cli.main([".#x", "--history", stub_graph, "--cores", "8,abc", "--json"])
assert exc.value.code != 0
assert "integer" in capsys.readouterr().err
# --- new feature surfaces (issues #12–#16) ---------------------------------
def test_backcompat_plain_estimate(stub_graph, capsys):
"""Plain `nix-estimate <attr> --json` still parses with no subcommand."""
out = _run_json(capsys, [".#foo", "--history", stub_graph, "--json"])
assert out["attr"] == ".#foo"
assert out["max_jobs"] == 1
assert out["node_ram_gb"] is None
def test_max_jobs_accepted_and_affects_output(stub_graph, capsys):
base = _run_json(capsys, [".#x", "--history", stub_graph, "--json"])
wide = _run_json(
capsys, [".#x", "--history", stub_graph, "--max-jobs", "4", "--json"]
)
assert base["max_jobs"] == 1 and wide["max_jobs"] == 4
# Sharing cores across 4 jobs changes at least one grid makespan.
assert wide["grid"] != base["grid"]
def test_node_ram_gb_accepted(stub_graph, capsys):
out = _run_json(
capsys,
[".#x", "--history", stub_graph, "--max-jobs", "4",
"--node-ram-gb", "8", "--json"],
)
assert out["node_ram_gb"] == 8.0
def test_timeline_renders_report(stub_graph, capsys):
assert cli.main([".#x", "--history", stub_graph, "--timeline"]) == 0
out = capsys.readouterr().out
assert "timeline — recommended shape" in out
assert "node0" in out
def test_timeline_json_carries_assignments(stub_graph, capsys):
out = _run_json(
capsys, [".#x", "--history", stub_graph, "--timeline", "--json"]
)
assert "timeline" in out
tl = out["timeline"]
assert tl["makespan_min"] > 0
# every scheduled task name appears in some lane
tasks = [t[0] for lane in tl["assignments"].values() for t in lane]
assert "sink" in tasks
def test_provision_emits_yaml_to_stdout(stub_graph, capsys):
assert cli.main([".#x", "--history", stub_graph, "--provision"]) == 0
cap = capsys.readouterr()
assert "kind: EphemeralBuilder" in cap.out
assert "apiVersion: builder-arbitrage.oleks.space/v1" in cap.out
assert "class: nix-builder" in cap.out
assert "replicas:" in cap.out and "cores:" in cap.out
# human report is diverted to stderr, not stdout
assert "RECOMMENDATION" in cap.err
assert "RECOMMENDATION" not in cap.out
def test_provision_custom_class_and_name(stub_graph, capsys):
assert cli.main(
[".#x", "--history", stub_graph,
"--provision", "runner", "--provision-name", "my-builder"]
) == 0
out = capsys.readouterr().out
assert "class: runner" in out
assert "name: my-builder" in out
def test_mine_subcommand_writes_json(monkeypatch, capsys, tmp_path):
"""`mine` runs mine_history (stubbed) and writes {name: minutes} JSON."""
calls = []
def fake_mine_history(installables, **kw):
calls.append(list(installables))
return {"hello": 2.0, "world": 5.0}
monkeypatch.setattr(cli.miner, "mine_history", fake_mine_history)
out_file = tmp_path / "hist.json"
rc = cli.main(["mine", ".#hello", ".#world", "--repeat", "3", "-o", str(out_file)])
assert rc == 0
assert len(calls) == 3 # repeated
data = json.loads(out_file.read_text())
assert data == {"hello": 2.0, "world": 5.0}
def test_mine_subcommand_to_stdout(monkeypatch, capsys):
monkeypatch.setattr(
cli.miner, "mine_history", lambda inst, **kw: {"pkg": 1.5}
)
assert cli.main(["mine", ".#pkg"]) == 0
assert json.loads(capsys.readouterr().out) == {"pkg": 1.5}