Compare commits

..

21 Commits

Author SHA1 Message Date
discountchubbs f6e2fb1c33 Merge remote-tracking branch 'origin/master' into rebuild-deep-rl 2026-10-14 08:04:54 -07:00
James Vecellio-Grant 9bebe42f73 default models: fix big (#2046)
* fix big

* whoops
2026-10-13 23:45:39 -04:00
James Vecellio-Grant f4996a0529 lint 2026-10-13 20:39:15 -07:00
James Vecellio-Grant 0414843e2f method() persistent_load 2026-10-13 20:09:00 -07:00
James Vecellio-Grant 5acb9d2c8e whoops 2026-10-13 20:01:23 -07:00
discountchubbs 2bf8f88978 bump 2026-10-13 14:39:33 -07:00
discountchubbs e6e5b1c165 bump tg to upstream tg/tg:master. no rebuild needed 2026-10-13 14:36:10 -07:00
discountchubbs b7d06aa440 publish faster 2026-10-13 11:21:32 -07:00
discountchubbs a9ccd8fedf zed compress 2026-10-13 10:44:41 -07:00
discountchubbs 6220734b2d parallel 2026-10-13 08:56:33 -07:00
discountchubbs f3f7296776 stream chunks for big models to save oom 2026-10-13 08:04:31 -07:00
discountchubbs 9f3a2607ee dedup chestnutty 2026-10-12 23:31:56 -07:00
discountchubbs b4c3a19db1 skip out on op modules 2026-10-12 23:22:02 -07:00
discountchubbs 04098893d9 run as bash: for when i wake up:
these models will fail, look at workflow, likely need env setup now that we use recursion
2026-10-12 23:14:11 -07:00
discountchubbs bc064e892f standin for chestnut linear_compile 2026-10-12 23:09:55 -07:00
discountchubbs be51bc793d 5ms speed up on qcom 2026-10-12 22:36:27 -07:00
discountchubbs f7c65f7d6a lint 2026-10-12 20:55:33 -07:00
discountchubbs bd27f3b54f offroad proc for cache warmup 2026-10-12 19:53:39 -07:00
discountchubbs d7c1827134 new revolution 2026-10-12 17:02:05 -07:00
discountchubbs 1344824746 modeld_v2: shes deeeeeep 2026-10-12 16:34:00 -07:00
discountchubbs 75b5926f4f models: fix qcom float16 overflow 2026-10-09 12:36:31 -07:00
30 changed files with 679 additions and 733 deletions
+41 -65
View File
@@ -47,6 +47,7 @@ jobs:
if [ "${{ inputs.target }}" = "big" ]; then
NAME=$(python3 -c "from openpilot.sunnypilot.models.model_name import DEFAULT_BIG_MODEL; print(DEFAULT_BIG_MODEL)")
HF_DEFAULTS_PATH="models/defaults/big"
ONNX_PATH="openpilot/selfdrive/modeld/models/big_driving_tinygrad.pkl"
elif [ "${{ inputs.target }}" = "dm" ]; then
ONNX_PATH="openpilot/selfdrive/modeld/models/dmonitoring_model.onnx"
HF_DEFAULTS_PATH="models/defaults/dm"
@@ -116,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 }}
@@ -141,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,19 +207,33 @@ jobs:
raise RuntimeError('Chestnut PCIe link not ready after 10 attempts')
"
- name: Chunk big pkl
- name: Compile big warp
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')
"
export PYTHONPATH="${PYTHONPATH}:${{ github.workspace }}/tinygrad_repo:${{ github.workspace }}"
CHESTNUT_TG_FLAGS="DEBUG=1 DEV=USB+AMD:LLVM FRAME_DEV=CPU FLOAT16=1 JIT_BATCH_SIZE=0 GMMU=0 TC_OPT=2 TC_MIN_GLOBALS=32"
COMPILER="${{ github.workspace }}/tinygrad_repo/examples/openpilot"
MODELS_DIR="${{ github.workspace }}/openpilot/selfdrive/modeld/models"
MODEL_SIZE=$(python3 -c "from openpilot.common.transformations.model import MEDMODEL_INPUT_SIZE as s; print(f'{s[0]}x{s[1]}')")
for res in $(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}')"); do
CAM_W=$(echo $res | cut -d x -f1)
CAM_H=$(echo $res | cut -d x -f2)
STRIDE_INFO=$(python3 -c "
import sys
from openpilot.system.camerad.cameras.nv12_info import get_nv12_info
s,y,u,sz = get_nv12_info(int(sys.argv[1]), int(sys.argv[2]))
print(f'{s},{y},{u},{sz}')
" $CAM_W $CAM_H)
WARP_PKL="${MODELS_DIR}/big_driving_warp_${res}_tinygrad.pkl"
env ${CHESTNUT_TG_FLAGS} python3 "${COMPILER}/compile_warp.py" \
--frame ${CAM_W},${CAM_H},${STRIDE_INFO} \
--warp-to ${MODEL_SIZE} \
--layout yuv420 \
--frames 2 \
--output "${WARP_PKL}"
done
- name: Prepare output
env:
@@ -246,10 +246,9 @@ 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=$(sha256sum "$MODELS_DIR/${PKL_BASE}" | cut -d' ' -f1)
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)
python3 "${{ github.workspace }}/release/ci/model_generator.py" \
--model-dir "$MODELS_DIR" \
@@ -258,6 +257,8 @@ jobs:
--upstream-branch "${{ needs.resolve.outputs.onnx_ref }}" \
--onnx-sha256 "$PKL_LFS_SHA256"
cp "$MODELS_DIR"/big_driving_warp_* "$OUTPUT_DIR/" || true
echo "model-${MODEL_NAME}-${{ github.run_number }}" > "$OUTPUT_DIR/artifact_name.txt"
- name: Upload big model artifact
@@ -335,6 +336,8 @@ jobs:
--hf-defaults-path "${{ needs.resolve.outputs.hf_defaults_path }}" \
--artifact-name "$ARTIFACT_NAME" \
--model-dir output \
--onnx-path "${{ needs.resolve.outputs.onnx_path }}" \
--onnx-ref "${{ needs.resolve.outputs.onnx_ref }}" \
--model-name "${{ needs.resolve.outputs.model_name }}" \
--tinygrad-ref "${{ needs.resolve.outputs.tinygrad_ref }}" \
--run-number "${{ github.run_number }}"
@@ -358,24 +361,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': [{
@@ -389,8 +382,7 @@ jobs:
'type': 'chunked',
'artifact': {
'file_name': pkl_name,
'download_uri': {'url': '', 'sha256': digest.hexdigest()},
'chunks': chunks
'download_uri': {'url': '', 'sha256': digest.hexdigest()}
}
}]
}]
@@ -453,20 +445,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
@@ -499,8 +477,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
@@ -514,4 +491,3 @@ jobs:
run: |
source ${UV_PROJECT_ENVIRONMENT}/bin/activate
PYTHONPATH=$PYTHONPATH:${{ github.workspace }}/ ${{ github.workspace }}/scripts/manage-powersave.py --enable
@@ -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:
@@ -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,26 +51,23 @@ 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
ENCODED_URL=$(python3 -c "import urllib.parse; print(urllib.parse.quote('${BASE_URL}/${warp}', safe=':/'))")
printf '%s\t%s\n' "$ENCODED_URL" "${DEST_DIR}/${warp}" >> "$DOWNLOAD_LIST"
done
elif [ "$CANONICAL" = "big_driving_tinygrad.pkl" ]; then
for warp in big_driving_warp_1928x1208_tinygrad.pkl big_driving_warp_1344x760_tinygrad.pkl; do
ENCODED_URL=$(python3 -c "import urllib.parse; print(urllib.parse.quote('${BASE_URL}/${warp}', safe=':/'))")
printf '%s\t%s\n' "$ENCODED_URL" "${DEST_DIR}/${warp}" >> "$DOWNLOAD_LIST"
done
fi
}
@@ -82,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#* }"
+22 -10
View File
@@ -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
@@ -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
+9 -25
View File
@@ -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
+1 -1
View File
@@ -55,7 +55,7 @@ def chestnut_action(command):
def compile_model(onnx_path, pkl_path):
onnx_path, target_pkl_path = File(onnx_path).abspath, File(pkl_path).abspath
cmd = (f'{tg_flags} {mac_brew_string} {taskset}python3 "{compiler}/compile_onnx.py" '
f'"{onnx_path}" "{target_pkl_path}" --device-input "*" --out-of-band --benchmark-runs 1')
f'"{onnx_path}" "{target_pkl_path}" --out-of-band --benchmark-runs 1')
lenv.Command(
target_pkl_path,
tinygrad_files + [onnx_path, Value(cmd)],
+10 -9
View File
@@ -1,9 +1,11 @@
import io
import pickle
import struct
import sys
from pathlib import Path
from openpilot.common.file_chunker import get_manifest_path
from openpilot.common.hardware import AGNOS
from openpilot.common.hardware.usb import CHESTNUT_USB_PRODUCT, USB_DEVICES_PATH, is_chestnut_usb_id
MODELS_DIR = Path(__file__).resolve().parent / 'models'
@@ -13,15 +15,14 @@ def modeld_pkl_path(chestnut: bool):
prefix = 'big_' if chestnut else ''
return MODELS_DIR / f'{prefix}driving_tinygrad.pkl'
def load_oob(f):
opcodes = f.read(struct.unpack('<q', f.read(8))[0])
def buffers():
while (h := f.read(8)):
pb = pickle.PickleBuffer(bytearray(struct.unpack('<q', h)[0]))
if f.readinto(pb) != pb.raw().nbytes:
raise EOFError("incomplete model buffer")
yield pb
return pickle.load(io.BytesIO(opcodes), buffers=buffers())
def load_oob(f, chestnut=False):
import numpy as np
from tinygrad import Context, Tensor
device = 'USB+AMD:LLVM' if chestnut else 'QCOM' if AGNOS else 'METAL' if sys.platform == 'darwin' else 'CPU:LLVM'
with Context(DEV=device):
opcodes = f.read(struct.unpack('<q', f.read(8))[0])
buffers = Tensor(np.frombuffer(f.read(), np.uint8)).realize().uop.buffer
return type("_Unpickler", (pickle.Unpickler,), {"persistent_load": lambda self, pid: buffers.view(*pid)})(io.BytesIO(opcodes)).load()
def chestnut_present() -> bool:
for d in USB_DEVICES_PATH.glob("*"):
+5 -4
View File
@@ -134,8 +134,8 @@ class FrameMeta:
def input_view(buffer: Buffer, shape: tuple[int, ...], dtype: DType, offset: int) -> Tensor:
view = buffer.view(math.prod(shape), dtype, offset).ensure_allocated()
return Tensor(UOp.from_buffer(view)).reshape(shape)
view = buffer.view(math.prod(shape) * dtype.itemsize, offset).ensure_allocated()
return Tensor(UOp.from_buffer(view, dtype)).reshape(shape)
class ModelState(ModelStateBase):
@@ -143,7 +143,7 @@ class ModelState(ModelStateBase):
def __init__(self, cam_w: int, cam_h: int, chestnut: bool):
ModelStateBase.__init__(self)
jits = load_oob(open_file_chunked(modeld_pkl_path(chestnut)))
jits = load_oob(open_file_chunked(modeld_pkl_path(chestnut)), chestnut)
self.model_device = jits['input_specs']['new_img'][2]
self.input_shapes = {name: (shape, np.dtype(dtype)) for name, (shape, dtype, _) in jits['input_specs'].items()}
self.state_pairs = {name: f'next_{name}' for name in self.input_shapes if f'next_{name}' in jits['metadata']['output_shapes']}
@@ -230,7 +230,8 @@ def main(demo=False):
CHESTNUT = chestnut_present() and chestnut_compiled()
if CHESTNUT:
os.environ['HCQDEV_WAIT_TIMEOUT_MS'] = '3000'
from tinygrad.runtime.ops_amd import AMDDevice
AMDDevice.wait_timeout_ms = 3000
params = Params()
params.put_bool("ChestnutLoading", CHESTNUT)
params.remove("ChestnutActive")
@@ -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)")
@@ -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)")
+246 -77
View File
@@ -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('<q', f.read(8))[0])
def buffers():
while (h := f.read(8)):
pb = pickle.PickleBuffer(bytearray(struct.unpack('<q', h)[0]))
f.readinto(pb)
yield pb
return DynamicTinygradUnpickler(io.BytesIO(opcodes), buffers=buffers()).load()
def flatten(value) -> 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('<q', len(data)))
tmp.write(data)
buffer.release()
return False
def _relocate(value, part_blob: memoryview, blob: Blob):
if isinstance(value, dict):
return {k: blob.add(part_blob[slice(*v)]) if k in ("lib", "init", "data") else _relocate(v, part_blob, blob) for k, v in value.items()}
return [_relocate(v, part_blob, blob) for v in value] if isinstance(value, list) else value
stream = io.BytesIO()
pickle.Pickler(stream, protocol=5, buffer_callback=buffer_cb).dump(obj)
opcodes = stream.getvalue()
f.write(struct.pack('<q', len(opcodes)))
f.write(opcodes)
tmp.seek(0)
shutil.copyfileobj(tmp, f)
def write_bundle(path: str, parts: dict, extras: dict | None = None):
blob = Blob()
headers = {key: _relocate(part_header, memoryview(part_blob), blob) for key, (part_header, part_blob) in parts.items()}
extras_location = blob.add(pickle.dumps(extras or {}, protocol=4))
entries = [{"key": list(key) if isinstance(key, tuple) else key, "blob": [0, blob.size], "header": header} for key, header in headers.items()]
header = json.dumps({"parts": entries, "extras": extras_location}).encode()
with open(path, "wb") as file:
file.write(MAGIC + struct.pack("<Q", len(header)) + header + bytes(-(len(MAGIC) + 8 + len(header)) % ALIGNMENT))
file.writelines(blob.chunks)
def read_parts(path: str) -> 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("<Q", data, len(MAGIC))[0]
header, blob = json.loads(bytes(data[len(MAGIC) + 8:start])), data[start + -start % ALIGNMENT:]
parts = {tuple(e["key"]) if isinstance(e["key"], list) else e["key"]: (e["header"], blob[slice(*e["blob"])]) for e in header["parts"]}
return parts, _MetadataUnpickler(io.BytesIO(bytes(blob[slice(*header["extras"])]))).load()
def load_model(path: str, device_map: dict | None = None, verify: bool = True, cache: bool = True) -> 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("<Q", file.read(8))[0]) == _cache_key(path, device_map)
def cache_is_fresh(path: str, device_map: dict | None = None) -> 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("<Q", len(key)) + key)
_Pickler(file, protocol=5).dump({name: _Jit(jit.captured, jit.input_names, jit.buffers) for name, jit in jits.items()})
os.replace(tmp_path, f"{path}.cache")
# Everything below this comment block calls tinygrad internals, so after a tinygrad bump this
# is the only code to update, and models are not recompiled. How to do this when discountchubbs is not available:
# 1. run modeld_v2/tests/test_plain_artifact.py on the new tinygrad and fix each failing call against the new API: Buffer,
# UOp.from_buffer/param (ParamArg), ProgramInfo, Target, KernelInfo, compile_linear, CapturedJit and
# _prepare_jit_inputs in tinygrad/engine/jit.py
# 2. load a real model with verify=True: ArtifactMismatch means the rebuilt graph still differs from the one that was compiled
# The boot cache is keyed on the tinygrad source, so it rebuilds itself. Only a change in how the device runs kernel binaries needs a recompile.
# Honestly, I would just have your clanker try it because you likely have zero idea what you are looking at lol..
def _dtype(text: str):
from tinygrad.dtype import dtypes
return eval(text, {"__builtins__": {}}, {"dtypes": dtypes})
def _device(recorded: str, device_map: dict | None) -> 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")
@@ -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:
+5 -6
View File
@@ -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)
@@ -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"")
@@ -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
@@ -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)))
@@ -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
@@ -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
@@ -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)
+42
View File
@@ -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()
+2 -30
View File
@@ -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.")
+6 -16
View File
@@ -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
+8 -84
View File
@@ -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)
@@ -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)
@@ -170,6 +170,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)),
+6 -43
View File
@@ -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)
-2
View File
@@ -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(