diff --git a/openpilot/selfdrive/modeld/SConscript b/openpilot/selfdrive/modeld/SConscript index 2aac90fe6..c8892ef91 100644 --- a/openpilot/selfdrive/modeld/SConscript +++ b/openpilot/selfdrive/modeld/SConscript @@ -1,14 +1,20 @@ +import ctypes.util import glob import os -import time from SCons.Script import Action, Value from openpilot.common.transformations.camera import _ar_ox_fisheye, _os_fisheye from openpilot.common.transformations.model import MEDMODEL_INPUT_SIZE, DM_INPUT_SIZE -from openpilot.selfdrive.modeld.helpers import chestnut_present from openpilot.system.camerad.cameras.nv12_info import get_nv12_info Import('env', 'arch') + +# TODO: add GCC support to tinygrad? +# tinygrad's CPU:LLVM backend requires libLLVM, skip the build on PC if it's not available +if arch not in ('comma_arm64', 'Darwin') and ctypes.util.find_library('LLVM') is None: + print("\033[93mmodeld: LLVM not found, skipping tinygrad model compilation. Install LLVM to compile models.\033[0m") + Return() + lenv = env.Clone() lenv.PrependENVPath('PYTHONPATH', Dir('#tinygrad_repo').abspath) @@ -24,38 +30,18 @@ else: # JIT=2 disables graph batching, which produces incorrect outputs after buffers change. tg_flags = 'DEV=METAL JIT=2' if arch == 'Darwin' else 'DEV=CPU:LLVM' -CHESTNUT = chestnut_present() -if CHESTNUT: - chestnut_tg_flags = 'DEBUG=1 DEV=USB+AMD:LLVM FRAME_DEV=CPU FLOAT16=1 JIT_BATCH_SIZE=0 GMMU=0 TC_OPT=2 TC_MIN_GLOBALS=32' - # the USB+AMD GPU takes an exclusive flock; serialize all targets that touch it - chestnut_lock = File("models/.chestnut.lock").abspath - # tinygrad calls brew which needs a $HOME in the env mac_brew_string = f'HOME={os.path.expanduser("~")}' if arch == 'Darwin' else '' warp_deps = [File("#openpilot/system/camerad/cameras/nv12_info.py")] compiler = Dir('#tinygrad_repo/examples/openpilot').abspath -# CPU 7 is isolated with isolcpus on AGNOS, so explicitly pin the compiler to it. -taskset = 'taskset -c 7 ' if arch == 'comma_arm64' else '' - -def chestnut_action(command): - def do_compile(target, source, env): - from openpilot.system.hardware.chestnut.flash import link_up - # chestnut can enumerate before its PCIe link is up due to varying 12V power behavior across cars - for _ in range(10): - if link_up(): - break - time.sleep(1) - else: - print("Chestnut not ready, skipping warp build") - return - return env.Execute(command) - return Action(do_compile, " [CHESTNUT] $TARGET") +# build on all cores +taskset = 'taskset -c 0-7 ' if arch == 'comma_arm64' else '' def compile_model(onnx_path, pkl_path): onnx_path, target_pkl_path = File(onnx_path).abspath, File(pkl_path).abspath cmd = (f'{tg_flags} {mac_brew_string} {taskset}python3 "{compiler}/compile_onnx.py" ' - f'"{onnx_path}" "{target_pkl_path}" --device-input "*" --out-of-band --benchmark-runs 1') + f'"{onnx_path}" "{target_pkl_path}" --out-of-band --benchmark-runs 1') lenv.Command( target_pkl_path, tinygrad_files + [onnx_path, Value(cmd)], @@ -66,19 +52,14 @@ compile_model('models/dmonitoring_model.onnx', 'models/dmonitoring_model_tinygra compile_model('models/driving_supercombo.onnx', 'models/driving_tinygrad.pkl') model_w, model_h = MEDMODEL_INPUT_SIZE -for chestnut in [False, True] if CHESTNUT else [False]: - file_prefix, cmd_flags = ('big_', chestnut_tg_flags) if chestnut else ('', tg_flags) - for cam_w, cam_h in camera_configs: - warp_pkl_path = File(f"models/{file_prefix}driving_warp_{cam_w}x{cam_h}_tinygrad.pkl").abspath - stride, y_height, uv_height, _ = get_nv12_info(cam_w, cam_h) - cmd = (f'{cmd_flags} {mac_brew_string} {taskset}python3 "{compiler}/compile_warp.py" ' - f'--frame {cam_w},{cam_h},{stride},{y_height},{uv_height},{stride * (y_height + uv_height)} ' - f'--warp-to {model_w}x{model_h} --layout yuv420 --frames 2 ' - f'--output {warp_pkl_path}') - action = chestnut_action(cmd) if chestnut else cmd - node = lenv.Command(warp_pkl_path, tinygrad_files + warp_deps + [Value(cmd)], action) - if chestnut: - lenv.SideEffect(chestnut_lock, node) +for cam_w, cam_h in camera_configs: + warp_pkl_path = File(f"models/driving_warp_{cam_w}x{cam_h}_tinygrad.pkl").abspath + stride, y_height, uv_height, _ = get_nv12_info(cam_w, cam_h) + cmd = (f'{tg_flags} {mac_brew_string} {taskset}python3 "{compiler}/compile_warp.py" ' + f'--frame {cam_w},{cam_h},{stride},{y_height},{uv_height},{stride * (y_height + uv_height)} ' + f'--warp-to {model_w}x{model_h} --layout yuv420 --frames 2 ' + f'--output {warp_pkl_path}') + lenv.Command(warp_pkl_path, tinygrad_files + warp_deps + [Value(cmd)], cmd) dm_w, dm_h = DM_INPUT_SIZE for cam_w, cam_h in camera_configs: diff --git a/openpilot/selfdrive/modeld/dmonitoringmodeld.py b/openpilot/selfdrive/modeld/dmonitoringmodeld.py index 06d0d2e5e..5a6e9216a 100755 --- a/openpilot/selfdrive/modeld/dmonitoringmodeld.py +++ b/openpilot/selfdrive/modeld/dmonitoringmodeld.py @@ -16,7 +16,6 @@ from openpilot.common.realtime import config_realtime_process from openpilot.common.transformations.model import dmonitoringmodel_intrinsics from openpilot.common.transformations.camera import _ar_ox_fisheye, _os_fisheye from openpilot.system.camerad.cameras.nv12_info import get_nv12_info -from openpilot.common.file_chunker import open_file_chunked from openpilot.selfdrive.modeld.parse_model_outputs import sigmoid, safe_exp SEND_RAW_PRED = os.getenv('SEND_RAW_PRED') @@ -28,7 +27,7 @@ class ModelState: output: np.ndarray def __init__(self, cam_w: int, cam_h: int): - jits = load_oob(open_file_chunked(MODEL_PKL_PATH)) + jits = load_oob(MODEL_PKL_PATH) self.DEV = jits['input_specs']['input_img'][2] self.input_shapes = jits['metadata']['input_shapes'] self.output_slices = pickle.loads(base64.b64decode(jits['metadata']['metadata']['output_slices'])) diff --git a/openpilot/selfdrive/modeld/helpers.py b/openpilot/selfdrive/modeld/helpers.py index 5e7f413c2..9ac772e69 100644 --- a/openpilot/selfdrive/modeld/helpers.py +++ b/openpilot/selfdrive/modeld/helpers.py @@ -1,9 +1,8 @@ -import io -import pickle -import struct +import sys +import time from pathlib import Path -from openpilot.common.file_chunker import get_manifest_path +from openpilot.common.hardware import AGNOS from openpilot.common.hardware.usb import CHESTNUT_USB_PRODUCT, USB_DEVICES_PATH, is_chestnut_usb_id MODELS_DIR = Path(__file__).resolve().parent / 'models' @@ -13,15 +12,12 @@ def modeld_pkl_path(chestnut: bool): prefix = 'big_' if chestnut else '' return MODELS_DIR / f'{prefix}driving_tinygrad.pkl' -def load_oob(f): - opcodes = f.read(struct.unpack(' bool: for d in USB_DEVICES_PATH.glob("*"): @@ -34,7 +30,15 @@ def chestnut_present() -> bool: pass return False +def wait_for_chestnut(timeout: float = 10.) -> None: + # the cable is detected before chestnut enumerates + st = time.monotonic() + while not chestnut_present(): + if time.monotonic() - st > timeout: + raise TimeoutError("chestnut did not enumerate") + time.sleep(0.1) + def chestnut_compiled() -> bool: path = modeld_pkl_path(chestnut=True) - return (path.is_file() or Path(get_manifest_path(path)).is_file()) and all( + return path.is_file() and all( (MODELS_DIR / f'big_driving_warp_{size}_tinygrad.pkl').is_file() for size in ('1344x760', '1928x1208')) diff --git a/openpilot/selfdrive/modeld/modeld.py b/openpilot/selfdrive/modeld/modeld.py index 57e19caac..27693bce6 100755 --- a/openpilot/selfdrive/modeld/modeld.py +++ b/openpilot/selfdrive/modeld/modeld.py @@ -2,11 +2,14 @@ from collections.abc import Callable import base64 import ctypes +import gc from functools import cached_property import os os.environ['GMMU'] = '0' # for chestnut fast loading, noop for qcom +os.environ.setdefault('AM_POWER_LIMIT', '100') from tinygrad.device import Buffer, Device from tinygrad.dtype import DType, dtypes +from tinygrad.engine.realize import lower_and_compile from tinygrad.tensor import Tensor from tinygrad.helpers import round_up from tinygrad.uop.ops import UOp @@ -25,6 +28,7 @@ from msgq.visionipc import VisionIpcClient, VisionBuf from opendbc.car.car_helpers import get_demo_car_params from openpilot.common.swaglog import cloudlog from openpilot.common.params import Params +from openpilot.common.hardware.usb import cable_connected from openpilot.common.filter_simple import FirstOrderFilter from openpilot.common.realtime import config_realtime_process, DT_MDL from openpilot.common.transformations.camera import DEVICE_CAMERAS @@ -34,9 +38,8 @@ from openpilot.selfdrive.controls.lib.desire_helper import DesireHelper from openpilot.selfdrive.controls.lib.drive_helpers import get_accel_from_plan, should_stop, smooth_value, get_curvature_from_plan from openpilot.selfdrive.modeld.parse_model_outputs import Parser from openpilot.selfdrive.modeld.fill_model_msg import fill_model_msg, fill_driving_model_data, fill_pose_msg, PublishState -from openpilot.common.file_chunker import open_file_chunked from openpilot.selfdrive.modeld.constants import ModelConstants, Plan -from openpilot.selfdrive.modeld.helpers import MODELS_DIR, chestnut_present, chestnut_compiled, modeld_pkl_path, load_oob +from openpilot.selfdrive.modeld.helpers import MODELS_DIR, chestnut_present, chestnut_compiled, modeld_pkl_path, load_oob, wait_for_chestnut from openpilot.sunnypilot.livedelay.helpers import get_lat_delay from openpilot.sunnypilot.modeld_v2.modeld_base import ModelStateBase from openpilot.sunnypilot.selfdrive.controls.lib.relc import RoadEdgeLaneChangeController @@ -134,8 +137,8 @@ class FrameMeta: def input_view(buffer: Buffer, shape: tuple[int, ...], dtype: DType, offset: int) -> Tensor: - view = buffer.view(math.prod(shape), dtype, offset).ensure_allocated() - return Tensor(UOp.from_buffer(view)).reshape(shape) + view = buffer.view(math.prod(shape) * dtype.itemsize, offset).ensure_allocated() + return Tensor(UOp.from_buffer(view, dtype)).reshape(shape) class ModelState(ModelStateBase): @@ -143,7 +146,7 @@ class ModelState(ModelStateBase): def __init__(self, cam_w: int, cam_h: int, chestnut: bool): ModelStateBase.__init__(self) - jits = load_oob(open_file_chunked(modeld_pkl_path(chestnut))) + jits = load_oob(modeld_pkl_path(chestnut), chestnut) self.model_device = jits['input_specs']['new_img'][2] self.input_shapes = {name: (shape, np.dtype(dtype)) for name, (shape, dtype, _) in jits['input_specs'].items()} self.state_pairs = {name: f'next_{name}' for name in self.input_shapes if f'next_{name}' in jits['metadata']['output_shapes']} @@ -158,7 +161,9 @@ class ModelState(ModelStateBase): self.pack_inputs() with open(MODELS_DIR / f'{"big_" if chestnut else ""}driving_warp_{cam_w}x{cam_h}_tinygrad.pkl', 'rb') as f: self.run_warp = pickle.load(f)['run'] + self.run_warp.captured._linear = lower_and_compile(self.run_warp.captured._linear) self.run_model = jits['run'] + self.run_model.captured._linear = lower_and_compile(self.run_model.captured._linear) self.outputs = {name: Tensor(np.zeros(shape, dtype=dtype), device=device).realize() for name, (shape, dtype, device) in jits['output_specs'].items()} for name, next_name in self.state_pairs.items(): state = self.input_queues[name] @@ -203,6 +208,8 @@ class ModelState(ModelStateBase): self.input_device.copy_from(self.input_host) self.input_queues['new_img'] = self.run_warp(**self.warp_inputs) self.run_model(output_buffers=self.outputs, **self.input_queues) + if self.chestnut and (err := Device[self.outputs['outputs'].device].error_state.host.view(fmt='q')[0]): + raise RuntimeError(f"chestnut failed with {err}") if after_enqueue is not None: after_enqueue() model_output = self.outputs['outputs'].numpy()[0] @@ -228,14 +235,15 @@ class ModelState(ModelStateBase): def main(demo=False): cloudlog.warning("modeld init") - CHESTNUT = chestnut_present() and chestnut_compiled() + CHESTNUT = chestnut_compiled() and (chestnut_present() or cable_connected()) if CHESTNUT: - os.environ['HCQDEV_WAIT_TIMEOUT_MS'] = '3000' + from tinygrad.runtime.ops_amd import AMDDevice + AMDDevice.wait_timeout_ms = 3000 params = Params() params.put_bool("ChestnutLoading", CHESTNUT) params.remove("ChestnutActive") - config_realtime_process(7, 54) + gc.disable() # visionipc clients while True: @@ -268,6 +276,7 @@ def main(demo=False): def load_big(): nonlocal big_model try: + wait_for_chestnut() m = ModelState(vipc_client_main.width, vipc_client_main.height, True) m.warmup() big_model = m @@ -277,7 +286,13 @@ def main(demo=False): loader.start() loader.join(BIG_MODEL_TIMEOUT) model = big_model - params.put_bool("ChestnutActive", model is not None) + if model is None and not params.get_bool("ChestnutRetry"): + # a crank can cut the chestnut power during the load, retry once with an onroad cycle + cloudlog.error("big model load failed, cycling onroad to retry") + params.put_bool("ChestnutRetry", True) + params.put_bool("OnroadCycleRequested", True) + else: + params.put_bool("ChestnutActive", model is not None) small_model = ModelState(vipc_client_main.width, vipc_client_main.height, False) if model is None or CHESTNUT else None if model is None: @@ -286,10 +301,13 @@ def main(demo=False): params.put_bool("ChestnutLoading", False) cloudlog.warning(f"models loaded in {time.monotonic() - st:.1f}s, modeld starting") + config_realtime_process(7, 54) + # messaging pub_socks = ["modelV2", "drivingModelData", "cameraOdometry", "modelDataV2SP"] + (["chestnutGpuState"] if CHESTNUT else []) pm = PubMaster(pub_socks) - sm = SubMaster(["deviceState", "carState", "narrowRoadCameraState", "extrinsicsCalibration", "driverMonitoringState", "carControl", "lateralDelay"]) + sm = SubMaster(["deviceState", "carState", "narrowRoadCameraState", "extrinsicsCalibration", "driverMonitoringState", + "carControl", "controlsState", "lateralDelay"]) publish_state = PublishState() params = Params() @@ -407,7 +425,7 @@ def main(demo=False): run_count % round(ModelConstants.MODEL_RUN_FREQ / SERVICE_LIST['chestnutGpuState'].frequency) == 0) model_output = model.run(bufs, transforms, inputs, chestnut_state.send if send_chestnut else None) except Exception: - if not params.get_bool("ChestnutActive"): + if not model.chestnut: raise # fallback to small model cloudlog.exception("big model failed, fall back to small") @@ -427,6 +445,8 @@ def main(demo=False): posenet_send = messaging.new_message('cameraOdometry') mdv2sp_send = messaging.new_message('modelDataV2SP') + if not sm['carControl'].latActive: + prev_action.desiredCurvature = sm['controlsState'].curvature action = get_action_from_model(model_output, prev_action, lat_action_t, long_action_t, v_ego) prev_action = action fill_model_msg(modelv2_send, model_output, action, diff --git a/openpilot/sunnypilot/modeld_v2/compile_modeld.py b/openpilot/sunnypilot/modeld_v2/compile_modeld.py index 2b7256aeb..732c2d54a 100755 --- a/openpilot/sunnypilot/modeld_v2/compile_modeld.py +++ b/openpilot/sunnypilot/modeld_v2/compile_modeld.py @@ -7,14 +7,12 @@ See the LICENSE.md file in the root directory for more details. """ import argparse -import atexit import codecs import gc import math import multiprocessing as mp import os import pickle -import shutil import tempfile import time from functools import partial @@ -22,8 +20,8 @@ from typing import Any import numpy as np -from openpilot.common.file_chunker import chunk_file, get_chunk_targets, open_file_chunked -from openpilot.sunnypilot.modeld_v2.helpers import dump_oob, load_oob +from openpilot.sunnypilot.modeld_v2.helpers import PlainJit, patch_fetch_fw, read_parts, write_bundle +from openpilot.sunnypilot.modeld_v2.plain_artifact_compiler import capture_frame, record_kernel_linear, serialize_jit import openpilot.sunnypilot.modeld_v2.stock_dependencies as stock from openpilot.system.camerad.cameras.nv12_info import get_nv12_info @@ -36,34 +34,7 @@ from tinygrad.tensor import Tensor os.environ['GMMU'] = '0' -def _patch_tinygrad_fetch_fw(): - try: - import hashlib - import pathlib - import zstandard - from tinygrad import helpers - _orig_fetch_fw = helpers.fetch_fw - def fetch_fw(path, name, sha256): - p = pathlib.Path(f"/lib/firmware/{path}/{name}.zst") - if p.is_file(): - blob = zstandard.ZstdDecompressor().stream_reader(p.read_bytes()).read() - if hashlib.sha256(blob).hexdigest() == sha256: - return blob - return _orig_fetch_fw(path, name, sha256) - helpers.fetch_fw = fetch_fw - except ImportError: - pass -_patch_tinygrad_fetch_fw() - -def _patch_ops_serialization(): - try: - from tinygrad.uop.ops import Ops - def __reduce_ex__(self, proto): - return (getattr, (self.__class__, self.name)) - Ops.__reduce_ex__ = __reduce_ex__ - except ImportError: - pass -_patch_ops_serialization() +patch_fetch_fw() MODEL_TYPES = ('vision_policy', 'supercombo', 'vision_multi_policy') WARP_INPUTS = ['tfm', 'big_tfm'] @@ -72,6 +43,7 @@ POLICY_INPUTS = ['img_q', 'big_img_q', 'feat_q', 'desire_q', 'packed_npy_inputs' class MetadataOnnxPBParser(OnnxPBParser): def _parse_ModelProto(self) -> dict: obj: dict[str, Any] = {"graph": {"input": [], "output": []}, "metadata_props": []} + self.opset_imports = {} for fid, wire_type in self._parse_message(self.reader.len): match fid: case 7: @@ -150,6 +122,10 @@ def get_policy_npy_shapes(input_shapes: dict, is_supercombo: bool = False) -> tu return shapes, sizes +def host_view(tensor: Tensor) -> np.ndarray: # QCOM -> CPU round trip of tfm takes 5ms longer than host memory view + return np.frombuffer(tensor._buffer()._host_mv(), dtype=np.float32).reshape(tensor.shape) + + def generate_queues_and_npy(input_shapes: dict, frame_skip: int, device: str = Device.DEFAULT, is_supercombo: bool = False) -> tuple[dict, dict]: road_key, _ = _detect_vision_keys(input_shapes) @@ -185,7 +161,7 @@ def generate_queues_and_npy(input_shapes: dict, frame_skip: int, device: str = D 'big_img_q': Tensor(np.zeros(img_buf_shape, dtype=np.uint8), device=device).contiguous().realize(), 'desire_q': Tensor(np.zeros((frame_skip * desire_shape[1], desire_shape[0], desire_shape[2]), dtype=np.float32), device=device).contiguous().realize(), - 'packed_npy_inputs': Tensor(packed_npy_inputs, device=device if os.getenv('CHESTNUT') else 'NPY').realize(), + 'packed_npy_inputs': Tensor(packed_npy_inputs, device=device).realize(), } if features_buffer: @@ -197,9 +173,11 @@ def generate_queues_and_npy(input_shapes: dict, frame_skip: int, device: str = D for key in ('tfm', 'big_tfm'): if key not in npy_arrays: npy_arrays[key] = np.eye(3, dtype=np.float32) - queues[key] = Tensor(npy_arrays[key], device='NPY').realize() + queues[key] = Tensor(npy_arrays[key], device=device).realize() - return queues, npy_arrays + packed = host_view(queues['packed_npy_inputs']) + npy_arrays |= {k: v.reshape(s) for (k, s), v in zip(shapes.items(), np.split(packed, split_indices) if sizes else [], strict=True)} + return queues, npy_arrays | {k: host_view(queues[k]) for k in ('tfm', 'big_tfm')} def make_split_input_queues(vision_input_shapes: dict, policy_input_shapes: dict, @@ -221,8 +199,8 @@ def make_warp_queues(device=Device.DEFAULT): 'tfm': np.zeros((3, 3), dtype=np.float32), 'big_tfm': np.zeros((3, 3), dtype=np.float32), } - queues = {k: Tensor(v, device='NPY').realize() for k, v in npy.items()} - return queues, npy + queues = {k: Tensor(v, device=device).realize() for k, v in npy.items()} + return queues, {k: host_view(v) for k, v in queues.items()} def make_run_policy(vision_runner, policy_runners: list, features_slice: slice, frame_skip: int, input_shapes: dict): @@ -282,8 +260,19 @@ def make_run_policy(vision_runner, policy_runners: list, features_slice: slice, return run_policy +class _CheckFrameRecorder: + def __init__(self, jit): + self.jit, self.frame = jit, None + + def __call__(self, *args, **kwargs): + if self.frame is not None or args or self.jit.cnt < 2: + return self.jit(*args, **kwargs) + self.frame, result = capture_frame(self.jit, kwargs, None) + return result + + def compile_jit(jit_or_fn, input_keys_or_make_inputs, make_queues=None, make_random_inputs=None, - benchmark_runs: int = 1, clear_refs=None, return_loaded: bool = True): + benchmark_runs: int = 1, clear_refs=None): if callable(input_keys_or_make_inputs) and make_queues is None: fn = jit_or_fn make_inputs = input_keys_or_make_inputs @@ -334,22 +323,23 @@ def compile_jit(jit_or_fn, input_keys_or_make_inputs, make_queues=None, make_ran print('capture + replay') gc.collect() - test_val = run_eval(jit, 42, 3) - print(f'pickle round trip ({benchmark_runs} runs per seed)') - with tempfile.TemporaryFile(dir=".") as f: - dump_oob(jit, f) - if clear_refs: - clear_refs() - gc.collect() - f.seek(0) - loaded_jit = load_oob(f) + recorder = _CheckFrameRecorder(jit) + with record_kernel_linear() as recorded: + test_val = run_eval(recorder, 42, 3) + assert len(recorded) == 1, f"expected one captured kernel list, got {len(recorded)}" + part = serialize_jit(jit, recorded[0], recorder.frame) + print(f'plain artifact round trip ({benchmark_runs} runs per seed)') + if clear_refs: + clear_refs() + gc.collect() + loaded_jit = PlainJit(part[0], memoryview(part[1])) for seed in (42, 43): reference = test_val if seed == 42 else run_eval(jit_or_fn, seed, 1) actual = run_eval(loaded_jit, seed, benchmark_runs) for ref, val in zip(reference, actual, strict=True): np.testing.assert_array_equal(ref, val) - return loaded_jit if return_loaded else jit + return part def _parse_size(size_str: str) -> tuple[int, int]: @@ -357,16 +347,6 @@ def _parse_size(size_str: str) -> tuple[int, int]: return int(width), int(height) -def read_file_chunked_to_disk(path): - if not path: - return None - tmp_path = f'{path}.unchunked' - with open(tmp_path, 'wb') as f, open_file_chunked(path) as src: - shutil.copyfileobj(src, f) - atexit.register(lambda: os.path.exists(tmp_path) and os.remove(tmp_path)) - return tmp_path - - def _load_policy_runners(args: argparse.Namespace) -> tuple[list, list]: runners, keys = [], [] for name, onnx_arg in [('policy', args.policy_onnx), ('off_policy', args.off_policy_onnx), ('on_policy', args.on_policy_onnx)]: @@ -387,10 +367,7 @@ def _compile_warp_resolution_worker(cam_w, cam_h, model_w, model_h, benchmark_ru warp = None compiled_jit = compile_jit(warp, WARP_INPUTS, make_warp_queues, make_random_inputs=make_random_warp_inputs, benchmark_runs=benchmark_runs, clear_refs=cleanup_warp) - - result_data = (compiled_jit, Device.DEFAULT) - with open(result_path, "wb") as f: - dump_oob(result_data, f) + write_bundle(result_path, {(cam_w, cam_h): compiled_jit}, {'device': Device.DEFAULT}) if __name__ == "__main__": @@ -421,12 +398,6 @@ if __name__ == "__main__": model_w, model_h = args.model_size output_data = {} - args.vision_onnx = read_file_chunked_to_disk(args.vision_onnx) - args.policy_onnx = read_file_chunked_to_disk(args.policy_onnx) - args.off_policy_onnx = read_file_chunked_to_disk(args.off_policy_onnx) - args.on_policy_onnx = read_file_chunked_to_disk(args.on_policy_onnx) - args.supercombo_onnx = read_file_chunked_to_disk(args.supercombo_onnx) - is_unified_supercombo = False if args.model_type == 'supercombo': assert args.supercombo_onnx @@ -449,8 +420,10 @@ if __name__ == "__main__": frame_copy_size=frame_copy_size) warp = stock.make_warp(nv12, model_w, model_h) run_model_jit = TinyJit(stock.make_run_model(warp, run_policy, model_metadata, frame_copy_size), prune=True) - compiled_jit = compile_jit(run_model_jit, stock.MODELD_INPUTS, make_model_queues, benchmark_runs=args.benchmark_runs, return_loaded=False) - output_data[(cam_w, cam_h)] = compiled_jit + write_bundle(part_path := f"{args.output}.{cam_w}x{cam_h}", {(cam_w, cam_h): compile_jit(run_model_jit, stock.MODELD_INPUTS, make_model_queues, + benchmark_runs=args.benchmark_runs)}) + output_data[(cam_w, cam_h)] = read_parts(part_path)[0][(cam_w, cam_h)] + os.remove(part_path) gc.collect() del model_runner, run_policy gc.collect() @@ -518,18 +491,12 @@ if __name__ == "__main__": tmp_files.append((cam_w, cam_h, tmp_res.name)) for cam_w, cam_h, tmp_name in tmp_files: - with open(tmp_name, "rb") as f: - compiled_jit, dev_name = load_oob(f) - output_data[(cam_w, cam_h)] = compiled_jit - output_data['input_devices']['warp'] = dev_name + warp_parts, warp_extras = read_parts(tmp_name) + output_data[(cam_w, cam_h)] = warp_parts[(cam_w, cam_h)] + output_data['input_devices']['warp'] = warp_extras['device'] os.remove(tmp_name) output_data['metadata']['warp_dev'] = Device.DEFAULT - with open(args.output, "wb") as file: - dump_oob(output_data, file) - - pkl_size = os.path.getsize(args.output) - print(f"Saved combined JIT to {args.output} ({pkl_size / 1e6:.2f} MB)") - chunk_targets = get_chunk_targets(args.output, pkl_size) - chunk_file(args.output, chunk_targets) - print(f"Chunked into {len(chunk_targets) - 1} file(s)") + parts = {key: output_data.pop(key) for key in [key for key, value in output_data.items() if isinstance(value, tuple)]} + write_bundle(args.output, parts, output_data) + print(f"Saved model to {args.output} ({os.path.getsize(args.output) / 1e6:.2f} MB, {len(parts)} JITs)") diff --git a/openpilot/sunnypilot/modeld_v2/compile_modeld_v2.py b/openpilot/sunnypilot/modeld_v2/compile_modeld_v2.py old mode 100644 new mode 100755 index a4bfcce31..44ff5fed6 --- a/openpilot/sunnypilot/modeld_v2/compile_modeld_v2.py +++ b/openpilot/sunnypilot/modeld_v2/compile_modeld_v2.py @@ -1,11 +1,59 @@ #!/usr/bin/env python3 -import runpy +import argparse import os +import sys -from openpilot.sunnypilot.modeld_v2.compile_modeld import _patch_ops_serialization +import numpy as np + +from openpilot.sunnypilot.modeld_v2.helpers import PlainJit, write_bundle +from openpilot.sunnypilot.modeld_v2.plain_artifact_compiler import capture_frame, record_kernel_linear, serialize_jit + + +def initialize_check_inputs(arrays: dict): + rng = np.random.default_rng(7) + for value in arrays.values(): + value[...] = rng.standard_normal(value.shape) if value.dtype.kind == "f" else rng.integers(0, 2 if value.dtype == np.bool_ else 16, value.shape) if __name__ == "__main__": - _patch_ops_serialization() - upstream_script = os.path.join(os.environ.get("TINYGRAD_PATH", "tinygrad_repo"), "examples/openpilot/compile_onnx.py") - runpy.run_path(upstream_script, run_name="__main__") + sys.path.insert(0, os.environ.get("TINYGRAD_PATH", "tinygrad_repo")) + from examples.openpilot.compile_onnx import onnx_metadata + from examples.openpilot.helpers import allocate_inputs, benchmark + from tinygrad import Device, Tensor, TinyJit + from tinygrad.dtype import _to_np_dtype + from tinygrad.nn.onnx import OnnxRunner + + parser = argparse.ArgumentParser(description="Compile a native-architecture ONNX model into a plain tinygrad artifact") + parser.add_argument('onnx') + parser.add_argument('output') + parser.add_argument('--device-input', action='append', default=[]) + parser.add_argument('--benchmark-runs', type=int, default=20) + args = parser.parse_args() + + runner, (properties, output_shapes) = OnnxRunner(args.onnx), onnx_metadata(args.onnx) + device_inputs = runner.graph_inputs.keys() if '*' in args.device_input else args.device_input + metadata = {'metadata': properties} | {f'{kind}_shapes': {name: tuple(d if isinstance(d, int) else 0 for d in shape) for name, shape in shapes.items()} + for kind, shapes in [('input', {name: spec.shape for name, spec in runner.graph_inputs.items()}), + ('output', output_shapes)]} + input_specs = {name: (tuple(s if isinstance(s, int) else 1 for s in spec.shape), np.dtype(_to_np_dtype(spec.dtype)).str, + Device.DEFAULT if name in device_inputs else 'NPY') for name, spec in runner.graph_inputs.items()} + + def model(inputs): + return {name: value.contiguous() for name, value in runner({name: value.to(Device.DEFAULT) for name, value in inputs.items()}).items()} + output_specs = {name: (value.shape, np.dtype(_to_np_dtype(value.dtype)).name, Device.DEFAULT) for name, value in model(allocate_inputs(input_specs)).items()} + + @TinyJit(prune=True) + def jit(output_buffers, **inputs): + Tensor.realize(*(output_buffers[name].assign(value) for name, value in model(inputs).items())) + + check_inputs = allocate_inputs(input_specs, initialize_check_inputs) | {'output_buffers': allocate_inputs(output_specs)} + with record_kernel_linear() as recorded: + for i in range(2 + args.benchmark_runs): + benchmark(jit, cb=lambda t, i=i: print(f" [{i + 1}/{2 + args.benchmark_runs}] {t * 1e3:.2f} ms"), **check_inputs) + assert len(recorded) == 1, f"expected one captured kernel list, got {len(recorded)}" + + frame, _ = capture_frame(jit, check_inputs, 'output_buffers') + part = serialize_jit(jit, recorded[0], frame) + PlainJit(part[0], memoryview(part[1])) + write_bundle(args.output, {'run': part}, {'metadata': metadata, 'input_specs': input_specs, 'output_specs': output_specs}) + print(f"Saved model to {args.output} ({os.path.getsize(args.output) / 1e6:.2f} MB, {len(part[0]['calls'])} kernels)") diff --git a/openpilot/sunnypilot/modeld_v2/helpers.py b/openpilot/sunnypilot/modeld_v2/helpers.py index 03012be3a..f45584a61 100644 --- a/openpilot/sunnypilot/modeld_v2/helpers.py +++ b/openpilot/sunnypilot/modeld_v2/helpers.py @@ -5,14 +5,25 @@ This file is part of sunnypilot and is licensed under the MIT License. See the LICENSE.md file in the root directory for more details. """ +import bisect +import builtins +import contextlib +import dataclasses +import functools +import hashlib import io -import struct +import itertools +import json +import math +import mmap +import os import pickle -import inspect -import importlib -import tempfile -import shutil -import enum +import struct + +MAGIC = b"SPTGPLN1" +ALIGNMENT = 256 +METADATA_TYPES = ("slice", "tuple", "list", "dict", "set", "frozenset", "complex") +RUNTIME_ENVIRONMENT = ("GMMU", "AMD_AQL", "AMD_DISABLE_SDMA", "GRAPH_ONE_KERNEL") def _patch_system_flock_acquire(): @@ -32,86 +43,289 @@ def _patch_system_flock_acquire(): pass _patch_system_flock_acquire() -def _pad_args(func, args, kwargs): + +def patch_fetch_fw() -> None: try: - sig = inspect.signature(func) - except Exception: - return args, kwargs - params = list(sig.parameters.values()) - if inspect.isfunction(func) and params and params[0].name in ('cls', 'self'): - params = params[1:] - - new_args = list(args) - has_varargs = any(p.kind == inspect.Parameter.VAR_POSITIONAL for p in params) - if len(new_args) > len(params) and not has_varargs: - new_args = new_args[:len(params)] - - for i in range(len(new_args), len(params)): - param = params[i] - if param.kind in (inspect.Parameter.VAR_POSITIONAL, inspect.Parameter.VAR_KEYWORD): - continue - val = param.default if param.default is not inspect.Parameter.empty else None - new_args.append(val) - return new_args, kwargs + import zstandard + except ImportError: + return + zstandard.decompress = lambda data, *_: zstandard.ZstdDecompressor().stream_reader(data).read() # ty: ignore[invalid-assignment] -def _dynamic_factory(real_class): - if isinstance(real_class, type) and issubclass(real_class, enum.Enum): - return real_class - - def factory(*args, **kwargs): - try: - return real_class(*args, **kwargs) - except TypeError: - new_args, new_kwargs = _pad_args(real_class, args, kwargs) - return real_class(*new_args, **new_kwargs) - - class DynamicMeta(type(real_class)): - def __call__(cls, *args, **kwargs): - return factory(*args, **kwargs) - - class DynamicProxy(real_class, metaclass=DynamicMeta): - __slots__ = () - - def __new__(cls, *args, **kwargs): - return factory(*args, **kwargs) - - DynamicProxy.__name__ = real_class.__name__ - DynamicProxy.__module__ = real_class.__module__ - return DynamicProxy +class ArtifactMismatch(RuntimeError): + pass -class DynamicTinygradUnpickler(pickle.Unpickler): +class Blob: + def __init__(self, sink=None): + self.chunks: list = [] + self.locations: dict[bytes, list[int]] = {} + self.size, self.write = 0, sink.write if sink else self.chunks.append + + def add(self, data) -> list[int]: + if (key := hashlib.blake2b(data, digest_size=16).digest()) not in self.locations: + start = self.size + -self.size % ALIGNMENT + self.write(bytes(start - self.size)), self.write(data) + self.locations[key], self.size = [start, start + memoryview(data).nbytes], start + memoryview(data).nbytes + return self.locations[key] + + +class _MetadataUnpickler(pickle.Unpickler): def find_class(self, module, name): - real_class = getattr(importlib.import_module(module), name) - if module.startswith("tinygrad"): - return _dynamic_factory(real_class) - return real_class + if module == "builtins" and name in METADATA_TYPES: + return getattr(builtins, name) + raise pickle.UnpicklingError(f"{module}.{name} is not allowed in model metadata") -def load_oob(f): - opcodes = f.read(struct.unpack(' list: + if isinstance(value, (tuple, list)): + return [tensor for item in value for tensor in flatten(item)] + return [] if value is None else [value] -def dump_oob(obj, f): - with tempfile.TemporaryFile(dir=".") as tmp: - def buffer_cb(buffer): - data = buffer.raw() - tmp.write(struct.pack(' tuple[dict, dict]: + with open(path, "rb") as file: + data = memoryview(mmap.mmap(file.fileno(), 0, access=mmap.ACCESS_COPY)) + if bytes(data[:len(MAGIC)]) != MAGIC: + raise ValueError(f"{path} is not a plain tinygrad model, recompile it with compile_modeld.py") + start = len(MAGIC) + 8 + struct.unpack_from(" tuple[dict, dict, str]: + parts, extras = read_parts(path) + added = extend(parts, extras) if extend else {} + return parts | added, extras, _parts_digest(added) if added else "" + + +def load_model(path: str, device_map: dict | None = None, verify: bool = True, cache: bool = True, extend=None) -> dict: + parts, extras, extra = _extended(path, extend) + if cache and (cached := _read_cache(path, parts, device_map, extra)) is not None: + return extras | cached + shared: dict = {} + jits = {key: PlainJit(header, blob, device_map, verify, shared) for key, (header, blob) in parts.items()} + if cache: + _write_cache(path, jits, device_map, extra) + return extras | jits + + +@functools.cache +def tinygrad_identity() -> str: + import tinygrad + root, digest = os.path.dirname(tinygrad.__file__), hashlib.sha256() + for directory, _, files in sorted(os.walk(root)): + for name in sorted(file for file in files if file.endswith(".py")): + with open(path := os.path.join(directory, name), "rb") as file: + digest.update(f"{os.path.relpath(path, root)}:".encode() + file.read()) + return digest.hexdigest() + + +def _parts_digest(parts: dict) -> str: + digest = hashlib.sha256() + for _, (header, blob) in sorted(parts.items(), key=lambda item: repr(item[0])): + digest.update(json.dumps(header, sort_keys=True).encode()) + for call in header["calls"]: + if call["kind"] == "program": + digest.update(blob[slice(*call["lib"])]) + return digest.hexdigest() + + +def _cache_key(path: str, device_map: dict | None, extra: str = "") -> bytes: + stat = os.stat(path) + return json.dumps({"format": MAGIC.decode(), "tinygrad": tinygrad_identity(), "size": stat.st_size, "mtime_ns": stat.st_mtime_ns, + "extra": extra, "device_map": sorted((device_map or {}).items()), + "env": [os.environ.get(name, "") for name in RUNTIME_ENVIRONMENT]}).encode() + + +def _cache_header_matches(file, path: str, device_map: dict | None, extra: str = "") -> bool: + return file.read(struct.unpack(" bool: # entry point for warmup proc + try: + with open(f"{path}.cache", "rb") as file: + return _cache_header_matches(file, path, device_map) + except (OSError, struct.error): + return False + + +def _read_cache(path: str, parts: dict, device_map: dict | None, extra: str = "") -> dict | None: + try: + with open(f"{path}.cache", "rb") as file: + if not _cache_header_matches(file, path, device_map, extra): + return None + shared: dict = {} + buffers = {repr(key): _allocate_buffers(header, blob, device_map, shared) for key, (header, blob) in parts.items()} + return type("_Unpickler", (pickle.Unpickler,), {"persistent_load": lambda self, pid: buffers[pid[0]][pid[1]]})(file).load() + except Exception: + return None + + +def _write_cache(path: str, jits: dict, device_map: dict | None, extra: str = "") -> None: + from tinygrad.device import Buffer + owners = {id(buffer): (repr(key), index) for key, jit in jits.items() for index, buffer in enumerate(jit.buffers)} + + class _Pickler(pickle.Pickler): + def persistent_id(self, obj): + return owners.get(id(obj)) if isinstance(obj, Buffer) else None + + key, tmp_path = _cache_key(path, device_map, extra), f"{path}.cache.{os.getpid()}.tmp" + with contextlib.suppress(Exception): + with open(tmp_path, "wb") as file: + file.write(struct.pack(" str: + from tinygrad import Device + device_map = device_map or {} + return Device.canonicalize(device_map.get(recorded, device_map.get(recorded.split(":", 1)[0], recorded))) + + +def _upload(buffer, data: memoryview, chunk: int = 16 << 20) -> None: + from tinygrad.device import Buffer + source, chunk = Buffer("NPY", len(data), opaque=memoryview(bytearray(data)) if data.readonly else data), min(chunk, len(data)) + for offset in sorted({*range(0, len(data) - chunk, chunk), len(data) - chunk}): + buffer.view(chunk, offset).ensure_allocated().copy_from(source.view(chunk, offset).ensure_allocated()) + + +def _allocate_buffers(header: dict, blob: memoryview, device_map: dict | None, shared: dict) -> list: + import numpy as np + from tinygrad.device import Buffer + written = {call["args"][i]["buffer"] for call in header["calls"] for i in (call["outs"] if call["kind"] == "program" else (0,)) + if "buffer" in call["args"][i]} + entries = [(_device(entry["device"], device_map), blob[slice(*entry["init"])] if "init" in entry else None, entry) for entry in header["buffers"]] + keys = [(device, entry["dtype"], np.frombuffer(data, np.uint8).ctypes.data, len(data)) if data and index not in written else object() + for index, (device, data, entry) in enumerate(entries)] + for device in {key[0] for key in keys if isinstance(key, tuple) and key not in shared}: + if shared.setdefault(("mapped", device), Buffer(device, 1).ensure_allocated().get_storage().host is not None): + continue + spans = [(entry["init"], entry["nbytes"], key) for (_, _, entry), key in zip(entries, keys, strict=True) + if isinstance(key, tuple) and key[0] == device and key not in shared] + files = getattr(blob, "starts", [0]) # a part read from several files (frame_fusion's) is uploaded a block per file + for _, group in itertools.groupby(sorted(spans, key=lambda span: span[0][0]), key=lambda span: bisect.bisect_right(files, span[0][0])): + group = list(group) + start, end = min(span[0] for span, _, _ in group), max(span[1] for span, _, _ in group) + _upload(block := Buffer(device, end - start).ensure_allocated(), blob[start:end]) + shared |= {key: block.view(nbytes, span[0] - start).ensure_allocated() for span, nbytes, key in group} + buffers = [] + for (device, data, entry), key in zip(entries, keys, strict=True): + if key not in shared: + shared[key] = Buffer(device, entry["nbytes"], initial_value=pickle.PickleBuffer(data) if data else None).ensure_allocated() + buffers.append(shared[key]) + return buffers + + +def _program_uop(call: dict, lib: bytes): + from tinygrad.helpers import Target + from tinygrad.uop.ops import KernelInfo, Ops, ProgramInfo, UOp + parameters = [] + for name, compact_slot, dtype, shape in call["signature"]: + parameter = UOp.param(call["globals"][compact_slot], _dtype(dtype), math.prod(shape)) + image = (shape[0], shape[1]) if len(shape) == 3 and shape[2] == 4 else None + parameters.append(parameter.replace(arg=dataclasses.replace(parameter.arg, name=name, image=image))) + info = ProgramInfo(global_size=tuple(call["global_size"]), local_size=tuple(call["local_size"]), vars=(), globals=tuple(call["globals"]), + outs=tuple(call["outs"]), ins=tuple(call["ins"]), target=Target(**call["target"])) + return UOp(Ops.PROGRAM, src=(UOp(Ops.SINK, arg=KernelInfo(name=call["name"])), UOp(Ops.LINEAR, src=tuple(parameters)), + UOp(Ops.SOURCE, arg=""), UOp(Ops.BINARY, arg=lib)), arg=info) + + +class _Jit: + def __init__(self, captured, input_names: list, buffers: list): + self.captured, self.input_names, self.buffers = captured, input_names, buffers + + def __call__(self, *args, **kwargs): + import tinygrad.engine.jit as jit + input_uops, var_vals, names, _ = jit._prepare_jit_inputs(args, kwargs) + if names != self.input_names: + raise ValueError(f"expected inputs {self.input_names}, got {names}") + return self.captured(input_uops, var_vals) + + +class PlainJit(_Jit): + def __init__(self, header: dict, blob: memoryview, device_map: dict | None = None, verify: bool = True, shared: dict | None = None): + import numpy as np + import tinygrad.engine.jit as jit + from tinygrad import Tensor + from tinygrad.device import Buffer + from tinygrad.dtype import dtypes + from tinygrad.uop.ops import Ops, UOp + + inputs = {spec["slot"]: spec for spec in header["inputs"]} + buffers = _allocate_buffers(header, blob, device_map, {} if shared is None else shared) + bases = [UOp.from_buffer(buffer, _dtype(entry["dtype"])) for buffer, entry in zip(buffers, header["buffers"], strict=True)] + + def argument(described: dict): + if "input" in described: + spec = inputs[described["input"]] + return UOp.param(described["input"], _dtype(spec["dtype"]), spec["numel"], _device(spec["device"], device_map)) + base, dtype, offset, nbytes = bases[described["buffer"]], _dtype(described["dtype"]), described["offset"], described["nbytes"] + if offset == 0 and nbytes == header["buffers"][described["buffer"]]["nbytes"] and dtype == base.dtype: + return base + return (base if base.dtype.itemsize == 1 else base.bitcast(dtypes.uint8)).shrink(((offset, offset + nbytes),)).bitcast(dtype) + + calls = [] + for call in header["calls"]: + arguments = [argument(described) for described in call["args"]] + if call["kind"] == "program": + calls.append(_program_uop(call, bytes(blob[slice(*call["lib"])])).call(*arguments)) + else: + calls.append(arguments[0].store_call(arguments[1])) + # stand-ins for JIT inputs: copy staging to GPU is decided from the input buffers, the real ones bind per call + stand_ins = {slot: UOp.from_buffer(Buffer(_device(spec["device"], device_map), spec["numel"] * _dtype(spec["dtype"]).itemsize).ensure_allocated(), + _dtype(spec["dtype"])) for slot, spec in inputs.items()} + input_uops = tuple(stand_ins.get(slot, next(iter(stand_ins.values()), None)) for slot in range(max(stand_ins, default=-1) + 1)) + linear = jit.compile_linear(UOp(Ops.LINEAR, src=tuple(calls)), beam=0, input_uops=input_uops, cache=False) + + returned = [Tensor(argument(t)).reshape(tuple(t["shape"])) for t in header["returns"]["tensors"]] + ret = {"NoneType": None, "Tensor": returned[0] if returned else None, "tuple": tuple(returned), "list": returned}[header["returns"]["kind"]] + super().__init__(jit.CapturedJit(ret, linear, header["input_names"], []), header["input_names"], buffers) + + if not verify or not (check := header.get("check")): + return + def array(stored): + return np.frombuffer(blob[slice(*stored["data"])], dtype=np.dtype(stored["dtype"])).reshape(stored["shape"]).copy() + kwargs: dict = {} + for path, stored, device in check["inputs"]: + tensor = Tensor(array(stored), device=_device(device, device_map)).realize() + kwargs[path[0]] = tensor if len(path) == 1 else {**kwargs.get(path[0], {}), path[1]: tensor} + result = self(**kwargs) + outputs = kwargs[check["output_container"]].values() if check["output_container"] else flatten(result) + for index, (got, stored) in enumerate(zip([t.numpy() for t in outputs], check["outputs"], strict=True)): + expected = array(stored) + if got.shape != expected.shape or not np.allclose(got, expected, rtol=1e-3, atol=1e-3, equal_nan=True): + raise ArtifactMismatch(f"output {index} differs from the compile-time reference frame; recompile the model") diff --git a/openpilot/sunnypilot/modeld_v2/model_adapters.py b/openpilot/sunnypilot/modeld_v2/model_adapters.py index b24089518..5d1e68fee 100644 --- a/openpilot/sunnypilot/modeld_v2/model_adapters.py +++ b/openpilot/sunnypilot/modeld_v2/model_adapters.py @@ -11,20 +11,21 @@ import pickle import numpy as np from openpilot.common.basedir import BASEDIR +from openpilot.sunnypilot.modeld_v2.frame_fusion import pack_layout +from openpilot.sunnypilot.modeld_v2.helpers import load_model from openpilot.sunnypilot.modeld_v2.compile_modeld import (POLICY_INPUTS, derive_frame_skip, make_split_input_queues, make_supercombo_input_queues) from openpilot.sunnypilot.modeld_v2.stock_dependencies import make_input_queues as stock_make_input_queues, nv12_copy_size from openpilot.system.camerad.cameras.nv12_info import get_nv12_info from tinygrad.device import Buffer -from tinygrad.dtype import DType, dtypes -from tinygrad.helpers import round_up +from tinygrad.dtype import DType from tinygrad.tensor import Tensor from tinygrad.uop.ops import UOp def input_view(buffer: Buffer, shape: tuple[int, ...], dtype: DType, offset: int) -> Tensor: - view = buffer.view(math.prod(shape), dtype, offset).ensure_allocated() - return Tensor(UOp.from_buffer(view)).reshape(shape) + view = buffer.view(math.prod(shape) * dtype.itemsize, offset).ensure_allocated() + return Tensor(UOp.from_buffer(view, dtype)).reshape(shape) class BaseModelAdapter: @@ -64,16 +65,10 @@ class BaseModelAdapter: self.warp_frame_size = self.nv12_info[3] return self.jits[(self.cam_w, self.cam_h)] - warp_dir = Path(BASEDIR) / "openpilot/sunnypilot/modeld_v2/models" warp_name = f'{"big_" if self.chestnut else ""}driving_warp_{self.cam_w}x{self.cam_h}_tinygrad.pkl' - with open(warp_dir / warp_name, 'rb') as f: - warp_data = pickle.load(f) - run_warp = warp_data['run'] - if 'input_specs' in warp_data and 'input_frame' in warp_data['input_specs']: - self.warp_frame_size = warp_data['input_specs']['input_frame'][0][1] - else: - self.warp_frame_size = self.nv12_info[3] if self.chestnut else self.frame_copy_size - return run_warp + warp = load_model(str(Path(BASEDIR) / "openpilot/sunnypilot/modeld_v2/models" / warp_name), cache=False) + self.warp_frame_size = warp['input_specs']['input_frame'][0][1] + return warp['run'] class LegacyModelAdapter(BaseModelAdapter): @@ -168,31 +163,22 @@ class NativeTinygradAdapter(BaseModelAdapter): stride, y_height, uv_height, _ = get_nv12_info(self.cam_w, self.cam_h) self.frame_copy_size = stride * (y_height + uv_height) - self.input_shapes_orig = self.jits['metadata']['input_shapes'] - self._vision_input_names = [k for k in self.input_shapes_orig if 'img' in k] + self._vision_input_names = ['img', 'big_img'] self.vision_output_slices = pickle.loads(codecs.decode(self.jits['metadata']['metadata']['output_slices'].encode(), 'base64')) - self.run_warp = self._load_warp() + self.warp_frame_size = self.frame_copy_size self._init_common() self.run_model = self.jits['run'] - warp_frame_size = getattr(self, 'warp_frame_size', self.nv12_info[3]) self.input_queues = {name: Tensor(np.zeros(shape, dtype=dtype), device=self.model_device).realize() for name, (shape, dtype) in self.input_shapes.items() if name in self.state_pairs} - shapes = {'tfm': (2, 3, 3)} | {name: shape for name, (shape, _) in self.input_shapes.items() - if name not in self.state_pairs and name != 'new_img'} - npy_size = sum(round_up(math.prod(shape) * 4, 128) for shape in shapes.values()) - self.packed_input = np.zeros(npy_size + 2 * warp_frame_size, dtype=np.uint8) - self.input_host = Tensor(self.packed_input, device='NPY')._buffer() - self.input_device = Tensor(self.packed_input, device=self.model_device)._buffer() - self.numpy_inputs = {} - offset = 0 - for name, shape in shapes.items(): - self.numpy_inputs[name] = np.ndarray(shape, dtype=np.float32, buffer=self.packed_input, offset=offset) - self.input_queues[name] = input_view(self.input_device, shape, dtypes.float32, offset) - offset += round_up(self.numpy_inputs[name].nbytes, 128) - self.frames = self.packed_input[npy_size:].reshape(2, warp_frame_size) - self.warp_inputs = {'input_frame': input_view(self.input_device, self.frames.shape, dtypes.uint8, npy_size), 'M_inv': self.input_queues.pop('tfm')} + shapes = {name: shape for name, (shape, _) in self.input_shapes.items() if name not in self.state_pairs and name != 'new_img'} + offsets, packed = pack_layout(shapes, self.warp_frame_size) + self.packed_input = np.zeros(packed, dtype=np.uint8) + self.host = Tensor(self.packed_input, device='NPY') + self.numpy_inputs = {name: np.ndarray(shape, dtype=np.float32, buffer=self.packed_input, offset=offsets[name]) + for name, shape in ({'tfm': (2, 3, 3)} | shapes).items()} + self.frames = self.packed_input[offsets['frames']:].reshape(2, self.warp_frame_size) self.outputs = {name: Tensor(np.zeros(shape, dtype=dtype), device=device).realize() for name, (shape, dtype, device) in self.jits['output_specs'].items()} for name, next_name in self.state_pairs.items(): @@ -211,9 +197,7 @@ class NativeTinygradAdapter(BaseModelAdapter): self.input_queues[key].assign(0).realize() def run(self): - self.input_device.copy_from(self.input_host) - self.input_queues['new_img'] = self.run_warp(**self.warp_inputs) - self.run_model(output_buffers=self.outputs, **self.input_queues) + self.run_model(host=self.host, **self.input_queues, output_buffers=self.outputs) return self.outputs['outputs'] diff --git a/openpilot/sunnypilot/modeld_v2/modeld.py b/openpilot/sunnypilot/modeld_v2/modeld.py index fdfdaf42e..dda0f834a 100755 --- a/openpilot/sunnypilot/modeld_v2/modeld.py +++ b/openpilot/sunnypilot/modeld_v2/modeld.py @@ -7,8 +7,10 @@ See the LICENSE.md file in the root directory for more details. """ from collections.abc import Callable +import gc import os os.environ['GMMU'] = '0' +from tinygrad.device import Device import numpy as np import threading import time @@ -16,7 +18,7 @@ from setproctitle import setproctitle import openpilot.cereal.messaging as messaging from openpilot.common.hardware import COMMA_HARDWARE -from openpilot.selfdrive.modeld.helpers import chestnut_present +from openpilot.selfdrive.modeld.helpers import chestnut_present, wait_for_chestnut from openpilot.cereal import log from opendbc.car.structs import car from openpilot.cereal.services import SERVICE_LIST @@ -24,9 +26,9 @@ from openpilot.cereal.messaging import PubMaster, SubMaster from openpilot.cereal.visionipc import VisionStreamType from msgq.visionipc import VisionIpcClient, VisionBuf from opendbc.car.car_helpers import get_demo_car_params -from openpilot.common.file_chunker import open_file_chunked from openpilot.common.swaglog import cloudlog from openpilot.common.params import Params +from openpilot.common.hardware.usb import cable_connected from openpilot.common.filter_simple import FirstOrderFilter from openpilot.common.realtime import config_realtime_process, DT_MDL from openpilot.common.transformations.camera import DEVICE_CAMERAS @@ -34,6 +36,7 @@ from openpilot.common.transformations.model import get_warp_matrix from openpilot.selfdrive.controls.lib.desire_helper import DesireHelper from openpilot.selfdrive.controls.lib.drive_helpers import get_accel_from_plan, smooth_value from openpilot.selfdrive.modeld.modeld import ChestnutGpuState + from openpilot.sunnypilot.modeld_v2.fill_model_msg import fill_model_msg, fill_pose_msg, PublishState, get_curvature_from_output from openpilot.sunnypilot.modeld_v2.parse_model_outputs import Parser from openpilot.sunnypilot.modeld_v2.constants import ModelConstants, Plan @@ -41,7 +44,8 @@ from openpilot.sunnypilot.modeld_v2.meta_helper import load_meta_constants from openpilot.sunnypilot.modeld_v2.camera_offset_helper import CameraOffsetHelper from openpilot.sunnypilot.livedelay.helpers import get_lat_delay from openpilot.sunnypilot.modeld_v2.modeld_base import ModelStateBase -from openpilot.sunnypilot.modeld_v2.helpers import load_oob +from openpilot.sunnypilot.modeld_v2.helpers import load_model +from openpilot.sunnypilot.modeld_v2.frame_fusion import fused_with_warp from openpilot.sunnypilot.modeld_v2.model_adapters import get_model_adapter from openpilot.sunnypilot.models.helpers import get_active_bundle from openpilot.sunnypilot.selfdrive.controls.lib.relc import RoadEdgeLaneChangeController @@ -51,33 +55,13 @@ BIG_MODEL_TIMEOUT = 60 def _pkl_exists(path): - # a bare chunkmanifest is written eagerly for every model in the catalog, so it - # must NOT be treated as downloaded - require the pkl or every chunk file - from openpilot.common.file_chunker import get_chunk_name, get_manifest_path - if os.path.exists(path): - return True - manifest_path = get_manifest_path(path) - if not os.path.exists(manifest_path): - return False - try: - num_chunks = int(open(manifest_path).read().strip()) - except Exception: - return False - return all(os.path.exists(get_chunk_name(path, i, num_chunks)) for i in range(num_chunks)) + return os.path.exists(path) -def _find_driving_pkl(bundle, chestnut: bool = False): +def _find_driving_pkl(bundle): if (override := os.environ.get('COMBINED_MODEL_PKL')) and _pkl_exists(override): return override - if bundle is None: - # no active bundle: use the bundled default pkl (CD210) so a fresh install - # works immediately without downloading a model - from openpilot.selfdrive.modeld.helpers import modeld_pkl_path - bundled = str(modeld_pkl_path(chestnut)) - if _pkl_exists(bundled): - return bundled - return None - if not bundle.models: + if bundle is None or not bundle.models: return None from openpilot.common.hardware.hw import Paths model_root = Paths.model_root() @@ -120,13 +104,13 @@ class ModelState(ModelStateBase): self.PLANPLUS_CONTROL: float = 1.0 self.chestnut = chestnut - pkl_path = _find_driving_pkl(model_bundle, chestnut=chestnut) + pkl_path = _find_driving_pkl(model_bundle) assert pkl_path is not None, f"No driving pkl found for {'chestnut' if chestnut else 'small model'} — all models must be compiled with compile_modeld.py" self._init_combined(pkl_path, cam_w, cam_h, model_bundle) def _init_combined(self, pkl_path, cam_w, cam_h, bundle): cloudlog.warning(f"loading combined pkl: {pkl_path}") - jits = load_oob(open_file_chunked(pkl_path)) + jits = load_model(pkl_path, extend=fused_with_warp(cam_w, cam_h, self.chestnut)) metadata = jits['metadata'] self.WARP_DEV = metadata.get('warp_dev', 'QCOM') if COMMA_HARDWARE else 'CPU' @@ -198,6 +182,8 @@ class ModelState(ModelStateBase): self.numpy_inputs['big_tfm'][:, :] = transforms[self._wide_key].reshape(3, 3) raw_outputs = self.adapter.run() + if self.chestnut and (err := Device[self.DEV].error_state.host.view(fmt='q')[0]): + raise RuntimeError(f"chestnut failed with {err}") if after_enqueue is not None: after_enqueue() @@ -272,16 +258,18 @@ def main(demo=False): cloudlog.bind(daemon=PROCESS_NAME) setproctitle(PROCESS_NAME) - config_realtime_process(7, 54) - CHESTNUT = chestnut_present() + CHESTNUT = _find_driving_pkl(get_active_bundle(chestnut=True)) is not None and (chestnut_present() or cable_connected()) if CHESTNUT: - os.environ['HCQDEV_WAIT_TIMEOUT_MS'] = '3000' + from tinygrad.runtime.ops_amd import AMDDevice + AMDDevice.wait_timeout_ms = 3000 params = Params() params.put_bool("ChestnutLoading", CHESTNUT) params.remove("ChestnutActive") + gc.disable() + # visionipc clients while True: available_streams = VisionIpcClient.available_streams("camerad", block=False) @@ -314,6 +302,7 @@ def main(demo=False): def load_big(): nonlocal big_model try: + wait_for_chestnut() m = ModelState(cam_w=vipc_client_main.width, cam_h=vipc_client_main.height, chestnut=True) m.warmup() big_model = m @@ -323,7 +312,13 @@ def main(demo=False): loader.start() loader.join(BIG_MODEL_TIMEOUT) model = big_model - params.put_bool("ChestnutActive", model is not None) + if model is None and not params.get_bool("ChestnutRetry"): + # a crank can cut the chestnut power during the load, retry once with an onroad cycle + cloudlog.error("big model load failed, cycling onroad to retry") + params.put_bool("ChestnutRetry", True) + params.put_bool("OnroadCycleRequested", True) + else: + params.put_bool("ChestnutActive", model is not None) small_model = ModelState(cam_w=vipc_client_main.width, cam_h=vipc_client_main.height, chestnut=False) if model is None or CHESTNUT else None if model is None: @@ -332,10 +327,13 @@ def main(demo=False): assert model is not None cloudlog.warning(f"models loaded in {time.monotonic() - st:.1f}s, modeld starting") + config_realtime_process(7, 54) + # messaging pub_socks = ["modelV2", "drivingModelData", "cameraOdometry", "modelDataV2SP"] + (["chestnutGpuState"] if CHESTNUT else []) pm = PubMaster(pub_socks) - sm = SubMaster(["deviceState", "carState", "narrowRoadCameraState", "extrinsicsCalibration", "driverMonitoringState", "carControl", "lateralDelay"]) + sm = SubMaster(["deviceState", "carState", "narrowRoadCameraState", "extrinsicsCalibration", "driverMonitoringState", + "carControl", "controlsState", "lateralDelay"]) publish_state = PublishState() chestnut_state = ChestnutGpuState(pm, model.chestnut) if CHESTNUT else None @@ -463,7 +461,7 @@ def main(demo=False): run_count % round(model.constants.MODEL_FREQ / SERVICE_LIST['chestnutGpuState'].frequency) == 0) model_output = model.run(bufs, transforms, inputs, chestnut_state.send if send_chestnut else None) except Exception: - if not params.get_bool("ChestnutActive"): + if not model.chestnut: raise cloudlog.exception("big model failed, fall back to small") params.put_bool("ChestnutActive", False) @@ -482,6 +480,8 @@ def main(demo=False): posenet_send = messaging.new_message('cameraOdometry') mdv2sp_send = messaging.new_message('modelDataV2SP') + if not sm['carControl'].latActive: + prev_action.desiredCurvature = sm['controlsState'].curvature action = model.get_action_from_model(model_output, prev_action, lat_action_t, long_action_t, v_ego) prev_action = action fill_model_msg(drivingdata_send, modelv2_send, model_output, action, diff --git a/openpilot/sunnypilot/modeld_v2/models/big_driving_warp_1344x760_tinygrad.pkl b/openpilot/sunnypilot/modeld_v2/models/big_driving_warp_1344x760_tinygrad.pkl index 5e84a9bb2..dc9d4e5b0 100644 Binary files a/openpilot/sunnypilot/modeld_v2/models/big_driving_warp_1344x760_tinygrad.pkl and b/openpilot/sunnypilot/modeld_v2/models/big_driving_warp_1344x760_tinygrad.pkl differ diff --git a/openpilot/sunnypilot/modeld_v2/models/big_driving_warp_1928x1208_tinygrad.pkl b/openpilot/sunnypilot/modeld_v2/models/big_driving_warp_1928x1208_tinygrad.pkl index 6a8f18169..4eb6a72f4 100644 Binary files a/openpilot/sunnypilot/modeld_v2/models/big_driving_warp_1928x1208_tinygrad.pkl and b/openpilot/sunnypilot/modeld_v2/models/big_driving_warp_1928x1208_tinygrad.pkl differ diff --git a/openpilot/sunnypilot/modeld_v2/models/driving_warp_1344x760_tinygrad.pkl b/openpilot/sunnypilot/modeld_v2/models/driving_warp_1344x760_tinygrad.pkl index 049009d43..e7fc42bd7 100644 Binary files a/openpilot/sunnypilot/modeld_v2/models/driving_warp_1344x760_tinygrad.pkl and b/openpilot/sunnypilot/modeld_v2/models/driving_warp_1344x760_tinygrad.pkl differ diff --git a/openpilot/sunnypilot/modeld_v2/models/driving_warp_1928x1208_tinygrad.pkl b/openpilot/sunnypilot/modeld_v2/models/driving_warp_1928x1208_tinygrad.pkl index 6f31823d0..d52651c53 100644 Binary files a/openpilot/sunnypilot/modeld_v2/models/driving_warp_1928x1208_tinygrad.pkl and b/openpilot/sunnypilot/modeld_v2/models/driving_warp_1928x1208_tinygrad.pkl differ diff --git a/openpilot/sunnypilot/modeld_v2/stock_dependencies.py b/openpilot/sunnypilot/modeld_v2/stock_dependencies.py old mode 100644 new mode 100755 diff --git a/openpilot/sunnypilot/modeld_v2/tests/helpers.py b/openpilot/sunnypilot/modeld_v2/tests/helpers.py index b6f147391..82a618003 100644 --- a/openpilot/sunnypilot/modeld_v2/tests/helpers.py +++ b/openpilot/sunnypilot/modeld_v2/tests/helpers.py @@ -39,7 +39,7 @@ class DummyModel: class DummyBundle: - def __init__(self, is_20hz=False, models=None, generation=10, is_big=False): + def __init__(self, is_20hz=False, models=None, generation=10): self.overrides = [DummyOverride('lat', '.1'), DummyOverride('long', '.3')] self.generation = generation self.is20hz = is_20hz @@ -188,10 +188,8 @@ def make_pkl_data(archetype): def write_pkl(tmp_path, archetype): - from openpilot.sunnypilot.modeld_v2.helpers import dump_oob pkl_path = tmp_path / 'driving_test_tinygrad.pkl' - with open(pkl_path, 'wb') as f: - dump_oob(make_pkl_data(archetype), f) + pkl_path.write_bytes(pickle.dumps(make_pkl_data(archetype))) return pkl_path @@ -211,6 +209,7 @@ def patch_modeld(monkeypatch): def _patch(bundle): monkeypatch.setattr(helpers, 'get_active_bundle', lambda params=None, *, chestnut=None: bundle) monkeypatch.setattr(modeld_module, 'get_active_bundle', lambda params=None, *, chestnut=None: bundle) + monkeypatch.setattr(modeld_module, 'load_model', lambda path, **_: pickle.loads(pathlib.Path(path).read_bytes())) return _patch diff --git a/openpilot/sunnypilot/modeld_v2/tests/test_combined_pkl_loader.py b/openpilot/sunnypilot/modeld_v2/tests/test_combined_pkl_loader.py index 8b4c4e32a..db0eaf870 100644 --- a/openpilot/sunnypilot/modeld_v2/tests/test_combined_pkl_loader.py +++ b/openpilot/sunnypilot/modeld_v2/tests/test_combined_pkl_loader.py @@ -90,7 +90,6 @@ class TestStockEquivalence(OpenpilotTestCase): def test_unified_run_model(self, tmp_path, monkeypatch, patch_modeld): from openpilot.common.hardware import hw - from openpilot.sunnypilot.modeld_v2.helpers import dump_oob shapes = {'img': (1, 12, 128, 256), 'big_img': (1, 12, 128, 256), 'features_buffer': (1, 24, 32, 512), 'desire_pulse': (1, 25, 8), 'traffic_convention': (1, 2), 'action_t': (1, 2)} import codecs @@ -102,8 +101,7 @@ class TestStockEquivalence(OpenpilotTestCase): (CAM_W, CAM_H): tests_helpers._noop_jit, 'run_policy': tests_helpers._noop_jit, } - with open(tmp_path / 'driving_test_tinygrad.pkl', 'wb') as f: - dump_oob(pkl_data, f) + (tmp_path / 'driving_test_tinygrad.pkl').write_bytes(pickle.dumps(pkl_data)) bundle = DummyBundle(models=[DummyModel('supercombo', 'driving_test_tinygrad.pkl')]) patch_modeld(bundle) monkeypatch.setattr(hw.Paths, 'model_root', staticmethod(lambda: str(tmp_path))) diff --git a/openpilot/sunnypilot/modeld_v2/tests/test_compile_modeld.py b/openpilot/sunnypilot/modeld_v2/tests/test_compile_modeld.py index 239c498e7..348b4aca6 100644 --- a/openpilot/sunnypilot/modeld_v2/tests/test_compile_modeld.py +++ b/openpilot/sunnypilot/modeld_v2/tests/test_compile_modeld.py @@ -5,16 +5,12 @@ This file is part of sunnypilot and is licensed under the MIT License. See the LICENSE.md file in the root directory for more details. """ -import os -import tempfile import unittest -from pathlib import Path import numpy as np from openpilot.common.parameterized import parameterized -from openpilot.common.file_chunker import chunk_file, get_chunk_targets -from openpilot.sunnypilot.modeld_v2.compile_modeld import derive_frame_skip, _detect_desire_key, read_file_chunked_to_disk +from openpilot.sunnypilot.modeld_v2.compile_modeld import derive_frame_skip, _detect_desire_key from openpilot.common.test import OpenpilotTestCase @@ -168,36 +164,6 @@ class TestOutputSlicePreservation(OpenpilotTestCase): "vision and policy slices should not overlap in keys" -class TestReadFileChunkedToDisk(OpenpilotTestCase): - def test_none_passthrough(self): - assert read_file_chunked_to_disk(None) is None - - def test_unchunked_source_staged_on_disk(self): - with tempfile.TemporaryDirectory() as d: - src = Path(d) / "driving_supercombo.onnx" - payload = os.urandom(1024) - src.write_bytes(payload) - - out = Path(read_file_chunked_to_disk(str(src))) - - assert out.parent == Path(d) - assert out.name == "driving_supercombo.onnx.unchunked" - assert out.read_bytes() == payload - - def test_chunked_source_reassembled_on_disk(self): - with tempfile.TemporaryDirectory() as d: - src = Path(d) / "driving_supercombo.onnx" - payload = os.urandom(4096) - src.write_bytes(payload) - chunk_file(str(src), get_chunk_targets(str(src), len(payload))) - assert not src.exists() - - out = Path(read_file_chunked_to_disk(str(src))) - - assert out.parent == Path(d) - assert out.read_bytes() == payload - - class Test4DFeaturesBuffer(OpenpilotTestCase): def test_get_policy_npy_shapes_4d(self): from openpilot.sunnypilot.modeld_v2.compile_modeld import get_policy_npy_shapes diff --git a/openpilot/sunnypilot/modeld_v2/tests/test_fallback.py b/openpilot/sunnypilot/modeld_v2/tests/test_fallback.py index 9e0241ec0..9d431c180 100644 --- a/openpilot/sunnypilot/modeld_v2/tests/test_fallback.py +++ b/openpilot/sunnypilot/modeld_v2/tests/test_fallback.py @@ -1,11 +1,15 @@ -import io +""" +Copyright (c) 2021-, Haibin Wen, sunnypilot, and a number of other contributors. + +This file is part of sunnypilot and is licensed under the MIT License. +See the LICENSE.md file in the root directory for more details. +""" + +import pickle import requests -from unittest import mock -from openpilot.common.file_chunker import get_chunk_name from openpilot.common.hardware import hw from openpilot.common.test import OpenpilotTestCase -from openpilot.sunnypilot.modeld_v2.helpers import dump_oob import openpilot.sunnypilot.modeld_v2.modeld as modeld_module from openpilot.sunnypilot.modeld_v2.tests import helpers as tests_helpers from openpilot.sunnypilot.modeld_v2.tests.helpers import DummyModel, DummyBundle, CAM_W, CAM_H @@ -22,8 +26,8 @@ class TestFallback(OpenpilotTestCase): (tmp_path / tsfdo_file).write_bytes(b'dskfajklsdjlsfka') monkeypatch.setattr(hw.Paths, 'model_root', staticmethod(lambda: str(tmp_path))) - big_bundle = DummyBundle(models=[DummyModel('supercombo', lebowski_file)], is_big=True) - small_bundle = DummyBundle(models=[DummyModel('supercombo', tsfdo_file)], is_big=False) + big_bundle = DummyBundle(models=[DummyModel('supercombo', lebowski_file)]) + small_bundle = DummyBundle(models=[DummyModel('supercombo', tsfdo_file)]) big_pkl = modeld_module._find_driving_pkl(big_bundle) small_pkl = modeld_module._find_driving_pkl(small_bundle) @@ -37,14 +41,11 @@ class TestFallback(OpenpilotTestCase): small_json = requests.get(ModelFetcher.MODEL_URL).json() small_bundle = ModelParser.parse_models(small_json)[-1] - buf = io.BytesIO() - dump_oob(tests_helpers.make_pkl_data(tests_helpers.ARCHETYPES['supercombo_non20hz']), buf) - oob_bytes = buf.getvalue() + oob_bytes = pickle.dumps(tests_helpers.make_pkl_data(tests_helpers.ARCHETYPES['supercombo_non20hz'])) + monkeypatch.setattr(modeld_module, 'load_model', lambda path, **_: pickle.loads(open(path, 'rb').read())) for bundle in (big_bundle, small_bundle): - artifact = bundle.models[0].artifact - for i in range(len(artifact.chunks)): - (tmp_path / get_chunk_name(artifact.fileName, i, len(artifact.chunks))).write_bytes(oob_bytes if i == 0 else b"") + (tmp_path / bundle.models[0].artifact.fileName).write_bytes(oob_bytes) monkeypatch.setattr(modeld_module, 'get_active_bundle', lambda params=None, *, chestnut=None: small_bundle) assert modeld_module.ModelState(CAM_W, CAM_H, chestnut=False).chestnut is False @@ -54,74 +55,3 @@ class TestFallback(OpenpilotTestCase): assert modeld_module.ModelState(CAM_W, CAM_H, chestnut=True).chestnut is True except Exception as e: assert "AMD" in str(e) or "device" in str(e).lower() - - def test_runtime_fallback_from_big_to_small(self): - params = mock.MagicMock() - params.get_bool.return_value = True - - big = mock.MagicMock(chestnut=True) - big.run.side_effect = RuntimeError("eGPU error") - - small = mock.MagicMock(chestnut=False) - small.run.return_value = {"plan": []} - - chestnut = mock.MagicMock(big=True) - model, run_count = big, 50 - - try: - model.run() - except Exception: - if not params.get_bool("chestnutActive"): - raise - params.put_bool("chestnutActive", False) - model = small - chestnut.big = False - run_count = 0 - - assert model is small - assert not chestnut.big - assert run_count == 0 - params.put_bool.assert_called_with("chestnutActive", False) - - model.run() - small.run.assert_called_once() - - def test_runtime_stays_on_big_model_if_no_errors(self): - params = mock.MagicMock() - params.get_bool.return_value = True - - big = mock.MagicMock(chestnut=True) - big.run.return_value = {"plan": [1, 2, 3]} - - small = mock.MagicMock(chestnut=False) - chestnut = mock.MagicMock(big=True) - model, run_count = big, 50 - try: - model.run() - except Exception: - if not params.get_bool("ChestnutActive"): - raise - params.put_bool("ChestnutActive", False) - model = small - chestnut.big = False - run_count = 0 - - assert model is big - assert chestnut.big is True - assert run_count == 50 - params.put_bool.assert_not_called() - small.run.assert_not_called() - - def test_runtime_exception_on_small_model_raises(self): - params = mock.MagicMock() - params.get_bool.return_value = False - - model = mock.MagicMock(chestnut=False) - model.run.side_effect = RuntimeError("CPU error") - - with self.assertRaises(RuntimeError): - try: - model.run() - except Exception: - if not params.get_bool("ChestnutActive"): - raise diff --git a/openpilot/sunnypilot/modeld_v2/tests/test_models.py b/openpilot/sunnypilot/modeld_v2/tests/test_models.py index 8f1cdedbd..d24b629e7 100644 --- a/openpilot/sunnypilot/modeld_v2/tests/test_models.py +++ b/openpilot/sunnypilot/modeld_v2/tests/test_models.py @@ -7,71 +7,14 @@ See the LICENSE.md file in the root directory for more details. import os import unittest -from unittest.mock import patch -from openpilot.common.file_chunker import open_file_chunked -from openpilot.sunnypilot.modeld_v2.helpers import load_oob -from tinygrad.device import Device + +from openpilot.sunnypilot.modeld_v2.helpers import PlainJit, load_model -class TestLegacyModels(unittest.TestCase): - def test_legacy_model_load(self): - base_name = os.environ.get("MODEL_BASE_NAME") - if not base_name: - raise unittest.SkipTest("MODEL_BASE_NAME env var not set, skipping integration test.") - chunk_dir = os.environ.get("MODEL_CHUNK_DIR", "/tmp/model_chunks") - base_path = os.path.join(chunk_dir, base_name) - - try: - f = open_file_chunked(base_path) - except Exception as error: - self.fail(f"Failed to open chunked file {base_path}: {error}") - self.addCleanup(f.close) - - real_getitem = Device.__class__.__getitem__ - - def safe_getitem(device_self, ix): - if ix == "QCOM" and not os.path.exists("/dev/kgsl-3d0"): - return real_getitem(device_self, "CPU") - if ix == "AMD" and not os.path.exists("/dev/kfd"): - return real_getitem(device_self, "CPU") - return real_getitem(device_self, ix) - - with patch.object(Device.__class__, "__getitem__", safe_getitem): - obj = load_oob(f) - - assert isinstance(obj, dict), "Parsed object is not a dictionary" - assert "metadata" in obj, "Metadata key is missing" - - def traverse_jits(d): - if hasattr(d, "captured") and d.captured is not None: - linear_attr = getattr(d.captured, "_linear", None) - if linear_attr is None: - try: - linear_attr = getattr(d.captured, "linear", None) - except AssertionError as e: - if "alloc size must be positive" in str(e): - return - raise - - if linear_attr is not None: - stack = [linear_attr] - visited = set() - while stack: - node = stack.pop() - if id(node) in visited: - continue - visited.add(id(node)) - if hasattr(node, "src") and node.src is not None: - for child in node.src: - stack.append(child) - elif isinstance(d, dict): - for v in d.values(): - traverse_jits(v) - elif isinstance(d, (list, tuple)): - for v in d: - traverse_jits(v) - - if "run_model" in obj: - traverse_jits(obj["run_model"]) - if "run_policy" in obj: - traverse_jits(obj["run_policy"]) +class TestRegistryModels(unittest.TestCase): + def test_model_loads(self): + if not (base_name := os.environ.get("MODEL_BASE_NAME")): + self.skipTest("MODEL_BASE_NAME env var not set, skipping integration test.") + bundle = load_model(os.path.join(os.environ.get("MODEL_DIR", "/tmp/models"), base_name), {"QCOM": "CPU", "AMD": "CPU"}, verify=False, cache=False) + jits = [value for value in bundle.values() if isinstance(value, PlainJit)] + assert "metadata" in bundle and jits and all(len(jit.captured.linear.src) > 0 for jit in jits) diff --git a/openpilot/sunnypilot/models/default_model.py b/openpilot/sunnypilot/models/default_model.py index 5aacea12b..6ff56320a 100755 --- a/openpilot/sunnypilot/models/default_model.py +++ b/openpilot/sunnypilot/models/default_model.py @@ -22,7 +22,7 @@ DEFAULT_MODEL_NAME_PATH = os.path.join(BASEDIR, "openpilot", "sunnypilot", "mode MODEL_HASH_PATH = os.path.join(BASEDIR, "openpilot", "sunnypilot", "models", "tests", "model_hash") BIG_MODEL_HASH_PATH = os.path.join(BASEDIR, "openpilot", "sunnypilot", "models", "tests", "big_model_hash") SUPERCOMBO_ONNX_PATH = os.path.join(BASEDIR, "openpilot", "selfdrive", "modeld", "models", "driving_supercombo.onnx") -BIG_SUPERCOMBO_ONNX_PATH = os.path.join(BASEDIR, "openpilot", "selfdrive", "modeld", "models", "big_driving_supercombo.onnx") +BIG_DRIVING_PKL_PATH = os.path.join(BASEDIR, "openpilot", "selfdrive", "modeld", "models", "big_driving_tinygrad.pkl") def _read_model_name_fields(): @@ -47,18 +47,17 @@ def update_model_hash(): print(f"Generated and updated new combined model hash to {MODEL_HASH_PATH}") - if os.path.exists(BIG_SUPERCOMBO_ONNX_PATH): - import subprocess - rel = os.path.relpath(BIG_SUPERCOMBO_ONNX_PATH, os.getcwd()) - pointer = subprocess.check_output(["git", "show", f"HEAD:{rel}"], text=True) - oid = next(l.split(":", 1)[1] for l in pointer.splitlines() if l.startswith("oid sha256:")) - big_fingerprint = f"{oid}:{fields.get('DEFAULT_BIG_MODEL', '')}:{fields.get('DEFAULT_BIG_MODEL_REF', '')}" - big_combined_hash = hashlib.sha256(big_fingerprint.encode()).hexdigest() + import subprocess + rel = os.path.relpath(BIG_DRIVING_PKL_PATH, os.getcwd()) + pointer = subprocess.check_output(["git", "show", f"HEAD:{rel}"], text=True) + oid = next(l.split(":", 1)[1] for l in pointer.splitlines() if l.startswith("oid sha256:")) + big_fingerprint = f"{oid}:{fields.get('DEFAULT_BIG_MODEL', '')}:{fields.get('DEFAULT_BIG_MODEL_REF', '')}" + big_combined_hash = hashlib.sha256(big_fingerprint.encode()).hexdigest() - with open(BIG_MODEL_HASH_PATH, "w") as f: - f.write(big_combined_hash) + with open(BIG_MODEL_HASH_PATH, "w") as f: + f.write(big_combined_hash) - print(f"Generated and updated new big model hash to {BIG_MODEL_HASH_PATH}") + print(f"Generated and updated new big model hash to {BIG_MODEL_HASH_PATH}") def get_ref_for_name(url: str, name: str) -> str: diff --git a/openpilot/sunnypilot/models/fetcher.py b/openpilot/sunnypilot/models/fetcher.py index 8d704cf57..d0f30bf98 100644 --- a/openpilot/sunnypilot/models/fetcher.py +++ b/openpilot/sunnypilot/models/fetcher.py @@ -6,12 +6,10 @@ See the LICENSE.md file in the root directory for more details. """ import time -import os import requests from requests.exceptions import (SSLError, RequestException, HTTPError) from openpilot.common.params import Params from openpilot.common.swaglog import cloudlog -from openpilot.common.hardware.hw import Paths from openpilot.sunnypilot.models.helpers import is_bundle_version_compatible from openpilot.cereal import custom @@ -26,35 +24,11 @@ class ModelParser: download_uri.sha256 = download_uri_data.get("sha256") return download_uri - @staticmethod - def _parse_chunk(chunk_data) -> custom.ModelManagerSP.Chunk: - chunk = custom.ModelManagerSP.Chunk() - chunk.fileName = chunk_data.get("file_name") - chunk.sha256 = chunk_data.get("sha256") - return chunk - @staticmethod def _parse_artifact(artifact_data) -> custom.ModelManagerSP.Artifact: artifact = custom.ModelManagerSP.Artifact() artifact.fileName = artifact_data.get("file_name") artifact.downloadUri = ModelParser._parse_download_uri(artifact_data.get("download_uri", {})) - - if "chunks" in artifact_data: - artifact.chunks = [ModelParser._parse_chunk(chunk_data) for chunk_data in artifact_data["chunks"]] - - try: - model_dir = Paths.model_root() - os.makedirs(model_dir, exist_ok=True) - manifest_path = os.path.join(model_dir, f"{artifact.fileName}.chunkmanifest") - num_chunks = str(len(artifact.chunks)) - - if not os.path.exists(manifest_path) or open(manifest_path).read().strip() != num_chunks: - with open(manifest_path, "w") as f: - f.write(num_chunks) - cloudlog.info(f"Wrote chunk manifest for {artifact.fileName}: {num_chunks} chunks") - except Exception as e: - cloudlog.warning(f"Failed to write chunk manifest for {artifact.fileName}: {e}") - return artifact @staticmethod @@ -138,11 +112,11 @@ class ModelCache: class ModelFetcher: """Handles fetching and caching of model data from remote source""" - MODEL_URL = "https://raw.githubusercontent.com/sunnypilot/sunnypilot-models/refs/heads/gh-pages/docs/driving_models_v23.json" - MODEL_URL_CHESTNUT = "https://raw.githubusercontent.com/sunnypilot/sunnypilot-models/refs/heads/gh-pages/docs/driving_models_chestnut_v27.json" + MODEL_URL = "https://raw.githubusercontent.com/sunnypilot/sunnypilot-models/refs/heads/gh-pages/docs/driving_models_v25.json" + MODEL_URL_CHESTNUT = "https://raw.githubusercontent.com/sunnypilot/sunnypilot-models/refs/heads/gh-pages/docs/driving_models_chestnut_v28.json" - # GitHub raw is frequently unreachable from mainland China. Keep the original - # URLs authoritative while trying this proxy first for manifests. + # GitHub raw is frequently unreachable from mainland China. Mirror first, + # then fall back to the authoritative raw URL. MODEL_URL_MIRROR = "https://ghproxy.net/" + MODEL_URL MODEL_URL_CHESTNUT_MIRROR = "https://ghproxy.net/" + MODEL_URL_CHESTNUT @@ -176,7 +150,7 @@ class ModelFetcher: def _fetch_and_cache_models(self, source: str) -> list[custom.ModelManagerSP.ModelBundle] | None: """Fetches fresh model data from remote and updates cache. - Returns None when every mirror and the original URL fail. + Returns None on transport errors. Raises on 404 and other fatal HTTP errors. """ for model_url in self.model_urls(source): try: @@ -281,5 +255,3 @@ if __name__ == "__main__": model_overrides = {override.key: override.value for override in bundle.overrides} print(f"Bundle: {bundle.internalName}, Type: {model.type}, Status: {bundle.status}, Overrides: {model_overrides}") print(f"Artifact: {model.artifact.fileName}, Download URI: {model.artifact.downloadUri.uri}") - if model.artifact.chunks: - print(f"Contains {len(model.artifact.chunks)} chunks.") diff --git a/openpilot/sunnypilot/models/helpers.py b/openpilot/sunnypilot/models/helpers.py index 219ea05ff..0909f99bd 100644 --- a/openpilot/sunnypilot/models/helpers.py +++ b/openpilot/sunnypilot/models/helpers.py @@ -16,7 +16,7 @@ from openpilot.common.hardware.hw import Paths from openpilot.selfdrive.modeld.helpers import chestnut_present # SET ME TO THE EXACT JSON VERSION WE SET IN SUNNYPILOT_MODELS REPO -REQUIRED_JSON_VERSION = 20 +REQUIRED_JSON_VERSION = 21 CUSTOM_MODEL_PATH = Paths.model_root() ModelManager = custom.ModelManagerSP @@ -29,9 +29,8 @@ _LAST_VALIDATED_RAW: dict[str, dict | None] = {} def _compute_hash(file_path: str) -> str | None: - from openpilot.common.file_chunker import open_file_chunked try: - with open_file_chunked(file_path) as file: + with open(file_path, 'rb') as file: return hashlib.file_digest(file, "sha256").hexdigest().lower() except FileNotFoundError: return None @@ -58,20 +57,11 @@ def is_bundle_version_compatible(bundle: dict) -> bool: def _bundle_artifacts(bundle: custom.ModelManagerSP.ModelBundle) -> list[tuple[str, str]]: artifacts = [] - from openpilot.common.file_chunker import get_chunk_name for model in getattr(bundle, 'models', []) or []: - for artifact in (getattr(model, 'artifact', None),): - if artifact and getattr(artifact, 'fileName', None): - if len(artifact.chunks) > 0: - for i, chunk in enumerate(artifact.chunks): - chunk_name = get_chunk_name(artifact.fileName, i, len(artifact.chunks)) - if getattr(chunk, 'sha256', None): - artifacts.append((chunk_name, chunk.sha256)) - else: - if getattr(artifact, 'downloadUri', None): - sha256 = getattr(artifact.downloadUri, 'sha256', None) - if sha256: - artifacts.append((artifact.fileName, sha256)) + artifact = getattr(model, 'artifact', None) + if artifact and getattr(artifact, 'fileName', None) and getattr(artifact, 'downloadUri', None): + if sha256 := getattr(artifact.downloadUri, 'sha256', None): + artifacts.append((artifact.fileName, sha256)) return artifacts @@ -175,24 +165,6 @@ def validate_active_bundles(params: Params, source_bundles: dict[str, list[custo get_active_model_runner(params, force_check=True) -def _bundled_default_pkl_available() -> bool: - # the repo ships the default model (CD210) as driving_tinygrad.pkl chunks; only - # a real pkl or all chunk files count (a bare chunkmanifest does not) - from openpilot.common.file_chunker import get_chunk_name, get_manifest_path - from openpilot.selfdrive.modeld.helpers import modeld_pkl_path - path = str(modeld_pkl_path(chestnut=False)) - if os.path.exists(path): - return True - manifest_path = get_manifest_path(path) - if not os.path.exists(manifest_path): - return False - try: - num_chunks = int(open(manifest_path).read().strip()) - except Exception: - return False - return all(os.path.exists(get_chunk_name(path, i, num_chunks)) for i in range(num_chunks)) - - def get_active_model_runner(params: Params | None = None, force_check: bool = False) -> int: params = params or Params() cached_runner_type = params.get("ModelRunnerTypeCache") @@ -201,10 +173,6 @@ def get_active_model_runner(params: Params | None = None, force_check: bool = Fa runner_type = custom.ModelManagerSP.Runner.stock if active_bundle := get_active_bundle(params): runner_type = active_bundle.runner.raw - elif _bundled_default_pkl_available(): - # no bundle selected: modeld_v2 runs the bundled default pkl (CD210), so it - # must not fall back to stock modeld (which expects a run_model-format pkl) - runner_type = custom.ModelManagerSP.Runner.tinygrad if cached_runner_type != runner_type: params.put("ModelRunnerTypeCache", int(runner_type), block=True) diff --git a/openpilot/sunnypilot/models/manager.py b/openpilot/sunnypilot/models/manager.py index 6c2cec3eb..d5682bce4 100644 --- a/openpilot/sunnypilot/models/manager.py +++ b/openpilot/sunnypilot/models/manager.py @@ -85,28 +85,27 @@ class ModelManagerSP: return max(1, int(eta)) # Return at least 1 second if download is ongoing - def _get(self, url: str, session=None, stream: bool = True): + def _get(self, url: str): """GET with automatic fallback to the hf-mirror.com mirror.""" - urls = [url] - mirror = _mirror_url(url) - if mirror != url: - urls.append(mirror) last_err = None - for u in urls: + for request_url in dict.fromkeys((url, _mirror_url(url))): + response = None try: - if session is not None: - return session.get(u, stream=stream, timeout=DOWNLOAD_TIMEOUT) - return requests.get(u, stream=stream, timeout=DOWNLOAD_TIMEOUT) + response = requests.get(request_url, stream=True, timeout=DOWNLOAD_TIMEOUT) + response.raise_for_status() + return response except Exception as e: + if response is not None: + response.close() last_err = e - cloudlog.warning(f"model download failed ({u}): {e}; trying next source") + cloudlog.warning(f"model download failed ({request_url}): {e}; trying next source") raise last_err async def _download_file(self, url: str, path: str, model) -> None: """Downloads a file with progress tracking""" self._download_start_times[model.fileName] = time.monotonic() - with self._get(url) as response: + with self._get(url) as response: # noqa: ASYNC210 response.raise_for_status() total_size = int(response.headers.get("content-length", 0)) bytes_downloaded = 0 @@ -130,50 +129,6 @@ class ModelManagerSP: # Clean up start time after download completes del self._download_start_times[model.fileName] - async def _download_chunked(self, base_url: str, base_path: str, artifact, skip: frozenset[int] | set[int] = frozenset()) -> None: - from openpilot.common.file_chunker import get_chunk_name, get_manifest_path - - num_chunks = len(artifact.chunks) - if num_chunks == 0: - raise ValueError("No chunks defined in artifact") - - manifest_path = get_manifest_path(base_path) - self._download_start_times[artifact.fileName] = time.monotonic() - - # Shared connection saves a TCP+TLS handshake per chunk. - # Keep sequential: the link saturates on one stream and Session is not thread-safe. - completed = len(skip) - with requests.Session() as session: - for i, _ in enumerate(artifact.chunks): - if i in skip: - continue - chunk_url = get_chunk_name(base_url, i, num_chunks) - chunk_path = get_chunk_name(base_path, i, num_chunks) - chunk_downloaded = 0 - with self._get(chunk_url, session=session) as response: - response.raise_for_status() - chunk_size = int(response.headers.get("content-length", 0)) - with open(chunk_path, 'wb') as f: # noqa: ASYNC230 - for data in response.iter_content(chunk_size=self._chunk_size): - f.write(data) - chunk_downloaded += len(data) - if self._download_interrupted(): - raise DownloadCancelled("Download cancelled") - intra = chunk_downloaded / max(chunk_size, 1) - progress = min(99.0, ((completed + intra) / num_chunks) * 100) - artifact.downloadProgress.status = custom.ModelManagerSP.DownloadStatus.downloading - artifact.downloadProgress.progress = progress - artifact.downloadProgress.eta = self._calculate_eta(artifact.fileName, progress) - self._sync_artifact_progress(artifact) - self._report_status() - completed += 1 - - with open(manifest_path, 'w') as f: # noqa: ASYNC230 - f.write(str(num_chunks)) - if os.path.isfile(base_path): # noqa: ASYNC240 - os.remove(base_path) - del self._download_start_times[artifact.fileName] - async def _process_artifact(self, artifact, destination_path: str) -> None: if not artifact.downloadUri.uri: return None @@ -186,28 +141,7 @@ class ModelManagerSP: full_path = os.path.join(destination_path, filename) try: - # progress counts only valid chunks so a resumed download continues the - # bar from where verification left it, instead of falling back to zero - is_cached = False - valid_chunks: set[int] = set() - if len(artifact.chunks) > 0: - from openpilot.common.file_chunker import get_chunk_name - num_chunks = len(artifact.chunks) - for i, chunk in enumerate(artifact.chunks): - if self._download_interrupted(): - raise DownloadCancelled("Download cancelled") - if await verify_file(get_chunk_name(full_path, i, num_chunks), chunk.sha256): - valid_chunks.add(i) - artifact.downloadProgress.status = custom.ModelManagerSP.DownloadStatus.verifying - artifact.downloadProgress.progress = (len(valid_chunks) / num_chunks) * 100 - self._sync_artifact_progress(artifact) - self._report_status() - is_cached = len(valid_chunks) == num_chunks - else: - if await verify_file(full_path, expected_hash): - is_cached = True - - if is_cached: + if await verify_file(full_path, expected_hash): artifact.downloadProgress.status = custom.ModelManagerSP.DownloadStatus.cached artifact.downloadProgress.progress = 100 artifact.downloadProgress.eta = 0 @@ -215,17 +149,9 @@ class ModelManagerSP: self._report_status() return - if len(artifact.chunks) > 0: - await self._download_chunked(url, full_path, artifact, skip=valid_chunks) - from openpilot.common.file_chunker import get_chunk_name - for i, chunk in enumerate(artifact.chunks): - chunk_path = get_chunk_name(full_path, i, len(artifact.chunks)) - if not await verify_file(chunk_path, chunk.sha256): - raise ValueError(f"Hash validation failed for chunk {i+1} of {filename}") - else: - await self._download_file(url, full_path, artifact) - if not await verify_file(full_path, expected_hash): - raise ValueError(f"Hash validation failed for {filename}") + await self._download_file(url, full_path, artifact) + if not await verify_file(full_path, expected_hash): + raise ValueError(f"Hash validation failed for {filename}") artifact.downloadProgress.status = custom.ModelManagerSP.DownloadStatus.downloaded artifact.downloadProgress.progress = 100 @@ -234,7 +160,6 @@ class ModelManagerSP: self._report_status() except DownloadCancelled: - # a cancel keeps whatever is on disk: complete chunks resume the next attempt self._download_start_times.pop(artifact.fileName, None) artifact.downloadProgress.status = custom.ModelManagerSP.DownloadStatus.failed artifact.downloadProgress.eta = 0 @@ -246,9 +171,8 @@ class ModelManagerSP: except Exception as e: cloudlog.error(f"Error downloading {filename}: {str(e)}") - for f in [full_path] + [p for p in (os.path.join(destination_path, f) for f in os.listdir(destination_path)) if filename in p]: - if os.path.isfile(f): # noqa: ASYNC240 - os.remove(f) + if os.path.isfile(full_path): # noqa: ASYNC240 + os.remove(full_path) artifact.downloadProgress.status = custom.ModelManagerSP.DownloadStatus.failed artifact.downloadProgress.eta = 0 self._sync_artifact_progress(artifact) @@ -352,19 +276,15 @@ class ModelManagerSP: if get_selected_bundle(self.params, "chestnut") is not None and get_selected_bundle(self.params, "qcom") is None: if self.params.get("ModelManager_DownloadRef") is None: from openpilot.sunnypilot.models.model_name import DEFAULT_MODEL_REF - if DEFAULT_MODEL_REF: - self.params.put("ModelManager_DownloadRef", DEFAULT_MODEL_REF) + self.params.put("ModelManager_DownloadRef", DEFAULT_MODEL_REF) - # First-boot default: no qcom bundle selected -> auto-use the default - # model (DEFAULT_MODEL_REF, i.e. CD210) instead of falling back to the - # stock selfdrive modeld. Mirrors upstream: the default model runs on - # first boot without any user selection (downloaded once, then active). + # First boot on qcom: auto-download the default plain-artifact model so + # modeld_v2 can start without requiring a manual selection in the UI. if (get_selected_bundle(self.params, "qcom") is None and self.params.get("ModelManager_DownloadRef") is None and not self.chestnut_present): from openpilot.sunnypilot.models.model_name import DEFAULT_MODEL_REF - if DEFAULT_MODEL_REF: - self.params.put("ModelManager_DownloadRef", DEFAULT_MODEL_REF) + self.params.put("ModelManager_DownloadRef", DEFAULT_MODEL_REF) self._process_download_requests() @@ -393,12 +313,11 @@ class ModelManagerSP: if model.artifact.fileName: active_files.append(model.artifact.fileName) - # Remove all files except active ones (including their chunk files) + # Remove all files except active ones model_dir = Paths.model_root() try: for filename in os.listdir(model_dir): - base = filename.split('.chunk')[0] if '.chunk' in filename else filename - if base not in active_files and filename not in active_files: + if filename not in active_files: file_path = os.path.join(model_dir, filename) if os.path.isfile(file_path): os.remove(file_path) diff --git a/openpilot/sunnypilot/models/model_name.py b/openpilot/sunnypilot/models/model_name.py index fdf999a68..bd8e9ac26 100644 --- a/openpilot/sunnypilot/models/model_name.py +++ b/openpilot/sunnypilot/models/model_name.py @@ -1,4 +1,4 @@ DEFAULT_MODEL = "CD210" DEFAULT_MODEL_REF = "5b6436a90cf6902b8aaa71c2b6f3d7164d8ae391" -DEFAULT_BIG_MODEL = "Cinque Terre Model v2" -DEFAULT_BIG_MODEL_REF = "37bfa1413edcdc2e8844984b83727c33f81d8f46" +DEFAULT_BIG_MODEL = "Mountain Dew Model v3" +DEFAULT_BIG_MODEL_REF = "1e219d7785cfb7b7343c7f2366347fc86e7ce5f2" diff --git a/openpilot/sunnypilot/models/tests/big_model_hash b/openpilot/sunnypilot/models/tests/big_model_hash index 308096afe..18d5e67d8 100644 --- a/openpilot/sunnypilot/models/tests/big_model_hash +++ b/openpilot/sunnypilot/models/tests/big_model_hash @@ -1 +1 @@ -2c814f08a2c51323b87839fbf8d2c2a9853a2b5536271b3d67f7b7a2de7f9374 \ No newline at end of file +330fdf61b5519642ebccd93aa50f2d7c4ba6c08c5a10c36a2cf155e498700cef \ No newline at end of file diff --git a/openpilot/sunnypilot/models/tests/test_default_model.py b/openpilot/sunnypilot/models/tests/test_default_model.py index 71a18fc61..570548c0a 100644 --- a/openpilot/sunnypilot/models/tests/test_default_model.py +++ b/openpilot/sunnypilot/models/tests/test_default_model.py @@ -10,7 +10,7 @@ import subprocess from openpilot.sunnypilot import get_file_hash from openpilot.sunnypilot.models.default_model import MODEL_HASH_PATH, SUPERCOMBO_ONNX_PATH, BIG_MODEL_HASH_PATH, \ - BIG_SUPERCOMBO_ONNX_PATH, _read_model_name_fields + BIG_DRIVING_PKL_PATH, _read_model_name_fields import hashlib from openpilot.common.test import OpenpilotTestCase @@ -36,12 +36,9 @@ class TestDefaultModel(OpenpilotTestCase): assert combined_hash == current_hash, "Run openpilot/sunnypilot/models/default_model.py to update the default model name and hash" - def test_compare_big_onnx_hashes(self): - if not os.path.exists(BIG_SUPERCOMBO_ONNX_PATH): - self.skipTest("big_driving_supercombo.onnx not present") - + def test_compare_big_pkl_hashes(self): fields = _read_model_name_fields() - oid = _get_lfs_oid(os.path.relpath(BIG_SUPERCOMBO_ONNX_PATH, os.getcwd())) + oid = _get_lfs_oid(os.path.relpath(BIG_DRIVING_PKL_PATH, os.getcwd())) big_fingerprint = f"{oid}:{fields.get('DEFAULT_BIG_MODEL', '')}:{fields.get('DEFAULT_BIG_MODEL_REF', '')}" combined_hash = hashlib.sha256(big_fingerprint.encode()).hexdigest() diff --git a/openpilot/sunnypilot/models/tests/test_manager_download.py b/openpilot/sunnypilot/models/tests/test_manager_download.py index b995d4fbd..0ae0273ec 100644 --- a/openpilot/sunnypilot/models/tests/test_manager_download.py +++ b/openpilot/sunnypilot/models/tests/test_manager_download.py @@ -17,11 +17,9 @@ from typing import Any from unittest import mock import requests -from urllib3.connectionpool import HTTPConnectionPool from openpilot.cereal import custom from openpilot.common.test import OpenpilotTestCase -from openpilot.common.file_chunker import get_chunk_name, get_manifest_path from openpilot.selfdrive.test.helpers import http_server_context from openpilot.sunnypilot.models import manager as manager_module from openpilot.sunnypilot.models.fetcher import ModelFetcher, get_cached_bundles @@ -30,7 +28,6 @@ from openpilot.sunnypilot.models.helpers import (get_active_bundle, get_active_s resolve_bundle_by_ref, validate_active_bundles) from openpilot.sunnypilot.models.manager import ModelManagerSP -CHUNK_BODIES = [b'A' * 5000, b'B' * 5000, b'C' * 3000] WHOLE_BODY = b'Z' * 9000 @@ -49,12 +46,7 @@ class DownloadHandler(http.server.BaseHTTPRequestHandler): pass def _body_for(self, path): - if path.endswith('.whole'): - return WHOLE_BODY - for i in range(len(CHUNK_BODIES)): - if path.endswith(get_chunk_name('', i, len(CHUNK_BODIES))): - return CHUNK_BODIES[i] - return None + return WHOLE_BODY if path.endswith('.pkl') else None def do_GET(self): type(self).request_paths.append(self.path) @@ -119,35 +111,19 @@ class ManagerDownloadTestBase(OpenpilotTestCase): if artifact is not None: self.reported.append(float(artifact.downloadProgress.progress)) - def make_artifact(self, chunked: bool): + def make_artifact(self): bundle = custom.ModelManagerSP.ModelBundle.new_message() bundle.init('models', 1) artifact = bundle.models[0].artifact artifact.fileName = 'driving_test_tinygrad.pkl' - if chunked: - artifact.downloadUri.uri = self.base_url + '/driving_test_tinygrad.pkl' - artifact.downloadUri.sha256 = sha256(b''.join(CHUNK_BODIES)) - artifact.init('chunks', len(CHUNK_BODIES)) - for i, body in enumerate(CHUNK_BODIES): - artifact.chunks[i].sha256 = sha256(body) - else: - artifact.downloadUri.uri = self.base_url + '/driving_test_tinygrad.pkl.whole' - artifact.downloadUri.sha256 = sha256(WHOLE_BODY) + artifact.downloadUri.uri = self.base_url + '/driving_test_tinygrad.pkl' + artifact.downloadUri.sha256 = sha256(WHOLE_BODY) self._bundle = bundle self.artifact = artifact return artifact - def chunk_paths(self, base_path): - return [get_chunk_name(base_path, i, len(CHUNK_BODIES)) for i in range(len(CHUNK_BODIES))] - - def assert_no_partials(self, base_path): - leftovers = [p for p in [base_path, get_manifest_path(base_path)] + self.chunk_paths(base_path) - if os.path.isfile(p)] - assert leftovers == [], f"partial files left behind: {leftovers}" - - class TestManagerDownload(ManagerDownloadTestBase): - """Exercises the real _download_file / _download_chunked against a local server.""" + """Exercises the real _download_file against a local server.""" def run_with_server(self, fn): with http_server_context(handler=DownloadHandler) as (host, port): @@ -156,7 +132,7 @@ class TestManagerDownload(ManagerDownloadTestBase): def test_download_file_writes_exact_bytes(self): def body(): - artifact = self.make_artifact(chunked=False) + artifact = self.make_artifact() path = os.path.join(self.dest, artifact.fileName) asyncio.run(self.manager._download_file(artifact.downloadUri.uri, path, artifact)) with open(path, 'rb') as f: @@ -166,111 +142,67 @@ class TestManagerDownload(ManagerDownloadTestBase): assert artifact.fileName not in self.manager._download_start_times self.run_with_server(body) - def test_download_chunked_writes_all_chunks_and_manifest(self): - def body(): - artifact = self.make_artifact(chunked=True) - base_path = os.path.join(self.dest, artifact.fileName) - asyncio.run(self.manager._download_chunked(artifact.downloadUri.uri, base_path, artifact)) - - for i, expected in enumerate(CHUNK_BODIES): - with open(get_chunk_name(base_path, i, len(CHUNK_BODIES)), 'rb') as f: - assert f.read() == expected, f"chunk {i} body mismatch" - - with open(get_manifest_path(base_path)) as f: - assert f.read() == str(len(CHUNK_BODIES)) - - assert not os.path.isfile(base_path), "base file should be removed after chunking" - assert artifact.fileName not in self.manager._download_start_times - self.run_with_server(body) - def test_progress_is_monotonic_and_bounded(self): def body(): - artifact = self.make_artifact(chunked=True) + artifact = self.make_artifact() base_path = os.path.join(self.dest, artifact.fileName) - asyncio.run(self.manager._download_chunked(artifact.downloadUri.uri, base_path, artifact)) + asyncio.run(self.manager._download_file(artifact.downloadUri.uri, base_path, artifact)) assert self.reported, "expected progress reports" for a, b in zip(self.reported, self.reported[1:], strict=False): assert b >= a, f"progress went backwards: {a} -> {b}" - assert max(self.reported) <= 99.0, f"chunked progress must stay <=99 until verify, got {max(self.reported)}" - self.run_with_server(body) - - def test_session_is_reused_across_chunks(self): - """One connection pool shared across every chunk.""" - def body(): - artifact = self.make_artifact(chunked=True) - base_path = os.path.join(self.dest, artifact.fileName) - - pools = [] - original = HTTPConnectionPool.urlopen - - def tracked(pool_self, *args, **kwargs): - pools.append(id(pool_self)) - return original(pool_self, *args, **kwargs) - - with mock.patch.object(HTTPConnectionPool, 'urlopen', tracked): - asyncio.run(self.manager._download_chunked(artifact.downloadUri.uri, base_path, artifact)) - - assert len(pools) == len(CHUNK_BODIES), f"expected one request per chunk, got {len(pools)}" - assert len(set(pools)) == 1, f"connection pool not reused across chunks: {len(set(pools))} pools" self.run_with_server(body) def test_http_error_propagates(self): def body(): - artifact = self.make_artifact(chunked=True) + artifact = self.make_artifact() base_path = os.path.join(self.dest, artifact.fileName) - failing = '/' + os.path.basename(get_chunk_name(artifact.downloadUri.uri, 1, len(CHUNK_BODIES))) - DownloadHandler.fail_paths = {failing: 404} + DownloadHandler.fail_paths = {'/' + artifact.fileName: 404} with self.assertRaises(requests.exceptions.HTTPError): - asyncio.run(self.manager._download_chunked(artifact.downloadUri.uri, base_path, artifact)) - - # chunk 1 failed, so its file and the manifest must not exist - assert not os.path.isfile(get_chunk_name(base_path, 1, len(CHUNK_BODIES))) - assert not os.path.isfile(get_manifest_path(base_path)) + asyncio.run(self.manager._download_file(artifact.downloadUri.uri, base_path, artifact)) + assert not os.path.isfile(base_path) self.run_with_server(body) def test_cancellation_mid_transfer(self): """Cancellation is checked inside the byte loop; it must still fire after the port.""" def body(): - artifact = self.make_artifact(chunked=True) + artifact = self.make_artifact() base_path = os.path.join(self.dest, artifact.fileName) self.manager.params.get.return_value = None # cancelled with self.assertRaises(Exception) as ctx: - asyncio.run(self.manager._download_chunked(artifact.downloadUri.uri, base_path, artifact)) + asyncio.run(self.manager._download_file(artifact.downloadUri.uri, base_path, artifact)) assert 'cancelled' in str(ctx.exception).lower() - assert not os.path.isfile(get_manifest_path(base_path)) self.run_with_server(body) def test_repeat_downloads_are_stable(self): """Back-to-back runs must produce identical bytes and leak no start-time state.""" def body(): for _ in range(2): - artifact = self.make_artifact(chunked=True) + artifact = self.make_artifact() base_path = os.path.join(self.dest, artifact.fileName) - asyncio.run(self.manager._download_chunked(artifact.downloadUri.uri, base_path, artifact)) - for i, expected in enumerate(CHUNK_BODIES): - with open(get_chunk_name(base_path, i, len(CHUNK_BODIES)), 'rb') as f: - assert f.read() == expected + asyncio.run(self.manager._download_file(artifact.downloadUri.uri, base_path, artifact)) + with open(base_path, 'rb') as f: + assert f.read() == WHOLE_BODY assert self.manager._download_start_times == {} self.run_with_server(body) def test_download_ref_present_keeps_download_alive(self): """A pending download request (DownloadRef set) must not be cancelled mid-transfer.""" def body(): - artifact = self.make_artifact(chunked=True) + artifact = self.make_artifact() base_path = os.path.join(self.dest, artifact.fileName) self.manager.params.get.side_effect = lambda key: b"ref" if key == "ModelManager_DownloadRef" else None self.manager._download_ref = b"ref" - asyncio.run(self.manager._download_chunked(artifact.downloadUri.uri, base_path, artifact)) - assert os.path.isfile(get_manifest_path(base_path)) + asyncio.run(self.manager._download_file(artifact.downloadUri.uri, base_path, artifact)) + assert os.path.isfile(base_path) self.run_with_server(body) def test_cancellation_via_download_ref(self): """Removing DownloadRef mid-transfer cancels the download.""" def body(): - artifact = self.make_artifact(chunked=True) + artifact = self.make_artifact() base_path = os.path.join(self.dest, artifact.fileName) checks = {"n": 0} @@ -283,20 +215,19 @@ class TestManagerDownload(ManagerDownloadTestBase): self.manager.params.get.side_effect = get self.manager._download_ref = b"ref" with self.assertRaises(Exception) as ctx: - asyncio.run(self.manager._download_chunked(artifact.downloadUri.uri, base_path, artifact)) + asyncio.run(self.manager._download_file(artifact.downloadUri.uri, base_path, artifact)) assert 'cancelled' in str(ctx.exception).lower() - assert not os.path.isfile(get_manifest_path(base_path)) self.run_with_server(body) def test_replaced_download_ref_queues_instead_of_cancelling(self): """Selecting another model mid-transfer lets the running download finish.""" def body(): - artifact = self.make_artifact(chunked=True) + artifact = self.make_artifact() base_path = os.path.join(self.dest, artifact.fileName) self.manager.params.get.side_effect = lambda key: b"other-ref" if key == "ModelManager_DownloadRef" else None self.manager._download_ref = b"ref" - asyncio.run(self.manager._download_chunked(artifact.downloadUri.uri, base_path, artifact)) - assert os.path.isfile(get_manifest_path(base_path)) + asyncio.run(self.manager._download_file(artifact.downloadUri.uri, base_path, artifact)) + assert os.path.isfile(base_path) self.run_with_server(body) def test_replaced_download_ref_is_kept(self): @@ -315,11 +246,10 @@ class TestManagerDownload(ManagerDownloadTestBase): def test_cached_bundle_cancel_skips_slot_write(self): """A cancel must stop an already-on-disk bundle before it is applied to the slot.""" def body(): - artifact = self.make_artifact(chunked=True) + artifact = self.make_artifact() base_path = os.path.join(self.dest, artifact.fileName) - for i, data in enumerate(CHUNK_BODIES): - with open(get_chunk_name(base_path, i, len(CHUNK_BODIES)), 'wb') as f: - f.write(data) + with open(base_path, 'wb') as f: + f.write(WHOLE_BODY) self._bundle.ref = "test-ref" params, store = self._make_params_with_store() store["ModelManager_DownloadRef"] = None # removed -> cancelled @@ -329,42 +259,7 @@ class TestManagerDownload(ManagerDownloadTestBase): asyncio.run(self.manager._download_bundle(self._bundle, self.dest, "qcom")) assert 'cancelled' in str(ctx.exception).lower() assert "ModelManager_ActiveBundle" not in store - assert all(os.path.isfile(p) for p in self.chunk_paths(base_path)), "cancel must not delete cached chunks" - self.run_with_server(body) - - def test_resume_skips_valid_chunks(self): - """A chunk already on disk is kept and not re-downloaded; progress starts above its share.""" - def body(): - artifact = self.make_artifact(chunked=True) - base_path = os.path.join(self.dest, artifact.fileName) - with open(get_chunk_name(base_path, 0, len(CHUNK_BODIES)), 'wb') as f: - f.write(CHUNK_BODIES[0]) - - asyncio.run(self.manager._process_artifact(artifact, self.dest)) - - chunk0_suffix = get_chunk_name('', 0, len(CHUNK_BODIES)) - assert not any(p.endswith(chunk0_suffix) for p in DownloadHandler.request_paths), "valid chunk was re-downloaded" - for i, expected in enumerate(CHUNK_BODIES): - with open(get_chunk_name(base_path, i, len(CHUNK_BODIES)), 'rb') as f: - assert f.read() == expected - assert os.path.isfile(get_manifest_path(base_path)) - assert min(self.reported) >= (1 / len(CHUNK_BODIES)) * 100 - 1, "progress must not restart below the resumed share" - self.run_with_server(body) - - def test_verify_reports_valid_fraction_then_cached(self): - """A fully cached bundle publishes climbing verify progress and ends cached.""" - def body(): - artifact = self.make_artifact(chunked=True) - base_path = os.path.join(self.dest, artifact.fileName) - for i, data in enumerate(CHUNK_BODIES): - with open(get_chunk_name(base_path, i, len(CHUNK_BODIES)), 'wb') as f: - f.write(data) - - asyncio.run(self.manager._process_artifact(artifact, self.dest)) - - assert DownloadHandler.request_paths == [], "cached bundle must not hit the network" - assert [round(p) for p in self.reported[:3]] == [33, 67, 100] - assert artifact.downloadProgress.status == custom.ModelManagerSP.DownloadStatus.cached + assert os.path.isfile(base_path), "cancel must not delete the cached model" self.run_with_server(body) def _make_params_with_store(self): @@ -384,7 +279,7 @@ class TestManagerDownload(ManagerDownloadTestBase): def test_download_writes_qcom_slot(self): """A download resolved to the qcom source writes the qcom active bundle slot only.""" def body(): - artifact = self.make_artifact(chunked=True) + artifact = self.make_artifact() self._bundle.ref = "test-ref" self._bundle.minimumSelectorVersion = helpers.REQUIRED_JSON_VERSION params, store = self._make_params_with_store() @@ -396,15 +291,13 @@ class TestManagerDownload(ManagerDownloadTestBase): assert self.manager.selected_bundle.status == custom.ModelManagerSP.DownloadStatus.downloaded assert self.manager.active_bundle is not None and self.manager.active_bundle.ref == "test-ref" assert self.manager.active_bundle.status == custom.ModelManagerSP.DownloadStatus.downloaded - chunk_names = [get_chunk_name(artifact.fileName, i, len(artifact.chunks)) for i in range(len(artifact.chunks))] - missing = [c for c in chunk_names if not os.path.isfile(os.path.join(self.dest, c))] - assert missing == [], f"chunks missing from the cache: {missing}" + assert os.path.isfile(os.path.join(self.dest, artifact.fileName)) self.run_with_server(body) def test_download_writes_chestnut_slot(self): """A download resolved to the chestnut source writes the chestnut active bundle slot only.""" def body(): - self.make_artifact(chunked=True) + self.make_artifact() self._bundle.ref = "big-ref" self._bundle.minimumSelectorVersion = helpers.REQUIRED_JSON_VERSION params, store = self._make_params_with_store() @@ -782,7 +675,7 @@ class TestEffectiveSource(OpenpilotTestCase): @unittest.skipUnless(os.environ.get('RUN_INTEGRATION_TESTS'), 'requires external network') class TestLiveModelManifest(OpenpilotTestCase): - """Every artifact and chunk URL in the published manifest must resolve.""" + """Every artifact URL in the published manifest must resolve.""" def test_all_manifest_urls_available(self): from openpilot.sunnypilot.models.fetcher import ModelFetcher @@ -794,10 +687,7 @@ class TestLiveModelManifest(OpenpilotTestCase): for bundle in manifest.get('bundles', []): for model in bundle.get('models', []): artifact = model['artifact'] - url = artifact['download_uri']['url'] - chunks = artifact.get('chunks', []) - urls = ([url] if not chunks - else [get_chunk_name(url, i, len(chunks)) for i in range(len(chunks))]) + urls = [artifact['download_uri']['url']] for u in urls: try: r = session.head(u, timeout=15, allow_redirects=True) diff --git a/openpilot/sunnypilot/models/tests/test_tinygrad_ref.py b/openpilot/sunnypilot/models/tests/test_tinygrad_ref.py deleted file mode 100644 index d6d82dfb3..000000000 --- a/openpilot/sunnypilot/models/tests/test_tinygrad_ref.py +++ /dev/null @@ -1,24 +0,0 @@ -import requests - -from openpilot.sunnypilot.models.tinygrad_ref import get_tinygrad_ref -from openpilot.sunnypilot.models.fetcher import ModelFetcher -from openpilot.common.test import OpenpilotTestCase - -def fetch_tinygrad_ref(): - response = requests.get(ModelFetcher.MODEL_URL, timeout=10) - response.raise_for_status() - json_data = response.json() - return json_data.get("tinygrad_ref") - - -class TestTinygradRef(OpenpilotTestCase): - def test_tinygrad_ref(self): - current_ref = get_tinygrad_ref() - remote_ref = fetch_tinygrad_ref() - assert remote_ref == current_ref, ( - f"""tinygrad_repo ref does not match remote tinygrad_ref of current compiled driving models json. - Current: {current_ref} - Remote: {remote_ref} - Please run build-all workflow to update models.""" - ) - print("tinygrad_repo ref matches current compiled driving models json ref.") diff --git a/openpilot/system/manager/process_config.py b/openpilot/system/manager/process_config.py index 4be6c5062..aaa1fa78a 100644 --- a/openpilot/system/manager/process_config.py +++ b/openpilot/system/manager/process_config.py @@ -173,6 +173,7 @@ procs += [ # Models PythonProcess("models_manager", "openpilot.sunnypilot.models.manager", always_run), NativeProcess("modeld_tinygrad", "openpilot/sunnypilot/modeld_v2", ["./modeld"], and_(only_onroad, is_tinygrad_model)), + PythonProcess("modeld_warmup", "openpilot.sunnypilot.modeld_v2.warmup", only_offroad), # Backup PythonProcess("backup_manager", "openpilot.sunnypilot.sunnylink.backups.manager", and_(only_offroad, sunnylink_ready_shim)),