datadb's picture
upload track a submission package
80dcfe9 verified
Raw
History Blame Contribute Delete
7.32 kB
#!/usr/bin/env python
"""Submission runner for Track A.
This wrapper keeps the packaged output in the Phase 3 guideline shape:
result/
traces.json
results.csv
runtime.json
It delegates prediction to the copied main.py, then normalizes filenames and
columns. It does not train or download anything.
"""
from __future__ import annotations
import argparse
import csv
import json
import os
import shutil
import subprocess
import sys
from pathlib import Path
ROOT = Path(__file__).resolve().parent
def _resolve(path_text: str) -> Path:
path = Path(path_text)
return path if path.is_absolute() else ROOT / path
def _api_key_names(spec: str) -> list[str]:
return [name.strip() for name in spec.split(",") if name.strip()]
def _ensure_local_vllm_key(env: dict[str, str], spec: str) -> None:
names = _api_key_names(spec)
if names and not any(env.get(name) for name in names):
env[names[0]] = "EMPTY"
def _write_results_csv(raw_result_csv: Path, output_csv: Path) -> None:
with raw_result_csv.open("r", encoding="utf-8", newline="") as f:
reader = csv.DictReader(f)
rows = list(reader)
output_csv.parent.mkdir(parents=True, exist_ok=True)
with output_csv.open("w", encoding="utf-8", newline="") as f:
writer = csv.DictWriter(f, fieldnames=["scenario_id", "prediction"])
writer.writeheader()
for row in rows:
scenario_id = row.get("scenario_id") or row.get("ID") or row.get("id") or ""
prediction = row.get("prediction")
if prediction is None:
prediction = row.get("Track A", "")
writer.writerow(
{
"scenario_id": str(scenario_id).strip(),
"prediction": str(prediction or "").strip(),
}
)
def _normalize_traces(traces_path: Path) -> list[dict]:
if not traces_path.exists():
raise FileNotFoundError(f"Missing traces file: {traces_path}")
with traces_path.open("r", encoding="utf-8") as f:
traces = json.load(f)
if not isinstance(traces, list):
raise ValueError("traces.json must contain a JSON list")
normalized = []
for i, row in enumerate(traces):
if not isinstance(row, dict):
row = {"completion": str(row)}
row.setdefault("scenario_id", row.get("ID", f"scenario_{i + 1}"))
row.setdefault("completion_id", 0)
row.setdefault("completion", "")
normalized.append(row)
with traces_path.open("w", encoding="utf-8") as f:
json.dump(normalized, f, indent=2, ensure_ascii=False)
return normalized
def _write_runtime_json(traces: list[dict], runtime_path: Path) -> None:
runtimes = []
for row in traces:
runtime_seconds = row.get("runtime_seconds", row.get("execution_time_seconds", 0.0))
try:
runtime_seconds = float(runtime_seconds)
except (TypeError, ValueError):
runtime_seconds = 0.0
runtimes.append(
{
"scenario_id": str(row.get("scenario_id", "")),
"runtime_seconds": runtime_seconds,
}
)
with runtime_path.open("w", encoding="utf-8") as f:
json.dump(runtimes, f, indent=2)
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser()
parser.add_argument(
"--input",
default=os.environ.get("TEST_PATH", "data/Phase_2/test.json"),
help="Path to the Track A test JSON file.",
)
parser.add_argument(
"--output",
default="result",
help="Directory where traces.json, results.csv, and runtime.json are written.",
)
parser.add_argument(
"--model_bundle",
default="models/model_v4_bundle.pkl",
help="Path to the auxiliary Track A model bundle.",
)
parser.add_argument(
"--server_url",
default=os.environ.get("TRACK_A_SERVER_URL", "https://localhost:8081/no"),
)
parser.add_argument(
"--model_url",
default=os.environ.get("OPENAI_BASE_URL", "http://localhost:8001/v1"),
)
parser.add_argument(
"--model_name",
default=os.environ.get("OPENAI_MODEL", "Qwen3.5-35B-A3B"),
)
parser.add_argument(
"--api_key_env",
default="OPENAI_API_KEY,AGENT_API_KEY,OPENROUTER_API_KEY",
)
parser.add_argument("--concurrency", type=int, default=1)
parser.add_argument("--max_steps", type=int, default=4)
parser.add_argument("--max_tool_calls", type=int, default=8)
parser.add_argument("--question_timeout", type=float, default=180.0)
parser.add_argument("--checkpoint_every", type=int, default=1)
parser.add_argument("--max_samples", type=int, default=None)
parser.add_argument("--no_agent", action="store_true")
args, extra = parser.parse_known_args()
args.extra_main_args = extra
return args
def main() -> None:
args = parse_args()
input_path = _resolve(args.input)
model_bundle = _resolve(args.model_bundle)
output_dir = _resolve(args.output)
if not input_path.exists():
raise FileNotFoundError(f"Input test file not found: {input_path}")
if not model_bundle.exists():
raise FileNotFoundError(f"Model bundle not found: {model_bundle}")
output_dir.mkdir(parents=True, exist_ok=True)
work_dir = output_dir / "_raw_track_a"
if work_dir.exists():
shutil.rmtree(work_dir)
work_dir.mkdir(parents=True)
raw_result_csv = work_dir / "result.csv"
debug_json = work_dir / "debug.json"
traces_json = output_dir / "traces.json"
command = [
sys.executable,
str(ROOT / "main.py"),
"--test_path",
str(input_path),
"--model_bundle",
str(model_bundle),
"--track_b_test",
"",
"--out",
str(raw_result_csv),
"--debug_out",
str(debug_json),
"--traces_out",
str(traces_json),
"--server_url",
args.server_url,
"--model_url",
args.model_url,
"--model_name",
args.model_name,
"--api_key_env",
args.api_key_env,
"--concurrency",
str(args.concurrency),
"--max_steps",
str(args.max_steps),
"--max_tool_calls",
str(args.max_tool_calls),
"--question_timeout",
str(args.question_timeout),
"--checkpoint_every",
str(args.checkpoint_every),
"--no_progress",
]
if args.max_samples is not None:
command.extend(["--max_samples", str(args.max_samples)])
if args.no_agent:
command.append("--no_agent")
command.extend(args.extra_main_args)
env = os.environ.copy()
_ensure_local_vllm_key(env, args.api_key_env)
subprocess.run(command, cwd=str(ROOT), env=env, check=True)
_write_results_csv(raw_result_csv, output_dir / "results.csv")
traces = _normalize_traces(traces_json)
_write_runtime_json(traces, output_dir / "runtime.json")
shutil.rmtree(work_dir)
if __name__ == "__main__":
main()