run_coverhunter_finetune.py 2.58 KB
#!/usr/bin/env python3
import argparse
import json
import subprocess
from datetime import datetime
from pathlib import Path


DEFAULT_PYTHON = "/usr/local/miniconda3/bin/python"


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--python", default=DEFAULT_PYTHON)
    parser.add_argument("--config", default="configs/coverhunter_finetune_4gb.yaml")
    parser.add_argument("--data", required=True)
    parser.add_argument("--output-root", default="data/training_runs")
    parser.add_argument("--run-name", default=None)
    parser.add_argument("--noise-root", action="append", default=[])
    parser.add_argument("--device", default="auto")
    parser.add_argument("--segment-strategy", default="hybrid")
    parser.add_argument("--resume", default=None)
    parser.add_argument("--dry-run", action="store_true")
    args = parser.parse_args()

    timestamp = datetime.utcnow().strftime("%Y%m%dT%H%M%SZ")
    run_name = args.run_name or f"coverhunter_finetune_{timestamp}"
    run_dir = Path(args.output_root) / run_name
    run_dir.mkdir(parents=True, exist_ok=True)

    command = [
        args.python,
        "train.py",
        "--config",
        args.config,
        "--data",
        args.data,
        "--output",
        str(run_dir),
        "--device",
        args.device,
        "--segment-strategy",
        args.segment_strategy,
    ]
    if args.resume:
        command.extend(["--resume", args.resume])
    if args.dry_run:
        command.append("--dry-run")
    for noise_root in args.noise_root:
        command.extend(["--noise-root", noise_root])

    metadata = {
        "run_name": run_name,
        "created_at": datetime.utcnow().isoformat() + "Z",
        "python": args.python,
        "command": command,
        "config": args.config,
        "data": args.data,
        "noise_roots": args.noise_root,
        "run_dir": str(run_dir),
    }
    with open(run_dir / "run_request.json", "w") as f:
        json.dump(metadata, f, indent=2)

    result = subprocess.run(command, cwd=Path(__file__).resolve().parents[1], text=True, capture_output=True)
    (run_dir / "stdout.log").write_text(result.stdout)
    (run_dir / "stderr.log").write_text(result.stderr)
    summary = {
        **metadata,
        "returncode": result.returncode,
        "completed_at": datetime.utcnow().isoformat() + "Z",
        "artifacts": sorted(path.name for path in run_dir.iterdir()),
    }
    with open(run_dir / "run_summary.json", "w") as f:
        json.dump(summary, f, indent=2)
    if result.returncode != 0:
        raise SystemExit(result.returncode)


if __name__ == "__main__":
    main()