foxhunt/crates/ml-explainability/tests/ig_diag_cli_integration.rs

//! End-to-end integration test for the `ig_diag` CLI.
//!
//! 1. Builds a fresh `DQN` with a small config,
//! 2. Saves its branching weights to a safetensors checkpoint with the
//!    same plain-name layout that the single-head loader expects,
//! 3. Writes a small states JSON file,
//! 4. Invokes the compiled `ig_diag` binary via `std::process::Command`,
//! 5. Parses the output JSON and asserts the schema + completeness axiom.
//!
//! Gated behind `#[cfg(feature = "cuda")] #[ignore]` because it requires a
//! CUDA runtime and the binary to be built with the `ig-diag-cli` feature.
//! Run with:
//! ```bash
//! cargo build -p ml-explainability --features ig-diag-cli --bin ig_diag
//! cargo test -p ml-explainability --features ig-diag-cli --test \
//!     ig_diag_cli_integration -- --ignored --nocapture
//! ```

#![cfg(all(feature = "cuda", feature = "ig-diag-cli"))]

use std::path::PathBuf;
use std::process::Command;

#[test]
#[ignore = "requires CUDA runtime + compiled ig_diag binary"]
fn test_ig_diag_cli_end_to_end() {
    // 1. Build a DQN, save a safetensors checkpoint.
    let dir = tempfile::tempdir().expect("tempdir");
    let ckpt_path = dir.path().join("ig_diag_test.safetensors");
    let states_path = dir.path().join("states.json");
    let output_path = dir.path().join("ig_report.json");

    save_dqn_checkpoint(&ckpt_path);

    // 2. Write a small states JSON file (3 random STATE_DIM=96 vectors).
    let state_dim = ml_core::state_layout::STATE_DIM;
    let states_json = serde_json::json!({
        "states": [
            (0..state_dim)
                .map(|i| ((i as f32 * 0.0123).sin() as f64) * 0.1)
                .collect::<Vec<f64>>(),
            (0..state_dim)
                .map(|i| ((i as f32 * 0.0456 + 1.0).cos() as f64) * 0.2)
                .collect::<Vec<f64>>(),
            (0..state_dim)
                .map(|i| ((i as f32 * 0.0789 - 0.5).sin() as f64) * 0.3)
                .collect::<Vec<f64>>(),
        ],
    });
    std::fs::write(&states_path, serde_json::to_string(&states_json).unwrap())
        .expect("write states.json");

    // 3. Invoke the ig_diag binary.
    let binary = binary_path();
    assert!(
        binary.exists(),
        "ig_diag binary not found at {} — run: \
         cargo build -p ml-explainability --features ig-diag-cli --bin ig_diag",
        binary.display(),
    );

    let status = Command::new(&binary)
        .arg("--checkpoint")
        .arg(&ckpt_path)
        .arg("--states")
        .arg(&states_path)
        .arg("--num-steps")
        .arg("64") // enough steps for <5% completeness on a non-linear DQN
        .arg("--output")
        .arg(&output_path)
        .status()
        .expect("spawn ig_diag");

    assert!(status.success(), "ig_diag exited non-zero: {status}");

    // 4. Parse + validate the output JSON.
    let raw = std::fs::read_to_string(&output_path).expect("read output JSON");
    let v: serde_json::Value = serde_json::from_str(&raw).expect("parse output JSON");

    // Schema checks.
    assert_eq!(v["schema_version"], 1);
    assert_eq!(v["state_dim"], state_dim);
    assert_eq!(v["num_states"], 3);
    assert!(
        v["features"].is_array(),
        "features should be array",
    );
    assert_eq!(
        v["features"].as_array().unwrap().len(),
        state_dim,
        "one entry per feature",
    );
    let feat0 = &v["features"][0];
    assert!(feat0["mean_abs"].is_number(), "mean_abs should be numeric");
    assert!(feat0["stddev"].is_number(), "stddev should be numeric");
    assert!(feat0["mean_signed"].is_number(), "mean_signed should be numeric");

    let top10 = v["top_10_by_mean_abs"].as_array().expect("top10 array");
    assert!(top10.len() <= 10);
    assert!(!top10.is_empty(), "top10 must be non-empty");

    // Completeness: worst relative error < 5%.
    //
    // For IG with num_steps=16 on a DQN forward (non-linear), 5% is a loose
    // upper bound (in practice it's much tighter). Tighten if needed.
    let worst = v["completeness"]["worst_relative_error"]
        .as_f64()
        .expect("worst_relative_error numeric");
    assert!(
        worst < 0.05,
        "completeness axiom violated: worst relative error {worst} >= 5%",
    );

    // Per-state completeness entries must be present.
    let per_state = v["completeness"]["per_state"]
        .as_array()
        .expect("per_state array");
    assert_eq!(per_state.len(), 3);

    println!("ig_diag CLI integration test OK — worst_rel_err={worst:.6}");
}

/// Return the path to the compiled `ig_diag` binary. CARGO_BIN_EXE_* is
/// injected by cargo for integration tests in the same crate; fall back to
/// target/debug if the env var is missing (e.g., when running manually).
fn binary_path() -> PathBuf {
    if let Ok(p) = std::env::var("CARGO_BIN_EXE_ig_diag") {
        return PathBuf::from(p);
    }
    // Fallback: walk up from CARGO_MANIFEST_DIR to find workspace/target/debug.
    let mut p = PathBuf::from(env!("CARGO_MANIFEST_DIR"));
    while p.pop() {
        let candidate = p.join("target").join("debug").join("ig_diag");
        if candidate.exists() {
            return candidate;
        }
    }
    PathBuf::from("target/debug/ig_diag")
}

/// Build a DQN with small hidden dims, disable NoisyNet, dump its branching
/// weights to a plain-name safetensors file. Mirrors the minimal save path
/// used by `DqnInferenceAdapter`'s round-trip test.
fn save_dqn_checkpoint(path: &std::path::Path) {
    use ml_dqn::{DQN, DQNConfig};

    let mut config = DQNConfig::emergency_safe_defaults();
    config.num_actions = 3;
    config.hidden_dims = vec![64, 64];

    let dqn = DQN::new(config.clone()).expect("DQN new");
    let stream = dqn.cuda_stream().clone();
    let branching = dqn
        .branching_q_network
        .as_ref()
        .expect("branching_q_network");

    // Download each tensor + serialize with safetensors using the plain-name
    // layout ({shared,value,branch}_*). The loader accepts this or the
    // regime-prefixed layout.
    let named = branching.named_weight_slices();
    let byte_vecs: Vec<(String, Vec<usize>, Vec<u8>)> = named
        .iter()
        .map(|(name, slice, shape)| {
            let mut host = vec![0.0_f32; slice.len()];
            stream.memcpy_dtoh(*slice, &mut host).expect("DtoH");
            let bytes: Vec<u8> = host.iter().flat_map(|v| v.to_le_bytes()).collect();
            (name.clone(), shape.clone(), bytes)
        })
        .collect();

    let mut tensor_map = std::collections::HashMap::new();
    for (name, shape, bytes) in &byte_vecs {
        let view = safetensors::tensor::TensorView::new(
            safetensors::Dtype::F32,
            shape.clone(),
            bytes,
        )
        .expect("TensorView");
        tensor_map.insert(name.as_str(), view);
    }

    let arch_meta = Some(config.checkpoint_metadata());
    safetensors::serialize_to_file(tensor_map, arch_meta, path)
        .expect("serialize_to_file");
}