open-standards/jouleclaw/jouleclaw-rs/crates/jouleclaw-harness-econ/examples/swebench_local_coder.rs

//! The genuine number, end to end: a real instruction-tuned **code model**
//! (Granite-4.0-h, pure-Rust loader-gguf) generates a fix, which is applied as
//! the whole corrected file and judged by a **real `python3` test run**
//! ([`WholeFileValidator`]), with the decode energy metered. Output is the econ
//! report's joules-and-$ per VERIFIED fix on an operationally-verified fix.
//!
//! The local model emits the corrected file (the protocol that works for local
//! models — they fumble unified-diff syntax), wrapped in Granite's chat
//! template. Verification is operational: the patch is written and `python3`
//! actually imports + asserts. Energy covers the decode alone,
//! which runs on this process's own threads (the loader's CPU kernels; built
//! without its GPU `fabric`), so the meter is `jouleclaw-energy`'s in-process meter
//! (`detect_in_process_meter`) at its own provenance: this process's own
//! CPU-energy counter (`process-cpu`, `HwShunt` over `ProcessCpu`) where it is
//! live, as on macOS 27.0.1 on an M5 Max and an M3 Ultra; else the host meter —
//! the IOReport counter (`HwShunt`) where its CPU rail is live, or the
//! load-derived estimator (`Estimator`). The `python3` judge is a subprocess and
//! runs outside the metered window. The report says which meter ran (decode time
//! + tokens are exact). A local model is priced energy-only.
//!
//! Run (defaults to the granite-4.0-h-micro snapshot if present):
//! ```text
//! cargo run --release -p jouleclaw-harness-econ --example swebench_local_coder
//! cargo run --release -p jouleclaw-harness-econ --example swebench_local_coder -- /path/to/snapshot
//! ```

use std::path::{Path, PathBuf};
use std::process::Command;
use std::sync::atomic::{AtomicU32, Ordering};
use std::sync::{Arc, Mutex};
use std::time::Instant;

use jouleclaw_cascade::{
    Answer, AnswerError, AnswerOutput, Cascade, ContextRef, ExecutionTrace, JouleBudget,
    L3ModelId, QualityFloor, Query, QueryInput, Runtime, Tier, TierEstimate, TierId,
    VerificationStatus,
};
use jouleclaw_energy::meter::detect_in_process_meter;
use jouleclaw_energy::Provenance;
use jouleclaw_harness_econ::swebench::exec::{ShellTestRunner, WholeFileValidator};
use jouleclaw_harness_econ::swebench::{PatchValidator, SweInstance};
use jouleclaw_harness_econ::{joules_obtained, window_joules, EconHarness, PriceModel, Task};
use jouleclaw_loader_gguf::bpe_tokenizer::BpeTokenizer;
use jouleclaw_loader_gguf::granite_h::{GraniteHConfig, GraniteHModel};
use jouleclaw_loader_gguf::safetensors_reader::read_safetensors_file;
use jouleclaw_skill::{ConstantSkillInducer, SkillTier};

const DEFAULT_SNAPSHOT: &str = "/Volumes/macos_4TB_external_1/hf_hub/models--ibm-granite--granite-4.0-h-micro/snapshots/d5f01a3ea75f088947be3aae039f4ad52837dfde";

struct DispatchStat {
    secs: f64,
    tokens: usize,
    joules: f64,
    /// The meter that metered the decode, and the provenance of its readings.
    meter: &'static str,
    provenance: Provenance,
    text: String,
}

/// A cascade tier backed by a real instruction-tuned code model. Wraps the
/// query in Granite's chat template, greedy-decodes until the end-of-text
/// token, and meters the decode window.
struct LocalCoderTier {
    model: GraniteHModel,
    tok: BpeTokenizer,
    max_tokens: usize,
    eot: Option<u32>,
    calls: Arc<AtomicU32>,
    log: Arc<Mutex<Vec<DispatchStat>>>,
}

impl LocalCoderTier {
    fn chat_prompt(&self, user: &str) -> Vec<u32> {
        // Granite 4.0 chat template; encode() splits on the added role tokens.
        let s = format!(
            "<|start_of_role|>user<|end_of_role|>{user}<|end_of_text|>\n\
             <|start_of_role|>assistant<|end_of_role|>"
        );
        self.tok.encode(&s)
    }
}

