diff --git a/scripts/model_compiler.py b/scripts/model_compiler.py index 4f3b658f5..72ce0ec06 100644 --- a/scripts/model_compiler.py +++ b/scripts/model_compiler.py @@ -301,6 +301,57 @@ def sha256_file(path: Path) -> str: return digest.hexdigest() +def _read_protobuf_varint(source) -> int: + value = 0 + for shift in range(0, 70, 7): + byte = source.read(1) + if not byte: + raise ValueError("unexpected end of file while reading protobuf varint") + value |= (byte[0] & 0x7F) << shift + if not byte[0] & 0x80: + return value + raise ValueError("protobuf varint is too long") + + +def validate_onnx_source(path: Path) -> None: + """Validate the top-level ONNX protobuf without materializing model weights.""" + size = path.stat().st_size + if size == 0: + raise ValueError(f"ONNX source is empty: {path}") + + with open(path, "rb") as source: + if source.read(128).startswith(b"version https://git-lfs.github.com/spec/v1"): + raise ValueError(f"ONNX source is a Git LFS pointer, not model data: {path}") + source.seek(0) + + while source.tell() < size: + tag = _read_protobuf_varint(source) + field, wire_type = tag >> 3, tag & 0x07 + if field == 7: # ModelProto.graph + if wire_type != 2: + raise ValueError(f"ONNX graph has invalid protobuf wire type {wire_type}: {path}") + graph_size = _read_protobuf_varint(source) + remaining = size - source.tell() + if graph_size <= 0: + raise ValueError(f"ONNX graph is empty: {path}") + if graph_size > remaining: + raise ValueError(f"ONNX source is truncated: graph needs {graph_size} bytes but only {remaining} remain: {path}") + return + + if wire_type == 0: + _read_protobuf_varint(source) + elif wire_type == 1: + source.seek(8, os.SEEK_CUR) + elif wire_type == 2: + source.seek(_read_protobuf_varint(source), os.SEEK_CUR) + elif wire_type == 5: + source.seek(4, os.SEEK_CUR) + else: + raise ValueError(f"ONNX source has invalid protobuf wire type {wire_type}: {path}") + + raise ValueError(f"ONNX ModelProto has no graph: {path}") + + def multipart_output_paths(artifact: Path, output_dir: Path | None = None) -> list[Path]: output_dir = output_dir or artifact.parent return [ @@ -552,6 +603,11 @@ def main() -> int: raise SystemExit(f"No staged ONNX files found for {model_key} in {args.input_dir}") input_format = select_input_format(args.input_format, files) + _, source_args = driving_compile_args(files, input_format) + for option, source in zip(source_args[::2], source_args[1::2], strict=True): + source_path = Path(source) + validate_onnx_source(source_path) + print(f" source {option.removeprefix('--')}: {source_path} ({source_path.stat().st_size} bytes)") version = infer_model_version(model_key, args.version) if not version and input_format == "supercombo": version = "v15" diff --git a/selfdrive/modeld/compile_modeld.py b/selfdrive/modeld/compile_modeld.py index 50f3df792..4dacdcd15 100644 --- a/selfdrive/modeld/compile_modeld.py +++ b/selfdrive/modeld/compile_modeld.py @@ -476,6 +476,9 @@ def _parse_size(value): def read_file_chunked_to_disk(path): from openpilot.common.file_chunker import open_file_chunked + if os.path.isfile(path): + return os.fspath(path) + temporary_path = f"{path}.unchunked" try: with open(temporary_path, "wb") as output, open_file_chunked(path) as source: diff --git a/starpilot/assets/tests/test_model_pipeline.py b/starpilot/assets/tests/test_model_pipeline.py index 71e6afa94..3448438c4 100644 --- a/starpilot/assets/tests/test_model_pipeline.py +++ b/starpilot/assets/tests/test_model_pipeline.py @@ -79,7 +79,15 @@ def test_requested_model_id_uses_only_staged_source(tmp_path): } -def test_fat_onnx_is_streamed_to_disk(tmp_path, monkeypatch): +def test_regular_fat_onnx_is_parsed_in_place(tmp_path, monkeypatch): + source = tmp_path / "big_driving_supercombo.onnx" + source.write_bytes(b"model") + monkeypatch.setattr(file_chunker, "open_file_chunked", lambda _: (_ for _ in ()).throw(AssertionError("must not copy"))) + + assert compile_modeld.read_file_chunked_to_disk(source) == str(source) + + +def test_chunked_fat_onnx_is_streamed_to_disk(tmp_path, monkeypatch): payload = b"fat model" * 1024 class StreamingOnly(io.BytesIO): @@ -97,6 +105,30 @@ def test_fat_onnx_is_streamed_to_disk(tmp_path, monkeypatch): Path(staged).unlink(missing_ok=True) +def test_onnx_preflight_accepts_graph_without_reading_weights(tmp_path): + source = tmp_path / "model.onnx" + source.write_bytes(b"\x08\x09\x3a\x02\x12\x00") + + model_compiler.validate_onnx_source(source) + + +def test_onnx_preflight_rejects_empty_truncated_and_lfs_sources(tmp_path): + empty = tmp_path / "empty.onnx" + empty.touch() + truncated = tmp_path / "truncated.onnx" + truncated.write_bytes(b"\x3a\x08bad") + pointer = tmp_path / "pointer.onnx" + pointer.write_text("version https://git-lfs.github.com/spec/v1\n") + + for source, message in ((empty, "empty"), (truncated, "truncated"), (pointer, "Git LFS pointer")): + try: + model_compiler.validate_onnx_source(source) + except ValueError as error: + assert message in str(error) + else: + raise AssertionError(f"{source} should have failed validation") + + def test_dropbox_urls_are_direct_downloads(): url = "https://www.dropbox.com/scl/fi/id/model.pkl?rlkey=key&st=value&dl=0" normalized = download_functions.normalize_download_url(url)