harness_lock_core/
tuner.rs

1//! The tuner: turns a manifest into an `harness.lock`.
2
3use std::collections::BTreeMap;
4use std::path::{Path, PathBuf};
5
6use serde_json::json;
7
8use crate::adapters::{HarnessAdapter, ReferenceAdapter};
9use crate::error::{Error, Result};
10use crate::lockfile::{Context, Decision, Evidence, Harness, Instructions, LLM, Lockfile, Policy, Provenance, Skill, Tool, Verify};
11use crate::manifest::{HarnessCandidate, Manifest};
12use crate::search::{select_winner, Stack, StackScore, TrialResult};
13use crate::store::Store;
14use crate::tasks::TaskPack;
15use crate::worker::{Backend, PythonAdapter};
16
17/// Aggregated metrics container produced by the tuner.
18
19pub struct Tuner {
20    pub manifest: Manifest,
21    pub workdir: PathBuf,
22    pub store: Store,
23    pub backend: Backend,
24    pack_by_task: BTreeMap<String, crate::tasks::TaskPack>,
25    adapter: Box<dyn HarnessAdapter>,
26    holdout: Vec<crate::tasks::Task>,
27}
28
29impl Tuner {
30    /// Create a tuner that runs on the local, in-process reference harness.
31    pub fn new(manifest: Manifest, workdir: impl Into<PathBuf>) -> Self {
32        Self::with_backend(manifest, workdir, Backend::Local)
33    }
34
35    /// Create a tuner for a specific backend. `devbox`/`cluster` delegate trials
36    /// to the Python worker (which may lift them to real Flyte tasks); `local`
37    /// keeps the deterministic offline reference harness.
38    pub fn with_backend(manifest: Manifest, workdir: impl Into<PathBuf>, backend: Backend) -> Self {
39        let workdir: PathBuf = workdir.into();
40        let store = Store::new(workdir.join("store"));
41        let adapter: Box<dyn HarnessAdapter> = if backend.needs_worker() {
42            Box::new(PythonAdapter::new())
43        } else {
44            Box::new(ReferenceAdapter::default())
45        };
46        Tuner {
47            manifest,
48            workdir,
49            store,
50            backend,
51            pack_by_task: BTreeMap::new(),
52            adapter,
53            holdout: Vec::new(),
54        }
55    }
56
57    /// Load the manifest's task packs and return the search-split tasks.
58    pub fn load_packs(&mut self) -> Vec<crate::tasks::Task> {
59        let mut search: Vec<crate::tasks::Task> = Vec::new();
60        let base = std::env::current_dir().unwrap_or_else(|_| PathBuf::from("."));
61        for ref_ in self.manifest.task_packs() {
62            match TaskPack::load(&ref_, Some(&base)) {
63                Ok(pack) => {
64                    let (s, h) = pack.split(self.manifest.search_frac(), self.manifest.budget.random_seed);
65                    for t in pack.tasks.clone() {
66                        self.pack_by_task.insert(t.id.clone(), pack.clone());
67                    }
68                    search.extend(s);
69                    self.holdout.extend(h);
70                }
71                Err(e) => eprintln!("warn: skipping task pack {ref_}: {e}"),
72            }
73        }
74        search
75    }
76
77    /// The holdout-split tasks (populated by [`load_packs`]).
78    pub fn holdout(&self) -> &[crate::tasks::Task] {
79        &self.holdout
80    }
81
82    /// Cartesian product over the manifest's candidate axes.
83    pub fn enumerate_stacks(&self) -> Vec<Stack> {
84        let c = &self.manifest.candidates;
85        let llm_combos = llm_combos(&c.llm);
86        let decisions = if c.decision.is_empty() {
87            vec!["none".to_string()]
88        } else {
89            c.decision.clone()
90        };
91        let harnesses = harness_combos(&c.harness);
92
93        let mut stacks = Vec::new();
94        for llm in &llm_combos {
95            for decision in &decisions {
96                for (harness, pattern) in &harnesses {
97                    stacks.push(Stack {
98                        llm: llm.clone(),
99                        harness: harness.clone(),
100                        pattern: pattern.clone(),
101                        decision: decision.clone(),
102                        tools: self.manifest.tools.keys().cloned().collect(),
103                        skills: self.manifest.skills.keys().cloned().collect(),
104                        instructions: instructions_files(&self.manifest.instructions),
105                        max_steps: 40,
106                        max_cost_usd_per_task: None,
107                        retries: policy_axis(&self.manifest.candidates.policy, "retries", 0),
108                        temperature: policy_axis_float(&self.manifest.candidates.policy, "temperature"),
109                        tool_timeout_seconds: None,
110                    });
111                }
112            }
113        }
114        if stacks.len() > 200 {
115            stacks.truncate(200);
116        }
117        stacks
118    }
119
120    /// Run one trial (copying fixtures, running the harness, grading).
121    pub fn run_trial(&self, stack: &Stack, task: &crate::tasks::Task, seed: u64, ws_root: &Path) -> TrialResult {
122        let ws = ws_root.join(format!("{}_{}_{}", &stack.hash()[..10.min(stack.hash().len())], task.id, seed));
123        let _ = std::fs::remove_dir_all(&ws);
124        let _ = std::fs::create_dir_all(&ws);
125
126        let outcome = self.adapter.run(stack, task, &ws, seed).unwrap_or_else(|_| {
127            crate::adapters::RunOutcome {
128                success: false,
129                cost_usd: 0.0,
130                latency_s: 0.0,
131                steps: 0,
132                trace: json!([]),
133            }
134        });
135        let grades = self
136            .pack_by_task
137            .get(&task.id)
138            .map(|p| p.grade(task, &ws))
139            .unwrap_or_default();
140        let success = outcome.success && grades.iter().all(|g| g.success);
141        let score = grades.iter().map(|g| g.score).sum::<f64>() / grades.len().max(1) as f64;
142
143        // persist the trace to the store (content-addressed)
144        let trace = json!({
145            "trial": {"stack": stack.hash(), "task": task.id, "seed": seed},
146            "events": outcome.trace,
147            "result": {"artifacts": task.expected_artifacts()},
148            "metrics": {"steps": outcome.steps, "cost_usd": outcome.cost_usd, "latency_s": outcome.latency_s},
149            "grades": grades.iter().map(|g| json!({"success": g.success, "score": g.score})).collect::<Vec<_>>(),
150        });
151        let _ = self.store.put_value(&trace);
152
153        TrialResult {
154            stack_hash: stack.hash(),
155            task_id: task.id.clone(),
156            seed,
157            success,
158            score,
159            cost_usd: outcome.cost_usd,
160            latency_s: outcome.latency_s,
161            steps: outcome.steps,
162        }
163    }
164
165    /// Tune: enumerate, run, aggregate, select, and write `harness.lock`.
166    pub fn tune(&mut self) -> Result<Lockfile> {
167        let search_tasks = self.load_packs();
168        let stacks = self.enumerate_stacks();
169        let seeds = self.manifest.seeds();
170        let ws_root = self.workdir.join("runs");
171        let _ = std::fs::create_dir_all(&ws_root);
172
173        let mut results_by_stack: BTreeMap<String, Vec<TrialResult>> = BTreeMap::new();
174        for stack in &stacks {
175            for task in &search_tasks {
176                for seed in 0..seeds {
177                    let tr = self.run_trial(stack, task, seed, &ws_root);
178                    results_by_stack.entry(stack.hash()).or_default().push(tr);
179                }
180            }
181        }
182        // respect max-trials globally
183        let mut n = 0usize;
184        let max = self.manifest.budget.max_trials.unwrap_or(usize::MAX);
185        for v in results_by_stack.values_mut() {
186            if n >= max {
187                v.clear();
188            } else {
189                v.truncate(max.saturating_sub(n));
190                n += v.len();
191            }
192        }
193
194        let mut scores: Vec<StackScore> = Vec::new();
195        for stack in &stacks {
196            if let Some(trials) = results_by_stack.get(&stack.hash()) {
197                if !trials.is_empty() {
198                    scores.push(StackScore {
199                        stack: stack.clone(),
200                        trials: trials.clone(),
201                        n: trials.len(),
202                    });
203                }
204            }
205        }
206
207        let winner = select_winner(
208            scores.clone(),
209            &self.manifest.objective.primary,
210            &self.manifest.objective.constraints,
211            &self.manifest.objective.select,
212        )
213        .ok_or_else(|| Error::Tuning("no viable stack found".into()))?;
214
215        let lock = self.build_resolve(&winner, &scores);
216        lock.save(std::path::Path::new(crate::LOCKFILE_PATH))?;
217        let _ = self.write_report(&lock, &scores);
218        Ok(lock)
219    }
220
221    fn build_resolve(&self, winner: &StackScore, _scores: &[StackScore]) -> Lockfile {
222        let s = &winner.stack;
223        let llms = s
224            .llm
225            .iter()
226            .map(|(role, spec)| {
227                let (provider, model) = match spec.split_once(':') {
228                    Some((p, m)) if !m.is_empty() => (p.to_string(), m.to_string()),
229                    _ => ("reference".to_string(), spec.clone()),
230                };
231                LLM {
232                    role: role.clone(),
233                    provider,
234                    model,
235                    resolved_from: Some(spec.clone()),
236                    params: BTreeMap::new(),
237                    fingerprint: Some(format!("sha256:{}", crate::hash::sha256_hex(spec.as_bytes()))),
238                }
239            })
240            .collect();
241
242        let mut success_rate = BTreeMap::new();
243        success_rate.insert("mean".to_string(), roundf(winner.success_rate()));
244        let mut cost = BTreeMap::new();
245        cost.insert("mean".to_string(), roundf(winner.mean_cost()));
246        let mut latency = BTreeMap::new();
247        latency.insert("p50".to_string(), roundf(winner.mean_latency()));
248        latency.insert("p95".to_string(), roundf(winner.p95_latency()));
249
250        let mut stack_map = BTreeMap::new();
251        stack_map.insert("name".to_string(), serde_json::json!(self.manifest.name()));
252        stack_map.insert("selection".to_string(), serde_json::json!(self.manifest.objective.select));
253        let mut search_map = BTreeMap::new();
254        search_map.insert("strategy".to_string(), serde_json::json!(self.manifest.budget.strategy));
255        search_map.insert("candidates".to_string(), serde_json::json!(winner.n));
256
257        Lockfile {
258            manifest_hash: Some(self.manifest.manifest_hash()),
259            stack: stack_map,
260            llm: llms,
261            decision: Decision {
262                kind: s.decision.clone(),
263                decisions: if s.decision == "none" { vec![] } else { vec!["retry".into(), "stop".into()] },
264                ..Default::default()
265            },
266            harness: Harness {
267                name: s.harness.clone(),
268                pattern: Some(s.pattern.clone()),
269                ..Default::default()
270            },
271            tool: s.tools.iter().map(|t| Tool { name: t.clone(), kind: "builtin".into(), ..Default::default() }).collect(),
272            skill: s.skills.iter().map(|sk| Skill { name: sk.clone(), origin: "first-party".into(), editable: true, evidence_digest: Some(String::new()), ..Default::default() }).collect(),
273            instructions: s.instructions.iter().map(|i| Instructions { name: i.clone(), scope: "project".into(), ..Default::default() }).collect(),
274            policy: Policy {
275                retries: 1,
276                max_steps: 40,
277                human_approval: "never".into(),
278                ..Default::default()
279            },
280            context: Context { isolation: self.manifest.context.get("isolation").and_then(serde_json::Value::as_str).unwrap_or("locked-only").into(), ..Default::default() },
281            evidence: Evidence {
282                task_pack: Some(self.manifest.task_packs().join(",")),
283                split: "search".into(),
284                trials: winner.n,
285                success_rate,
286                cost_usd_per_task: cost,
287                latency_s_per_task: latency,
288                ..Default::default()
289            },
290            verify: Verify::default(),
291            provenance: Provenance {
292                backend: self.manifest.backend_target(),
293                search: search_map,
294                ..Default::default()
295            },
296            ..Default::default()
297        }
298    }
299
300    fn write_report(&self, lock: &Lockfile, _scores: &[StackScore]) -> Result<()> {
301        let report = self.workdir.join("report.md");
302        std::fs::write(&report, format!("# harness-lock report\n\n{}\n", lock.summary()))?;
303        Ok(())
304    }
305}
306
307fn roundf(v: f64) -> f64 {
308    (v * 1e4).round() / 1e4
309}
310
311fn llm_combos(cands: &BTreeMap<String, crate::manifest::LLMCandidate>) -> Vec<BTreeMap<String, String>> {
312    if cands.is_empty() {
313        let mut m = BTreeMap::new();
314        m.insert("planner".to_string(), "reference:qwen-7b".to_string());
315        return vec![m];
316    }
317    let roles: Vec<String> = cands.keys().cloned().collect();
318    let mut per_role: Vec<Vec<(String, String)>> = Vec::new();
319    for role in &roles {
320        let opts = cands.get(role).map(|c| c.any_of.clone()).unwrap_or_default();
321        let opts = if opts.is_empty() {
322            vec!["reference:qwen-7b".to_string()]
323        } else {
324            opts
325        };
326        per_role.push(opts.into_iter().map(|o| (role.clone(), o)).collect());
327    }
328    // cartesian product
329    let mut combos: Vec<BTreeMap<String, String>> = vec![BTreeMap::new()];
330    for options in per_role {
331        let mut next = Vec::new();
332        for c in &combos {
333            for (role, model) in &options {
334                let mut cc = c.clone();
335                cc.insert(role.clone(), model.clone());
336                next.push(cc);
337            }
338        }
339        combos = next;
340    }
341    combos
342}
343
344fn harness_combos(cands: &[HarnessCandidate]) -> Vec<(String, String)> {
345    if cands.is_empty() {
346        return vec![("reference".to_string(), "single-agent".to_string())];
347    }
348    let mut out = Vec::new();
349    for h in cands {
350        let patterns = if h.pattern.is_empty() { vec!["single-agent".to_string()] } else { h.pattern.clone() };
351        for p in patterns {
352            out.push((h.name.clone(), p));
353        }
354    }
355    out
356}
357
358fn instructions_files(ins: &BTreeMap<String, serde_json::Value>) -> Vec<String> {
359    ins.get("files")
360        .and_then(serde_json::Value::as_array)
361        .map(|a| a.iter().filter_map(|v| v.as_str().map(|s| s.to_string())).collect())
362        .unwrap_or_default()
363}
364
365fn policy_axis(policy: &BTreeMap<String, Vec<serde_json::Value>>, key: &str, default: u32) -> u32 {
366    policy
367        .get(key)
368        .and_then(|v| v.first())
369        .and_then(serde_json::Value::as_u64)
370        .map(|n| n as u32)
371        .unwrap_or(default)
372}
373
374fn policy_axis_float(policy: &BTreeMap<String, Vec<serde_json::Value>>, key: &str) -> Option<f64> {
375    policy
376        .get(key)
377        .and_then(|v| v.first())
378        .and_then(serde_json::Value::as_f64)
379}
harness_lock_core/tuner.rs

harness_lock_core/
tuner.rs
