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

harness_lock_core/
tuner.rs
