"""Offline graph-partition and release-integrity tests using small toy graphs.""" import importlib.util import json from pathlib import Path import pytest onnx = pytest.importorskip("onnx") np = pytest.importorskip("numpy") spec = importlib.util.spec_from_file_location( "vons_shared_export", Path(__file__).parents[1] / "tools/export_shared_bundle.py", ) assert spec is not None and spec.loader is not None exporter = importlib.util.module_from_spec(spec) spec.loader.exec_module(exporter) digest = exporter.digest export_bundle = exporter.export_bundle partition_graphs = exporter.partition_graphs verify_manifest = exporter.verify_manifest def toy_model(backend: str, *, weight: float = 1.0): helper = onnx.helper tensor = onnx.TensorProto inputs = [helper.make_tensor_value_info(name, tensor.INT64, ["batch", "options", 2]) for name in ("input_ids", "attention_mask", "token_type_ids")] inputs.append(helper.make_tensor_value_info("option_mask", tensor.BOOL, ["batch", "options"])) values = [helper.make_tensor_value_info("view_19", tensor.FLOAT, ["batch", "options", 2]), helper.make_tensor_value_info("div", tensor.FLOAT, ["batch", 2])] weights = [onnx.numpy_helper.from_array(np.eye(2, dtype=np.float32) * weight, name="encoder.weight"), onnx.numpy_helper.from_array(np.array([1], np.int64), name="axis1"), onnx.numpy_helper.from_array(np.array([2], np.int64), name="axis2")] nodes = [helper.make_node("Cast", ["input_ids"], ["cast"], to=tensor.FLOAT), helper.make_node("MatMul", ["cast", "encoder.weight"], ["view_19"]), helper.make_node("ReduceMean", ["view_19", "axis1"], ["div"], keepdims=0), helper.make_node("ReduceSum", ["div", "axis1"], ["answerability"], keepdims=0)] output = "logits" if backend == "direct" else "scores" if backend == "direct": nodes.append(helper.make_node("ReduceSum", ["view_19", "axis2"], [output], keepdims=0)) else: inputs.append(helper.make_tensor_value_info("initial_noise", tensor.FLOAT, ["batch", "options"])) nodes.extend([ helper.make_node("Shape", ["token_type_ids"], ["batch_shape"], start=0, end=1), helper.make_node("Expand", ["answerability", "batch_shape"], ["expanded"]), helper.make_node("Unsqueeze", ["expanded", "axis1"], ["scalar"]), helper.make_node("Add", ["initial_noise", "scalar"], [output]), ]) outputs = [helper.make_tensor_value_info(output, tensor.FLOAT, ["batch", "options"]), helper.make_tensor_value_info("answerability", tensor.FLOAT, ["batch"])] model = helper.make_model(helper.make_graph(nodes, "toy", inputs, outputs, weights, value_info=values), opset_imports=[helper.make_opsetid("", 18)]) model.ir_version = 10 return model def test_partition_keeps_outputs_and_removes_token_dependency(): parts = partition_graphs(toy_model("direct"), toy_model("diffusion")) for model in parts.values(): onnx.checker.check_model(model) assert [x.name for x in parts["encoder"].graph.output] == ["candidate_embeddings", "pooled"] assert all("token_type_ids" not in n.input for n in parts["diffusion"].graph.node) def test_different_encoders_cannot_be_silently_shared(): with pytest.raises(ValueError, match="parameter tensors differ"): partition_graphs(toy_model("direct"), toy_model("diffusion", weight=2)) def test_partition_matches_full_graph_with_nonzero_noise(): ort = pytest.importorskip("onnxruntime") original = {name: toy_model(name) for name in ("direct", "diffusion")} parts = partition_graphs(original["direct"], original["diffusion"]) options = ort.SessionOptions() options.intra_op_num_threads = 1 sessions = {key: ort.InferenceSession(value.SerializeToString(), options, providers=["CPUExecutionProvider"]) for key, value in parts.items()} feed = {name: np.array([[[1, 2], [3, 4], [5, 6]]], np.int64) for name in ("input_ids", "attention_mask", "token_type_ids")} feed["option_mask"] = np.array([[True, True, True]]) encoded = sessions["encoder"].run(None, feed) for backend, model in original.items(): shared = {"candidate_embeddings": encoded[0], "pooled": encoded[1], "option_mask": feed["option_mask"]} full_feed = dict(feed) if backend == "diffusion": shared.pop("candidate_embeddings") shared["initial_noise"] = np.array([[0.2, -0.7, 1.1]], np.float32) full_feed["initial_noise"] = shared["initial_noise"] full = ort.InferenceSession(model.SerializeToString(), options, providers=["CPUExecutionProvider"]).run(None, full_feed) split = sessions[backend].run(None, shared) for expected, actual in zip(full, split, strict=True): np.testing.assert_allclose(actual, expected, atol=1e-6) def make_bundle(tmp_path: Path): for backend in ("direct", "diffusion"): onnx.save(toy_model(backend), tmp_path / f"{backend}.onnx") tokenizer = tmp_path / "tokenizer" tokenizer.mkdir() (tokenizer / "tokenizer.json").write_text("{}") (tokenizer / "tokenizer_config.json").write_text("{}") target = tmp_path / "bundle" manifest = export_bundle(tmp_path / "direct.onnx", tmp_path / "diffusion.onnx", tokenizer, target) return target, manifest def test_asset_count_includes_config_calibration_and_both_heads(tmp_path): target, manifest = make_bundle(tmp_path) assert {r["role"] for r in manifest["files"]} >= { "encoder_graph", "direct_graph", "diffusion_graph", "calibration", "config", "tokenizer", } assert manifest["model_asset_bytes"] == sum((target / r["path"]).stat().st_size for r in manifest["files"]) assert verify_manifest(target / "manifest.json")["pass"] with pytest.raises(ValueError, match="new or empty"): export_bundle(tmp_path / "direct.onnx", tmp_path / "diffusion.onnx", tmp_path / "tokenizer", target) def test_tamper_and_external_expected_digest_are_rejected(tmp_path): target, _ = make_bundle(tmp_path) path = target / "manifest.json" with pytest.raises(ValueError, match="release manifest digest"): verify_manifest(path, expected_digest="0" * 64) (target / "config.json").write_text("tampered") with pytest.raises(ValueError, match="asset digest"): verify_manifest(path) def test_referenced_assets_cannot_be_omitted(tmp_path): target, manifest = make_bundle(tmp_path) manifest["files"] = [r for r in manifest["files"] if r["path"] != "calibration.json"] manifest["model_asset_bytes"] = sum(r["bytes"] for r in manifest["files"]) path = target / "manifest.json" path.write_text(json.dumps(manifest)) path.with_suffix(".sha256").write_text(digest(path)) with pytest.raises(ValueError, match="omits required"): verify_manifest(path)