runner.py

"""Tecpatl command-line entry point.

    python runner.py tasks/counter_4bit   run one task
    python runner.py sweep                run every task on every available provider
    python runner.py stats                summarise the trace dataset
    python runner.py failures             show each failure and the repair that followed
    python runner.py replay               re-ask old compile failures, both feedback styles
    python runner.py preflight            check there is disk room for a local model
    python runner.py tt-emit <task>       package a verified design as a Tiny Tapeout project

Running a task means: generate Verilog, check it behaves (cocotb), check it
can be built (Yosys), feed any failure back to the model, and record every
attempt to a JSONL trace.
"""

from __future__ import annotations

import argparse
import sys
from pathlib import Path

from agent.provider import ProviderError, available_providers, provider_status
from runner.classify import FAILURE_KINDS
from runner.failures import main as failures_main
from runner.orchestrator import DEFAULT_MAX_ATTEMPTS, Task, run_task
from runner.preflight import DEFAULT_HEADROOM_GB
from runner.preflight import main as preflight_main
from runner.replay import main as replay_main
from runner.stats import main as stats_main
from runner.ttproject import TEMPLATE_SOURCE, TEMPLATE_TAG, emit, verified_design
from runner.toolchain import ToolchainError, activate

REPO_ROOT = Path(__file__).resolve().parent
SUBCOMMANDS = ("stats", "failures", "sweep", "replay", "preflight", "tt-emit")


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        prog="runner.py",
        description="Generate and verify a Verilog design for one task.",
    )
    parser.add_argument(
        "task",
        type=Path,
        nargs="?",
        help="Path to a task directory, or one of: " + ", ".join(SUBCOMMANDS),
    )
    parser.add_argument(
        "--provider",
        choices=available_providers(),
        default=None,
        help="LLM backend (default: $TECPATL_PROVIDER, else claude)",
    )
    parser.add_argument(
        "--model",
        default=None,
        help="Override the backend's default model",
    )
    parser.add_argument(
        "--spec-level",
        default="full",
        help='Which specification to give the model (default: "full")',
    )
    parser.add_argument(
        "--max-attempts",
        type=int,
        default=DEFAULT_MAX_ATTEMPTS,
        help=f"Attempt budget before giving up (default: {DEFAULT_MAX_ATTEMPTS})",
    )
    parser.add_argument(
        "--check-toolchain",
        action="store_true",
        help="Report where the hardware tools were found, then exit",
    )
    parser.add_argument(
        "--tt-build",
        action="store_true",
        help=(
            "In gate 3, run the real LibreLane build rather than stopping at "
            "the contract and area tiers. Needs a PDK and a container runtime."
        ),
    )
    return parser


def build_stats_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        prog="runner.py stats",
        description="Summarise the JSONL traces produced so far.",
    )
    parser.add_argument(
        "--traces",
        type=Path,
        default=REPO_ROOT / "traces",
        help="Directory of traces to read (default: traces/)",
    )
    parser.add_argument(
        "--by-task",
        action="store_true",
        help="Break results down by task (the default view)",
    )
    parser.add_argument(
        "--by-provider",
        action="store_true",
        help="Break results down by provider and model",
    )
    parser.add_argument(
        "--by-spec-level",
        action="store_true",
        help="Break results down by how much the specification said",
    )
    return parser


def build_failures_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        prog="runner.py failures",
        description="Show each failed attempt and the repair that followed it.",
    )
    parser.add_argument(
        "--traces",
        type=Path,
        default=REPO_ROOT / "traces",
        help="Directory of traces to read (default: traces/)",
    )
    parser.add_argument("--task", default=None, help="Only this task")
    parser.add_argument("--provider", default=None, help="Only this provider")
    parser.add_argument(
        "--verdict",
        choices=("SIM_FAIL", "SYNTH_FAIL", "NO_DESIGN"),
        default=None,
        help="Only failures with this verdict",
    )
    parser.add_argument(
        "--kind",
        choices=FAILURE_KINDS,
        default=None,
        help="Only failures with this label",
    )
    parser.add_argument(
        "--spec-level",
        default=None,
        help="Only failures from runs at this spec level",
    )
    parser.add_argument(
        "--diff-lines",
        type=int,
        default=40,
        help="Diff lines to show per repair (default: 40)",
    )
    parser.add_argument(
        "--json",
        action="store_true",
        dest="as_json",
        help="Emit the error/repair pairs as JSON",
    )
    return parser


def build_replay_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        prog="runner.py replay",
        description=(
            "Re-ask every recorded compile failure twice: once described by "
            "the Python exception the old code sent, once by iverilog's own "
            "diagnostics. Measures whether the reply compiles. Writes no "
            "traces."
        ),
    )
    parser.add_argument(
        "--traces",
        type=Path,
        default=REPO_ROOT / "traces",
        help="Directory of traces to read (default: traces/)",
    )
    parser.add_argument(
        "--provider",
        choices=available_providers(),
        default=None,
        help="LLM backend to re-ask (default: $TECPATL_PROVIDER, else claude)",
    )
    parser.add_argument("--model", default=None, help="Override the backend's model")
    parser.add_argument("--task", default=None, help="Only failures from this task")
    parser.add_argument(
        "--repeats",
        type=int,
        default=1,
        help="Times to ask each prompt, to see past one sample (default: 1)",
    )
    return parser


