Merge pull request #127 from WecoAI/feature/run-subcommands

Add weco run subcommands for inspecting and managing optimization runs

Author: aliroberts

README.md

@@ -186,14 +186,50 @@ For more advanced examples, including [Triton](/examples/triton/README.md), [CUD
 
 | Command | Description | When to Use |
 |---------|-------------|-------------|
-| `weco run [options]` | Direct optimization execution | **For advanced users** - When you know exactly what to optimize and how |
+| `weco run [options]` | Start a new optimization | When you know what to optimize and how |
 | `weco resume <run-id>` | Resume an interrupted run | Continue from the last completed step |
 | `weco login` | Authenticate with Weco | First-time setup or switching accounts |
-| `weco logout` | Clear authentication credentials | To switch accounts or troubleshoot authentication issues |
+| `weco logout` | Clear authentication credentials | Switch accounts or troubleshoot auth |
 | `weco credits balance` | Check your current credit balance | Monitor usage |
 | `weco credits topup [amount]` | Purchase additional credits | When you need more credits (default: 10) |
 | `weco credits autotopup` | Configure automatic top-up | Set up automatic credit replenishment |
 
+### Run Subcommands
+
+Inspect and manage optimization runs. All output is JSON, designed for programmatic access (AI coding agents, scripts).
+
+| Command | Description |
+|---------|-------------|
+| `weco run status <run-id>` | Run progress, pending nodes, review mode flag |
+| `weco run results <run-id>` | Results sorted by metric |
+| `weco run show <run-id> --step <N\|best>` | Single node detail with code |
+| `weco run diff <run-id> --step <N\|best>` | Unified code diff between steps |
+| `weco run stop <run-id>` | Graceful termination (tree preserved) |
+| `weco run instruct <run-id> "<text>"` | Update instructions mid-run |
+| `weco run review <run-id>` | List pending approval nodes (review mode) |
+| `weco run revise <run-id> --node <id> --source <file>` | Replace a node's code |
+| `weco run submit <run-id> --node <id>` | Evaluate and submit a node |
+
+```bash
+# Check progress
+weco run status 0002e071-1b67-411f-a514-36947f0c4b31
+
+# Top 5 results as JSON
+weco run results 0002e071-1b67-411f-a514-36947f0c4b31 --top 5
+
+# Diff best solution against baseline
+weco run diff 0002e071-1b67-411f-a514-36947f0c4b31 --step best
+
+# Review mode: inspect, optionally edit, and submit
+weco run review 0002e071-1b67-411f-a514-36947f0c4b31
+weco run submit 0002e071-1b67-411f-a514-36947f0c4b31 --node <node-id>
+
+# Submit with your own code (explicit path mapping)
+weco run submit <run-id> --node <id> --source module.py=./my_version.py
+```
+
+**Source path mapping:** When using `--source` with `revise` or `submit`, you can map local files to the run's source paths using `target_path=local_path` syntax (e.g., `--source module.py=./optimized.py`). Without an explicit mapping, files are matched positionally to the run's original source paths.
+
 ### Observe Commands
 
 Track experiments from your own optimization loop (LLM agents, custom scripts, manual experiments) in the Weco dashboard:

weco/cli.py

@@ -176,6 +176,72 @@ def configure_run_parser(run_parser: argparse.ArgumentParser) -> None:
         _load_backend(backend_name).register_args(run_parser)
 
 
+def _configure_run_subcommands(run_parser: argparse.ArgumentParser) -> None:
+    """Add subcommands under ``weco run`` for inspecting and managing existing runs."""
+    subs = run_parser.add_subparsers(dest="run_subcommand")
+
+    # weco run status <run-id>
+    p = subs.add_parser("status", help="Show run status and progress (JSON)")
+    p.add_argument("run_id", type=str, help="Run UUID")
+
+    # weco run results <run-id>
+    p = subs.add_parser("results", help="Show results sorted by metric")
+    p.add_argument("run_id", type=str, help="Run UUID")
+    p.add_argument("--top", type=int, default=None, help="Show only the top N results")
+    p.add_argument("--format", type=str, choices=["json", "table", "csv"], default="json", help="Output format")
+    p.add_argument("--plot", action="store_true", help="Show ASCII metric trajectory")
+    p.add_argument("--include-code", action="store_true", help="Include full source code")
+
+    # weco run show <run-id> --step N
+    p = subs.add_parser("show", help="Show details for a specific step")
+    p.add_argument("run_id", type=str, help="Run UUID")
+    p.add_argument("--step", type=str, required=True, help="Step number or 'best'")
+
+    # weco run diff <run-id> --step N
+    p = subs.add_parser("diff", help="Show code diff between steps")
+    p.add_argument("run_id", type=str, help="Run UUID")
+    p.add_argument("--step", type=str, required=True, help="Step number or 'best'")
+    p.add_argument("--against", type=str, default="baseline", help="'baseline' (default), 'parent', or step number")
+
+    # weco run stop <run-id>
+    p = subs.add_parser("stop", help="Terminate a running optimization")
+    p.add_argument("run_id", type=str, help="Run UUID")
+
+    # weco run instruct <run-id> <instructions>
+    p = subs.add_parser("instruct", help="Update additional instructions for an active run")
+    p.add_argument("run_id", type=str, help="Run UUID")
+    p.add_argument("instructions", type=str, help="New instructions (text or path to file)")
+
+    # weco run review <run-id>
+    p = subs.add_parser("review", help="Show pending approval nodes (review mode)")
+    p.add_argument("run_id", type=str, help="Run UUID")
+
+    # weco run revise <run-id> --node <id> --source <file>
+    p = subs.add_parser("revise", help="Replace a pending node's code with a new revision")
+    p.add_argument("run_id", type=str, help="Run UUID")
+    p.add_argument("--node", type=str, required=True, help="Node ID to revise")
+    revise_source = p.add_mutually_exclusive_group(required=True)
+    revise_source.add_argument("-s", "--source", type=str, help="Path to a single source file")
+    revise_source.add_argument("--sources", nargs="+", type=str, help="Paths to multiple source files")
+
+    # weco run submit <run-id> --node <id>
+    p = subs.add_parser("submit", help="Submit a pending node for evaluation (review mode)")
+    p.add_argument("run_id", type=str, help="Run UUID")
+    p.add_argument("--node", type=str, required=True, help="Node ID to submit")
+    submit_source = p.add_mutually_exclusive_group()
+    submit_source.add_argument("-s", "--source", type=str, help="Optional: source file (creates revision before submitting)")
+    submit_source.add_argument(
+        "--sources", nargs="+", type=str, help="Optional: source files (creates revision before submitting)"
+    )
+    p.add_argument(
+        "-c",
+        "--eval-command",
+        type=str,
+        default=None,
+        help="Override the eval command (use when the stored command doesn't work in this environment)",
+    )
+
+
 def configure_credits_parser(credits_parser: argparse.ArgumentParser) -> None:
     """Configure the credits command parser and all its subcommands."""
     credits_subparsers = credits_parser.add_subparsers(dest="credits_command", help="Credit management commands")
@@ -278,6 +344,52 @@ def configure_resume_parser(resume_parser: argparse.ArgumentParser) -> None:
     )
 
 
