diff --git a/.github/workflows/sunnypilot-build-model.yaml b/.github/workflows/sunnypilot-build-model.yaml index 9414a7fd0..ff09489b9 100644 --- a/.github/workflows/sunnypilot-build-model.yaml +++ b/.github/workflows/sunnypilot-build-model.yaml @@ -176,6 +176,15 @@ jobs: DEV=QCOM FLOAT16=1 NOLOCALS=1 JIT_BATCH_SIZE=0 python3 "${{ env.MODELS_DIR }}/../get_model_metadata.py" "$onnx_file" || true done + - name: Validate Model Outputs + run: | + source /etc/profile + export UV_PROJECT_ENVIRONMENT=${HOME}/venv + export VIRTUAL_ENV=$UV_PROJECT_ENVIRONMENT + python3 "${{ github.workspace }}/release/ci/model_generator.py" \ + --validate-only \ + --model-dir "${{ env.MODELS_DIR }}" + - name: Prepare Output run: | sudo rm -rf ${{ env.OUTPUT_DIR }} diff --git a/release/ci/model_generator.py b/release/ci/model_generator.py index 96352254b..da6b93303 100755 --- a/release/ci/model_generator.py +++ b/release/ci/model_generator.py @@ -1,4 +1,5 @@ import os +import pickle import sys import hashlib import json @@ -6,6 +7,41 @@ import re from pathlib import Path from datetime import datetime, UTC +REQUIRED_OUTPUT_KEYS = frozenset({ + "plan", + "lane_lines", + "road_edges", + "lead", + "desire_state", + "desire_pred", + "meta", + "lead_prob", + "lane_lines_prob", + "pose", + "wide_from_device_euler", + "road_transform", + "hidden_state", +}) +OPTIONAL_OUTPUT_KEYS = frozenset({ + "planplus", + "sim_pose", + "desired_curvature", +}) + + +def validate_model_outputs(metadata_paths: list[Path]) -> None: + combined_keys: set[str] = set() + for path in metadata_paths: + with open(path, "rb") as f: + metadata = pickle.load(f) + combined_keys.update(metadata.get("output_slices", {}).keys()) + missing = REQUIRED_OUTPUT_KEYS - combined_keys + if missing: + raise ValueError(f"Combined model metadata is missing required output keys: {sorted(missing)}") + detected_optional = sorted(OPTIONAL_OUTPUT_KEYS & combined_keys) + if detected_optional: + print(f"Optional output keys detected: {detected_optional}") + def create_short_name(full_name): # Remove parentheses and extract alphanumeric words @@ -124,9 +160,19 @@ if __name__ == "__main__": parser.add_argument("--output-dir", default="./output", help="Output directory for metadata") parser.add_argument("--custom-name", help="Custom display name for the model") parser.add_argument("--is-20hz", action="store_true", help="Whether this is a 20Hz model") + parser.add_argument("--validate-only", action="store_true") parser.add_argument("--upstream-branch", default="unknown", help="Upstream branch name") args = parser.parse_args() + if args.validate_only: + metadata_paths = glob.glob(os.path.join(args.model_dir, "*_metadata.pkl")) + if not metadata_paths: + print(f"No metadata files found in {args.model_dir}", file=sys.stderr) + sys.exit(1) + validate_model_outputs([Path(p) for p in metadata_paths]) + print(f"Validated {len(metadata_paths)} metadata files successfully.") + sys.exit(0) + # Find all ONNX files in the given directory model_paths = glob.glob(os.path.join(args.model_dir, "*.onnx")) if not model_paths: