smoosh smoosh

2026-08-05 16:26:06 +08:00 · 2026-05-26 20:10:16 -05:00
parent 0fec51a5d9
commit e9935624f7
17 changed files with 2406 additions and 1286 deletions
@@ -3,6 +3,7 @@ import argparse
 import codecs
 import os
 import pickle
+import json
 import re
 import shutil
 import subprocess
@@ -10,11 +11,18 @@ import sys

 from pathlib import Path

-
 REPO_ROOT = Path(__file__).resolve().parents[1]
+if str(REPO_ROOT) not in sys.path:
+  sys.path.insert(0, str(REPO_ROOT))
+
+from openpilot.selfdrive.modeld.constants import ModelConstants
+from openpilot.starpilot.common.model_versions import uses_combined_driving_artifacts
+
 DEFAULT_INPUT_ROOT = Path("/data/openpilot/uncompiledmodels")
 DEFAULT_OUTPUT_ROOT = Path("/data/openpilot/compiledmodels")
 COMPILE_SCRIPT = REPO_ROOT / "tinygrad_repo/examples/openpilot/compile3.py"
+COMBINED_COMPILE_SCRIPT = REPO_ROOT / "selfdrive/modeld/compile_modeld.py"
+MODEL_VERSIONS_CACHE = Path("/data/models/.model_versions.json")
 DM_MODEL_KEY = "dm"
 DM_MODEL_NAME = "dmonitoring_model"
 DM_TARGET_ALIASES = {DM_MODEL_KEY, "dmonitoring", DM_MODEL_NAME}
@@ -26,7 +34,33 @@ COMPONENT_ALIASES = {
  "driving_policy": ("driving_policy", "policy"),
  "driving_vision": ("driving_vision", "vision"),
 }
-REQUIRED_COMPONENTS = {"driving_policy", "driving_vision"}
+MEDMODEL_INPUT_SIZE = (512, 256)
+DEFAULT_CAMERA_RESOLUTIONS = (
+  (1928, 1208),
+  (1344, 760),
+)
+
+
+def build_compile_env(*, combined: bool = False) -> dict[str, str]:
+  env = os.environ.copy()
+  existing_pythonpath = env.get("PYTHONPATH", "")
+  env["PYTHONPATH"] = f"{REPO_ROOT}:{existing_pythonpath}" if existing_pythonpath else str(REPO_ROOT)
+
+  numeric_defaults = {
+    "DEBUG": "0",
+    "FLOAT16": "1",
+    "IMAGE": "2",
+    "JIT_BATCH_SIZE": "0",
+    "NOLOCALS": "1",
+  }
+  for key, default in numeric_defaults.items():
+    value = env.get(key)
+    try:
+      int(str(value), 0)
+    except (TypeError, ValueError):
+      env[key] = default
+
+  return env


 def parse_args() -> argparse.Namespace:
@@ -37,6 +71,7 @@ def parse_args() -> argparse.Namespace:
  parser.add_argument("--dm", action="store_true", help="Compile the driver monitoring model into dmonitoring_model_tinygrad.pkl.")
  parser.add_argument("--input-dir", type=Path, default=DEFAULT_INPUT_ROOT, help="Directory containing staged ONNX files. Flat root files like driving_policy.onnx are preferred.")
  parser.add_argument("--output-dir", type=Path, default=DEFAULT_OUTPUT_ROOT, help="Directory for compiled tinygrad pkls and metadata.")
+  parser.add_argument("--version", help="Model version. v16+ uses the combined driving_tinygrad artifact path. If omitted, split-policy staged models default to the combined build.")
  parser.add_argument("--list", action="store_true", help="List detected staged models and exit.")
  parser.add_argument("--force", action="store_true", help="Legacy no-op. Compiled outputs are always cleared before a build.")

@@ -66,14 +101,6 @@ def detect_component(path: Path) -> str | None:
  return None


-def normalize_model_files(model_files: dict[str, Path]) -> dict[str, Path]:
-  normalized = dict(model_files)
-  on_policy_path = normalized.pop("driving_on_policy", None)
-  if on_policy_path is not None and "driving_policy" not in normalized and "driving_off_policy" in normalized:
-    normalized["driving_policy"] = on_policy_path
-  return normalized
-
-
 def find_staged_dm(input_root: Path) -> Path | None:
  if not input_root.is_dir():
    return None
@@ -107,7 +134,6 @@ def find_staged_models(input_root: Path) -> dict[str, dict[str, Path]]:
      component = detect_component(onnx_file)
      if component:
        model_files[component] = onnx_file
-    model_files = normalize_model_files(model_files)
    if model_files:
      found[child.name] = model_files

@@ -137,7 +163,7 @@ def find_staged_models(input_root: Path) -> dict[str, dict[str, Path]]:
      flat_root_files[component] = onnx_file

  if flat_root_files:
-    found["_root"] = normalize_model_files(flat_root_files)
+    found["_root"] = flat_root_files

  return found

@@ -156,7 +182,7 @@ def resolve_model_files(input_root: Path, model_key: str) -> dict[str, Path]:
    component = detect_component(onnx_file)
    if component:
      prefixed_files[component] = onnx_file
-  return normalize_model_files(prefixed_files)
+  return prefixed_files


 def get_metadata_value_by_name(model, name: str):
@@ -190,17 +216,94 @@ def write_metadata(onnx_path: Path, output_path: Path) -> None:


 def compile_component(onnx_path: Path, output_path: Path) -> None:
