tenant.openagents/omega

1from __future__ import annotations
2
3import argparse
4import subprocess
5import sys
6import tarfile
7import time
8from pathlib import Path
9
10from . import run_index, source
11from .common import (
12    default_namespace,
13    deployed_function,
14    print_json,
15    print_table,
16    resolve_run_location,
17    run_command,
18    safe_extract_archive,
19)
20
21RECENT_RUN_COLUMNS = [
22    ("created_at", "created"),
23    ("run_id", "run"),
24    ("experiment_name", "experiment"),
25    ("namespace", "namespace"),
26    ("agent_model", "model"),
27    ("judge_preset", "judge"),
28    ("build_id", "build"),
29    ("suite_id", "suite"),
30]
31REMOTE_RUN_COLUMNS = [
32    ("updated_at", "updated"),
33    ("namespace", "namespace"),
34    ("experiment_name", "experiment"),
35    ("run_id", "run"),
36    ("status", "status"),
37    ("agent_model", "model"),
38    ("judge_preset", "judge"),
39    ("build_id", "build"),
40    ("suite_id", "suite"),
41]
42BUILD_COLUMNS = [
43    ("built_at_utc", "built"),
44    ("build_id", "build"),
45    ("ready", "ready"),
46    ("base_sha", "base"),
47    ("patch_sha256", "patch"),
48]
49SUITE_COLUMNS = [
50    ("part", "part"),
51    ("experiment_name", "experiment"),
52    ("run_id", "run"),
53    ("status", "status"),
54    ("updated_at", "updated"),
55    ("trial_count", "trials"),
56]
57
58
59def volume_get(
60    args: argparse.Namespace, remote_path: str, local_path: str
61) -> subprocess.CompletedProcess[str]:
62    return run_command(
63        ["modal", "volume", "get", args.volume, remote_path, local_path],
64        capture=local_path == "-",
65    )
66
67
68def read_volume_text(args: argparse.Namespace, remote_path: str) -> str:
69    result = volume_get(args, remote_path, "-")
70    return result.stdout
71
72
73def build_ready_on_volume(args: argparse.Namespace, build_id: str) -> bool:
74    try:
75        read_volume_text(args, f"builds/{build_id}/READY")
76        return True
77    except (subprocess.CalledProcessError, OSError):
78        return False
79
80
81def ensure_volume_exists(args: argparse.Namespace) -> int:
82    print(f"\nEnsuring Modal volume '{args.volume}' exists...")
83    try:
84        result = subprocess.run(
85            ["modal", "volume", "create", args.volume],
86            text=True,
87            capture_output=True,
88            check=True,
89        )
90        if result.stdout:
91            print(result.stdout, end="")
92    except subprocess.CalledProcessError as error:
93        output = (error.stdout or "") + (error.stderr or "")
94        if "already exists" in output.lower():
95            print(f"Volume '{args.volume}' already exists; continuing.")
96        else:
97            if output:
98                print(output, end="", file=sys.stderr)
99            return error.returncode
100    return 0
101
102
103def print_rows(
104    rows: list[dict], columns: list[tuple[str, str]], *, as_json: bool = False
105) -> None:
106    if as_json:
107        print_json(rows)
108    else:
109        print_table(rows, columns)
110
111
112def run_remote_prefix(args: argparse.Namespace, run_id: str) -> str:
113    namespace, experiment_name = resolve_run_location(args, run_id)
114    return f"runs/{namespace}/{experiment_name}/{run_id}"
115
116
117def read_run_file(args: argparse.Namespace, file_name: str) -> str:
118    run_id = resolve_run_id(args)
119    return read_volume_text(args, f"{run_remote_prefix(args, run_id)}/{file_name}")
120
121
122def suite_rows(args: argparse.Namespace) -> list[dict]:
123    namespace = default_namespace(args)
124    return deployed_function(args, "suite_status").remote(namespace, args.suite_id)
125
126
127def suite_member_args(args: argparse.Namespace, row: dict) -> argparse.Namespace:
128    run_args = argparse.Namespace(**vars(args))
129    run_args.experiment_name = row["experiment_name"]
130    run_args.run_id = row["run_id"]
131    return run_args
132
133
134def resolve_run_id(args: argparse.Namespace) -> str:
135    """Return the run id to act on, defaulting to the most recent local run.
136
137    Commands like `status`/`logs` accept no run id at all, in which case we use
138    the most recently launched run recorded in the local index so a bare
139    `zed-eval status` answers "how's my latest run doing?".
140    """
141    run_id = getattr(args, "run_id", None)
142    if run_id:
143        return run_id
144    entry = run_index.most_recent()
145    if not entry:
146        raise ValueError(
147            "no run id given and the local run index is empty. Launch a run "
148            "first, or pass a run id explicitly."
149        )
150    args.run_id = entry["run_id"]
151    print(
152        f"(no run id given; using most recent: {entry['run_id']} / "
153        f"{entry.get('experiment_name')})",
154        file=sys.stderr,
155    )
156    return entry["run_id"]
157
158
159def command_runs(args: argparse.Namespace) -> int:
160    entries = run_index.recent(args.limit)
161    if getattr(args, "json", False):
162        print_json(entries)
163        return 0
164    if not entries:
165        print(
166            "No runs recorded locally yet. Launch a run, or use "
167            "`zed-eval list --details` to query runs on the volume."
168        )
169        return 0
170    print_table(entries, RECENT_RUN_COLUMNS)
171    return 0
172
173
174def command_list(args: argparse.Namespace) -> int:
175    namespace = None if args.all_namespaces else default_namespace(args)
176    experiment_name = getattr(args, "experiment_name", None)
177    if experiment_name:
178        experiment_name = source.sanitize_namespace(experiment_name)
179    if args.details or args.json:
180        rows = deployed_function(args, "list_runs").remote(
181            namespace, experiment_name, args.limit
182        )
183        print_rows(rows, REMOTE_RUN_COLUMNS, as_json=args.json)
184        return 0
185    remote_path = (
186        f"runs/{namespace}/{experiment_name}/"
187        if experiment_name
188        else f"runs/{namespace}/"
189    )
190    run_command(["modal", "volume", "ls", args.volume, remote_path])
191    return 0
192
193
194def command_builds(args: argparse.Namespace) -> int:
195    if args.details or args.json:
196        rows = deployed_function(args, "list_builds").remote(args.limit)
197        print_rows(rows, BUILD_COLUMNS, as_json=args.json)
198        return 0
199    run_command(["modal", "volume", "ls", args.volume, "builds/"])
200    return 0
201
202
203def command_status(args: argparse.Namespace) -> int:
204    # One-shot status ping: the controller's state.json only carries a coarse
205    # status (pending -> running -> completed/failed), so there is no per-trial
206    # progress worth following. Print it once and return.
207    print(read_run_file(args, "state.json"), end="")
208    return 0
209
210
211def command_logs(args: argparse.Namespace) -> int:
212    print(read_run_file(args, "controller.log"), end="")
213    return 0
214
215
216def command_fetch(args: argparse.Namespace) -> int:
217    jobs_dir = Path(args.jobs_dir).expanduser()
218    jobs_dir.mkdir(parents=True, exist_ok=True)
219    temporary_archive = jobs_dir / f".{args.run_id}.tar.gz"
220    remote_path = f"{run_remote_prefix(args, args.run_id)}/harbor-job.tar.gz"
221    print(f"Fetching {args.volume}:/{remote_path} -> {temporary_archive}")
222    volume_get(args, remote_path, str(temporary_archive))
223    with tarfile.open(temporary_archive, "r:gz") as archive:
224        safe_extract_archive(archive, jobs_dir)
225    temporary_archive.unlink(missing_ok=True)
226    print(f"Extracted Harbor job under {jobs_dir / args.run_id}")
227    return 0
228
229
230def command_suite_status(args: argparse.Namespace) -> int:
231    print_rows(suite_rows(args), SUITE_COLUMNS, as_json=args.json)
232    return 0
233
234
235def command_suite_logs(args: argparse.Namespace) -> int:
236    rows = suite_rows(args)
237    while True:
238        for row in rows:
239            print(
240                f"\n=== {row['part']} / {row['experiment_name']} / {row['run_id']} ==="
241            )
242            run_args = suite_member_args(args, row)
243            run_args.follow = False
244            command_logs(run_args)
245        if not args.follow:
246            return 0
247        time.sleep(args.interval)
248
249
250def command_suite_fetch(args: argparse.Namespace) -> int:
251    result = 0
252    for row in suite_rows(args):
253        result = command_fetch(suite_member_args(args, row)) or result
254    return result
255