wafer-core 0.1.37__py3-none-any.whl → 0.1.39__py3-none-any.whl

wafer_core/rollouts/progress_app.py

@@ -0,0 +1,395 @@
+"""Eval progress display built on pytui's Elm-style App.
+
+Tails events.jsonl (written by evaluate() via logging), derives display state,
+renders progress in alternate screen via a subprocess.
+
+Usage:
+    with progress_display(output_dir=output_dir):
+        await evaluate(dataset, config)
+
+Or detached mode (two terminals):
+    # Terminal 1: Run eval (writes events.jsonl via logging)
+    python run_eval.py --config config.py
+
+    # Terminal 2: Watch progress
+    python -m wafer_core.rollouts.progress_app /path/to/events.jsonl
+"""
+
+from __future__ import annotations
+
+import json
+import os
+import signal
+import subprocess
+import sys
+import time
+from collections.abc import Generator
+from contextlib import contextmanager
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Any
+
+# ── ANSI colors ──────────────────────────────────────────────────────────────
+
+DIM = "\x1b[90m"
+RESET = "\x1b[0m"
+GREEN = "\x1b[32m"
+RED = "\x1b[31m"
+YELLOW = "\x1b[33m"
+CYAN = "\x1b[36m"
+BOLD = "\x1b[1m"
+
+# Event types the progress display understands.
+# The "message" field in the JSONL record is the event type discriminator.
+# Records with unrecognised message values are silently skipped.
+KNOWN_EVENT_TYPES = frozenset({
+    "eval_start",
+    "eval_end",
+    "sample_start",
+    "sample_end",
+    "turn",
+    "modal_progress",
+    "gepa_iteration",
+    "gepa_accepted",
+    "gepa_rejected",
+    "sample_retry",
+})
+
+
+# ── State ────────────────────────────────────────────────────────────────────
+
+
+@dataclass
+class SampleState:
+    """State of a single sample."""
+
+    id: str
+    name: str = ""
+    turn: int = 0
+    phase: str = ""
+    score: float | None = None
+    status: str = "started"  # started, complete, retry
+    retry_attempt: int = 0
+    start_time: float = field(default_factory=time.time)
+    last_update: float = field(default_factory=time.time)
+
+
+@dataclass
+class RenderState:
+    """Display state derived from events."""
+
+    eval_name: str = ""
+    total: int = 0
+    completed: int = 0
+    samples: dict[str, SampleState] = field(default_factory=dict)
+    start_time: float = field(default_factory=time.time)
+    gepa_iter: int | None = None
+    gepa_total: int | None = None
+    gepa_best: float | None = None
+    scores: list[float] = field(default_factory=list)
+
+
+# ── Pure functions ───────────────────────────────────────────────────────────
+
+
+def _format_time(seconds: float) -> str:
+    """Format seconds as H:MM:SS or M:SS."""
+    h = int(seconds) // 3600
+    m = int(seconds) % 3600 // 60
+    s = int(seconds) % 60
+    if h:
+        return f"{h}:{m:02d}:{s:02d}"
+    return f"{m}:{s:02d}"
+
+
+def _format_bar(progress: float, width: int) -> str:
+    """Render a progress bar with unicode blocks."""
+    if width <= 0:
+        return ""
+    filled = progress * width
+    full_blocks = int(filled)
+    partial_idx = int(8 * filled) % 8
+    partial = " ▏▎▍▌▋▊▉"[partial_idx] if partial_idx else ""
+    bar = "█" * full_blocks + partial
+    return bar.ljust(width)
+
+
+def derive_state(events: list[dict[str, Any]]) -> RenderState:
+    """Derive current display state from event stream. Stateless."""
+    state = RenderState()
+
+    for event in events:
+        event_type = event.get("message")
+        if event_type is None or event_type not in KNOWN_EVENT_TYPES:
+            continue
+
+        sample_id = event.get("sample_id", "")
+
+        if event_type == "eval_start":
+            state.eval_name = event.get("eval_name", "eval")
+            state.total = event.get("total", 0)
+            if ts := event.get("timestamp"):
+                from datetime import datetime
+
+                state.start_time = datetime.fromisoformat(ts.replace("Z", "+00:00")).timestamp()
+            else:
+                state.start_time = time.time()
+
+        elif event_type == "sample_start":
+            state.samples[sample_id] = SampleState(
+                id=sample_id,
+                name=event.get("sample_name", sample_id),
+                last_update=time.time(),
+                phase="starting",
+            )
+
+        elif event_type == "turn":
+            if sample_id in state.samples:
+                if "turn" in event:
+                    state.samples[sample_id].turn = event["turn"]
+                state.samples[sample_id].last_update = time.time()
+                status = event.get("status", "running")
+                state.samples[sample_id].phase = status
+
+        elif event_type == "modal_progress":
+            if sample_id in state.samples:
+                state.samples[sample_id].phase = event.get("phase", "")
+                state.samples[sample_id].last_update = time.time()
+
+        elif event_type == "sample_end":
+            if sample_id in state.samples:
+                state.samples[sample_id].status = "complete"
+                score = event.get("score")
+                if score is not None:
+                    state.samples[sample_id].score = score
+                    state.scores.append(score)
+                state.completed += 1
+
+        elif event_type == "sample_retry":
+            if sample_id in state.samples:
+                state.samples[sample_id].status = "retry"
+                state.samples[sample_id].retry_attempt = event.get("attempt", 1)
+
+        elif event_type == "gepa_iteration":
+            state.gepa_iter = event.get("iter")
+            state.gepa_total = event.get("total")
+            state.gepa_best = event.get("best")
+
+        elif event_type == "eval_end":
+            pass  # Handled by update() -> Cmd.quit()
+
+    return state
+
+
+def _format_sample_row(sample: SampleState, width: int) -> str:
+    """Format a single sample row."""
+    name = sample.name[:25].ljust(25)
+    turn_info = f"T:{sample.turn}"
+
+    if sample.status == "complete":
+        if sample.score is not None:
+            icon = f"{GREEN}✓{RESET}" if sample.score > 0.5 else f"{RED}✗{RESET}"
+            return f"  {name} {turn_info:>4}  {icon} score={sample.score:.2f}"
+        else:
+            return f"  {name} {turn_info:>4}  {GREEN}✓{RESET} done"
+    elif sample.status == "retry":
+        return f"  {name} {turn_info:>4}  {YELLOW}⟳{RESET} retry (attempt {sample.retry_attempt})"
+    else:
+        phase = sample.phase or "running"
+        return f"  {name} {turn_info:>4}  {CYAN}{phase}...{RESET}"
+
+
+# ── Render function ──────────────────────────────────────────────────────────
+
+
+def render(state: RenderState, width: int, height: int) -> list[str]:
+    """Render state to list of lines. Pure function."""
+    lines: list[str] = []
+
+    # Header
+    if state.gepa_iter is not None:
+        header = (
+            f"GEPA iter {state.gepa_iter}/{state.gepa_total or '?'} │ best: {state.gepa_best:.0%}"
+            if state.gepa_best
+            else f"GEPA iter {state.gepa_iter}/{state.gepa_total or '?'}"
+        )
+    else:
+        elapsed = time.time() - state.start_time
+        if state.total > 0:
+            progress = state.completed / state.total
+            bar_width = min(30, width - 50)
+            bar = _format_bar(progress, bar_width)
+            eta = ""
+            if state.completed > 0 and progress < 1.0:
+                remaining = (elapsed / progress) - elapsed
+                eta = f" <{_format_time(remaining)}"
+            header = (
+                f"{state.eval_name}: {BOLD}{state.completed}/{state.total}{RESET} "
+                f"|{bar}| {100 * progress:3.0f}% "
+                f"[{_format_time(elapsed)}{eta}]"
+            )
+        else:
+            header = f"{state.eval_name}: {state.completed} samples [{_format_time(elapsed)}]"
+
+    lines.append(header[:width])
+
+    # Active samples, sorted by most recently updated
+    active = [s for s in state.samples.values() if s.status != "complete" and s.phase]
+    active.sort(key=lambda s: s.last_update, reverse=True)
+
+    max_samples = max(1, height - 4)
+    for sample in active[:max_samples]:
+        lines.append(_format_sample_row(sample, width))
+
+    total_in_flight = sum(1 for s in state.samples.values() if s.status != "complete")
+    hidden = total_in_flight - len(active[:max_samples])
+    if hidden > 0:
+        lines.append(f"{DIM}  ... and {hidden} more in flight{RESET}")
+
+    if state.scores:
+        mean_score = sum(state.scores) / len(state.scores)
+        lines.append(f"{DIM}score: {mean_score:.1%}{RESET}")
+
+    return lines
+
+
+# ── Elm architecture (subprocess entry point) ────────────────────────────────
+
+
+@dataclass(frozen=True)
+class Model:
+    """Immutable model for the pytui App."""
+
+    events: tuple[dict[str, Any], ...] = ()
+    done: bool = False
+
+
+@dataclass(frozen=True)
+class NewEvent:
+    line: str
+
+
+def _parse_event(line: str) -> dict[str, Any] | None:
+    try:
+        return json.loads(line)
+    except json.JSONDecodeError:
+        return None
+
+
+def update(model: Model, msg: object) -> tuple[Model, Any]:
+    from ._pytui import Cmd, KeyPress
+
+    match msg:
+        case KeyPress(key="q" | "\x03"):
+            return model, Cmd.quit()
+        case NewEvent(line=line):
+            event = _parse_event(line)
+            if event is None:
+                return model, Cmd.none()
+            new_events = model.events + (event,)
+            if event.get("message") == "eval_end":
+                return Model(events=new_events, done=True), Cmd.quit()
+            return Model(events=new_events, done=model.done), Cmd.none()
+    return model, Cmd.none()
+
+
+def view(model: Model, width: int, height: int) -> list[str]:
+    state = derive_state(list(model.events))
+    lines = render(state, width, height)
+
+    # Pad to fill screen, add footer
+    footer = f"{DIM}q: quit{RESET}"
+    while len(lines) < height - 1:
+        lines.append("")
+    lines.append(footer)
+    return lines[:height]
+
+
+# ── Context manager (spawns subprocess) ──────────────────────────────────────
+
+
+@contextmanager
+def progress_display(
+    output_dir: Path | str,
+    disable: bool = False,
+) -> Generator[Path, None, None]:
+    """Context manager that spawns a pytui progress display in a subprocess.
+
+    The subprocess owns the terminal (alternate screen, raw mode via /dev/tty).
+    The parent process writes events.jsonl via logging as usual.
+
+    Args:
+        output_dir: Directory containing events.jsonl
+        disable: If True, skip progress display (verbose mode)
+    """
+    if disable:
+        yield Path(output_dir)
+        return
+
+    output_dir = Path(output_dir)
+    output_dir.mkdir(parents=True, exist_ok=True)
+    events_file = output_dir / "events.jsonl"
+
+    proc = subprocess.Popen(
+        [sys.executable, "-m", "wafer_core.rollouts.progress_app", str(events_file)],
+        # Don't touch stdout/stderr — child inherits the terminal.
+        # pytui opens /dev/tty directly for input, writes to stdout for rendering.
+        #
+        # start_new_session: put subprocess in its own session so it doesn't
+        # receive SIGINT from Ctrl+C. The parent gets SIGINT and terminates
+        # the subprocess in the finally block. Without this, pytui's raw mode
+        # swallows Ctrl+C (terminal driver doesn't generate SIGINT in raw mode)
+        # and the parent never sees it.
+        stdin=subprocess.DEVNULL,
+        start_new_session=True,
+    )
+
+    try:
+        yield output_dir
+    finally:
+        # Give the subprocess a chance to clean up its terminal state.
+        proc.send_signal(signal.SIGTERM)
+        try:
+            proc.wait(timeout=5.0)
+        except subprocess.TimeoutExpired:
+            proc.kill()
+            proc.wait()
+
+        # Defensive: restore terminal to sane state in case the subprocess
+        # didn't clean up (e.g. got SIGKILL, or crashed before restoring).
+        # stty sane resets raw mode, re-enables echo, etc.
+        tty_fd = os.open("/dev/tty", os.O_RDWR)
+        try:
+            subprocess.run(["stty", "sane"], stdin=tty_fd)
+        finally:
+            os.close(tty_fd)
+
+
+# ── CLI entry point (subprocess runs this) ────────────────────────────────────
+
+if __name__ == "__main__":
+    if len(sys.argv) != 2:
+        print("Usage: python -m wafer_core.rollouts.progress_app <events.jsonl>", file=sys.stderr)
+        sys.exit(1)
+
+    from ._pytui import App, Cmd, Sub
+
+    events_file = sys.argv[1]
+
+    # Model needs the path for subscriptions. We use a non-frozen wrapper
+    # to thread the path through — subscriptions() needs it but it's static.
+    _events_path = events_file
+
+    # Patch subscriptions to close over the path
+    def _subscriptions(model: Model) -> Sub:
+        if model.done:
+            return Sub.none()
+        return Sub.file_tail(_events_path, lambda line: NewEvent(line=line))
+
+    App(
+        init=(Model(), Cmd.none()),
+        update=update,
+        view=view,
+        subscriptions=_subscriptions,
+        alternate_screen=False,
+    ).run()