forked from IQ.Lvbs/IQ.Pilot
141 lines
4.3 KiB
Python
141 lines
4.3 KiB
Python
#!/usr/bin/env python3
|
|
import functools
|
|
import hashlib
|
|
import json
|
|
import re
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
MODELD_DIR = Path(__file__).resolve().parent
|
|
MODELS_DIR = MODELD_DIR / 'models'
|
|
BASEDIR = MODELD_DIR.parents[2]
|
|
METADATA_SCRIPT = MODELD_DIR / 'get_model_metadata.py'
|
|
PYPROJECT = BASEDIR / 'pyproject.toml'
|
|
TINYGRAD_REVISION_FILE = BASEDIR / 'artifacts/package_sources/tinygrad/.iqpilot-revision'
|
|
|
|
MODEL_NAMES = ['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_revision() -> str:
|
|
match = re.search(r'"tinygrad @ git\+https://[^@]+@([0-9a-f]{40})"', PYPROJECT.read_text())
|
|
if match is not None:
|
|
return match.group(1)
|
|
|
|
try:
|
|
revision = TINYGRAD_REVISION_FILE.read_text().strip()
|
|
except OSError as e:
|
|
raise RuntimeError("missing pinned tinygrad revision") from e
|
|
if re.fullmatch(r'[0-9a-f]{40}', revision) is None:
|
|
raise RuntimeError("invalid pinned tinygrad revision")
|
|
return revision
|
|
|
|
|
|
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_revision().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 packaged_prebuilt_matches(model_name: str) -> bool:
|
|
return not (MODELS_DIR / f'{model_name}.onnx').is_file() and outputs_match(model_name)
|
|
|
|
|
|
def verify_prebuilt(model_name: str, flags: str) -> bool:
|
|
if not (MODELS_DIR / f'{model_name}.onnx').is_file():
|
|
return False
|
|
data = _load_checks().get(model_name, {})
|
|
if data.get('signature') != compute_signature(model_name, flags):
|
|
return False
|
|
return outputs_match(model_name)
|
|
|
|
|
|
def verification_details(model_name: str, flags: str) -> list[str]:
|
|
data = _load_checks().get(model_name, {})
|
|
details = [
|
|
f'signature expected={data.get("signature", "missing")} actual={compute_signature(model_name, flags)}',
|
|
f'tinygrad={_tinygrad_revision()}',
|
|
f'onnx={_file_sha256(MODELS_DIR / f"{model_name}.onnx")}',
|
|
f'metadata_script={_file_sha256(METADATA_SCRIPT)}',
|
|
]
|
|
outputs = data.get('outputs', {})
|
|
for fn in _output_names(model_name):
|
|
path = MODELS_DIR / fn
|
|
actual = _file_sha256(path) if path.is_file() else 'missing'
|
|
details.append(f'{fn} expected={outputs.get(fn, "missing")} actual={actual}')
|
|
return details
|
|
|
|
|
|
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:
|
|
valid = verify_prebuilt(name, flags)
|
|
print(f'{name}: {"OK" if valid else "STALE"}')
|
|
if not valid:
|
|
for detail in verification_details(name, flags):
|
|
print(f' {detail}')
|