StarPilot/scripts/model_rebuild.py

#!/usr/bin/env python3
"""Rebuild a model generation on an explicitly selected compiler device."""

from __future__ import annotations

import argparse
import hashlib
import json
from pathlib import Path
import shlex
import subprocess
import sys
import time
from urllib.request import urlopen
from urllib.error import HTTPError

REPO = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(REPO))
from openpilot.starpilot.models.catalog import ARTIFACT_ABI, CATALOG_PATH, COMPILER_REVISION, GENERATION
from openpilot.starpilot.models.manager import atomic_json, validate_manifest

BUCKET = "StarPilot-Driving/StarPilot-Resources"
REMOTE_WORK = "/data/starpilot-v26-rebuild"
CHUNK_BYTES = 45 * 1024 * 1024


def digest(path: Path) -> str:
  with path.open("rb") as stream:
    return hashlib.file_digest(stream, "sha256").hexdigest()


def run(args: list[str], *, output=None) -> None:
  subprocess.run(args, check=True, stdout=output, stderr=subprocess.STDOUT if output is not None else None)


def remote(device: str, command: list[str], *, output=None) -> None:
  run(["ssh", "-o", "BatchMode=yes", "-o", "ConnectTimeout=10", "-o", "ServerAliveInterval=30",
       device, shlex.join(command)], output=output)


def inventory() -> list[dict]:
  result = []
  url = f"https://huggingface.co/api/buckets/{BUCKET}/tree?recursive=true&limit=1000"
  while url:
    with urlopen(url, timeout=40) as response:
      result.extend(json.load(response))
      links = response.headers.get("Link", "")
    url = next((p.split(">")[0].strip("< ") for p in links.split(",") if 'rel="next"' in p), "")
  return result


def source_plan(model_id: str, source_map: dict, objects: list[dict]) -> list[dict]:
  source_id = source_map[model_id]["source_id"]
  files = [obj for obj in objects if obj["type"] == "file" and obj["path"].startswith(f"onnx/{source_id}/") and
           obj["path"].endswith(".onnx")]
  if not files:
    raise ValueError(f"No archived ONNX source for {model_id} ({source_id})")
  return files


def prepare_sources(model_id: str, files: list[dict], workspace: Path, hf: str) -> Path:
  directory = workspace / "onnx" / model_id
  directory.mkdir(parents=True, exist_ok=True)
  for obj in files:
    name = Path(obj["path"]).name
    if "driving_" not in name:
      raise ValueError(f"Unrecognized model source filename: {name}")
    name = model_id + "_" + name[name.index("driving_"):]
    destination = directory / name
    if not destination.is_file() or destination.stat().st_size != obj["size"]:
      run([hf, "buckets", "cp", f"hf://buckets/{BUCKET}/{obj['path']}", str(destination), "--format", "quiet"])
    if destination.stat().st_size != obj["size"]:
      raise ValueError("Source download size mismatch")
  return directory


def split_artifact(artifact: Path, directory: Path) -> tuple[list[Path], int]:
  directory.mkdir(parents=True, exist_ok=True)
  count = (artifact.stat().st_size + CHUNK_BYTES - 1) // CHUNK_BYTES
  if count == 1:
    return [artifact], 0
  parts = []
  with artifact.open("rb") as stream:
    for index in range(1, count + 1):
      part = directory / f"{artifact.name}.chunk{index:02d}of{count:02d}"
      part.write_bytes(stream.read(CHUNK_BYTES))
      parts.append(part)
  return parts, count


def publish_artifact(model_id: str, artifact: Path, receipt: dict, workspace: Path, hf: str) -> dict:
  if receipt["compiler_revision"] != COMPILER_REVISION or receipt["artifact_abi"] != ARTIFACT_ABI or digest(artifact) != receipt["artifact_sha256"]:
    raise ValueError("Compiled artifact identity changed")
  parts, count = split_artifact(artifact, workspace / "release" / model_id)
  for part in parts:
    run([hf, "buckets", "cp", str(part), f"hf://buckets/{BUCKET}/models/{GENERATION}/{model_id}/{part.name}", "--format", "quiet"])
  if count:
    marker = workspace / "release" / model_id / f"{artifact.name}.chunkmanifest"
    marker.write_text(str(count))
    run([hf, "buckets", "cp", str(marker), f"hf://buckets/{BUCKET}/models/{GENERATION}/{model_id}/{marker.name}", "--format", "quiet"])
  return {"artifact_format": ARTIFACT_ABI, "artifact_size": artifact.stat().st_size,
          "artifact_sha256": receipt["artifact_sha256"], "artifact_chunk_count": count,
          "source_sha256": receipt["source_sha256"]}


