Reapply chunker (#37292)

* Reapply chunker

* good size

* rm glob

* cleaner

* back to 45mb

* warp need not be fixed

* add manifest path

* lil cleaner
This commit is contained in:
Harald Schäfer
2026-02-23 16:49:48 -08:00
committed by GitHub
parent 76d084d877
commit 16dda06a0c
6 changed files with 60 additions and 68 deletions
+1 -2
View File
@@ -64,8 +64,7 @@ flycheck_*
cppcheck_report.txt
comma*.sh
selfdrive/modeld/models/*.pkl
selfdrive/modeld/models/*.pkl.*
selfdrive/modeld/models/*.pkl*
# openpilot log files
*.bz2
+37
View File
@@ -0,0 +1,37 @@
import math
import os
from pathlib import Path
CHUNK_SIZE = 45 * 1024 * 1024 # 45MB, under GitHub's 50MB limit
def get_chunk_name(name, idx, num_chunks):
return f"{name}.chunk{idx+1:02d}of{num_chunks:02d}"
def get_manifest_path(name):
return f"{name}.chunkmanifest"
def get_chunk_paths(path, file_size):
num_chunks = math.ceil(file_size / CHUNK_SIZE)
return [get_manifest_path(path)] + [get_chunk_name(path, i, num_chunks) for i in range(num_chunks)]
def chunk_file(path, targets):
manifest_path, *chunk_paths = targets
with open(path, 'rb') as f:
data = f.read()
actual_num_chunks = max(1, math.ceil(len(data) / CHUNK_SIZE))
assert len(chunk_paths) >= actual_num_chunks, f"Allowed {len(chunk_paths)} chunks but needs at least {actual_num_chunks}, for path {path}"
for i, chunk_path in enumerate(chunk_paths):
with open(chunk_path, 'wb') as f:
f.write(data[i * CHUNK_SIZE:(i + 1) * CHUNK_SIZE])
Path(manifest_path).write_text(str(len(chunk_paths)))
os.remove(path)
def read_file_chunked(path):
manifest_path = get_manifest_path(path)
if os.path.isfile(manifest_path):
num_chunks = int(Path(manifest_path).read_text().strip())
return b''.join(Path(get_chunk_name(path, i, num_chunks)).read_bytes() for i in range(num_chunks))
if os.path.isfile(path):
return Path(path).read_bytes()
raise FileNotFoundError(path)
+17 -23
View File
@@ -1,14 +1,18 @@
import os
import glob
from openpilot.common.file_chunker import chunk_file, get_chunk_paths
Import('env', 'arch')
chunker_file = File("#common/file_chunker.py")
lenv = env.Clone()
CHUNK_BYTES = int(os.environ.get("TG_CHUNK_BYTES", str(45 * 1024 * 1024)))
tinygrad_root = env.Dir("#").abspath
tinygrad_files = ["#"+x for x in glob.glob(env.Dir("#tinygrad_repo").relpath + "/**", recursive=True, root_dir=tinygrad_root)
if 'pycache' not in x and os.path.isfile(os.path.join(tinygrad_root, x))]
def estimate_pickle_max_size(onnx_size):
return 1.2 * onnx_size + 10 * 1024 * 1024 # 20% + 10MB is plenty
# Get model metadata
for model_name in ['driving_vision', 'driving_policy', 'dmonitoring_model']:
fn = File(f"models/{model_name}").abspath
@@ -26,39 +30,29 @@ image_flag = {
'larch64': 'IMAGE=2',
}.get(arch, 'IMAGE=0')
script_files = [File(Dir("#selfdrive/modeld").File("compile_warp.py").abspath)]
cmd = f'{tg_flags} python3 {Dir("#selfdrive/modeld").abspath}/compile_warp.py '
compile_warp_cmd = f'{tg_flags} python3 {Dir("#selfdrive/modeld").abspath}/compile_warp.py '
from openpilot.common.transformations.camera import _ar_ox_fisheye, _os_fisheye
warp_targets = []
for cam in [_ar_ox_fisheye, _os_fisheye]:
w, h = cam.width, cam.height
warp_targets += [File(f"models/warp_{w}x{h}_tinygrad.pkl").abspath, File(f"models/dm_warp_{w}x{h}_tinygrad.pkl").abspath]
lenv.Command(warp_targets, tinygrad_files + script_files, cmd)
lenv.Command(warp_targets, tinygrad_files + script_files, compile_warp_cmd)
def tg_compile(flags, model_name):
pythonpath_string = 'PYTHONPATH="${PYTHONPATH}:' + env.Dir("#tinygrad_repo").abspath + '"'
fn = File(f"models/{model_name}").abspath
out = fn + "_tinygrad.pkl"
full = out + ".full"
parts = out + ".parts"
full_node = lenv.Command(
full,
[fn + ".onnx"] + tinygrad_files,
f'{pythonpath_string} {flags} {image_flag} python3 {Dir("#tinygrad_repo").abspath}/examples/openpilot/compile3.py {fn}.onnx {full}'
pkl = fn + "_tinygrad.pkl"
onnx_path = fn + ".onnx"
chunk_targets = get_chunk_paths(pkl, estimate_pickle_max_size(os.path.getsize(onnx_path)))
def do_chunk(target, source, env):
chunk_file(pkl, chunk_targets)
return lenv.Command(
chunk_targets,
[onnx_path] + tinygrad_files + [chunker_file],
[f'{pythonpath_string} {flags} {image_flag} python3 {Dir("#tinygrad_repo").abspath}/examples/openpilot/compile3.py {fn}.onnx {pkl}',
do_chunk]
)
split_script = File(Dir("#selfdrive/modeld").File("external_pickle.py").abspath)
parts_node = lenv.Command(
parts,
[full_node, split_script, Value(str(CHUNK_BYTES))],
[f'python3 {split_script.abspath} {full} {out} {CHUNK_BYTES}', Delete(full)],
)
lenv.NoCache(parts_node)
lenv.AlwaysBuild(parts_node)
return parts_node
# Compile small models
for model_name in ['driving_vision', 'driving_policy', 'dmonitoring_model']:
tg_compile(tg_flags, model_name)
+2 -2
View File
@@ -16,8 +16,8 @@ from openpilot.common.realtime import config_realtime_process
from openpilot.common.transformations.model import dmonitoringmodel_intrinsics
from openpilot.common.transformations.camera import _ar_ox_fisheye, _os_fisheye
from openpilot.system.camerad.cameras.nv12_info import get_nv12_info
from openpilot.common.file_chunker import read_file_chunked
from openpilot.selfdrive.modeld.parse_model_outputs import sigmoid, safe_exp
from openpilot.selfdrive.modeld.external_pickle import load_external_pickle
PROCESS_NAME = "selfdrive.modeld.dmonitoringmodeld"
SEND_RAW_PRED = os.getenv('SEND_RAW_PRED')
@@ -45,7 +45,7 @@ class ModelState:
self.tensor_inputs = {k: Tensor(v, device='NPY').realize() for k,v in self.numpy_inputs.items()}
self._blob_cache : dict[int, Tensor] = {}
self.image_warp = None
self.model_run = load_external_pickle(MODEL_PKL_PATH)
self.model_run = pickle.loads(read_file_chunked(str(MODEL_PKL_PATH)))
def run(self, buf: VisionBuf, calib: np.ndarray, transform: np.ndarray) -> tuple[np.ndarray, float]:
self.numpy_inputs['calib'][0,:] = calib
-38
View File
@@ -1,38 +0,0 @@
#!/usr/bin/env python3
import hashlib
import pickle
import sys
from pathlib import Path
def split_pickle(full_path: Path, out_prefix: Path, chunk_bytes: int) -> None:
data = full_path.read_bytes()
out_dir = out_prefix.parent
for p in out_dir.glob(f"{out_prefix.name}.data-*"):
p.unlink()
total = (len(data) + chunk_bytes - 1) // chunk_bytes
names = []
for i in range(0, len(data), chunk_bytes):
name = f"{out_prefix.name}.data-{(i // chunk_bytes) + 1:04d}-of-{total:04d}"
(out_dir / name).write_bytes(data[i:i + chunk_bytes])
names.append(name)
manifest = hashlib.sha256(data).hexdigest() + "\n" + "\n".join(names) + "\n"
(out_dir / (out_prefix.name + ".parts")).write_text(manifest)
def load_external_pickle(prefix: Path):
parts = prefix.parent / (prefix.name + ".parts")
lines = parts.read_text().splitlines()
expected_hash, chunk_names = lines[0], lines[1:]
data = bytearray()
for name in chunk_names:
data += (prefix.parent / name).read_bytes()
if hashlib.sha256(data).hexdigest() != expected_hash:
raise RuntimeError(f"hash mismatch loading {prefix}")
return pickle.loads(data)
if __name__ == "__main__":
split_pickle(Path(sys.argv[1]), Path(sys.argv[2]), int(sys.argv[3]))
+3 -3
View File
@@ -27,8 +27,8 @@ from openpilot.selfdrive.controls.lib.desire_helper import DesireHelper
from openpilot.selfdrive.controls.lib.drive_helpers import get_accel_from_plan, smooth_value, get_curvature_from_plan
from openpilot.selfdrive.modeld.parse_model_outputs import Parser
from openpilot.selfdrive.modeld.fill_model_msg import fill_model_msg, fill_pose_msg, PublishState
from openpilot.common.file_chunker import read_file_chunked
from openpilot.selfdrive.modeld.constants import ModelConstants, Plan
from openpilot.selfdrive.modeld.external_pickle import load_external_pickle
PROCESS_NAME = "selfdrive.modeld.modeld"
@@ -178,8 +178,8 @@ class ModelState:
self.parser = Parser()
self.frame_buf_params : dict[str, tuple[int, int, int, int]] = {}
self.update_imgs = None
self.vision_run = load_external_pickle(VISION_PKL_PATH)
self.policy_run = load_external_pickle(POLICY_PKL_PATH)
self.vision_run = pickle.loads(read_file_chunked(str(VISION_PKL_PATH)))
self.policy_run = pickle.loads(read_file_chunked(str(POLICY_PKL_PATH)))
def slice_outputs(self, model_outputs: np.ndarray, output_slices: dict[str, slice]) -> dict[str, np.ndarray]:
parsed_model_outputs = {k: model_outputs[np.newaxis, v] for k,v in output_slices.items()}