wai-neural-e2e: make the job link and run, and fix the decoder mismatch it was hiding #101

Merged
dcharlot merged 3 commits from ci/neural-e2e-guard into main 2026-09-30 19:50:20 -04:00
2 changed files with 93 additions and 14 deletions

View file

@ -30,6 +30,20 @@ jobs:
neural-e2e:
name: synthetic fixtures → Rust decode → reference (±4)
runs-on: linux
# The Rust decoders link `ort`, whose `download-binaries` feature pulls a
# prebuilt static onnxruntime (1.28.0 for ort 2.0.0-rc.13). That library was
# built against glibc >= 2.38 and GCC 13's libstdc++. The runner's default
# job image, catthehacker/ubuntu:act-22.04, has glibc 2.35 and GCC 11, so the
# test binary fails to LINK with five undefined symbols:
# __isoc23_strtol __isoc23_strtoll __isoc23_strtoull (glibc 2.38+)
# basic_string<char>::_M_replace_cold basic_string<wchar_t>::_M_replace_cold
# (libstdc++, GCC 13+)
# That link error is why this job failed at ~1 minute on every run: the
# Python steps pass on a warm cache and the job dies at its last step. Checked
# in both images: act-22.04 lacks all five, act-24.04 (glibc 2.39, g++ 13)
# provides all five.
container:
image: catthehacker/ubuntu:act-24.04
steps:
- uses: actions/checkout@v5
- uses: actions/setup-python@v5
@ -60,4 +74,32 @@ jobs:
working-directory: wai/wai-rs
# --no-default-features: the neural decoders need only `ort`, NOT the C
# codec chain (AV1/JXL/opus), so the job needs no system codec deps.
run: cargo test --no-default-features --features neural --test neural_native_decode -- _matches_reference --nocapture
#
# The grep is the assertion, not decoration. Every `*_matches_reference`
# test opens with `if !wai.exists() || !ref.exists() || !model.exists() {
# eprintln!("skipping ..."); return; }` — so a missing fixture makes the
# test SKIP, and a skip is reported as a PASS. Verified by running this
# exact command on a machine with no /tmp fixtures: it prints
# "6 passed; 0 failed" while comparing nothing.
#
# That means a green `cargo test` here does NOT mean the decoders were
# checked against their references. If the Python fixture step ever
# succeeds while writing nothing, this job goes green having verified
# nothing at all — the failure mode is silent and looks like success.
# Failing on "skipping" is what makes the job's name true.
#
# One test is excluded by name, not tolerated as a skip:
# `snac_real_model_matches_reference` needs real pretrained SNAC weights,
# which this synthetic-fixture job never builds, so it would skip on every
# run and the guard would fail a job whose five in-scope decoders all
# match. Excluding it keeps the guard strict for everything in scope.
run: |
set -o pipefail
cargo test --no-default-features --features neural --test neural_native_decode \
-- _matches_reference --skip snac_real_model_matches_reference --nocapture \
2>&1 | tee /tmp/neural_e2e.log
if grep -q "skipping" /tmp/neural_e2e.log; then
echo "::error::a reference comparison was SKIPPED — its fixtures are missing, so this job verified nothing"
grep "skipping" /tmp/neural_e2e.log
exit 1
fi

View file