def build_preflight_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        prog="runner.py preflight",
        description=(
            "Report where Ollama stores its models and whether there is room "
            "for another one. Reports only; downloads nothing."
        ),
    )
    parser.add_argument("--model", default=None, help="Name of the model to be pulled")
    parser.add_argument(
        "--size-gb",
        type=float,
        default=0.0,
        help="Download size in GB, from the model's page on ollama.com",
    )
    parser.add_argument(
        "--headroom-gb",
        type=float,
        default=DEFAULT_HEADROOM_GB,
        help=f"Space to leave free afterwards (default: {DEFAULT_HEADROOM_GB})",
    )
    return parser


def build_ttemit_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        prog="runner.py tt-emit",
        description=(
            "Package a design that already cleared all three gates as a "
            "Tiny Tapeout project directory, ready to be built by the "
            "shuttle's own flow. The Verilog is taken from the trace "
            "unmodified. This emits a project; it does not submit one."
        ),
    )
    parser.add_argument("task", type=Path, help="Task directory to package")
    parser.add_argument(
        "--out",
        type=Path,
        default=None,
        help="Directory to write the project into (default: build/tt-<task>)",
    )
    parser.add_argument(
        "--traces",
        type=Path,
        default=REPO_ROOT / "traces",
        help="Where to look for the verified design (default: traces/)",
    )
    return parser


def build_sweep_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        prog="runner.py sweep",
        description=(
            "Run every task on every provider that has working credentials. "
            "Providers without credentials are skipped with a reason."
        ),
    )
    parser.add_argument(
        "--providers",
        nargs="+",
        choices=available_providers(),
        default=None,
        help="Limit to these providers (default: every available one)",
    )
    parser.add_argument(
        "--tasks",
        nargs="+",
        default=None,
        help="Limit to these task names (default: every task in tasks/)",
    )
    parser.add_argument(
        "--spec-levels",
        nargs="+",
        default=None,
        help="Limit to these spec levels (default: every level each task offers)",
    )
    parser.add_argument(
        "--max-attempts",
        type=int,
        default=DEFAULT_MAX_ATTEMPTS,
        help=f"Attempt budget per run (default: {DEFAULT_MAX_ATTEMPTS})",
    )
    return parser


def run_sweep(args: argparse.Namespace) -> int:
    """Run the cross-product of providers, tasks and spec levels.

    A task that offers only a full spec contributes one run per provider; a
    task with a terse variant contributes two, so the same circuit is measured
    against two amounts of prose with the same reference testbench.
    """
    task_dirs = sorted(
        p for p in (REPO_ROOT / "tasks").iterdir() if (p / "task.json").is_file()
    )
    if args.tasks:
        wanted = set(args.tasks)
        task_dirs = [p for p in task_dirs if p.name in wanted]

    status = provider_status()
    candidates = args.providers or available_providers()
    usable = [name for name in candidates if status[name].available]
    skipped = [name for name in candidates if not status[name].available]

    print("Sweep")
    print("=" * 72)
    print(f"  tasks      : {', '.join(p.name for p in task_dirs)}")
    print(f"  providers  : {', '.join(usable) if usable else '(none available)'}")
    for name in skipped:
        print(f"  skipping {name}: {status[name].reason}")
    print()

    if not usable or not task_dirs:
        print("Nothing to run.")
        return 1

    # Build the full matrix up front so its size is visible before it runs.
    matrix: list[tuple[str, Path, str]] = []
    for provider_name in usable:
        for task_dir in task_dirs:
            levels = Task.levels(task_dir)
            if args.spec_levels:
                levels = [lv for lv in levels if lv in set(args.spec_levels)]
            for level in levels:
                matrix.append((provider_name, task_dir, level))

    print(f"  runs       : {len(matrix)}")
    print()
    if not matrix:
        print("Nothing to run.")
        return 1

    results: list[tuple[str, str, str, bool]] = []
    for provider_name, task_dir, level in matrix:
        print("=" * 72)
        try:
            passed = run_task(
                task_dir,
                spec_level=level,
                provider_name=provider_name,
                max_attempts=args.max_attempts,
            )
        except ProviderError as exc:
            # One backend failing must not abandon the rest of the sweep.
            print(f"Provider error on {task_dir.name}: {exc}", file=sys.stderr)
            passed = False
        results.append((provider_name, task_dir.name, level, passed))
        print()

    print("=" * 72)
    print("Sweep summary")
    for provider_name, task_name, level, passed in results:
        print(
            f"  {'PASS' if passed else 'FAIL'}  {provider_name:<8} "
            f"{task_name:<18} {level}"
        )
    verified = sum(1 for *_, ok in results if ok)
    print(f"  {verified}/{len(results)} verified")
    print()
    print("  python runner.py stats --by-spec-level    to compare the levels")
    print("  python runner.py failures                 to read the failures")
    return 0 if verified == len(results) else 1


