mirror of
https://gitlvb.teallvbs.xyz/IQ.Lvbs/IQ.Pilot.git
synced 2026-07-19 14:42:06 +08:00
112 lines
3.3 KiB
Python
112 lines
3.3 KiB
Python
#!/usr/bin/env python3
|
|
import functools
|
|
import hashlib
|
|
import json
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
MODELD_DIR = Path(__file__).resolve().parent
|
|
MODELS_DIR = MODELD_DIR / 'models'
|
|
BASEDIR = MODELD_DIR.parents[1]
|
|
TINYGRAD_DIR = BASEDIR / 'tinygrad_repo'
|
|
METADATA_SCRIPT = MODELD_DIR / 'get_model_metadata.py'
|
|
|
|
MODEL_NAMES = ['driving_vision', 'driving_policy', 'dmonitoring_model']
|
|
|
|
|
|
def _hash_file(h, path: Path) -> None:
|
|
with open(path, 'rb') as f:
|
|
for chunk in iter(lambda: f.read(1024 * 1024), b''):
|
|
h.update(chunk)
|
|
|
|
|
|
def _file_sha256(path: Path) -> str:
|
|
h = hashlib.sha256()
|
|
_hash_file(h, path)
|
|
return h.hexdigest()
|
|
|
|
|
|
@functools.lru_cache(maxsize=1)
|
|
def _tinygrad_digest() -> str:
|
|
h = hashlib.sha256()
|
|
# mirror the SConscript glob: no hidden files/dirs (also excludes .git, which differs per clone), no pycache
|
|
def _included(p: Path) -> bool:
|
|
rel = p.relative_to(TINYGRAD_DIR).parts
|
|
return p.is_file() and '__pycache__' not in rel and not any(part.startswith('.') for part in rel)
|
|
files = sorted(p for p in TINYGRAD_DIR.rglob('*') if _included(p))
|
|
for p in files:
|
|
h.update(str(p.relative_to(BASEDIR)).encode())
|
|
_hash_file(h, p)
|
|
return h.hexdigest()
|
|
|
|
|
|
CHECK_PATH = MODELS_DIR / 'prebuilt_check.json'
|
|
|
|
|
|
def _load_checks() -> dict:
|
|
try:
|
|
data = json.loads(CHECK_PATH.read_text())
|
|
except (json.JSONDecodeError, OSError):
|
|
return {}
|
|
return data if isinstance(data, dict) else {}
|
|
|
|
|
|
def _output_names(model_name: str) -> list[str]:
|
|
return [f'{model_name}_tinygrad.pkl', f'{model_name}_metadata.pkl']
|
|
|
|
|
|
def compute_signature(model_name: str, flags: str) -> str:
|
|
h = hashlib.sha256()
|
|
h.update(flags.encode())
|
|
h.update(_tinygrad_digest().encode())
|
|
_hash_file(h, METADATA_SCRIPT)
|
|
_hash_file(h, MODELS_DIR / f'{model_name}.onnx')
|
|
return h.hexdigest()
|
|
|
|
|
|
def outputs_match(model_name: str) -> bool:
|
|
"""The committed artifacts on disk are exactly the ones the check file pins."""
|
|
data = _load_checks().get(model_name, {})
|
|
outputs = data.get('outputs', {})
|
|
if set(outputs) != set(_output_names(model_name)):
|
|
return False
|
|
for fn, expected in outputs.items():
|
|
p = MODELS_DIR / fn
|
|
if not p.is_file() or _file_sha256(p) != expected:
|
|
return False
|
|
return True
|
|
|
|
|
|
def verify_prebuilt(model_name: str, flags: str) -> bool:
|
|
data = _load_checks().get(model_name, {})
|
|
if data.get('signature') != compute_signature(model_name, flags):
|
|
return False
|
|
return outputs_match(model_name)
|
|
|
|
|
|
def write_check(model_name: str, flags: str) -> None:
|
|
outputs = {}
|
|
for fn in _output_names(model_name):
|
|
p = MODELS_DIR / fn
|
|
if not p.is_file():
|
|
raise FileNotFoundError(f'missing build output: {p}')
|
|
outputs[fn] = _file_sha256(p)
|
|
checks = _load_checks()
|
|
checks[model_name] = {'signature': compute_signature(model_name, flags), 'outputs': outputs}
|
|
CHECK_PATH.write_text(json.dumps(checks, indent=2, sort_keys=True) + '\n')
|
|
|
|
|
|
def _larch64_flags() -> str:
|
|
return "DEV=QCOM IMAGE=2 FLOAT16=1 NOLOCALS=1 JIT_BATCH_SIZE=0 OPENPILOT_HACKS=1"
|
|
|
|
|
|
if __name__ == '__main__':
|
|
mode = sys.argv[1] if len(sys.argv) > 1 else 'verify'
|
|
flags = _larch64_flags()
|
|
for name in MODEL_NAMES:
|
|
if mode == 'write':
|
|
write_check(name, flags)
|
|
print(f'{name}: check written')
|
|
else:
|
|
print(f'{name}: {"OK" if verify_prebuilt(name, flags) else "STALE"}')
|