Compare commits

..

8 Commits

Author SHA1 Message Date
github-actions[bot] 8cb737c241 modeld_v2: spatial features (PR-1934) 2026-08-23 08:12:05 +00:00
Jason Wen 211f990f6b models: fix sunnylink default model display and false big model re-downloading (#1941)
* big needs small

* no download

* actually

* send it
2026-08-23 04:04:46 -04:00
Jason Wen 97468e4fa4 [TIZI/TICI] ui: remove calibration reset dialog on model change (#1942) 2026-08-23 03:48:52 -04:00
Jason Wen 6c6fba9a14 ci: fix flaky LLK test (#1940) 2026-08-23 02:57:02 -04:00
Jason Wen 34621cf816 ci: refactor big model chunk handling (#1939) 2026-08-23 02:48:10 -04:00
Jason Wen 086530b7c6 [TIZI/TICI] ui: fix path width during gas and steering override (#1938) 2026-08-22 21:47:38 -04:00
Jason Wen 4f46433e2b alerts: add branch metadata to chestnut offroad warning (#1936)
* alerts: add branch metadata to chestnut offroad warning

* all branches
2026-08-22 10:20:06 -04:00
Jason Wen 5a8567e3e7 ci: chestnut prebuilt branches (#1935)
* ci: chestnut prebuilt branches

* fix

* nope

* big

* try again

* diff

* malformed

* auth

* more
2026-08-22 03:41:48 -04:00
12 changed files with 271 additions and 60 deletions
+29 -24
View File
@@ -8,17 +8,35 @@ env:
HF_DEFAULTS_PATH: models/defaults/big
jobs:
resolve_name:
runs-on: ubuntu-24.04
outputs:
model_name: ${{ steps.name.outputs.model_name }}
onnx_ref: ${{ steps.name.outputs.onnx_ref }}
steps:
- uses: actions/checkout@v4
- id: name
run: |
NAME=$(PYTHONPATH=${{ github.workspace }} python3 -c "from openpilot.sunnypilot.models.model_name import DEFAULT_BIG_MODEL; print(DEFAULT_BIG_MODEL)")
ONNX_REF=$(git log -1 --format='%H' -- openpilot/selfdrive/modeld/models/big_driving_supercombo.onnx)
echo "model_name=${NAME}" >> $GITHUB_OUTPUT
echo "onnx_ref=$ONNX_REF" >> $GITHUB_OUTPUT
build_model:
needs: resolve_name
uses: ./.github/workflows/sunnypilot-build-model.yaml
with:
upstream_branch: ${{ github.sha }}
custom_name: default-big-model
upstream_branch: ${{ needs.resolve_name.outputs.onnx_ref }}
custom_name: ${{ needs.resolve_name.outputs.model_name }}
target_hardware: usbgpu
secrets: inherit
upload_defaults:
needs: build_model
needs: [ resolve_name, build_model ]
runs-on: ubuntu-24.04
permissions:
id-token: write
contents: write
steps:
- uses: actions/checkout@v4
with:
@@ -31,7 +49,7 @@ jobs:
- name: Download artifact name
uses: actions/download-artifact@v4
with:
name: artifact-name-default-big-model
name: artifact-name-${{ needs.resolve_name.outputs.model_name }}
path: artifact_name
- name: Read artifact name
@@ -46,33 +64,20 @@ jobs:
name: ${{ steps.artifact.outputs.artifact_name }}
path: output
- name: Upload model to HF defaults
- name: Upload to HF and update default_models.json
env:
HF_OIDC_RESOURCE: datasets/${{ env.HF_REPO }}
ARTIFACT_NAME: ${{ steps.artifact.outputs.artifact_name }}
run: |
rm -f output/artifact_name.txt
hf upload ${{ env.HF_REPO }} \
output/ \
"${HF_DEFAULTS_PATH}/${ARTIFACT_NAME}/" \
--repo-type=dataset
- name: Get tinygrad ref and ONNX hash
id: meta
run: |
export PYTHONPATH=$(pwd)
echo "tinygrad_ref=$(python3 openpilot/sunnypilot/models/tinygrad_ref.py)" >> $GITHUB_OUTPUT
echo "onnx_sha256=$(sha256sum openpilot/selfdrive/modeld/models/big_driving_supercombo.onnx | cut -d' ' -f1)" >> $GITHUB_OUTPUT
- name: Update default_models.json on HF
env:
HF_OIDC_RESOURCE: datasets/${{ env.HF_REPO }}
ARTIFACT_NAME: ${{ steps.artifact.outputs.artifact_name }}
run: |
python3 release/ci/upload_default_model.py \
--hf-repo "${{ env.HF_REPO }}" \
--hf-defaults-path "${{ env.HF_DEFAULTS_PATH }}" \
--artifact-name "$ARTIFACT_NAME" \
--metadata-path "output/metadata.json" \
--onnx-sha256 "${{ steps.meta.outputs.onnx_sha256 }}" \
--tinygrad-ref "${{ steps.meta.outputs.tinygrad_ref }}"
--model-dir output \
--onnx-path "openpilot/selfdrive/modeld/models/big_driving_supercombo.onnx" \
--onnx-ref "${{ needs.resolve_name.outputs.onnx_ref }}" \
--model-name "${{ needs.resolve_name.outputs.model_name }}" \
--tinygrad-ref "$(python3 openpilot/sunnypilot/models/tinygrad_ref.py)" \
--run-number "${{ github.run_number }}"
@@ -36,6 +36,7 @@ jobs:
publish_concurrency_group: ${{ steps.strategy.outputs.publish_concurrency_group }}
is_stable_branch: ${{ steps.strategy.outputs.is_stable_branch }}
build: ${{ steps.strategy.outputs.build }}
include_big_model: ${{ steps.strategy.outputs.include_big_model }}
steps:
- uses: actions/checkout@v4
- name: Extract deploy strategy
@@ -78,6 +79,9 @@ jobs:
stable_version=$(cat openpilot/sunnypilot/common/version.h | grep SUNNYPILOT_VERSION | sed -e 's/[^0-9|.]//g');
echo "version=$([ "$is_stable_branch" = "true" ] && echo "$stable_version" || echo "$BUILD")" >> $GITHUB_OUTPUT
echo "extra_version_identifier=${environment}" >> $GITHUB_OUTPUT
include_big_model="$(echo "$CONFIG" | jq -r '.include_big_model // false')";
echo "include_big_model=$include_big_model" >> $GITHUB_OUTPUT
fi
echo "build=$BUILD" >> $GITHUB_OUTPUT
cat $GITHUB_OUTPUT
@@ -203,6 +207,74 @@ jobs:
source ${UV_PROJECT_ENVIRONMENT}/bin/activate
PYTHONPATH=$PYTHONPATH:${{ github.workspace }}/ ${{ github.workspace }}/scripts/manage-powersave.py --enable
prepare_chestnut:
needs: [ prepare_strategy ]
runs-on: ubuntu-24.04
if: ${{ needs.prepare_strategy.outputs.include_big_model == 'true' }}
outputs:
onnx_sha256: ${{ steps.resolve.outputs.onnx_sha256 }}
env:
HF_REPO: sunnypilot/sunnypilot_models_v1
HF_DEFAULTS_PATH: models/defaults/big
steps:
- uses: actions/checkout@v4
with:
ref: ${{ github.head_ref || github.ref_name }}
- run: git lfs pull -I "openpilot/selfdrive/modeld/models/big_driving_supercombo.onnx"
- name: Check HF defaults and build if needed
id: resolve
run: |
ACTUAL_ONNX_HASH=$(sha256sum "openpilot/selfdrive/modeld/models/big_driving_supercombo.onnx" | cut -d' ' -f1)
echo "Repo ONNX hash: $ACTUAL_ONNX_HASH"
echo "onnx_sha256=$ACTUAL_ONNX_HASH" >> $GITHUB_OUTPUT
JSON_URL="https://huggingface.co/datasets/${HF_REPO}/resolve/main/${HF_DEFAULTS_PATH}/default_models.json"
check_hash() {
DEFAULTS=$(curl -fsSL "$JSON_URL" 2>/dev/null) || return 1
BUNDLE=$(echo "$DEFAULTS" | jq --arg hash "$ACTUAL_ONNX_HASH" '.bundles[] | select(.onnx_sha256 == $hash)' 2>/dev/null)
[ -n "$BUNDLE" ] && [ "$BUNDLE" != "null" ]
}
if check_hash; then
echo "HF defaults match repo ONNX"
else
echo "No matching model on HF — triggering build"
gh workflow run build-default-big-model.yaml --ref "${{ github.head_ref || github.ref_name }}"
echo "Waiting for build to start..."
sleep 120
RUN_ID=$(gh run list --workflow=build-default-big-model.yaml --branch="${{ github.head_ref || github.ref_name }}" --limit=1 --json databaseId --jq '.[0].databaseId')
if [ -z "$RUN_ID" ] || [ "$RUN_ID" = "null" ]; then
echo "::error::Failed to find build-default-big-model run"
exit 1
fi
echo "Waiting for run $RUN_ID..."
gh run watch "$RUN_ID"
CONCLUSION=$(gh run view "$RUN_ID" --json conclusion --jq '.conclusion')
if [ "$CONCLUSION" != "success" ]; then
echo "::error::build-default-big-model failed: $CONCLUSION"
exit 1
fi
if ! check_hash; then
echo "::error::HF defaults still don't match after build"
exit 1
fi
fi
env:
GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
- name: Cancel run on failure
if: failure()
run: gh run cancel ${{ github.run_id }}
env:
GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
publish:
concurrency:
@@ -211,14 +283,20 @@ jobs:
# Otherwise, if a job is waiting to be published due to environment wait time, it would be canceled by a new commit and restart the wait time.
group: ${{ needs.prepare_strategy.outputs.publish_concurrency_group }}
cancel-in-progress: ${{ needs.prepare_strategy.outputs.cancel_publish_in_progress == 'true' }}
if: ${{ (always() && !cancelled() && !failure()) && needs.build.result == 'success' && needs.prepare_strategy.result == 'success' && (!contains(github.event_name, 'pull_request') || (github.event.action == 'labeled' && github.event.label.name == 'prebuilt')) }}
needs: [ build, prepare_strategy ]
if: ${{
always() && !cancelled() &&
needs.build.result == 'success' &&
needs.prepare_strategy.result == 'success' &&
(!contains(github.event_name, 'pull_request') || (github.event.action == 'labeled' && github.event.label.name == 'prebuilt')) &&
(needs.prepare_strategy.outputs.include_big_model != 'true' || needs.prepare_chestnut.result == 'success')
}}
needs: [ build, prepare_strategy, prepare_chestnut ]
runs-on: ubuntu-24.04
environment: ${{ needs.prepare_strategy.outputs.environment }}
steps:
- uses: actions/checkout@v4
- name: Download build artifacts
- name: Download prebuilt artifact
uses: actions/download-artifact@v4
with:
name: prebuilt
@@ -228,6 +306,44 @@ jobs:
mkdir -p ${{ env.OUTPUT_DIR }}
tar xzf prebuilt.tar.gz -C ${{ env.OUTPUT_DIR }}
- name: Prepare chestnut output
if: ${{ needs.prepare_chestnut.result == 'success' }}
run: |
mkdir -p "${{ github.workspace }}/chestnut_output"
tar xzf prebuilt.tar.gz -C "${{ github.workspace }}/chestnut_output"
- name: Download big model chunks from HF
if: ${{ needs.prepare_chestnut.result == 'success' }}
env:
HF_REPO: sunnypilot/sunnypilot_models_v1
HF_DEFAULTS_PATH: models/defaults/big
run: |
ONNX_HASH="${{ needs.prepare_chestnut.outputs.onnx_sha256 }}"
JSON_URL="https://huggingface.co/datasets/${HF_REPO}/resolve/main/${HF_DEFAULTS_PATH}/default_models.json"
DEFAULTS=$(curl -fsSL "$JSON_URL")
BUNDLE=$(echo "$DEFAULTS" | jq --arg hash "$ONNX_HASH" '.bundles[] | select(.onnx_sha256 == $hash)')
mkdir -p big_model_chunks
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')
CANONICAL="big_driving_tinygrad.pkl"
echo "$ARTIFACT" | jq -r '.chunks[].file_name' | while read CHUNK_NAME; do
CHUNK_IDX=$(echo "$CHUNK_NAME" | grep -oP 'chunk\K[0-9]+of[0-9]+')
CANONICAL_CHUNK="${CANONICAL}.chunk${CHUNK_IDX}"
ENCODED_URL=$(python3 -c "import urllib.parse; print(urllib.parse.quote('${BASE_URL}/${CHUNK_NAME}', safe=':/'))")
echo "Downloading $CHUNK_NAME -> $CANONICAL_CHUNK"
curl -fsSL -o "big_model_chunks/${CANONICAL_CHUNK}" "$ENCODED_URL"
done
echo "$NUM_CHUNKS" > "big_model_chunks/${CANONICAL}.chunkmanifest"
- name: Inject big model into chestnut
if: ${{ needs.prepare_chestnut.result == 'success' }}
run: |
cp big_model_chunks/* "${{ github.workspace }}/chestnut_output/openpilot/selfdrive/modeld/models/"
- name: Configure Git
run: |
git config --global user.email "github-actions[bot]@users.noreply.github.com"
@@ -248,6 +364,22 @@ jobs:
"https://x-access-token:${{github.token}}@github.com/sunnypilot/sunnypilot.git" \
"${{ needs.prepare_strategy.outputs.extra_version_identifier }}"
- name: Publish chestnut branch
if: ${{ needs.prepare_chestnut.result == 'success' }}
env:
GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
run: |
CHESTNUT_BRANCH="${{ needs.prepare_strategy.outputs.new_branch }}-chestnut"
CHESTNUT_DIR="${{ github.workspace }}/chestnut_output"
${{ env.CI_DIR }}/publish.sh \
"${{ github.workspace }}" \
"$CHESTNUT_DIR" \
"$CHESTNUT_BRANCH" \
"${{ needs.prepare_strategy.outputs.version }}" \
"https://x-access-token:${{github.token}}@github.com/sunnypilot/sunnypilot.git" \
"${{ needs.prepare_strategy.outputs.extra_version_identifier }}"
- name: Tag ${{ needs.prepare_strategy.outputs.environment }}
if: ${{ needs.prepare_strategy.outputs.is_stable_branch == 'true' && (github.event_name != 'push' || !startsWith(github.ref, 'refs/tags/')) }}
run: |
@@ -260,6 +392,7 @@ jobs:
- prepare_strategy
- build
- publish
- prepare_chestnut
runs-on: ubuntu-24.04
if: ${{ (always() && !cancelled() && !failure())
&& needs.publish.result == 'success'
@@ -279,6 +412,7 @@ jobs:
export commit_short_sha="${commit_short_sha:0:7}"
export extra_version_identifier="${{ needs.prepare_strategy.outputs.extra_version_identifier || github.run_number }}"
export PUBLIC_REPO_URL="${{ env.PUBLIC_REPO_URL }}"
export chestnut_branch="${{ needs.prepare_chestnut.result == 'success' && format('{0}-chestnut', needs.prepare_strategy.outputs.new_branch) || '' }}"
MESSAGE=$(cat << 'EOF' | envsubst
${{ vars.DISCOURSE_GENERAL_UPDATE_NOTICE }}
+9
View File
@@ -16,6 +16,15 @@ MASTER_SP_BRANCHES = ['master']
RELEASE_BRANCHES = ['release-tizi-staging', 'release-mici-staging', 'release-tizi', 'release-mici', 'nightly']
TESTED_BRANCHES = RELEASE_BRANCHES + ['devel-staging', 'nightly-dev'] + RELEASE_SP_BRANCHES + TESTED_SP_BRANCHES
CHESTNUT_BRANCHES = {
"staging": "staging-chestnut",
"dev": "dev-chestnut",
"release-mici": "release-chestnut",
"release-tizi": "release-chestnut",
"release-mici-staging": "release-chestnut-staging",
"release-tizi-staging": "release-chestnut-staging",
}
SP_BRANCH_MIGRATIONS = {
("tici", "staging-c3-new"): "staging-tici",
("tici", "dev-c3-new"): "staging-tici",
@@ -18,7 +18,7 @@
"_comment": "Set extra field to the failed reason."
},
"Offroad_ChestnutBranch": {
"text": "Chestnut detected! Switch to the release-chestnut branch to use chestnut-class models.",
"text": "Chestnut detected! Switch to the %1 branch to use chestnut-class models.",
"severity": 0
},
"Offroad_UnregisteredHardware": {
@@ -178,27 +178,14 @@ class ModelsLayout(Widget):
# circled_slash is authored grey; tinting it again only darkens it
return {"name": name, "text_color": rl.GRAY, "icon": "icons/circled_slash.png", "icon_color": rl.WHITE}
@staticmethod
def _show_reset_params_dialog():
def _callback(response):
if response == DialogResult.CONFIRM:
ui_state.params.remove("CalibrationParams")
ui_state.params.remove("LiveTorqueParameters")
msg = tr("Model download has started in the background. We suggest resetting calibration. Would you like to do that now?")
dialog = ConfirmDialog(msg, tr("Reset Calibration"), callback=_callback)
gui_app.push_widget(dialog)
def _on_model_selected(self, result):
if result != DialogResult.CONFIRM:
return
selected_ref = self.model_dialog.selection_ref
if selected_ref == "Default":
ui_state.params.remove("ModelManager_ActiveBundle")
self._show_reset_params_dialog()
elif selected_bundle := next((bundle for bundle in self.model_manager.availableBundles if bundle.ref == selected_ref), None):
ui_state.params.put("ModelManager_DownloadIndex", selected_bundle.index)
if self.model_manager.activeBundle and selected_bundle.generation != self.model_manager.activeBundle.generation:
self._show_reset_params_dialog()
self.model_dialog = None
@staticmethod
@@ -8,6 +8,7 @@ from openpilot.common.filter_simple import FirstOrderFilter
from openpilot.selfdrive.ui.ui_state import ui_state, UIStatus
from openpilot.selfdrive.ui.sunnypilot.onroad.chevron_metrics import ChevronMetrics
from openpilot.selfdrive.ui.sunnypilot.onroad.rainbow_path import RainbowPath
from openpilot.selfdrive.ui.sunnypilot.ui_state import MADSState
from openpilot.system.ui.lib.application import gui_app
@@ -19,6 +20,11 @@ class ModelRendererSP:
@property
def _lateral_active(self) -> bool:
sm = ui_state.sm
if sm.valid["selfdriveStateSP"]:
mads = sm["selfdriveStateSP"].mads
if mads.available:
return mads.enabled and mads.state != MADSState.paused
return ui_state.status in (UIStatus.ENGAGED, UIStatus.LAT_ONLY)
def _get_path_half_width(self) -> float:
+12 -3
View File
@@ -143,13 +143,17 @@ class ModelManagerSP:
is_cached = False
if len(artifact.chunks) > 0:
from openpilot.common.file_chunker import get_chunk_name
num_chunks = len(artifact.chunks)
chunks_valid = True
for i, chunk in enumerate(artifact.chunks):
chunk_path = get_chunk_name(full_path, i, len(artifact.chunks))
chunk_path = get_chunk_name(full_path, i, num_chunks)
if not await verify_file(chunk_path, chunk.sha256):
chunks_valid = False
break
if chunks_valid and len(artifact.chunks) > 0:
artifact.downloadProgress.progress = ((i + 1) / num_chunks) * 100
self._sync_artifact_progress(artifact)
self._report_status()
if chunks_valid and num_chunks > 0:
is_cached = True
else:
if await verify_file(full_path, expected_hash):
@@ -216,6 +220,9 @@ class ModelManagerSP:
"""Downloads all models in a bundle"""
self.selected_bundle = model_bundle
self.selected_bundle.status = custom.ModelManagerSP.DownloadStatus.downloading
for model in self.selected_bundle.models:
model.artifact.downloadProgress.status = custom.ModelManagerSP.DownloadStatus.downloading
self._report_status()
os.makedirs(destination_path, exist_ok=True)
try:
@@ -260,7 +267,9 @@ class ModelManagerSP:
self.active_bundle = get_active_bundle(self.params)
if (index_to_download := self.params.get("ModelManager_DownloadIndex")) is not None:
if model_to_download := next((model for model in self.available_models if model.index == index_to_download), None):
if self.active_bundle and self.active_bundle.index == index_to_download:
self.params.remove("ModelManager_DownloadIndex")
elif model_to_download := next((model for model in self.available_models if model.index == index_to_download), None):
try:
self.download(model_to_download, Paths.model_root())
except Exception as e:
@@ -83,9 +83,14 @@ class TestLocationdProc(OpenpilotTestCase):
self.pm.send(msg.which(), msg)
if msg.which() == "cameraOdometry":
self.pm.wait_for_readers_to_update(msg.which(), timeout=1, dt=0.005)
time.sleep(1) # wait for async params write
for _ in range(50):
val = self.params.get('LastGPSPositionLLK')
if val is not None:
break
time.sleep(0.1)
lastGPS = json.loads(self.params.get('LastGPSPositionLLK'))
self.assertIsNotNone(val, "LastGPSPositionLLK not written within 5s")
lastGPS = json.loads(val)
self.assertAlmostEqual(lastGPS['latitude'], self.lat, delta=0.001)
self.assertAlmostEqual(lastGPS['longitude'], self.lon, delta=0.001)
self.assertAlmostEqual(lastGPS['altitude'], self.alt, delta=0.001)
@@ -28,7 +28,8 @@ from websocket import (ABNF, WebSocket, WebSocketException, WebSocketTimeoutExce
create_connection, WebSocketConnectionClosedException)
import openpilot.cereal.messaging as messaging
from openpilot.sunnypilot.models.default_model import get_default_model
from openpilot.selfdrive.modeld.helpers import usbgpu_present, usbgpu_compiled
from openpilot.sunnypilot.models.model_name import DEFAULT_MODEL, DEFAULT_BIG_MODEL
from openpilot.sunnypilot.selfdrive.car.sync_sunnylink_params import update_car_list_param
from openpilot.sunnypilot.sunnylink.api import SunnylinkApi
from openpilot.sunnypilot.sunnylink.utils import sunnylink_need_register, sunnylink_ready, get_param_as_byte, save_param_from_base64_encoded_string
@@ -181,7 +182,10 @@ def getParamsMetadata() -> str:
schema = generate_schema()
schema["capabilities"] = generate_capabilities()
schema["capability_labels"] = CAPABILITY_LABELS
schema["default_model"] = get_default_model()
# mirrors get_default_model() — ui_state unavailable in sunnylinkd process
show_big = (usbgpu_present() and usbgpu_compiled()
and (params.get_bool("UsbGpuActive") or params.get_bool("UsbGpuLoading") or params.get_bool("IsOffroad")))
schema["default_model"] = DEFAULT_BIG_MODEL if show_big else DEFAULT_MODEL
raw = json.dumps(schema, separators=(",", ":")).encode("utf-8")
return base64.b64encode(gzip.compress(raw)).decode("utf-8")
except Exception:
+6 -2
View File
@@ -27,7 +27,7 @@ from openpilot.common.swaglog import cloudlog
from openpilot.sunnypilot.system.statsd import statlog
from openpilot.system.hardware.power_monitoring import PowerMonitoring
from openpilot.system.hardware.fan_controller import FanController
from openpilot.common.version import terms_version, training_version, get_build_metadata, terms_version_sp
from openpilot.common.version import terms_version, training_version, get_build_metadata, terms_version_sp, CHESTNUT_BRANCHES
ThermalStatus = log.DeviceState.ThermalStatus
@@ -301,7 +301,11 @@ def hardware_thread(end_event, hw_queue) -> None:
set_usb_state(msg.deviceState, last_hw_state.usb_state)
chestnut.update(started_ts is None, last_hw_state.usb_state)
set_offroad_alert_if_changed("Offroad_ChestnutBranch", msg.deviceState.chestnutPresent and not big_model_available)
current_channel = get_build_metadata().channel
chestnut_target = CHESTNUT_BRANCHES.get(current_channel)
chestnut_needs_switch = msg.deviceState.chestnutPresent and not big_model_available and chestnut_target is not None
set_offroad_alert_if_changed("Offroad_ChestnutBranch", chestnut_needs_switch,
extra_text=chestnut_target if chestnut_needs_switch else None)
# this subset is only used for offroad
temp_sources = [
+20 -2
View File
@@ -90,6 +90,18 @@ def _rename_pkl_with_chunks(old_pkl: Path, new_pkl: Path) -> Path:
return old_pkl.rename(new_pkl)
def _hash_onnx_files(model_dir: Path) -> str | None:
onnx_files = sorted(model_dir.glob("*.onnx"))
if not onnx_files:
return None
digest = hashlib.sha256()
for f in onnx_files:
with f.open('rb') as fh:
while block := fh.read(1024 * 1024):
digest.update(block)
return digest.hexdigest()
def generate_chunked_model(driving_pkl: Path) -> dict:
tinygrad_hash = _hash_pkl(driving_pkl)
@@ -123,7 +135,8 @@ def generate_chunked_model(driving_pkl: Path) -> dict:
}
def create_metadata_json(models: list, output_dir: Path, custom_name=None, short_name=None, is_20hz=False, upstream_branch="unknown") -> None:
def create_metadata_json(models: list, output_dir: Path, custom_name=None, short_name=None, is_20hz=False, upstream_branch="unknown",
onnx_sha256=None) -> None:
bundle_json = {
"short_name": short_name,
"display_name": custom_name or upstream_branch,
@@ -139,6 +152,9 @@ def create_metadata_json(models: list, output_dir: Path, custom_name=None, short
"models": models,
}
if onnx_sha256:
bundle_json["onnx_sha256"] = onnx_sha256
# Write metadata to output_dir
metadata_json = {
"bundles": [bundle_json]
@@ -178,4 +194,6 @@ if __name__ == "__main__":
_driving_pkl = new_pkl
_model_metadata = generate_chunked_model(_driving_pkl)
create_metadata_json([_model_metadata], _output_dir, args.custom_name, _short_name, args.is_20hz, args.upstream_branch)
_onnx_sha256 = _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)
+38 -8
View File
@@ -7,35 +7,66 @@ See the LICENSE.md file in the root directory for more details.
"""
import argparse
import hashlib
import json
import sys
import tempfile
from huggingface_hub import HfApi, hf_hub_download
def hash_file(path: str) -> str:
digest = hashlib.sha256()
with open(path, 'rb') as f:
while block := f.read(1024 * 1024):
digest.update(block)
return digest.hexdigest()
def main():
parser = argparse.ArgumentParser()
parser.add_argument("--hf-repo", required=True)
parser.add_argument("--hf-defaults-path", required=True)
parser.add_argument("--artifact-name", required=True)
parser.add_argument("--metadata-path", required=True)
parser.add_argument("--onnx-sha256", required=True)
parser.add_argument("--model-dir", required=True)
parser.add_argument("--onnx-path", required=True)
parser.add_argument("--onnx-ref", required=True)
parser.add_argument("--model-name", required=True)
parser.add_argument("--tinygrad-ref", required=True)
parser.add_argument("--run-number", required=True)
args = parser.parse_args()
with open(args.metadata_path) as f:
api = HfApi()
onnx_sha256 = hash_file(args.onnx_path)
short_ref = args.onnx_ref[:8]
folder_name = f"model-{args.model_name}-{short_ref}-{args.run_number}"
print(f"ONNX hash: {onnx_sha256}")
print(f"ONNX ref: {args.onnx_ref} (short: {short_ref})")
print(f"Folder: {folder_name}")
metadata_path = f"{args.model_dir}/metadata.json"
with open(metadata_path) as f:
metadata = json.load(f)
bundle = metadata['bundles'][0]
bundle['onnx_sha256'] = args.onnx_sha256
bundle['display_name'] = args.model_name
bundle['onnx_sha256'] = onnx_sha256
bundle['onnx_ref'] = args.onnx_ref
artifact = bundle['models'][0]['artifact']
hf_base = f"https://huggingface.co/datasets/{args.hf_repo}/resolve/main/{args.hf_defaults_path}/{args.artifact_name}"
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(
folder_path=args.model_dir,
path_in_repo=f"{args.hf_defaults_path}/{folder_name}",
repo_id=args.hf_repo,
repo_type="dataset",
)
json_filename = f"{args.hf_defaults_path}/default_models.json"
try:
local_path = hf_hub_download(repo_id=args.hf_repo, repo_type='dataset', filename=json_filename)
@@ -47,7 +78,7 @@ def main():
defaults_json['tinygrad_ref'] = args.tinygrad_ref
existing_idx = next((i for i, b in enumerate(defaults_json['bundles'])
if b.get('display_name') == bundle.get('display_name')), None)
if b.get('onnx_sha256') == onnx_sha256), None)
if existing_idx is not None:
defaults_json['bundles'][existing_idx] = bundle
else:
@@ -55,7 +86,6 @@ def main():
print(json.dumps(defaults_json, indent=2))
api = HfApi()
with tempfile.NamedTemporaryFile(mode='w', suffix='.json', delete=False) as f:
json.dump(defaults_json, f, indent=2)
tmp_path = f.name