impl Tier for LocalCoderTier {
    fn id(&self) -> TierId {
        TierId::L3(L3ModelId(0))
    }
    fn estimate_cost(&self, _q: &Query) -> Option<TierEstimate> {
        Some(TierEstimate {
            joules: 30.0,
            latency: std::time::Duration::from_secs(20),
            confidence_floor: 0.8,
        })
    }
    fn try_answer(&mut self, q: &Query, _budget: f64) -> Result<Answer, AnswerError> {
        let user = match &q.input {
            QueryInput::Text(s) => s.clone(),
            _ => {
                return Err(AnswerError::TierFailed { tier: self.id(), cause: "non-text".into() })
            }
        };
        let ids = self.chat_prompt(&user);

        // The decode runs on this process's threads, so the in-process meter
        // covers it. A window the meter put no figure on stops the run
        // (`window_joules`): a failed read is not a zero-energy reading, and a
        // counter that did not advance did not see the work.
        let meter = detect_in_process_meter();
        let e0 = meter.read();
        let t = Instant::now();

        let mut state = self.model.new_state();
        let mut logits = self.model.forward_prompt(&ids, &mut state);
        let mut generated: Vec<u32> = Vec::new();
        for step in 0..self.max_tokens {
            let l = if step == 0 {
                std::mem::take(&mut logits)
            } else {
                self.model.forward_token(*generated.last().unwrap(), &mut state)
            };
            let (arg, _) = l.iter().enumerate().fold((0usize, f32::NEG_INFINITY), |(bi, bv), (i, &v)| {
                if v > bv {
                    (i, v)
                } else {
                    (bi, bv)
                }
            });
            let tok = arg as u32;
            if Some(tok) == self.eot {
                break;
            }
            generated.push(tok);
        }

        let secs = t.elapsed().as_secs_f64();
        let (joules, provenance) = window_joules(e0, meter.read())
            .unwrap_or_else(|why| panic!("no energy figure for this dispatch from `{}`: {why}", meter.active_name()));
        let text = self.tok.decode(&generated);

        self.calls.fetch_add(1, Ordering::SeqCst);
        self.log.lock().unwrap().push(DispatchStat {
            secs,
            tokens: generated.len(),
            joules,
            meter: meter.active_name(),
            provenance,
            text: text.clone(),
        });

        Ok(Answer {
            output: AnswerOutput::Text(text),
            tier_used: self.id(),
            joules_spent: joules,
            confidence: 0.8,
            trace: ExecutionTrace::default(),
            verification: VerificationStatus::Resolved,
        })
    }
}

fn git(dir: &Path, args: &[&str]) -> std::process::Output {
    Command::new("git").current_dir(dir).args(args).output().expect("git runs")
}

