Oleks
38aaa76a11
- --max-jobs / --node-ram-gb pass through to estimate(); shown in report header
and JSON.
- --timeline recomputes the recommended shape's schedule and renders a compact
per-node text Gantt (also emitted under 'timeline' in --json).
- --provision [CLASS] (+ --provision-name) renders the EphemeralBuilder YAML to
stdout while the human/JSON report goes to stderr, so it pipes to kubectl.
- new 'mine' subcommand: nix-estimate mine <installables...> [--repeat N] [-o]
runs mine_history, merges repeats via merge_histories, writes {name: minutes}.
Default estimate path preserved: 'mine' is dispatched on the leading token,
so 'nix-estimate <attr> [flags]' is unchanged.
327 lines
12 KiB
Python
327 lines
12 KiB
Python
"""Command-line entrypoint: ``nix-estimate <flake-attr>`` (+ the ``mine`` subcommand)."""
|
||
|
||
from __future__ import annotations
|
||
|
||
import argparse
|
||
import json
|
||
import sys
|
||
|
||
from . import __version__, costmodel, estimate as estimate_mod, graph, miner, provision
|
||
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, cold: bool = False) -> str:
|
||
L = []
|
||
L.append(f"nix-estimator {__version__} — {attr}")
|
||
L.append("=" * 64)
|
||
L.append(f"derivations in closure : {est.nodes_evaluated}")
|
||
label = "to build (COLD; whole closure)" if cold else "to build (cache-miss)"
|
||
L.append(f"{label:<22} : {est.to_build}")
|
||
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)")
|
||
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(f"critical path: {len(est.longest_chain)} derivations (heaviest on it):")
|
||
for name, mins in sorted(est.longest_chain, key=lambda x: -x[1])[:6]:
|
||
L.append(f" {mins:6.1f} min {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 _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")
|
||
ap.add_argument(
|
||
"--nix-arg",
|
||
action="append",
|
||
default=[],
|
||
help="extra arg passed through to nix (repeatable)",
|
||
)
|
||
ap.add_argument(
|
||
"--cold",
|
||
"--full",
|
||
dest="cold",
|
||
action="store_true",
|
||
help="estimate a from-scratch build: ignore the local/binary "
|
||
"cache and cost the whole closure (upper bound)",
|
||
)
|
||
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(
|
||
"--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")
|
||
|
||
|
||
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`.
|
||
vals = set()
|
||
for x in raw.split(","):
|
||
x = x.strip()
|
||
try:
|
||
v = int(x)
|
||
except ValueError:
|
||
ap.error(f"--{label}: {x!r} is not an integer")
|
||
if v <= 0:
|
||
ap.error(f"--{label}: {v} is not a positive integer")
|
||
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")
|
||
|
||
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
|
||
if args.cold:
|
||
to_build = None # cost the whole closure regardless of cache state
|
||
else:
|
||
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)
|
||
|
||
history = _load_history(args.history)
|
||
est = estimate(
|
||
closure,
|
||
preds,
|
||
nodes,
|
||
history=history,
|
||
node_grid=node_grid,
|
||
core_grid=core_grid,
|
||
max_jobs=args.max_jobs,
|
||
node_ram_gb=args.node_ram_gb,
|
||
)
|
||
|
||
# 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:
|
||
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())
|