-  env = os.environ.copy()
-  existing_pythonpath = env.get("PYTHONPATH", "")
-  env["PYTHONPATH"] = f"{REPO_ROOT}:{existing_pythonpath}" if existing_pythonpath else str(REPO_ROOT)
  subprocess.run(
    [sys.executable, str(COMPILE_SCRIPT), str(onnx_path), str(output_path)],
    cwd=REPO_ROOT,
-    env=env,
+    env=build_compile_env(combined=False),
    check=True,
  )


+def compile_combined_model(component_paths: dict[str, Path], output_path: Path) -> None:
+  vision_path = component_paths["driving_vision"]
+  off_policy_path = component_paths["driving_off_policy"]
+  on_policy_path = component_paths.get("driving_on_policy") or component_paths.get("driving_policy")
+  if on_policy_path is None:
+    raise ValueError("Combined compile requires driving_on_policy.onnx (or driving_policy.onnx) alongside driving_off_policy.onnx")
+
+  frame_skip = ModelConstants.MODEL_RUN_FREQ // ModelConstants.MODEL_CONTEXT_FREQ
+  camera_resolutions = [f"{width}x{height}" for width, height in DEFAULT_CAMERA_RESOLUTIONS]
+  subprocess.run(
+    [
+      sys.executable,
+      str(COMBINED_COMPILE_SCRIPT),
+      "--model-size",
+      f"{MEDMODEL_INPUT_SIZE[0]}x{MEDMODEL_INPUT_SIZE[1]}",
+      "--camera-resolutions",
+      *camera_resolutions,
+      "--vision-onnx",
+      str(vision_path),
+      "--off-policy-onnx",
+      str(off_policy_path),
+      "--on-policy-onnx",
+      str(on_policy_path),
+      "--output",
+      str(output_path),
+      "--frame-skip",
+      str(frame_skip),
+    ],
+    cwd=REPO_ROOT,
+    env=build_compile_env(combined=True),
+    check=True,
+  )
+
+
+def infer_model_version(model_key: str, explicit_version: str | None) -> str:
+  if explicit_version:
+    return explicit_version.strip()
+
+  if MODEL_VERSIONS_CACHE.is_file():
+    try:
+      version_map = json.loads(MODEL_VERSIONS_CACHE.read_text())
+      version = version_map.get(model_key)
+      if isinstance(version, str) and version.strip():
+        return version.strip()
+    except Exception:
+      pass
+
+  return ""
+
+
+def should_use_combined_artifacts(model_version: str, model_files: dict[str, Path]) -> bool:
+  if uses_combined_driving_artifacts(model_version):
+    return True
+  if model_version.strip():
+    return False
+
+  has_vision = "driving_vision" in model_files
+  has_off_policy = "driving_off_policy" in model_files
+  has_on_policy = "driving_on_policy" in model_files or "driving_policy" in model_files
+  return has_vision and has_off_policy and has_on_policy
+
+
+def resolve_split_component_inputs(model_files: dict[str, Path]) -> dict[str, Path]:
+  resolved: dict[str, Path] = {}
+
+  vision_path = model_files.get("driving_vision")
+  if vision_path is not None:
+    resolved["driving_vision"] = vision_path
+
+  policy_path = model_files.get("driving_policy") or model_files.get("driving_on_policy")
+  if policy_path is not None:
+    resolved["driving_policy"] = policy_path
+
+  off_policy_path = model_files.get("driving_off_policy")
+  if off_policy_path is not None:
+    resolved["driving_off_policy"] = off_policy_path
+
+  return resolved
+
+
 def clear_existing_outputs(output_dir: Path) -> list[Path]:
  removed = []
  for existing in sorted(output_dir.iterdir()):
@@ -277,18 +380,38 @@ def main() -> int:
      f"or optionally {args.input_dir / model_key}/*.onnx"
    )

-  missing = sorted(REQUIRED_COMPONENTS - set(files))
-  if missing:
-    raise SystemExit(f"Missing required ONNX files for {model_key}: {', '.join(missing)}")
+  model_version = infer_model_version(model_key, args.version)
+  use_combined_artifacts = should_use_combined_artifacts(model_version, files)

  args.output_dir.mkdir(parents=True, exist_ok=True)
-  print(f"Compiling {model_key} from {args.input_dir} -> {args.output_dir}")
+  mode_label = "combined" if use_combined_artifacts else "split"
+  version_label = model_version or ("auto-combined" if use_combined_artifacts else "legacy-default")
+  print(f"Compiling {model_key} ({version_label}, {mode_label}) from {args.input_dir} -> {args.output_dir}")

  removed = clear_existing_outputs(args.output_dir)
  if removed:
    print(f"  cleared {len(removed)} existing output entries")

-  for component, onnx_path in sorted(files.items()):
+  if use_combined_artifacts:
+    required_components = {"driving_vision", "driving_off_policy"}
+    if not (files.get("driving_on_policy") or files.get("driving_policy")):
+      required_components.add("driving_on_policy")
+    missing = sorted(component for component in required_components if component not in files)
+    if missing:
+      raise SystemExit(f"Missing required ONNX files for combined compile of {model_key}: {', '.join(missing)}")
+
+    output_pkl = args.output_dir / f"{model_key}_driving_tinygrad.pkl"
+    compile_combined_model(files, output_pkl)
+    print(f"  saved {output_pkl.name}")
+    print("Done.")
+    return 0
+
+  split_components = resolve_split_component_inputs(files)
+  missing = sorted(component for component in ("driving_policy", "driving_vision") if component not in split_components)
+  if missing:
+    raise SystemExit(f"Missing required ONNX files for {model_key}: {', '.join(missing)}")
+
+  for component, onnx_path in sorted(split_components.items()):
    output_pkl = args.output_dir / f"{model_key}_{component}_tinygrad.pkl"
    output_metadata = args.output_dir / f"{model_key}_{component}_metadata.pkl"