qualia_cli/llm_raw_bench/
command.rs

1use std::path::Path;
2use std::time::{Duration, SystemTime};
3
4use qualia_core_db::inference::inference_bench::raw_decode::{
5    run_raw_decode_blocking, RawDecodeConfig,
6};
7use qualia_core_db::inference::runtime::{
8    cleanup_stale_runs, ArtifactCleanupCounters, ArtifactRetention, BackendKind, RunArtifactDir,
9};
10
11pub struct CommandConfig<'a> {
12    pub model: &'a Path,
13    pub steps: u32,
14    pub warmups: u16,
15    pub runs: u16,
16    pub quantization: &'a str,
17    pub prompt: &'a str,
18    pub target_prompt_tokens: Option<u32>,
19    pub retain_artifacts: Option<&'a Path>,
20}
21
22pub fn run(command: CommandConfig<'_>) -> Result<(), String> {
23    let mut config = RawDecodeConfig::new(
24        command
25            .model
26            .file_name()
27            .and_then(|name| name.to_str())
28            .unwrap_or("model"),
29        command.model.to_string_lossy(),
30        command.quantization,
31        command.prompt,
32    );
33    config.decode_steps = command.steps;
34    config.warmup_runs = command.warmups;
35    config.measured_runs = command.runs;
36    config.target_prompt_tokens = command.target_prompt_tokens;
37    config.requested_backend = BackendKind::Unknown;
38
39    let mut result = run_raw_decode_blocking(&config)?;
40    let retained = if let Some(target) = command.retain_artifacts {
41        let scratch_parent = target
42            .parent()
43            .filter(|parent| !parent.as_os_str().is_empty())
44            .unwrap_or_else(|| Path::new("."));
45        if scratch_parent.is_dir() {
46            let cutoff = SystemTime::now()
47                .checked_sub(Duration::from_secs(24 * 60 * 60))
48                .unwrap_or(SystemTime::UNIX_EPOCH);
49            match cleanup_stale_runs(scratch_parent, cutoff) {
50                Ok(cleanup) => {
51                    result.manifest.receipt.artifacts.temp_removed_bytes = result
52                        .manifest
53                        .receipt
54                        .artifacts
55                        .temp_removed_bytes
56                        .saturating_add(cleanup.removed_bytes);
57                    result.manifest.receipt.artifacts.temp_cleanup_failures = result
58                        .manifest
59                        .receipt
60                        .artifacts
61                        .temp_cleanup_failures
62                        .saturating_add(cleanup.failures);
63                }
64                Err(_) => {
65                    result.manifest.receipt.artifacts.temp_cleanup_failures = result
66                        .manifest
67                        .receipt
68                        .artifacts
69                        .temp_cleanup_failures
70                        .saturating_add(1);
71                }
72            }
73        }
74        let mut artifacts = RunArtifactDir::new_in(
75            scratch_parent,
76            "raw-decode",
77            8 * 1024 * 1024,
78            ArtifactRetention::RetainTo(target.to_path_buf()),
79        )
80        .map_err(|e| e.to_string())?;
81
82        let manifest_json = manifest_json_with_predicted_stats(&mut result)?;
83        artifacts
84            .write_bounded("manifest.json", manifest_json.as_bytes())
85            .map_err(|e| e.to_string())?;
86        let tokens_json =
87            serde_json::to_vec_pretty(&result.generated_token_ids).map_err(|e| e.to_string())?;
88        artifacts
89            .write_bounded("generated-token-ids.json", &tokens_json)
90            .map_err(|e| e.to_string())?;
91        let token_bytes_json =
92            serde_json::to_vec_pretty(&result.generated_token_bytes).map_err(|e| e.to_string())?;
93        artifacts
94            .write_bounded("generated-token-bytes.json", &token_bytes_json)
95            .map_err(|e| e.to_string())?;
96        artifacts
97            .write_bounded("generated-text.txt", result.generated_text.as_bytes())
98            .map_err(|e| e.to_string())?;
99        let finish = artifacts.finish().map_err(|e| e.to_string())?;
100        finish.retained_path
101    } else {
102        None
103    };
104
105    println!(
106        "RAW_DECODE median_tok_s={:.4} p95_ms_per_token={:.4} backend={:?} steps={} warmups={} runs={} dispatches={} fences={} d2h_bytes={} fallback_count={}",
107        result.manifest.median_tok_s,
108        result.manifest.p95_ms_per_token,
109        result.manifest.receipt.executed_backend,
110        result.manifest.decode_steps_executed,
111        result.manifest.warmup_runs,
112        result.manifest.measured_runs,
113        result.manifest.receipt.counters.compute_dispatches,
114        result.manifest.receipt.counters.device_fences,
115        result.manifest.receipt.counters.device_to_host_bytes,
116        result.manifest.receipt.counters.fallback_count,
117    );
118    println!(
119        "RAW_DECODE_TEXT {}",
120        serde_json::to_string(&result.generated_text).map_err(|e| e.to_string())?
121    );
122    println!(
123        "{}",
124        serde_json::to_string_pretty(&result.manifest).map_err(|e| e.to_string())?
125    );
126    if let Some(path) = retained {
127        eprintln!("RAW_DECODE_ARTIFACTS {}", path.display());
128    }
129    Ok(())
130}
131
132fn manifest_json_with_predicted_stats(
133    result: &mut qualia_core_db::inference::inference_bench::raw_decode::RawDecodeResult,
134) -> Result<String, String> {
135    let token_bytes = serde_json::to_vec_pretty(&result.generated_token_ids)
136        .map_err(|e| e.to_string())?
137        .len() as u64;
138    let token_piece_bytes = serde_json::to_vec_pretty(&result.generated_token_bytes)
139        .map_err(|e| e.to_string())?
140        .len() as u64;
141    let text_bytes = result.generated_text.len() as u64;
142    let mut previous_len = 0u64;
143    for _ in 0..4 {
144        let json = serde_json::to_string_pretty(&result.manifest).map_err(|e| e.to_string())?;
145        let total = json.len() as u64 + token_bytes + token_piece_bytes + text_bytes;
146        result.manifest.receipt.artifacts = ArtifactCleanupCounters {
147            temp_created_bytes: total,
148            temp_removed_bytes: 0,
149            temp_retained_bytes: total,
150            temp_cleanup_failures: 0,
151        };
152        if total == previous_len {
153            return serde_json::to_string_pretty(&result.manifest).map_err(|e| e.to_string());
154        }
155        previous_len = total;
156    }
157    serde_json::to_string_pretty(&result.manifest).map_err(|e| e.to_string())
158}
qualia_cli/llm_raw_bench/command.rs

qualia_cli/llm_raw_bench/
command.rs
