Compare commits

..

8 Commits

Author SHA1 Message Date
discountchubbs 0a4540dbe8 set var 2026-10-21 13:46:20 -07:00
discountchubbs 6f1f132ba2 set env 2026-10-21 13:18:40 -07:00
discountchubbs bba6db5dbf ew 2026-10-21 12:51:52 -07:00
discountchubbs 2cf321b9ce rm container 2026-10-21 12:44:48 -07:00
discountchubbs f22ed86eed models: self healing test for tg 2026-10-21 12:36:26 -07:00
James Vecellio-Grant 1aa3a430fb models: pin tg to @02e6981 (#2071) 2026-10-21 11:11:24 -07:00
Amy Jeanes ddd7b14ebc models: download whole-file artifacts with parallel byte ranges (#2051)
* models: download whole-file artifacts with parallel byte ranges

HuggingFace's Xet CAS throttles each connection to roughly 1-2 MB/s but does
not limit connection count, so the single-stream _download_file left the link
idle most of the time. Replace it with a downloader that probes the size with
a one-byte Range request, creates the file at its final size and streams each
8 MB piece straight to its offset from one of 12 worker threads. The file is
never copied or assembled and only ever occupies its own size on disk. A
server that answers the probe with 200 gets one plain stream instead.
Every request sends Accept-Encoding: identity, because byte ranges count
encoded bytes while requests hands back decoded ones, so a compressed
response would land at the wrong offsets and fail as a short read.

Workers only write bytes; a reporter coroutine on the event loop publishes
progress every 0.5 s (instead of every 128 KB), re-raises the first worker
failure and polls the cancel ref. Transient errors (connection, timeout,
truncated body, HTTP 408/429/500/502/503/504) give a piece up to three
attempts; any other HTTP error fails the download at once. The whole-file
path now also publishes "verifying" while hashing before and after the
transfer, and the download is hashed against the manifest exactly as before.

A cancel has to wait for in-flight pieces to unwind (a blocked socket read
cannot be interrupted), so the row shows "canceling..." through a new
DownloadProgress.canceling flag instead of a frozen bar, and the mici cancel
button hides once the ref is gone, as it already does on the big UI.

DownloadProgress gains speed, canceling and resuming, appended to the capnp so
older readers are unaffected. All three are declared here so the wire format
lands in one place; speed and resuming are filled in by the follow-up speed
and resume changes.

Builds on #2050, which drops chunking and publishes models as whole files;
this makes that whole-file download parallel.

Measured on a comma 3X on WiFi: ~2 MB/s single stream, 17-24 MB/s steady and
30-60 MB/s peak with 12 range connections, the count picked from an on-device
sweep. Disk is not the bottleneck (eMMC writes at 56-67 MB/s with fdatasync).

Tests drive the real downloader against a local threaded range server: piece
tiling, actual parallelism (a stalled piece does not block the others), the
no-range fallback, per-piece retry, 4xx, cancel via ModelManager_DownloadRef
and mid-piece, verify/cached/downloaded state, a probe failure moving
past verifying, and hash mismatch cleanup.

Also drops two chunk-era leftovers in the tests: a stale "no chunks" note on
manifest_bundle and a one-item URL loop in the live manifest check.

Assisted-by: Claude Code (Claude Opus 4.8)
Assisted-by: Claude Code (Claude Opus 5.5)

* remove unneeded comments

* models: trim the manager tests to the regressions they catch

Each remaining test is there because a deliberate break in manager.py,
helpers.py or fetcher.py fails it. Tests that only re-checked what another
already catches are merged into a happy-path download or a table, and
checks of constants rather than behaviour are dropped. The fake server
now starts per test, keeps its state in one dataclass, and can cut a
response short.

Against 66 deliberate breaks, every one the old tests caught still fails,
and a 404 being retried, a short read and its retry, and the display
flags choosing the big model without a GPU are now caught too.

Assisted-by: Claude Code (Claude Opus 5.5)

* Update test_manager_download.py

* Update manager.py

---------

Co-authored-by: James Vecellio-Grant <159560811+Discountchubbs@users.noreply.github.com>
2026-10-21 09:35:43 -07:00
James Vecellio-Grant 922513f312 ci: lfs pointers survive gc (#2066)
* ci: lfs pointers survive gc

* comma got gc'ed

* format
2026-10-18 22:41:23 -07:00
19 changed files with 896 additions and 647 deletions
+10 -9
View File
@@ -33,7 +33,7 @@ jobs:
- name: LFS Fetch
run: |
git lfs fetch
git lfs fetch || echo "::warning::some objects are missing on commaai/openpilot-lfs; mirroring the rest"
- name: Set up Git
run: |
@@ -68,6 +68,7 @@ jobs:
git config lfs.activitytimeout 300
git config lfs.transfer.maxretries 5
git config lfs.concurrenttransfers 4
git config lfs.allowincompletepush true
- name: Push LFS
id: sync-and-commit
@@ -108,19 +109,19 @@ jobs:
p = subprocess.Popen(["git", "-C", repo, "cat-file", "--batch=%(objectsize)"],
stdin=subprocess.PIPE, stdout=subprocess.PIPE)
def feed():
for sha in small:
p.stdin.write(sha.encode() + b"\n")
p.stdin.close()
for sha in small:
p.stdin.write(sha.encode() + b"\n")
p.stdin.close()
threading.Thread(target=feed, daemon=True).start()
rx = re.compile(rb"^version https://git-lfs\.github\.com/spec/v1\noid sha256:([0-9a-f]{64})\nsize ([0-9]+)\n")
seen = {}
for _ in small:
n = int(p.stdout.readline()); body = p.stdout.read(n); p.stdout.read(1)
m = rx.match(body)
if m:
seen.setdefault(m.group(1).decode(), m.group(2).decode())
n = int(p.stdout.readline()); body = p.stdout.read(n); p.stdout.read(1)
m = rx.match(body)
if m:
seen.setdefault(m.group(1).decode(), m.group(2).decode())
for oid, size in sorted(seen.items()):
print(oid, size)
print(oid, size)
EOF
for repo in . sunnypilot; do
python3 scan_pointers.py "$repo"
+53
View File
@@ -3,9 +3,62 @@ on:
pull_request:
paths:
- 'tinygrad_repo'
- 'openpilot/sunnypilot/modeld_v2/**'
workflow_dispatch:
env:
CI: 1
GIT_CONFIG_COUNT: 1
GIT_CONFIG_KEY_0: lfs.fetchexclude
GIT_CONFIG_VALUE_0: openpilot/selfdrive/modeld/models/big_driving_tinygrad.pkl
jobs:
plain-artifact:
name: Plain artifact round trip and pickle converter
runs-on: ubuntu-24.04
env:
PYTHONPATH: ${{ github.workspace }}/tinygrad_repo:${{ github.workspace }}
TINYGRAD_TESTS: 1
steps:
- uses: actions/checkout@v7
- run: ./tools/op.sh setup
- name: Build openpilot
run: scons -j$(nproc)
- name: Compile, load, cache and fuse a model with the pinned tinygrad
env:
DEV: CPU
run: tools/op.sh test -v openpilot/sunnypilot/modeld_v2/tests/test_plain_artifact.py
- name: Install LLVM 20 and ROCm comgr for the amd mock, as tinygrad_repo/.github/actions/setup-tinygrad does
run: |
. /etc/os-release
curl -fsSL https://apt.llvm.org/llvm-snapshot.gpg.key | sudo tee /etc/apt/trusted.gpg.d/apt.llvm.org.asc > /dev/null
echo "deb http://apt.llvm.org/$VERSION_CODENAME/ llvm-toolchain-$VERSION_CODENAME-20 main" | sudo tee /etc/apt/sources.list.d/llvm.list
sudo mkdir -p /etc/apt/keyrings
curl -fsSL https://repo.radeon.com/rocm/rocm.gpg.key | gpg --dearmor | sudo tee /etc/apt/keyrings/rocm.gpg > /dev/null
echo "deb [arch=amd64 signed-by=/etc/apt/keyrings/rocm.gpg] https://repo.radeon.com/rocm/apt/7.1 $VERSION_CODENAME main" | sudo tee /etc/apt/sources.list.d/rocm.list
printf 'Package: *\nPin: release o=repo.radeon.com\nPin-Priority: 600\n' | sudo tee /etc/apt/preferences.d/rocm-pin-600
sudo apt-get -qq update && sudo apt-get -qq install -y --no-install-recommends libllvm20 clang-20 lld-20 comgr
printf '/opt/rocm/lib\n/opt/rocm/lib64\n' | sudo tee /etc/ld.so.conf.d/rocm.conf && sudo ldconfig
echo "/usr/lib/llvm-20/bin" >> "$GITHUB_PATH"
- name: Convert a pickled USB+AMD jit and replay it under tinygrad's amd mock
env:
DEV: MOCKUSB+AMD
CHECK_OOB: 0
GMMU: 0
run: tools/op.sh test -v openpilot/sunnypilot/modeld_v2/tests/test_convert_upstream_pkl.py
- name: Where tinygrad moved, for the helper or failure point named above
if: failure() && github.event_name == 'pull_request'
run: |
git fetch -q --depth=1 origin "${{ github.base_ref }}"
old=$(git rev-parse "origin/${{ github.base_ref }}:tinygrad_repo") new=$(git rev-parse HEAD:tinygrad_repo)
git -C tinygrad_repo fetch -q origin "$old" || true
echo "tinygrad_repo $old -> $new, in what helpers.py and plain_artifact_compiler.py read:"
git -C tinygrad_repo diff "$old" "$new" -- tinygrad/uop/ops.py tinygrad/engine/jit.py tinygrad/engine/realize.py tinygrad/device.py \
tinygrad/helpers.py tinygrad/dtype.py tinygrad/renderer/__init__.py
echo "and in what convert_upstream_pkl.py decodes:"
git -C tinygrad_repo diff --stat "$old" "$new" -- tinygrad/runtime/autogen/am/pm4_nv.py tinygrad/runtime/autogen/amd_gpu.py
git -C tinygrad_repo diff "$old" "$new" -- tinygrad/runtime/support/hcq2.py tinygrad/runtime/support/usb.py tinygrad/runtime/ops_amd.py
generate-matrix:
runs-on: ubuntu-latest
outputs:
+3
View File
@@ -138,6 +138,9 @@ struct ModelManagerSP @0xaedffd8f31e7b55d {
status @0 :DownloadStatus;
progress @1 :Float32;
eta @2 :UInt32;
speed @3 :Float32;
canceling @4 :Bool;
resuming @5 :Bool;
}
struct Chunk {
-1
View File
@@ -9,7 +9,6 @@ os.environ['GMMU'] = '0' # for chestnut fast loading, noop for qcom
os.environ.setdefault('AM_POWER_LIMIT', '100')
os.environ.setdefault('AM_PERF_LEVEL', 'high')
from tinygrad.device import Buffer, Device
from tinygrad.engine.realize import lower_and_compile
from tinygrad.dtype import DType, dtypes
from tinygrad.engine.realize import lower_and_compile
from tinygrad.tensor import Tensor
-34
View File
@@ -1,7 +1,4 @@
import time
from typing import Any
import numpy as np
def latest_gps_source(sm) -> str | None:
@@ -10,34 +7,3 @@ def latest_gps_source(sm) -> str | None:
sources = [s for s in ('gpsLocation', 'gpsLocationExternal')
if sm.valid[s] and sm[s].hasFix and 0 <= now - sm.recv_time[s] < 2]
return max(sources, key=lambda s: sm.logMonoTime[s]) if sources else None
# convert mapbox steps to valhalla enums the model was trained on
VALHALLA_COMMAND_COUNT = 46
VALHALLA_SIDED_TYPES = { # (straight, right, left)
'depart': (1, 2, 3), 'arrive': (4, 5, 6),
'on ramp': (17, 18, 19), 'off ramp': (20, 20, 21), 'fork': (22, 23, 24),
'roundabout': (26, 26, 26), 'rotary': (26, 26, 26), 'roundabout turn': (26, 26, 26),
'exit roundabout': (27, 27, 27), 'exit rotary': (27, 27, 27), 'notification': (28, 28, 28),
}
VALHALLA_TURN_TYPES = {'straight': 8, 'slight right': 9, 'right': 10, 'sharp right': 11, 'uturn': 13,
'sharp left': 14, 'left': 15, 'slight left': 16} # turn, continue, new name, end of road
def valhalla_type(maneuver_type: str, modifier: str) -> int | None:
if maneuver_type == 'merge':
return None # Valhalla gives no maneuver at a merge, so the model never saw one
if maneuver_type in VALHALLA_SIDED_TYPES:
straight, right, left = VALHALLA_SIDED_TYPES[maneuver_type]
return right if 'right' in modifier else left if 'left' in modifier else straight
return VALHALLA_TURN_TYPES.get(modifier, 8)
def navigation_input(maneuvers: Any, traveled: float) -> np.ndarray:
navigation = np.zeros(VALHALLA_COMMAND_COUNT + 1, dtype=np.float32)
for maneuver in maneuvers:
command, distance = valhalla_type(maneuver.type, maneuver.modifier), maneuver.distance - traveled
if command is not None and command > 6 and distance >= 0:
navigation[command], navigation[-1] = 1, distance / 100
break
return navigation
@@ -235,6 +235,8 @@ class ModelsLayout(Widget):
progress = sum(p.progress for p in progresses) / len(progresses)
ds = custom.ModelManagerSP.DownloadStatus
if any(p.canceling for p in progresses):
return {"name": name, "downloading": True, "progress": progress, "status_text": tr("canceling...")}
if ds.failed in statuses:
# close.png is authored black and a tint cannot lift it, hence close2
return {"name": name, "status_text": tr("download failed"), "text_color": rl.RED, "icon": "icons/close2.png"}
@@ -251,11 +251,12 @@ class ModelsLayoutMici(NavScroller):
self.current_model_info.info_text.set_text(tr("download failed"))
elif manager.selectedBundle and manager.selectedBundle.status == custom.ModelManagerSP.DownloadStatus.downloading:
self.cancel_download_btn.set_visible(True)
self.cancel_download_btn.set_visible(ui_state.params.get("ModelManager_DownloadRef") is not None)
device.set_override_interactive_timeout(5)
progress = 0.0
count = 0
verifying = False
canceling = False
for model in manager.selectedBundle.models:
count += 1
p = model.artifact.downloadProgress
@@ -263,11 +264,13 @@ class ModelsLayoutMici(NavScroller):
custom.ModelManagerSP.DownloadStatus.verifying):
progress += p.progress
verifying = verifying or p.status == custom.ModelManagerSP.DownloadStatus.verifying
canceling = canceling or p.canceling
elif p.status in (custom.ModelManagerSP.DownloadStatus.downloaded,
custom.ModelManagerSP.DownloadStatus.cached):
progress += 100.0
self.current_model_info.current_model_header.set_text(tr("verifying") if verifying else tr("downloading"))
header = tr("canceling...") if canceling else tr("verifying") if verifying else tr("downloading")
self.current_model_info.current_model_header.set_text(header)
self.cancel_download_btn.set_text(tr("cancel verification") if verifying else tr("cancel download"))
self.current_model_info.current_model_header._shimmer = True
name_text = manager.selectedBundle.internalName.lower()
@@ -7,6 +7,7 @@ See the LICENSE.md file in the root directory for more details.
"""
import argparse
import dataclasses
import inspect
import itertools
import mmap
import re
@@ -45,6 +46,11 @@ def _address(v):
return _root(v.src[0]), add + start, size
def _buffer(u):
from tinygrad.uop.ops import Ops
return u.buffer if u.op is Ops.BUFFER else getattr(u.arg, "buffer", None)
def _patched(buf, topo) -> tuple[bytearray, dict]:
from tinygrad.uop.ops import Ops
blob, rows = bytearray(buf.max_numel() * buf.dtype.itemsize), {}
@@ -72,7 +78,7 @@ def _evaluate(u, env):
table, index = _root(u.src[0].src[0]), _evaluate(u.src[0].src[1], env)
return struct.unpack_from(u.dtype.fmt, table.arg, index * u.dtype.itemsize)[0] if table.op is Ops.BINARY and index is not None else None
if u.op in GroupOp.ALU and None not in (values := [_evaluate(x, env) for x in u.src]):
return exec_alu(u.op, u.dtype, values, False)
return exec_alu(u.op, u.dtype, values, *([False] if "truncate_output" in inspect.signature(exec_alu).parameters else []))
return None
@@ -112,7 +118,7 @@ def convert(jit, kwargs: dict, frame: dict, target) -> tuple[dict, memoryview]:
from tinygrad.uop.ops import Ops
submit = jit.captured._linear.src[0].without_after
kernels, topo = submit.arg.aux.kernels, list(jit.captured._linear.toposort(enter_calls=True))
allocs = {str(u.tag): u for u in topo if u.op is Ops.ALLOC and u.tag is not None}
allocs = {str(name): u for u in topo if u.op is Ops.ALLOC and (name := getattr(u, "tag", None) or u.arg.name) is not None}
(commands, patched), kernargs = _patched(allocs["cmdbuf_compute_0"], topo), allocs["kernargs_compute_0"]
argument_words, pokes = _patched(kernargs, topo)[1], _input_pokes(submit, kernargs)
sh = {name: getattr(amd_gpu, f"reg{name}") + amd_gpu.GC_BASE__INST0_SEG0 - pm4.PACKET3_SET_SH_REG_START
@@ -158,7 +164,7 @@ def convert(jit, kwargs: dict, frame: dict, target) -> tuple[dict, memoryview]:
args.append({"input": pokes[slot][0]})
elif slot in argument_words:
base, at, nbytes = _address(argument_words[slot])
args.append({"buffer": buffers.setdefault(base, len(buffers)), "offset": at, "nbytes": nbytes or base.arg.buffer.nbytes, "dtype": u8})
args.append({"buffer": buffers.setdefault(base, len(buffers)), "offset": at, "nbytes": nbytes or base.max_numel() * base.dtype.itemsize, "dtype": u8})
outs, ins = kernel[6]
calls.append({"kind": "program", "name": re.sub(r"\x1b\[[0-9;]*m", "", kernel[1]), "lib": programs[image_alloc], "target": dataclasses.asdict(target),
"signature": [[f"data{i}", i, u8, [a.get("nbytes", 1)]] for i, a in enumerate(args)], "global_size": grid, "local_size": local,
@@ -166,12 +172,13 @@ def convert(jit, kwargs: dict, frame: dict, target) -> tuple[dict, memoryview]:
needs_init |= {args[i]["buffer"] for i in ins if "buffer" in args[i] and args[i]["buffer"] not in written}
written |= {args[i]["buffer"] for i in outs if "buffer" in args[i]}
assert all(_buffer(b) is not None for b, i in buffers.items() if i in needs_init), "a kernel reads a scratch before any kernel writes it"
names = list(jit.captured.expected_names)
inputs = [{"slot": slot, "name": names[slot] if slot < len(names) else None, "dtype": repr(u.dtype), "numel": int(u.max_numel()), "device": u.device}
for slot, u in enumerate(_prepare_jit_inputs((), kwargs)[0])]
header = {"input_names": names, "calls": calls, "returns": {"kind": "NoneType", "tensors": []}, "inputs": inputs,
"buffers": [{"nbytes": b.arg.buffer.nbytes, "device": b.arg.buffer.device, "dtype": u8}
| ({"init": blob.add(b.arg.buffer.ensure_allocated().as_memoryview())} if i in needs_init else {}) for b, i in buffers.items()],
"buffers": [{"nbytes": bb.nbytes if (bb := _buffer(b)) else b.max_numel() * b.dtype.itemsize, "device": bb.device if bb else b.device, "dtype": u8}
| ({"init": blob.add(bb.ensure_allocated().as_memoryview())} if i in needs_init else {}) for b, i in buffers.items()],
"check": {"inputs": [[path, store(array), device] for path, array, device in frame["inputs"]],
"outputs": [store(array) for array in frame["outputs"]], "output_container": frame["output_container"]}}
with sink:
@@ -9,16 +9,14 @@ import itertools
import math
from functools import partial
import numpy as np
from openpilot.common.basedir import BASEDIR
from openpilot.sunnypilot.modeld_v2.helpers import _dtype, read_parts
def pack_layout(specs: dict[str, tuple], frame_size: int) -> tuple[dict[str, int], int]:
def pack_layout(shapes: dict[str, tuple], frame_size: int) -> tuple[dict[str, int], int]:
offsets, offset = {}, 0
for name, (shape, dtype) in ({"tfm": ((2, 3, 3), np.float32)} | specs).items():
offsets[name], offset = offset, offset + -(-math.prod(shape) * np.dtype(dtype).itemsize // 128) * 128
for name, shape in ({"tfm": (2, 3, 3)} | shapes).items():
offsets[name], offset = offset, offset + -(-math.prod(shape) * 4 // 128) * 128
return offsets | {"frames": offset}, offset + 2 * frame_size
@@ -40,7 +38,7 @@ def frame_part(warp_path: str, parts: dict, extras: dict) -> dict:
return {}
(model, model_blob), (warp, warp_blob) = parts["run"], read_parts(warp_path)[0]["run"]
states = sorted(name for name in extras["input_specs"] if f"next_{name}" in extras["metadata"]["output_shapes"])
host = {name: (shape, dtype) for name, (shape, dtype, _) in extras["input_specs"].items() if name not in states and name != "new_img"}
host = {name: shape for name, (shape, _, _) in extras["input_specs"].items() if name not in states and name != "new_img"}
frame_size = next(s["numel"] for s in warp["inputs"] if s["name"] == "input_frame") // 2
offsets, packed = pack_layout(host, frame_size)
shift, upload, u8 = len(model_blob), len(model["buffers"]) + len(warp["buffers"]), repr(_dtype("dtypes.uint8"))
+9 -3
View File
@@ -250,6 +250,12 @@ def _allocate_buffers(header: dict, blob: memoryview, device_map: dict | None, s
return buffers
def _fields(what, /, **values): # a tinygrad dataclass built or replaced with the fields it has today: a dropped field is fine, a new required one is not
known = {field.name for field in dataclasses.fields(what)}
values = {name: value for name, value in values.items() if name in known}
return what(**values) if isinstance(what, type) else dataclasses.replace(what, **values)
def _program_uop(call: dict, lib: bytes):
from tinygrad.helpers import Target
from tinygrad.uop.ops import KernelInfo, Ops, ProgramInfo, UOp
@@ -257,9 +263,9 @@ def _program_uop(call: dict, lib: bytes):
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"]))
parameters.append(parameter.replace(arg=_fields(parameter.arg, name=name, image=image)))
info = _fields(ProgramInfo, global_size=tuple(call["global_size"]), local_size=tuple(call["local_size"]), globals=tuple(call["globals"]),
outs=tuple(call["outs"]), ins=tuple(call["ins"]), target=_fields(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)
@@ -172,12 +172,12 @@ class NativeTinygradAdapter(BaseModelAdapter):
self.input_queues = {name: Tensor(np.zeros(shape, dtype=dtype), device=self.model_device).realize()
for name, (shape, dtype) in self.input_shapes.items() if name in self.state_pairs}
specs = {name: spec for name, spec in self.input_shapes.items() if name not in self.state_pairs and name != 'new_img'}
offsets, packed = pack_layout(specs, self.warp_frame_size)
shapes = {name: shape for name, (shape, _) in self.input_shapes.items() if name not in self.state_pairs and name != 'new_img'}
offsets, packed = pack_layout(shapes, self.warp_frame_size)
self.packed_input = np.zeros(packed, dtype=np.uint8)
self.host = Tensor(self.packed_input, device='NPY')
self.numpy_inputs = {name: np.ndarray(shape, dtype=dtype, buffer=self.packed_input, offset=offsets[name])
for name, (shape, dtype) in ({'tfm': ((2, 3, 3), np.dtype(np.float32))} | specs).items()}
self.numpy_inputs = {name: np.ndarray(shape, dtype=np.float32, buffer=self.packed_input, offset=offsets[name])
for name, shape in ({'tfm': (2, 3, 3)} | shapes).items()}
self.frames = self.packed_input[offsets['frames']:].reshape(2, self.warp_frame_size)
self.outputs = {name: Tensor(np.zeros(shape, dtype=dtype), device=device).realize() for name, (shape, dtype, device) in self.jits['output_specs'].items()}
+1 -22
View File
@@ -35,7 +35,6 @@ from openpilot.common.transformations.camera import DEVICE_CAMERAS
from openpilot.common.transformations.model import get_warp_matrix
from openpilot.selfdrive.controls.lib.desire_helper import DesireHelper
from openpilot.selfdrive.controls.lib.drive_helpers import get_accel_from_plan, smooth_value
from openpilot.selfdrive.navigation.helpers import navigation_input
from openpilot.selfdrive.modeld.modeld import ChestnutGpuState
from openpilot.sunnypilot.modeld_v2.fill_model_msg import fill_model_msg, fill_pose_msg, PublishState, get_curvature_from_output
@@ -173,9 +172,6 @@ class ModelState(ModelStateBase):
for key in ('traffic_convention', 'lateral_control_params', 'action_t'):
if key in self.numpy_inputs and key in inputs:
self.numpy_inputs[key][:] = inputs[key]
for name in ('nav_map', 'navigation'):
if name in self.numpy_inputs:
self.numpy_inputs[name][:] = inputs.get(name, 0)
if self.adapter.is_native:
for i, key in enumerate(self._vision_input_names):
@@ -257,10 +253,6 @@ class ModelState(ModelStateBase):
return log.ModelDataV2.Action(desiredCurvature=float(desired_curvature), desiredAcceleration=float(desired_accel), shouldStop=bool(stop))
def takes_nav(*models) -> bool:
return any(name in m.numpy_inputs for m in models if m is not None for name in ('nav_map', 'navigation'))
def main(demo=False):
cloudlog.warning("modeld init")
@@ -290,7 +282,6 @@ def main(demo=False):
vipc_client_main_stream = VisionStreamType.VISION_STREAM_WIDE_ROAD if main_wide_camera else VisionStreamType.VISION_STREAM_NARROW_ROAD
vipc_client_main = VisionIpcClient("camerad", vipc_client_main_stream, True)
vipc_client_extra = VisionIpcClient("camerad", VisionStreamType.VISION_STREAM_WIDE_ROAD, False)
vipc_client_map = VisionIpcClient("mapsd", VisionStreamType.VISION_STREAM_MAP, True)
cloudlog.warning(f"vision stream set up, main_wide_camera: {main_wide_camera}, use_extra_client: {use_extra_client}")
while not vipc_client_main.connect(False):
@@ -341,9 +332,8 @@ def main(demo=False):
# messaging
pub_socks = ["modelV2", "drivingModelData", "cameraOdometry", "modelDataV2SP"] + (["chestnutGpuState"] if CHESTNUT else [])
pm = PubMaster(pub_socks)
nav = takes_nav(model, small_model)
sm = SubMaster(["deviceState", "carState", "narrowRoadCameraState", "extrinsicsCalibration", "driverMonitoringState",
"carControl", "controlsState", "lateralDelay"] + (["navInstruction"] if nav else []))
"carControl", "controlsState", "lateralDelay"])
publish_state = PublishState()
chestnut_state = ChestnutGpuState(pm, model.chestnut) if CHESTNUT else None
@@ -376,7 +366,6 @@ def main(demo=False):
DH = DesireHelper()
meta_constants = load_meta_constants()
RELC = RoadEdgeLaneChangeController()
nav_map, nav_map_time = None, 0.
while True:
# Keep receiving frames until we are at least 1 frame ahead of previous extra frame
@@ -466,16 +455,6 @@ def main(demo=False):
if 'action_t' in model.numpy_inputs:
inputs['action_t'] = np.array([lat_action_t, long_action_t], dtype=np.float32)
now = time.monotonic()
if nav and (vipc_client_map.is_connected() or vipc_client_map.connect(False)):
if (buf_map := vipc_client_map.recv(0)) is not None:
nav_map = np.frombuffer(buf_map.data, dtype=np.uint8, count=buf_map.width * buf_map.height).reshape(buf_map.height, buf_map.width).copy()
nav_map_time = now
if nav_map is not None and now - nav_map_time < 0.5: # map timeout
inputs['nav_map'] = nav_map
if nav and sm.valid['navInstruction'] and now - sm.recv_time['navInstruction'] < 2:
inputs['navigation'] = navigation_input(sm['navInstruction'].allManeuvers, v_ego * (now - sm.recv_time['navInstruction']))
mt1 = time.perf_counter()
try:
send_chestnut = (chestnut_state is not None and
@@ -23,7 +23,7 @@ def record_kernel_linear():
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]):
with mock.patch.object(realize, "hcq_compile", lambda linear, *a, **k: (linear.src and recorded.append(linear), hcq_compile(linear, *a, **k))[1]):
return compile_linear(*args, **kwargs)
with mock.patch.object(jit, "compile_linear", recording_compile_linear):
@@ -38,6 +38,18 @@ def capture_frame(jit, kwargs: dict, output_container: str | None = None):
return {"inputs": inputs, "outputs": [tensor.numpy().copy() for tensor in outputs], "output_container": output_container}, result
def kernel_signature(ast, elf) -> list:
from tinygrad.dtype import AddrSpace
if hasattr(ast, "kernel_params"):
if any(p.arg.addrspace is AddrSpace.ALU for p in ast.kernel_params):
raise ValueError(f"{elf.name} takes a variable, which plain artifacts don't support")
compact = {slot: i for i, slot in enumerate(ast.arg.globals)}
return [[p.arg.name, compact[p.arg.slot], repr(p.dtype), [int(x) for x in p._shape]] for p in ast.kernel_params]
if not all(isinstance(slot, int) for _, slot, _, _ in elf.signature):
raise ValueError("TinyELF.signature lost its slots and this tinygrad has no UOp.kernel_params: update kernel_signature")
return [[name, slot, repr(dtype), [int(x) for x in shape]] for name, slot, dtype, shape in elf.signature]
def serialize_jit(jit, kernel_linear, check: dict | None = None) -> tuple[dict, memoryview]:
from tinygrad.uop.ops import Ops
sink = tempfile.TemporaryFile(dir=".")
@@ -51,6 +63,12 @@ def serialize_jit(jit, kernel_linear, check: dict | None = None) -> tuple[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}
root, offset = uop, 0
while root.op in (Ops.AFTER, Ops.BITCAST, Ops.SHRINK):
offset, root = offset + (root.src[1].arg * root.src[0].dtype.itemsize if root.op is Ops.SHRINK else 0), root.src[0]
if root.op is Ops.ALLOC:
base = bases.setdefault(id(root), (len(bases), root, root.dtype))
return {"buffer": base[0], "offset": offset, "nbytes": int(uop.max_numel()) * uop.dtype.itemsize, "dtype": repr(uop.dtype)}
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))
@@ -63,7 +81,7 @@ def serialize_jit(jit, kernel_linear, check: dict | None = None) -> tuple[dict,
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],
"signature": kernel_signature(ast, elf),
"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
@@ -77,7 +95,9 @@ def serialize_jit(jit, kernel_linear, check: dict | None = None) -> tuple[dict,
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)}
from tinygrad.device import Buffer
assert not [base for index, base, _ in bases.values() if index in needs_init and not isinstance(base, Buffer)], "a kernel reads a scratch before writing it"
buffers = [{"nbytes": base.nbytes if isinstance(base, Buffer) else int(base.max_numel()) * dtype.itemsize, "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:
@@ -0,0 +1,121 @@
"""
Copyright (c) 2021-, Haibin Wen, sunnypilot, and a number of other contributors.
This file is part of sunnypilot and is licensed under the MIT License.
See the LICENSE.md file in the root directory for more details.
"""
import os
import re
import sys
import tempfile
import unittest
from unittest import mock
import numpy as np
os.environ.setdefault("CHECK_OOB", "0")
os.environ.setdefault("GMMU", "0")
if sys.platform == "darwin" and os.environ.get("DEV", "").startswith("MOCKUSB"): # the mock's VRAM is a memfd, which macOS lacks: an unlinked file is the same
from tinygrad.runtime.autogen import libc
def _memfd_create(name, flags):
fd, path = tempfile.mkstemp()
os.unlink(path)
return fd
libc.memfd_create = _memfd_create
from openpilot.common.test import OpenpilotTestCase
from openpilot.sunnypilot.modeld_v2 import convert_upstream_pkl
from openpilot.sunnypilot.modeld_v2.convert_upstream_pkl import convert
from openpilot.sunnypilot.modeld_v2.helpers import PlainJit, load_model
from openpilot.sunnypilot.modeld_v2.plain_artifact_compiler import capture_frame
IMG, DESIRE, STATE, OUT = (6, 8, 8), (8,), (4, 16), (1, 32)
class TestConvertUpstreamPkl(OpenpilotTestCase):
@classmethod
def setUpClass(cls):
super().setUpClass()
if not os.environ.get("TINYGRAD_TESTS"):
raise unittest.SkipTest("TINYGRAD_TESTS not set. skipping converter test")
from tinygrad import Device, Tensor
assert Device.DEFAULT == "AMD" and getattr(Device["AMD"], "is_usb", False), f"{Device.DEFAULT} is not an AMD device behind the USB bridge"
assert (Tensor([1, 2, 3]) * 2).tolist() == [2, 4, 6]
cls.tmp = tempfile.TemporaryDirectory()
cls.rng = np.random.default_rng(0)
cls.weights = Tensor(cls.rng.standard_normal((16, 32)).astype(np.float32)).realize()
cls.bias = Tensor(cls.rng.standard_normal((1, 32)).astype(np.float32)).realize()
cls.input_specs = {"input_imgs": (IMG, "uint8", Device.DEFAULT), "desire": (DESIRE, "float32", Device.DEFAULT),
"features": (STATE, "float32", Device.DEFAULT)}
cls.output_specs = {"outputs": (OUT, "float32", Device.DEFAULT), "next_features": (STATE, "float32", Device.DEFAULT)}
@classmethod
def tearDownClass(cls):
cls.tmp.cleanup()
super().tearDownClass()
def model(self, input_imgs, desire, features, output_buffers):
from tinygrad import Tensor
feat = features.mean(axis=0).reshape(1, 16) + input_imgs.float().mean() + desire.sum()
Tensor.realize(output_buffers["outputs"].assign((feat @ self.weights + self.bias).relu()),
output_buffers["next_features"].assign(features[1:].cat(feat * 2)))
def inputs(self, seed: int) -> dict:
from examples.openpilot.helpers import allocate_inputs
rng = np.random.default_rng(seed)
def initialize(arrays: dict):
for value in arrays.values():
value[...] = rng.standard_normal(value.shape) if value.dtype.kind == "f" else rng.integers(0, 256, value.shape)
return allocate_inputs(self.input_specs, initialize) | {"output_buffers": allocate_inputs(self.output_specs)}
def jit(self):
from tinygrad import TinyJit
run = TinyJit(self.model, prune=True)
kwargs = self.inputs(1)
run(**kwargs)
run(**kwargs)
return run
def assert_matches_eager(self, plain, seed: int):
fresh, expected = self.inputs(seed), self.inputs(seed)
plain(**fresh)
self.model(**expected)
for name in self.output_specs:
assert np.array_equal(fresh["output_buffers"][name].numpy(), expected["output_buffers"][name].numpy()), name
def convert(self, jit) -> tuple[dict, memoryview]:
from tinygrad import Device
kwargs = self.inputs(2)
frame, _ = capture_frame(jit, kwargs, "output_buffers")
header, blob = convert(jit, kwargs, frame, Device[Device.DEFAULT].renderer.target)
kernels = jit.captured._linear.src[0].without_after.arg.aux.kernels
assert [call["name"] for call in header["calls"]] == [re.sub(r"\x1b\[[0-9;]*m", "", kernel[1]) for kernel in kernels]
assert all(re.fullmatch(r"\w+", call["name"]) for call in header["calls"]), "kernel names: failure point 1"
assert [(tuple(call["outs"]), tuple(call["ins"])) for call in header["calls"]] == [tuple(kernel[6]) for kernel in kernels], "outs/ins: failure point 1"
bound = {argument["input"] for call in header["calls"] for argument in call["args"] if "input" in argument}
assert bound == {spec["slot"] for spec in header["inputs"]}, f"input slots bound by pokes {sorted(bound)}: failure point 4"
assert [spec["name"] for spec in header["inputs"]][:len(self.input_specs)] == sorted(self.input_specs)
return header, blob
def test_converted_jit_matches_eager(self): # the submit decoded into plain calls, replayed against its own frame
jit = self.jit()
plain = PlainJit(*self.convert(jit))
self.assert_matches_eager(plain, 3)
self.assert_matches_eager(plain, 4)
def test_pickle_round_trip_through_main(self):
from examples.openpilot.helpers import dump_pickle, load_pickle
from tinygrad.engine.realize import lower_and_compile
pkl, out = os.path.join(self.tmp.name, "tiny_tinygrad.pkl"), os.path.join(self.tmp.name, "tiny_plain.pkl")
dump_pickle({"run": self.jit(), "input_specs": self.input_specs, "output_specs": self.output_specs, "metadata": {"name": "tiny"}}, pkl)
artifact = load_pickle(pkl, out_of_band=True)
artifact["run"].captured._linear = lower_and_compile(artifact["run"].captured._linear)
self.assert_matches_eager(PlainJit(*self.convert(artifact["run"])), 5)
with mock.patch.object(sys, "argv", ["convert_upstream_pkl", pkl, out]):
convert_upstream_pkl.main()
bundle = load_model(out, cache=False, verify=True)
assert bundle["metadata"] == {"name": "tiny"} and bundle["input_specs"] == self.input_specs and bundle["output_specs"] == self.output_specs
self.assert_matches_eager(bundle["run"], 6)
@@ -0,0 +1,162 @@
"""
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 functools
import inspect
import os
import sys
import tempfile
from unittest import mock
import numpy as np
import unittest
import openpilot.sunnypilot.modeld_v2.helpers as helpers
from openpilot.common.test import OpenpilotTestCase
from openpilot.sunnypilot.modeld_v2 import compile_warp
from openpilot.sunnypilot.modeld_v2.frame_fusion import frame_part, pack_layout
from openpilot.sunnypilot.modeld_v2.helpers import PlainJit, cache_is_fresh, load_model, read_parts, write_bundle
from openpilot.sunnypilot.modeld_v2.plain_artifact_compiler import capture_frame, record_kernel_linear, serialize_jit
CAM_W, CAM_H, MODEL_W, MODEL_H = 32, 16, 16, 8
FRAME = f"{CAM_W},{CAM_H},{CAM_W},{CAM_H},{CAM_H // 2},{CAM_W * CAM_H * 3 // 2}"
FRAME_SIZE = CAM_W * CAM_H * 3 // 2
HOST = {"desire": (8,), "traffic_convention": (2,)}
STATE, OUT = (4, 16), (1, 32)
def compile_jit(fn, inputs: dict, output_container: str | None = None) -> tuple[dict, memoryview]:
from tinygrad import TinyJit
run = TinyJit(fn, prune=True)
with record_kernel_linear() as recorded:
run(**inputs)
run(**inputs)
check, _ = capture_frame(run, inputs, output_container)
assert len(recorded) == 1, f"expected one captured kernel list, got {len(recorded)}"
return serialize_jit(run, recorded[0], check)
class TestPlainArtifact(OpenpilotTestCase):
@classmethod
def setUpClass(cls):
super().setUpClass()
if not os.environ.get("TINYGRAD_TESTS"):
raise unittest.SkipTest("TINYGRAD_TESTS env var not set, skipping artifact test.")
from tinygrad import Tensor
cls.tmp = tempfile.TemporaryDirectory()
cls.rng = np.random.default_rng(0)
cls.weights = Tensor(cls.rng.standard_normal((16, 32)).astype(np.float32)).realize()
cls.bias = Tensor(cls.rng.standard_normal((1, 32)).astype(np.float32)).realize()
@classmethod
def tearDownClass(cls):
cls.tmp.cleanup()
super().tearDownClass()
def model(self, x, h): # a two input, one weight model: the shape of a policy without the size
return ((x @ self.weights + self.bias).relu() + h.mean()).contiguous()
def tensors(self, seed: int):
from tinygrad import Tensor
rng = np.random.default_rng(seed)
return {"x": Tensor(rng.standard_normal((1, 16)).astype(np.float32)).realize(), "h": Tensor(rng.standard_normal((4, 16)).astype(np.float32)).realize()}
def artifact(self, name: str = "model") -> str:
path = os.path.join(self.tmp.name, f"{name}_tinygrad.pkl")
if not os.path.exists(path):
write_bundle(path, {"run": compile_jit(self.model, self.tensors(1))}, {"metadata": {"name": name}})
return path
def test_round_trip_matches_jit(self): # compiler -> bundle -> PlainJit, replay verified, then fresh inputs bit exact with eager tinygrad
bundle = load_model(self.artifact(), cache=False, verify=True)
assert bundle["metadata"] == {"name": "model"} and isinstance(bundle["run"], PlainJit)
fresh = self.tensors(2)
assert np.array_equal(bundle["run"](**fresh).numpy(), self.model(**fresh).numpy())
def test_replay_check_catches_a_changed_kernel(self):
parts, extras = read_parts(self.artifact())
header, blob = parts["run"]
poisoned = dict(header, check=dict(header["check"], outputs=[header["check"]["inputs"][0][1]]))
with self.assertRaises(helpers.ArtifactMismatch):
PlainJit(poisoned, blob)
def test_cache_round_trip_holds_no_weights_and_is_rebuilt_when_stale(self):
path = self.artifact("cached")
expected = load_model(path, cache=False, verify=False)["run"](**self.tensors(3)).numpy()
load_model(path)
assert cache_is_fresh(path)
with open(f"{path}.cache", "rb") as file:
cache = file.read()
assert self.weights.numpy().tobytes() not in cache and len(cache) < os.path.getsize(path)
with mock.patch.object(helpers, "PlainJit", side_effect=AssertionError("a fresh cache must not rebuild")):
cached = load_model(path)
assert np.array_equal(cached["run"](**self.tensors(3)).numpy(), expected)
with mock.patch.object(helpers, "tinygrad_identity", return_value="another tinygrad"):
assert not cache_is_fresh(path)
assert isinstance(load_model(path)["run"], PlainJit)
assert cache_is_fresh(path)
assert not cache_is_fresh(path) # keyed on the tinygrad that wrote it, so the next boot after a bump rebuilds once
os.utime(path)
with mock.patch.object(helpers, "tinygrad_identity", return_value="another tinygrad"):
assert not cache_is_fresh(path)
def test_fused_with_warp_matches_warp_then_model(self):
from tinygrad import Tensor
warp_path = os.path.join(self.tmp.name, "warp_tinygrad.pkl")
with mock.patch.object(sys, "argv", ["compile_warp", "--frame", FRAME, "--warp-to", f"{MODEL_W}x{MODEL_H}", "--output", warp_path]):
compile_warp.main()
warp = load_model(warp_path, cache=False)["run"]
def native(new_img, desire, traffic_convention, state_feat_q, output_buffers):
feat = state_feat_q.mean(axis=0).reshape(1, 16) + new_img.float().mean() + desire.sum() + traffic_convention[0]
Tensor.realize(output_buffers["outputs"].assign((feat @ self.weights + self.bias).relu()),
output_buffers["next_state_feat_q"].assign(state_feat_q[1:].cat(feat * 2)))
def model_inputs(new_img):
return {"new_img": new_img} | {name: Tensor(np.ones(shape, np.float32) * 0.5).realize() for name, shape in HOST.items()} | \
{"state_feat_q": Tensor(self.rng.standard_normal(STATE).astype(np.float32)).realize(),
"output_buffers": {"outputs": Tensor.zeros(OUT).contiguous().realize(), "next_state_feat_q": Tensor.zeros(STATE).contiguous().realize()}}
frames = Tensor(self.rng.integers(0, 256, (2, FRAME_SIZE), dtype=np.uint8)).realize()
tfm = Tensor(np.tile(np.eye(3, dtype=np.float32), (2, 1, 1))).realize()
path = os.path.join(self.tmp.name, "native_tinygrad.pkl")
part = compile_jit(native, model_inputs(warp(input_frame=frames, M_inv=tfm).realize()), "output_buffers")
extras = {"input_specs": {"new_img": ((2, 6, MODEL_H // 2, MODEL_W // 2), "|u1", "CPU"), "state_feat_q": (STATE, "<f4", "CPU")}
| {name: (shape, "<f4", "CPU") for name, shape in HOST.items()},
"output_specs": {"outputs": (OUT, "<f4", "CPU"), "next_state_feat_q": (STATE, "<f4", "CPU")},
"metadata": {"output_shapes": {"outputs": OUT, "next_state_feat_q": STATE}}}
write_bundle(path, {"run": part}, extras)
fused = load_model(path, cache=False, extend=functools.partial(frame_part, warp_path))
assert fused["run"].input_names == ["host", "state_feat_q"]
offsets, packed = pack_layout(HOST, FRAME_SIZE)
host = np.zeros(packed, np.uint8)
inputs = model_inputs(None)
for name, shape in ({"tfm": (2, 3, 3)} | HOST).items():
np.ndarray(shape, np.float32, host, offsets[name])[...] = tfm.numpy() if name == "tfm" else inputs[name].numpy()
host[offsets["frames"]:] = frames.numpy().ravel()
fused["run"](host=Tensor(host, device="NPY"), state_feat_q=inputs["state_feat_q"], output_buffers=inputs["output_buffers"])
expected = model_inputs(warp(input_frame=frames, M_inv=tfm).realize()) | {k: inputs[k] for k in ("state_feat_q",)}
native(**expected)
for name in ("outputs", "next_state_feat_q"):
assert np.array_equal(inputs["output_buffers"][name].numpy(), expected["output_buffers"][name].numpy()), name
def test_tinygrad_api_the_loader_reads(self):
import tinygrad.engine.jit as jit
from tinygrad.device import Buffer, TinyELF
from tinygrad.helpers import Target
from tinygrad.uop.ops import ParamArg, ProgramInfo, UOp
fields = {cls: {f.name for f in dataclasses.fields(cls)} for cls in (ProgramInfo, ParamArg, Target)}
assert fields[ProgramInfo] >= {"global_size", "local_size", "globals", "outs", "ins", "target"}, "ProgramInfo changed: update helpers._program_uop"
assert fields[ParamArg] >= {"slot", "dtype", "name", "image"}, "ParamArg changed: update helpers._program_uop"
assert fields[Target] >= {"device", "renderer", "arch"}, "Target changed: update the 'target' key in plain_artifact_compiler.serialize_jit"
assert "signature" in getattr(TinyELF, "__annotations__", {}) or hasattr(UOp, "kernel_params"), \
"kernel signatures moved: update plain_artifact_compiler.serialize_jit"
assert callable(getattr(UOp, "param", None)) and callable(getattr(UOp, "from_buffer", None)), "UOp.param/from_buffer: update helpers.PlainJit"
assert callable(getattr(UOp, "store_call", None)) and callable(getattr(UOp, "call", None)), "UOp.call/store_call: update helpers.PlainJit"
assert callable(getattr(Buffer, "view", None)) and callable(getattr(Buffer, "ensure_allocated", None)), "Buffer.view: update helpers._upload"
assert callable(getattr(jit, "compile_linear", None)) and callable(getattr(jit, "_prepare_jit_inputs", None)), "engine.jit: update helpers._Jit"
assert "expected_names" in inspect.signature(jit.CapturedJit).parameters, "CapturedJit: update helpers.PlainJit and plain_artifact_compiler.serialize_jit"
+146 -48
View File
@@ -7,7 +7,10 @@ See the LICENSE.md file in the root directory for more details.
import asyncio
import os
import re
import threading
import time
from concurrent.futures import ThreadPoolExecutor
import requests
from openpilot.common.params import Params
@@ -22,12 +25,99 @@ from openpilot.sunnypilot.models.helpers import (ACTIVE_BUNDLE_KEYS, get_active_
# (connect, read) seconds. read is per-request inactivity, not a total cap
DOWNLOAD_TIMEOUT = (30, 30)
# HuggingFace throttles each connection to ~1-2 MB/s but not the connection count; 12 saturates the device link
MAX_CONCURRENT_PIECES = 12
PIECE_SIZE = 8 * 1024 * 1024
PIECE_ATTEMPTS = 3
PIECE_BACKOFF = 1.0 # seconds, growing linearly: waits 1 s then 2 s between a piece's three attempts
TRANSIENT_HTTP_STATUS = frozenset({408, 429, 500, 502, 503, 504})
REPORT_INTERVAL = 0.5 # seconds between progress publications
# byte ranges count encoded bytes but iter_content yields decoded ones, so ask for an unencoded body
NO_COMPRESSION = {"Accept-Encoding": "identity"}
class DownloadCancelled(Exception):
pass
def _is_transient(e: BaseException) -> bool:
if isinstance(e, requests.HTTPError):
return e.response is not None and e.response.status_code in TRANSIENT_HTTP_STATUS
return isinstance(e, (requests.ConnectionError, requests.Timeout, requests.exceptions.ChunkedEncodingError))
class _Progress:
def __init__(self):
self._bytes = 0
self._lock = threading.Lock()
def add_bytes(self, n: int) -> None:
with self._lock:
self._bytes += n
def snapshot(self) -> int:
with self._lock:
return self._bytes
def _piece_ranges(total: int, piece_size: int) -> list[tuple[int, int]]:
return [(start, min(start + piece_size, total)) for start in range(0, total, piece_size)]
def _prepare_target(path: str, total: int) -> None:
with open(path, "wb") as f:
f.truncate(total)
def _content_range_total(response: requests.Response) -> int | None:
if response.status_code != 206:
return None
match = re.fullmatch(r"bytes \d+-\d+/(\d+)", response.headers.get("Content-Range", ""))
return int(match.group(1)) if match else None
def _probe_size(url: str) -> tuple[int, bool]:
with requests.get(url, headers={**NO_COMPRESSION, "Range": "bytes=0-0"}, stream=True, timeout=DOWNLOAD_TIMEOUT) as response:
response.raise_for_status()
total = _content_range_total(response)
if total is not None:
return total, True
return int(response.headers.get("Content-Length", 0)), False
def _fetch_piece(url: str, path: str, index: int, start: int, end: int | None, progress: _Progress,
cancel: threading.Event, block_size: int) -> None:
headers = {**NO_COMPRESSION, "Range": f"bytes={start}-{end - 1}"} if end is not None else NO_COMPRESSION
label = f"{os.path.basename(path)} piece {index}"
for attempt in range(PIECE_ATTEMPTS):
written = 0
try:
with requests.get(url, headers=headers, stream=True, timeout=DOWNLOAD_TIMEOUT) as response:
response.raise_for_status()
if end is not None and response.status_code != 206:
raise ValueError(f"server ignored byte range for {label}")
with open(path, "r+b") as f:
f.seek(start)
for block in response.iter_content(chunk_size=block_size):
if cancel.is_set():
raise DownloadCancelled("Download cancelled")
f.write(block)
written += len(block)
progress.add_bytes(len(block))
if end is None:
f.truncate() # unknown length: the body defines the file size
if end is not None and written != end - start:
raise requests.exceptions.ConnectionError(f"short read: {written} of {end - start} bytes")
return
except requests.RequestException as e:
progress.add_bytes(-written) # the piece restarts from scratch
if not _is_transient(e) or attempt == PIECE_ATTEMPTS - 1:
raise
cloudlog.warning(f"retrying {label} after {type(e).__name__}: {e}")
if cancel.wait(PIECE_BACKOFF * (1 + attempt)):
raise DownloadCancelled("Download cancelled") from None
class ModelManagerSP:
"""Manages model downloads and status reporting"""
@@ -41,7 +131,7 @@ class ModelManagerSP:
self.source_models: dict[str, list[custom.ModelManagerSP.ModelBundle]] = {}
self.selected_bundle: custom.ModelManagerSP.ModelBundle = None
self.active_bundle: custom.ModelManagerSP.ModelBundle = get_active_bundle(self.params, chestnut=self.chestnut_present)
self._chunk_size = 128 * 1000 # 128 KB chunks
self._block_size = 128 * 1000 # 128 KB network read blocks
self._download_start_times: dict[str, float] = {} # Track start time per model
self._download_ref: bytes | str | None = None
@@ -56,7 +146,6 @@ class ModelManagerSP:
self._download_ref = None
def _sync_artifact_progress(self, source_artifact) -> None:
"""Mirror download progress to all artifacts sharing the same filename in the selected bundle."""
if not self.selected_bundle:
return
for model in self.selected_bundle.models:
@@ -67,7 +156,6 @@ class ModelManagerSP:
artifact.downloadProgress.eta = source_artifact.downloadProgress.eta
def _calculate_eta(self, filename: str, progress: float) -> int:
"""Calculate ETA based on elapsed time and current progress"""
if filename not in self._download_start_times or progress <= 0:
return 60 # Default ETA for new downloads
@@ -75,39 +163,58 @@ class ModelManagerSP:
if elapsed_time <= 0:
return 60
# If we're at X% after Y seconds, we can estimate total time as (Y / X) * 100
total_estimated_time = (elapsed_time / progress) * 100
eta = total_estimated_time - elapsed_time
return max(1, int(eta)) # Return at least 1 second if download is ongoing
async def _download_file(self, url: str, path: str, model) -> None:
"""Downloads a file with progress tracking"""
self._download_start_times[model.fileName] = time.monotonic()
def _set_progress(self, artifact, status, progress: float, eta: int = 0) -> None:
artifact.downloadProgress.status = status
artifact.downloadProgress.progress = progress
artifact.downloadProgress.eta = eta
self._sync_artifact_progress(artifact)
self._report_status()
with requests.get(url, stream=True, timeout=DOWNLOAD_TIMEOUT) as response: # noqa: ASYNC210
response.raise_for_status()
total_size = int(response.headers.get("content-length", 0))
bytes_downloaded = 0
def _publish_progress(self, artifact, done_bytes: int, total: int) -> None:
# held at 99 until the file passes its hash check
progress = min(99.0, done_bytes / total * 100) if total > 0 else 0.0
eta = self._calculate_eta(artifact.fileName, progress)
self._set_progress(artifact, custom.ModelManagerSP.DownloadStatus.downloading, progress, eta)
with open(path, 'wb') as f: # noqa: ASYNC230
for chunk in response.iter_content(chunk_size=self._chunk_size): # type: bytes
f.write(chunk)
bytes_downloaded += len(chunk)
async def _report_until_done(self, tasks: list[asyncio.Future], artifact, progress: _Progress, total: int) -> None:
pending = set(tasks)
while pending:
done, pending = await asyncio.wait(pending, timeout=REPORT_INTERVAL)
for task in done:
task.result() # re-raises a worker failure
if self._download_interrupted():
raise DownloadCancelled("Download cancelled")
self._publish_progress(artifact, progress.snapshot(), total)
if self._download_interrupted():
raise DownloadCancelled("Download cancelled")
async def _download_file(self, url: str, path: str, artifact) -> None:
self._download_start_times[artifact.fileName] = time.monotonic()
loop = asyncio.get_running_loop()
if total_size > 0:
progress = (bytes_downloaded / total_size) * 100
model.downloadProgress.status = custom.ModelManagerSP.DownloadStatus.downloading
model.downloadProgress.progress = progress
model.downloadProgress.eta = self._calculate_eta(model.fileName, progress)
self._sync_artifact_progress(model)
self._report_status()
total, ranged = await loop.run_in_executor(None, _probe_size, url)
pieces: list[tuple[int, int | None]] = list(_piece_ranges(total, PIECE_SIZE)) if ranged else [(0, None)]
await loop.run_in_executor(None, _prepare_target, path, total)
progress = _Progress()
cancel = threading.Event()
# Clean up start time after download completes
del self._download_start_times[model.fileName]
with ThreadPoolExecutor(max_workers=MAX_CONCURRENT_PIECES) as pool:
tasks = [loop.run_in_executor(pool, _fetch_piece, url, path, i, start, end, progress, cancel, self._block_size)
for i, (start, end) in enumerate(pieces)]
try:
await self._report_until_done(tasks, artifact, progress, total)
except BaseException as e:
cancel.set()
if isinstance(e, DownloadCancelled):
artifact.downloadProgress.canceling = True
self._set_progress(artifact, custom.ModelManagerSP.DownloadStatus.downloading, artifact.downloadProgress.progress)
for task in tasks:
task.cancel() # drops queued pieces; running ones see the event at their next block
await asyncio.gather(*tasks, return_exceptions=True)
raise
del self._download_start_times[artifact.fileName]
async def _process_artifact(self, artifact, destination_path: str) -> None:
if not artifact.downloadUri.uri:
@@ -119,33 +226,30 @@ class ModelManagerSP:
expected_hash = artifact.downloadUri.sha256
filename = artifact.fileName
full_path = os.path.join(destination_path, filename)
status = custom.ModelManagerSP.DownloadStatus
artifact.downloadProgress.canceling = False
try:
self._set_progress(artifact, status.verifying, 0)
if await verify_file(full_path, expected_hash):
artifact.downloadProgress.status = custom.ModelManagerSP.DownloadStatus.cached
artifact.downloadProgress.progress = 100
artifact.downloadProgress.eta = 0
self._sync_artifact_progress(artifact)
self._report_status()
self._set_progress(artifact, status.cached, 100)
return
self._set_progress(artifact, status.downloading, 0)
await self._download_file(url, full_path, artifact)
self._set_progress(artifact, status.verifying, 99)
if not await verify_file(full_path, expected_hash):
raise ValueError(f"Hash validation failed for {filename}")
artifact.downloadProgress.status = custom.ModelManagerSP.DownloadStatus.downloaded
artifact.downloadProgress.progress = 100
artifact.downloadProgress.eta = 0
self._sync_artifact_progress(artifact)
self._report_status()
self._set_progress(artifact, status.downloaded, 100)
except DownloadCancelled:
self._download_start_times.pop(artifact.fileName, None)
artifact.downloadProgress.status = custom.ModelManagerSP.DownloadStatus.failed
artifact.downloadProgress.status = status.failed
artifact.downloadProgress.eta = 0
self._sync_artifact_progress(artifact)
if self.selected_bundle:
self.selected_bundle.status = custom.ModelManagerSP.DownloadStatus.failed
self.selected_bundle.status = status.failed
self._report_status()
raise
@@ -153,11 +257,11 @@ class ModelManagerSP:
cloudlog.error(f"Error downloading {filename}: {str(e)}")
if os.path.isfile(full_path): # noqa: ASYNC240
os.remove(full_path)
artifact.downloadProgress.status = custom.ModelManagerSP.DownloadStatus.failed
artifact.downloadProgress.status = status.failed
artifact.downloadProgress.eta = 0
self._sync_artifact_progress(artifact)
if self.selected_bundle:
self.selected_bundle.status = custom.ModelManagerSP.DownloadStatus.failed
self.selected_bundle.status = status.failed
self._report_status()
self._download_start_times.pop(artifact.fileName, None)
raise
@@ -272,12 +376,7 @@ class ModelManagerSP:
rk.keep_time()
def clear_model_cache(self) -> None:
"""
Clears the model cache directory of all files except those in the active model bundle.
"""
# Get list of files used by both slots' selected bundles (either may become
# the truly active bundle depending on hardware availability)
""" Clears the model cache directory of all files except those in the active model bundle."""
active_files = []
for source in ACTIVE_BUNDLE_KEYS:
if selected_bundle := get_selected_bundle(self.params, source):
@@ -285,7 +384,6 @@ class ModelManagerSP:
if model.artifact.fileName:
active_files.append(model.artifact.fileName)
# Remove all files except active ones
model_dir = Paths.model_root()
try:
for filename in os.listdir(model_dir):
File diff suppressed because it is too large Load Diff
@@ -57,7 +57,7 @@ class DownloadStatusAction(ItemAction):
self.icon: str | None = None
self.icon_color: rl.Color | None = None
self._font = gui_app.font(FontWeight.NORMAL)
# raw progress arrives in steps, one per 128KB chunk the manager publishes
# raw progress arrives in steps, one per progress report from the manager
self._progress = FirstOrderFilter(0.0, 0.5, 1 / gui_app.target_fps)
# integrated per frame; (t * speed) % span jumps whenever the fill width changes
self._sweep = 0.0