+def _dispatch_run_subcommand(sub: str, args: argparse.Namespace) -> None:
+    """Dispatch ``weco run <subcommand>`` to the appropriate handler."""
+    from .commands.run import status, results, show, diff, stop, instruct, review, revise, submit
+
+    def _collect_source_paths() -> list[str] | None:
+        if getattr(args, "sources", None):
+            return args.sources
+        if getattr(args, "source", None):
+            return [args.source]
+        return None
+
+    handlers = {
+        "status": lambda: status.handle(run_id=args.run_id, console=console),
+        "results": lambda: results.handle(
+            run_id=args.run_id,
+            top=args.top,
+            format=args.format,
+            plot=args.plot,
+            include_code=args.include_code,
+            console=console,
+        ),
+        "show": lambda: show.handle(run_id=args.run_id, step=args.step, console=console),
+        "diff": lambda: diff.handle(run_id=args.run_id, step=args.step, against=args.against, console=console),
+        "stop": lambda: stop.handle(run_id=args.run_id, console=console),
+        "instruct": lambda: instruct.handle(run_id=args.run_id, instructions=args.instructions, console=console),
+        "review": lambda: review.handle(run_id=args.run_id, console=console),
+        "revise": lambda: revise.handle(
+            run_id=args.run_id, node_id=args.node, source_paths=_collect_source_paths(), console=console
+        ),
+        "submit": lambda: submit.handle(
+            run_id=args.run_id,
+            node_id=args.node,
+            source_paths=_collect_source_paths(),
+            eval_command_override=getattr(args, "eval_command", None),
+            console=console,
+        ),
+    }
+
+    handler = handlers.get(sub)
+    if handler is None:
+        console.print(f"[bold red]Unknown run subcommand: {sub}[/]")
+        sys.exit(1)
+    handler()
+    sys.exit(0)
+
+
 def execute_run_command(args: argparse.Namespace) -> None:
     """Execute the 'weco run' command with all its logic."""
     from .optimizer import optimize
@@ -419,9 +531,10 @@ def _main() -> None:
 
     # --- Run Command Parser Setup ---
     run_parser = subparsers.add_parser(
-        "run", help="Run code optimization", formatter_class=argparse.RawTextHelpFormatter, allow_abbrev=False
+        "run", help="Run and manage optimizations", formatter_class=argparse.RawTextHelpFormatter, allow_abbrev=False
     )
