chunk tinygrad pkl below GitHub max size - NoCache and AlwaysBuild (#37194)

* nocache

* +

* fixes

* lint

* not split

* use pathlib

* cleanup

* better

* even better
This commit is contained in:
YassineYousfi
2026-02-13 10:14:24 -08:00
committed by GitHub
parent 9b7bf4a101
commit 2ba6df2506
6 changed files with 65 additions and 12 deletions
+20 -3
View File
@@ -3,6 +3,7 @@ import glob
Import('env', 'arch')
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)
@@ -36,12 +37,28 @@ lenv.Command(warp_targets, tinygrad_files + script_files, cmd)
def tg_compile(flags, model_name):
pythonpath_string = 'PYTHONPATH="${PYTHONPATH}:' + env.Dir("#tinygrad_repo").abspath + '"'
fn = File(f"models/{model_name}").abspath
return lenv.Command(
fn + "_tinygrad.pkl",
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 {fn}_tinygrad.pkl'
f'{pythonpath_string} {flags} {image_flag} python3 {Dir("#tinygrad_repo").abspath}/examples/openpilot/compile3.py {fn}.onnx {full}'
)
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)