fn main() {
    if Command::new("python3").arg("--version").output().is_err() {
        eprintln!("python3 not found — this demo runs a real python test");
        return;
    }
    let dir = PathBuf::from(std::env::args().nth(1).unwrap_or_else(|| DEFAULT_SNAPSHOT.into()));
    if !dir.join("config.json").is_file() {
        eprintln!("model snapshot not found at {}", dir.display());
        return;
    }

    // ---- build a real repo with a buggy python file ----
    static C: AtomicU32 = AtomicU32::new(0);
    let repo = std::env::temp_dir()
        .join(format!("jc-coder-{}-{}", std::process::id(), C.fetch_add(1, Ordering::SeqCst)));
    std::fs::create_dir_all(&repo).unwrap();
    git(&repo, &["init", "-q"]);
    git(&repo, &["config", "user.email", "[email protected]"]);
    git(&repo, &["config", "user.name", "t"]);
    std::fs::write(repo.join("calc.py"), "def add(a, b):\n    return a - b\n").unwrap();
    git(&repo, &["add", "-A"]);
    git(&repo, &["commit", "-q", "-m", "base"]);
    let base = String::from_utf8(git(&repo, &["rev-parse", "HEAD"]).stdout).unwrap().trim().to_string();

    let instruction = "Fix the bug in the Python file `calc.py`. The function `add` must return \
        the sum a + b, but it currently returns a - b. Output ONLY the complete corrected \
        contents of calc.py inside a single ```python code block.\n\ncalc.py:\n```python\n\
        def add(a, b):\n    return a - b\n```";

    let inst = SweInstance {
        instance_id: "local__calc-add".into(),
        repo: "local/calc".into(),
        base_commit: base,
        problem_statement: instruction.into(),
        patch: String::new(),
        test_patch: String::new(),
        fail_to_pass: vec!["t::add".into()],
        pass_to_pass: vec![],
        version: None,
        environment_setup_commit: None,
    };

    let runner = ShellTestRunner::new(|node: &str| match node {
        "t::add" => "python3 -c 'import calc; assert calc.add(2,3)==5 and calc.add(10,4)==14'".to_string(),
        _ => "false".to_string(),
    });
    let validator: Arc<dyn PatchValidator> =
        Arc::new(WholeFileValidator::new(&repo, |_: &SweInstance| PathBuf::from("calc.py"), runner));

    // ---- load the real code model ----
    let load = Instant::now();
    let config: GraniteHConfig =
        serde_json::from_slice(&std::fs::read(dir.join("config.json")).unwrap()).unwrap();
    let mut shard_paths: Vec<PathBuf> = std::fs::read_dir(&dir)
        .unwrap()
        .filter_map(|e| e.ok().map(|e| e.path()))
        .filter(|p| {
            p.extension().is_some_and(|e| e == "safetensors")
                && p.file_name().is_some_and(|n| n.to_string_lossy().starts_with("model"))
        })
        .collect();
    shard_paths.sort();
    let shards: Vec<_> = shard_paths
        .iter()
        .map(|p| read_safetensors_file(p).unwrap_or_else(|e| panic!("{p:?}: {e:?}")))
        .collect();
    let model = GraniteHModel::bind_sharded(config, shards).expect("bind");
    let tok = BpeTokenizer::from_file(dir.join("tokenizer.json")).expect("tokenizer");
    let eot = tok.special_id("<|end_of_text|>");
    println!("loaded granite in {:.2}s (eot={eot:?})", load.elapsed().as_secs_f64());

    let calls = Arc::new(AtomicU32::new(0));
    let log = Arc::new(Mutex::new(Vec::<DispatchStat>::new()));
    let tier = LocalCoderTier { model, tok, max_tokens: 200, eot, calls: calls.clone(), log: log.clone() };
    let slot = Mutex::new(Some(tier));

    let tasks = vec![Task::text(
        inst.instance_id.clone(),
        "py-fix",
        Query {
            input: QueryInput::Text(inst.problem_statement.clone()),
            budget: JouleBudget::expensive(),
            quality: QualityFloor::any(),
            context: ContextRef::fresh(),
            deadline: None,
        },
        {
            let v = validator.clone();
            let inst = inst.clone();
            move |o| matches!(o, AnswerOutput::Text(t) if v.validate(&inst, t))
        },
    )];

    let mut harness = EconHarness::new(Box::new(ConstantSkillInducer), 3);
    harness.price = PriceModel { usd_per_kwh: 0.15, l3_call_usd: 0.0, l4_call_usd: 0.0 };

    let report = harness.run(
        move |store| {
            let mut casc = Cascade::new();
            casc.register(Box::new(SkillTier::new(store)));
            casc.register(Box::new(slot.lock().unwrap().take().expect("built once")));
            Runtime::new_without_l0(casc)
        },
        &tasks,
        1,
    );

    println!("\n=== Local code-model SWE-style fix — REAL model, REAL python test, metered ===\n");
    let stats = log.lock().unwrap();
    for s in stats.iter() {
        println!(
            "model: {} tok in {:.1}s ({:.1} tok/s), {:.3} J {} by {} ({})",
            s.tokens,
            s.secs,
            if s.secs > 0.0 { s.tokens as f64 / s.secs } else { 0.0 },
            s.joules,
            joules_obtained(s.provenance),
            s.meter,
            s.provenance
        );
        println!("--- model output ---\n{}\n--------------------", s.text.trim());
    }
    print!("{}", report.to_table());
    println!(
        "\nresolved (python test really ran): {}/{}",
        report.verified, report.total
    );
    if report.verified > 0 {
        println!(
            "joules per VERIFIED fix : {:.3e} J   ({:?})",
            report.joules_per_verified_fix, report.energy_provenance
        );
        println!("dollars per VERIFIED fix: ${:.8}  (local → energy-only)", report.usd_per_verified_fix);
    } else {
        println!(
            "0 verified — the local 3B model's patch did not pass the test this run.\n\
             That is an honest cascade signal: insufficient tier → escalate. The energy\n\
             was still really spent ({:.3} J) and is accounted; it just earns no fix credit.",
            stats.iter().map(|s| s.joules).sum::<f64>()
        );
    }

    drop(stats);
    let _ = std::fs::remove_dir_all(&repo);
}