Support multipart model artifacts

This commit is contained in:
firestar5683
2026-06-23 00:17:15 -05:00
parent d97100bd14
commit 6ee00d3f60
6 changed files with 390 additions and 34 deletions
+95
View File
@@ -1,6 +1,7 @@
#!/usr/bin/env python3
import argparse
import codecs
import hashlib
import json
import os
import pickle
@@ -36,6 +37,8 @@ MEDMODEL_INPUT_SIZE = (512, 256)
DM_INPUT_SIZE = (1440, 960)
MODEL_RUN_FREQ = 20
MODEL_CONTEXT_FREQ = 5
REPOSITORY_FILE_LIMIT = 100 * 1024 * 1024
DEFAULT_MULTIPART_SIZE = 95 * 1024 * 1024
def build_compile_env() -> dict[str, str]:
@@ -77,6 +80,8 @@ def parse_args() -> argparse.Namespace:
)
parser.add_argument("--list", action="store_true", help="List staged models and exit.")
parser.add_argument("--force", action="store_true", help="Accepted for compatibility; selected outputs are always replaced.")
parser.add_argument("--split-artifact", type=Path, help="Split an existing oversized PKL without compiling.")
parser.add_argument("--chunk-size-mib", type=int, default=95, help="Multipart size in MiB; must be below 100.")
args, unknown = parser.parse_known_args()
dynamic_flags = [value[2:] for value in unknown if value.startswith("--")]
@@ -93,6 +98,10 @@ def parse_args() -> argparse.Namespace:
args.model = None
if args.dm and args.model:
parser.error("Use either --dm or a driving model ID.")
if args.split_artifact and (args.dm or args.model):
parser.error("--split-artifact cannot be combined with --dm or a model ID.")
if not 1 <= args.chunk_size_mib < 100:
parser.error("--chunk-size-mib must be between 1 and 99.")
return args
@@ -264,11 +273,81 @@ def remove_paths(paths: list[Path]) -> int:
return count
def sha256_file(path: Path) -> str:
digest = hashlib.sha256()
with open(path, "rb") as artifact_file:
for chunk in iter(lambda: artifact_file.read(1024 * 1024), b""):
digest.update(chunk)
return digest.hexdigest()
def multipart_output_paths(artifact: Path, output_dir: Path | None = None) -> list[Path]:
output_dir = output_dir or artifact.parent
return [
*sorted(output_dir.glob(f"{artifact.name}.p[0-9][0-9]")),
output_dir / f"{artifact.name}.sha256",
]
def split_oversized_artifact(
artifact: Path,
output_dir: Path | None = None,
chunk_size: int = DEFAULT_MULTIPART_SIZE,
force: bool = False,
) -> list[Path]:
artifact = artifact.resolve()
output_dir = (output_dir or artifact.parent).resolve()
if not artifact.is_file():
raise FileNotFoundError(artifact)
if chunk_size <= 0 or chunk_size >= REPOSITORY_FILE_LIMIT:
raise ValueError("Multipart chunk size must be between 1 byte and 100 MiB.")
remove_paths(multipart_output_paths(artifact, output_dir))
if artifact.stat().st_size <= REPOSITORY_FILE_LIMIT and not force:
return []
output_dir.mkdir(parents=True, exist_ok=True)
digest = hashlib.sha256()
part_paths: list[Path] = []
with open(artifact, "rb") as source:
for index in range(100):
part_path = output_dir / f"{artifact.name}.p{index:02d}"
part_size = 0
with open(part_path, "wb") as part_file:
while part_size < chunk_size:
chunk = source.read(min(1024 * 1024, chunk_size - part_size))
if not chunk:
break
part_file.write(chunk)
digest.update(chunk)
part_size += len(chunk)
if part_size == 0:
part_path.unlink()
break
part_paths.append(part_path)
if not part_paths:
raise ValueError(f"Artifact is empty: {artifact}")
checksum_path = output_dir / f"{artifact.name}.sha256"
checksum_path.write_text(f"{digest.hexdigest()} {artifact.name}\n")
verify_digest = hashlib.sha256()
for part_path in part_paths:
with open(part_path, "rb") as part_file:
for chunk in iter(lambda: part_file.read(1024 * 1024), b""):
verify_digest.update(chunk)
if verify_digest.hexdigest() != digest.hexdigest():
remove_paths([*part_paths, checksum_path])
raise RuntimeError("Split artifact failed checksum verification.")
return [*part_paths, checksum_path]
def compile_driving(model_key: str, files: dict[str, Path], input_format: str, version: str, output_dir: Path) -> Path:
model_type, source_args = driving_compile_args(files, input_format)
output_path = output_dir / f"{model_key}_driving_tinygrad.pkl"
removed = remove_paths([
output_path,
*multipart_output_paths(output_path, output_dir),
*output_dir.glob(f"{model_key}_driving_*_tinygrad.pkl"),
*output_dir.glob(f"{model_key}_driving_*_metadata.pkl"),
])
@@ -349,6 +428,17 @@ def list_models(staged: dict[str, dict[str, Path]], input_root: Path) -> int:
def main() -> int:
args = parse_args()
if args.split_artifact:
outputs = split_oversized_artifact(
args.split_artifact,
args.output_dir,
args.chunk_size_mib * 1024 * 1024,
force=True,
)
for output in outputs:
print(f" saved {output.name} ({output.stat().st_size} bytes)")
return 0
staged = find_staged_models(args.input_dir)
if args.list:
return list_models(staged, args.input_dir)
@@ -379,6 +469,11 @@ def main() -> int:
print(f"Compiling {model_key} ({input_format}, {version_label}) from {args.input_dir} -> {args.output_dir}")
output = compile_driving(model_key, files, input_format, version, args.output_dir)
print(f" saved {output.name}")
multipart_outputs = split_oversized_artifact(output)
if multipart_outputs:
print(" artifact exceeds 100 MiB; created repository-safe multipart files:")
for multipart_output in multipart_outputs:
print(f" {multipart_output.name} ({multipart_output.stat().st_size} bytes)")
print("Done.")
return 0
+28 -13
View File
@@ -11,6 +11,11 @@ import sys
from pathlib import Path
REPO_ROOT = Path(__file__).resolve().parents[1]
if str(REPO_ROOT / "scripts") not in sys.path:
sys.path.insert(0, str(REPO_ROOT / "scripts"))
from model_compiler import split_oversized_artifact
DEFAULT_OPENPILOT = Path.home() / "openpilot"
DEFAULT_WORKSPACE = Path("/Volumes/T5/StarPilot-Model-Rebuild-2026-06-22")
DEFAULT_SOURCE_MAP = REPO_ROOT / "scripts/model_source_map_v22.json"
@@ -204,9 +209,22 @@ def remote(command: str, *, capture: bool = False):
return result
def stage_ready_artifact(artifact: Path, workspace: Path) -> None:
ready_path = workspace / "ready-for-resources" / artifact.name
multipart_outputs = split_oversized_artifact(artifact, ready_path.parent)
if multipart_outputs:
ready_path.unlink(missing_ok=True)
for multipart_output in multipart_outputs:
multipart_output.chmod(0o644)
else:
shutil.copyfile(artifact, ready_path)
ready_path.chmod(0o644)
def compile_model(model_id: str, source: dict, version: str, workspace: Path, force: bool) -> dict:
local_output = workspace / "compiled" / f"{model_id}_driving_tinygrad.pkl"
if local_output.is_file() and not force:
stage_ready_artifact(local_output, workspace)
return artifact_result(model_id, local_output, "skipped")
source_dir = workspace / "onnx" / model_id
@@ -230,13 +248,7 @@ def compile_model(model_id: str, source: dict, version: str, workspace: Path, fo
run(["rsync", "-az", f"{REMOTE}:{remote_output}", str(local_output)])
local_output.chmod(0o644)
ready_path = workspace / "ready-for-resources" / local_output.name
shutil.copyfile(local_output, ready_path)
ready_path.chmod(0o644)
if local_output.stat().st_size > 100 * 1024 * 1024:
external_path = workspace / "external-upload" / local_output.name
shutil.copyfile(local_output, external_path)
external_path.chmod(0o644)
stage_ready_artifact(local_output, workspace)
result = artifact_result(model_id, local_output, "compiled")
(workspace / "results" / f"{model_id}_artifact.json").write_text(json.dumps(result, indent=2) + "\n")
return result
@@ -249,7 +261,7 @@ def artifact_result(model_id: str, path: Path, status: str) -> dict:
"path": str(path),
"size": path.stat().st_size,
"sha256": sha256_file(path),
"external_upload": path.stat().st_size > 100 * 1024 * 1024,
"multipart": path.stat().st_size > 100 * 1024 * 1024,
}
@@ -288,7 +300,7 @@ def update_manifest(base_manifest: Path, workspace: Path) -> dict:
"released": "2026-06-17",
"community_favorite": False,
})
external_handoff = []
multipart_handoff = []
for model in models:
artifact = workspace / "compiled" / f"{model['id']}_driving_tinygrad.pkl"
model.pop("artifact_format", None)
@@ -298,19 +310,22 @@ def update_manifest(base_manifest: Path, workspace: Path) -> dict:
if not artifact.is_file() or artifact.stat().st_size <= 100 * 1024 * 1024:
model.pop("artifact_url", None)
else:
external_handoff.append({
multipart_handoff.append({
"id": model["id"],
"filename": artifact.name,
"size": artifact.stat().st_size,
"sha256": sha256_file(artifact),
"artifact_url": model.get("artifact_url", ""),
"parts": [
path.name
for path in sorted((workspace / "ready-for-resources").glob(f"{artifact.name}.p[0-9][0-9]"))
],
})
output = {"models": models}
output_path = workspace / "manifests/model_names_v22.json"
output_path.parent.mkdir(parents=True, exist_ok=True)
output_path.write_text(json.dumps(output, indent=2, ensure_ascii=False) + "\n")
(workspace / "external-upload" / "handoff.json").write_text(
json.dumps(external_handoff, indent=2) + "\n",
(workspace / "ready-for-resources" / "multipart.json").write_text(
json.dumps(multipart_handoff, indent=2) + "\n",
)
return output