def ttemit_main(args) -> int:
    """Package the verified design for one task as a Tiny Tapeout project."""
    if not args.task.is_dir():
        print(f"Not a task directory: {args.task}", file=sys.stderr)
        return 2

    task = Task.load(args.task)
    if not task.tinytapeout:
        print(
            f"Task {task.name} declares no 'tinytapeout' block in task.json, "
            "so there is no project metadata to emit.",
            file=sys.stderr,
        )
        return 2

    found = verified_design(args.traces, task.name)
    if found is None:
        print(
            f"No verified design for {task.name} in {args.traces}. Run the "
            f"task first: python runner.py {args.task}",
            file=sys.stderr,
        )
        return 1
    verilog, trace_path = found

    out_dir = args.out or REPO_ROOT / "build" / f"tt-{task.name}"
    test_source = None
    test_file = args.task / f"{task.test_module}.py"
    if test_file.is_file():
        test_source = test_file.read_text(encoding="utf-8")

    project = emit(
        out_dir,
        verilog=verilog,
        top_module=task.top_module,
        config=task.tinytapeout,
        spec=task.spec,
        test_source=test_source,
    )

    print("Tiny Tapeout project emitted")
    print("=" * 72)
    print()
    print(f"  task         {task.name}")
    print(f"  top module   {project.top_module}")
    print(f"  tiles        {project.tiles}")
    print(f"  design from  {trace_path.name}")
    print(f"  template     {TEMPLATE_SOURCE} @ {TEMPLATE_TAG}")
    print(f"  written to   {project.directory}")
    print()
    for name in project.files:
        print(f"    {name}")
    print()
    print("  Push this to GitHub and its gds workflow will build the layout,")
    print("  run the Tiny Tapeout precheck and re-test the gate-level netlist.")
    print("  Emitting a project is not submitting one: this enters no shuttle.")
    return 0


def main(argv: list[str] | None = None) -> int:
    argv = list(sys.argv[1:] if argv is None else argv)

    # These read files the runner already wrote, or drive their own runs, so
    # they are dispatched before the single-task argument parsing.
    if argv and argv[0] == "stats":
        args = build_stats_parser().parse_args(argv[1:])
        return stats_main(
            args.traces,
            REPO_ROOT / "tasks",
            by_provider=args.by_provider,
            by_task=args.by_task,
            by_spec_level=args.by_spec_level,
        )

    if argv and argv[0] == "failures":
        args = build_failures_parser().parse_args(argv[1:])
        return failures_main(
            args.traces,
            task=args.task,
            provider=args.provider,
            verdict=args.verdict,
            kind=args.kind,
            spec_level=args.spec_level,
            as_json=args.as_json,
            diff_lines=args.diff_lines,
        )

    if argv and argv[0] == "tt-emit":
        args = build_ttemit_parser().parse_args(argv[1:])
        return ttemit_main(args)

    if argv and argv[0] == "preflight":
        args = build_preflight_parser().parse_args(argv[1:])
        return preflight_main(
            model=args.model,
            size_gb=args.size_gb,
            headroom_gb=args.headroom_gb,
        )

    if argv and argv[0] == "replay":
        # Needs iverilog, since the measurement is whether the reply compiles.
        try:
            activate()
        except ToolchainError as exc:
            print(f"Toolchain error: {exc}", file=sys.stderr)
            return 2
        args = build_replay_parser().parse_args(argv[1:])
        try:
            return replay_main(
                args.traces,
                provider_name=args.provider,
                model=args.model,
                repeats=args.repeats,
                task=args.task,
            )
        except ProviderError as exc:
            print(f"Provider error: {exc}", file=sys.stderr)
            return 2

    if argv and argv[0] == "sweep":
        try:
            activate()
        except ToolchainError as exc:
            print(f"Toolchain error: {exc}", file=sys.stderr)
            return 2
        return run_sweep(build_sweep_parser().parse_args(argv[1:]))

    args = build_parser().parse_args(argv)

    # Fail early and legibly if the simulator is missing, rather than after
    # spending a model call on a design we cannot check.
    try:
        report = activate()
    except ToolchainError as exc:
        print(f"Toolchain error: {exc}", file=sys.stderr)
        return 2

    if args.check_toolchain:
        for key, value in report.items():
            print(f"{key:>14}: {value}")
        return 0

    if args.task is None:
        print("A task directory is required.", file=sys.stderr)
        return 2

    if not args.task.is_dir():
        print(f"Not a task directory: {args.task}", file=sys.stderr)
        return 2

    try:
        passed = run_task(
            args.task,
            spec_level=args.spec_level,
            provider_name=args.provider,
            model=args.model,
            max_attempts=args.max_attempts,
            tt_build=args.tt_build,
        )
    except ProviderError as exc:
        print(f"Provider error: {exc}", file=sys.stderr)
        return 2

    return 0 if passed else 1


if __name__ == "__main__":
    sys.exit(main())