From 51edcebb2f7436fa18ca0c3d195e9efc643ab7ab Mon Sep 17 00:00:00 2001 From: Milind Srivastava Date: Mon, 2 Nov 2026 11:55:31 -0500 Subject: [PATCH] devtools: print a stage_pipeline document as text print_stages.py prints each stage's DAG and the memo after each step (groups with children, parents and queries; candidates with origin, reads and pane-window parameters). --enumerated prints every memo choice as a logical DAG. The Stage Viewer shows the same document graphically. Co-Authored-By: Claude Opus 5.5 --- .github/workflows/rust.yml | 2 +- tools/dag-viewer/README.md | 13 +- tools/dag-viewer/print_stages.py | 189 ++++++++++++++++++++++++++ tools/dag-viewer/test_print_stages.py | 77 +++++++++++ 4 files changed, 278 insertions(+), 3 deletions(-) create mode 100644 tools/dag-viewer/print_stages.py create mode 100644 tools/dag-viewer/test_print_stages.py diff --git a/.github/workflows/rust.yml b/.github/workflows/rust.yml index d45d197f5..ccba9e4c9 100644 --- a/.github/workflows/rust.yml +++ b/.github/workflows/rust.yml @@ -103,7 +103,7 @@ jobs: working-directory: tools/dag-viewer run: | python3 -m pip install py-mini-racer==0.6.0 - python3 -m unittest -v test_viewer test_server + python3 -m unittest -v test_viewer test_server test_print_stages - name: Run the DAG viewer's Node tests run: node --test tools/dag-viewer/stages.test.mjs tools/dag-viewer/enumeration.test.mjs tools/dag-viewer/memo.test.mjs diff --git a/tools/dag-viewer/README.md b/tools/dag-viewer/README.md index bb448c31f..fa6e8d656 100644 --- a/tools/dag-viewer/README.md +++ b/tools/dag-viewer/README.md @@ -19,6 +19,13 @@ shows: To run it, see [RUNNING.md](RUNNING.md). +For the same document as text, run `python3 tools/dag-viewer/print_stages.py +run.json`. It prints each stage's DAG and the memo after each step: groups +with their children, parents and queries, and each group's candidates. Add +`--enumerated` to print every memo choice as a logical DAG instead +(`stage_pipeline --max-enumerated N` sets how many choices the document +lists). + ## The page - **Examples**: one tab per entry of `examples.json`: #509 Examples 1, 3a @@ -478,6 +485,8 @@ sections are still shown. If a check fails, it shows the problems instead. `ratio-of-quantiles-5m.json` (one ratio of two equal per-series quantiles plus a third, every minute, ε = 0.01, δ = 0.001) shows the `memo` section: the memo after each step. +- `print_stages.py`: prints a document as text. `test_print_stages.py`: its + tests, over a document planned by the built `stage_pipeline`. - `cytoscape.min.js`, `dagre.min.js`, `cytoscape-dagre.js`: vendored, so the page works offline. @@ -486,7 +495,7 @@ sections are still shown. If a check fails, it shows the problems instead. From `tools/dag-viewer`: ```bash -python3 -m unittest test_viewer test_server +python3 -m unittest test_viewer test_server test_print_stages ``` `stages.js` and `memo.js` also have Node tests, which load the elements into the vendored @@ -498,7 +507,7 @@ node --test tools/dag-viewer/stages.test.mjs tools/dag-viewer/enumeration.test.m The JavaScript runs in V8 through `py_mini_racer` (`pip install py-mini-racer==0.6.0`) against a stub DOM, so no browser is needed; without -it those tests are skipped. The scalar-query and `/api/plan` tests need a +it those tests are skipped. The scalar-query, `/api/plan` and `print_stages.py` tests need a built `stage_pipeline` (`cargo build -p asap-devtools --bin stage_pipeline`). The Rust side: diff --git a/tools/dag-viewer/print_stages.py b/tools/dag-viewer/print_stages.py new file mode 100644 index 000000000..281376b08 --- /dev/null +++ b/tools/dag-viewer/print_stages.py @@ -0,0 +1,189 @@ +"""Print a stage_pipeline document as text: each stage's DAG and the memo after each step. + + cargo run -q -p asap-devtools --bin stage_pipeline -- --workload w.json --out run.json + python3 tools/dag-viewer/print_stages.py run.json # stages and memo steps + python3 tools/dag-viewer/print_stages.py run.json --enumerated # every enumerated memo choice + +`--max-enumerated N` on stage_pipeline raises how many memo choices it lists.""" +import json, sys +from collections import Counter + +d = json.load(open(sys.argv[1])) +CHILD_KEYS = {"child", "children", "inputs", "left", "right", "summary_input"} +DROP = {"schema", "output_names"} + + +def compact(v): + return json.dumps(v, separators=(",", ":")) + + +def op_label(op): + # {"NonASAP": {"Aggregate": {...}}} -> ("NonASAP.Aggregate", params, children) + kind = next(iter(op)); inner = op[kind] + if not isinstance(inner, dict): + return f"{kind}.{inner}", "", [] + name = next(iter(inner)); body = inner[name] + if not isinstance(body, dict): + return f"{kind}.{name}", compact(body), [] + kids = [x for k in CHILD_KEYS & body.keys() for x in (body[k] if isinstance(body[k], list) else [body[k]])] + params = {k: v for k, v in body.items() if k not in CHILD_KEYS | DROP and v not in (None, [], {})} + return f"{kind}.{name}", compact(params), kids + + +def schema_cols(s): + return ",".join(f["name"] for f in (s or {}).get("fields", [])) + + +def print_dag(dag, indent=" "): + print(f"{indent}roots: {compact(dag['roots'])}") + for i, n in enumerate(dag["nodes"]): + op = n.get("operator") or n.get("payload") + name, params, kids = op_label(op) + mat = n.get("materialization") + mat = "" if mat in (None, "NotMaterialized") else f" mat={compact(mat)}" + cols = schema_cols(n.get("schema") or n.get("output_schema")) + print(f"{indent}#{n.get('id', i)} {name} <- {kids} [{cols}]{mat}") + if params != "{}": + print(f"{indent} {params}") + + +# stage_pipeline --no-memo leaves the memo out. +memo = d.get("memo", {"steps": [], "enumeration": None}) + + +def window_params(node): + """PaneChain/SummaryMerge window parameters, which the memo labels omit.""" + name, body = next(iter(node["operator"].get("ASAP", {}).items()), (None, None)) + if name == "PaneChain": + return name, f"[{body['pane_count']}x{body['pane_width_ms'] // 1000}s]" + if name == "SummaryMerge": + # A missing newest_panes means every input pane is merged. + return name, f"[newest {body.get('newest_panes') or 'all'}]" + return None + + +def with_window_params(label, root, stop): + # Own nodes only: stop at nodes that belong to the input groups. + found, stack, seen = [], [root], set() + while stack: + i = stack.pop() + if i in seen or str(i) in stop: + continue + seen.add(i) + params = window_params(memo["nodes"][i]) + if params: + found.append(params) + _, kids = op_label(memo["nodes"][i]["operator"])[1:] + stack.extend(reversed([k for k in kids if isinstance(k, int)])) + pos = 0 + for name, suffix in found: + pos = label.find(name, pos) + if pos < 0: + break + pos += len(name) + label = label[:pos] + suffix + label[pos:] + return label + + +def short_node(i): + """One line per memo-table node: operator plus the parameters that tell plans apart.""" + name, _, kids = op_label(memo["nodes"][i]["operator"]) + op = memo["nodes"][i]["operator"] + body = next(iter(op[next(iter(op))].values())) + extra = [] + if isinstance(body, dict): + if "family" in body: + sketch = body["family"]["Sketch"][0] + extra.append(f"{sketch['algorithm']}{compact(next(iter(sketch['params'].values())))}") + if "range" in body: + extra.append(f"{body['range']['secs']}s") + q = body.get("query", {}).get("Quantile") if isinstance(body.get("query"), dict) else None + if q: + extra.append(f"q={q['q']}") + for m in body.get("measures", []): + extra.append(f"{m['kind']}(q={m.get('q')})") + params = window_params(memo["nodes"][i]) + label = name.split(".")[1] + (params[1] if params else "") + return f"{label} {' '.join(extra)}".strip(), [k for k in kids if isinstance(k, int)] + + +def print_memo_dag(roots): + reachable, stack = set(), list(roots) + while stack: + i = stack.pop() + if i not in reachable: + reachable.add(i) + stack.extend(short_node(i)[1]) + # The memo node table lists children before parents. + for i in sorted(reachable): + text, kids = short_node(i) + root = f" <- ROOT q{roots.index(i)}" if i in roots else "" + print(f" n{i} {text} <- {['n%d' % k for k in kids]}{root}") + + +if "--enumerated" in sys.argv: + e = memo["enumeration"] + if e is None: + sys.exit("no memo enumeration: the document was written with --no-memo") + print(f"Every choice of one candidate per reached group, at memo step {e['step']}: " + f"{len(e['choices'])} listed, {compact(e['per_step'])}") + for c in e["choices"]: + mark = " <== SELECTED" if c["selected"] else "" + choice = " ".join(f"G{g}:c{i}" for g, i in c["choice"].items()) + print(f"\n{c['id']} {choice}{mark}") + print_memo_dag(c["roots"]) if c["roots"] is not None else print(f" error: {c['error']}") + sys.exit() + +print("=" * 80, "\nSTAGE 0: frontends' DAG") +print_dag(d["stage0_query_language_frontends"]["dag"]) + +for step in memo["steps"]: + print("=" * 80, f"\nMEMO after {step['step']}: {step['title']}") + for g in step["groups"]: + roots = f" ROOT of q{g['roots_of']}" if g["roots_of"] else "" + label = with_window_params(g["label"], g["node"], set()) + print(f" G{g['id']} {label} children={g['children']} parents={g['parents']} queries={g['queries']}{roots}") + for c in g["candidates"]: + new = " NEW" if c["added"] and step["step"] != "stage0" else "" + label = with_window_params(c["label"], c["node"], c["group_of_node"]) + print(f" cand {c['index']} [{c['origin']}{new}] {label} reads={c['children']} " + f"implements={c['implements']} budget_parts={c['accuracy_budget_parts']}") + if g["candidates_capped"]: + print(f" ... {g['candidate_count']} total") +if "selected_choice" in memo: + print(f" selected choice (group -> cand): {compact(memo['selected_choice'])}") + +s1 = d["stage1_logical_asap_optimization"] +print("=" * 80, f"\nSTAGE 1: logical candidates ({s1['combinations']} enumerated, capped={s1['capped']})") +for p in s1["pruned_candidates"]: + print(f" pass1 pruned: {compact(p)}") +for c in s1["candidates"]: + print(f" {c.get('id')} {c.get('display_name', '')}") + print_dag(c["dag"]) if "dag" in c else print(f" error: {c.get('error')}") + +s2 = d["stage2_physical_asap_optimization"] +print("=" * 80, "\nSTAGE 2: physical candidates") +for c in s2["candidates"]: + print(f" {c['id']} from {c['from_logical']} {c['display_name']}") + print_dag(c["dag"]) +for p in s2["pruned"]: + print(f" pruned {compact(p)}") +pruned = Counter((p["reason"], p["timing"]) for p in s2["pruned_candidates"]) +for (reason, timing), n in sorted(pruned.items()): + print(f" pruned {n} {timing} candidates: {reason}") + +s3 = d["stage3_plan_selection"] +print("=" * 80, "\nSTAGE 3: plan selection") +if "error" in s3: + sys.exit(f" error: {s3['error']}") +r = s3["search_report"] +print(f" search {r['status']}: {r['milp_variables']} MILP variables, {r['solver_ms']} ms solving") +sel = s3["selected_dag"] +print(f" selected Stage 2 candidates per node: {sel['candidate_of_node']}, outputs: {sel['node_of_output_reference']}") +for e in sel["edges"]: + print(f" node {e['producer']} -> node {e['consumer']} input {e['input']}") +print(f" cost: {s3['cost']['cost_per_second']:.6g} cost/s") +for node, nc in s3["cost"]["per_node"].items(): + print(f" #{node} {nc['operator']}: {nc['cost_per_second']:.6g} {nc['detail']}") +for q, a in enumerate(s3["accepted_accuracy_per_query"]): + print(f" q{q} accuracy: {a['summary']}") diff --git a/tools/dag-viewer/test_print_stages.py b/tools/dag-viewer/test_print_stages.py new file mode 100644 index 000000000..de4600eff --- /dev/null +++ b/tools/dag-viewer/test_print_stages.py @@ -0,0 +1,77 @@ +"""Tests for print_stages.py on a document stage_pipeline writes now. + +Run from tools/dag-viewer: `python3 -m unittest test_print_stages`. Needs a +built stage_pipeline (`cargo build -p asap-devtools --bin stage_pipeline`). +""" + +from __future__ import annotations + +import subprocess +import sys +import tempfile +import unittest +from pathlib import Path + +import server + +HERE = Path(__file__).resolve().parent + + +@unittest.skipUnless(server.BINARY.exists(), "needs a built stage_pipeline") +class PrintStagesTest(unittest.TestCase): + @classmethod + def setUpClass(cls) -> None: + cls.tmp = tempfile.TemporaryDirectory() + cls.doc = Path(cls.tmp.name) / "run.json" + subprocess.run( + [ + str(server.BINARY), + "--promql", "quantile_over_time(0.5, data[5m])", + "--promql", "quantile_over_time(0.9, data[5m])", + "--series-count", "100", "--sample-interval-ms", "15000", + "--max-error", "0.01", "--failure-probability", "0.001", + "--milp-time-limit-s", "5", + "--out", str(cls.doc), + ], + check=True, + capture_output=True, + ) + + @classmethod + def tearDownClass(cls) -> None: + cls.tmp.cleanup() + + def run_printer(self, *flags: str) -> str: + return subprocess.run( + [sys.executable, str(HERE / "print_stages.py"), str(self.doc), *flags], + check=True, + capture_output=True, + text=True, + ).stdout + + def test_prints_every_stage_and_memo_step(self) -> None: + """Each stage and each memo step gets a section; the two queries share one group.""" + out = self.run_printer() + for heading in [ + "STAGE 0: frontends' DAG", + "MEMO after stage0", + "MEMO after pass1", + "MEMO after pass2", + "MEMO after concrete_states", + "STAGE 1: logical candidates", + "STAGE 2: physical candidates", + "STAGE 3: plan selection", + ]: + self.assertIn(heading, out) + self.assertIn("queries=[0, 1]", out) + self.assertIn("cost/s", out) + + def test_enumerated_lists_memo_choices(self) -> None: + """--enumerated prints one DAG per memo choice, each with its query roots.""" + out = self.run_printer("--enumerated") + self.assertIn("E1 G0:c0", out) + self.assertIn("<- ROOT q1", out) + + +if __name__ == "__main__": + unittest.main()