""" Dynamic progress tracker for the quality-check runner. Maintains in-memory per-task and per-template state, atomically rewrites PROGRESS.md on every update and appends to a per-task CSV. PROGRESS.md is designed to be safe to open at any time; it lists running status, ETA, a worst-first per-template table, and any templates that were skipped because not enough matching data files exist in the data pool. """ from __future__ import annotations import csv import json import os import shutil import time from collections import defaultdict from datetime import datetime from pathlib import Path from typing import Any, Dict, List, Optional _TASK_CSV_FIELDS = [ "chart_name", "input", "ok", "elapsed_s", "final_svg", "final_svg_bytes", "chart_svg", "chart_svg_fallback_png", "n_shapes", "n_text", "err", ] def _atomic_write(path: Path, content: str) -> None: tmp = path.with_suffix(path.suffix + ".tmp") tmp.write_text(content, encoding="utf-8") tmp.replace(path) class ProgressTracker: """Track quality-check progress and keep PROGRESS.md / _tasks.csv current.""" def __init__( self, plan: Dict[str, Any], output_dir: Path, match_csv_path: Optional[Path] = None, min_data_for_plan: Optional[int] = None, ): self.plan = plan self.output_dir = Path(output_dir) self.output_dir.mkdir(parents=True, exist_ok=True) self.progress_md = self.output_dir / "PROGRESS.md" self.tasks_csv = self.output_dir / "_tasks.csv" self.summary_csv = self.output_dir / "_summary.csv" # Build per-template plan map: chart_name -> target (number of jobs). self.target_per_tpl: Dict[str, int] = {} self.chart_type_per_tpl: Dict[str, str] = {} self.engine_per_tpl: Dict[str, str] = {} for t in plan["templates"]: self.target_per_tpl[t["chart_name"]] = len(t["picked_data_files"]) self.chart_type_per_tpl[t["chart_name"]] = t["chart_type"] self.engine_per_tpl[t["chart_name"]] = t["engine"] self.total_tasks = sum(self.target_per_tpl.values()) self.n_templates = len(self.target_per_tpl) # Per-task records accumulated in memory (also written incrementally). self.records: List[Dict[str, Any]] = [] # Per-template rollups (computed on flush, but cached for show). self._rollups_cache: Optional[List[Dict[str, Any]]] = None self.t_start = time.time() self.last_flush = 0.0 self._csv_fh = None self._csv_writer = None self._completed_keys: set = set() self._is_closed = False # Optional context: skipped templates (filtered by min-data threshold). self.skipped: List[Dict[str, Any]] = [] if match_csv_path and min_data_for_plan is not None and match_csv_path.exists(): self._load_skipped(match_csv_path, min_data_for_plan) def _load_skipped(self, match_csv_path: Path, min_data: int) -> None: engine = self.plan.get("engine", "") included = set(self.target_per_tpl.keys()) with open(match_csv_path, "r") as fh: for r in csv.DictReader(fh): if engine and r["engine"] != engine: continue if r["chart_name"] in included: continue n = int(r["num_compatible_data"]) if n < min_data: self.skipped.append({ "engine": r["engine"], "chart_type": r["chart_type"], "chart_name": r["chart_name"], "num_compatible_data": n, }) self.skipped.sort(key=lambda r: (r["num_compatible_data"], r["chart_name"])) def open_csv(self, resume: bool = False) -> None: """Open the tasks CSV. If resume=True and the file exists, load previous rows so we can skip them. Otherwise truncate.""" if resume and self.tasks_csv.exists(): with open(self.tasks_csv, "r") as fh: reader = csv.DictReader(fh) for row in reader: self.records.append(row) self._completed_keys.add((row["chart_name"], row["input"])) mode = "a" write_header = False else: if self.tasks_csv.exists(): self.tasks_csv.unlink() mode = "w" write_header = True self._csv_fh = open(self.tasks_csv, mode, newline="", encoding="utf-8") self._csv_writer = csv.DictWriter(self._csv_fh, fieldnames=_TASK_CSV_FIELDS) if write_header: self._csv_writer.writeheader() self._csv_fh.flush() def already_done(self, chart_name: str, input_basename: str) -> bool: return (chart_name, input_basename) in self._completed_keys def add(self, result: Dict[str, Any]) -> None: """Append a result and update on-disk state.""" self.records.append(result) self._completed_keys.add((result["chart_name"], result["input"])) if self._csv_writer: # CSV reader will read empty strings as is; coerce non-strings. self._csv_writer.writerow({k: result.get(k, "") for k in _TASK_CSV_FIELDS}) self._csv_fh.flush() self._rollups_cache = None # Throttle PROGRESS.md updates to once every ~2s to avoid IO churn. now = time.time() if now - self.last_flush >= 2.0 or len(self.records) == self.total_tasks: self.flush() self.last_flush = now def _rollups(self) -> List[Dict[str, Any]]: if self._rollups_cache is not None: return self._rollups_cache by_tpl: Dict[str, List[Dict[str, Any]]] = defaultdict(list) for r in self.records: by_tpl[r["chart_name"]].append(r) rows = [] for chart_name, target in self.target_per_tpl.items(): rs = by_tpl.get(chart_name, []) done = len(rs) n_ok = sum(1 for r in rs if str(r.get("ok")) == "True") n_fail = done - n_ok n_fb = sum(1 for r in rs if str(r.get("chart_svg_fallback_png")) == "True") n_empty = sum( 1 for r in rs if str(r.get("ok")) == "True" and (int(r.get("n_shapes") or 0) + int(r.get("n_text") or 0)) < 8 ) mean_size = ( sum(int(r.get("final_svg_bytes") or 0) for r in rs if str(r.get("ok")) == "True") / n_ok if n_ok else 0 ) mean_shapes = ( sum(int(r.get("n_shapes") or 0) for r in rs if str(r.get("ok")) == "True") / n_ok if n_ok else 0 ) mean_t = ( sum(float(r.get("elapsed_s") or 0) for r in rs) / done if done else 0 ) if done == 0: status = "pending" elif done < target: status = "running" elif n_fail > 0: status = "done (failures)" elif n_fb > 0: status = "done (fallback)" elif n_empty > 0: status = "done (warn)" else: status = "done" rows.append({ "chart_name": chart_name, "chart_type": self.chart_type_per_tpl[chart_name], "engine": self.engine_per_tpl[chart_name], "target": target, "done": done, "n_success": n_ok, "n_fail": n_fail, "n_fallback_png": n_fb, "n_empty": n_empty, "mean_shapes": round(mean_shapes, 1), "mean_size_kb": round(mean_size / 1024, 1), "mean_elapsed_s": round(mean_t, 1), "status": status, }) # Worst-first sort: failures > fallback > empty > slow > everything else. def _sortkey(r): return ( -r["n_fail"], -r["n_fallback_png"], -r["n_empty"], -r["mean_elapsed_s"], r["chart_name"], ) rows.sort(key=_sortkey) self._rollups_cache = rows return rows def write_summary_csv(self) -> None: rows = self._rollups() if not rows: return # Write to a tmp file then atomic-rename so a concurrent reader # (e.g. build_quality_preview run while the driver is still going) # never sees a half-written file. tmp = self.summary_csv.with_suffix(self.summary_csv.suffix + ".tmp") with open(tmp, "w", newline="", encoding="utf-8") as fh: w = csv.DictWriter(fh, fieldnames=list(rows[0].keys())) w.writeheader() for r in rows: w.writerow(r) tmp.replace(self.summary_csv) def flush(self) -> None: rows = self._rollups() self.write_summary_csv() md = self._render_md(rows) _atomic_write(self.progress_md, md) def close(self) -> None: self._is_closed = True self.flush() if self._csv_fh: self._csv_fh.close() self._csv_fh = None self._csv_writer = None # ---------------------------------------------------------------- markdown def _render_md(self, rows: List[Dict[str, Any]]) -> str: done = len(self.records) total = self.total_tasks ok = sum(1 for r in self.records if str(r.get("ok")) == "True") fail = done - ok fb = sum(1 for r in self.records if str(r.get("chart_svg_fallback_png")) == "True") empty = sum( 1 for r in self.records if str(r.get("ok")) == "True" and (int(r.get("n_shapes") or 0) + int(r.get("n_text") or 0)) < 8 ) elapsed = time.time() - self.t_start avg_t = ( sum(float(r.get("elapsed_s") or 0) for r in self.records) / done if done else 0 ) # If the tracker was rebuilt from an existing CSV the wall-clock # elapsed is bogus (it's just the rebuild time); detect that and # fall back to a single-thread approximation so the reported # throughput / ETA don't go to infinity. if elapsed < max(avg_t, 1.0): elapsed_display = "n/a (rebuilt from CSV)" rate = 0.0 eta_s = 0 else: elapsed_display = f"{elapsed/60:.1f} min" rate = done / max(elapsed, 1e-6) eta_s = (total - done) / max(rate, 1e-6) if done < total else 0 if done >= total: status = "finished" elif self._is_closed: status = "stopped" else: status = "running" lines: List[str] = [] lines.append("# Chart Template Quality Check Progress") lines.append("") lines.append( f"**Last updated**: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')} " f"  ·   **Status**: `{status}`" ) lines.append("") eta_display = "-" if rate == 0 else f"{eta_s/60:.1f} min" lines.append( f"**Overall**: {done}/{total} tasks " f"({(done/total*100 if total else 0):.1f}%) · " f"OK {ok} · FAIL **{fail}** · " f"fallback_png **{fb}** · empty **{empty}** · " f"elapsed {elapsed_display} · " f"avg-task {avg_t:.1f}s · " f"ETA {eta_display}" ) lines.append("") # ---- summary stats by engine ---- lines.append("## Summary") lines.append("") lines.append("| metric | value |") lines.append("|---|---|") lines.append(f"| plan templates | {self.n_templates} |") lines.append(f"| total tasks | {total} |") lines.append(f"| completed | {done} ({(done/total*100 if total else 0):.1f}%) |") lines.append(f"| succeeded | {ok} |") lines.append(f"| failed | **{fail}** |") lines.append(f"| fallback_png | **{fb}** |") lines.append(f"| empty (shapes+text<8) | **{empty}** |") lines.append(f"| elapsed | {elapsed_display} |") throughput_display = "-" if rate == 0 else f"{rate*60:.1f} tasks/min" lines.append(f"| throughput | {throughput_display} |") lines.append(f"| avg task time | {avg_t:.1f} s |") lines.append("") # ---- per-template table ---- lines.append("## Per-template progress (worst-first)") lines.append("") lines.append( "| chart_name | chart_type | done/target | ok | fail | fb_png | empty | " "mean_shapes | mean_size | mean_t (s) | status |" ) lines.append( "|---|---|---|---|---|---|---|---|---|---|---|" ) for r in rows: status_md = r["status"] if r["n_fail"] or r["n_fallback_png"]: status_md = f"**{status_md}**" lines.append( f"| `{r['chart_name']}` " f"| {r['chart_type']} " f"| {r['done']}/{r['target']} " f"| {r['n_success']} " f"| {r['n_fail']} " f"| {r['n_fallback_png']} " f"| {r['n_empty']} " f"| {r['mean_shapes']} " f"| {r['mean_size_kb']} KB " f"| {r['mean_elapsed_s']} " f"| {status_md} |" ) lines.append("") # ---- skipped templates ---- if self.skipped: lines.append(f"## Templates skipped from this run ({len(self.skipped)} total)") lines.append("") lines.append( "These templates exist in the registry but didn't have enough matching " "data files in the pool to be included in this run. Their `chart_name` " "is shown along with how many data files in the pool matched their " "`requirements`. To cover them, expand the data pool (e.g. adapt more " "files from `/data/lizhen/resources/converted/`) and re-run " "`scripts/match_templates_to_data.py`." ) lines.append("") lines.append("| chart_type | chart_name | matched data files |") lines.append("|---|---|---|") for r in self.skipped: lines.append( f"| {r['chart_type']} | `{r['chart_name']}` " f"| {r['num_compatible_data']} |" ) lines.append("") return "\n".join(lines) + "\n"