From 36a9c02b8f2e7c9f96898e859a5fe219082162f4 Mon Sep 17 00:00:00 2001 From: Jason Wen Date: Wed, 3 Jun 2026 01:46:14 -0400 Subject: [PATCH] sync: post-merge fixes for tinygrad bump --- .../workflows/build-all-tinygrad-models.yaml | 2 + .../build-single-tinygrad-model.yaml | 12 +- .github/workflows/sunnypilot-build-model.yaml | 54 +- SConstruct | 2 +- release/ci/model_generator.py | 96 +++- sunnypilot/modeld_v2/SConscript | 118 +++-- sunnypilot/modeld_v2/compile_modeld.py | 480 ++++++++++++++++++ sunnypilot/modeld_v2/modeld.py | 254 +++++---- sunnypilot/modeld_v2/tests/conftest.py | 204 ++++++++ .../tests/test_buffer_logic_inspect.py | 263 ---------- .../tests/test_combined_pkl_loader.py | 263 ++++++++++ .../modeld_v2/tests/test_compile_modeld.py | 161 ++++++ sunnypilot/modeld_v2/tests/test_warp.py | 3 +- sunnypilot/modeld_v2/warp.py | 53 +- sunnypilot/models/fetcher.py | 6 +- sunnypilot/models/helpers.py | 44 +- sunnypilot/models/manager.py | 99 +++- sunnypilot/models/model_name.py | 2 +- sunnypilot/models/tests/model_hash | 2 +- 19 files changed, 1638 insertions(+), 480 deletions(-) create mode 100755 sunnypilot/modeld_v2/compile_modeld.py create mode 100644 sunnypilot/modeld_v2/tests/conftest.py delete mode 100644 sunnypilot/modeld_v2/tests/test_buffer_logic_inspect.py create mode 100644 sunnypilot/modeld_v2/tests/test_combined_pkl_loader.py create mode 100644 sunnypilot/modeld_v2/tests/test_compile_modeld.py diff --git a/.github/workflows/build-all-tinygrad-models.yaml b/.github/workflows/build-all-tinygrad-models.yaml index 412676e5f..1e341cb8e 100644 --- a/.github/workflows/build-all-tinygrad-models.yaml +++ b/.github/workflows/build-all-tinygrad-models.yaml @@ -120,6 +120,7 @@ jobs: with: upstream_branch: ${{ matrix.model.ref }} custom_name: ${{ matrix.model.display_name }} + is_20hz: ${{ matrix.model.is_20hz }} recompiled_dir: ${{ needs.setup.outputs.recompiled_dir }} json_version: ${{ needs.setup.outputs.json_version }} secrets: inherit @@ -157,6 +158,7 @@ jobs: with: upstream_branch: ${{ matrix.model.ref }} custom_name: ${{ matrix.model.display_name }} + is_20hz: ${{ matrix.model.is_20hz }} recompiled_dir: ${{ needs.setup.outputs.recompiled_dir }} json_version: ${{ needs.setup.outputs.json_version }} artifact_suffix: -retry diff --git a/.github/workflows/build-single-tinygrad-model.yaml b/.github/workflows/build-single-tinygrad-model.yaml index e7e3b67b5..f10f1b71a 100644 --- a/.github/workflows/build-single-tinygrad-model.yaml +++ b/.github/workflows/build-single-tinygrad-model.yaml @@ -24,6 +24,11 @@ on: required: false type: string default: '' + is_20hz: + description: 'Is this a 20Hz model' + required: false + type: boolean + default: true bypass_push: description: 'Bypass pushing to GitLab for build-all' required: false @@ -39,6 +44,11 @@ on: description: 'Custom name for the model (no date, only name)' required: false type: string + is_20hz: + description: 'Is this a 20Hz model' + required: false + type: boolean + default: true recompiled_dir: description: 'Existing recompiled directory number (e.g. 3 for recompiled3)' required: true @@ -82,7 +92,7 @@ jobs: with: upstream_branch: ${{ inputs.upstream_branch }} custom_name: ${{ inputs.custom_name || inputs.upstream_branch }} - is_20hz: true + is_20hz: ${{ inputs.is_20hz }} artifact_suffix: ${{ inputs.artifact_suffix }} secrets: inherit diff --git a/.github/workflows/sunnypilot-build-model.yaml b/.github/workflows/sunnypilot-build-model.yaml index ff09489b9..3c5554fcc 100644 --- a/.github/workflows/sunnypilot-build-model.yaml +++ b/.github/workflows/sunnypilot-build-model.yaml @@ -164,18 +164,54 @@ jobs: source /etc/profile export UV_PROJECT_ENVIRONMENT=${HOME}/venv export VIRTUAL_ENV=$UV_PROJECT_ENVIRONMENT - export PYTHONPATH="${PYTHONPATH}:${{ env.TINYGRAD_PATH }}" + export PYTHONPATH="${PYTHONPATH}:${{ env.TINYGRAD_PATH }}:${{ github.workspace }}" - # Loop through all .onnx files + COMPILE_MODELD="${{ github.workspace }}/sunnypilot/modeld_v2/compile_modeld.py" + MODEL_SIZE=$(python3 -c "from openpilot.common.transformations.model import MEDMODEL_INPUT_SIZE as s; print(f'{s[0]}x{s[1]}')") + CAMERA_RES=$(python3 -c "from openpilot.common.transformations.camera import _ar_ox_fisheye as a, _os_fisheye as o; print(f'{a.width}x{a.height} {o.width}x{o.height}')") + TG_FLAGS="DEV=QCOM IMAGE=1 FLOAT16=1 NOLOCALS=1 JIT_BATCH_SIZE=0 OPENPILOT_HACKS=1" + + # Generate metadata for all ONNX files find "${{ env.MODELS_DIR }}" -maxdepth 1 -name '*.onnx' | while IFS= read -r onnx_file; do - base_name=$(basename "$onnx_file" .onnx) - output_file="${{ env.MODELS_DIR }}/${base_name}_tinygrad.pkl" - - echo "Compiling: $onnx_file -> $output_file" - QCOM=1 python3 "${{ env.TINYGRAD_PATH }}/examples/openpilot/compile3.py" "$onnx_file" "$output_file" - DEV=QCOM FLOAT16=1 NOLOCALS=1 JIT_BATCH_SIZE=0 python3 "${{ env.MODELS_DIR }}/../get_model_metadata.py" "$onnx_file" || true + echo "Generating metadata: $onnx_file" + env ${TG_FLAGS} python3 "${{ env.MODELS_DIR }}/../get_model_metadata.py" "$onnx_file" || true done + # Detect model type and build compile args + VISION_ONNX="${{ env.MODELS_DIR }}/driving_vision.onnx" + POLICY_ONNX="${{ env.MODELS_DIR }}/driving_policy.onnx" + OFF_POLICY_ONNX="${{ env.MODELS_DIR }}/driving_off_policy.onnx" + ON_POLICY_ONNX="${{ env.MODELS_DIR }}/driving_on_policy.onnx" + SUPERCOMBO_ONNX="${{ env.MODELS_DIR }}/supercombo.onnx" + + MODEL_TYPE="" ONNX_ARGS="" OUTPUT_NAME="" + if [ -f "$VISION_ONNX" ]; then + ONNX_ARGS="--vision-onnx $VISION_ONNX" + if [ -f "$ON_POLICY_ONNX" ] && [ -f "$OFF_POLICY_ONNX" ]; then + MODEL_TYPE=vision_multi_policy + ONNX_ARGS="$ONNX_ARGS --off-policy-onnx $OFF_POLICY_ONNX --on-policy-onnx $ON_POLICY_ONNX" + elif [ -f "$OFF_POLICY_ONNX" ] && [ -f "$POLICY_ONNX" ]; then + MODEL_TYPE=vision_multi_policy + ONNX_ARGS="$ONNX_ARGS --policy-onnx $POLICY_ONNX --off-policy-onnx $OFF_POLICY_ONNX" + elif [ -f "$POLICY_ONNX" ]; then + MODEL_TYPE=vision_policy + ONNX_ARGS="$ONNX_ARGS --policy-onnx $POLICY_ONNX" + fi + elif [ -f "$SUPERCOMBO_ONNX" ]; then + MODEL_TYPE=supercombo + ONNX_ARGS="--supercombo-onnx $SUPERCOMBO_ONNX" + fi + + if [ -n "$MODEL_TYPE" ]; then + echo "Detected: $MODEL_TYPE -> driving_tinygrad.pkl" + env ${TG_FLAGS} python3 "$COMPILE_MODELD" \ + --model-type $MODEL_TYPE \ + --model-size $MODEL_SIZE \ + --camera-resolutions $CAMERA_RES \ + $ONNX_ARGS \ + --output "${{ env.MODELS_DIR }}/driving_tinygrad.pkl" + fi + - name: Validate Model Outputs run: | source /etc/profile @@ -194,6 +230,8 @@ jobs: rsync -avm \ --include='*.dlc' \ --include='*.pkl' \ + --include='*.chunk*' \ + --include='*.chunkmanifest' \ --include='*.onnx' \ --exclude='*' \ --delete-excluded \ diff --git a/SConstruct b/SConstruct index 5eb25ff50..e714571d8 100644 --- a/SConstruct +++ b/SConstruct @@ -190,7 +190,7 @@ else: np_version = SCons.Script.Value(np.__version__) Export('envCython', 'np_version') -Export('env', 'arch', 'acados') +Export('env', 'arch', 'acados', 'release') # Setup cache dir default_cache_dir = '/data/scons_cache' if arch == "larch64" else '/tmp/scons_cache' diff --git a/release/ci/model_generator.py b/release/ci/model_generator.py index feeb80095..afee782be 100755 --- a/release/ci/model_generator.py +++ b/release/ci/model_generator.py @@ -1,3 +1,10 @@ +""" +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 os import pickle import sys @@ -32,6 +39,9 @@ OPTIONAL_OUTPUT_KEYS = frozenset({ def validate_model_outputs(metadata_paths: list[Path]) -> None: combined_keys: set[str] = set() for path in metadata_paths: + if path.stat().st_size == 0: + print(f"skipping empty metadata: {path}") + continue with open(path, "rb") as f: metadata = pickle.load(f) combined_keys.update(metadata.get("output_slices", {}).keys()) @@ -78,38 +88,65 @@ def create_short_name(full_name): return result[:8] -def generate_metadata(model_path: Path, output_dir: Path, short_name: str): - model_path = model_path - output_path = output_dir +def _read_pkl_bytes(pkl_path: Path) -> bytes: + manifest = Path(f"{pkl_path}.chunkmanifest") + if manifest.exists(): + num_chunks = int(manifest.read_text().strip()) + parts = [] + for i in range(num_chunks): + chunk = Path(f"{pkl_path}.chunk{i + 1:02d}of{num_chunks:02d}") + parts.append(chunk.read_bytes()) + return b''.join(parts) + return pkl_path.read_bytes() + + +def _find_driving_pkl(output_path: Path) -> Path | None: + for pattern in ('driving_tinygrad.pkl', 'driving_*_tinygrad.pkl'): + matches = sorted(output_path.glob(pattern)) + if matches: + return matches[0] + for pattern in ('driving_tinygrad.pkl.chunkmanifest', 'driving_*_tinygrad.pkl.chunkmanifest'): + matches = sorted(output_path.glob(pattern)) + if matches: + return Path(str(matches[0]).removesuffix('.chunkmanifest')) + return None + + +def _rename_pkl_with_chunks(old_pkl: Path, new_pkl: Path) -> Path: + manifest = Path(f"{old_pkl}.chunkmanifest") + if manifest.exists(): + for f in sorted(old_pkl.parent.glob(f"{old_pkl.name}.chunk*")): + f.rename(old_pkl.parent / f.name.replace(old_pkl.name, new_pkl.name, 1)) + return new_pkl + return old_pkl.rename(new_pkl) + + +def generate_metadata(model_path: Path, output_dir: Path, short_name: str, driving_pkl: Path): base = model_path.stem + metadata_file = output_dir / f"{base}_metadata.pkl" - # Define output files for tinygrad and metadata - tinygrad_file = output_path / f"{base}_tinygrad.pkl" - metadata_file = output_path / f"{base}_metadata.pkl" + if short_name: + renamed_meta = output_dir / f"{base}_{short_name.lower()}_metadata.pkl" + if metadata_file.exists() and not renamed_meta.exists(): + metadata_file = metadata_file.rename(renamed_meta) + elif renamed_meta.exists(): + metadata_file = renamed_meta - if not tinygrad_file.exists() or not metadata_file.exists(): - print(f"Error: Missing files for model {base} ({tinygrad_file} or {metadata_file})", file=sys.stderr) + if not metadata_file.exists(): + print(f"Warning: Missing metadata for {base} ({metadata_file}), skipping", file=sys.stderr) return - # Calculate the sha256 hashes - with open(tinygrad_file, 'rb') as f: - tinygrad_hash = hashlib.sha256(f.read()).hexdigest() + tinygrad_hash = hashlib.sha256(_read_pkl_bytes(driving_pkl)).hexdigest() with open(metadata_file, 'rb') as f: metadata_hash = hashlib.sha256(f.read()).hexdigest() - # Rename the files if a custom file name is provided - if short_name: - tinygrad_file = tinygrad_file.rename(output_path / f"{base}_{short_name.lower()}_tinygrad.pkl") - metadata_file = metadata_file.rename(output_path / f"{base}_{short_name.lower()}_metadata.pkl") - - # Build the metadata structure model_type = "offPolicy" if "off_policy" in base else "onPolicy" if "on_policy" in base else base.split("_")[-1] - model_metadata = { + return { "type": model_type, "artifact": { - "file_name": tinygrad_file.name, + "file_name": driving_pkl.name, "download_uri": { "url": "https://gitlab.com/sunnypilot/public/docs.sunnypilot.ai/-/raw/main/", "sha256": tinygrad_hash @@ -124,9 +161,6 @@ def generate_metadata(model_path: Path, output_dir: Path, short_name: str): } } - # Return model metadata - return model_metadata - def create_metadata_json(models: list, output_dir: Path, custom_name=None, short_name=None, is_20hz=False, upstream_branch="unknown"): metadata_json = { @@ -181,14 +215,28 @@ if __name__ == "__main__": _output_dir = Path(args.output_dir) _output_dir.mkdir(exist_ok=True, parents=True) + _short_name = create_short_name(args.custom_name) if args.custom_name else None + + _driving_pkl = _find_driving_pkl(_output_dir) + if not _driving_pkl: + print(f"No driving_tinygrad.pkl found in {_output_dir}", file=sys.stderr) + sys.exit(1) + + if _short_name: + new_pkl = _output_dir / f"driving_{_short_name.lower()}_tinygrad.pkl" + if not new_pkl.exists(): + _driving_pkl = _rename_pkl_with_chunks(_driving_pkl, new_pkl) + else: + _driving_pkl = new_pkl + _models = [] for _model_path in model_paths: - _model_metadata = generate_metadata(Path(_model_path), _output_dir, create_short_name(args.custom_name)) + _model_metadata = generate_metadata(Path(_model_path), _output_dir, _short_name, _driving_pkl) if _model_metadata: _models.append(_model_metadata) if _models: - create_metadata_json(_models, _output_dir, args.custom_name, create_short_name(args.custom_name), args.is_20hz, args.upstream_branch) + create_metadata_json(_models, _output_dir, args.custom_name, _short_name, args.is_20hz, args.upstream_branch) else: print("No models processed.", file=sys.stderr) diff --git a/sunnypilot/modeld_v2/SConscript b/sunnypilot/modeld_v2/SConscript index 8526fd272..a5b0bad86 100644 --- a/sunnypilot/modeld_v2/SConscript +++ b/sunnypilot/modeld_v2/SConscript @@ -1,12 +1,88 @@ import os import glob -Import('env', 'arch') +from openpilot.common.transformations.camera import _ar_ox_fisheye, _os_fisheye +from openpilot.common.transformations.model import MEDMODEL_INPUT_SIZE +from openpilot.system.hardware import HARDWARE, PC + +Import('env', 'arch', 'release') lenv = env.Clone() tinygrad_files = ["#"+x for x in glob.glob(env.Dir("#tinygrad_repo").relpath + "/**", recursive=True, root_dir=env.Dir("#").abspath) if 'pycache' not in x] -# Get model metadata -PC = not os.path.isfile('/TICI') + +def get_camera_configs(): + DEVICE_RESOLUTIONS = { + "tici": (_ar_ox_fisheye.width, _ar_ox_fisheye.height), + "tizi": (_ar_ox_fisheye.width, _ar_ox_fisheye.height), + "mici": (_os_fisheye.width, _os_fisheye.height), + } + if release or PC or 'CI' in os.environ: + return set(DEVICE_RESOLUTIONS.values()) + return [DEVICE_RESOLUTIONS[HARDWARE.get_device_type()]] + +CAMERA_CONFIGS = get_camera_configs() + +tg_flags = { + 'larch64': 'DEV=QCOM FLOAT16=1 NOLOCALS=1 JIT_BATCH_SIZE=0', + 'Darwin': f'DEV=CPU HOME={os.path.expanduser("~")}', +}.get(arch, 'DEV=CPU:LLVM') + +image_flag = { + 'larch64': 'IMAGE=2', +}.get(arch, 'IMAGE=0') + +model_w, model_h = MEDMODEL_INPUT_SIZE +from openpilot.selfdrive.modeld.constants import ModelConstants +frame_skip = ModelConstants.MODEL_RUN_FREQ // ModelConstants.MODEL_CONTEXT_FREQ +camera_res_args = ' '.join(f'{cw}x{ch}' for cw, ch in CAMERA_CONFIGS) + +pythonpath_string = 'PYTHONPATH="${PYTHONPATH}:' + env.Dir("#tinygrad_repo").abspath + ':' + env.Dir("#").abspath + '"' +compile_modeld_script = File("compile_modeld.py").abspath +upstream_compile_script = File(Dir("#selfdrive/modeld").File("compile_modeld.py").abspath) +script_deps = [File("compile_modeld.py"), upstream_compile_script] + +def compile_combined(model_type, onnx_args, output_name): + output_pkl = File(f"models/{output_name}").abspath + cmd = (f'{pythonpath_string} {tg_flags} {image_flag} python3 {compile_modeld_script} ' + f'--model-type {model_type} ' + f'--model-size {model_w}x{model_h} ' + f'--camera-resolutions {camera_res_args} ' + f'{onnx_args} ' + f'--frame-skip {frame_skip} ' + f'--output {output_pkl}') + onnx_files = [f for f in onnx_args.split() if f.endswith('.onnx')] + return lenv.Command(output_pkl, tinygrad_files + script_deps + [File(f) for f in onnx_files if os.path.isfile(f)], cmd) + +# Vision + Policy (stock default model) +vision_onnx = File("models/driving_vision.onnx").abspath +policy_onnx = File("models/driving_policy.onnx").abspath +if os.path.isfile(vision_onnx) and os.path.isfile(policy_onnx): + compile_combined('vision_policy', + f'--vision-onnx {vision_onnx} --policy-onnx {policy_onnx}', + 'driving_combined_tinygrad.pkl') + +# Vision + Off-Policy +off_policy_onnx = File("models/driving_off_policy.onnx").abspath +if os.path.isfile(vision_onnx) and os.path.isfile(off_policy_onnx): + policy_arg = f'--policy-onnx {policy_onnx}' if os.path.isfile(policy_onnx) else '' + compile_combined('vision_multi_policy', + f'--vision-onnx {vision_onnx} {policy_arg} --off-policy-onnx {off_policy_onnx}', + 'driving_combined_multi_tinygrad.pkl') + +# Vision + On-Policy + Off-Policy +on_policy_onnx = File("models/driving_on_policy.onnx").abspath +if os.path.isfile(vision_onnx) and os.path.isfile(on_policy_onnx) and os.path.isfile(off_policy_onnx): + compile_combined('vision_multi_policy', + f'--vision-onnx {vision_onnx} --off-policy-onnx {off_policy_onnx} --on-policy-onnx {on_policy_onnx}', + 'driving_combined_tri_tinygrad.pkl') + +# Supercombo +supercombo_onnx = File("models/supercombo.onnx").abspath +if os.path.isfile(supercombo_onnx): + compile_combined('supercombo', + f'--supercombo-onnx {supercombo_onnx}', + 'driving_combined_supercombo_tinygrad.pkl') + if PC: inputs = tinygrad_files + [File(Dir("#sunnypilot/modeld_v2").File("install_models_pc.py").abspath)] outputs = [] @@ -21,39 +97,3 @@ if PC: if outputs: lenv.Command(outputs, inputs, cmd) -tg_flags = { - 'larch64': 'DEV=QCOM FLOAT16=1 NOLOCALS=1 JIT_BATCH_SIZE=0', - 'Darwin': f'DEV=CPU THREADS=0 HOME={os.path.expanduser("~")}', -}.get(arch, 'DEV=CPU CPU_LLVM=1 THREADS=0') - -image_flag = { - 'larch64': 'IMAGE=2', -}.get(arch, 'IMAGE=0') - -def tg_compile(flags, model_name): - pythonpath_string = 'PYTHONPATH="${PYTHONPATH}:' + env.Dir("#tinygrad_repo").abspath + '"' - fn = File(f"models/{model_name}").abspath - out = fn + "_tinygrad.pkl" - - return lenv.Command( - out, - [fn + ".onnx"] + tinygrad_files, - f'{pythonpath_string} {flags} {image_flag} python3 {Dir("#tinygrad_repo").abspath}/examples/openpilot/compile3.py {fn}.onnx {out}' - ) - -# Compile models -for model_name in ['supercombo', 'driving_vision', 'driving_off_policy', 'driving_on_policy', 'driving_policy']: - if File(f"models/{model_name}.onnx").exists(): - tg_compile(tg_flags, model_name) - -script_files = [File("warp.py"), File(Dir("#selfdrive/modeld").File("compile_warp.py").abspath)] -pythonpath_string = 'PYTHONPATH="${PYTHONPATH}:' + env.Dir("#tinygrad_repo").abspath + ':' + env.Dir("#").abspath + '"' -compile_warp_cmd = f'{pythonpath_string} {tg_flags} python3 -m sunnypilot.modeld_v2.warp' - -from openpilot.common.transformations.camera import _ar_ox_fisheye, _os_fisheye -warp_targets = [] -for cam in [_ar_ox_fisheye, _os_fisheye]: - w, h = cam.width, cam.height - for bl in [2, 5]: - warp_targets.append(File(f"models/warp_{w}x{h}_b{bl}_tinygrad.pkl").abspath) -lenv.Command(warp_targets, tinygrad_files + script_files, compile_warp_cmd) diff --git a/sunnypilot/modeld_v2/compile_modeld.py b/sunnypilot/modeld_v2/compile_modeld.py new file mode 100755 index 000000000..387d3378e --- /dev/null +++ b/sunnypilot/modeld_v2/compile_modeld.py @@ -0,0 +1,480 @@ +#!/usr/bin/env python3 +""" +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 argparse +import os +import pickle +import time +from functools import partial +from collections import defaultdict + +import numpy as np +from tinygrad.tensor import Tensor +from tinygrad.device import Device +from tinygrad.engine.jit import TinyJit + +from openpilot.selfdrive.modeld.compile_modeld import ( + NV12Frame, make_frame_prepare, + shift_and_sample, sample_skip, sample_desire, +) + +MODEL_TYPES = ('vision_policy', 'supercombo', 'vision_multi_policy') + + +def _detect_desire_key(policy_input_shapes): + for k in policy_input_shapes: + if k.startswith('desire'): + return k + return None + + +def _detect_vision_keys(vision_input_shapes): + img_keys = sorted([k for k in vision_input_shapes if 'img' in k]) + road_key = next((k for k in img_keys if 'big' not in k), None) + wide_key = next((k for k in img_keys if 'big' in k), None) + if road_key is None or wide_key is None: + raise ValueError(f"Cannot determine road/wide image keys from {list(vision_input_shapes.keys())}") + return road_key, wide_key + + +def make_split_input_queues(vision_input_shapes, policy_input_shapes, frame_skip, device): + road_key, _ = _detect_vision_keys(vision_input_shapes) + img = vision_input_shapes[road_key] + n_frames = img[1] // 6 + img_buf_shape = (frame_skip * (n_frames - 1) + 1, 6, img[2], img[3]) + + fb = policy_input_shapes['features_buffer'] + desire_key = _detect_desire_key(policy_input_shapes) + dp = policy_input_shapes[desire_key] + tc = policy_input_shapes.get('traffic_convention', (1, 2)) + + npy = { + 'desire': np.zeros(dp[2], dtype=np.float32), + 'traffic_convention': np.zeros(tc, dtype=np.float32), + 'tfm': np.zeros((3, 3), dtype=np.float32), + 'big_tfm': np.zeros((3, 3), dtype=np.float32), + } + + handled = {'features_buffer', desire_key, 'traffic_convention'} + for key, shape in policy_input_shapes.items(): + if key in handled: + continue + npy[key] = np.zeros(shape, dtype=np.float32) + + input_queues = { + 'img_q': Tensor(np.zeros(img_buf_shape, dtype=np.uint8), device=device).contiguous().realize(), + 'big_img_q': Tensor(np.zeros(img_buf_shape, dtype=np.uint8), device=device).contiguous().realize(), + 'feat_q': Tensor(np.zeros((frame_skip * (fb[1] - 1) + 1, fb[0], fb[2]), dtype=np.float32), device=device).contiguous().realize(), + 'desire_q': Tensor(np.zeros((frame_skip * dp[1], dp[0], dp[2]), dtype=np.float32), device=device).contiguous().realize(), + **{k: Tensor(v, device='NPY').realize() for k, v in npy.items()}, + } + return input_queues, npy + + +def make_run_split_policy(vision_runner, policy_runner, nv12: NV12Frame, model_w, model_h, + vision_features_slice, frame_skip, desire_key, extra_policy_keys, + vision_road_key, vision_wide_key, prepare_only=False): + frame_prepare = make_frame_prepare(nv12, model_w, model_h) + sample_skip_fn = partial(sample_skip, frame_skip=frame_skip) + sample_desire_fn = partial(sample_desire, frame_skip=frame_skip) + + def run_policy(img_q, big_img_q, feat_q, desire_q, desire, traffic_convention, tfm, big_tfm, frame, big_frame, **extra): + npy_tensors = [tfm.to(Device.DEFAULT), big_tfm.to(Device.DEFAULT), + desire.to(Device.DEFAULT), traffic_convention.to(Device.DEFAULT)] + extra_device = {k: extra[k].to(Device.DEFAULT) for k in extra_policy_keys} + Tensor.realize(*npy_tensors, *extra_device.values()) + tfm, big_tfm, desire, traffic_convention = npy_tensors + + img = shift_and_sample(img_q, frame_prepare(frame, tfm).unsqueeze(0), sample_skip_fn) + big_img = shift_and_sample(big_img_q, frame_prepare(big_frame, big_tfm).unsqueeze(0), sample_skip_fn) + + if prepare_only: + return img, big_img + + vision_out = next(iter(vision_runner({vision_road_key: img, vision_wide_key: big_img}).values())).cast('float32') + + new_feat = vision_out[:, vision_features_slice].reshape(1, -1).unsqueeze(0) + feat_buf = shift_and_sample(feat_q, new_feat, sample_skip_fn) + desire_buf = shift_and_sample(desire_q, desire.reshape(1, 1, -1), sample_desire_fn) + + inputs = {'features_buffer': feat_buf, desire_key: desire_buf, 'traffic_convention': traffic_convention, **extra_device} + policy_out = next(iter(policy_runner(inputs).values())).cast('float32') + + return vision_out, policy_out + return run_policy + + +def compile_split_policy(nv12: NV12Frame, model_w, model_h, prepare_only, frame_skip, + vision_runner, policy_runner, vision_metadata, policy_metadata): + print(f"Compiling combined policy JIT for {nv12.width}x{nv12.height} (prepare_only={prepare_only})...") + + vision_features_slice = vision_metadata['output_slices']['hidden_state'] + vision_input_shapes = vision_metadata['input_shapes'] + policy_input_shapes = policy_metadata['input_shapes'] + desire_key = _detect_desire_key(policy_input_shapes) + extra_policy_keys = [k for k in policy_input_shapes if k not in ('features_buffer', desire_key, 'traffic_convention')] + vision_road_key, vision_wide_key = _detect_vision_keys(vision_input_shapes) + + _run = make_run_split_policy(vision_runner, policy_runner, nv12, model_w, model_h, + vision_features_slice, frame_skip, desire_key, extra_policy_keys, + vision_road_key, vision_wide_key, prepare_only) + run_policy_jit = TinyJit(_run, prune=True) + + SEED = 42 + + def random_inputs_run_fn(fn, seed, test_val=None, test_buffers=None, expect_match=True): + input_queues, npy = make_split_input_queues(vision_input_shapes, policy_input_shapes, frame_skip, Device.DEFAULT) + np.random.seed(seed) + Tensor.manual_seed(seed) + + testing = test_val is not None or test_buffers is not None + n_runs = 1 if testing else 3 + + for i in range(n_runs): + frame = Tensor.randint(nv12.size, low=0, high=256, dtype='uint8').realize() + big_frame = Tensor.randint(nv12.size, low=0, high=256, dtype='uint8').realize() + for v in npy.values(): + v[:] = np.random.randn(*v.shape).astype(v.dtype) + Device.default.synchronize() + st = time.perf_counter() + outs = fn(**input_queues, frame=frame, big_frame=big_frame) + mt = time.perf_counter() + Device.default.synchronize() + et = time.perf_counter() + print(f" [{i+1}/{n_runs}] enqueue {(mt-st)*1e3:6.2f} ms -- total {(et-st)*1e3:6.2f} ms") + + if i == 0: + val = [np.copy(v.numpy()) for v in outs] + buffers = [np.copy(v.numpy().copy()) for v in input_queues.values()] + + if test_val is not None: + match = all(np.array_equal(a, b) for a, b in zip(val, test_val, strict=True)) + assert match == expect_match, f"outputs {'differ from' if expect_match else 'match'} baseline (seed={seed})" + if test_buffers is not None: + match = all(np.array_equal(a, b) for a, b in zip(buffers, test_buffers, strict=True)) + assert match == expect_match, f"buffers {'differ from' if expect_match else 'match'} baseline (seed={seed})" + return fn, val, buffers + + print('capture + replay') + run_policy_jit, test_val, test_buffers = random_inputs_run_fn(run_policy_jit, SEED) + + print('pickle round trip') + run_policy_jit = pickle.loads(pickle.dumps(run_policy_jit)) + random_inputs_run_fn(run_policy_jit, SEED, test_val, test_buffers, expect_match=True) + random_inputs_run_fn(run_policy_jit, SEED+1, test_val, test_buffers, expect_match=False) + return run_policy_jit + + +def derive_frame_skip(vision_input_shapes, policy_input_shapes): + fb = policy_input_shapes.get('features_buffer') + if fb is None: + return 1 + fb_history = fb[1] + if fb_history >= 99: + return 1 + return 4 + + +def make_supercombo_input_queues(input_shapes, frame_skip, device): + img_shape = input_shapes.get('img', input_shapes.get('input_imgs')) + if img_shape is None: + raise ValueError("No img input found in model shapes") + + n_frames = img_shape[1] // 6 + img_buf_shape = (frame_skip * (n_frames - 1) + 1, 6, img_shape[2], img_shape[3]) + + npy_keys = {} + queue_keys = {} + + for key, shape in input_shapes.items(): + if 'img' in key: + continue + if len(shape) == 3 and shape[1] > 1: + if key.startswith('desire'): + npy_keys[key] = np.zeros(shape[2], dtype=np.float32) + queue_keys[f'{key}_q'] = Tensor( + np.zeros((frame_skip * shape[1], shape[0], shape[2]), dtype=np.float32), + device=device).contiguous().realize() + elif key == 'features_buffer': + queue_keys['feat_q'] = Tensor( + np.zeros((frame_skip * (shape[1] - 1) + 1, shape[0], shape[2]), dtype=np.float32), + device=device).contiguous().realize() + else: + npy_keys[key] = np.zeros(shape, dtype=np.float32) + elif len(shape) == 2: + npy_keys[key] = np.zeros(shape, dtype=np.float32) + + if 'traffic_convention' not in npy_keys: + tc_shape = input_shapes.get('traffic_convention', (1, 2)) + npy_keys['traffic_convention'] = np.zeros(tc_shape, dtype=np.float32) + + npy_keys['tfm'] = np.zeros((3, 3), dtype=np.float32) + npy_keys['big_tfm'] = np.zeros((3, 3), dtype=np.float32) + + input_queues = { + 'img_q': Tensor(np.zeros(img_buf_shape, dtype=np.uint8), device=device).contiguous().realize(), + 'big_img_q': Tensor(np.zeros(img_buf_shape, dtype=np.uint8), device=device).contiguous().realize(), + **queue_keys, + **{k: Tensor(v, device='NPY').realize() for k, v in npy_keys.items()}, + } + return input_queues, npy_keys + + +def make_run_supercombo(model_runner, nv12: NV12Frame, model_w, model_h, + features_slice, frame_skip, input_shapes, prepare_only=False): + frame_prepare = make_frame_prepare(nv12, model_w, model_h) + sample_skip_fn = partial(sample_skip, frame_skip=frame_skip) + sample_desire_fn = partial(sample_desire, frame_skip=frame_skip) + + desire_key = _detect_desire_key(input_shapes) + if desire_key is None: + raise ValueError(f"No desire* key found in input_shapes: {list(input_shapes.keys())}") + road_img_key, wide_img_key = _detect_vision_keys(input_shapes) + extra_policy_keys = [k for k in input_shapes + if k not in (desire_key, 'features_buffer', 'traffic_convention') + and 'img' not in k] + + def run_supercombo(img_q, big_img_q, feat_q, desire_q, + frame, big_frame, **kwargs): + desire = kwargs.get(desire_key) + traffic_convention = kwargs.get('traffic_convention') + tfm = kwargs['tfm'] + big_tfm = kwargs['big_tfm'] + + tfm = tfm.to(Device.DEFAULT) + big_tfm = big_tfm.to(Device.DEFAULT) + desire = desire.to(Device.DEFAULT) + traffic_convention = traffic_convention.to(Device.DEFAULT) + Tensor.realize(tfm, big_tfm, desire, traffic_convention) + + img = shift_and_sample(img_q, frame_prepare(frame, tfm).unsqueeze(0), sample_skip_fn) + big_img = shift_and_sample(big_img_q, frame_prepare(big_frame, big_tfm).unsqueeze(0), sample_skip_fn) + + if prepare_only: + return img, big_img + + desire_buf = shift_and_sample(desire_q, desire.reshape(1, 1, -1), sample_desire_fn) + feat_buf = sample_skip_fn(feat_q) + + inputs = {road_img_key: img, wide_img_key: big_img, + desire_key: desire_buf, 'features_buffer': feat_buf, + 'traffic_convention': traffic_convention} + for k in extra_policy_keys: + if k in kwargs: + inputs[k] = kwargs[k].to(Device.DEFAULT) + + model_out = next(iter(model_runner(inputs).values())).cast('float32') + + new_feat = model_out[:, features_slice].reshape(1, -1).unsqueeze(0) + shift_and_sample(feat_q, new_feat, sample_skip_fn) + + return model_out + + return run_supercombo + + +def make_run_vision_multi_policy(vision_runner, policy_runners, nv12: NV12Frame, model_w, model_h, + vision_features_slice, frame_skip, desire_key, extra_policy_keys, + vision_road_key, vision_wide_key, prepare_only=False): + frame_prepare = make_frame_prepare(nv12, model_w, model_h) + sample_skip_fn = partial(sample_skip, frame_skip=frame_skip) + sample_desire_fn = partial(sample_desire, frame_skip=frame_skip) + + def run_multi_policy(img_q, big_img_q, feat_q, desire_q, desire, + traffic_convention, tfm, big_tfm, frame, big_frame, **extra): + npy_tensors = [tfm.to(Device.DEFAULT), big_tfm.to(Device.DEFAULT), + desire.to(Device.DEFAULT), traffic_convention.to(Device.DEFAULT)] + extra_device = {k: extra[k].to(Device.DEFAULT) for k in extra_policy_keys} + Tensor.realize(*npy_tensors, *extra_device.values()) + tfm, big_tfm, desire, traffic_convention = npy_tensors + + img = shift_and_sample(img_q, frame_prepare(frame, tfm).unsqueeze(0), sample_skip_fn) + big_img = shift_and_sample(big_img_q, frame_prepare(big_frame, big_tfm).unsqueeze(0), sample_skip_fn) + + if prepare_only: + return img, big_img + + vision_out = next(iter(vision_runner({vision_road_key: img, vision_wide_key: big_img}).values())).cast('float32') + + new_feat = vision_out[:, vision_features_slice].reshape(1, -1).unsqueeze(0) + feat_buf = shift_and_sample(feat_q, new_feat, sample_skip_fn) + desire_buf = shift_and_sample(desire_q, desire.reshape(1, 1, -1), sample_desire_fn) + + inputs = {'features_buffer': feat_buf, desire_key: desire_buf, 'traffic_convention': traffic_convention, **extra_device} + + policy_outputs = [] + for runner in policy_runners: + policy_out = next(iter(runner(inputs).values())).cast('float32') + policy_outputs.append(policy_out) + + return (vision_out, *policy_outputs) + + return run_multi_policy + + +def _warmup_and_serialize(run_jit, input_queues, npy, nv12): + for i in range(3): + np.random.seed(42 + i) + frame = Tensor.randint(nv12.size, low=0, high=256, dtype='uint8').realize() + big_frame = Tensor.randint(nv12.size, low=0, high=256, dtype='uint8').realize() + for v in npy.values(): + v[:] = np.random.randn(*v.shape).astype(v.dtype) + Device.default.synchronize() + st = time.perf_counter() + run_jit(**input_queues, frame=frame, big_frame=big_frame) + mt = time.perf_counter() + Device.default.synchronize() + et = time.perf_counter() + print(f" [{i + 1}/3] enqueue {(mt - st) * 1e3:6.2f} ms -- total {(et - st) * 1e3:6.2f} ms") + return pickle.loads(pickle.dumps(run_jit)) + + +def compile_supercombo(nv12: NV12Frame, model_w, model_h, prepare_only, frame_skip, + model_runner, metadata): + print(f"Compiling combined supercombo JIT for {nv12.width}x{nv12.height} (prepare_only={prepare_only})...") + + features_slice = metadata['output_slices']['hidden_state'] + input_shapes = metadata['input_shapes'] + + _run = make_run_supercombo(model_runner, nv12, model_w, model_h, + features_slice, frame_skip, input_shapes, prepare_only) + run_jit = TinyJit(_run, prune=True) + + input_queues, npy = make_supercombo_input_queues(input_shapes, frame_skip, Device.DEFAULT) + + run_jit = _warmup_and_serialize(run_jit, input_queues, npy, nv12) + return run_jit + + +def compile_multi_policy(nv12: NV12Frame, model_w, model_h, prepare_only, frame_skip, + vision_runner, policy_runners, vision_metadata, policy_metadata): + print(f"Compiling combined multi-policy JIT for {nv12.width}x{nv12.height} (prepare_only={prepare_only})...") + + vision_features_slice = vision_metadata['output_slices']['hidden_state'] + vision_input_shapes = vision_metadata['input_shapes'] + policy_input_shapes = policy_metadata['input_shapes'] + desire_key = _detect_desire_key(policy_input_shapes) + extra_policy_keys = [k for k in policy_input_shapes if k not in ('features_buffer', desire_key, 'traffic_convention')] + vision_road_key, vision_wide_key = _detect_vision_keys(vision_input_shapes) + + _run = make_run_vision_multi_policy(vision_runner, policy_runners, nv12, model_w, model_h, + vision_features_slice, frame_skip, desire_key, extra_policy_keys, + vision_road_key, vision_wide_key, prepare_only) + run_jit = TinyJit(_run, prune=True) + + input_queues, npy = make_split_input_queues(vision_input_shapes, policy_input_shapes, frame_skip, Device.DEFAULT) + + run_jit = _warmup_and_serialize(run_jit, input_queues, npy, nv12) + return run_jit + + +def _parse_size(s): + w, h = s.lower().split('x') + return int(w), int(h) + + +if __name__ == "__main__": + from tinygrad.nn.onnx import OnnxRunner + from openpilot.system.camerad.cameras.nv12_info import get_nv12_info + from openpilot.selfdrive.modeld.get_model_metadata import make_metadata_dict + + p = argparse.ArgumentParser(description="Compile combined JIT pkl for sunnypilot modeld_v2") + p.add_argument('--model-type', choices=MODEL_TYPES, required=True) + p.add_argument('--model-size', type=_parse_size, required=True, help='model input WxH') + p.add_argument('--camera-resolutions', type=_parse_size, nargs='+', required=True) + p.add_argument('--frame-skip', type=int, default=None, help='frame skip value (auto-derived if not provided)') + p.add_argument('--output', required=True) + + p.add_argument('--vision-onnx', help='vision ONNX (for split models)') + p.add_argument('--policy-onnx', help='policy ONNX (for vision_policy)') + p.add_argument('--off-policy-onnx', help='off-policy ONNX (for vision_multi_policy)') + p.add_argument('--on-policy-onnx', help='on-policy ONNX (for vision_multi_policy)') + p.add_argument('--supercombo-onnx', help='supercombo ONNX (for supercombo)') + + args = p.parse_args() + out = defaultdict(dict) + + if args.model_type == 'vision_policy': + assert args.vision_onnx and args.policy_onnx + vision_runner = OnnxRunner(args.vision_onnx) + policy_runner = OnnxRunner(args.policy_onnx) + out['metadata']['vision'] = make_metadata_dict(args.vision_onnx) + out['metadata']['policy'] = make_metadata_dict(args.policy_onnx) + + frame_skip = args.frame_skip if args.frame_skip is not None else derive_frame_skip(out['metadata']['vision']['input_shapes'], + out['metadata']['policy']['input_shapes']) + + for cam_w, cam_h in args.camera_resolutions: + nv12 = NV12Frame(cam_w, cam_h, *get_nv12_info(cam_w, cam_h)) + model_w, model_h = args.model_size + out[(cam_w, cam_h)] = { + name: compile_split_policy(nv12, model_w, model_h, prepare_only, frame_skip, + vision_runner, policy_runner, + out['metadata']['vision'], out['metadata']['policy']) + for name, prepare_only in [('warp_enqueue', True), ('run_policy', False)] + } + + elif args.model_type == 'supercombo': + assert args.supercombo_onnx + model_runner = OnnxRunner(args.supercombo_onnx) + out['metadata']['model'] = make_metadata_dict(args.supercombo_onnx) + + frame_skip = args.frame_skip if args.frame_skip is not None else derive_frame_skip({}, out['metadata']['model']['input_shapes']) + + for cam_w, cam_h in args.camera_resolutions: + nv12 = NV12Frame(cam_w, cam_h, *get_nv12_info(cam_w, cam_h)) + model_w, model_h = args.model_size + out[(cam_w, cam_h)] = { + name: compile_supercombo(nv12, model_w, model_h, prepare_only, frame_skip, + model_runner, out['metadata']['model']) + for name, prepare_only in [('warp_enqueue', True), ('run_policy', False)] + } + + elif args.model_type == 'vision_multi_policy': + assert args.vision_onnx + vision_runner = OnnxRunner(args.vision_onnx) + out['metadata']['vision'] = make_metadata_dict(args.vision_onnx) + + policy_runners = [] + policy_onnxes = [] + if args.policy_onnx: + policy_onnxes.append(('policy', args.policy_onnx)) + if args.off_policy_onnx: + policy_onnxes.append(('off_policy', args.off_policy_onnx)) + if args.on_policy_onnx: + policy_onnxes.append(('on_policy', args.on_policy_onnx)) + + for name, onnx_path in policy_onnxes: + runner = OnnxRunner(onnx_path) + policy_runners.append(runner) + out['metadata'][name] = make_metadata_dict(onnx_path) + + first_policy_key = policy_onnxes[0][0] + frame_skip = args.frame_skip if args.frame_skip is not None else derive_frame_skip(out['metadata']['vision']['input_shapes'], + out['metadata'][first_policy_key]['input_shapes']) + + for cam_w, cam_h in args.camera_resolutions: + nv12 = NV12Frame(cam_w, cam_h, *get_nv12_info(cam_w, cam_h)) + model_w, model_h = args.model_size + out[(cam_w, cam_h)] = { + name: compile_multi_policy(nv12, model_w, model_h, prepare_only, frame_skip, + vision_runner, policy_runners, + out['metadata']['vision'], out['metadata'][first_policy_key]) + for name, prepare_only in [('warp_enqueue', True), ('run_policy', False)] + } + + with open(args.output, "wb") as f: + pickle.dump(out, f) + pkl_size = os.path.getsize(args.output) + print(f"Saved combined JIT to {args.output} ({pkl_size / 1e6:.2f} MB)") + + from openpilot.common.file_chunker import chunk_file, get_chunk_targets + chunk_targets = get_chunk_targets(args.output, pkl_size) + chunk_file(args.output, chunk_targets) + num_chunks = len(chunk_targets) - 1 + print(f"Chunked into {num_chunks} file(s)") diff --git a/sunnypilot/modeld_v2/modeld.py b/sunnypilot/modeld_v2/modeld.py index f86228618..dfff2e6c2 100755 --- a/sunnypilot/modeld_v2/modeld.py +++ b/sunnypilot/modeld_v2/modeld.py @@ -1,4 +1,11 @@ #!/usr/bin/env python3 +""" +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 os from openpilot.system.hardware import TICI os.environ['DEV'] = 'QCOM' if TICI else 'CPU' @@ -6,6 +13,7 @@ USBGPU = "USBGPU" in os.environ if USBGPU: os.environ['DEV'] = 'AMD' os.environ['AMD_IFACE'] = 'USB' +import pickle import time import numpy as np import cereal.messaging as messaging @@ -26,18 +34,40 @@ from openpilot.selfdrive.controls.lib.drive_helpers import get_accel_from_plan, 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.constants import Plan -from openpilot.sunnypilot.modeld_v2.warp import Warp 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.models.helpers import get_active_bundle -from openpilot.sunnypilot.models.runners.helpers import get_model_runner PROCESS_NAME = "selfdrive.modeld.modeld_tinygrad" +def _pkl_exists(path): + from openpilot.common.file_chunker import get_manifest_path + return os.path.exists(path) or os.path.exists(get_manifest_path(path)) + + +def _find_driving_pkl(bundle): + if (override := os.environ.get('COMBINED_MODEL_PKL')) and _pkl_exists(override): + return override + if bundle is None or not bundle.models: + return None + from openpilot.system.hardware.hw import Paths + model_root = Paths.model_root() + + pkl_name = bundle.models[0].artifact.fileName + pkl_path = os.path.join(model_root, pkl_name) + if _pkl_exists(pkl_path): + return pkl_path + + fallback = os.path.join(model_root, 'driving_tinygrad.pkl') + if _pkl_exists(fallback): + return fallback + return None + + class FrameMeta: frame_id: int = 0 timestamp_sof: int = 0 @@ -49,117 +79,170 @@ class FrameMeta: class ModelState(ModelStateBase): - frames: dict[str, Warp] inputs: dict[str, np.ndarray] - prev_desire: np.ndarray # for tracking the rising edge of the pulse - temporal_idxs: slice | np.ndarray + prev_desire: np.ndarray - def __init__(self): + def __init__(self, cam_w: int, cam_h: int): ModelStateBase.__init__(self) - try: - self.model_runner = get_model_runner() - self.constants = self.model_runner.constants - except Exception as e: - cloudlog.exception(f"Failed to initialize model runner: {str(e)}") - raise - model_bundle = get_active_bundle() + env_pkl = os.environ.get('COMBINED_MODEL_PKL') + if env_pkl and os.path.exists(env_pkl): + model_bundle = None + else: + model_bundle = get_active_bundle() self.generation = model_bundle.generation if model_bundle is not None else None - overrides = {override.key: override.value for override in model_bundle.overrides} + overrides = {override.key: override.value for override in model_bundle.overrides} if model_bundle else {} self.LAT_SMOOTH_SECONDS = float(overrides.get('lat', ".0")) self.LONG_SMOOTH_SECONDS = float(overrides.get('long', ".0")) self.MIN_LAT_CONTROL_SPEED = 0.3 self.PLANPLUS_CONTROL: float = 1.0 - buffer_length = 5 if self.model_runner.is_20hz else 2 - self.warp = Warp(buffer_length) + pkl_path = _find_driving_pkl(model_bundle) + assert pkl_path is not None, "No driving pkl found — 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): + from tinygrad.tensor import Tensor + from openpilot.system.camerad.cameras.nv12_info import get_nv12_info + from openpilot.sunnypilot.modeld_v2.compile_modeld import derive_frame_skip, make_split_input_queues + from tinygrad.device import Device + + from openpilot.common.file_chunker import read_file_chunked + + cloudlog.warning(f"loading combined pkl: {pkl_path}") + jits = pickle.loads(read_file_chunked(pkl_path)) + + self.DEV = Device.DEFAULT + + metadata = jits['metadata'] + if 'model' in metadata: + model_metadata = metadata['model'] + self.vision_output_slices = model_metadata['output_slices'] + self.policy_output_slices = {} + self._policy_slices_list = [] + self._combined_model_type = 'supercombo' + self._vision_input_names = [k for k in model_metadata['input_shapes'] if 'img' in k] + from openpilot.sunnypilot.modeld_v2.compile_modeld import make_supercombo_input_queues + frame_skip = derive_frame_skip({}, model_metadata['input_shapes']) + self.input_queues, self.npy = make_supercombo_input_queues(model_metadata['input_shapes'], frame_skip, device=self.DEV) + else: + vision_metadata = metadata['vision'] + policy_keys = [k for k in metadata if k != 'vision'] + if policy_keys == ['policy']: + self._combined_model_type = 'split' + else: + self._combined_model_type = 'multi_policy' + self.vision_output_slices = vision_metadata['output_slices'] + self._policy_keys = policy_keys + self._policy_slices_list = [metadata[k]['output_slices'] for k in policy_keys] + self.policy_output_slices = self._policy_slices_list[0] + self._has_on_policy = any('on' in k.lower() for k in policy_keys) + first_policy_metadata = metadata[policy_keys[0]] + vision_input_shapes = vision_metadata['input_shapes'] + policy_input_shapes = first_policy_metadata['input_shapes'] + self._vision_input_names = [k for k in vision_input_shapes if 'img' in k] + frame_skip = derive_frame_skip(vision_input_shapes, policy_input_shapes) + self.input_queues, self.npy = make_split_input_queues(vision_input_shapes, policy_input_shapes, frame_skip, device=self.DEV) + + from openpilot.sunnypilot.modeld_v2.parse_model_outputs_split import Parser as SplitParser + from openpilot.sunnypilot.modeld_v2.parse_model_outputs import Parser as CombinedParser + self.parser = SplitParser() if self._combined_model_type != 'supercombo' else CombinedParser() + + is_20hz = bundle.is20hz if bundle else self._combined_model_type in ('split', 'multi_policy') + if is_20hz: + from openpilot.sunnypilot.models.split_model_constants import SplitModelConstants + self.constants = SplitModelConstants() + else: + from openpilot.sunnypilot.modeld_v2.constants import ModelConstants + self.constants = ModelConstants() + self.prev_desire = np.zeros(self.constants.DESIRE_LEN, dtype=np.float32) - self.numpy_inputs = {} - self.temporal_buffers = {} - self.temporal_idxs_map = {} + self.full_frames: dict = {} + self._blob_cache: dict = {} + nv12_info = get_nv12_info(cam_w, cam_h) + self.frame_buf_params = dict.fromkeys(self._vision_input_names, nv12_info) - for key, shape in self.model_runner.input_shapes.items(): - if key not in self.model_runner.vision_input_names: # Policy inputs - self.numpy_inputs[key] = np.zeros(shape, dtype=np.float32) + self._run_policy = jits[(cam_w, cam_h)]['run_policy'] + self._warp_enqueue = jits[(cam_w, cam_h)]['warp_enqueue'] + road_name = next(k for k in self._vision_input_names if 'big' not in k) + yuv_size = self.frame_buf_params[road_name][3] + self._warp_enqueue( + **self.input_queues, + frame=Tensor(np.zeros(yuv_size, dtype=np.uint8), device=self.DEV).contiguous().realize(), + big_frame=Tensor(np.zeros(yuv_size, dtype=np.uint8), device=self.DEV).contiguous().realize()) - # Temporal input: shape is [batch, history, features] - if len(shape) == 3 and shape[1] > 1: - buffer_history_len = shape[1] * 4 if shape[1] < 99 else shape[1] # Allow for higher history buffers in the future - feature_len = shape[2] - features_buffer_shape = self.model_runner.input_shapes.get('features_buffer') - if shape[1] in (24, 25) and features_buffer_shape is not None and features_buffer_shape[1] == 24: # 20Hz - buffer_history_len = (features_buffer_shape[1] + 1) * 4 - step = int(-buffer_history_len / shape[1]) - self.temporal_idxs_map[key] = np.arange(step, step * (shape[1] + 1), step)[::-1] - elif shape[1] == 25: # Split - skip = buffer_history_len // shape[1] - self.temporal_idxs_map[key] = np.arange(buffer_history_len)[-1 - (skip * (shape[1] - 1))::skip] - elif shape[1] >= 99: # non20hz - self.temporal_idxs_map[key] = np.arange(shape[1]) - self.temporal_buffers[key] = np.zeros((1, buffer_history_len, feature_len), dtype=np.float32) @property def mlsim(self) -> bool: return bool(self.generation is not None and self.generation >= 11) + @property + def vision_input_names(self) -> list[str]: + return self._vision_input_names + @property def desire_key(self) -> str: - return next(key for key in self.numpy_inputs if key.startswith('desire')) + return next(k for k in self.npy if k.startswith('desire')) def run(self, bufs: dict[str, VisionBuf], transforms: dict[str, np.ndarray], inputs: dict[str, np.ndarray], prepare_only: bool) -> dict[str, np.ndarray] | None: - # Model decides when action is completed, so desire input is just a pulse triggered on rising edge - inputs[self.desire_key][0] = 0 - new_desire = np.where(inputs[self.desire_key] - self.prev_desire > .99, inputs[self.desire_key], 0) - self.prev_desire[:] = inputs[self.desire_key] - self.temporal_buffers[self.desire_key][0,:-1] = self.temporal_buffers[self.desire_key][0,1:] - self.temporal_buffers[self.desire_key][0,-1] = new_desire + from tinygrad.tensor import Tensor - # Roll buffer and assign based on desire.shape[1] value - if self.temporal_buffers[self.desire_key].shape[1] > self.numpy_inputs[self.desire_key].shape[1]: - skip = self.temporal_buffers[self.desire_key].shape[1] // self.numpy_inputs[self.desire_key].shape[1] - self.numpy_inputs[self.desire_key][:] = (self.temporal_buffers[self.desire_key][0].reshape( - self.numpy_inputs[self.desire_key].shape[0], self.numpy_inputs[self.desire_key].shape[1], skip, -1).max(axis=2)) - else: - self.numpy_inputs[self.desire_key][:] = self.temporal_buffers[self.desire_key][0, self.temporal_idxs_map[self.desire_key]] + for key in bufs.keys(): + ptr = np.frombuffer(bufs[key].data, dtype=np.uint8).ctypes.data + yuv_size = self.frame_buf_params[key][3] + cache_key = (key, ptr) + if cache_key not in self._blob_cache: + self._blob_cache[cache_key] = Tensor.from_blob(ptr, (yuv_size,), dtype='uint8', device=self.DEV) + self.full_frames[key] = self._blob_cache[cache_key] - for key in self.numpy_inputs: - if key in inputs and key not in [self.desire_key]: - self.numpy_inputs[key][:] = inputs[key] + desire_key = self.desire_key + inputs[desire_key][0] = 0 + self.npy[desire_key][:] = np.where(inputs[desire_key] - self.prev_desire > .99, inputs[desire_key], 0) + self.prev_desire[:] = inputs[desire_key] + for key in ('traffic_convention', 'lateral_control_params'): + if key in self.npy and key in inputs: + self.npy[key][:] = inputs[key] - imgs_tensors = self.warp.process(bufs, transforms) - for name, tensor in imgs_tensors.items(): - self.model_runner.inputs[name] = tensor - self.model_runner.prepare_inputs(self.numpy_inputs) + road_key = next(n for n in bufs if 'big' not in n) + wide_key = next(n for n in bufs if 'big' in n) + self.npy['tfm'][:, :] = transforms[road_key].reshape(3, 3) + self.npy['big_tfm'][:, :] = transforms[wide_key].reshape(3, 3) if prepare_only: + self._warp_enqueue(**self.input_queues, frame=self.full_frames[road_key], big_frame=self.full_frames[wide_key]) return None - # Run model inference - outputs = self.model_runner.run_model() + raw_outputs = self._run_policy(**self.input_queues, frame=self.full_frames[road_key], big_frame=self.full_frames[wide_key]) - # Update features_buffer - self.temporal_buffers['features_buffer'][0, :-1] = self.temporal_buffers['features_buffer'][0, 1:] - self.temporal_buffers['features_buffer'][0, -1] = outputs['hidden_state'][0, :] - self.numpy_inputs['features_buffer'][:] = self.temporal_buffers['features_buffer'][0, self.temporal_idxs_map['features_buffer']] + if self._combined_model_type == 'supercombo': + model_output = raw_outputs.numpy().flatten() + sliced = {k: model_output[np.newaxis, v] for k, v in self.vision_output_slices.items()} + outputs = self.parser.parse_outputs(sliced) + else: + vision_output = raw_outputs[0].numpy().flatten() + vision_sliced = {k: vision_output[np.newaxis, v] for k, v in self.vision_output_slices.items()} + outputs = self.parser.parse_vision_outputs(vision_sliced) - if "desired_curvature" in outputs: - input_name_prev = None - if "prev_desired_curv" in self.numpy_inputs.keys(): - input_name_prev = 'prev_desired_curv' - if input_name_prev and input_name_prev in self.temporal_buffers: - self.process_desired_curvature(outputs, input_name_prev) + for i, policy_slices in enumerate(self._policy_slices_list): + policy_output = raw_outputs[i + 1].numpy().flatten() + policy_sliced = {k: policy_output[np.newaxis, v] for k, v in policy_slices.items()} + parsed = self.parser.parse_policy_outputs(policy_sliced) + if 'off' in self._policy_keys[i] and self._has_on_policy: + parsed.pop('plan', None) + outputs.update(parsed) + + if 'planplus' in outputs and 'plan' in outputs: + outputs['plan'] = outputs['plan'] + outputs['planplus'] + + if 'desired_curvature' in outputs and 'prev_desired_curv' in self.npy: + buf = self.npy['prev_desired_curv'] + buf[0, :-1] = buf[0, 1:] + buf[0, -1, :] = outputs['desired_curvature'][0, :] if not self.mlsim else 0 return outputs - def process_desired_curvature(self, outputs, input_name_prev): - self.temporal_buffers[input_name_prev][0,:-1] = self.temporal_buffers[input_name_prev][0,1:] - self.temporal_buffers[input_name_prev][0,-1,:] = outputs['desired_curvature'][0, :] - self.numpy_inputs[input_name_prev][:] = self.temporal_buffers[input_name_prev][0, self.temporal_idxs_map[input_name_prev]] - if self.mlsim: - self.numpy_inputs[input_name_prev][:] = 0*self.temporal_buffers[input_name_prev][0, self.temporal_idxs_map[input_name_prev]] - def get_action_from_model(self, model_output: dict[str, np.ndarray], prev_action: log.ModelDataV2.Action, lat_action_t: float, long_action_t: float, v_ego: float) -> log.ModelDataV2.Action: plan = model_output['plan'][0] @@ -186,10 +269,6 @@ def main(demo=False): setproctitle(PROCESS_NAME) config_realtime_process(7, 54) - cloudlog.warning("loading model") - model = ModelState() - cloudlog.warning("models loaded, modeld starting") - # visionipc clients while True: available_streams = VisionIpcClient.available_streams("camerad", block=False) @@ -213,6 +292,10 @@ def main(demo=False): if use_extra_client: cloudlog.warning(f"connected extra cam with buffer size: {vipc_client_extra.buffer_len} ({vipc_client_extra.width} x {vipc_client_extra.height})") + cloudlog.warning("loading model") + model = ModelState(cam_w=vipc_client_main.width, cam_h=vipc_client_main.height) + cloudlog.warning("models loaded, modeld starting") + # messaging pm = PubMaster(["modelV2", "drivingModelData", "cameraOdometry", "modelDataV2SP"]) sm = SubMaster(["deviceState", "carState", "roadCameraState", "liveCalibration", "driverMonitoringState", "carControl", "liveDelay"]) @@ -246,6 +329,7 @@ def main(demo=False): prev_action = log.ModelDataV2.Action() DH = DesireHelper() + meta_constants = load_meta_constants() while True: # Keep receiving frames until we are at least 1 frame ahead of previous extra frame @@ -318,14 +402,14 @@ def main(demo=False): if prepare_only: cloudlog.error(f"skipping model eval. Dropped {vipc_dropped_frames} frames") - bufs = {name: buf_extra if 'big' in name else buf_main for name in model.model_runner.vision_input_names} - transforms = {name: model_transform_extra if 'big' in name else model_transform_main for name in model.model_runner.vision_input_names} + bufs = {name: buf_extra if 'big' in name else buf_main for name in model.vision_input_names} + transforms = {name: model_transform_extra if 'big' in name else model_transform_main for name in model.vision_input_names} inputs:dict[str, np.ndarray] = { model.desire_key: vec_desire, 'traffic_convention': traffic_convention, } - if "lateral_control_params" in model.numpy_inputs.keys(): + if 'lateral_control_params' in model.npy: inputs['lateral_control_params'] = np.array([v_ego, lat_delay], dtype=np.float32) mt1 = time.perf_counter() @@ -343,7 +427,7 @@ def main(demo=False): prev_action = action fill_model_msg(drivingdata_send, modelv2_send, model_output, action, publish_state, meta_main.frame_id, meta_extra.frame_id, frame_id, - frame_drop_ratio, meta_main.timestamp_eof, model_execution_time, live_calib_seen, load_meta_constants()) + frame_drop_ratio, meta_main.timestamp_eof, model_execution_time, live_calib_seen, meta_constants) desire_state = modelv2_send.modelV2.meta.desireState l_lane_change_prob = desire_state[log.Desire.laneChangeLeft] diff --git a/sunnypilot/modeld_v2/tests/conftest.py b/sunnypilot/modeld_v2/tests/conftest.py new file mode 100644 index 000000000..89e2f2d4c --- /dev/null +++ b/sunnypilot/modeld_v2/tests/conftest.py @@ -0,0 +1,204 @@ +""" +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 pytest + +import openpilot.sunnypilot.models.helpers as helpers +import openpilot.sunnypilot.modeld_v2.modeld as modeld_module +from openpilot.sunnypilot.modeld_v2.constants import ModelConstants +from openpilot.sunnypilot.models.split_model_constants import SplitModelConstants + + +class DummyOverride: + def __init__(self, key, value): + self.key = key + self.value = value + + +class DummyArtifact: + def __init__(self, file_name): + self.fileName = file_name + + +class DummyModelType: + def __init__(self, raw): + self.raw = raw + + +class DummyModel: + def __init__(self, type_str, artifact_file): + self.type = DummyModelType(type_str) + self.artifact = DummyArtifact(artifact_file) + + +class DummyBundle: + 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 + self.models = models or [] + + +class Archetype: + def __init__(self, name, metadata_structure, model_stubs, is_20hz, + expected_model_type, expected_constants_class, expected_parser_module, + expected_desire_key): + self.name = name + self.metadata_structure = metadata_structure + self.model_stubs = model_stubs + self.is_20hz = is_20hz + self.expected_model_type = expected_model_type + self.expected_constants_class = expected_constants_class + self.expected_parser_module = expected_parser_module + self.expected_desire_key = expected_desire_key + + +def _noop_jit(**kwargs): + pass + + +def _make_vision_policy_metadata(vision_input_shapes, policy_input_shapes, + vision_output_slices, policy_output_slices): + return { + 'vision': {'input_shapes': vision_input_shapes, 'output_slices': vision_output_slices}, + 'policy': {'input_shapes': policy_input_shapes, 'output_slices': policy_output_slices}, + } + + +def _make_multi_policy_metadata(vision_input_shapes, policy_input_shapes, + vision_output_slices, policy_output_slices): + return { + 'vision': {'input_shapes': vision_input_shapes, 'output_slices': vision_output_slices}, + 'offPolicy': {'input_shapes': policy_input_shapes, 'output_slices': policy_output_slices}, + } + + +def _make_tri_policy_metadata(vision_input_shapes, policy_input_shapes, + vision_output_slices, policy_output_slices): + return { + 'vision': {'input_shapes': vision_input_shapes, 'output_slices': vision_output_slices}, + 'onPolicy': {'input_shapes': policy_input_shapes, 'output_slices': policy_output_slices}, + 'offPolicy': {'input_shapes': policy_input_shapes, 'output_slices': policy_output_slices}, + } + + +def _make_supercombo_metadata(input_shapes, output_slices): + return {'model': {'input_shapes': input_shapes, 'output_slices': output_slices}} + + +SPLIT_VISION_INPUT_SHAPES = {'img': (1, 12, 128, 256), 'big_img': (1, 12, 128, 256)} +SPLIT_POLICY_INPUT_SHAPES = {'features_buffer': (1, 25, 512), 'desire_pulse': (1, 25, 8), 'traffic_convention': (1, 2)} +SPLIT_VISION_SLICES = {'hidden_state': slice(0, 512), 'pose': slice(512, 524)} +SPLIT_POLICY_SLICES = {'plan': slice(0, 495), 'meta': slice(495, 550)} + +SUPERCOMBO_INPUT_SHAPES = { + 'img': (1, 12, 128, 256), 'big_img': (1, 12, 128, 256), + 'desire': (1, 100, 8), 'features_buffer': (1, 99, 512), + 'lateral_control_params': (1, 2), 'prev_desired_curv': (1, 100, 1), + 'traffic_convention': (1, 2), +} +SUPERCOMBO_SLICES = {'plan': slice(0, 495), 'hidden_state': slice(495, 1007), 'meta': slice(1007, 1062)} + +CAM_W, CAM_H = 1928, 1208 + +ARCHETYPES = { + 'vision_policy_split': Archetype( + name='vision_policy_split', + metadata_structure=_make_vision_policy_metadata( + SPLIT_VISION_INPUT_SHAPES, SPLIT_POLICY_INPUT_SHAPES, + SPLIT_VISION_SLICES, SPLIT_POLICY_SLICES), + model_stubs=[DummyModel('vision', 'driving_test_tinygrad.pkl'), + DummyModel('policy', 'driving_test_tinygrad.pkl')], + is_20hz=True, + expected_model_type='split', + expected_constants_class=SplitModelConstants, + expected_parser_module='parse_model_outputs_split', + expected_desire_key='desire', + ), + 'vision_multi_policy': Archetype( + name='vision_multi_policy', + metadata_structure=_make_multi_policy_metadata( + SPLIT_VISION_INPUT_SHAPES, SPLIT_POLICY_INPUT_SHAPES, + SPLIT_VISION_SLICES, SPLIT_POLICY_SLICES), + model_stubs=[DummyModel('vision', 'driving_test_tinygrad.pkl'), + DummyModel('offPolicy', 'driving_test_tinygrad.pkl')], + is_20hz=True, + expected_model_type='multi_policy', + expected_constants_class=SplitModelConstants, + expected_parser_module='parse_model_outputs_split', + expected_desire_key='desire', + ), + 'tri_policy': Archetype( + name='tri_policy', + metadata_structure=_make_tri_policy_metadata( + SPLIT_VISION_INPUT_SHAPES, SPLIT_POLICY_INPUT_SHAPES, + SPLIT_VISION_SLICES, SPLIT_POLICY_SLICES), + model_stubs=[DummyModel('vision', 'driving_test_tinygrad.pkl'), + DummyModel('onPolicy', 'driving_test_tinygrad.pkl'), + DummyModel('offPolicy', 'driving_test_tinygrad.pkl')], + is_20hz=True, + expected_model_type='multi_policy', + expected_constants_class=SplitModelConstants, + expected_parser_module='parse_model_outputs_split', + expected_desire_key='desire', + ), + 'supercombo_non20hz': Archetype( + name='supercombo_non20hz', + metadata_structure=_make_supercombo_metadata(SUPERCOMBO_INPUT_SHAPES, SUPERCOMBO_SLICES), + model_stubs=[DummyModel('supercombo', 'driving_test_tinygrad.pkl')], + is_20hz=False, + expected_model_type='supercombo', + expected_constants_class=ModelConstants, + expected_parser_module='parse_model_outputs', + expected_desire_key='desire', + ), +} + + +def make_pkl_data(archetype): + return { + 'metadata': archetype.metadata_structure, + (CAM_W, CAM_H): {'run_policy': _noop_jit, 'warp_enqueue': _noop_jit}, + } + + +def write_pkl(tmp_path, archetype): + pkl_path = tmp_path / 'driving_test_tinygrad.pkl' + with open(pkl_path, 'wb') as f: + pickle.dump(make_pkl_data(archetype), f) + return pkl_path + + +def make_bundle(archetype): + return DummyBundle( + models=archetype.model_stubs, + is_20hz=archetype.is_20hz, + ) + + +@pytest.fixture +def patch_modeld(monkeypatch): + def _patch(bundle): + monkeypatch.setattr(helpers, 'get_active_bundle', lambda params=None: bundle, raising=False) + monkeypatch.setattr(modeld_module, 'get_active_bundle', lambda params=None: bundle, raising=False) + + return _patch + + +@pytest.fixture +def model_state_factory(tmp_path, monkeypatch, patch_modeld): + from openpilot.system.hardware import hw + + def _create(archetype): + write_pkl(tmp_path, archetype) + bundle = make_bundle(archetype) + patch_modeld(bundle) + monkeypatch.setattr(hw.Paths, 'model_root', staticmethod(lambda: str(tmp_path))) + return modeld_module.ModelState(cam_w=CAM_W, cam_h=CAM_H) + + return _create diff --git a/sunnypilot/modeld_v2/tests/test_buffer_logic_inspect.py b/sunnypilot/modeld_v2/tests/test_buffer_logic_inspect.py deleted file mode 100644 index 15009c94d..000000000 --- a/sunnypilot/modeld_v2/tests/test_buffer_logic_inspect.py +++ /dev/null @@ -1,263 +0,0 @@ -import numpy as np -import pytest -from typing import Any - -import openpilot.sunnypilot.models.helpers as helpers -import openpilot.sunnypilot.models.runners.helpers as runner_helpers -import openpilot.sunnypilot.modeld_v2.modeld as modeld_module - -ModelState = modeld_module.ModelState - -# These are the shapes extracted/loaded from the model onnx -SHAPE_MODE_PARAMS = [ - ({'desire': (1, 100, 8), 'features_buffer': (1, 99, 512), "nav_features": (1, 256), "nav_instructions": (1, 150)}, 'non20hz'), # Optimus Prime - ({'desire': (1, 100, 8), 'features_buffer': (1, 99, 512), "lat_planner_state": (1, 4),}, 'non20hz'), # farmville - ({'desire': (1, 100, 8), 'features_buffer': (1, 99, 512), "lateral_control_params": (1, 2), "prev_desired_curv": (1, 100, 1)}, 'non20hz'), # wd40 - ({'desire': (1, 100, 8), 'features_buffer': (1, 99, 512), 'prev_desired_curv': (1, 100, 1), "lateral_control_params": (1, 2),}, 'non20hz'), # NTS - ({'desire': (1, 25, 8), 'features_buffer': (1, 24, 512)}, '20hz'), # NPR - ({'desire': (1, 100, 8), 'features_buffer': (1, 99, 512), 'prev_desired_curv': (1, 100, 1), "lateral_control_params": (1, 2),}, 'non20hz'), # NTS - ({'desire': (1, 25, 8), 'features_buffer': (1, 25, 512)}, 'split'), # Steam Powered v2 - ({'desire_pulse': (1, 25, 8), 'features_buffer': (1, 25, 512)}, 'split'), # desire rename -] - - -# This creates a dummy runner, override, and bundle instance for the tests to run, without actually trying to load a physical model. -class DummyOverride: - def __init__(self, key: str, value: str) -> None: - self.key = key - self.value = value - - -class DummyBundle: - def __init__(self) -> None: - self.overrides = [DummyOverride('lat', '.1'), DummyOverride('long', '.3')] - self.generation = 10 # default to non-mlsim for buffer-update tests, as raising to 11 here will zero curvature buffer - - -class DummyModelRunner: - def __init__(self, input_shapes: dict[str, tuple[int, int, int]], constants: Any = None) -> None: - self.input_shapes = input_shapes - self.constants = constants or type('C', (), { - 'FULL_HISTORY_BUFFER_LEN': 100, - 'FEATURE_LEN': 512, - 'DESIRE_LEN': 8, - 'PREV_DESIRED_CURV_LEN': 1, - 'INPUT_HISTORY_BUFFER_LEN': 25, - 'TEMPORAL_SKIP': 4, - })() - self.vision_input_names: list[str] = [] - shape = input_shapes.get('desire', (1, 0, 0)) # [batch, history, features] - if shape[1] == 25: - self.is_20hz = True - else: - self.is_20hz = False - - # Minimal prepare/run methods so ModelState can be run without actually running the model - def prepare_inputs(self, numpy_inputs): - return None - - def run_model(self): - return { - 'hidden_state': np.zeros((1, self.constants.FEATURE_LEN), dtype=np.float32), - 'desired_curvature': np.zeros((1, 1), dtype=np.float32), - } - - -@pytest.fixture -def shapes(request): - return request.param - - -@pytest.fixture -def bundle() -> DummyBundle: - return DummyBundle() - - -@pytest.fixture -def runner(shapes) -> DummyModelRunner: - return DummyModelRunner(shapes) - - -@pytest.fixture -def apply_patches(monkeypatch: pytest.MonkeyPatch, bundle: DummyBundle, runner: DummyModelRunner): - monkeypatch.setattr(helpers, 'get_active_bundle', lambda params=None: bundle, raising=False) - monkeypatch.setattr(runner_helpers, 'get_model_runner', lambda: runner, raising=False) - monkeypatch.setattr(modeld_module, 'get_model_runner', lambda: runner, raising=False) - monkeypatch.setattr(modeld_module, 'get_active_bundle', lambda params=None: bundle, raising=False) - - -# These are expected shapes and indices based on the time the model was presented -def get_expected_indices(shape, constants, mode, key=None): - if mode == 'split': - start = -1 - (constants.TEMPORAL_SKIP * (constants.INPUT_HISTORY_BUFFER_LEN - 1)) - arr = np.arange(constants.FULL_HISTORY_BUFFER_LEN) - idxs = arr[start::constants.TEMPORAL_SKIP] - return idxs - elif mode == '20hz': - num_elements = shape[1] - step_size = int(-100 / num_elements) - idxs = np.arange(step_size, step_size * (num_elements + 1), step_size)[::-1] - return idxs - elif mode == 'non20hz': - return np.arange(shape[1]) - return None - - -@pytest.mark.parametrize("shapes,mode", SHAPE_MODE_PARAMS, indirect=["shapes"]) -def test_buffer_shapes_and_indices(shapes, mode, apply_patches): - state = ModelState() - constants = DummyModelRunner(shapes).constants - for key in shapes: - buf = state.temporal_buffers.get(key, None) - idxs = state.temporal_idxs_map.get(key, None) - if buf is None: - continue # not all shapes are 3D, and the non-3D ones are not buffered - # Buffer shape logic - if mode == 'split': - expected_shape = (1, constants.FULL_HISTORY_BUFFER_LEN, shapes[key][2]) - expected_idxs = get_expected_indices(shapes[key], constants, 'split', key) - elif mode == '20hz': - expected_shape = (1, constants.FULL_HISTORY_BUFFER_LEN, shapes[key][2]) - expected_idxs = get_expected_indices(shapes[key], constants, '20hz', key) - elif mode == 'non20hz': - expected_shape = (1, shapes[key][1], shapes[key][2]) - expected_idxs = get_expected_indices(shapes[key], constants, 'non20hz', key) - - assert buf is not None, f"{key}: buffer not found" - assert buf.shape == expected_shape, f"{key}: buffer shape {buf.shape} != expected {expected_shape}" - if expected_idxs is not None: - assert np.all(idxs == expected_idxs), f"{key}: buffer idxs {idxs} != expected {expected_idxs}" - else: - assert idxs is None or idxs.size == 0, f"{key}: buffer idxs should be None or empty" - - -def legacy_buffer_update(buf, new_val, mode, key, constants, idxs, input_shape, prev_desire=None): - # This is what we compare the new dynamic logic to, to ensure it does the same thing - if mode == 'split': - if key == 'desire' or key.startswith('desire'): - buf[0,:-1] = buf[0,1:] - buf[0,-1] = new_val - return buf.reshape((1, constants.INPUT_HISTORY_BUFFER_LEN, constants.TEMPORAL_SKIP, -1)).max(axis=2) - elif key == 'features_buffer': - buf[0,:-1] = buf[0,1:] - buf[0,-1] = new_val - return buf[0, idxs] - elif key == 'prev_desired_curv': - buf[0,:-1] = buf[0,1:] - buf[0,-1,:] = new_val - return buf[0, idxs] - elif mode == '20hz': - if key == 'desire': - buf[:-1] = buf[1:] - buf[-1] = new_val - reshape_dims = (1, buf.shape[1], -1, buf.shape[2]) - reshaped = buf.reshape(reshape_dims).max(axis=2) - # Slice to last shape[1] elements to match model input shape - input_len = reshaped.shape[1] - model_input_len = 25 # For 20hz mode, desire shape[1] is 25 - if input_len > model_input_len: - reshaped = reshaped[:, -model_input_len:, :] - return reshaped - elif key == 'features_buffer': - buffer_history_len = buf.shape[1] - legacy_buf = np.zeros((buffer_history_len, buf.shape[2]), dtype=np.float32) - legacy_buf[:] = buf[0] - legacy_buf[:-1] = legacy_buf[1:] - legacy_buf[-1] = new_val - return legacy_buf[idxs] - elif key == 'prev_desired_curv': - buffer_history_len = buf.shape[1] - legacy_buf = np.zeros((buffer_history_len, buf.shape[2]), dtype=np.float32) - legacy_buf[:] = buf[0] - legacy_buf[:-1] = legacy_buf[1:] - legacy_buf[-1,:] = new_val - return legacy_buf[idxs] - elif mode == 'non20hz': - if key == 'desire': - desire_len = constants.DESIRE_LEN - if prev_desire is None: - prev_desire = np.zeros(desire_len, dtype=np.float32) - # Set first element to zero - new_val = new_val.copy() - new_val[0] = 0 - # Shift buffer by desire len - buf[0][:-desire_len] = buf[0][desire_len:] - # Only insert new desire if rising edge - buf[0][-desire_len:] = np.where(new_val - prev_desire > 0.99, new_val, 0) - prev_desire[:] = new_val - return buf[0] - elif key == 'features_buffer': - buf[0, :-1] = buf[0, 1:] - buf[0, -1] = new_val - return buf[0, -input_shape[1]:] # (99, 512) - elif key == 'prev_desired_curv': - length = new_val.shape[0] - buf[0,:-length,0] = buf[0,length:,0] - buf[0,-length:,0] = new_val[:length] - return buf[0] - return None - - -def dynamic_buffer_update(state, key, new_val, mode): - if key == 'desire' or key.startswith('desire'): - inputs = {k: np.zeros(v[2], dtype=np.float32) if len(v) == 3 else np.zeros(v[1], dtype=np.float32) - for k, v in state.model_runner.input_shapes.items() if k != key} - inputs[key] = new_val.copy() - # ModelState.run expects desire as a pulse, so we zero the first element. - inputs[key][0] = 0 - state.run({}, {}, inputs, prepare_only=False) - return state.numpy_inputs[key] - - if key == 'features_buffer': - inputs = {k: np.zeros(v[2], dtype=np.float32) if len(v) == 3 else np.zeros(v[1], dtype=np.float32) - for k, v in state.model_runner.input_shapes.items() if k != 'features_buffer'} - def run_model_stub(): - return { - 'hidden_state': np.asarray(new_val, dtype=np.float32).reshape(1, -1), - } - state.model_runner.run_model = run_model_stub - state.run({}, {}, inputs, prepare_only=False) - return state.numpy_inputs['features_buffer'][0] - - if key == 'prev_desired_curv': - inputs = {k: np.zeros(v[2], dtype=np.float32) if len(v) == 3 else np.zeros(v[1], dtype=np.float32) - for k, v in state.model_runner.input_shapes.items() if k != 'prev_desired_curv'} - def run_model_stub(): - return { - 'hidden_state': np.zeros((1, state.constants.FEATURE_LEN), dtype=np.float32), - 'desired_curvature': np.asarray(new_val, dtype=np.float32).reshape(1, -1), - } - state.model_runner.run_model = run_model_stub - state.run({}, {}, inputs, prepare_only=False) - return state.numpy_inputs['prev_desired_curv'][0] - return None - - -@pytest.mark.parametrize("shapes,mode", SHAPE_MODE_PARAMS, indirect=["shapes"]) -@pytest.mark.parametrize("key", ["desire", "features_buffer", "prev_desired_curv"]) -def test_buffer_update_equivalence(shapes, mode, key, apply_patches): - state = ModelState() - if key == "desire": - desire_keys = [k for k in shapes.keys() if k.startswith('desire')] - if desire_keys: - actual_key = desire_keys[0] # Use the first (and likely only) desire key - else: - actual_key = key - - if actual_key not in state.numpy_inputs: - pytest.skip() - - constants = DummyModelRunner(shapes).constants - buf = state.temporal_buffers.get(actual_key, None) - idxs = state.temporal_idxs_map.get(actual_key, None) - input_shape = shapes[actual_key] - prev_desire = np.zeros(constants.DESIRE_LEN, dtype=np.float32) if key == 'desire' else None - - for step in range(20): # multiple steps to ensure history is built up - new_val = np.full((input_shape[2],), step, dtype=np.float32) - expected = legacy_buffer_update(buf, new_val, mode, actual_key, constants, idxs, input_shape, prev_desire) - actual = dynamic_buffer_update(state, actual_key, new_val, mode) - if expected is not None and actual is not None and expected.shape != actual.shape: - if expected.ndim == 2 and actual.ndim == 2 and expected.shape[1] == actual.shape[1]: - expected = expected[-actual.shape[0]:] - assert np.allclose(actual, expected), f"{mode} {actual_key}: dynamic buffer update does not match legacy logic" diff --git a/sunnypilot/modeld_v2/tests/test_combined_pkl_loader.py b/sunnypilot/modeld_v2/tests/test_combined_pkl_loader.py new file mode 100644 index 000000000..d4ef0d476 --- /dev/null +++ b/sunnypilot/modeld_v2/tests/test_combined_pkl_loader.py @@ -0,0 +1,263 @@ +""" +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 pytest + +import openpilot.sunnypilot.models.helpers as helpers +import openpilot.sunnypilot.modeld_v2.modeld as modeld_module +from openpilot.sunnypilot.modeld_v2.modeld import _find_driving_pkl +from openpilot.sunnypilot.modeld_v2.tests.conftest import DummyModel, DummyBundle, ARCHETYPES, CAM_W, CAM_H, \ + SPLIT_VISION_INPUT_SHAPES, SPLIT_POLICY_INPUT_SHAPES + +ModelState = modeld_module.ModelState + + +# Pkl discovery + +class TestFindDrivingPkl: + def test_returns_none_when_no_bundle(self): + assert _find_driving_pkl(None) is None + + def test_returns_none_when_no_models(self): + bundle = DummyBundle(models=[]) + assert _find_driving_pkl(bundle) is None + + def test_returns_none_when_pkl_not_on_disk(self): + bundle = DummyBundle(models=[ + DummyModel('vision', 'driving_fof_tinygrad.pkl'), + DummyModel('policy', 'driving_fof_tinygrad.pkl'), + ]) + assert _find_driving_pkl(bundle) is None + + def test_finds_pkl_by_artifact_name(self, tmp_path, monkeypatch): + (tmp_path / 'driving_fof_tinygrad.pkl').write_bytes(b'fake') + from openpilot.system.hardware import hw + monkeypatch.setattr(hw.Paths, 'model_root', staticmethod(lambda: str(tmp_path))) + + bundle = DummyBundle(models=[ + DummyModel('vision', 'driving_fof_tinygrad.pkl'), + DummyModel('policy', 'driving_fof_tinygrad.pkl'), + ]) + result = _find_driving_pkl(bundle) + assert result is not None + assert 'driving_fof_tinygrad.pkl' in result + + def test_finds_fallback_driving_tinygrad(self, tmp_path, monkeypatch): + (tmp_path / 'driving_tinygrad.pkl').write_bytes(b'fake') + from openpilot.system.hardware import hw + monkeypatch.setattr(hw.Paths, 'model_root', staticmethod(lambda: str(tmp_path))) + + bundle = DummyBundle(models=[DummyModel('vision', 'nonexistent.pkl')]) + result = _find_driving_pkl(bundle) + assert result is not None + assert 'driving_tinygrad.pkl' in result + + +# Init — assertion guard + +class TestModelStateCombinedInit: + def test_asserts_when_no_pkl(self, monkeypatch): + bundle = DummyBundle(models=[], is_20hz=True) + monkeypatch.setattr(helpers, 'get_active_bundle', lambda params=None: bundle, raising=False) + monkeypatch.setattr(modeld_module, 'get_active_bundle', lambda params=None: bundle, raising=False) + with pytest.raises(AssertionError, match="No driving pkl found"): + ModelState(cam_w=CAM_W, cam_h=CAM_H) + + +class TestStockEquivalence: + + def test_split_queue_keys_match_stock(self, model_state_factory): + from openpilot.selfdrive.modeld.compile_modeld import make_input_queues + from openpilot.sunnypilot.modeld_v2.compile_modeld import derive_frame_skip + + state = model_state_factory(ARCHETYPES['vision_policy_split']) + + frame_skip = derive_frame_skip(SPLIT_VISION_INPUT_SHAPES, SPLIT_POLICY_INPUT_SHAPES) + stock_queues, stock_npy = make_input_queues(SPLIT_VISION_INPUT_SHAPES, SPLIT_POLICY_INPUT_SHAPES, frame_skip, + device='NPY') + + # TODO-SP: remove action_t skip once SP adds prerequisite for deep models (action_t input queue) + skip_keys = {'action_t'} + assert set(state.input_queues.keys()) == set(stock_queues.keys()) - skip_keys, \ + f"Queue keys differ: v2={set(state.input_queues.keys())}, stock={set(stock_queues.keys())}" + assert set(state.npy.keys()) == set(stock_npy.keys()) - skip_keys, \ + f"Npy keys differ: v2={set(state.npy.keys())}, stock={set(stock_npy.keys())}" + + def test_split_queue_keys_work_with_desire_key(self, model_state_factory): + from openpilot.sunnypilot.modeld_v2.compile_modeld import derive_frame_skip, make_split_input_queues + + policy_shapes_desire = {'features_buffer': (1, 25, 512), 'desire': (1, 25, 8), 'traffic_convention': (1, 2)} + frame_skip = derive_frame_skip(SPLIT_VISION_INPUT_SHAPES, policy_shapes_desire) + queues, npy = make_split_input_queues(SPLIT_VISION_INPUT_SHAPES, policy_shapes_desire, frame_skip, device='NPY') + + assert 'desire_q' in queues + assert 'desire' in npy + assert 'img_q' in queues + assert 'feat_q' in queues + + def test_split_vision_input_names_match_stock(self, model_state_factory): + state = model_state_factory(ARCHETYPES['vision_policy_split']) + assert state.vision_input_names == ['img', 'big_img'] + + def test_split_output_slices_preserved(self, model_state_factory): + arch = ARCHETYPES['vision_policy_split'] + state = model_state_factory(arch) + assert state.vision_output_slices == arch.metadata_structure['vision']['output_slices'] + assert state.policy_output_slices == arch.metadata_structure['policy']['output_slices'] + + +ARCHETYPE_NAMES = list(ARCHETYPES.keys()) + + +class TestModelTypeDetection: + @pytest.mark.parametrize("archetype_name", ARCHETYPE_NAMES) + def test_combined_model_type(self, archetype_name, model_state_factory): + arch = ARCHETYPES[archetype_name] + state = model_state_factory(arch) + assert state._combined_model_type == arch.expected_model_type, \ + f"{arch.name}: got {state._combined_model_type}, expected {arch.expected_model_type}" + + +class TestConstantsSelection: + @pytest.mark.parametrize("archetype_name", ARCHETYPE_NAMES) + def test_constants_class(self, archetype_name, model_state_factory): + arch = ARCHETYPES[archetype_name] + state = model_state_factory(arch) + assert type(state.constants) is arch.expected_constants_class, \ + f"{arch.name}: got {type(state.constants).__name__}, expected {arch.expected_constants_class.__name__}" + + +class TestParserSelection: + @pytest.mark.parametrize("archetype_name", ARCHETYPE_NAMES) + def test_parser_module(self, archetype_name, model_state_factory): + arch = ARCHETYPES[archetype_name] + state = model_state_factory(arch) + parser_module = type(state.parser).__module__ + assert parser_module.endswith(arch.expected_parser_module), \ + f"{arch.name}: parser from {parser_module}, expected module ending with {arch.expected_parser_module}" + + +class TestDesireKeyDetection: + @pytest.mark.parametrize("archetype_name", ARCHETYPE_NAMES) + def test_desire_key(self, archetype_name, model_state_factory): + arch = ARCHETYPES[archetype_name] + state = model_state_factory(arch) + assert state.desire_key == arch.expected_desire_key, \ + f"{arch.name}: got {state.desire_key}, expected {arch.expected_desire_key}" + + +class TestVisionInputNames: + @pytest.mark.parametrize("archetype_name", ARCHETYPE_NAMES) + def test_vision_names_contain_img(self, archetype_name, model_state_factory): + arch = ARCHETYPES[archetype_name] + state = model_state_factory(arch) + assert len(state.vision_input_names) >= 1 + for name in state.vision_input_names: + assert 'img' in name, f"{arch.name}: vision input name '{name}' missing 'img'" + + +class TestOutputSlices: + @pytest.mark.parametrize("archetype_name", ARCHETYPE_NAMES) + def test_vision_slices_populated(self, archetype_name, model_state_factory): + arch = ARCHETYPES[archetype_name] + state = model_state_factory(arch) + assert len(state.vision_output_slices) > 0, f"{arch.name}: vision_output_slices empty" + + @pytest.mark.parametrize("archetype_name", ARCHETYPE_NAMES) + def test_policy_slices_match_type(self, archetype_name, model_state_factory): + arch = ARCHETYPES[archetype_name] + state = model_state_factory(arch) + if arch.expected_model_type == 'supercombo': + assert state.policy_output_slices == {}, f"{arch.name}: supercombo should have empty policy slices" + else: + assert len(state.policy_output_slices) > 0, f"{arch.name}: split/multi should have policy slices" + + +class TestInputQueueCreation: + @pytest.mark.parametrize("archetype_name", ARCHETYPE_NAMES) + def test_queues_not_empty(self, archetype_name, model_state_factory): + arch = ARCHETYPES[archetype_name] + state = model_state_factory(arch) + assert len(state.input_queues) > 0, f"{arch.name}: input_queues empty" + + @pytest.mark.parametrize("archetype_name", ARCHETYPE_NAMES) + def test_npy_contains_transforms(self, archetype_name, model_state_factory): + arch = ARCHETYPES[archetype_name] + state = model_state_factory(arch) + assert 'tfm' in state.npy, f"{arch.name}: 'tfm' missing from npy" + assert 'big_tfm' in state.npy, f"{arch.name}: 'big_tfm' missing from npy" + assert state.npy['tfm'].shape == (3, 3) + assert state.npy['big_tfm'].shape == (3, 3) + + @pytest.mark.parametrize("archetype_name", ARCHETYPE_NAMES) + def test_npy_contains_desire(self, archetype_name, model_state_factory): + arch = ARCHETYPES[archetype_name] + state = model_state_factory(arch) + assert arch.expected_desire_key in state.npy, \ + f"{arch.name}: '{arch.expected_desire_key}' missing from npy" + + +class TestFrameBufferParams: + @pytest.mark.parametrize("archetype_name", ARCHETYPE_NAMES) + def test_frame_buf_params_per_vision_input(self, archetype_name, model_state_factory): + arch = ARCHETYPES[archetype_name] + state = model_state_factory(arch) + for name in state.vision_input_names: + assert name in state.frame_buf_params, f"{arch.name}: frame_buf_params missing '{name}'" + nv12_info = state.frame_buf_params[name] + assert len(nv12_info) >= 4, f"{arch.name}: nv12_info for '{name}' too short" + + +class TestBundleOverrides: + @pytest.mark.parametrize("archetype_name", ARCHETYPE_NAMES) + def test_smoothing_params_from_overrides(self, archetype_name, model_state_factory): + arch = ARCHETYPES[archetype_name] + state = model_state_factory(arch) + assert state.LAT_SMOOTH_SECONDS == 0.1 + assert state.LONG_SMOOTH_SECONDS == 0.3 + + @pytest.mark.parametrize("archetype_name", ARCHETYPE_NAMES) + def test_generation_from_bundle(self, archetype_name, model_state_factory): + arch = ARCHETYPES[archetype_name] + state = model_state_factory(arch) + assert state.generation == 10 + + +class TestMlsimProperty: + def test_mlsim_false_for_gen10(self, model_state_factory): + state = model_state_factory(ARCHETYPES['supercombo_non20hz']) + assert state.mlsim is False + + def test_mlsim_true_for_gen11(self, tmp_path, monkeypatch, patch_modeld): + from openpilot.system.hardware import hw + from openpilot.sunnypilot.modeld_v2.tests.conftest import write_pkl, ARCHETYPES as A + + arch = A['supercombo_non20hz'] + write_pkl(tmp_path, arch) + bundle = DummyBundle(models=arch.model_stubs, is_20hz=arch.is_20hz, generation=11) + patch_modeld(bundle) + monkeypatch.setattr(hw.Paths, 'model_root', staticmethod(lambda: str(tmp_path))) + + state = ModelState(cam_w=CAM_W, cam_h=CAM_H) + assert state.mlsim is True + + +class TestCrossArchetypeMismatch: + def test_wrong_is_20hz_changes_constants(self, tmp_path, monkeypatch, patch_modeld): + from openpilot.system.hardware import hw + from openpilot.sunnypilot.modeld_v2.tests.conftest import write_pkl + from openpilot.sunnypilot.modeld_v2.constants import ModelConstants + + arch = ARCHETYPES['vision_policy_split'] + write_pkl(tmp_path, arch) + bundle = DummyBundle(models=arch.model_stubs, is_20hz=False) + patch_modeld(bundle) + monkeypatch.setattr(hw.Paths, 'model_root', staticmethod(lambda: str(tmp_path))) + + state = ModelState(cam_w=CAM_W, cam_h=CAM_H) + assert type(state.constants) is ModelConstants, \ + "Wrong is_20hz should produce wrong constants class" diff --git a/sunnypilot/modeld_v2/tests/test_compile_modeld.py b/sunnypilot/modeld_v2/tests/test_compile_modeld.py new file mode 100644 index 000000000..f40467851 --- /dev/null +++ b/sunnypilot/modeld_v2/tests/test_compile_modeld.py @@ -0,0 +1,161 @@ +""" +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 numpy as np +import pytest + +from openpilot.sunnypilot.modeld_v2.compile_modeld import derive_frame_skip, _detect_desire_key + + +class TestDeriveFrameSkip: + def test_non20hz_supercombo(self): + vision = {} + policy = {'features_buffer': (1, 99, 512), 'desire': (1, 100, 8)} + assert derive_frame_skip(vision, policy) == 1 + + def test_20hz_supercombo(self): + vision = {} + policy = {'features_buffer': (1, 24, 512), 'desire': (1, 25, 8)} + assert derive_frame_skip(vision, policy) == 4 + + def test_split_vision_policy(self): + vision = {'img': (1, 12, 128, 256)} + policy = {'features_buffer': (1, 25, 512), 'desire_pulse': (1, 25, 8)} + assert derive_frame_skip(vision, policy) == 4 + + def test_no_features_buffer(self): + assert derive_frame_skip({}, {}) == 1 + + +class TestFrameSkipBufferLengthEquivalence: + @pytest.mark.parametrize("frame_skip,expected_buffer_length", [ + (1, 2), + (4, 5), + ]) + def test_img_buffer_size_matches_warp_buffer_length(self, frame_skip, expected_buffer_length): + n_frames = 2 + img_buf_dim0 = frame_skip * (n_frames - 1) + 1 + assert img_buf_dim0 == expected_buffer_length, \ + f"frame_skip={frame_skip}: img_buf[0]={img_buf_dim0}, expected {expected_buffer_length}" + + @pytest.mark.parametrize("is_20hz,expected_frame_skip,expected_buffer_length", [ + (False, 1, 2), + (True, 4, 5), + ]) + def test_is_20hz_to_frame_skip_to_buffer_length(self, is_20hz, expected_frame_skip, expected_buffer_length): + if is_20hz: + policy_shapes = {'features_buffer': (1, 24, 512)} + else: + policy_shapes = {'features_buffer': (1, 99, 512)} + frame_skip = derive_frame_skip({}, policy_shapes) + assert frame_skip == expected_frame_skip + + n_frames = 2 + img_buf_dim0 = frame_skip * (n_frames - 1) + 1 + assert img_buf_dim0 == expected_buffer_length + + +class TestTemporalSamplingEquivalence: + def test_non20hz_desire_sampling_identity(self): + buf = np.random.randn(100, 1, 8).astype(np.float32) + frame_skip = 1 + sampled = buf[::frame_skip].reshape(-1, 8) + assert sampled.shape == (100, 8) + np.testing.assert_array_equal(sampled, buf[:, 0, :]) + + def test_20hz_desire_sampling_max(self): + buf = np.zeros((100, 1, 8), dtype=np.float32) + buf[99, 0, 3] = 1.0 + frame_skip = 4 + reshaped = buf.reshape(-1, frame_skip, 1, 8).max(axis=1) + sampled = reshaped.reshape(-1, 8) + assert sampled.shape == (25, 8) + assert sampled[24, 3] == 1.0 + assert sampled[23, 3] == 0.0 + + def test_split_features_buffer_sampling_skip(self): + buf = np.arange(100 * 512, dtype=np.float32).reshape(100, 1, 512) + frame_skip = 4 + sampled = buf[::frame_skip].reshape(-1, 512) + assert sampled.shape == (25, 512) + np.testing.assert_array_equal(sampled[0], buf[0, 0]) + np.testing.assert_array_equal(sampled[1], buf[4, 0]) + np.testing.assert_array_equal(sampled[24], buf[96, 0]) + + def test_non20hz_features_buffer_sampling_identity(self): + buf = np.arange(99 * 512, dtype=np.float32).reshape(99, 1, 512) + frame_skip = 1 + sampled = buf[::frame_skip].reshape(-1, 512) + assert sampled.shape == (99, 512) + np.testing.assert_array_equal(sampled, buf[:, 0, :]) + + +class TestTemporalIdxEquivalence: + @pytest.mark.parametrize("mode,desire_shape,fb_shape,frame_skip", [ + ('non20hz', (1, 100, 8), (1, 99, 512), 1), + ('20hz', (1, 25, 8), (1, 24, 512), 4), + ('split', (1, 25, 8), (1, 25, 512), 4), + ]) + def test_features_buffer_idx_equivalence(self, mode, desire_shape, fb_shape, frame_skip): + history = fb_shape[1] + + if mode == 'non20hz': + modelstate_idxs = np.arange(history) + buf_len = history + elif mode == '20hz': + buf_len = (history + 1) * 4 + step = int(-buf_len / history) + modelstate_idxs = np.arange(step, step * (history + 1), step)[::-1] + elif mode == 'split': + buf_len = history * 4 + skip = buf_len // history + modelstate_idxs = np.arange(buf_len)[-1 - (skip * (history - 1))::skip] + + assert len(modelstate_idxs) == fb_shape[1], \ + f"{mode}: ModelState idx count {len(modelstate_idxs)} != input shape {fb_shape[1]}" + + @pytest.mark.parametrize("mode,desire_shape,fb_shape,frame_skip", [ + ('non20hz', (1, 100, 8), (1, 99, 512), 1), + ('20hz', (1, 25, 8), (1, 24, 512), 4), + ('split', (1, 25, 8), (1, 25, 512), 4), + ]) + def test_desire_idx_equivalence(self, mode, desire_shape, fb_shape, frame_skip): + history = desire_shape[1] + + compile_desire_buf_len = frame_skip * history if mode != 'non20hz' else history + compile_sampled_count = compile_desire_buf_len // frame_skip if frame_skip > 1 else compile_desire_buf_len + assert compile_sampled_count == history, \ + f"{mode}: compile desire samples {compile_sampled_count} != model input {history}" + + +class TestDetectDesireKey: + def test_finds_desire(self): + shapes = {'features_buffer': (1, 99, 512), 'desire': (1, 100, 8), 'traffic_convention': (1, 2)} + assert _detect_desire_key(shapes) == 'desire' + + def test_finds_desire_pulse(self): + shapes = {'features_buffer': (1, 25, 512), 'desire_pulse': (1, 25, 8), 'traffic_convention': (1, 2)} + assert _detect_desire_key(shapes) == 'desire_pulse' + + def test_returns_none_when_no_desire(self): + shapes = {'features_buffer': (1, 99, 512), 'traffic_convention': (1, 2)} + assert _detect_desire_key(shapes) is None + + +class TestOutputSlicePreservation: + def test_vision_hidden_state_slice_used_for_features(self): + mock_slices = {'hidden_state': slice(0, 512), 'plan': slice(512, 1024)} + features_slice = mock_slices['hidden_state'] + fake_output = np.random.randn(1, 1024).astype(np.float32) + features = fake_output[:, features_slice] + assert features.shape == (1, 512) + + def test_policy_output_slices_independent(self): + vision_slices = {'hidden_state': slice(0, 512)} + policy_slices = {'plan': slice(0, 495), 'meta': slice(495, 550)} + assert set(vision_slices.keys()) & set(policy_slices.keys()) == set(), \ + "vision and policy slices should not overlap in keys" diff --git a/sunnypilot/modeld_v2/tests/test_warp.py b/sunnypilot/modeld_v2/tests/test_warp.py index daf0dd528..49dc634a4 100644 --- a/sunnypilot/modeld_v2/tests/test_warp.py +++ b/sunnypilot/modeld_v2/tests/test_warp.py @@ -2,7 +2,8 @@ import os os.environ['DEV'] = 'CPU' import pytest import numpy as np -from openpilot.selfdrive.modeld.compile_warp import get_nv12_info, CAMERA_CONFIGS +from openpilot.system.camerad.cameras.nv12_info import get_nv12_info +from openpilot.sunnypilot.modeld_v2.warp import CAMERA_CONFIGS from openpilot.sunnypilot.modeld_v2.warp import Warp, MODEL_W, MODEL_H VISION_NAME_PAIRS = [ # needed to account for supercombos input_imgs diff --git a/sunnypilot/modeld_v2/warp.py b/sunnypilot/modeld_v2/warp.py index fd8be4683..f91e456c0 100644 --- a/sunnypilot/modeld_v2/warp.py +++ b/sunnypilot/modeld_v2/warp.py @@ -7,10 +7,43 @@ from tinygrad.engine.jit import TinyJit from tinygrad.device import Device from openpilot.system.camerad.cameras.nv12_info import get_nv12_info -from openpilot.selfdrive.modeld.compile_warp import ( - CAMERA_CONFIGS, MEDMODEL_INPUT_SIZE, make_frame_prepare, make_update_both_imgs, - warp_pkl_path, -) +from openpilot.common.transformations.model import MEDMODEL_INPUT_SIZE +from openpilot.common.transformations.camera import _ar_ox_fisheye, _os_fisheye +from openpilot.selfdrive.modeld.compile_modeld import NV12Frame, make_frame_prepare as _make_frame_prepare + +CAMERA_CONFIGS = [ + (_ar_ox_fisheye.width, _ar_ox_fisheye.height), + (_os_fisheye.width, _os_fisheye.height), +] + + +def make_frame_prepare(cam_w, cam_h, model_w, model_h): + nv12 = NV12Frame(cam_w, cam_h, *get_nv12_info(cam_w, cam_h)) + return _make_frame_prepare(nv12, model_w, model_h) + + +def warp_pkl_path(w, h): + from openpilot.selfdrive.modeld.helpers import MODELS_DIR + return MODELS_DIR / f'warp_{w}x{h}_tinygrad.pkl' + + +def make_update_img_input(frame_prepare, model_w, model_h): + def update_img_input_tinygrad(tensor, frame, M_inv): + M_inv = M_inv.to(Device.DEFAULT) + new_img = frame_prepare(frame, M_inv) + tensor.assign(tensor[6:].cat(new_img, dim=0).contiguous()) + return Tensor.cat(tensor[:6], tensor[-6:], dim=0).contiguous().reshape(1, 12, model_h//2, model_w//2) + return update_img_input_tinygrad + + +def make_update_both_imgs(frame_prepare, model_w, model_h): + update_img = make_update_img_input(frame_prepare, model_w, model_h) + def update_both_imgs_tinygrad(calib_img_buffer, new_img, M_inv, + calib_big_img_buffer, new_big_img, M_inv_big): + calib_img_pair = update_img(calib_img_buffer, new_img, M_inv) + calib_big_img_pair = update_img(calib_big_img_buffer, new_big_img, M_inv_big) + return calib_img_pair, calib_big_img_pair + return update_both_imgs_tinygrad MODELS_DIR = Path(__file__).parent / 'models' MODEL_W, MODEL_H = MEDMODEL_INPUT_SIZE @@ -58,7 +91,17 @@ def compile_v2_warp(cam_w, cam_h, buffer_length): print(f" Saved to {pkl_path}") jit = pickle.load(open(pkl_path, "rb")) - jit(*inputs) + verify_frame = np.random.randint(0, 256, yuv_size, dtype=np.uint8) + verify_big_frame = np.random.randint(0, 256, yuv_size, dtype=np.uint8) + fresh_inputs = [ + Tensor.zeros(img_buffer_shape, dtype='uint8').contiguous().realize(), + Tensor.from_blob(verify_frame.ctypes.data, (yuv_size,), dtype='uint8').realize(), + Tensor(Tensor.randn(3, 3).mul(8).realize().numpy(), device='NPY'), + Tensor.zeros(img_buffer_shape, dtype='uint8').contiguous().realize(), + Tensor.from_blob(verify_big_frame.ctypes.data, (yuv_size,), dtype='uint8').realize(), + Tensor(Tensor.randn(3, 3).mul(8).realize().numpy(), device='NPY'), + ] + jit(*fresh_inputs) class Warp: diff --git a/sunnypilot/models/fetcher.py b/sunnypilot/models/fetcher.py index 0b6853da8..6484f2b44 100644 --- a/sunnypilot/models/fetcher.py +++ b/sunnypilot/models/fetcher.py @@ -110,13 +110,13 @@ class ModelCache: def set(self, data: dict) -> None: """Updates the cache with new model data""" - self.params.put(self._CACHE_KEY, data) - self.params.put(self._LAST_SYNC_KEY, int(time.monotonic() * 1e9)) + self.params.put(self._CACHE_KEY, data, block=True) + self.params.put(self._LAST_SYNC_KEY, int(time.monotonic() * 1e9), block=True) 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_v16.json" + MODEL_URL = "https://raw.githubusercontent.com/sunnypilot/sunnypilot-models/refs/heads/gh-pages/docs/driving_models_v17.json" def __init__(self, params: Params): self.params = params diff --git a/sunnypilot/models/helpers.py b/sunnypilot/models/helpers.py index 562708031..c9344f54e 100644 --- a/sunnypilot/models/helpers.py +++ b/sunnypilot/models/helpers.py @@ -6,7 +6,6 @@ See the LICENSE.md file in the root directory for more details. """ import hashlib -import os import pickle import numpy as np @@ -28,16 +27,12 @@ ModelManager = custom.ModelManagerSP async def verify_file(file_path: str, expected_hash: str) -> bool: - """Verifies file hash against expected hash""" - if not os.path.exists(file_path): + from openpilot.common.file_chunker import read_file_chunked + try: + data = read_file_chunked(file_path) + except FileNotFoundError: return False - - sha256_hash = hashlib.sha256() - with open(file_path, "rb") as file: - for chunk in iter(lambda: file.read(4096), b""): - sha256_hash.update(chunk) - - return sha256_hash.hexdigest().lower() == expected_hash.lower() + return hashlib.sha256(data).hexdigest().lower() == expected_hash.lower() def is_bundle_version_compatible(bundle: dict) -> bool: @@ -75,32 +70,13 @@ def get_active_bundle(params: Params = None) -> custom.ModelManagerSP.ModelBundl return None -def get_active_model_runner(params: Params = None, force_check=False) -> custom.ModelManagerSP.Runner: - """ - Determines and returns the active model runner type, based on provided parameters. - The function utilizes caching to prevent redundant calculations and checks. - - If the cached "ModelRunnerTypeCache" exists in the provided parameters and `force_check` - is set to False, the cached value is directly returned. Otherwise, the function determines - the runner type based on the active model bundle. If a model bundle containing a drive - model exists, the runner type is derived based on the filename of the drive model. - Finally, it updates the cache with the determined runner type, if needed. - - :param params: The parameter set used to retrieve caching and runner details. If `None`, - a default `Params` instance is created internally. - :type params: Params - :param force_check: A flag indicating whether to bypass cached results and always - re-determine the runner type. Defaults to `False`. - :type force_check: bool - :return: The determined or cached model runner type. - :rtype: custom.ModelManagerSP.Runner - """ +def get_active_model_runner(params: Params = None, force_check=False) -> int: if params is None: params = Params() - if (cached_runner_type := params.get("ModelRunnerTypeCache")) and not force_check: - if isinstance(cached_runner_type, str) and cached_runner_type.isdigit(): - return int(cached_runner_type) + cached_runner_type = params.get("ModelRunnerTypeCache") + if cached_runner_type is not None and not force_check: + return cached_runner_type runner_type = custom.ModelManagerSP.Runner.stock @@ -108,7 +84,7 @@ def get_active_model_runner(params: Params = None, force_check=False) -> custom. runner_type = active_bundle.runner.raw if cached_runner_type != runner_type: - params.put("ModelRunnerTypeCache", int(runner_type)) + params.put("ModelRunnerTypeCache", int(runner_type), block=True) return runner_type diff --git a/sunnypilot/models/manager.py b/sunnypilot/models/manager.py index 8fee0798b..b5ccb736c 100644 --- a/sunnypilot/models/manager.py +++ b/sunnypilot/models/manager.py @@ -33,6 +33,17 @@ class ModelManagerSP: self._chunk_size = 128 * 1000 # 128 KB chunks self._download_start_times: dict[str, float] = {} # Track start time per model + def _sync_artifact_progress(self, source_artifact) -> None: + """Mirror download progress to all artifacts sharing the same filename in the selected bundle.""" + if not self.selected_bundle: + return + for model in self.selected_bundle.models: + for artifact in (model.artifact, model.metadata): + if artifact is not source_artifact and artifact.fileName == source_artifact.fileName: + artifact.downloadProgress.status = source_artifact.downloadProgress.status + artifact.downloadProgress.progress = source_artifact.downloadProgress.progress + artifact.downloadProgress.eta = source_artifact.downloadProgress.eta + def _calculate_eta(self, filename: str, progress: float) -> int: """Calculate ETA based on elapsed time and current progress""" if filename not in self._download_start_times or progress <= 0: @@ -63,7 +74,7 @@ class ModelManagerSP: f.write(chunk) bytes_downloaded += len(chunk) - if not self.params.get("ModelManager_DownloadIndex"): + if self.params.get("ModelManager_DownloadIndex") is None: raise Exception("Download cancelled") if total_size > 0: @@ -71,13 +82,55 @@ class ModelManagerSP: model.downloadProgress.status = custom.ModelManagerSP.DownloadStatus.downloading model.downloadProgress.progress = progress model.downloadProgress.eta = self._calculate_eta(model.fileName, progress) + self._sync_artifact_progress(model) self._report_status() # 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) -> None: + from openpilot.common.file_chunker import get_manifest_path, get_chunk_name + manifest_url = get_manifest_path(base_url) + manifest_path = get_manifest_path(base_path) + + async with aiohttp.ClientSession() as session: + async with session.get(manifest_url) as resp: + if resp.status == 404: + raise FileNotFoundError + resp.raise_for_status() + num_chunks = int((await resp.read()).strip()) + + self._download_start_times[artifact.fileName] = time.monotonic() + + for i in range(num_chunks): + chunk_url = get_chunk_name(base_url, i, num_chunks) + chunk_path = get_chunk_name(base_path, i, num_chunks) + chunk_downloaded = 0 + async with aiohttp.ClientSession() as session: + async with session.get(chunk_url) as response: + response.raise_for_status() + chunk_size = int(response.headers.get("content-length", 0)) + with open(chunk_path, 'wb') as f: + async for data in response.content.iter_chunked(self._chunk_size): + f.write(data) + chunk_downloaded += len(data) + if self.params.get("ModelManager_DownloadIndex") is None: + raise Exception("Download cancelled") + intra = chunk_downloaded / max(chunk_size, 1) + progress = min(99, (i + 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() + + with open(manifest_path, 'w') as f: + f.write(str(num_chunks)) + if os.path.isfile(base_path): + os.remove(base_path) + del self._download_start_times[artifact.fileName] + async def _process_artifact(self, artifact, destination_path: str) -> None: - """Processes a single model download including verification""" if not artifact.downloadUri.uri: return None @@ -87,32 +140,38 @@ class ModelManagerSP: full_path = os.path.join(destination_path, filename) try: - # Check existing file - if os.path.exists(full_path) and await verify_file(full_path, expected_hash): + if await verify_file(full_path, expected_hash): artifact.downloadProgress.status = custom.ModelManagerSP.DownloadStatus.cached artifact.downloadProgress.progress = 100 artifact.downloadProgress.eta = 0 + self._sync_artifact_progress(artifact) self._report_status() return - # Download and verify - await self._download_file(url, full_path, artifact) + try: + await self._download_chunked(url, full_path, artifact) + except (FileNotFoundError, aiohttp.ClientResponseError): + 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 artifact.downloadProgress.eta = 0 + self._sync_artifact_progress(artifact) self._report_status() except Exception as e: cloudlog.error(f"Error downloading {filename}: {str(e)}") - if os.path.exists(full_path): - os.remove(full_path) + 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): + os.remove(f) artifact.downloadProgress.status = custom.ModelManagerSP.DownloadStatus.failed artifact.downloadProgress.eta = 0 + self._sync_artifact_progress(artifact) self.selected_bundle.status = custom.ModelManagerSP.DownloadStatus.failed self._report_status() - # Clean up start time if it exists self._download_start_times.pop(artifact.fileName, None) raise @@ -144,11 +203,22 @@ class ModelManagerSP: os.makedirs(destination_path, exist_ok=True) try: - tasks = [self._process_model(model, destination_path) for model in self.selected_bundle.models] - await asyncio.gather(*tasks) + seen_artifacts: set[str] = set() + for model in self.selected_bundle.models: + for artifact in (model.metadata, model.artifact): + if not artifact.fileName: + continue + if artifact.fileName in seen_artifacts: + artifact.downloadProgress.status = custom.ModelManagerSP.DownloadStatus.cached + artifact.downloadProgress.progress = 100 + artifact.downloadProgress.eta = 0 + else: + seen_artifacts.add(artifact.fileName) + await self._process_artifact(artifact, destination_path) + self.active_bundle = self.selected_bundle self.active_bundle.status = custom.ModelManagerSP.DownloadStatus.downloaded - self.params.put("ModelManager_ActiveBundle", self.active_bundle.to_dict()) + self.params.put("ModelManager_ActiveBundle", self.active_bundle.to_dict(), block=True) self.selected_bundle = None except Exception: @@ -206,11 +276,12 @@ class ModelManagerSP: if hasattr(model, 'metadata') and model.metadata.fileName: active_files.append(model.metadata.fileName) - # Remove all files except active ones + # Remove all files except active ones (including their chunk files) model_dir = Paths.model_root() try: for filename in os.listdir(model_dir): - if filename not in active_files: + base = filename.split('.chunk')[0] if '.chunk' in filename else filename + if base not in active_files and 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/sunnypilot/models/model_name.py b/sunnypilot/models/model_name.py index 2d9c54976..02a6c2bac 100644 --- a/sunnypilot/models/model_name.py +++ b/sunnypilot/models/model_name.py @@ -1 +1 @@ -DEFAULT_MODEL = "POP model" +DEFAULT_MODEL = "CD210" diff --git a/sunnypilot/models/tests/model_hash b/sunnypilot/models/tests/model_hash index eaf923358..f363f8309 100644 --- a/sunnypilot/models/tests/model_hash +++ b/sunnypilot/models/tests/model_hash @@ -1 +1 @@ -5d4d21f1899de21137f69d74a4602c44cc5a6b04cf4e4aa9d0ec9206f8c30350 \ No newline at end of file +32f57bdc91f910df1f48ddae7c59aaf6e751f9df6756da481a210577dbce8bcf \ No newline at end of file