diff --git a/.github/workflows/build-default-models.yaml b/.github/workflows/build-default-models.yaml index c51bf5cd74..f75c826851 100644 --- a/.github/workflows/build-default-models.yaml +++ b/.github/workflows/build-default-models.yaml @@ -117,20 +117,6 @@ jobs: ${{ github.workspace }}/${{ env.SMALL_PKL }} \ --device-input "*" --out-of-band --benchmark-runs 1 - - name: Chunk small pkl - run: | - source ${UV_PROJECT_ENVIRONMENT}/bin/activate - export PYTHONPATH=${{ github.workspace }} - python3 -c " - from openpilot.common.file_chunker import chunk_file, get_chunk_targets - import os - pkl = '${{ github.workspace }}/${{ env.SMALL_PKL }}' - size = os.path.getsize(pkl) - targets = get_chunk_targets(pkl, size) - chunk_file(pkl, targets) - print(f'Chunked into {len(targets)} files') - " - - name: Prepare output env: MODEL_NAME: ${{ needs.resolve.outputs.safe_model_name }} @@ -142,8 +128,7 @@ jobs: PKL_BASE="driving_tinygrad.pkl" mkdir -p "$OUTPUT_DIR" - cp "$MODELS_DIR/${PKL_BASE}".chunk* "$OUTPUT_DIR/" - cp "$MODELS_DIR/${PKL_BASE}.chunkmanifest" "$OUTPUT_DIR/" + cp "$MODELS_DIR/${PKL_BASE}" "$OUTPUT_DIR/" ONNX_SHA256=$(git cat-file blob "$(git ls-files -s "${{ github.workspace }}/${{ env.SMALL_ONNX }}" | awk '{print $2}')" | grep '^oid sha256:' | cut -d: -f2) @@ -221,20 +206,6 @@ jobs: raise RuntimeError('Chestnut PCIe link not ready after 10 attempts') " - - name: Chunk big pkl - run: | - source ${UV_PROJECT_ENVIRONMENT}/bin/activate - export PYTHONPATH=${{ github.workspace }} - python3 -c " - from openpilot.common.file_chunker import chunk_file, get_chunk_targets - import os - pkl = '${{ github.workspace }}/${{ env.BIG_PKL }}' - size = os.path.getsize(pkl) - targets = get_chunk_targets(pkl, size) - chunk_file(pkl, targets) - print(f'Chunked into {len(targets)} files') - " - - name: Compile big warp run: | source ${UV_PROJECT_ENVIRONMENT}/bin/activate @@ -274,8 +245,7 @@ jobs: PKL_BASE="big_driving_tinygrad.pkl" mkdir -p "$OUTPUT_DIR" - cp "$MODELS_DIR/${PKL_BASE}".chunk* "$OUTPUT_DIR/" - cp "$MODELS_DIR/${PKL_BASE}.chunkmanifest" "$OUTPUT_DIR/" + cp "$MODELS_DIR/${PKL_BASE}" "$OUTPUT_DIR/" PKL_LFS_SHA256=$(git cat-file blob "$(git ls-files -s "${{ github.workspace }}/${{ env.BIG_PKL }}" | awk '{print $2}')" | grep '^oid sha256:' | cut -d: -f2) @@ -390,24 +360,14 @@ jobs: from datetime import datetime, UTC dm_dir = Path('dm_output') - manifest = list(dm_dir.glob('*.chunkmanifest')) - assert manifest, 'No chunkmanifest found' - pkl_name = manifest[0].name.removesuffix('.chunkmanifest') - num_chunks = int(manifest[0].read_text().strip()) - - chunks = [] - for i in range(num_chunks): - chunk = dm_dir / f'{pkl_name}.chunk{i+1:02d}of{num_chunks:02d}' - chunks.append({ - 'file_name': chunk.name, - 'sha256': hashlib.sha256(chunk.read_bytes()).hexdigest() - }) + matches = sorted(dm_dir.glob('dmonitoring_*_tinygrad.pkl')) + assert len(matches) == 1, f'expected one DM pkl in {dm_dir}, found {matches}' + pkl_name = matches[0].name digest = hashlib.sha256() - for c in chunks: - with open(dm_dir / c['file_name'], 'rb') as f: - while block := f.read(1024*1024): - digest.update(block) + with open(dm_dir / pkl_name, 'rb') as f: + while block := f.read(1024*1024): + digest.update(block) metadata = { 'bundles': [{ @@ -421,8 +381,7 @@ jobs: 'type': 'chunked', 'artifact': { 'file_name': pkl_name, - 'download_uri': {'url': '', 'sha256': digest.hexdigest()}, - 'chunks': chunks + 'download_uri': {'url': '', 'sha256': digest.hexdigest()} } }] }] @@ -485,20 +444,6 @@ jobs: ${{ github.workspace }}/${{ env.DM_ONNX }} \ ${{ github.workspace }}/${{ env.DM_PKL }} - - name: Chunk DM pkl - run: | - source ${UV_PROJECT_ENVIRONMENT}/bin/activate - export PYTHONPATH=${{ github.workspace }} - python3 -c " - from openpilot.common.file_chunker import chunk_file, get_chunk_targets - import os - pkl = '${{ github.workspace }}/${{ env.DM_PKL }}' - size = os.path.getsize(pkl) - targets = get_chunk_targets(pkl, size) - chunk_file(pkl, targets) - print(f'Chunked {pkl} into {len(targets)} chunks') - " - - name: Compile DM warp run: | source ${UV_PROJECT_ENVIRONMENT}/bin/activate @@ -531,8 +476,7 @@ jobs: - name: Prepare DM output run: | mkdir -p dm_output - cp ${{ github.workspace }}/${{ env.DM_PKL }}.chunk* dm_output/ - cp ${{ github.workspace }}/${{ env.DM_PKL }}.chunkmanifest dm_output/ + cp ${{ github.workspace }}/${{ env.DM_PKL }} dm_output/ cp ${{ github.workspace }}/openpilot/selfdrive/modeld/models/dm_warp_* dm_output/ - name: Upload DM artifact diff --git a/.github/workflows/build-single-tinygrad-model.yaml b/.github/workflows/build-single-tinygrad-model.yaml index 4badd9a224..077077a714 100644 --- a/.github/workflows/build-single-tinygrad-model.yaml +++ b/.github/workflows/build-single-tinygrad-model.yaml @@ -137,6 +137,7 @@ jobs: needs: build_model runs-on: ubuntu-latest permissions: + actions: read id-token: write contents: write steps: @@ -179,18 +180,21 @@ jobs: echo "artifact_name=$ARTIFACT_NAME" >> $GITHUB_OUTPUT - name: Download and extract model artifact - uses: actions/download-artifact@v4 - with: - name: ${{ steps.read-artifact-name.outputs.artifact_name }} - path: output + env: + GH_TOKEN: ${{ github.token }} + ARTIFACT_NAME: ${{ steps.read-artifact-name.outputs.artifact_name }} + run: | + id=$(gh api -X GET "repos/$GITHUB_REPOSITORY/actions/runs/$GITHUB_RUN_ID/artifacts" -f name="$ARTIFACT_NAME" --jq '.artifacts[0].id') + url=$(curl -sf -o /dev/null -w '%{redirect_url}' -H "Authorization: Bearer $GH_TOKEN" "https://api.github.com/repos/$GITHUB_REPOSITORY/actions/artifacts/$id/zip") + sudo apt-get install -y -qq aria2 > /dev/null + aria2c -q -x16 -s16 -k8M --max-tries=5 --file-allocation=none -o model.zip "$url" && unzip -q model.zip -d output && rm model.zip - name: Create models folder env: ARTIFACT_NAME: ${{ steps.read-artifact-name.outputs.artifact_name }} run: | mkdir -p "local_models/${RECOMPILED_DIR}/${ARTIFACT_NAME}" - cp -r output/* "local_models/${RECOMPILED_DIR}/${ARTIFACT_NAME}/" - rm -f "local_models/${RECOMPILED_DIR}/${ARTIFACT_NAME}/artifact_name.txt" + cp output/metadata.json "local_models/${RECOMPILED_DIR}/${ARTIFACT_NAME}/" - name: Upload to Hugging Face env: diff --git a/.github/workflows/download-hf-model-chunks/action.yml b/.github/workflows/download-hf-model/action.yml similarity index 74% rename from .github/workflows/download-hf-model-chunks/action.yml rename to .github/workflows/download-hf-model/action.yml index 2eb87cfbb7..863d81c89f 100644 --- a/.github/workflows/download-hf-model-chunks/action.yml +++ b/.github/workflows/download-hf-model/action.yml @@ -1,5 +1,5 @@ -name: Download HF model chunks -description: Resolve and download model chunks from HuggingFace in parallel +name: Download HF model +description: Resolve and download models from HuggingFace in parallel inputs: hf_repo: @@ -19,7 +19,7 @@ inputs: runs: using: composite steps: - - name: Download model chunks + - name: Download models shell: bash env: HF_REPO: ${{ inputs.hf_repo }} @@ -33,7 +33,7 @@ runs: resolve_chunks() { local HF_PATH="$1" ONNX_HASH="$2" CANONICAL="$3" DEST_DIR="$4" local JSON_URL="https://huggingface.co/datasets/${HF_REPO}/resolve/main/${HF_PATH}/default_models.json" - local DEFAULTS BUNDLE ARTIFACT BASE_URL NUM_CHUNKS + local DEFAULTS BUNDLE ARTIFACT FILE_URL BASE_URL DEFAULTS=$(curl -fsSL "$JSON_URL") if [ -n "$TINYGRAD_REF" ]; then BUNDLE=$(echo "$DEFAULTS" | jq --arg hash "$ONNX_HASH" --arg tg "$TINYGRAD_REF" ' @@ -51,20 +51,12 @@ runs: fi ARTIFACT=$(echo "$BUNDLE" | jq -r '.models[0].artifact') - BASE_URL=$(echo "$ARTIFACT" | jq -r '.download_uri.url' | sed 's|/[^/]*$||') - NUM_CHUNKS=$(echo "$ARTIFACT" | jq -r '.chunks | length') + FILE_URL=$(echo "$ARTIFACT" | jq -r '.download_uri.url') + BASE_URL=$(echo "$FILE_URL" | sed 's|/[^/]*$||') mkdir -p "$DEST_DIR" - while IFS= read -r CHUNK_NAME; do - CHUNK_IDX=$(echo "$CHUNK_NAME" | grep -oP 'chunk\K[0-9]+of[0-9]+' || true) - if [ -z "$CHUNK_IDX" ]; then - echo "::error::Failed to parse chunk index from: $CHUNK_NAME" - return 1 - fi - ENCODED_URL=$(python3 -c "import urllib.parse; print(urllib.parse.quote('${BASE_URL}/${CHUNK_NAME}', safe=':/'))") - printf '%s\t%s\n' "$ENCODED_URL" "${DEST_DIR}/${CANONICAL}.chunk${CHUNK_IDX}" >> "$DOWNLOAD_LIST" - done < <(echo "$ARTIFACT" | jq -r '.chunks[].file_name') - echo "$NUM_CHUNKS" > "${DEST_DIR}/${CANONICAL}.chunkmanifest" + ENCODED_URL=$(python3 -c "import urllib.parse; print(urllib.parse.quote('${FILE_URL}', safe=':/'))") + printf '%s\t%s\n' "$ENCODED_URL" "${DEST_DIR}/${CANONICAL}" >> "$DOWNLOAD_LIST" if [ "$CANONICAL" = "dmonitoring_model_tinygrad.pkl" ]; then for warp in dm_warp_1928x1208_tinygrad.pkl dm_warp_1344x760_tinygrad.pkl; do @@ -87,7 +79,7 @@ runs: done TOTAL=$(wc -l < "$DOWNLOAD_LIST") - echo "Downloading $TOTAL chunks with 8 parallel connections..." + echo "Downloading $TOTAL files with 8 parallel connections..." xargs -P8 -d'\n' -I{} bash -c ' URL="${1%% *}" DEST="${1#* }" diff --git a/.github/workflows/sunnypilot-build-model.yaml b/.github/workflows/sunnypilot-build-model.yaml index 199c63f7ae..cf28033a58 100644 --- a/.github/workflows/sunnypilot-build-model.yaml +++ b/.github/workflows/sunnypilot-build-model.yaml @@ -149,7 +149,17 @@ jobs: git lfs pull -X "**/selfdrive/modeld/models/big_*.onnx,**/selfdrive/modeld/models/dmonitoring_*.onnx,**/selfdrive/modeld/models/big_*.pkl,**/selfdrive/modeld/models/dmonitoring_*.pkl" rm -f selfdrive/modeld/models/big_*.onnx selfdrive/modeld/models/dmonitoring_*.onnx selfdrive/modeld/models/big_*.pkl selfdrive/modeld/models/dmonitoring_*.pkl else - git lfs pull -I "**/selfdrive/modeld/models/big_*.onnx,**/selfdrive/modeld/models/big_*.pkl" -X "" + sudo apt-get install -y -qq aria2 > /dev/null # Parallel ranges because gh is flaky + LFS_URL=$(git config -f "$(git rev-parse --show-toplevel)/.lfsconfig" lfs.url) + for f in selfdrive/modeld/models/big_*.onnx selfdrive/modeld/models/big_*.pkl; do + grep -qs "^version https://git-lfs" "$f" || continue + OID=$(awk -F'sha256:' '/^oid/{print $2}' "$f"); SIZE=$(awk '/^size/{print $2}' "$f") + HREF=$(curl -fsSL --retry 3 -X POST -H "Accept: application/vnd.git-lfs+json" -H "Content-Type: application/vnd.git-lfs+json" "$LFS_URL/objects/batch" \ + -d "{\"operation\":\"download\",\"transfer\":[\"basic\"],\"objects\":[{\"oid\":\"$OID\",\"size\":$SIZE}]}" | jq -r '.objects[0].actions.download.href') + aria2c -x 16 -s 16 -k 16M --max-tries=5 --retry-wait=5 --file-allocation=none --summary-interval=30 --allow-overwrite=true \ + -d "$(dirname "$f")" -o "$(basename "$f")" "$HREF" + echo "$OID $f" | sha256sum -c - + done find selfdrive/modeld/models -type f \( -name "*.onnx" -o -name "*.pkl" \) ! -name "big_*.onnx" ! -name "big_*.pkl" -delete fi fi @@ -162,6 +172,7 @@ jobs: uses: actions/upload-artifact@v4 with: name: models-${{ env.REF }}${{ inputs.artifact_suffix }} + compression-level: 0 path: | ${{ github.workspace }}/openpilot/openpilot/selfdrive/modeld/models/*.onnx ${{ github.workspace }}/openpilot/openpilot/selfdrive/modeld/models/*.pkl @@ -224,11 +235,11 @@ jobs: 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_QCOM="DEV=QCOM IMAGE=1 FLOAT16=1 NOLOCALS=1 JIT_BATCH_SIZE=0 OPENPILOT_HACKS=1" + TG_FLAGS_QCOM="DEV=QCOM IMAGE=1 FLOAT16=1 NOLOCALS=1 OPENPILOT_HACKS=1" if [ "${{ inputs.target_hardware }}" == "chestnut" ]; then echo "CHESTNUT build" export CHESTNUT=1 - TG_FLAGS="DEBUG=1 DEV=USB+AMD:LLVM FRAME_DEV=CPU FLOAT16=1 JIT_BATCH_SIZE=0 GMMU=0 TC_OPT=2 TC_MIN_GLOBALS=32" + TG_FLAGS="DEBUG=1 DEV=USB+AMD:LLVM FRAME_DEV=CPU FLOAT16=1 GMMU=0 TC_OPT=2 TC_MIN_GLOBALS=32" OUTPUT_PKL="${{ env.MODELS_DIR }}/big_driving_tinygrad.pkl" else echo "QCOM build" @@ -240,14 +251,15 @@ jobs: if [ -f "${{ env.MODELS_DIR }}/new_arch.txt" ]; then echo "Detected new arch" - DOWNLOADED_ONNX=$(find "${{ env.MODELS_DIR }}" -maxdepth 1 -name "*driving*.onnx" -print -quit) + if [ "${{ inputs.target_hardware }}" == "chestnut" ]; then + DOWNLOADED_ONNX=$(find "${{ env.MODELS_DIR }}" -maxdepth 1 -name "big_*driving*.onnx" -print -quit) + else + DOWNLOADED_ONNX=$(find "${{ env.MODELS_DIR }}" -maxdepth 1 -name "*driving*.onnx" ! -name "big_*" -print -quit) + fi if [ -n "$DOWNLOADED_ONNX" ]; then echo "Compiling ONNX: $DOWNLOADED_ONNX" env ${TG_FLAGS} taskset -c 7 python3 openpilot/sunnypilot/modeld_v2/compile_modeld_v2.py \ - "$DOWNLOADED_ONNX" "$OUTPUT_PKL" --device-input "*" --out-of-band --benchmark-runs 1 - - echo "Chunking pkl" - python3 -c "from openpilot.common.file_chunker import chunk_file, get_chunk_targets; import os; p='$OUTPUT_PKL'; chunk_file(p, get_chunk_targets(p, os.path.getsize(p)))" + "$DOWNLOADED_ONNX" "$OUTPUT_PKL" --device-input "*" --benchmark-runs 1 else echo "Failed to find downloaded ONNX." exit 1 @@ -317,8 +329,7 @@ jobs: # Copy the model files rsync -avm \ --include='*.dlc' \ - --include='*.chunk*' \ - --include='*.chunkmanifest' \ + --include='/${{ inputs.target_hardware == 'chestnut' && 'big_' || '' }}driving_tinygrad.pkl' \ --exclude='*' \ --delete-excluded \ --chown=comma:comma \ @@ -340,6 +351,7 @@ jobs: with: name: model-${{ env.MODEL_NAME }}${{ inputs.artifact_suffix }}-${{ github.run_number }} path: ${{ env.OUTPUT_DIR }} + compression-level: 0 - name: Upload artifact name file uses: actions/upload-artifact@v4 diff --git a/.github/workflows/sunnypilot-build-prebuilt.yaml b/.github/workflows/sunnypilot-build-prebuilt.yaml index 7fdea327b9..78305f3fb4 100644 --- a/.github/workflows/sunnypilot-build-prebuilt.yaml +++ b/.github/workflows/sunnypilot-build-prebuilt.yaml @@ -503,8 +503,8 @@ jobs: mkdir -p ${{ env.OUTPUT_DIR }} tar xzf prebuilt.tar.gz -C ${{ env.OUTPUT_DIR }} - - name: Download model chunks from HF - uses: ./.github/workflows/download-hf-model-chunks + - name: Download models from HF + uses: ./.github/workflows/download-hf-model with: hf_repo: sunnypilot/sunnypilot_models_v1 dest_dir: ${{ env.OUTPUT_DIR }}/openpilot/selfdrive/modeld/models @@ -573,8 +573,8 @@ jobs: mkdir -p ${{ env.OUTPUT_DIR }} tar xzf prebuilt.tar.gz -C ${{ env.OUTPUT_DIR }} - - name: Download model chunks from HF - uses: ./.github/workflows/download-hf-model-chunks + - name: Download models from HF + uses: ./.github/workflows/download-hf-model with: hf_repo: sunnypilot/sunnypilot_models_v1 dest_dir: ${{ env.OUTPUT_DIR }}/openpilot/selfdrive/modeld/models diff --git a/.github/workflows/test_models.yml b/.github/workflows/test_models.yml index 769fca1051..2345385fca 100644 --- a/.github/workflows/test_models.yml +++ b/.github/workflows/test_models.yml @@ -47,38 +47,22 @@ jobs: with: submodules: true - - name: Download Model Chunks in Parallel + - name: Download model env: HF_TOKEN: ${{ secrets.HF_TOKEN }} + MODEL_URL: ${{ matrix.artifact.download_uri.url }} + MODEL_FILE: ${{ matrix.artifact.file_name }} + shell: bash run: | - mkdir -p /tmp/model_chunks - echo '${{ toJson(matrix.artifact.chunks) }}' > chunks.json - - BASE_URL="${{ matrix.artifact.download_uri.url }}" - export BASE_DIR=$(dirname "$BASE_URL") - - python3 -c ' - import json, os - with open("chunks.json") as f: - chunks = json.load(f) - manifest_path = f"/tmp/model_chunks/${{ matrix.artifact.file_name }}.chunkmanifest" - with open(manifest_path, "w") as f: - f.write(str(len(chunks))) - base_dir = os.environ["BASE_DIR"] - hf_token = os.environ.get("HF_TOKEN", "") - with open("/tmp/curl_config.txt", "w") as f: - for c in chunks: - fn = c["file_name"] - f.write(f"url = \"{base_dir}/{fn}\"\noutput = \"/tmp/model_chunks/{fn}\"\n") - if hf_token: - f.write(f"header = \"Authorization: Bearer {hf_token}\"\n") - ' - curl -Z --parallel-immediate --parallel-max 12 --retry 3 --retry-all-errors -s -S -f -L -K /tmp/curl_config.txt + mkdir -p /tmp/models + AUTH=() + if [ -n "$HF_TOKEN" ]; then AUTH=(-H "Authorization: Bearer $HF_TOKEN"); fi + curl --retry 3 --retry-all-errors -s -S -f -L "${AUTH[@]}" -o "/tmp/models/$MODEL_FILE" "$MODEL_URL" - name: Run Model Compatibility Test env: MODEL_BASE_NAME: ${{ matrix.artifact.file_name }} - MODEL_CHUNK_DIR: "/tmp/model_chunks" + MODEL_DIR: "/tmp/models" run: | export PYTHONPATH="$GITHUB_WORKSPACE/tinygrad_repo:$GITHUB_WORKSPACE:$PYTHONPATH" python3 -m pip uninstall -y tinygrad || uv pip uninstall -y tinygrad --system || true diff --git a/openpilot/sunnypilot/modeld_v2/compile_modeld.py b/openpilot/sunnypilot/modeld_v2/compile_modeld.py index 2b7256aeb9..ab338f3ad5 100755 --- a/openpilot/sunnypilot/modeld_v2/compile_modeld.py +++ b/openpilot/sunnypilot/modeld_v2/compile_modeld.py @@ -7,14 +7,12 @@ See the LICENSE.md file in the root directory for more details. """ import argparse -import atexit import codecs import gc import math import multiprocessing as mp import os import pickle -import shutil import tempfile import time from functools import partial @@ -22,8 +20,8 @@ from typing import Any import numpy as np -from openpilot.common.file_chunker import chunk_file, get_chunk_targets, open_file_chunked -from openpilot.sunnypilot.modeld_v2.helpers import dump_oob, load_oob +from openpilot.sunnypilot.modeld_v2.helpers import PlainJit, read_parts, write_bundle +from openpilot.sunnypilot.modeld_v2.plain_artifact_compiler import capture_frame, record_kernel_linear, serialize_jit import openpilot.sunnypilot.modeld_v2.stock_dependencies as stock from openpilot.system.camerad.cameras.nv12_info import get_nv12_info @@ -55,16 +53,6 @@ def _patch_tinygrad_fetch_fw(): pass _patch_tinygrad_fetch_fw() -def _patch_ops_serialization(): - try: - from tinygrad.uop.ops import Ops - def __reduce_ex__(self, proto): - return (getattr, (self.__class__, self.name)) - Ops.__reduce_ex__ = __reduce_ex__ - except ImportError: - pass -_patch_ops_serialization() - MODEL_TYPES = ('vision_policy', 'supercombo', 'vision_multi_policy') WARP_INPUTS = ['tfm', 'big_tfm'] POLICY_INPUTS = ['img_q', 'big_img_q', 'feat_q', 'desire_q', 'packed_npy_inputs'] @@ -72,6 +60,7 @@ POLICY_INPUTS = ['img_q', 'big_img_q', 'feat_q', 'desire_q', 'packed_npy_inputs' class MetadataOnnxPBParser(OnnxPBParser): def _parse_ModelProto(self) -> dict: obj: dict[str, Any] = {"graph": {"input": [], "output": []}, "metadata_props": []} + self.opset_imports = {} for fid, wire_type in self._parse_message(self.reader.len): match fid: case 7: @@ -150,6 +139,10 @@ def get_policy_npy_shapes(input_shapes: dict, is_supercombo: bool = False) -> tu return shapes, sizes +def host_view(tensor: Tensor) -> np.ndarray: # QCOM -> CPU round trip of tfm takes 5ms longer than host memory view + return np.frombuffer(tensor._buffer()._host_mv(), dtype=np.float32).reshape(tensor.shape) + + def generate_queues_and_npy(input_shapes: dict, frame_skip: int, device: str = Device.DEFAULT, is_supercombo: bool = False) -> tuple[dict, dict]: road_key, _ = _detect_vision_keys(input_shapes) @@ -185,7 +178,7 @@ def generate_queues_and_npy(input_shapes: dict, frame_skip: int, device: str = D 'big_img_q': Tensor(np.zeros(img_buf_shape, dtype=np.uint8), device=device).contiguous().realize(), 'desire_q': Tensor(np.zeros((frame_skip * desire_shape[1], desire_shape[0], desire_shape[2]), dtype=np.float32), device=device).contiguous().realize(), - 'packed_npy_inputs': Tensor(packed_npy_inputs, device=device if os.getenv('CHESTNUT') else 'NPY').realize(), + 'packed_npy_inputs': Tensor(packed_npy_inputs, device=device).realize(), } if features_buffer: @@ -197,9 +190,11 @@ def generate_queues_and_npy(input_shapes: dict, frame_skip: int, device: str = D for key in ('tfm', 'big_tfm'): if key not in npy_arrays: npy_arrays[key] = np.eye(3, dtype=np.float32) - queues[key] = Tensor(npy_arrays[key], device='NPY').realize() + queues[key] = Tensor(npy_arrays[key], device=device).realize() - return queues, npy_arrays + packed = host_view(queues['packed_npy_inputs']) + npy_arrays |= {k: v.reshape(s) for (k, s), v in zip(shapes.items(), np.split(packed, split_indices) if sizes else [], strict=True)} + return queues, npy_arrays | {k: host_view(queues[k]) for k in ('tfm', 'big_tfm')} def make_split_input_queues(vision_input_shapes: dict, policy_input_shapes: dict, @@ -221,8 +216,8 @@ def make_warp_queues(device=Device.DEFAULT): 'tfm': np.zeros((3, 3), dtype=np.float32), 'big_tfm': np.zeros((3, 3), dtype=np.float32), } - queues = {k: Tensor(v, device='NPY').realize() for k, v in npy.items()} - return queues, npy + queues = {k: Tensor(v, device=device).realize() for k, v in npy.items()} + return queues, {k: host_view(v) for k, v in queues.items()} def make_run_policy(vision_runner, policy_runners: list, features_slice: slice, frame_skip: int, input_shapes: dict): @@ -282,8 +277,19 @@ def make_run_policy(vision_runner, policy_runners: list, features_slice: slice, return run_policy +class _CheckFrameRecorder: + def __init__(self, jit): + self.jit, self.frame = jit, None + + def __call__(self, *args, **kwargs): + if self.frame is not None or args or self.jit.cnt < 2: + return self.jit(*args, **kwargs) + self.frame, result = capture_frame(self.jit, kwargs, None) + return result + + def compile_jit(jit_or_fn, input_keys_or_make_inputs, make_queues=None, make_random_inputs=None, - benchmark_runs: int = 1, clear_refs=None, return_loaded: bool = True): + benchmark_runs: int = 1, clear_refs=None): if callable(input_keys_or_make_inputs) and make_queues is None: fn = jit_or_fn make_inputs = input_keys_or_make_inputs @@ -334,22 +340,23 @@ def compile_jit(jit_or_fn, input_keys_or_make_inputs, make_queues=None, make_ran print('capture + replay') gc.collect() - test_val = run_eval(jit, 42, 3) - print(f'pickle round trip ({benchmark_runs} runs per seed)') - with tempfile.TemporaryFile(dir=".") as f: - dump_oob(jit, f) - if clear_refs: - clear_refs() - gc.collect() - f.seek(0) - loaded_jit = load_oob(f) + recorder = _CheckFrameRecorder(jit) + with record_kernel_linear() as recorded: + test_val = run_eval(recorder, 42, 3) + assert len(recorded) == 1, f"expected one captured kernel list, got {len(recorded)}" + part = serialize_jit(jit, recorded[0], recorder.frame) + print(f'plain artifact round trip ({benchmark_runs} runs per seed)') + if clear_refs: + clear_refs() + gc.collect() + loaded_jit = PlainJit(part[0], memoryview(part[1])) for seed in (42, 43): reference = test_val if seed == 42 else run_eval(jit_or_fn, seed, 1) actual = run_eval(loaded_jit, seed, benchmark_runs) for ref, val in zip(reference, actual, strict=True): np.testing.assert_array_equal(ref, val) - return loaded_jit if return_loaded else jit + return part def _parse_size(size_str: str) -> tuple[int, int]: @@ -357,16 +364,6 @@ def _parse_size(size_str: str) -> tuple[int, int]: return int(width), int(height) -def read_file_chunked_to_disk(path): - if not path: - return None - tmp_path = f'{path}.unchunked' - with open(tmp_path, 'wb') as f, open_file_chunked(path) as src: - shutil.copyfileobj(src, f) - atexit.register(lambda: os.path.exists(tmp_path) and os.remove(tmp_path)) - return tmp_path - - def _load_policy_runners(args: argparse.Namespace) -> tuple[list, list]: runners, keys = [], [] for name, onnx_arg in [('policy', args.policy_onnx), ('off_policy', args.off_policy_onnx), ('on_policy', args.on_policy_onnx)]: @@ -387,10 +384,7 @@ def _compile_warp_resolution_worker(cam_w, cam_h, model_w, model_h, benchmark_ru warp = None compiled_jit = compile_jit(warp, WARP_INPUTS, make_warp_queues, make_random_inputs=make_random_warp_inputs, benchmark_runs=benchmark_runs, clear_refs=cleanup_warp) - - result_data = (compiled_jit, Device.DEFAULT) - with open(result_path, "wb") as f: - dump_oob(result_data, f) + write_bundle(result_path, {(cam_w, cam_h): compiled_jit}, {'device': Device.DEFAULT}) if __name__ == "__main__": @@ -421,12 +415,6 @@ if __name__ == "__main__": model_w, model_h = args.model_size output_data = {} - args.vision_onnx = read_file_chunked_to_disk(args.vision_onnx) - args.policy_onnx = read_file_chunked_to_disk(args.policy_onnx) - args.off_policy_onnx = read_file_chunked_to_disk(args.off_policy_onnx) - args.on_policy_onnx = read_file_chunked_to_disk(args.on_policy_onnx) - args.supercombo_onnx = read_file_chunked_to_disk(args.supercombo_onnx) - is_unified_supercombo = False if args.model_type == 'supercombo': assert args.supercombo_onnx @@ -449,8 +437,10 @@ if __name__ == "__main__": frame_copy_size=frame_copy_size) warp = stock.make_warp(nv12, model_w, model_h) run_model_jit = TinyJit(stock.make_run_model(warp, run_policy, model_metadata, frame_copy_size), prune=True) - compiled_jit = compile_jit(run_model_jit, stock.MODELD_INPUTS, make_model_queues, benchmark_runs=args.benchmark_runs, return_loaded=False) - output_data[(cam_w, cam_h)] = compiled_jit + write_bundle(part_path := f"{args.output}.{cam_w}x{cam_h}", {(cam_w, cam_h): compile_jit(run_model_jit, stock.MODELD_INPUTS, make_model_queues, + benchmark_runs=args.benchmark_runs)}) + output_data[(cam_w, cam_h)] = read_parts(part_path)[0][(cam_w, cam_h)] + os.remove(part_path) gc.collect() del model_runner, run_policy gc.collect() @@ -518,18 +508,12 @@ if __name__ == "__main__": tmp_files.append((cam_w, cam_h, tmp_res.name)) for cam_w, cam_h, tmp_name in tmp_files: - with open(tmp_name, "rb") as f: - compiled_jit, dev_name = load_oob(f) - output_data[(cam_w, cam_h)] = compiled_jit - output_data['input_devices']['warp'] = dev_name + warp_parts, warp_extras = read_parts(tmp_name) + output_data[(cam_w, cam_h)] = warp_parts[(cam_w, cam_h)] + output_data['input_devices']['warp'] = warp_extras['device'] os.remove(tmp_name) output_data['metadata']['warp_dev'] = Device.DEFAULT - with open(args.output, "wb") as file: - dump_oob(output_data, file) - - pkl_size = os.path.getsize(args.output) - print(f"Saved combined JIT to {args.output} ({pkl_size / 1e6:.2f} MB)") - chunk_targets = get_chunk_targets(args.output, pkl_size) - chunk_file(args.output, chunk_targets) - print(f"Chunked into {len(chunk_targets) - 1} file(s)") + parts = {key: output_data.pop(key) for key in [key for key, value in output_data.items() if isinstance(value, tuple)]} + write_bundle(args.output, parts, output_data) + print(f"Saved model to {args.output} ({os.path.getsize(args.output) / 1e6:.2f} MB, {len(parts)} JITs)") diff --git a/openpilot/sunnypilot/modeld_v2/compile_modeld_v2.py b/openpilot/sunnypilot/modeld_v2/compile_modeld_v2.py index a4bfcce313..44ff5fed67 100755 --- a/openpilot/sunnypilot/modeld_v2/compile_modeld_v2.py +++ b/openpilot/sunnypilot/modeld_v2/compile_modeld_v2.py @@ -1,11 +1,59 @@ #!/usr/bin/env python3 -import runpy +import argparse import os +import sys -from openpilot.sunnypilot.modeld_v2.compile_modeld import _patch_ops_serialization +import numpy as np + +from openpilot.sunnypilot.modeld_v2.helpers import PlainJit, write_bundle +from openpilot.sunnypilot.modeld_v2.plain_artifact_compiler import capture_frame, record_kernel_linear, serialize_jit + + +def initialize_check_inputs(arrays: dict): + rng = np.random.default_rng(7) + for value in arrays.values(): + value[...] = rng.standard_normal(value.shape) if value.dtype.kind == "f" else rng.integers(0, 2 if value.dtype == np.bool_ else 16, value.shape) if __name__ == "__main__": - _patch_ops_serialization() - upstream_script = os.path.join(os.environ.get("TINYGRAD_PATH", "tinygrad_repo"), "examples/openpilot/compile_onnx.py") - runpy.run_path(upstream_script, run_name="__main__") + sys.path.insert(0, os.environ.get("TINYGRAD_PATH", "tinygrad_repo")) + from examples.openpilot.compile_onnx import onnx_metadata + from examples.openpilot.helpers import allocate_inputs, benchmark + from tinygrad import Device, Tensor, TinyJit + from tinygrad.dtype import _to_np_dtype + from tinygrad.nn.onnx import OnnxRunner + + parser = argparse.ArgumentParser(description="Compile a native-architecture ONNX model into a plain tinygrad artifact") + parser.add_argument('onnx') + parser.add_argument('output') + parser.add_argument('--device-input', action='append', default=[]) + parser.add_argument('--benchmark-runs', type=int, default=20) + args = parser.parse_args() + + runner, (properties, output_shapes) = OnnxRunner(args.onnx), onnx_metadata(args.onnx) + device_inputs = runner.graph_inputs.keys() if '*' in args.device_input else args.device_input + metadata = {'metadata': properties} | {f'{kind}_shapes': {name: tuple(d if isinstance(d, int) else 0 for d in shape) for name, shape in shapes.items()} + for kind, shapes in [('input', {name: spec.shape for name, spec in runner.graph_inputs.items()}), + ('output', output_shapes)]} + input_specs = {name: (tuple(s if isinstance(s, int) else 1 for s in spec.shape), np.dtype(_to_np_dtype(spec.dtype)).str, + Device.DEFAULT if name in device_inputs else 'NPY') for name, spec in runner.graph_inputs.items()} + + def model(inputs): + return {name: value.contiguous() for name, value in runner({name: value.to(Device.DEFAULT) for name, value in inputs.items()}).items()} + output_specs = {name: (value.shape, np.dtype(_to_np_dtype(value.dtype)).name, Device.DEFAULT) for name, value in model(allocate_inputs(input_specs)).items()} + + @TinyJit(prune=True) + def jit(output_buffers, **inputs): + Tensor.realize(*(output_buffers[name].assign(value) for name, value in model(inputs).items())) + + check_inputs = allocate_inputs(input_specs, initialize_check_inputs) | {'output_buffers': allocate_inputs(output_specs)} + with record_kernel_linear() as recorded: + for i in range(2 + args.benchmark_runs): + benchmark(jit, cb=lambda t, i=i: print(f" [{i + 1}/{2 + args.benchmark_runs}] {t * 1e3:.2f} ms"), **check_inputs) + assert len(recorded) == 1, f"expected one captured kernel list, got {len(recorded)}" + + frame, _ = capture_frame(jit, check_inputs, 'output_buffers') + part = serialize_jit(jit, recorded[0], frame) + PlainJit(part[0], memoryview(part[1])) + write_bundle(args.output, {'run': part}, {'metadata': metadata, 'input_specs': input_specs, 'output_specs': output_specs}) + print(f"Saved model to {args.output} ({os.path.getsize(args.output) / 1e6:.2f} MB, {len(part[0]['calls'])} kernels)") diff --git a/openpilot/sunnypilot/modeld_v2/helpers.py b/openpilot/sunnypilot/modeld_v2/helpers.py index 03012be3a9..8d0fc0e7fd 100644 --- a/openpilot/sunnypilot/modeld_v2/helpers.py +++ b/openpilot/sunnypilot/modeld_v2/helpers.py @@ -5,14 +5,23 @@ 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 builtins +import contextlib +import dataclasses +import functools +import hashlib import io -import struct +import json +import math +import mmap +import os import pickle -import inspect -import importlib -import tempfile -import shutil -import enum +import struct + +MAGIC = b"SPTGPLN1" +ALIGNMENT = 256 +METADATA_TYPES = ("slice", "tuple", "list", "dict", "set", "frozenset", "complex") +RUNTIME_ENVIRONMENT = ("GMMU", "AMD_AQL", "AMD_DISABLE_SDMA", "GRAPH_ONE_KERNEL") def _patch_system_flock_acquire(): @@ -32,86 +41,246 @@ def _patch_system_flock_acquire(): pass _patch_system_flock_acquire() -def _pad_args(func, args, kwargs): - try: - sig = inspect.signature(func) - except Exception: - return args, kwargs - params = list(sig.parameters.values()) - if inspect.isfunction(func) and params and params[0].name in ('cls', 'self'): - params = params[1:] - new_args = list(args) - has_varargs = any(p.kind == inspect.Parameter.VAR_POSITIONAL for p in params) - if len(new_args) > len(params) and not has_varargs: - new_args = new_args[:len(params)] - - for i in range(len(new_args), len(params)): - param = params[i] - if param.kind in (inspect.Parameter.VAR_POSITIONAL, inspect.Parameter.VAR_KEYWORD): - continue - val = param.default if param.default is not inspect.Parameter.empty else None - new_args.append(val) - return new_args, kwargs +class ArtifactMismatch(RuntimeError): + pass -def _dynamic_factory(real_class): - if isinstance(real_class, type) and issubclass(real_class, enum.Enum): - return real_class +class Blob: + def __init__(self, sink=None): + self.chunks: list = [] + self.locations: dict[bytes, list[int]] = {} + self.size, self.write = 0, sink.write if sink else self.chunks.append - def factory(*args, **kwargs): - try: - return real_class(*args, **kwargs) - except TypeError: - new_args, new_kwargs = _pad_args(real_class, args, kwargs) - return real_class(*new_args, **new_kwargs) - - class DynamicMeta(type(real_class)): - def __call__(cls, *args, **kwargs): - return factory(*args, **kwargs) - - class DynamicProxy(real_class, metaclass=DynamicMeta): - __slots__ = () - - def __new__(cls, *args, **kwargs): - return factory(*args, **kwargs) - - DynamicProxy.__name__ = real_class.__name__ - DynamicProxy.__module__ = real_class.__module__ - return DynamicProxy + def add(self, data) -> list[int]: + if (key := hashlib.blake2b(data, digest_size=16).digest()) not in self.locations: + start = self.size + -self.size % ALIGNMENT + self.write(bytes(start - self.size)), self.write(data) + self.locations[key], self.size = [start, start + memoryview(data).nbytes], start + memoryview(data).nbytes + return self.locations[key] -class DynamicTinygradUnpickler(pickle.Unpickler): +class _MetadataUnpickler(pickle.Unpickler): def find_class(self, module, name): - real_class = getattr(importlib.import_module(module), name) - if module.startswith("tinygrad"): - return _dynamic_factory(real_class) - return real_class + if module == "builtins" and name in METADATA_TYPES: + return getattr(builtins, name) + raise pickle.UnpicklingError(f"{module}.{name} is not allowed in model metadata") -def load_oob(f): - opcodes = f.read(struct.unpack(' list: + if isinstance(value, (tuple, list)): + return [tensor for item in value for tensor in flatten(item)] + return [] if value is None else [value] -def dump_oob(obj, f): - with tempfile.TemporaryFile(dir=".") as tmp: - def buffer_cb(buffer): - data = buffer.raw() - tmp.write(struct.pack(' tuple[dict, dict]: + with open(path, "rb") as file: + data = memoryview(mmap.mmap(file.fileno(), 0, access=mmap.ACCESS_READ)) + if bytes(data[:len(MAGIC)]) != MAGIC: + raise ValueError(f"{path} is not a plain tinygrad model, recompile it with compile_modeld.py") + start = len(MAGIC) + 8 + struct.unpack_from(" dict: + parts, extras = read_parts(path) + if cache and (cached := _read_cache(path, parts, device_map)) is not None: + return extras | cached + shared: dict = {} + jits = {key: PlainJit(header, blob, device_map, verify, shared) for key, (header, blob) in parts.items()} + if cache: + _write_cache(path, jits, device_map) + return extras | jits + + +@functools.cache +def tinygrad_identity() -> str: + import tinygrad + root, digest = os.path.dirname(tinygrad.__file__), hashlib.sha256() + for directory, _, files in sorted(os.walk(root)): + for name in sorted(file for file in files if file.endswith(".py")): + stat = os.stat(path := os.path.join(directory, name)) + digest.update(f"{os.path.relpath(path, root)}:{stat.st_size}:{stat.st_mtime_ns};".encode()) + return digest.hexdigest() + + +def _cache_key(path: str, device_map: dict | None) -> bytes: + stat = os.stat(path) + return json.dumps({"format": MAGIC.decode(), "tinygrad": tinygrad_identity(), "size": stat.st_size, "mtime_ns": stat.st_mtime_ns, + "device_map": sorted((device_map or {}).items()), "env": [os.environ.get(name, "") for name in RUNTIME_ENVIRONMENT]}).encode() + + +def _cache_header_matches(file, path: str, device_map: dict | None) -> bool: + return file.read(struct.unpack(" bool: # entry point for warmup proc + try: + with open(f"{path}.cache", "rb") as file: + return _cache_header_matches(file, path, device_map) + except (OSError, struct.error): + return False + + +def _read_cache(path: str, parts: dict, device_map: dict | None) -> dict | None: + try: + with open(f"{path}.cache", "rb") as file: + if not _cache_header_matches(file, path, device_map): + return None + shared: dict = {} + buffers = {repr(key): _allocate_buffers(header, blob, device_map, shared) for key, (header, blob) in parts.items()} + return type("_Unpickler", (pickle.Unpickler,), {"persistent_load": lambda self, pid: buffers[pid[0]][pid[1]]})(file).load() + except Exception: + return None + + +def _write_cache(path: str, jits: dict, device_map: dict | None) -> None: + from tinygrad.device import Buffer + owners = {id(buffer): (repr(key), index) for key, jit in jits.items() for index, buffer in enumerate(jit.buffers)} + + class _Pickler(pickle.Pickler): + def persistent_id(self, obj): + return owners.get(id(obj)) if isinstance(obj, Buffer) else None + + key, tmp_path = _cache_key(path, device_map), f"{path}.cache.{os.getpid()}.tmp" + with contextlib.suppress(Exception): + with open(tmp_path, "wb") as file: + file.write(struct.pack(" str: + from tinygrad import Device + device_map = device_map or {} + return Device.canonicalize(device_map.get(recorded, device_map.get(recorded.split(":", 1)[0], recorded))) + + +def _allocate_buffers(header: dict, blob: memoryview, device_map: dict | None, shared: dict) -> list: + import numpy as np + from tinygrad.device import Buffer + written = {call["args"][i]["buffer"] for call in header["calls"] for i in (call["outs"] if call["kind"] == "program" else (0,)) + if "buffer" in call["args"][i]} + buffers = [] + for index, entry in enumerate(header["buffers"]): + device, data = _device(entry["device"], device_map), blob[slice(*entry["init"])] if "init" in entry else None + # read-only data the parts share in the file (weights common to every camera resolution) is uploaded once to the nut + key = (device, entry["dtype"], np.frombuffer(data, np.uint8).ctypes.data, len(data)) if data and index not in written else object() + if key not in shared: + shared[key] = Buffer(device, entry["nbytes"], initial_value=pickle.PickleBuffer(data) if data else None).ensure_allocated() + buffers.append(shared[key]) + return buffers + + +def _program_uop(call: dict, lib: bytes): + from tinygrad.helpers import Target + from tinygrad.uop.ops import KernelInfo, Ops, ProgramInfo, UOp + parameters = [] + for name, compact_slot, dtype, shape in call["signature"]: + parameter = UOp.param(call["globals"][compact_slot], _dtype(dtype), math.prod(shape)) + image = (shape[0], shape[1]) if len(shape) == 3 and shape[2] == 4 else None + parameters.append(parameter.replace(arg=dataclasses.replace(parameter.arg, name=name, image=image))) + info = ProgramInfo(global_size=tuple(call["global_size"]), local_size=tuple(call["local_size"]), vars=(), globals=tuple(call["globals"]), + outs=tuple(call["outs"]), ins=tuple(call["ins"]), target=Target(**call["target"])) + return UOp(Ops.PROGRAM, src=(UOp(Ops.SINK, arg=KernelInfo(name=call["name"])), UOp(Ops.LINEAR, src=tuple(parameters)), + UOp(Ops.SOURCE, arg=""), UOp(Ops.BINARY, arg=lib)), arg=info) + + +class _Jit: + def __init__(self, captured, input_names: list, buffers: list): + self.captured, self.input_names, self.buffers = captured, input_names, buffers + + def __call__(self, *args, **kwargs): + import tinygrad.engine.jit as jit + input_uops, var_vals, names, _ = jit._prepare_jit_inputs(args, kwargs) + if names != self.input_names: + raise ValueError(f"expected inputs {self.input_names}, got {names}") + return self.captured(input_uops, var_vals) + + +class PlainJit(_Jit): + def __init__(self, header: dict, blob: memoryview, device_map: dict | None = None, verify: bool = True, shared: dict | None = None): + import numpy as np + import tinygrad.engine.jit as jit + from tinygrad import Tensor + from tinygrad.device import Buffer + from tinygrad.dtype import dtypes + from tinygrad.uop.ops import Ops, UOp + + inputs = {spec["slot"]: spec for spec in header["inputs"]} + buffers = _allocate_buffers(header, blob, device_map, {} if shared is None else shared) + bases = [UOp.from_buffer(buffer, _dtype(entry["dtype"])) for buffer, entry in zip(buffers, header["buffers"], strict=True)] + + def argument(described: dict): + if "input" in described: + spec = inputs[described["input"]] + return UOp.param(described["input"], _dtype(spec["dtype"]), spec["numel"], _device(spec["device"], device_map)) + base, dtype, offset, nbytes = bases[described["buffer"]], _dtype(described["dtype"]), described["offset"], described["nbytes"] + if offset == 0 and nbytes == header["buffers"][described["buffer"]]["nbytes"] and dtype == base.dtype: + return base + return (base if base.dtype.itemsize == 1 else base.bitcast(dtypes.uint8)).shrink(((offset, offset + nbytes),)).bitcast(dtype) + + calls = [] + for call in header["calls"]: + arguments = [argument(described) for described in call["args"]] + if call["kind"] == "program": + calls.append(_program_uop(call, bytes(blob[slice(*call["lib"])])).call(*arguments)) + else: + calls.append(arguments[0].store_call(arguments[1])) + # stand-ins for JIT inputs: copy staging to GPU is decided from the input buffers, the real ones bind per call + stand_ins = {slot: UOp.from_buffer(Buffer(_device(spec["device"], device_map), spec["numel"] * _dtype(spec["dtype"]).itemsize).ensure_allocated(), + _dtype(spec["dtype"])) for slot, spec in inputs.items()} + input_uops = tuple(stand_ins.get(slot, next(iter(stand_ins.values()), None)) for slot in range(max(stand_ins, default=-1) + 1)) + linear = jit.compile_linear(UOp(Ops.LINEAR, src=tuple(calls)), beam=0, input_uops=input_uops, cache=False) + + returned = [Tensor(argument(t)).reshape(tuple(t["shape"])) for t in header["returns"]["tensors"]] + ret = {"NoneType": None, "Tensor": returned[0] if returned else None, "tuple": tuple(returned), "list": returned}[header["returns"]["kind"]] + super().__init__(jit.CapturedJit(ret, linear, header["input_names"], []), header["input_names"], buffers) + + if not verify or not (check := header.get("check")): + return + def array(stored): + return np.frombuffer(blob[slice(*stored["data"])], dtype=np.dtype(stored["dtype"])).reshape(stored["shape"]).copy() + kwargs: dict = {} + for path, stored, device in check["inputs"]: + tensor = Tensor(array(stored), device=_device(device, device_map)).realize() + kwargs[path[0]] = tensor if len(path) == 1 else {**kwargs.get(path[0], {}), path[1]: tensor} + result = self(**kwargs) + outputs = kwargs[check["output_container"]].values() if check["output_container"] else flatten(result) + for index, (got, stored) in enumerate(zip([t.numpy() for t in outputs], check["outputs"], strict=True)): + expected = array(stored) + if got.shape != expected.shape or not np.allclose(got, expected, rtol=1e-3, atol=1e-3, equal_nan=True): + raise ArtifactMismatch(f"output {index} differs from the compile-time reference frame; recompile the model") diff --git a/openpilot/sunnypilot/modeld_v2/model_adapters.py b/openpilot/sunnypilot/modeld_v2/model_adapters.py index b240895180..e5f045d0ee 100644 --- a/openpilot/sunnypilot/modeld_v2/model_adapters.py +++ b/openpilot/sunnypilot/modeld_v2/model_adapters.py @@ -23,8 +23,8 @@ from tinygrad.uop.ops import UOp def input_view(buffer: Buffer, shape: tuple[int, ...], dtype: DType, offset: int) -> Tensor: - view = buffer.view(math.prod(shape), dtype, offset).ensure_allocated() - return Tensor(UOp.from_buffer(view)).reshape(shape) + view = buffer.view(math.prod(shape) * dtype.itemsize, offset).ensure_allocated() + return Tensor(UOp.from_buffer(view, dtype)).reshape(shape) class BaseModelAdapter: @@ -168,8 +168,7 @@ class NativeTinygradAdapter(BaseModelAdapter): stride, y_height, uv_height, _ = get_nv12_info(self.cam_w, self.cam_h) self.frame_copy_size = stride * (y_height + uv_height) - self.input_shapes_orig = self.jits['metadata']['input_shapes'] - self._vision_input_names = [k for k in self.input_shapes_orig if 'img' in k] + self._vision_input_names = ['img', 'big_img'] self.vision_output_slices = pickle.loads(codecs.decode(self.jits['metadata']['metadata']['output_slices'].encode(), 'base64')) self.run_warp = self._load_warp() diff --git a/openpilot/sunnypilot/modeld_v2/modeld.py b/openpilot/sunnypilot/modeld_v2/modeld.py index 2fd1baf999..b669868d6f 100755 --- a/openpilot/sunnypilot/modeld_v2/modeld.py +++ b/openpilot/sunnypilot/modeld_v2/modeld.py @@ -24,7 +24,6 @@ from openpilot.cereal.messaging import PubMaster, SubMaster from openpilot.cereal.visionipc import VisionStreamType from msgq.visionipc import VisionIpcClient, VisionBuf from opendbc.car.car_helpers import get_demo_car_params -from openpilot.common.file_chunker import open_file_chunked from openpilot.common.swaglog import cloudlog from openpilot.common.params import Params from openpilot.common.filter_simple import FirstOrderFilter @@ -42,7 +41,7 @@ from openpilot.sunnypilot.modeld_v2.meta_helper import load_meta_constants from openpilot.sunnypilot.modeld_v2.camera_offset_helper import CameraOffsetHelper from openpilot.sunnypilot.livedelay.helpers import get_lat_delay from openpilot.sunnypilot.modeld_v2.modeld_base import ModelStateBase -from openpilot.sunnypilot.modeld_v2.helpers import load_oob +from openpilot.sunnypilot.modeld_v2.helpers import load_model from openpilot.sunnypilot.modeld_v2.model_adapters import get_model_adapter from openpilot.sunnypilot.models.helpers import get_active_bundle from openpilot.sunnypilot.selfdrive.controls.lib.relc import RoadEdgeLaneChangeController @@ -52,8 +51,7 @@ BIG_MODEL_TIMEOUT = 60 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)) + return os.path.exists(path) def _find_driving_pkl(bundle): @@ -108,7 +106,7 @@ class ModelState(ModelStateBase): def _init_combined(self, pkl_path, cam_w, cam_h, bundle): cloudlog.warning(f"loading combined pkl: {pkl_path}") - jits = load_oob(open_file_chunked(pkl_path)) + jits = load_model(pkl_path) metadata = jits['metadata'] self.WARP_DEV = metadata.get('warp_dev', 'QCOM') if COMMA_HARDWARE else 'CPU' @@ -258,7 +256,8 @@ def main(demo=False): CHESTNUT = chestnut_present() if CHESTNUT: - os.environ['HCQDEV_WAIT_TIMEOUT_MS'] = '3000' + from tinygrad.runtime.ops_amd import AMDDevice + AMDDevice.wait_timeout_ms = 3000 params = Params() params.put_bool("ChestnutLoading", CHESTNUT) diff --git a/openpilot/sunnypilot/modeld_v2/plain_artifact_compiler.py b/openpilot/sunnypilot/modeld_v2/plain_artifact_compiler.py new file mode 100644 index 0000000000..05001cd670 --- /dev/null +++ b/openpilot/sunnypilot/modeld_v2/plain_artifact_compiler.py @@ -0,0 +1,90 @@ +""" +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 dataclasses +import mmap +import tempfile +from contextlib import contextmanager +from unittest import mock + +import numpy as np + +from openpilot.sunnypilot.modeld_v2.helpers import Blob, flatten + +# Compile time only/ see discountchubbs 4 assistance: turns a captured TinyJit into plain data that helpers.load_model rebuilds. It reads internals of the +# tinygrad that compiles the model, so a tinygrad bump that breaks it blocks new builds, never loading already built models. +@contextmanager +def record_kernel_linear(): + import tinygrad.engine.jit as jit + import tinygrad.engine.realize as realize + recorded, hcq_compile, compile_linear = [], realize.hcq_compile, jit.compile_linear + + def recording_compile_linear(*args, **kwargs): + with mock.patch.object(realize, "hcq_compile", lambda linear, *a, **k: (recorded.append(linear), hcq_compile(linear, *a, **k))[1]): + return compile_linear(*args, **kwargs) + + with mock.patch.object(jit, "compile_linear", recording_compile_linear): + yield recorded + + +def capture_frame(jit, kwargs: dict, output_container: str | None = None): + inputs = [([name] if key is None else [name, key], tensor.numpy().copy(), tensor.device) + for name, value in kwargs.items() for key, tensor in (value.items() if isinstance(value, dict) else [(None, value)])] + result = jit(**kwargs) + outputs = kwargs[output_container].values() if output_container else flatten(result) + return {"inputs": inputs, "outputs": [tensor.numpy().copy() for tensor in outputs], "output_container": output_container}, result + + +def serialize_jit(jit, kernel_linear, check: dict | None = None) -> tuple[dict, memoryview]: + from tinygrad.uop.ops import Ops + sink = tempfile.TemporaryFile(dir=".") + blob, bases, input_specs = Blob(sink), {}, {} + + def store(array) -> dict: + array = np.ascontiguousarray(array) + return {"data": blob.add(array.tobytes()), "dtype": array.dtype.str, "shape": list(array.shape)} + + def describe(uop) -> dict: + if uop.op is Ops.PARAM: + input_specs.setdefault(uop.arg.slot, {"dtype": repr(uop.dtype), "numel": int(uop.max_numel()), "device": uop.device}) + return {"input": uop.arg.slot} + buffer = uop.buffer + root = buffer._base if buffer._base is not None else buffer + base = bases.setdefault(id(root), (len(bases), root, uop.base.dtype if buffer._base is None else uop.dtype)) + return {"buffer": base[0], "offset": buffer.offset if buffer._base is not None else 0, "nbytes": buffer.nbytes, "dtype": repr(uop.dtype)} + + calls, written, needs_init = [], set(), set() + for call in kernel_linear.src: + ast, arguments = call.src[0], [describe(argument) for argument in call.src[1:]] + if ast.op is Ops.PROGRAM: + elf, info = ast.to_elf(), ast.arg + global_size, local_size = info.launch_dims({}) + calls.append({"kind": "program", "name": elf.name, "lib": blob.add(bytes(elf.lib)), "target": dataclasses.asdict(elf.target), + "signature": [[name, slot, repr(dtype), [int(x) for x in shape]] for name, slot, dtype, shape in elf.signature], + "global_size": [int(x) for x in global_size], "local_size": [int(x) for x in local_size], + "globals": list(info.globals), "outs": list(info.outs), "ins": list(info.ins), "args": arguments}) + outs, ins = info.outs, info.ins + elif ast.op is Ops.COPY: + calls.append({"kind": "copy", "args": arguments}) + outs, ins = (0,), (1,) + else: + raise ValueError(f"unsupported call {ast.op} {ast.arg}") + needs_init |= {arguments[i]["buffer"] for i in ins if "buffer" in arguments[i] and arguments[i]["buffer"] not in written} + written |= {arguments[i]["buffer"] for i in outs if "buffer" in arguments[i]} + + ret = jit.captured.ret + returns = {"kind": type(ret).__name__, "tensors": [describe(t.uop.base) | {"shape": [int(x) for x in t.shape]} for t in flatten(ret)]} + buffers = [{"nbytes": base.nbytes, "device": base.device, "dtype": repr(dtype)} + | ({"init": blob.add(base.as_memoryview())} if index in needs_init else {}) for index, base, dtype in bases.values()] + names = list(jit.captured.expected_names) + if check is not None: + check = {"inputs": [[path, store(array), device] for path, array, device in check["inputs"]], "outputs": [store(array) for array in check["outputs"]], + "output_container": check["output_container"]} + header = {"input_names": names, "buffers": buffers, "calls": calls, "returns": returns, "check": check, + "inputs": [{"slot": slot, "name": names[slot] if slot < len(names) else None} | input_specs[slot] for slot in sorted(input_specs)]} + with sink: + sink.flush() + return header, memoryview(mmap.mmap(sink.fileno(), 0, access=mmap.ACCESS_READ) if blob.size else b"") diff --git a/openpilot/sunnypilot/modeld_v2/tests/helpers.py b/openpilot/sunnypilot/modeld_v2/tests/helpers.py index 2dde6c19c6..9ab6af7afa 100644 --- a/openpilot/sunnypilot/modeld_v2/tests/helpers.py +++ b/openpilot/sunnypilot/modeld_v2/tests/helpers.py @@ -188,10 +188,8 @@ def make_pkl_data(archetype): def write_pkl(tmp_path, archetype): - from openpilot.sunnypilot.modeld_v2.helpers import dump_oob pkl_path = tmp_path / 'driving_test_tinygrad.pkl' - with open(pkl_path, 'wb') as f: - dump_oob(make_pkl_data(archetype), f) + pkl_path.write_bytes(pickle.dumps(make_pkl_data(archetype))) return pkl_path @@ -211,6 +209,7 @@ def patch_modeld(monkeypatch): def _patch(bundle): monkeypatch.setattr(helpers, 'get_active_bundle', lambda params=None, *, chestnut=None: bundle) monkeypatch.setattr(modeld_module, 'get_active_bundle', lambda params=None, *, chestnut=None: bundle) + monkeypatch.setattr(modeld_module, 'load_model', lambda path: pickle.loads(pathlib.Path(path).read_bytes())) return _patch diff --git a/openpilot/sunnypilot/modeld_v2/tests/test_combined_pkl_loader.py b/openpilot/sunnypilot/modeld_v2/tests/test_combined_pkl_loader.py index 8b4c4e32a1..db0eaf8703 100644 --- a/openpilot/sunnypilot/modeld_v2/tests/test_combined_pkl_loader.py +++ b/openpilot/sunnypilot/modeld_v2/tests/test_combined_pkl_loader.py @@ -90,7 +90,6 @@ class TestStockEquivalence(OpenpilotTestCase): def test_unified_run_model(self, tmp_path, monkeypatch, patch_modeld): from openpilot.common.hardware import hw - from openpilot.sunnypilot.modeld_v2.helpers import dump_oob shapes = {'img': (1, 12, 128, 256), 'big_img': (1, 12, 128, 256), 'features_buffer': (1, 24, 32, 512), 'desire_pulse': (1, 25, 8), 'traffic_convention': (1, 2), 'action_t': (1, 2)} import codecs @@ -102,8 +101,7 @@ class TestStockEquivalence(OpenpilotTestCase): (CAM_W, CAM_H): tests_helpers._noop_jit, 'run_policy': tests_helpers._noop_jit, } - with open(tmp_path / 'driving_test_tinygrad.pkl', 'wb') as f: - dump_oob(pkl_data, f) + (tmp_path / 'driving_test_tinygrad.pkl').write_bytes(pickle.dumps(pkl_data)) bundle = DummyBundle(models=[DummyModel('supercombo', 'driving_test_tinygrad.pkl')]) patch_modeld(bundle) monkeypatch.setattr(hw.Paths, 'model_root', staticmethod(lambda: str(tmp_path))) diff --git a/openpilot/sunnypilot/modeld_v2/tests/test_compile_modeld.py b/openpilot/sunnypilot/modeld_v2/tests/test_compile_modeld.py index 239c498e7d..348b4aca6a 100644 --- a/openpilot/sunnypilot/modeld_v2/tests/test_compile_modeld.py +++ b/openpilot/sunnypilot/modeld_v2/tests/test_compile_modeld.py @@ -5,16 +5,12 @@ This file is part of sunnypilot and is licensed under the MIT License. See the LICENSE.md file in the root directory for more details. """ -import os -import tempfile import unittest -from pathlib import Path import numpy as np from openpilot.common.parameterized import parameterized -from openpilot.common.file_chunker import chunk_file, get_chunk_targets -from openpilot.sunnypilot.modeld_v2.compile_modeld import derive_frame_skip, _detect_desire_key, read_file_chunked_to_disk +from openpilot.sunnypilot.modeld_v2.compile_modeld import derive_frame_skip, _detect_desire_key from openpilot.common.test import OpenpilotTestCase @@ -168,36 +164,6 @@ class TestOutputSlicePreservation(OpenpilotTestCase): "vision and policy slices should not overlap in keys" -class TestReadFileChunkedToDisk(OpenpilotTestCase): - def test_none_passthrough(self): - assert read_file_chunked_to_disk(None) is None - - def test_unchunked_source_staged_on_disk(self): - with tempfile.TemporaryDirectory() as d: - src = Path(d) / "driving_supercombo.onnx" - payload = os.urandom(1024) - src.write_bytes(payload) - - out = Path(read_file_chunked_to_disk(str(src))) - - assert out.parent == Path(d) - assert out.name == "driving_supercombo.onnx.unchunked" - assert out.read_bytes() == payload - - def test_chunked_source_reassembled_on_disk(self): - with tempfile.TemporaryDirectory() as d: - src = Path(d) / "driving_supercombo.onnx" - payload = os.urandom(4096) - src.write_bytes(payload) - chunk_file(str(src), get_chunk_targets(str(src), len(payload))) - assert not src.exists() - - out = Path(read_file_chunked_to_disk(str(src))) - - assert out.parent == Path(d) - assert out.read_bytes() == payload - - class Test4DFeaturesBuffer(OpenpilotTestCase): def test_get_policy_npy_shapes_4d(self): from openpilot.sunnypilot.modeld_v2.compile_modeld import get_policy_npy_shapes diff --git a/openpilot/sunnypilot/modeld_v2/tests/test_fallback.py b/openpilot/sunnypilot/modeld_v2/tests/test_fallback.py index 08054210f1..7967bf2f97 100644 --- a/openpilot/sunnypilot/modeld_v2/tests/test_fallback.py +++ b/openpilot/sunnypilot/modeld_v2/tests/test_fallback.py @@ -5,13 +5,11 @@ 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 io +import pickle import requests -from openpilot.common.file_chunker import get_chunk_name from openpilot.common.hardware import hw from openpilot.common.test import OpenpilotTestCase -from openpilot.sunnypilot.modeld_v2.helpers import dump_oob import openpilot.sunnypilot.modeld_v2.modeld as modeld_module from openpilot.sunnypilot.modeld_v2.tests import helpers as tests_helpers from openpilot.sunnypilot.modeld_v2.tests.helpers import DummyModel, DummyBundle, CAM_W, CAM_H @@ -43,14 +41,11 @@ class TestFallback(OpenpilotTestCase): small_json = requests.get(ModelFetcher.MODEL_URL).json() small_bundle = ModelParser.parse_models(small_json)[-1] - buf = io.BytesIO() - dump_oob(tests_helpers.make_pkl_data(tests_helpers.ARCHETYPES['supercombo_non20hz']), buf) - oob_bytes = buf.getvalue() + oob_bytes = pickle.dumps(tests_helpers.make_pkl_data(tests_helpers.ARCHETYPES['supercombo_non20hz'])) + monkeypatch.setattr(modeld_module, 'load_model', lambda path: pickle.loads(open(path, 'rb').read())) for bundle in (big_bundle, small_bundle): - artifact = bundle.models[0].artifact - for i in range(len(artifact.chunks)): - (tmp_path / get_chunk_name(artifact.fileName, i, len(artifact.chunks))).write_bytes(oob_bytes if i == 0 else b"") + (tmp_path / bundle.models[0].artifact.fileName).write_bytes(oob_bytes) monkeypatch.setattr(modeld_module, 'get_active_bundle', lambda params=None, *, chestnut=None: small_bundle) assert modeld_module.ModelState(CAM_W, CAM_H, chestnut=False).chestnut is False diff --git a/openpilot/sunnypilot/modeld_v2/tests/test_models.py b/openpilot/sunnypilot/modeld_v2/tests/test_models.py index 8f1cdedbd6..d24b629e77 100644 --- a/openpilot/sunnypilot/modeld_v2/tests/test_models.py +++ b/openpilot/sunnypilot/modeld_v2/tests/test_models.py @@ -7,71 +7,14 @@ See the LICENSE.md file in the root directory for more details. import os import unittest -from unittest.mock import patch -from openpilot.common.file_chunker import open_file_chunked -from openpilot.sunnypilot.modeld_v2.helpers import load_oob -from tinygrad.device import Device + +from openpilot.sunnypilot.modeld_v2.helpers import PlainJit, load_model -class TestLegacyModels(unittest.TestCase): - def test_legacy_model_load(self): - base_name = os.environ.get("MODEL_BASE_NAME") - if not base_name: - raise unittest.SkipTest("MODEL_BASE_NAME env var not set, skipping integration test.") - chunk_dir = os.environ.get("MODEL_CHUNK_DIR", "/tmp/model_chunks") - base_path = os.path.join(chunk_dir, base_name) - - try: - f = open_file_chunked(base_path) - except Exception as error: - self.fail(f"Failed to open chunked file {base_path}: {error}") - self.addCleanup(f.close) - - real_getitem = Device.__class__.__getitem__ - - def safe_getitem(device_self, ix): - if ix == "QCOM" and not os.path.exists("/dev/kgsl-3d0"): - return real_getitem(device_self, "CPU") - if ix == "AMD" and not os.path.exists("/dev/kfd"): - return real_getitem(device_self, "CPU") - return real_getitem(device_self, ix) - - with patch.object(Device.__class__, "__getitem__", safe_getitem): - obj = load_oob(f) - - assert isinstance(obj, dict), "Parsed object is not a dictionary" - assert "metadata" in obj, "Metadata key is missing" - - def traverse_jits(d): - if hasattr(d, "captured") and d.captured is not None: - linear_attr = getattr(d.captured, "_linear", None) - if linear_attr is None: - try: - linear_attr = getattr(d.captured, "linear", None) - except AssertionError as e: - if "alloc size must be positive" in str(e): - return - raise - - if linear_attr is not None: - stack = [linear_attr] - visited = set() - while stack: - node = stack.pop() - if id(node) in visited: - continue - visited.add(id(node)) - if hasattr(node, "src") and node.src is not None: - for child in node.src: - stack.append(child) - elif isinstance(d, dict): - for v in d.values(): - traverse_jits(v) - elif isinstance(d, (list, tuple)): - for v in d: - traverse_jits(v) - - if "run_model" in obj: - traverse_jits(obj["run_model"]) - if "run_policy" in obj: - traverse_jits(obj["run_policy"]) +class TestRegistryModels(unittest.TestCase): + def test_model_loads(self): + if not (base_name := os.environ.get("MODEL_BASE_NAME")): + self.skipTest("MODEL_BASE_NAME env var not set, skipping integration test.") + bundle = load_model(os.path.join(os.environ.get("MODEL_DIR", "/tmp/models"), base_name), {"QCOM": "CPU", "AMD": "CPU"}, verify=False, cache=False) + jits = [value for value in bundle.values() if isinstance(value, PlainJit)] + assert "metadata" in bundle and jits and all(len(jit.captured.linear.src) > 0 for jit in jits) diff --git a/openpilot/sunnypilot/modeld_v2/warmup.py b/openpilot/sunnypilot/modeld_v2/warmup.py new file mode 100755 index 0000000000..c4e12b83c1 --- /dev/null +++ b/openpilot/sunnypilot/modeld_v2/warmup.py @@ -0,0 +1,42 @@ +#!/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 +os.environ['GMMU'] = '0' +import multiprocessing +import time +from openpilot.common.hardware.hw import Paths +from openpilot.common.params import Params +from openpilot.common.swaglog import cloudlog +from openpilot.selfdrive.modeld.helpers import chestnut_present +from openpilot.sunnypilot.modeld_v2.helpers import cache_is_fresh, load_model +from openpilot.sunnypilot.models.helpers import get_selected_bundle + + +def main(): + """ + The main entry point for offroad cache creation. + Runs every 10 seconds and checks for stale model caches. + Chestnut cache needs the gpu to build, so that slot waits until it is plugged in, or simply runs during modeld startup. + """ + params, failed, spawn = Params(), set(), multiprocessing.get_context("spawn") + while True: + for bundle in [get_selected_bundle(params, source) for source in (("qcom", "chestnut") if chestnut_present() else ("qcom",))]: + path = os.path.join(Paths.model_root(), bundle.models[0].artifact.fileName) if bundle and bundle.models else "" + if not os.path.isfile(path) or cache_is_fresh(path) or (attempt := (path, os.stat(path).st_mtime_ns)) in failed: + continue + (worker := spawn.Process(target=load_model, args=(path,), daemon=True)).start() + worker.join() + if not cache_is_fresh(path): + cloudlog.error(f"model cache warmup failed for {path} with exit code {worker.exitcode}") + failed.add(attempt) + time.sleep(10) + + +if __name__ == "__main__": + main() diff --git a/openpilot/sunnypilot/models/fetcher.py b/openpilot/sunnypilot/models/fetcher.py index 2a3fc14cb2..0c40c9e598 100644 --- a/openpilot/sunnypilot/models/fetcher.py +++ b/openpilot/sunnypilot/models/fetcher.py @@ -6,12 +6,10 @@ See the LICENSE.md file in the root directory for more details. """ import time -import os import requests from requests.exceptions import (SSLError, RequestException, HTTPError) from openpilot.common.params import Params from openpilot.common.swaglog import cloudlog -from openpilot.common.hardware.hw import Paths from openpilot.sunnypilot.models.helpers import is_bundle_version_compatible from openpilot.cereal import custom @@ -26,35 +24,11 @@ class ModelParser: download_uri.sha256 = download_uri_data.get("sha256") return download_uri - @staticmethod - def _parse_chunk(chunk_data) -> custom.ModelManagerSP.Chunk: - chunk = custom.ModelManagerSP.Chunk() - chunk.fileName = chunk_data.get("file_name") - chunk.sha256 = chunk_data.get("sha256") - return chunk - @staticmethod def _parse_artifact(artifact_data) -> custom.ModelManagerSP.Artifact: artifact = custom.ModelManagerSP.Artifact() artifact.fileName = artifact_data.get("file_name") artifact.downloadUri = ModelParser._parse_download_uri(artifact_data.get("download_uri", {})) - - if "chunks" in artifact_data: - artifact.chunks = [ModelParser._parse_chunk(chunk_data) for chunk_data in artifact_data["chunks"]] - - try: - model_dir = Paths.model_root() - os.makedirs(model_dir, exist_ok=True) - manifest_path = os.path.join(model_dir, f"{artifact.fileName}.chunkmanifest") - num_chunks = str(len(artifact.chunks)) - - if not os.path.exists(manifest_path) or open(manifest_path).read().strip() != num_chunks: - with open(manifest_path, "w") as f: - f.write(num_chunks) - cloudlog.info(f"Wrote chunk manifest for {artifact.fileName}: {num_chunks} chunks") - except Exception as e: - cloudlog.warning(f"Failed to write chunk manifest for {artifact.fileName}: {e}") - return artifact @staticmethod @@ -138,8 +112,8 @@ class ModelCache: class ModelFetcher: """Handles fetching and caching of model data from remote source""" - MODEL_URL = "https://raw.githubusercontent.com/sunnypilot/sunnypilot-models/refs/heads/gh-pages/docs/driving_models_v23.json" - MODEL_URL_CHESTNUT = "https://raw.githubusercontent.com/sunnypilot/sunnypilot-models/refs/heads/gh-pages/docs/driving_models_chestnut_v27.json" + MODEL_URL = "https://raw.githubusercontent.com/sunnypilot/sunnypilot-models/refs/heads/gh-pages/docs/driving_models_v25.json" + MODEL_URL_CHESTNUT = "https://raw.githubusercontent.com/sunnypilot/sunnypilot-models/refs/heads/gh-pages/docs/driving_models_chestnut_v28.json" MODEL_SOURCES = { "qcom": (MODEL_URL, ""), @@ -270,5 +244,3 @@ if __name__ == "__main__": model_overrides = {override.key: override.value for override in bundle.overrides} print(f"Bundle: {bundle.internalName}, Type: {model.type}, Status: {bundle.status}, Overrides: {model_overrides}") print(f"Artifact: {model.artifact.fileName}, Download URI: {model.artifact.downloadUri.uri}") - if model.artifact.chunks: - print(f"Contains {len(model.artifact.chunks)} chunks.") diff --git a/openpilot/sunnypilot/models/helpers.py b/openpilot/sunnypilot/models/helpers.py index 49539c5745..0909f99bd1 100644 --- a/openpilot/sunnypilot/models/helpers.py +++ b/openpilot/sunnypilot/models/helpers.py @@ -16,7 +16,7 @@ from openpilot.common.hardware.hw import Paths from openpilot.selfdrive.modeld.helpers import chestnut_present # SET ME TO THE EXACT JSON VERSION WE SET IN SUNNYPILOT_MODELS REPO -REQUIRED_JSON_VERSION = 20 +REQUIRED_JSON_VERSION = 21 CUSTOM_MODEL_PATH = Paths.model_root() ModelManager = custom.ModelManagerSP @@ -29,9 +29,8 @@ _LAST_VALIDATED_RAW: dict[str, dict | None] = {} def _compute_hash(file_path: str) -> str | None: - from openpilot.common.file_chunker import open_file_chunked try: - with open_file_chunked(file_path) as file: + with open(file_path, 'rb') as file: return hashlib.file_digest(file, "sha256").hexdigest().lower() except FileNotFoundError: return None @@ -58,20 +57,11 @@ def is_bundle_version_compatible(bundle: dict) -> bool: def _bundle_artifacts(bundle: custom.ModelManagerSP.ModelBundle) -> list[tuple[str, str]]: artifacts = [] - from openpilot.common.file_chunker import get_chunk_name for model in getattr(bundle, 'models', []) or []: - for artifact in (getattr(model, 'artifact', None),): - if artifact and getattr(artifact, 'fileName', None): - if len(artifact.chunks) > 0: - for i, chunk in enumerate(artifact.chunks): - chunk_name = get_chunk_name(artifact.fileName, i, len(artifact.chunks)) - if getattr(chunk, 'sha256', None): - artifacts.append((chunk_name, chunk.sha256)) - else: - if getattr(artifact, 'downloadUri', None): - sha256 = getattr(artifact.downloadUri, 'sha256', None) - if sha256: - artifacts.append((artifact.fileName, sha256)) + artifact = getattr(model, 'artifact', None) + if artifact and getattr(artifact, 'fileName', None) and getattr(artifact, 'downloadUri', None): + if sha256 := getattr(artifact.downloadUri, 'sha256', None): + artifacts.append((artifact.fileName, sha256)) return artifacts diff --git a/openpilot/sunnypilot/models/manager.py b/openpilot/sunnypilot/models/manager.py index b313bccbc2..66b0fe28ab 100644 --- a/openpilot/sunnypilot/models/manager.py +++ b/openpilot/sunnypilot/models/manager.py @@ -109,50 +109,6 @@ class ModelManagerSP: # Clean up start time after download completes del self._download_start_times[model.fileName] - async def _download_chunked(self, base_url: str, base_path: str, artifact, skip: frozenset[int] | set[int] = frozenset()) -> None: - from openpilot.common.file_chunker import get_chunk_name, get_manifest_path - - num_chunks = len(artifact.chunks) - if num_chunks == 0: - raise ValueError("No chunks defined in artifact") - - manifest_path = get_manifest_path(base_path) - self._download_start_times[artifact.fileName] = time.monotonic() - - # Shared connection saves a TCP+TLS handshake per chunk. - # Keep sequential: the link saturates on one stream and Session is not thread-safe. - completed = len(skip) - with requests.Session() as session: - for i, _ in enumerate(artifact.chunks): - if i in skip: - continue - chunk_url = get_chunk_name(base_url, i, num_chunks) - chunk_path = get_chunk_name(base_path, i, num_chunks) - chunk_downloaded = 0 - with session.get(chunk_url, stream=True, timeout=DOWNLOAD_TIMEOUT) as response: - response.raise_for_status() - chunk_size = int(response.headers.get("content-length", 0)) - with open(chunk_path, 'wb') as f: # noqa: ASYNC230 - for data in response.iter_content(chunk_size=self._chunk_size): - f.write(data) - chunk_downloaded += len(data) - if self._download_interrupted(): - raise DownloadCancelled("Download cancelled") - intra = chunk_downloaded / max(chunk_size, 1) - progress = min(99.0, ((completed + intra) / num_chunks) * 100) - artifact.downloadProgress.status = custom.ModelManagerSP.DownloadStatus.downloading - artifact.downloadProgress.progress = progress - artifact.downloadProgress.eta = self._calculate_eta(artifact.fileName, progress) - self._sync_artifact_progress(artifact) - self._report_status() - completed += 1 - - with open(manifest_path, 'w') as f: # noqa: ASYNC230 - f.write(str(num_chunks)) - if os.path.isfile(base_path): # noqa: ASYNC240 - os.remove(base_path) - del self._download_start_times[artifact.fileName] - async def _process_artifact(self, artifact, destination_path: str) -> None: if not artifact.downloadUri.uri: return None @@ -165,28 +121,7 @@ class ModelManagerSP: full_path = os.path.join(destination_path, filename) try: - # progress counts only valid chunks so a resumed download continues the - # bar from where verification left it, instead of falling back to zero - is_cached = False - valid_chunks: set[int] = set() - if len(artifact.chunks) > 0: - from openpilot.common.file_chunker import get_chunk_name - num_chunks = len(artifact.chunks) - for i, chunk in enumerate(artifact.chunks): - if self._download_interrupted(): - raise DownloadCancelled("Download cancelled") - if await verify_file(get_chunk_name(full_path, i, num_chunks), chunk.sha256): - valid_chunks.add(i) - artifact.downloadProgress.status = custom.ModelManagerSP.DownloadStatus.verifying - artifact.downloadProgress.progress = (len(valid_chunks) / num_chunks) * 100 - self._sync_artifact_progress(artifact) - self._report_status() - is_cached = len(valid_chunks) == num_chunks - else: - if await verify_file(full_path, expected_hash): - is_cached = True - - if is_cached: + if await verify_file(full_path, expected_hash): artifact.downloadProgress.status = custom.ModelManagerSP.DownloadStatus.cached artifact.downloadProgress.progress = 100 artifact.downloadProgress.eta = 0 @@ -194,17 +129,9 @@ class ModelManagerSP: self._report_status() return - if len(artifact.chunks) > 0: - await self._download_chunked(url, full_path, artifact, skip=valid_chunks) - from openpilot.common.file_chunker import get_chunk_name - for i, chunk in enumerate(artifact.chunks): - chunk_path = get_chunk_name(full_path, i, len(artifact.chunks)) - if not await verify_file(chunk_path, chunk.sha256): - raise ValueError(f"Hash validation failed for chunk {i+1} of {filename}") - else: - await self._download_file(url, full_path, artifact) - if not await verify_file(full_path, expected_hash): - raise ValueError(f"Hash validation failed for {filename}") + await self._download_file(url, full_path, artifact) + if not await verify_file(full_path, expected_hash): + raise ValueError(f"Hash validation failed for {filename}") artifact.downloadProgress.status = custom.ModelManagerSP.DownloadStatus.downloaded artifact.downloadProgress.progress = 100 @@ -213,7 +140,6 @@ class ModelManagerSP: self._report_status() except DownloadCancelled: - # a cancel keeps whatever is on disk: complete chunks resume the next attempt self._download_start_times.pop(artifact.fileName, None) artifact.downloadProgress.status = custom.ModelManagerSP.DownloadStatus.failed artifact.downloadProgress.eta = 0 @@ -225,9 +151,8 @@ class ModelManagerSP: except Exception as e: cloudlog.error(f"Error downloading {filename}: {str(e)}") - for f in [full_path] + [p for p in (os.path.join(destination_path, f) for f in os.listdir(destination_path)) if filename in p]: - if os.path.isfile(f): # noqa: ASYNC240 - os.remove(f) + if os.path.isfile(full_path): # noqa: ASYNC240 + os.remove(full_path) artifact.downloadProgress.status = custom.ModelManagerSP.DownloadStatus.failed artifact.downloadProgress.eta = 0 self._sync_artifact_progress(artifact) @@ -360,12 +285,11 @@ class ModelManagerSP: if model.artifact.fileName: active_files.append(model.artifact.fileName) - # Remove all files except active ones (including their chunk files) + # Remove all files except active ones model_dir = Paths.model_root() try: for filename in os.listdir(model_dir): - base = filename.split('.chunk')[0] if '.chunk' in filename else filename - if base not in active_files and filename not in active_files: + if filename not in active_files: file_path = os.path.join(model_dir, filename) if os.path.isfile(file_path): os.remove(file_path) diff --git a/openpilot/sunnypilot/models/tests/test_manager_download.py b/openpilot/sunnypilot/models/tests/test_manager_download.py index b995d4fbd1..0ae0273ec4 100644 --- a/openpilot/sunnypilot/models/tests/test_manager_download.py +++ b/openpilot/sunnypilot/models/tests/test_manager_download.py @@ -17,11 +17,9 @@ from typing import Any from unittest import mock import requests -from urllib3.connectionpool import HTTPConnectionPool from openpilot.cereal import custom from openpilot.common.test import OpenpilotTestCase -from openpilot.common.file_chunker import get_chunk_name, get_manifest_path from openpilot.selfdrive.test.helpers import http_server_context from openpilot.sunnypilot.models import manager as manager_module from openpilot.sunnypilot.models.fetcher import ModelFetcher, get_cached_bundles @@ -30,7 +28,6 @@ from openpilot.sunnypilot.models.helpers import (get_active_bundle, get_active_s resolve_bundle_by_ref, validate_active_bundles) from openpilot.sunnypilot.models.manager import ModelManagerSP -CHUNK_BODIES = [b'A' * 5000, b'B' * 5000, b'C' * 3000] WHOLE_BODY = b'Z' * 9000 @@ -49,12 +46,7 @@ class DownloadHandler(http.server.BaseHTTPRequestHandler): pass def _body_for(self, path): - if path.endswith('.whole'): - return WHOLE_BODY - for i in range(len(CHUNK_BODIES)): - if path.endswith(get_chunk_name('', i, len(CHUNK_BODIES))): - return CHUNK_BODIES[i] - return None + return WHOLE_BODY if path.endswith('.pkl') else None def do_GET(self): type(self).request_paths.append(self.path) @@ -119,35 +111,19 @@ class ManagerDownloadTestBase(OpenpilotTestCase): if artifact is not None: self.reported.append(float(artifact.downloadProgress.progress)) - def make_artifact(self, chunked: bool): + def make_artifact(self): bundle = custom.ModelManagerSP.ModelBundle.new_message() bundle.init('models', 1) artifact = bundle.models[0].artifact artifact.fileName = 'driving_test_tinygrad.pkl' - if chunked: - artifact.downloadUri.uri = self.base_url + '/driving_test_tinygrad.pkl' - artifact.downloadUri.sha256 = sha256(b''.join(CHUNK_BODIES)) - artifact.init('chunks', len(CHUNK_BODIES)) - for i, body in enumerate(CHUNK_BODIES): - artifact.chunks[i].sha256 = sha256(body) - else: - artifact.downloadUri.uri = self.base_url + '/driving_test_tinygrad.pkl.whole' - artifact.downloadUri.sha256 = sha256(WHOLE_BODY) + artifact.downloadUri.uri = self.base_url + '/driving_test_tinygrad.pkl' + artifact.downloadUri.sha256 = sha256(WHOLE_BODY) self._bundle = bundle self.artifact = artifact return artifact - def chunk_paths(self, base_path): - return [get_chunk_name(base_path, i, len(CHUNK_BODIES)) for i in range(len(CHUNK_BODIES))] - - def assert_no_partials(self, base_path): - leftovers = [p for p in [base_path, get_manifest_path(base_path)] + self.chunk_paths(base_path) - if os.path.isfile(p)] - assert leftovers == [], f"partial files left behind: {leftovers}" - - class TestManagerDownload(ManagerDownloadTestBase): - """Exercises the real _download_file / _download_chunked against a local server.""" + """Exercises the real _download_file against a local server.""" def run_with_server(self, fn): with http_server_context(handler=DownloadHandler) as (host, port): @@ -156,7 +132,7 @@ class TestManagerDownload(ManagerDownloadTestBase): def test_download_file_writes_exact_bytes(self): def body(): - artifact = self.make_artifact(chunked=False) + artifact = self.make_artifact() path = os.path.join(self.dest, artifact.fileName) asyncio.run(self.manager._download_file(artifact.downloadUri.uri, path, artifact)) with open(path, 'rb') as f: @@ -166,111 +142,67 @@ class TestManagerDownload(ManagerDownloadTestBase): assert artifact.fileName not in self.manager._download_start_times self.run_with_server(body) - def test_download_chunked_writes_all_chunks_and_manifest(self): - def body(): - artifact = self.make_artifact(chunked=True) - base_path = os.path.join(self.dest, artifact.fileName) - asyncio.run(self.manager._download_chunked(artifact.downloadUri.uri, base_path, artifact)) - - for i, expected in enumerate(CHUNK_BODIES): - with open(get_chunk_name(base_path, i, len(CHUNK_BODIES)), 'rb') as f: - assert f.read() == expected, f"chunk {i} body mismatch" - - with open(get_manifest_path(base_path)) as f: - assert f.read() == str(len(CHUNK_BODIES)) - - assert not os.path.isfile(base_path), "base file should be removed after chunking" - assert artifact.fileName not in self.manager._download_start_times - self.run_with_server(body) - def test_progress_is_monotonic_and_bounded(self): def body(): - artifact = self.make_artifact(chunked=True) + artifact = self.make_artifact() base_path = os.path.join(self.dest, artifact.fileName) - asyncio.run(self.manager._download_chunked(artifact.downloadUri.uri, base_path, artifact)) + asyncio.run(self.manager._download_file(artifact.downloadUri.uri, base_path, artifact)) assert self.reported, "expected progress reports" for a, b in zip(self.reported, self.reported[1:], strict=False): assert b >= a, f"progress went backwards: {a} -> {b}" - assert max(self.reported) <= 99.0, f"chunked progress must stay <=99 until verify, got {max(self.reported)}" - self.run_with_server(body) - - def test_session_is_reused_across_chunks(self): - """One connection pool shared across every chunk.""" - def body(): - artifact = self.make_artifact(chunked=True) - base_path = os.path.join(self.dest, artifact.fileName) - - pools = [] - original = HTTPConnectionPool.urlopen - - def tracked(pool_self, *args, **kwargs): - pools.append(id(pool_self)) - return original(pool_self, *args, **kwargs) - - with mock.patch.object(HTTPConnectionPool, 'urlopen', tracked): - asyncio.run(self.manager._download_chunked(artifact.downloadUri.uri, base_path, artifact)) - - assert len(pools) == len(CHUNK_BODIES), f"expected one request per chunk, got {len(pools)}" - assert len(set(pools)) == 1, f"connection pool not reused across chunks: {len(set(pools))} pools" self.run_with_server(body) def test_http_error_propagates(self): def body(): - artifact = self.make_artifact(chunked=True) + artifact = self.make_artifact() base_path = os.path.join(self.dest, artifact.fileName) - failing = '/' + os.path.basename(get_chunk_name(artifact.downloadUri.uri, 1, len(CHUNK_BODIES))) - DownloadHandler.fail_paths = {failing: 404} + DownloadHandler.fail_paths = {'/' + artifact.fileName: 404} with self.assertRaises(requests.exceptions.HTTPError): - asyncio.run(self.manager._download_chunked(artifact.downloadUri.uri, base_path, artifact)) - - # chunk 1 failed, so its file and the manifest must not exist - assert not os.path.isfile(get_chunk_name(base_path, 1, len(CHUNK_BODIES))) - assert not os.path.isfile(get_manifest_path(base_path)) + asyncio.run(self.manager._download_file(artifact.downloadUri.uri, base_path, artifact)) + assert not os.path.isfile(base_path) self.run_with_server(body) def test_cancellation_mid_transfer(self): """Cancellation is checked inside the byte loop; it must still fire after the port.""" def body(): - artifact = self.make_artifact(chunked=True) + artifact = self.make_artifact() base_path = os.path.join(self.dest, artifact.fileName) self.manager.params.get.return_value = None # cancelled with self.assertRaises(Exception) as ctx: - asyncio.run(self.manager._download_chunked(artifact.downloadUri.uri, base_path, artifact)) + asyncio.run(self.manager._download_file(artifact.downloadUri.uri, base_path, artifact)) assert 'cancelled' in str(ctx.exception).lower() - assert not os.path.isfile(get_manifest_path(base_path)) self.run_with_server(body) def test_repeat_downloads_are_stable(self): """Back-to-back runs must produce identical bytes and leak no start-time state.""" def body(): for _ in range(2): - artifact = self.make_artifact(chunked=True) + artifact = self.make_artifact() base_path = os.path.join(self.dest, artifact.fileName) - asyncio.run(self.manager._download_chunked(artifact.downloadUri.uri, base_path, artifact)) - for i, expected in enumerate(CHUNK_BODIES): - with open(get_chunk_name(base_path, i, len(CHUNK_BODIES)), 'rb') as f: - assert f.read() == expected + asyncio.run(self.manager._download_file(artifact.downloadUri.uri, base_path, artifact)) + with open(base_path, 'rb') as f: + assert f.read() == WHOLE_BODY assert self.manager._download_start_times == {} self.run_with_server(body) def test_download_ref_present_keeps_download_alive(self): """A pending download request (DownloadRef set) must not be cancelled mid-transfer.""" def body(): - artifact = self.make_artifact(chunked=True) + artifact = self.make_artifact() base_path = os.path.join(self.dest, artifact.fileName) self.manager.params.get.side_effect = lambda key: b"ref" if key == "ModelManager_DownloadRef" else None self.manager._download_ref = b"ref" - asyncio.run(self.manager._download_chunked(artifact.downloadUri.uri, base_path, artifact)) - assert os.path.isfile(get_manifest_path(base_path)) + asyncio.run(self.manager._download_file(artifact.downloadUri.uri, base_path, artifact)) + assert os.path.isfile(base_path) self.run_with_server(body) def test_cancellation_via_download_ref(self): """Removing DownloadRef mid-transfer cancels the download.""" def body(): - artifact = self.make_artifact(chunked=True) + artifact = self.make_artifact() base_path = os.path.join(self.dest, artifact.fileName) checks = {"n": 0} @@ -283,20 +215,19 @@ class TestManagerDownload(ManagerDownloadTestBase): self.manager.params.get.side_effect = get self.manager._download_ref = b"ref" with self.assertRaises(Exception) as ctx: - asyncio.run(self.manager._download_chunked(artifact.downloadUri.uri, base_path, artifact)) + asyncio.run(self.manager._download_file(artifact.downloadUri.uri, base_path, artifact)) assert 'cancelled' in str(ctx.exception).lower() - assert not os.path.isfile(get_manifest_path(base_path)) self.run_with_server(body) def test_replaced_download_ref_queues_instead_of_cancelling(self): """Selecting another model mid-transfer lets the running download finish.""" def body(): - artifact = self.make_artifact(chunked=True) + artifact = self.make_artifact() base_path = os.path.join(self.dest, artifact.fileName) self.manager.params.get.side_effect = lambda key: b"other-ref" if key == "ModelManager_DownloadRef" else None self.manager._download_ref = b"ref" - asyncio.run(self.manager._download_chunked(artifact.downloadUri.uri, base_path, artifact)) - assert os.path.isfile(get_manifest_path(base_path)) + asyncio.run(self.manager._download_file(artifact.downloadUri.uri, base_path, artifact)) + assert os.path.isfile(base_path) self.run_with_server(body) def test_replaced_download_ref_is_kept(self): @@ -315,11 +246,10 @@ class TestManagerDownload(ManagerDownloadTestBase): def test_cached_bundle_cancel_skips_slot_write(self): """A cancel must stop an already-on-disk bundle before it is applied to the slot.""" def body(): - artifact = self.make_artifact(chunked=True) + artifact = self.make_artifact() base_path = os.path.join(self.dest, artifact.fileName) - for i, data in enumerate(CHUNK_BODIES): - with open(get_chunk_name(base_path, i, len(CHUNK_BODIES)), 'wb') as f: - f.write(data) + with open(base_path, 'wb') as f: + f.write(WHOLE_BODY) self._bundle.ref = "test-ref" params, store = self._make_params_with_store() store["ModelManager_DownloadRef"] = None # removed -> cancelled @@ -329,42 +259,7 @@ class TestManagerDownload(ManagerDownloadTestBase): asyncio.run(self.manager._download_bundle(self._bundle, self.dest, "qcom")) assert 'cancelled' in str(ctx.exception).lower() assert "ModelManager_ActiveBundle" not in store - assert all(os.path.isfile(p) for p in self.chunk_paths(base_path)), "cancel must not delete cached chunks" - self.run_with_server(body) - - def test_resume_skips_valid_chunks(self): - """A chunk already on disk is kept and not re-downloaded; progress starts above its share.""" - def body(): - artifact = self.make_artifact(chunked=True) - base_path = os.path.join(self.dest, artifact.fileName) - with open(get_chunk_name(base_path, 0, len(CHUNK_BODIES)), 'wb') as f: - f.write(CHUNK_BODIES[0]) - - asyncio.run(self.manager._process_artifact(artifact, self.dest)) - - chunk0_suffix = get_chunk_name('', 0, len(CHUNK_BODIES)) - assert not any(p.endswith(chunk0_suffix) for p in DownloadHandler.request_paths), "valid chunk was re-downloaded" - for i, expected in enumerate(CHUNK_BODIES): - with open(get_chunk_name(base_path, i, len(CHUNK_BODIES)), 'rb') as f: - assert f.read() == expected - assert os.path.isfile(get_manifest_path(base_path)) - assert min(self.reported) >= (1 / len(CHUNK_BODIES)) * 100 - 1, "progress must not restart below the resumed share" - self.run_with_server(body) - - def test_verify_reports_valid_fraction_then_cached(self): - """A fully cached bundle publishes climbing verify progress and ends cached.""" - def body(): - artifact = self.make_artifact(chunked=True) - base_path = os.path.join(self.dest, artifact.fileName) - for i, data in enumerate(CHUNK_BODIES): - with open(get_chunk_name(base_path, i, len(CHUNK_BODIES)), 'wb') as f: - f.write(data) - - asyncio.run(self.manager._process_artifact(artifact, self.dest)) - - assert DownloadHandler.request_paths == [], "cached bundle must not hit the network" - assert [round(p) for p in self.reported[:3]] == [33, 67, 100] - assert artifact.downloadProgress.status == custom.ModelManagerSP.DownloadStatus.cached + assert os.path.isfile(base_path), "cancel must not delete the cached model" self.run_with_server(body) def _make_params_with_store(self): @@ -384,7 +279,7 @@ class TestManagerDownload(ManagerDownloadTestBase): def test_download_writes_qcom_slot(self): """A download resolved to the qcom source writes the qcom active bundle slot only.""" def body(): - artifact = self.make_artifact(chunked=True) + artifact = self.make_artifact() self._bundle.ref = "test-ref" self._bundle.minimumSelectorVersion = helpers.REQUIRED_JSON_VERSION params, store = self._make_params_with_store() @@ -396,15 +291,13 @@ class TestManagerDownload(ManagerDownloadTestBase): assert self.manager.selected_bundle.status == custom.ModelManagerSP.DownloadStatus.downloaded assert self.manager.active_bundle is not None and self.manager.active_bundle.ref == "test-ref" assert self.manager.active_bundle.status == custom.ModelManagerSP.DownloadStatus.downloaded - chunk_names = [get_chunk_name(artifact.fileName, i, len(artifact.chunks)) for i in range(len(artifact.chunks))] - missing = [c for c in chunk_names if not os.path.isfile(os.path.join(self.dest, c))] - assert missing == [], f"chunks missing from the cache: {missing}" + assert os.path.isfile(os.path.join(self.dest, artifact.fileName)) self.run_with_server(body) def test_download_writes_chestnut_slot(self): """A download resolved to the chestnut source writes the chestnut active bundle slot only.""" def body(): - self.make_artifact(chunked=True) + self.make_artifact() self._bundle.ref = "big-ref" self._bundle.minimumSelectorVersion = helpers.REQUIRED_JSON_VERSION params, store = self._make_params_with_store() @@ -782,7 +675,7 @@ class TestEffectiveSource(OpenpilotTestCase): @unittest.skipUnless(os.environ.get('RUN_INTEGRATION_TESTS'), 'requires external network') class TestLiveModelManifest(OpenpilotTestCase): - """Every artifact and chunk URL in the published manifest must resolve.""" + """Every artifact URL in the published manifest must resolve.""" def test_all_manifest_urls_available(self): from openpilot.sunnypilot.models.fetcher import ModelFetcher @@ -794,10 +687,7 @@ class TestLiveModelManifest(OpenpilotTestCase): for bundle in manifest.get('bundles', []): for model in bundle.get('models', []): artifact = model['artifact'] - url = artifact['download_uri']['url'] - chunks = artifact.get('chunks', []) - urls = ([url] if not chunks - else [get_chunk_name(url, i, len(chunks)) for i in range(len(chunks))]) + urls = [artifact['download_uri']['url']] for u in urls: try: r = session.head(u, timeout=15, allow_redirects=True) diff --git a/openpilot/system/manager/process_config.py b/openpilot/system/manager/process_config.py index c95b110539..8959016844 100644 --- a/openpilot/system/manager/process_config.py +++ b/openpilot/system/manager/process_config.py @@ -175,6 +175,7 @@ procs += [ # Models PythonProcess("models_manager", "openpilot.sunnypilot.models.manager", only_offroad), NativeProcess("modeld_tinygrad", "openpilot/sunnypilot/modeld_v2", ["./modeld"], and_(only_onroad, is_tinygrad_model)), + PythonProcess("modeld_warmup", "openpilot.sunnypilot.modeld_v2.warmup", only_offroad), # Backup PythonProcess("backup_manager", "openpilot.sunnypilot.sunnylink.backups.manager", and_(only_offroad, sunnylink_ready_shim)), diff --git a/release/ci/model_generator.py b/release/ci/model_generator.py index 3b5603c4ae..5d0ca73644 100755 --- a/release/ci/model_generator.py +++ b/release/ci/model_generator.py @@ -54,18 +54,10 @@ def create_pkl_name(full_name: str) -> str: def _hash_pkl(pkl_path: Path) -> str: - manifest = Path(f"{pkl_path}.chunkmanifest") - if manifest.exists(): - num_chunks = int(manifest.read_text().strip()) - paths = [Path(f"{pkl_path}.chunk{i + 1:02d}of{num_chunks:02d}") for i in range(num_chunks)] - else: - paths = [pkl_path] - digest = hashlib.sha256() - for path in paths: - with path.open('rb') as f: - while block := f.read(1024 * 1024): - digest.update(block) + with pkl_path.open('rb') as f: + while block := f.read(1024 * 1024): + digest.update(block) return digest.hexdigest() @@ -74,22 +66,9 @@ def _find_driving_pkl(output_path: Path) -> Path | None: 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 _hash_onnx_files(model_dir: Path) -> str | None: onnx_files = [f for f in sorted(model_dir.glob("*.onnx")) if f.stat().st_size > 1024] if not onnx_files: @@ -102,22 +81,9 @@ def _hash_onnx_files(model_dir: Path) -> str | None: return digest.hexdigest() -def generate_chunked_model(driving_pkl: Path) -> dict: +def generate_model(driving_pkl: Path) -> dict: tinygrad_hash = _hash_pkl(driving_pkl) - chunks_config = [] - manifest_file = Path(f"{driving_pkl}.chunkmanifest") - if manifest_file.exists(): - num_chunks = int(manifest_file.read_text().strip()) - for i in range(num_chunks): - chunk_path = Path(f"{driving_pkl}.chunk{i + 1:02d}of{num_chunks:02d}") - if chunk_path.exists(): - chunk_hash = hashlib.sha256(chunk_path.read_bytes()).hexdigest() - chunks_config.append({ - "file_name": chunk_path.name, - "sha256": chunk_hash - }) - artifact_data = { "file_name": driving_pkl.name, "download_uri": { @@ -126,9 +92,6 @@ def generate_chunked_model(driving_pkl: Path) -> dict: } } - if chunks_config: - artifact_data["chunks"] = chunks_config - return { "type": "chunked", "artifact": artifact_data, @@ -193,11 +156,11 @@ if __name__ == "__main__": if _pkl: new_pkl = _output_dir / f"driving_{_pkl}_tinygrad.pkl" if not new_pkl.exists(): - _driving_pkl = _rename_pkl_with_chunks(_driving_pkl, new_pkl) + _driving_pkl = _driving_pkl.rename(new_pkl) else: _driving_pkl = new_pkl - _model_metadata = generate_chunked_model(_driving_pkl) + _model_metadata = generate_model(_driving_pkl) _onnx_sha256 = args.onnx_sha256 or _hash_onnx_files(Path(args.model_dir)) create_metadata_json([_model_metadata], _output_dir, args.custom_name, _short_name, args.is_20hz, args.upstream_branch, onnx_sha256=_onnx_sha256, is_big=is_big) diff --git a/release/ci/upload_default_model.py b/release/ci/upload_default_model.py index 2abd5087c0..8fce46a381 100644 --- a/release/ci/upload_default_model.py +++ b/release/ci/upload_default_model.py @@ -79,8 +79,6 @@ def main(): artifact = bundle['models'][0]['artifact'] hf_base = f"https://huggingface.co/datasets/{args.hf_repo}/resolve/main/{args.hf_defaults_path}/{folder_name}" artifact['download_uri']['url'] = f"{hf_base}/{artifact['file_name']}" - for chunk in artifact.get('chunks', []): - chunk['url'] = f"{hf_base}/{chunk['file_name']}" print(f"Uploading model to {args.hf_defaults_path}/{folder_name}/") api.upload_folder( diff --git a/tinygrad_repo b/tinygrad_repo index fe5d3169ba..69fd4d34d4 160000 --- a/tinygrad_repo +++ b/tinygrad_repo @@ -1 +1 @@ -Subproject commit fe5d3169ba4f41d0947ad174925f413cbea9d056 +Subproject commit 69fd4d34d4430036378aba2b43f9d339cdd6ea73