@ -120,6 +120,31 @@ def write_wai(path, cap, kind, header, z, y, H, W):
path.write_bytes(b"WAI1" + struct.pack("<I", len(mb)) + mb + struct.pack("<I", len(payload)) + payload)
def _session(path):
"""Open an ONNX model exactly as the Rust decoder does.
`wai-rs/src/ort_session.rs::load_deterministic` evaluates every model with one
intra-op thread, one inter-op thread and full graph optimisation, because float
ONNX inference is only reproducible that way. The fixtures must be ENCODED under
the same settings. With onnxruntime's defaults (every core) a convolution's
reduction is split across threads, the summation order changes and the float
output shifts. Where a scale then straddles a threshold of the Gaussian scale
table, encoder and decoder select adjacent CDFs, and the arithmetic-coded stream
desyncs from that symbol onward.
Measured, not hypothetical: under default sessions the MLIC++ fixture diverged at
symbol 416 of 12,288 -- encoded with scale 0.78787208 against CDF 16, decoded
with 0.78788924 against CDF 17, either side of scale_table[16] = 0.78788900 --
and decoded to max|delta|=255. Under these settings it decodes exactly, and
DCVC's residual max|delta|=1 goes to 0.
"""
so = ort.SessionOptions()
so.intra_op_num_threads = 1
so.inter_op_num_threads = 1
so.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
return ort.InferenceSession(str(path), sess_options=so, providers=["CPUExecutionProvider"])
def onx(sess, t):
return torch.from_numpy(sess.run(None, {sess.get_inputs()[0].name: t.numpy().astype(np.float32)})[0])
@ -132,9 +157,9 @@ def roundtrip_elic(net, x, mdir, size):
export_onnx(net.g_s, torch.zeros(1, 48, 16, 16), mdir / "synthesis.onnx")
(mdir / "hyper_cdfs_q3.json").write_text(json.dumps(fblob(net.entropy_bottleneck), separators=(",", ":")))
(mdir / "gauss_cdfs_q3.json").write_text(json.dumps(gblob(gc), separators=(",", ":")))
hs = ort.InferenceSession(str(mdir / "hyper_synthesis.onnx"))
cx = ort.InferenceSession(str(mdir / "context.onnx"))
gs = ort.InferenceSession(str(mdir / "synthesis.onnx"))
hs = _session(mdir / "hyper_synthesis.onnx")
cx = _session(mdir / "context.onnx")
gs = _session(mdir / "synthesis.onnx")
with torch.no_grad():
y = net.g_a(x); z = net.h_a(y)
zs = net.entropy_bottleneck.compress(z); z_hat = net.entropy_bottleneck.decompress(zs, z.size()[-2:])
@ -174,9 +199,9 @@ def roundtrip_mlicpp(net, x, mdir, size, K=4):
export_onnx(net.g_s, torch.zeros(1, 48, 16, 16), mdir / "synthesis.onnx")
(mdir / "hyper_cdfs_q3.json").write_text(json.dumps(fblob(net.entropy_bottleneck), separators=(",", ":")))
(mdir / "gauss_cdfs_q3.json").write_text(json.dumps(gblob(gc), separators=(",", ":")))
hs = ort.InferenceSession(str(mdir / "hyper_synthesis.onnx"))
en = ort.InferenceSession(str(mdir / "entropy.onnx"))
gs = ort.InferenceSession(str(mdir / "synthesis.onnx"))
hs = _session(mdir / "hyper_synthesis.onnx")
en = _session(mdir / "entropy.onnx")
gs = _session(mdir / "synthesis.onnx")
with torch.no_grad():
y = net.g_a(x); z = net.h_a(y)
zs = net.entropy_bottleneck.compress(z); z_hat = net.entropy_bottleneck.decompress(zs, z.size()[-2:])
@ -230,7 +255,7 @@ def roundtrip_snac(mdir: Path):
input_names=names, output_names=["audio_values"],
dynamic_axes={**{f"codes_{i}": {1: f"t{i}"} for i in range(3)},
"audio_values": {2: "s"}}, opset_version=17, dynamo=False)
sess = ort.InferenceSession(str(mdir / "decoder.onnx"))
sess = _session(mdir / "decoder.onnx")
ref = sess.run(None, {f"codes_{i}": codes[i].numpy() for i in range(3)})[0]
hdr = struct.pack("<III", sr, nsamp, 3) + b"".join(struct.pack("<I", L) for L in lengths)
body = b"".join(codes[i][0].numpy().astype(np.uint16).tobytes() for i in range(3))
@ -276,8 +301,8 @@ def roundtrip_dcvc(net, x, mdir: Path, size, Cf=48):
opset_version=17, dynamo=False)
(mdir / "hyper_cdfs_q3.json").write_text(json.dumps(fblob(net.entropy_bottleneck), separators=(",", ":")))
(mdir / "gauss_cdfs_q3.json").write_text(json.dumps(gblob(gc), separators=(",", ":")))
hs = ort.InferenceSession(str(mdir / "hyper_synthesis.onnx"))
gs = ort.InferenceSession(str(mdir / "temporal_decoder.onnx"))
hs = _session(mdir / "hyper_synthesis.onnx")
gs = _session(mdir / "temporal_decoder.onnx")
frames = [np.roll(np.asarray(x[0].permute(1, 2, 0)), shift=8 * k, axis=1) for k in range(3)]
F = torch.zeros(1, Cf, 16, 16); fcodes = []; ref = []
for fr in frames:
@ -343,8 +368,8 @@ def roundtrip_glc(net, mdir: Path, size, K=16, Cz=32, Sz=4, Sy=16):
input_names=["i"], output_names=["o"], dynamic_axes={"i": {2: "h", 3: "w"}, "o": {2: "h", 3: "w"}}, opset_version=17, dynamo=False)
export_onnx(net.synthesis, torch.zeros(1, 1, Sy, Sy), mdir / "synthesis.onnx")
(mdir / "hyper_cdfs_q3.json").write_text(json.dumps(fblob(net.entropy_bottleneck), separators=(",", ":")))
he = ort.InferenceSession(str(mdir / "hyper_entropy.onnx")); ce = ort.InferenceSession(str(mdir / "context_entropy.onnx"))
gs = ort.InferenceSession(str(mdir / "synthesis.onnx"))
he = _session(mdir / "hyper_entropy.onnx"); ce = _session(mdir / "context_entropy.onnx")
gs = _session(mdir / "synthesis.onnx")
rng = np.random.RandomState(7)
sm = lambda l: (lambda e: e / e.sum())(np.exp(l - l.max()))
cdf_of = lambda l: pmf_to_quantized_cdf(torch.tensor(sm(l), dtype=torch.float32), 16).tolist()
@ -379,7 +404,7 @@ def main() -> int:
# Each harness is independent; a failure in one must not block the others
# (e.g. a synthetic model that only traces at certain sizes). Report a
# per-decoder PASS/FAIL summary and exit non-zero only if ALL failed.
# per-decoder PASS/FAIL summary, and exit non-zero unless EVERY harness succeeded.
harnesses = [
("elic", lambda: roundtrip_elic(build(MiniElic, 0), x, Path("/tmp/wai_elic"), args.size)),
("mlicpp", lambda: roundtrip_mlicpp(build(MiniMlic, 1), x, Path("/tmp/wai_mlic"), args.size)),
@ -397,7 +422,19 @@ def main() -> int:
print(f" ✗ {name}: {type(e).__name__}: {e}")
print(f"generated {ok}/{len(harnesses)} synthetic neural fixtures into /tmp")
print("now: cargo test --features neural --test neural_native_decode -- _matches_reference")
return 0 if ok else 1
# Every harness must succeed, not merely one of them.
#
# `0 if ok else 1` reported success whenever a single harness worked, and the
# loop above deliberately catches per-harness exceptions to keep going. So a
# 1-of-5 run exited 0, the CI step went green, and the four decoders whose
# fixtures were never written then SKIPPED in the Rust suite — where a skip is
# reported as a pass. The job came out fully green having checked one decoder
# out of five, and nothing in the output said so.
#
# Partial generation is still useful interactively: the per-harness ✓/✗ lines
# are printed either way, and a caller who wants one fixture can ignore the
# exit code. What is not acceptable is CI reading a partial run as a pass.
return 0 if ok == len(harnesses) else 1
if __name__ == "__main__":