diff --git a/recipe/cua_s1/native.md b/recipe/cua_s1/native.md index 35336704..78d14b1a 100644 --- a/recipe/cua_s1/native.md +++ b/recipe/cua_s1/native.md @@ -25,7 +25,7 @@ CUA_S1_MODEL=weights/cua-s1-4b-0.2-text-merged target/release/omni-cua-s1-native For the local CUDA Graph experiment, also set `CUA_S1_GRAPH=1`. The first use of each exact ordered prompt-length vector warms the GEMM plans and captures the forward pass; -later requests replay it with freshly uploaded token ids. The text worker keeps a 64-entry FIFO cache. Growing the scratch allocation clears the captures before freeing +later requests replay it with freshly uploaded token ids. Each language mode keeps a separate 64-entry FIFO cache. Growing the scratch allocation clears the captures before freeing their buffers. Capture adds first-use latency; leave the variable unset to use the eager control. Rebuild both the worker and CUDA library together (ABI 6). If capture fails, the worker returns the completed eager result and disables diff --git a/recipe/cua_s1/native_vision_graph.md b/recipe/cua_s1/native_vision_graph.md index e6fcc71c..f21413d0 100644 --- a/recipe/cua_s1/native_vision_graph.md +++ b/recipe/cua_s1/native_vision_graph.md @@ -4,8 +4,8 @@ The `omni-cua-s1-vision` screenshot worker can capture the complete vision encoder, including all 24 blocks, unmerged FP32 LoRA branches and the merger. Set `CUA_S1_VISION_GRAPH=1` before loading the worker to enable it. Leave it unset or set it to `0` for eager execution. This switch is independent of -`CUA_S1_GRAPH`, which controls language execution; this change does not enable -multimodal language Graph replay. +`CUA_S1_GRAPH`, which controls text and multimodal language execution. Enable +each switch separately; keep both unset for eager execution. Both modes retain one scratch allocation for the latest exact `[T,H,W]` patch grid. Equal patch counts with different heights or widths replace the scratch @@ -57,9 +57,9 @@ Export the same `CUA_S1_BASE`, `CUA_S1_VISION_ADAPTER`, `CUA_S1_MODEL`, and ```sh cargo test --release --locked -p omni-cua-s1-native --lib vision_graph_ \ -- --ignored --test-threads=1 --nocapture -CUA_S1_VISION_GRAPH=0 cargo run --release --locked -p omni-cua-s1-native \ +CUA_S1_GRAPH=0 CUA_S1_VISION_GRAPH=0 cargo run --release --locked -p omni-cua-s1-native \ --example vision_graph_bench -- /tmp/vision-eager.json -CUA_S1_VISION_GRAPH=1 cargo run --release --locked -p omni-cua-s1-native \ +CUA_S1_GRAPH=0 CUA_S1_VISION_GRAPH=1 cargo run --release --locked -p omni-cua-s1-native \ --example vision_graph_bench -- /tmp/vision-graph.json ``` diff --git a/src/models/qwen3_5/native/src/model.rs b/src/models/qwen3_5/native/src/model.rs index 304c3083..1e65fdb8 100644 --- a/src/models/qwen3_5/native/src/model.rs +++ b/src/models/qwen3_5/native/src/model.rs @@ -598,9 +598,10 @@ pub struct Model { gemm: *mut c_void, /// Buffers for the largest packed token count so far; grows as needed. scratch: Option, - /// Opt-in replay with at most 64 captures keyed by ordered sequence lengths. + /// Opt-in replay with at most 64 captures per mode, keyed by sequence lengths. graph_enabled: bool, graphs: VecDeque<(Vec, cuda::Graph)>, + mm_graphs: VecDeque<(Vec, cuda::Graph)>, } // SAFETY: the raw pointers are device addresses and a cuBLASLt handle owned by the @@ -610,6 +611,7 @@ unsafe impl Send for Model {} impl Drop for Model { fn drop(&mut self) { self.graphs.clear(); + self.mm_graphs.clear(); // SAFETY: created by cs1_gemm_create and not destroyed before. unsafe { (cuda::api().cs1_gemm_destroy)(self.gemm) }; } @@ -692,6 +694,7 @@ impl Model { scratch: None, graph_enabled: std::env::var("CUA_S1_GRAPH").as_deref() == Ok("1"), graphs: VecDeque::new(), + mm_graphs: VecDeque::new(), }; Ok(model) } @@ -745,6 +748,7 @@ impl Model { cuda::set_device(0)?; if self.scratch.as_ref().is_none_or(|s| t > s.cap) { self.graphs.clear(); + self.mm_graphs.clear(); self.scratch = None; self.scratch = Some(Scratch::new( &self.cfg, @@ -793,28 +797,21 @@ impl Model { if self.graph_enabled { if let Some((_, graph)) = self.graphs.iter().find(|(shape, _)| *shape == lengths) { graph.launch(self.stream)?; + if std::env::var("CUA_S1_GRAPH_TRACE").as_deref() == Ok("1") { + eprintln!("CUA_S1_GRAPH mode=text event=replay tokens={t}"); + } } else { // Warm plans and keep the eager result: run() advances the residual // in place, so replaying on this cache miss would advance it twice. self.run(s, &lengths, false)?; cuda::synchronize(self.stream)?; - match cuda::Graph::capture(self.stream, || self.run(s, &lengths, false)) { - Ok(graph) => { - if self.graphs.len() == 64 { - self.graphs.pop_front(); - } - self.graphs.push_back((lengths.clone(), graph)); - } - Err(error) => { - eprintln!("CUDA Graph capture failed; using eager execution: {error:#}"); - self.graph_enabled = false; - self.graphs.clear(); - } - } + let captured = cuda::Graph::capture(self.stream, || self.run(s, &lengths, false)); + self.cache_graph(&lengths, false, captured); } } else { self.run(s, &lengths, false)?; } + let s = self.scratch.as_ref().unwrap(); let mut end = 0; lengths .iter() @@ -845,8 +842,58 @@ impl Model { self.upload_positions(s, input.position_ids)?; self.embed_tokens(s, input.token_ids)?; self.overwrite_image_rows(s, input, 0)?; - self.run(s, &[t], true)?; - self.last_hidden(s, t) + if self.graph_enabled { + if let Some((_, graph)) = self + .mm_graphs + .iter() + .find(|(shape, _)| shape.as_slice() == [t]) + { + graph.launch(self.stream)?; + if std::env::var("CUA_S1_GRAPH_TRACE").as_deref() == Ok("1") { + eprintln!("CUA_S1_GRAPH mode=multimodal event=replay tokens={t}"); + } + } else { + // Uploads and embedding are outside capture. Keep the warmed eager + // result; launching now would advance the residual a second time. + self.run(s, &[t], true)?; + cuda::synchronize(self.stream)?; + let captured = cuda::Graph::capture(self.stream, || self.run(s, &[t], true)); + self.cache_graph(&[t], true, captured); + } + } else { + self.run(s, &[t], true)?; + } + self.last_hidden(self.scratch.as_ref().unwrap(), t) + } + + /// Both modes retain the eager miss result; capture records without executing. + fn cache_graph(&mut self, lengths: &[usize], multimodal: bool, captured: Result) { + let mode = if multimodal { "multimodal" } else { "text" }; + match captured { + Ok(graph) => { + let cache = if multimodal { + &mut self.mm_graphs + } else { + &mut self.graphs + }; + if cache.len() == 64 { + cache.pop_front(); + } + cache.push_back((lengths.to_vec(), graph)); + if std::env::var("CUA_S1_GRAPH_TRACE").as_deref() == Ok("1") { + let t: usize = lengths.iter().sum(); + eprintln!("CUA_S1_GRAPH mode={mode} event=capture tokens={t}"); + } + } + Err(error) => { + eprintln!( + "CUDA Graph capture failed in {mode} mode; using eager execution: {error:#}" + ); + self.graph_enabled = false; + self.graphs.clear(); + self.mm_graphs.clear(); + } + } } fn upload_positions(&self, s: &Scratch, positions: [&[i64]; 3]) -> Result<()> { @@ -1590,3 +1637,6 @@ impl Model { Ok(()) } } +#[cfg(test)] +#[path = "../../../../../tests/qwen3_5/multimodal_graph.rs"] +mod graph_tests; diff --git a/tests/qwen3_5/README.md b/tests/qwen3_5/README.md new file mode 100644 index 00000000..c8a8376f --- /dev/null +++ b/tests/qwen3_5/README.md @@ -0,0 +1,47 @@ +# Native multimodal language Graph regressions + +These private-item tests live under root `tests/` and are wired into the +`omni-qwen3-5-native` library with `#[cfg(test)]`. Cargo compiles and discovers +all three cases without a generated source copy, a preparation script or manifest +changes. The wiring is excluded from release builds. + +Ordinary `cargo test --workspace --locked` leaves the three cases ignored: +standard CI has no CUDA device or model weights. It still compiles their bodies. +Actual execution requires a real NVIDIA CUDA device, a supported BF16 language +checkpoint with `image_token_id` and at least 2048 positions, and the Qwen native +CUDA library. The base checkpoint's F32 tensors need the existing native export; +see [the multimodal language exporter](../../recipe/cua_s1/export_multimodal_language.py) +and [native build instructions](../../recipe/cua_s1/native.md). + +```sh +export CUA_S1_TEST_MODEL=/absolute/path/to/merged-bf16-language +export CUA_S1_CUDA_LIB=/absolute/path/to/libqwen3_5_cuda.so +python3 tests/qwen3_5/run_graph_regressions.py \ + --log /tmp/native-mm-graph-gpu.log --result /tmp/native-mm-graph-result.json +``` + +The runner uses locked, offline, release Cargo execution. Fetch dependencies +beforehand (`cargo fetch --locked`) when a populated Cargo cache/vendor directory +is unavailable. Tests run serially in one process, with `CUA_S1_GRAPH_TRACE` +unset, to avoid shared CUDA device/API state interference. It requires exactly +three passing cases and one unconditional capture-failure diagnostic per mode. +It returns nonzero on missing hardware/weights, Cargo failure or a failed gate; +none of those are treated as skips. + +Coverage: + +- Capture misses and replay agree exactly with finite eager hidden states; + changed tokens, BF16 image rows and distinct T/H/W positions are uploaded; + text/multimodal interleave and scratch growth permit fresh captures. Text + singleton `[4]` and packed `[1, 3]` / `[3, 1]` keep distinct ordered cache + keys despite equal total tokens, and replay with changed token IDs. +- 65 lengths exercise the separate 64-entry multimodal FIFO, unchanged scratch + addresses, hits without insertion-order refresh, eviction/recapture and a + surviving text replay. +- A real CUDA recording closure returns an injected error. The production + capture-result handler preserves the completed eager result, clears both + caches, disables Graph, emits diagnostics and allows repeated public eager + forwards in both modes. This does not simulate every CUDA failure class. + +These are executor regressions, not a screenshot-encoder or agent-accuracy +benchmark. No latency, capture-cost or memory claim is made by these tests. diff --git a/tests/qwen3_5/multimodal_graph.rs b/tests/qwen3_5/multimodal_graph.rs new file mode 100644 index 00000000..e3979e15 --- /dev/null +++ b/tests/qwen3_5/multimodal_graph.rs @@ -0,0 +1,362 @@ +use super::*; +use half::bf16; + +fn load_model() -> Result { + let dir = std::env::var("CUA_S1_TEST_MODEL").context("CUA_S1_TEST_MODEL")?; + let lib = std::env::var("CUA_S1_CUDA_LIB").context("CUA_S1_CUDA_LIB")?; + let mut model = Model::load(Path::new(&dir), Path::new(&lib))?; + model.graph_enabled = false; + Ok(model) +} + +fn text_tokens(model: &Model, count: usize) -> Result> { + let vocabulary: Vec = (0..model.embed.shape[0]) + .filter_map(|id| u32::try_from(id).ok()) + .filter(|id| Some(*id) != model.cfg.image_token_id) + .take(4) + .collect(); + ensure!(vocabulary.len() == 4, "test needs four non-image token IDs"); + Ok(vocabulary.into_iter().cycle().take(count).collect()) +} + +fn assert_hidden_eq(actual: &[f32], expected: &[f32], context: &str) { + assert!( + expected.iter().all(|x| x.is_finite()), + "{context}: eager hidden state must be finite" + ); + assert_eq!(actual, expected, "{context}"); +} + +#[test] +#[ignore = "requires the Cua-S1 language checkpoint and real CUDA library; run serially"] +fn multimodal_graph_replays_current_inputs_and_clears_on_growth() -> Result<()> { + let mut model = load_model()?; + ensure!(model.cfg.max_positions >= 2048, "test needs 2048 positions"); + let image_token = model.cfg.image_token_id.context("image_token_id")?; + let text = text_tokens(&model, 4)?; + let ids = [text[0], image_token, image_token, text[1]]; + let indices = [1, 2]; + let positions = [0, 1, 1, 2]; + let zeros = vec![bf16::ZERO; 2 * model.cfg.hidden]; + let input = MultimodalInput { + token_ids: &ids, + image_token_indices: &indices, + image_embeddings: &zeros, + position_ids: [&positions; 3], + }; + let eager = model.forward_multimodal(&input)?; + model.graph_enabled = true; + assert_hidden_eq( + &model.forward_multimodal(&input)?, + &eager, + "capture miss must retain the eager result", + ); + assert_eq!(model.mm_graphs.len(), 1, "miss must cache a graph"); + assert!(model.graphs.is_empty(), "text cache must remain separate"); + assert!(model.graph_enabled, "capture must remain enabled"); + assert_hidden_eq( + &model.forward_multimodal(&input)?, + &eager, + "replay must equal eager exactly", + ); + + let changed_ids = [text[2], image_token, image_token, text[3]]; + let changed_features: Vec = (0..2 * model.cfg.hidden) + .map(|i| bf16::from_f32((i % 17) as f32 / 16.0 - 0.5)) + .collect(); + let temporal = [0, 3, 3, 5]; + let height = [0, 2, 4, 5]; + let width = [0, 4, 2, 5]; + let changed = MultimodalInput { + token_ids: &changed_ids, + image_token_indices: &indices, + image_embeddings: &changed_features, + position_ids: [&temporal, &height, &width], + }; + model.graph_enabled = false; + let changed_eager = model.forward_multimodal(&changed)?; + assert_ne!( + changed_eager, eager, + "changed inputs must exercise new data" + ); + model.graph_enabled = true; + assert_hidden_eq( + &model.forward_multimodal(&changed)?, + &changed_eager, + "replay must upload current IDs, image rows and T/H/W positions", + ); + assert_eq!(model.mm_graphs.len(), 1, "same length must reuse its graph"); + + model.graph_enabled = false; + let text_eager = model.forward(&text)?; + model.graph_enabled = true; + assert_hidden_eq(&model.forward(&text)?, &text_eager, "text capture miss"); + assert_hidden_eq(&model.forward(&text)?, &text_eager, "text replay"); + assert_eq!(model.graphs.len(), 1); + assert_eq!(model.mm_graphs.len(), 1); + assert_hidden_eq( + &model.forward_multimodal(&changed)?, + &changed_eager, + "multimodal replay after text interleave", + ); + // Equal packed totals must not alias singleton or differently ordered + // sequence boundaries. Reuse every shape again with changed token IDs. + let changed_text = [text[3], text[2], text[1], text[0]]; + let shapes = [vec![4], vec![1, 3], vec![3, 1]]; + for tokens in [text.as_slice(), changed_text.as_slice()] { + for lengths in &shapes { + let mut begin = 0; + let inputs: Vec<&[u32]> = lengths + .iter() + .map(|&length| { + let end = begin + length; + let ids = &tokens[begin..end]; + begin = end; + ids + }) + .collect(); + model.graph_enabled = false; + let eager = model.forward_batch(&inputs)?; + model.graph_enabled = true; + for context in ["packed text capture/hit", "packed text replay"] { + let actual = model.forward_batch(&inputs)?; + assert_eq!(actual.len(), eager.len(), "{context}: sequence count"); + for (actual, expected) in actual.iter().zip(&eager) { + assert_hidden_eq(actual, expected, context); + } + } + assert!(model.graph_enabled, "packed shape capture must succeed"); + assert!(model.graphs.iter().any(|(shape, _)| shape == lengths)); + assert_eq!(model.mm_graphs.len(), 1, "text must retain MM graph"); + assert_hidden_eq( + &model.forward_multimodal(&changed)?, + &changed_eager, + "multimodal replay after packed text interleave", + ); + } + } + assert_eq!( + model + .graphs + .iter() + .map(|(shape, _)| shape.clone()) + .collect::>(), + shapes.to_vec(), + "singleton, split and reversed splits must keep distinct FIFO keys", + ); + + let next_cap = model.scratch.as_ref().unwrap().cap + 1; + model.prepare_scratch(next_cap)?; + assert!(model.graphs.is_empty(), "growth must clear text graphs"); + assert!(model.mm_graphs.is_empty(), "growth must clear MM graphs"); + assert!(model.graph_enabled, "growth must allow fresh captures"); + assert_hidden_eq(&model.forward(&text)?, &text_eager, "text after growth"); + assert_hidden_eq( + &model.forward_multimodal(&changed)?, + &changed_eager, + "multimodal after growth", + ); + assert_eq!(model.graphs.len(), 1); + assert_eq!(model.mm_graphs.len(), 1); + model.synchronize()?; + Ok(()) +} + +#[test] +#[ignore = "requires the Cua-S1 language checkpoint and real CUDA library; run serially"] +fn multimodal_graph_fifo_evicts_and_recaptures_without_touching_text() -> Result<()> { + let mut model = load_model()?; + ensure!(model.cfg.max_positions >= 1024, "test needs 1024 positions"); + let text = text_tokens(&model, 4)?; + let text_eager = model.forward(&text)?; + model.graph_enabled = true; + assert_hidden_eq(&model.forward(&text)?, &text_eager, "text capture miss"); + assert_eq!(model.graphs.len(), 1); + let scratch_address = model.scratch.as_ref().unwrap().buf.at(0); + assert_eq!(model.scratch.as_ref().unwrap().cap, 1024); + + let mut first_eager = Vec::new(); + for length in 1..=65 { + let ids = text_tokens(&model, length)?; + let positions: Vec = (0..length as i64).collect(); + let input = MultimodalInput { + token_ids: &ids, + image_token_indices: &[], + image_embeddings: &[], + position_ids: [&positions; 3], + }; + model.graph_enabled = false; + let eager = model.forward_multimodal(&input)?; + if length == 1 { + first_eager = eager.clone(); + } + model.graph_enabled = true; + assert_hidden_eq( + &model.forward_multimodal(&input)?, + &eager, + &format!("multimodal capture miss at length {length}"), + ); + assert!(model.graph_enabled, "capture failed at length {length}"); + assert_eq!(model.mm_graphs.len(), length.min(64)); + let expected: Vec = (length.saturating_sub(63).max(1)..=length).collect(); + assert_eq!( + model + .mm_graphs + .iter() + .map(|(shape, _)| shape[0]) + .collect::>(), + expected, + "multimodal cache must evict in insertion order" + ); + assert_eq!(model.scratch.as_ref().unwrap().cap, 1024); + assert_eq!(model.scratch.as_ref().unwrap().buf.at(0), scratch_address); + assert_eq!(model.graphs.len(), 1, "MM capture must keep text graph"); + assert_eq!(model.graphs.front().unwrap().0, vec![text.len()]); + } + assert!( + !model + .mm_graphs + .iter() + .any(|(shape, _)| shape.as_slice() == [1]) + ); + + // A hit must not refresh FIFO order. Length 2 remains the oldest capture. + let ids = text_tokens(&model, 2)?; + let positions = [0, 1]; + let input = MultimodalInput { + token_ids: &ids, + image_token_indices: &[], + image_embeddings: &[], + position_ids: [&positions; 3], + }; + model.graph_enabled = false; + let eager = model.forward_multimodal(&input)?; + model.graph_enabled = true; + assert_hidden_eq(&model.forward_multimodal(&input)?, &eager, "oldest hit"); + assert_eq!(model.mm_graphs.front().unwrap().0, vec![2]); + assert_eq!(model.mm_graphs.back().unwrap().0, vec![65]); + + let ids = text_tokens(&model, 1)?; + let positions = [0]; + let input = MultimodalInput { + token_ids: &ids, + image_token_indices: &[], + image_embeddings: &[], + position_ids: [&positions; 3], + }; + assert_hidden_eq( + &model.forward_multimodal(&input)?, + &first_eager, + "recapturing evicted length must retain its eager result", + ); + assert_eq!(model.mm_graphs.len(), 64); + let expected: Vec = (3..=65).chain(std::iter::once(1)).collect(); + assert_eq!( + model + .mm_graphs + .iter() + .map(|(shape, _)| shape[0]) + .collect::>(), + expected + ); + assert_hidden_eq( + &model.forward_multimodal(&input)?, + &first_eager, + "recaptured length replay", + ); + assert_eq!(model.graphs.len(), 1, "text graph must survive eviction"); + assert_eq!(model.graphs.front().unwrap().0, vec![text.len()]); + assert_hidden_eq(&model.forward(&text)?, &text_eager, "surviving text replay"); + assert!(model.graph_enabled); + assert_eq!(model.scratch.as_ref().unwrap().cap, 1024); + assert_eq!(model.scratch.as_ref().unwrap().buf.at(0), scratch_address); + model.synchronize()?; + Ok(()) +} + +#[test] +#[ignore = "requires the Cua-S1 language checkpoint and real CUDA library; run serially"] +fn graph_recording_failure_keeps_eager_result_and_disables_both_modes() -> Result<()> { + for multimodal in [false, true] { + let mut model = load_model()?; + let image_token = model.cfg.image_token_id.context("image_token_id")?; + let text = text_tokens(&model, 4)?; + let ids = [text[0], image_token, image_token, text[1]]; + let indices = [1, 2]; + let temporal = [0, 1, 1, 3]; + let height = [0, 1, 2, 3]; + let width = [0, 2, 1, 3]; + let features: Vec = (0..2 * model.cfg.hidden) + .map(|i| bf16::from_f32((i % 11) as f32 / 10.0 - 0.5)) + .collect(); + let input = MultimodalInput { + token_ids: &ids, + image_token_indices: &indices, + image_embeddings: &features, + position_ids: [&temporal, &height, &width], + }; + let text_eager = model.forward(&text)?; + let mm_eager = model.forward_multimodal(&input)?; + model.graph_enabled = true; + assert_hidden_eq(&model.forward(&text)?, &text_eager, "seed text graph"); + assert_hidden_eq( + &model.forward_multimodal(&input)?, + &mm_eager, + "seed multimodal graph", + ); + assert!(model.graph_enabled, "both seed captures must succeed"); + assert_eq!(model.graphs.len(), 1); + assert_eq!(model.mm_graphs.len(), 1); + + // Leave a completed eager hidden state in scratch before a real CUDA + // capture begins and its recording closure returns an error. + model.graph_enabled = false; + let prior_eager = if multimodal { + model.forward_multimodal(&input)? + } else { + model.forward(&text)? + }; + model.graph_enabled = true; + let captured = + cuda::Graph::capture(model.stream, || anyhow::bail!("injected record failure")); + assert!( + captured + .as_ref() + .err() + .context("recording failure unexpectedly succeeded")? + .to_string() + .contains("injected record failure"), + "must exercise the actual recording-error path" + ); + model.cache_graph(&[text.len()], multimodal, captured); + assert_hidden_eq( + &model.last_hidden(model.scratch.as_ref().unwrap(), text.len())?, + &prior_eager, + "recording failure must preserve the earlier eager hidden state", + ); + assert!( + !model.graph_enabled, + "recording failure must disable replay" + ); + assert!(model.graphs.is_empty(), "failure must clear text graphs"); + assert!(model.mm_graphs.is_empty(), "failure must clear MM graphs"); + + for _ in 0..2 { + assert_hidden_eq( + &model.forward(&text)?, + &text_eager, + "public text forward after capture failure", + ); + assert_hidden_eq( + &model.forward_multimodal(&input)?, + &mm_eager, + "public multimodal forward after capture failure", + ); + assert!(!model.graph_enabled, "eager fallback must stay enabled"); + assert!(model.graphs.is_empty(), "text must not recapture"); + assert!(model.mm_graphs.is_empty(), "multimodal must not recapture"); + } + model.synchronize()?; + } + Ok(()) +} diff --git a/tests/qwen3_5/run_graph_regressions.py b/tests/qwen3_5/run_graph_regressions.py new file mode 100644 index 00000000..4f16717c --- /dev/null +++ b/tests/qwen3_5/run_graph_regressions.py @@ -0,0 +1,87 @@ +#!/usr/bin/env python3 +"""Run the real-GPU cases and require unconditional capture-failure diagnostics.""" + +import argparse +import json +import os +from pathlib import Path +import re +import subprocess +import sys + + +def diagnostics(output: str) -> dict: + pattern = re.compile( + r"CUDA Graph capture failed in (text|multimodal) mode; " + r"using eager execution: ([^\r\n]*)", + ) + matches = pattern.findall(output) + counts = {mode: sum(m == mode for m, _ in matches) for mode in ("text", "multimodal")} + valid = ( + len(matches) == 2 + and counts == {"text": 1, "multimodal": 1} + and all("injected record failure" in error for _, error in matches) + ) + return { + "diagnostic_count": len(matches), + "diagnostics_by_mode": counts, + "failure_diagnostics": [{"mode": mode, "error": error} for mode, error in matches], + "unconditional_diagnostics_passed": valid, + } + + +def main() -> int: + test_dir = Path(__file__).resolve().parent + repo = test_dir.parents[1] + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--log", required=True, type=Path, help="full Cargo output") + parser.add_argument("--result", type=Path, help="JSON result; defaults to LOG.json") + args = parser.parse_args() + log_path = args.log.resolve() + result_path = ( + args.result.resolve() if args.result else log_path.with_name(log_path.name + ".json") + ) + if log_path == result_path: + parser.error("--log and --result must be different paths") + log_path.parent.mkdir(parents=True, exist_ok=True) + result_path.parent.mkdir(parents=True, exist_ok=True) + env = os.environ.copy() + env.pop("CUA_S1_GRAPH_TRACE", None) + command = [ + "cargo", "test", "--offline", "--release", "--locked", + "-p", "omni-qwen3-5-native", "--lib", "model::graph_tests::", "--", + "--ignored", "--nocapture", "--test-threads=1", + ] + with log_path.open("wb") as log: + tested = subprocess.run( + command, cwd=repo, env=env, stdout=log, stderr=subprocess.STDOUT, + check=False, + ) + cargo_returncode = tested.returncode + summary = diagnostics(log_path.read_bytes().decode("utf-8", errors="replace")) + counts = re.findall( + r"test result: ok\. (\d+) passed; (\d+) failed; (\d+) ignored", + log_path.read_text(encoding="utf-8", errors="replace"), + ) + test_counts = [tuple(map(int, count)) for count in counts] + passed = ( + cargo_returncode == 0 + and test_counts == [(3, 0, 0)] + and summary["unconditional_diagnostics_passed"] + ) + result = { + "passed": passed, + "test_counts": test_counts, + "command": command, + "cargo_returncode": cargo_returncode, + "graph_trace_unset": "CUA_S1_GRAPH_TRACE" not in env, + "full_log": str(log_path), + **summary, + } + result_path.write_text(json.dumps(result, indent=2) + "\n", encoding="utf-8") + print(json.dumps(result, indent=2)) + return 0 if passed else 1 + + +if __name__ == "__main__": + sys.exit(main())