Skip to content

Commit bf0f190

Browse files
authored
fix: verify exported inference artifacts (#317)
1 parent 0475e30 commit bf0f190

8 files changed

Lines changed: 863 additions & 36 deletions

File tree

Lines changed: 42 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,42 @@
1+
name: Verify artifact execution receipt
2+
3+
on:
4+
push:
5+
branches: [main]
6+
paths:
7+
- "src/tether/runtime/verify_inference.py"
8+
- "src/tether/verify.py"
9+
- "src/tether/smoke.py"
10+
- "scripts/verify_artifact_receipt.py"
11+
- ".github/workflows/verify-artifact-receipt.yml"
12+
workflow_dispatch:
13+
14+
permissions:
15+
contents: read
16+
17+
jobs:
18+
execute-export:
19+
if: >-
20+
github.ref == 'refs/heads/main' &&
21+
(github.event_name == 'push' || github.event_name == 'workflow_dispatch')
22+
runs-on: ubuntu-latest
23+
steps:
24+
- uses: actions/checkout@v4
25+
- uses: actions/setup-python@v5
26+
with:
27+
python-version: "3.12"
28+
cache: pip
29+
cache-dependency-path: pyproject.toml
30+
- name: Install runtime
31+
run: |
32+
python -m pip install --upgrade pip
33+
pip install -e ".[onnx,serve]"
34+
- name: Execute generated export
35+
run: python scripts/verify_artifact_receipt.py --output verify-artifact-receipt.json
36+
- name: Upload provenance receipt
37+
uses: actions/upload-artifact@v4
38+
with:
39+
name: verify-artifact-receipt-${{ github.sha }}
40+
path: verify-artifact-receipt.json
41+
if-no-files-found: error
42+
retention-days: 90

scripts/verify_artifact_receipt.py

Lines changed: 103 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,103 @@
1+
#!/usr/bin/env python3
2+
"""Execute a real generated export and write a provenance receipt for #267."""
3+
4+
from __future__ import annotations
5+
6+
import argparse
7+
from datetime import datetime, timezone
8+
import hashlib
9+
import json
10+
import os
11+
from pathlib import Path
12+
import platform
13+
import tempfile
14+
15+
import numpy as np
16+
17+
from tether.runtime.verify_inference import load_verification_inference
18+
from tether.smoke import create_smoke_export
19+
20+
21+
def _sha256(path: Path) -> str:
22+
digest = hashlib.sha256()
23+
with path.open("rb") as stream:
24+
for chunk in iter(lambda: stream.read(1024 * 1024), b""):
25+
digest.update(chunk)
26+
return digest.hexdigest()
27+
28+
29+
def _validate_smoke_actions(actions: np.ndarray, noise: np.ndarray) -> None:
30+
"""Fail the workflow unless the generated Constant export truly executed."""
31+
expected = np.zeros_like(noise)
32+
if not np.array_equal(actions, expected):
33+
raise RuntimeError("smoke export actions do not match its declared constant graph")
34+
35+
36+
def build_receipt(output: Path) -> dict[str, object]:
37+
event_name = os.environ.get("GITHUB_EVENT_NAME")
38+
ref = os.environ.get("GITHUB_REF")
39+
if os.environ.get("GITHUB_ACTIONS") == "true":
40+
if event_name not in {"push", "workflow_dispatch"}:
41+
raise RuntimeError(f"receipt event is not allowlisted: {event_name!r}")
42+
if ref != "refs/heads/main":
43+
raise RuntimeError(f"receipt must execute from protected main, got {ref!r}")
44+
if os.environ.get("GITHUB_REF_PROTECTED") != "true":
45+
raise RuntimeError("main is not reported as a protected ref")
46+
with tempfile.TemporaryDirectory(prefix="tether-verify-artifact-") as temporary:
47+
export_dir = create_smoke_export(Path(temporary) / "export")
48+
inference = load_verification_inference(export_dir, device="cpu")
49+
noise = np.arange(50 * 32, dtype=np.float32).reshape(1, 50, 32)
50+
image = np.zeros((1, 3, 512, 512), dtype=np.float32)
51+
mask = np.ones((1,), dtype=np.bool_)
52+
actions = inference.predict_action_chunk(
53+
img_base=image,
54+
img_wrist_l=image,
55+
img_wrist_r=image,
56+
mask_base=mask,
57+
mask_wrist_l=mask,
58+
mask_wrist_r=mask,
59+
lang_tokens=np.zeros((1, 16), dtype=np.int64),
60+
lang_masks=np.ones((1, 16), dtype=np.bool_),
61+
noise=noise,
62+
state=np.zeros((1, 32), dtype=np.float32),
63+
episode_id="receipt",
64+
)
65+
if actions.shape != noise.shape:
66+
raise RuntimeError(f"unexpected action shape {actions.shape}; expected {noise.shape}")
67+
_validate_smoke_actions(actions, noise)
68+
receipt: dict[str, object] = {
69+
"schema_version": 1,
70+
"kind": "tether.verify_artifact_execution",
71+
"source_sha": os.environ.get("GITHUB_SHA", "local"),
72+
"workflow": os.environ.get("GITHUB_WORKFLOW", "local"),
73+
"workflow_ref": os.environ.get("GITHUB_WORKFLOW_REF", "local"),
74+
"repository": os.environ.get("GITHUB_REPOSITORY", "local"),
75+
"event_name": event_name or "local",
76+
"ref": ref or "local",
77+
"run_id": os.environ.get("GITHUB_RUN_ID", "local"),
78+
"run_attempt": os.environ.get("GITHUB_RUN_ATTEMPT", "local"),
79+
"generated_at": datetime.now(timezone.utc).isoformat(),
80+
"model_type": "smolvla",
81+
"export_kind": "monolithic_onnx",
82+
"backend": inference.get_stats().get("backend"),
83+
"model_sha256": _sha256(export_dir / "model.onnx"),
84+
"config_sha256": _sha256(export_dir / "tether_config.json"),
85+
"actions_sha256": hashlib.sha256(actions.tobytes()).hexdigest(),
86+
"python": platform.python_version(),
87+
"passed": True,
88+
}
89+
output.parent.mkdir(parents=True, exist_ok=True)
90+
output.write_text(json.dumps(receipt, indent=2, sort_keys=True) + "\n")
91+
return receipt
92+
93+
94+
def main() -> int:
95+
parser = argparse.ArgumentParser()
96+
parser.add_argument("--output", type=Path, default=Path("verify-artifact-receipt.json"))
97+
args = parser.parse_args()
98+
print(json.dumps(build_receipt(args.output), indent=2, sort_keys=True))
99+
return 0
100+
101+
102+
if __name__ == "__main__":
103+
raise SystemExit(main())

src/tether/eval/libero_rollout.py

Lines changed: 32 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -141,7 +141,6 @@ def run_libero_rollout(
141141
# Lazy imports — LIBERO + mujoco only needed at rollout time, not at module load.
142142
import collections
143143
import math
144-
import time
145144
import traceback
146145
from pathlib import Path
147146

@@ -450,12 +449,38 @@ def _build_batch(obs, task_description):
450449
return results
451450

452451

452+
def _load_exact_reference_policy(model_type: str, checkpoint: str) -> Any:
453+
"""Load the declared native family from the exact export provenance ref."""
454+
import importlib
455+
456+
targets = {
457+
"pi05": ("lerobot.policies.pi05.modeling_pi05", "PI05Policy"),
458+
"pi05_decomposed": ("lerobot.policies.pi05.modeling_pi05", "PI05Policy"),
459+
"pi0": ("lerobot.policies.pi0.modeling_pi0", "PI0Policy"),
460+
"smolvla": ("lerobot.policies.smolvla.modeling_smolvla", "SmolVLAPolicy"),
461+
}
462+
if model_type == "pi05_decomposed_student":
463+
module = importlib.import_module("tether.distill.snapflow_pi0_model")
464+
return module.load_snapflow_student(checkpoint)
465+
try:
466+
module_name, class_name = targets[model_type]
467+
except KeyError as exc:
468+
raise ValueError(
469+
f"unsupported native verification model_type={model_type!r}"
470+
) from exc
471+
policy_class = getattr(importlib.import_module(module_name), class_name)
472+
return policy_class.from_pretrained(checkpoint)
473+
474+
453475
def load_pi05_policy_and_processors(
454476
*,
455477
student_checkpoint: str,
456478
decomposed_dir: str,
457479
preprocessor_ref: str | None = None,
458480
force_teacher: bool = False,
481+
model_type: str = "pi05",
482+
require_exact_checkpoint: bool = False,
483+
device: str = "cuda",
459484
) -> tuple[Any, Any, Any]:
460485
"""Load PyTorch policy (for config + _preprocess_images) + processor pipelines.
461486
@@ -480,7 +505,10 @@ def load_pi05_policy_and_processors(
480505
)
481506

482507
student_ckpt_path = Path(student_checkpoint)
483-
if not force_teacher and (student_ckpt_path / "model.safetensors").exists():
508+
if require_exact_checkpoint:
509+
print(f"[load] Loading exact {model_type} reference from {student_checkpoint}")
510+
policy = _load_exact_reference_policy(model_type, student_checkpoint)
511+
elif not force_teacher and (student_ckpt_path / "model.safetensors").exists():
484512
print(f"[load] Loading SnapFlow student from {student_checkpoint}")
485513
from tether.distill.snapflow_pi0_model import load_snapflow_student
486514
policy = load_snapflow_student(student_checkpoint)
@@ -497,7 +525,7 @@ def load_pi05_policy_and_processors(
497525
f"inference still runs through decomposed ONNX)"
498526
)
499527
policy = PI05Policy.from_pretrained(fallback)
500-
policy.eval().to("cuda").to(torch.float32)
528+
policy.eval().to(device).to(torch.float32)
501529

502530
# Student-distillation checkpoints don't always ship the processor JSONs —
503531
# fall back to the teacher HF repo for baseline preprocessor + normalizer.
@@ -511,7 +539,7 @@ def load_pi05_policy_and_processors(
511539
config_filename="policy_preprocessor.json",
512540
to_transition=batch_to_transition,
513541
to_output=transition_to_batch,
514-
overrides={"device_processor": {"device": "cuda"}},
542+
overrides={"device_processor": {"device": device}},
515543
)
516544
postprocessor = PolicyProcessorPipeline.from_pretrained(
517545
pretrained_model_name_or_path=proc_ref,
@@ -532,7 +560,6 @@ def load_pi05_policy_and_processors(
532560
)
533561
if is_state_out_export:
534562
from tether.distill.pi05_state_out_processor import swap_prepare_step_in_pipeline
535-
from lerobot.utils.constants import ACTION
536563
max_state_dim = policy.config.max_action_dim # pi0.5: 32
537564
swap_prepare_step_in_pipeline(preprocessor, max_state_dim=max_state_dim)
538565
print(

src/tether/runtime/pi0_onnx_server.py

Lines changed: 10 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -177,6 +177,7 @@ def predict(
177177
noise: np.ndarray | None = None,
178178
lang_tokens: np.ndarray | None = None,
179179
lang_masks: np.ndarray | None = None,
180+
image_masks: list[np.ndarray] | None = None,
180181
) -> dict[str, Any]:
181182
"""Run one pi0 forward pass.
182183
@@ -220,7 +221,12 @@ def _prep_img(img: np.ndarray) -> np.ndarray:
220221
img_wrist_l = _prep_img(images_list[1])
221222
img_wrist_r = _prep_img(images_list[2])
222223

223-
mask = np.ones((1,), dtype=np.bool_)
224+
masks_list = list(image_masks or [])
225+
while len(masks_list) < 3:
226+
masks_list.append(np.ones((1,), dtype=np.bool_))
227+
mask_base, mask_wrist_l, mask_wrist_r = [
228+
np.asarray(value, dtype=np.bool_).reshape(-1) for value in masks_list[:3]
229+
]
224230

225231
# Lang: take externally-supplied tokens or tokenize the instruction
226232
if lang_tokens is None:
@@ -260,9 +266,9 @@ def _prep_img(img: np.ndarray) -> np.ndarray:
260266
"img_base": img_base,
261267
"img_wrist_l": img_wrist_l,
262268
"img_wrist_r": img_wrist_r,
263-
"mask_base": mask,
264-
"mask_wrist_l": mask,
265-
"mask_wrist_r": mask,
269+
"mask_base": mask_base,
270+
"mask_wrist_l": mask_wrist_l,
271+
"mask_wrist_r": mask_wrist_r,
266272
"lang_tokens": lang_tokens,
267273
"lang_masks": lang_masks,
268274
"state": state_arr,

src/tether/runtime/smolvla_onnx_server.py

Lines changed: 10 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -175,6 +175,7 @@ def predict(
175175
noise: np.ndarray | None = None,
176176
lang_tokens: np.ndarray | None = None,
177177
lang_masks: np.ndarray | None = None,
178+
image_masks: list[np.ndarray] | None = None,
178179
) -> dict[str, Any]:
179180
"""Run one SmolVLA forward pass. Accepts a single image or a list of 3."""
180181
if not self._ready:
@@ -207,7 +208,12 @@ def _prep_img(img: np.ndarray) -> np.ndarray:
207208
img_cam2 = _prep_img(images_list[1])
208209
img_cam3 = _prep_img(images_list[2])
209210

210-
mask = np.ones((1,), dtype=np.bool_)
211+
masks_list = list(image_masks or [])
212+
while len(masks_list) < 3:
213+
masks_list.append(np.ones((1,), dtype=np.bool_))
214+
mask_cam1, mask_cam2, mask_cam3 = [
215+
np.asarray(value, dtype=np.bool_).reshape(-1) for value in masks_list[:3]
216+
]
211217

212218
# Lang: tokenize (SmolLM2 vocab ~49152) or use provided tokens
213219
if lang_tokens is None:
@@ -251,9 +257,9 @@ def _prep_img(img: np.ndarray) -> np.ndarray:
251257
"img_cam1": img_cam1,
252258
"img_cam2": img_cam2,
253259
"img_cam3": img_cam3,
254-
"mask_cam1": mask,
255-
"mask_cam2": mask,
256-
"mask_cam3": mask,
260+
"mask_cam1": mask_cam1,
261+
"mask_cam2": mask_cam2,
262+
"mask_cam3": mask_cam3,
257263
"lang_tokens": lang_tokens,
258264
"lang_masks": lang_masks,
259265
"state": state_arr,

0 commit comments

Comments
 (0)