def compile_one(model: dict, args, source_map: dict, objects: list[dict]) -> dict:
  mid = model["id"]
  work = args.workspace
  result_path = work / "results" / f"{mid}.json"
  output_dir = work / "compiled" / mid
  artifact = output_dir / f"{mid}_driving_tinygrad.pkl"
  if result_path.is_file():
    previous = json.loads(result_path.read_text())
    if (previous.get("status") == "passed" and previous.get("compiler_revision") == COMPILER_REVISION and artifact.is_file() and
        digest(artifact) == previous["artifact_sha256"] and (not args.publish or previous.get("published"))):
      return previous
  sources = prepare_sources(mid, source_plan(mid, source_map, objects), work, args.hf)
  remote_dir = f"{REMOTE_WORK}/work/{mid}"
  remote(args.device, ["mkdir", "-p", f"{remote_dir}/onnx", f"{remote_dir}/compiled"])
  run(["rsync", "-az", str(sources) + "/", f"{args.device}:{remote_dir}/onnx/"])
  command = ["/usr/local/venv/bin/python", "/data/openpilot/scripts/model_compiler.py", "--model", mid,
             "--version", model["version"], "--input-format", source_map[mid]["input_format"],
             "--input-dir", f"{remote_dir}/onnx", "--output-dir", f"{remote_dir}/compiled"]
  if model.get("uses_external_gpu"):
    command.append("--gpu")
  with (work / "logs" / f"{mid}.log").open("ab") as log:
    remote(args.device, ["env", "PYTHONUNBUFFERED=1", *command], output=log)
    validate = ["/usr/local/venv/bin/python", "/data/openpilot/scripts/validate_driving_model.py",
                f"{remote_dir}/compiled/{artifact.name}", "--version", model["version"],
                "--source-dir", f"{remote_dir}/onnx", "--report", f"{remote_dir}/compiled/validation.json", "--max-p95-ms", "50"]
    if model.get("uses_external_gpu"):
      validate.append("--gpu")
    remote(args.device, validate, output=log)
  output_dir.mkdir(parents=True, exist_ok=True)
  run(["rsync", "-az", f"{args.device}:{remote_dir}/compiled/", str(output_dir) + "/"])
  receipt = json.loads(artifact.with_suffix(".build.json").read_text())
  validation = json.loads((output_dir / "validation.json").read_text())
  if (digest(artifact) != receipt["artifact_sha256"] or not validation.get("passed") or
      validation.get("artifact_sha256") != receipt["artifact_sha256"]):
    raise ValueError("Compiled model failed runtime validation")
  result = {**receipt, "status": "passed", "published": False, "validation": validation}
  if args.publish:
    result["release"] = publish_artifact(mid, artifact, receipt, work, args.hf)
    result["published"] = True
  atomic_json(result_path, result)
  # Only this invocation's sources and output are removed after the verified host copy exists.
  remote(args.device, ["/usr/local/venv/bin/python", "-c", "import shutil,sys; shutil.rmtree(sys.argv[1])", remote_dir])
  return result


def publish_manifest(workspace: Path, hf: str) -> None:
  manifest = json.loads(CATALOG_PATH.read_text())
  try:
    with urlopen(f"https://huggingface.co/buckets/{BUCKET}/resolve/manifests/model_names_{GENERATION}.json", timeout=30) as response:
      live = json.load(response)
    existing = validate_manifest(live)
    for row in manifest["models"]:
      if row["id"] in existing:
        row.update(existing[row["id"]])
  except HTTPError as error:
    if error.code != 404:
      raise
  for model in manifest["models"]:
    path = workspace / "results" / f"{model['id']}.json"
    if path.is_file():
      result = json.loads(path.read_text())
      if result.get("status") == "passed" and result.get("published") and result.get("compiler_revision") == COMPILER_REVISION:
        model.update(result["release"])
  path = workspace / "manifests" / f"model_names_{GENERATION}.json"
  atomic_json(path, manifest)
  run([hf, "buckets", "cp", str(path), f"hf://buckets/{BUCKET}/manifests/{path.name}", "--format", "quiet"])


def main():
  parser = argparse.ArgumentParser(description=__doc__)
  parser.add_argument("--device", required=True, help="Explicit SSH target, for example comma@comma-device.local")
  parser.add_argument("--workspace", required=True, type=Path)
  parser.add_argument("--models", nargs="+", help="Stable model IDs; defaults to the full catalog, favorites first")
  parser.add_argument("--hf", default=str(Path.home()/".local/bin/hf"))
  parser.add_argument("--publish", action="store_true", help="Upload validated v26 artifacts and manifest to Hugging Face")
  args = parser.parse_args()
  args.workspace = args.workspace.resolve()
  for directory in ("results", "logs", "manifests"):
    (args.workspace / directory).mkdir(parents=True, exist_ok=True)
  rows = json.loads(CATALOG_PATH.read_text())["models"]
  source_map = json.loads((REPO / "scripts/model_source_map_v26.json").read_text())
  if args.models:
    by_id = {row["id"]: row for row in rows}
    rows = [by_id[mid] for mid in args.models]
  else:
    rows.sort(key=lambda row: (not row.get("community_favorite", False), row["id"] != "cinquev3"))
  objects = inventory()
  atomic_json(args.workspace / "source-inventory.json", {"objects": objects})
  failures = []
  for model in rows:
    mid = model["id"]
    print(f"START {mid}", flush=True)
    started = time.monotonic()
    try:
      compile_one(model, args, source_map, objects)
      if args.publish:
        publish_manifest(args.workspace, args.hf)
      print(f"PASS {mid} {time.monotonic()-started:.0f}s", flush=True)
    except (OSError, ValueError, subprocess.CalledProcessError) as error:
      failures.append(mid)
      atomic_json(args.workspace / "results" / f"{mid}.failure.json", {"model_id": mid, "error": str(error)})
      print(f"FAIL {mid}: {error}", flush=True)
  return int(bool(failures))


if __name__ == "__main__":
  raise SystemExit(main())