-    configure_run_parser(run_parser)  # Use the helper to add arguments
+    configure_run_parser(run_parser)  # Flags for starting a new optimization
+    _configure_run_subcommands(run_parser)  # Subcommands for inspecting/managing runs
 
     # --- Login Command Parser Setup ---
     _ = subparsers.add_parser("login", help="Log in to Weco and save your API key.")
@@ -500,7 +613,11 @@ def _main() -> None:
         clear_api_key()
         sys.exit(0)
     elif args.command == "run":
-        execute_run_command(args)
+        sub = getattr(args, "run_subcommand", None)
+        if sub is not None:
+            _dispatch_run_subcommand(sub, args)
+        else:
+            execute_run_command(args)
     elif args.command == "credits":
         from .credits import handle_credits_command
 

weco/commands/__init__.py

@@ -0,0 +1,108 @@
+"""Shared helpers for ``weco run`` subcommand handlers."""
+
+import json
+import pathlib
+import sys
+
+from rich.console import Console
+
+from ..core.api import WecoClient
+from ..auth import handle_authentication
+
+
+def make_client(console: Console) -> WecoClient:
+    """Authenticate and return a ``WecoClient``, or exit on failure."""
+    _, auth_headers = handle_authentication(console)
+    if not auth_headers:
+        sys.exit(1)
+    return WecoClient(auth_headers)
+
+
+def fetch_run(client: WecoClient, run_id: str, include_history: bool = True) -> dict:
+    """Fetch full run data via ``GET /runs/{run_id}``, or exit on error."""
+    try:
+        return client.get_run_status(run_id, include_history=include_history)
+    except Exception as e:
+        print(json.dumps({"error": f"Failed to fetch run {run_id}: {e}"}))
+        sys.exit(1)
+
+
+def fetch_nodes(
+    client: WecoClient,
+    run_id: str,
+    *,
+    step: int | None = None,
+    status: str | None = None,
+    top: int | None = None,
+    sort: str | None = None,
+    include_code: bool = True,
+) -> tuple[dict, list[dict]]:
+    """Fetch nodes via ``GET /runs/{run_id}/nodes``, or exit on error.
+
+    Returns:
+        A ``(run_metadata, nodes)`` tuple where ``run_metadata`` contains
+        lightweight run info (status, metric_name, best_metric, etc.) and
+        ``nodes`` is the filtered/sorted list of node dicts.
+    """
+    try:
+        data = client.list_nodes(run_id, step=step, status=status, top=top, sort=sort, include_code=include_code)
+        return data.get("run", {}), data.get("nodes", [])
+    except Exception as e:
+        print(json.dumps({"error": f"Failed to fetch nodes for run {run_id}: {e}"}))
+        sys.exit(1)
+
+
+def read_source_code(source_args: list[str], run_code_keys: list[str] | None = None) -> dict[str, str]:
+    """Read source files and map them to run code keys.
+
+    Each entry in *source_args* is either:
+
+    * ``target_path=local_path`` — explicit mapping (store content of
+      *local_path* under the key *target_path*)
+    * ``local_path`` — no mapping; falls back to positional matching
+      against *run_code_keys* if available, otherwise uses the local path.
+
+    Args:
+        source_args: Raw ``--source`` / ``--sources`` values from argparse.
+        run_code_keys: The run's original source file keys (from baseline
+            node).  Used for positional fallback when no explicit mapping.
+
+    Returns:
+        Dict mapping target path → file content.
+    """
+    explicit: dict[str, str] = {}
+    unmapped_paths: list[str] = []
+    unmapped_contents: list[str] = []
+
+    for arg in source_args:
+        if "=" in arg:
+            target, local = arg.split("=", 1)
+            target, local = target.strip(), local.strip()
+            if not target or not local:
+                print(json.dumps({"error": f"Invalid source mapping: '{arg}'. Use target_path=local_path"}))
+                sys.exit(1)
+            p = pathlib.Path(local)
+            if not p.exists():
+                print(json.dumps({"error": f"Source file not found: {local}"}))
+                sys.exit(1)
+            explicit[target] = p.read_text()
+        else:
+            p = pathlib.Path(arg)
+            if not p.exists():
+                print(json.dumps({"error": f"Source file not found: {arg}"}))
+                sys.exit(1)
+            unmapped_paths.append(arg)
+            unmapped_contents.append(p.read_text())
+
+    # Resolve unmapped files
+    if unmapped_contents:
+        if run_code_keys and len(unmapped_contents) == len(run_code_keys):
+            # Positional match against the run's source keys
+            for key, content in zip(run_code_keys, unmapped_contents):
+                explicit[key] = content
+        else:
+            # Use local paths as-is
+            for path, content in zip(unmapped_paths, unmapped_contents):
+                explicit[str(pathlib.Path(path))] = content
+
+    return explicit

weco/commands/run/__init__.py

@@ -0,0 +1 @@
+"""Commands under ``weco run <subcommand>``."""