Skip to main content

qualia_core_db/inference/lab/
experiment.rs

1//! Experiment execution engine — wraps the existing `run_bench` harness with
2//! quality verification, thermal sampling, and structured result collection.
3//!
4//! Each experiment produces an `ExperimentResult` containing:
5//! - `BenchResult` (latency, tok/s, phase split)
6//! - `VerifiedTurn` (post-turn quality checks)
7//! - Thermal snapshot (GPU temp/power before + after)
8//! - VRAM usage from the GPU context ledger
9//! - CBOR config snapshot for reproduction
10
11use std::path::Path;
12use std::time::{SystemTime, UNIX_EPOCH};
13
14use serde::{Deserialize, Serialize};
15
16use crate::gpu_context::global_vram_ledger;
17use crate::inference_bench::{
18    run_bench, set_attention_o_fuse, set_attention_preproject, set_coop_gemv,
19    set_decode_budget_override, set_ffn_fusion, set_gpu_topk, set_kv_dict, set_kv_int8,
20    set_resident_decode, set_resident_prefill, set_resident_weights, set_spec_decode,
21    set_ternary_ffn, BenchConfig, BenchResult, LlmPhaseSnapshot,
22};
23use crate::post_turn_verify::{verify_and_heal_turn, VerifiedTurn};
24use crate::thermal_telemetry::{sample_gpu_thermal, GpuThermalSample};
25
26use super::config_space::{Configuration, ConfigurationSpace};
27
28/// Thermal reading before and after an experiment.
29#[derive(Debug, Clone, Copy, Serialize, Deserialize, Default)]
30pub struct ThermalSnapshot {
31    pub before_temp_c: u32,
32    pub before_power_w: f64,
33    pub after_temp_c: u32,
34    pub after_power_w: f64,
35    /// Energy estimate in joules: avg_power_w * duration_s.
36    pub energy_j: f64,
37}
38
39impl ThermalSnapshot {
40    fn capture(
41        before: Option<GpuThermalSample>,
42        after: Option<GpuThermalSample>,
43        duration_s: f64,
44    ) -> Self {
45        let (before_temp_c, before_power_w) =
46            before.map(|s| (s.temp_c, s.power_w)).unwrap_or((0, 0.0));
47        let (after_temp_c, after_power_w) =
48            after.map(|s| (s.temp_c, s.power_w)).unwrap_or((0, 0.0));
49        let avg_power = (before_power_w + after_power_w) / 2.0;
50        Self {
51            before_temp_c,
52            before_power_w,
53            after_temp_c,
54            after_power_w,
55            energy_j: avg_power * duration_s,
56        }
57    }
58}
59
60/// Serializable subset of `BenchResult` (the original only derives `Serialize`).
61#[derive(Debug, Clone, Serialize, Deserialize, Default)]
62pub struct BenchResultSerde {
63    pub label: String,
64    pub quantization: String,
65    pub prompt_tokens: u64,
66    pub output_tokens: u64,
67    pub cold_ttft_ms: f64,
68    pub cold_total_ms: f64,
69    pub warm_ttft_ms: f64,
70    pub warm_total_ms: f64,
71    pub load_ms: f64,
72    pub prefill_ms: f64,
73    pub prefill_tok_s: f64,
74    pub decode_ms: f64,
75    pub decode_tok_s: f64,
76    pub mapped_bytes: u64,
77    pub kv_cache_bytes: u64,
78}
79
80impl From<&BenchResult> for BenchResultSerde {
81    fn from(r: &BenchResult) -> Self {
82        Self {
83            label: r.label.clone(),
84            quantization: r.quantization.clone(),
85            prompt_tokens: r.prompt_tokens,
86            output_tokens: r.output_tokens,
87            cold_ttft_ms: r.cold_ttft_ms,
88            cold_total_ms: r.cold_total_ms,
89            warm_ttft_ms: r.warm_ttft_ms,
90            warm_total_ms: r.warm_total_ms,
91            load_ms: r.load_ms,
92            prefill_ms: r.prefill_ms,
93            prefill_tok_s: r.prefill_tok_s,
94            decode_ms: r.decode_ms,
95            decode_tok_s: r.decode_tok_s,
96            mapped_bytes: r.model.mapped_bytes,
97            kv_cache_bytes: r.model.kv_cache_bytes,
98        }
99    }
100}
101
102/// The full result of one experiment run.
103#[derive(Debug, Clone, Serialize, Deserialize)]
104pub struct ExperimentResult {
105    /// FNV-1a hash of the configuration CBOR (dedup key).
106    pub config_hash: u64,
107    /// Links to the hypothesis backlog (if any).
108    pub hypothesis_id: Option<String>,
109    /// The benchmark result (latency, tok/s, phase split).
110    pub bench: Option<BenchResultSerde>,
111    /// Intra-decode phase breakdown.
112    pub phase: PhaseSnapshotSerde,
113    /// Post-turn quality checks.
114    pub quality: QualityScore,
115    /// GPU thermal before/after.
116    pub thermal: ThermalSnapshot,
117    /// VRAM used (bytes) from the GPU context ledger.
118    pub vram_used: u64,
119    /// CBOR config snapshot for reproduction.
120    pub config_cbor: Vec<u8>,
121    /// Unix epoch timestamp (ms).
122    pub timestamp_ms: u64,
123    /// Sampler seed for reproducibility.
124    pub seed: u64,
125    /// Error message if the experiment failed.
126    pub error: Option<String>,
127}
128
129/// Serializable phase snapshot (the atomic counters are copied into plain u64s).
130#[derive(Debug, Clone, Copy, Serialize, Deserialize, Default)]
131pub struct PhaseSnapshotSerde {
132    pub load_ns: u64,
133    pub prefill_ns: u64,
134    pub prefill_tokens: u64,
135    pub decode_ns: u64,
136    pub decode_tokens: u64,
137    pub decode_forward_ns: u64,
138    pub decode_output_ns: u64,
139}
140
141impl From<LlmPhaseSnapshot> for PhaseSnapshotSerde {
142    fn from(s: LlmPhaseSnapshot) -> Self {
143        Self {
144            load_ns: s.load_ns,
145            prefill_ns: s.prefill_ns,
146            prefill_tokens: s.prefill_tokens,
147            decode_ns: s.decode_ns,
148            decode_tokens: s.decode_tokens,
149            decode_forward_ns: s.decode_forward_ns,
150            decode_output_ns: s.decode_output_ns,
151        }
152    }
153}
154
155/// Quality score derived from `VerifiedTurn` checks.
156#[derive(Debug, Clone, Serialize, Deserialize, Default)]
157pub struct QualityScore {
158    /// Fraction of checks that passed (0.0 to 1.0).
159    pub pass_rate: f64,
160    /// Total number of checks.
161    pub total_checks: usize,
162    /// Number of checks that passed.
163    pub passed: usize,
164    /// Whether the draft was repaired by the graph.
165    pub repaired: bool,
166    /// Final text length (chars).
167    pub text_len: usize,
168}
169
170impl QualityScore {
171    pub fn from_verified(v: &VerifiedTurn) -> Self {
172        let total = v.checks.len();
173        let passed = v.checks.iter().filter(|c| c.ok).count();
174        Self {
175            pass_rate: if total > 0 {
176                passed as f64 / total as f64
177            } else {
178                1.0
179            },
180            total_checks: total,
181            passed,
182            repaired: v.repaired,
183            text_len: v.final_text.len(),
184        }
185    }
186
187    /// A composite quality score in [0, 1]: pass_rate weighted by whether repair was needed.
188    pub fn composite(&self) -> f64 {
189        let base = self.pass_rate;
190        if self.repaired {
191            base * 0.9
192        } else {
193            base
194        }
195    }
196}
197
198/// Configuration for one experiment run.
199#[derive(Debug, Clone)]
200pub struct ExperimentConfig {
201    /// The configuration space this experiment belongs to.
202    pub space: ConfigurationSpace,
203    /// The concrete configuration to test.
204    pub config: Configuration,
205    /// Model path for the benchmark.
206    pub model_path: String,
207    /// Quantization label.
208    pub quantization: String,
209    /// Prompt to run.
210    pub prompt: String,
211    /// Decode token budget (0 = production default).
212    pub decode_tokens: u32,
213    /// Warm repeats.
214    pub warm_repeats: u32,
215    /// Sampler seed.
216    pub seed: u64,
217    /// Linked hypothesis ID.
218    pub hypothesis_id: Option<String>,
219}
220
221impl ExperimentConfig {
222    /// Build a `BenchConfig` from this experiment config.
223    pub fn to_bench_config(&self) -> BenchConfig {
224        BenchConfig {
225            label: format!("{}_h{:016x}", self.space.name, self.config.hash()),
226            model_path: self.model_path.clone(),
227            quantization: self.quantization.clone(),
228            prompt: self.prompt.clone(),
229            decode_tokens: self.decode_tokens,
230            warm_repeats: self.warm_repeats,
231        }
232    }
233
234    /// Apply the configuration's toggle values to the runtime.
235    /// This reads known parameter names from the config and sets the corresponding toggles.
236    pub fn apply_toggles(&self) {
237        if let Some(v) = self.config.get_bool("gpu_topk") {
238            set_gpu_topk(v);
239        }
240        if let Some(v) = self.config.get_bool("ternary_ffn") {
241            set_ternary_ffn(v);
242        }
243        if let Some(v) = self.config.get_bool("kv_dict") {
244            set_kv_dict(v);
245        }
246        if let Some(v) = self.config.get_bool("attention_preproject") {
247            set_attention_preproject(v);
248        }
249        if let Some(v) = self.config.get_bool("attention_o_fuse") {
250            set_attention_o_fuse(v);
251        }
252        if let Some(v) = self.config.get_bool("spec_decode") {
253            set_spec_decode(v);
254        }
255        if let Some(v) = self.config.get_bool("ffn_fusion") {
256            set_ffn_fusion(v);
257        }
258        if let Some(v) = self.config.get_bool("coop_gemv") {
259            set_coop_gemv(v);
260        }
261        if let Some(v) = self.config.get_bool("kv_int8") {
262            set_kv_int8(v);
263        }
264        if let Some(v) = self.config.get_bool("resident_decode") {
265            set_resident_decode(v);
266        }
267        if let Some(v) = self.config.get_bool("resident_prefill") {
268            set_resident_prefill(v);
269        }
270        if let Some(v) = self.config.get_bool("resident_weights") {
271            set_resident_weights(v);
272        }
273    }
274}
275
276/// Run one experiment end-to-end: apply toggles, run benchmark, verify quality,
277/// sample thermal, collect results.
278pub fn run_experiment(cfg: &ExperimentConfig) -> ExperimentResult {
279    let config_hash = cfg.config.hash();
280    let config_cbor = cfg.config.to_cbor();
281    let timestamp_ms = SystemTime::now()
282        .duration_since(UNIX_EPOCH)
283        .map(|d| d.as_millis() as u64)
284        .unwrap_or(0);
285
286    // Apply toggle configuration
287    cfg.apply_toggles();
288    set_decode_budget_override(cfg.decode_tokens);
289
290    // Thermal sample before
291    let thermal_before = sample_gpu_thermal();
292    let t_start = std::time::Instant::now();
293
294    // Run the benchmark
295    let bench_result = run_bench(&cfg.to_bench_config());
296
297    let elapsed = t_start.elapsed().as_secs_f64();
298
299    // Thermal sample after
300    let thermal_after = sample_gpu_thermal();
301    let thermal = ThermalSnapshot::capture(thermal_before, thermal_after, elapsed);
302
303    // Collect phase metrics
304    let phase = crate::inference_bench::phase_snapshot();
305    let phase_serde = PhaseSnapshotSerde::from(phase);
306
307    // VRAM usage from the global VRAM ledger
308    let vram_used = global_vram_ledger().used_bytes();
309
310    // Quality check: run post-turn verify on the output
311    let quality = match &bench_result {
312        Ok(br) => {
313            // We don't have the actual generated text from run_bench (it only returns metrics),
314            // so we use a quality proxy from the bench result itself.
315            // In a full implementation, we'd run decode_with_metrics to get the text.
316            // For now, we score based on whether the benchmark succeeded and produced tokens.
317            let has_output = br.output_tokens > 0;
318            QualityScore {
319                pass_rate: if has_output { 1.0 } else { 0.0 },
320                total_checks: 1,
321                passed: if has_output { 1 } else { 0 },
322                repaired: false,
323                text_len: 0,
324            }
325        }
326        Err(_) => QualityScore {
327            pass_rate: 0.0,
328            total_checks: 1,
329            passed: 0,
330            repaired: false,
331            text_len: 0,
332        },
333    };
334
335    // Restore defaults
336    set_decode_budget_override(0);
337
338    let (bench, error) = match bench_result {
339        Ok(r) => (Some(BenchResultSerde::from(&r)), None),
340        Err(e) => (None, Some(e)),
341    };
342
343    ExperimentResult {
344        config_hash,
345        hypothesis_id: cfg.hypothesis_id.clone(),
346        bench,
347        phase: phase_serde,
348        quality,
349        thermal,
350        vram_used,
351        config_cbor,
352        timestamp_ms,
353        seed: cfg.seed,
354        error,
355    }
356}
357
358/// Run an experiment with quality verification (uses `decode_with_metrics` to get
359/// the actual generated text, then runs `verify_and_heal_turn` on it).
360pub fn run_experiment_with_quality(cfg: &ExperimentConfig) -> ExperimentResult {
361    use crate::inference_bench::{decode_with_metrics, reset_phase_metrics};
362
363    let config_hash = cfg.config.hash();
364    let config_cbor = cfg.config.to_cbor();
365    let timestamp_ms = SystemTime::now()
366        .duration_since(UNIX_EPOCH)
367        .map(|d| d.as_millis() as u64)
368        .unwrap_or(0);
369
370    cfg.apply_toggles();
371
372    let thermal_before = sample_gpu_thermal();
373    let t_start = std::time::Instant::now();
374
375    reset_phase_metrics();
376
377    let decode_tokens = cfg.decode_tokens.max(8).min(128);
378    let result = decode_with_metrics(&cfg.model_path, &cfg.prompt, decode_tokens);
379
380    let elapsed = t_start.elapsed().as_secs_f64();
381    let thermal_after = sample_gpu_thermal();
382    let thermal = ThermalSnapshot::capture(thermal_before, thermal_after, elapsed);
383
384    let phase = crate::inference_bench::phase_snapshot();
385    let phase_serde = PhaseSnapshotSerde::from(phase);
386
387    let vram_used = global_vram_ledger().used_bytes();
388
389    let (bench, quality, error) = match result {
390        Ok((text, tok_s)) => {
391            // Run post-turn verification on the generated text
392            let verified = verify_and_heal_turn(&cfg.prompt, &text);
393            let qs = QualityScore::from_verified(&verified);
394
395            // Build a serializable bench result from the decode metrics
396            let bench = BenchResultSerde {
397                label: format!("{}_h{:016x}", cfg.space.name, config_hash),
398                quantization: cfg.quantization.clone(),
399                prompt_tokens: phase_serde.prefill_tokens,
400                output_tokens: phase_serde.decode_tokens,
401                cold_ttft_ms: 0.0,
402                cold_total_ms: 0.0,
403                warm_ttft_ms: 0.0,
404                warm_total_ms: elapsed * 1000.0,
405                load_ms: phase_serde.load_ns as f64 / 1_000_000.0,
406                prefill_ms: phase_serde.prefill_ns as f64 / 1_000_000.0,
407                prefill_tok_s: if phase_serde.prefill_ns > 0 && phase_serde.prefill_tokens > 0 {
408                    phase_serde.prefill_tokens as f64 / (phase_serde.prefill_ns as f64 / 1e9)
409                } else {
410                    0.0
411                },
412                decode_ms: phase_serde.decode_ns as f64 / 1_000_000.0,
413                decode_tok_s: tok_s,
414                mapped_bytes: 0,
415                kv_cache_bytes: 0,
416            };
417            (Some(bench), qs, None)
418        }
419        Err(e) => (None, QualityScore::default(), Some(e)),
420    };
421
422    set_decode_budget_override(0);
423
424    ExperimentResult {
425        config_hash,
426        hypothesis_id: cfg.hypothesis_id.clone(),
427        bench,
428        phase: phase_serde,
429        quality,
430        thermal,
431        vram_used,
432        config_cbor,
433        timestamp_ms,
434        seed: cfg.seed,
435        error,
436    }
437}
438
439/// Append-only JSONL experiment log.
440pub fn append_experiment_jsonl(path: &Path, result: &ExperimentResult) -> Result<(), String> {
441    if let Some(parent) = path.parent() {
442        std::fs::create_dir_all(parent).map_err(|e| e.to_string())?;
443    }
444    let line = serde_json::to_string(result).map_err(|e| e.to_string())?;
445    let mut file = std::fs::OpenOptions::new()
446        .create(true)
447        .append(true)
448        .open(path)
449        .map_err(|e| format!("open {}: {e}", path.display()))?;
450    use std::io::Write;
451    writeln!(file, "{line}").map_err(|e| e.to_string())?;
452    Ok(())
453}
454
455/// Load all experiment results from a JSONL log.
456pub fn load_experiment_log(path: &Path) -> Result<Vec<ExperimentResult>, String> {
457    if !path.exists() {
458        return Ok(Vec::new());
459    }
460    let content = std::fs::read_to_string(path).map_err(|e| e.to_string())?;
461    let mut results = Vec::new();
462    for line in content.lines() {
463        if line.trim().is_empty() {
464            continue;
465        }
466        match serde_json::from_str::<ExperimentResult>(line) {
467            Ok(r) => results.push(r),
468            Err(e) => log::warn!("experiment_log|skip_line|{e}"),
469        }
470    }
471    Ok(results)
472}
473
474#[cfg(test)]
475mod tests {
476    use super::*;
477    use crate::inference::lab::config_space::{ConfigurationSpace, ParameterDef};
478
479    #[test]
480    fn quality_score_from_verified() {
481        use crate::post_turn_verify::{VerifiedTurn, VerifyCheck};
482        let v = VerifiedTurn {
483            final_text: "Paris".into(),
484            display_html: String::new(),
485            cml_turtle: String::new(),
486            repaired: false,
487            checks: vec![
488                VerifyCheck {
489                    id: "a".into(),
490                    ok: true,
491                    detail: "ok".into(),
492                },
493                VerifyCheck {
494                    id: "b".into(),
495                    ok: true,
496                    detail: "ok".into(),
497                },
498            ],
499            grounding_reason: None,
500        };
501        let qs = QualityScore::from_verified(&v);
502        assert_eq!(qs.pass_rate, 1.0);
503        assert_eq!(qs.total_checks, 2);
504        assert!(!qs.repaired);
505    }
506
507    #[test]
508    fn quality_score_with_repair() {
509        use crate::post_turn_verify::{VerifiedTurn, VerifyCheck};
510        let v = VerifiedTurn {
511            final_text: "Paris".into(),
512            display_html: String::new(),
513            cml_turtle: String::new(),
514            repaired: true,
515            checks: vec![
516                VerifyCheck {
517                    id: "a".into(),
518                    ok: true,
519                    detail: "ok".into(),
520                },
521                VerifyCheck {
522                    id: "b".into(),
523                    ok: false,
524                    detail: "fail".into(),
525                },
526            ],
527            grounding_reason: Some("capital".into()),
528        };
529        let qs = QualityScore::from_verified(&v);
530        assert_eq!(qs.pass_rate, 0.5);
531        assert!(qs.repaired);
532        assert!((qs.composite() - 0.45).abs() < 1e-9);
533    }
534
535    #[test]
536    fn experiment_config_builds_bench_config() {
537        let space = ConfigurationSpace::new("test").with("gpu_topk", ParameterDef::Bool);
538        let cfg = space.build_from_normalized(&[1.0]);
539        let ec = ExperimentConfig {
540            space,
541            config: cfg,
542            model_path: "/dev/null".into(),
543            quantization: "Q8_0".into(),
544            prompt: "Hello".into(),
545            decode_tokens: 16,
546            warm_repeats: 1,
547            seed: 42,
548            hypothesis_id: Some("H-001".into()),
549        };
550        let bc = ec.to_bench_config();
551        assert_eq!(bc.quantization, "Q8_0");
552        assert_eq!(bc.decode_tokens, 16);
553        assert!(bc.label.contains("test_"));
554    }
555
556    #[test]
557    fn jsonl_log_roundtrip() {
558        let tmp =
559            std::env::temp_dir().join(format!("qualia_lab_test_{}.jsonl", std::process::id()));
560        let result = ExperimentResult {
561            config_hash: 42,
562            hypothesis_id: Some("H-001".into()),
563            bench: None,
564            phase: PhaseSnapshotSerde::default(),
565            quality: QualityScore::default(),
566            thermal: ThermalSnapshot::default(),
567            vram_used: 0,
568            config_cbor: vec![0xA1, 0x01],
569            timestamp_ms: 12345,
570            seed: 99,
571            error: Some("test".into()),
572        };
573        append_experiment_jsonl(&tmp, &result).unwrap();
574        let loaded = load_experiment_log(&tmp).unwrap();
575        assert_eq!(loaded.len(), 1);
576        assert_eq!(loaded[0].config_hash, 42);
577        assert_eq!(loaded[0].seed, 99);
578        let _ = std::fs::remove_file(&tmp);
579    }
580}