harness_lock_core/
worker.rs

1//! Bridge from the Rust core to the Python worker (`harness_lock.tasks.worker`).
2//!
3//! The Rust core stays the orchestrator and aggregator — manifests, search
4//! planning, trial bookkeeping, evidence, and lock construction. The
5//! ecosystem- and compute-heavy *work items* of `hlock optimize` that must run
6//! where the Python toolchain and GPUs live (real LLM invocation, user-trace
7//! loading with pandas/polars, synthetic-data generation, model grading, prompt
8//! search, LoRA fine-tuning) are handled by a small Python worker that the core
9//! shells out to over a versioned JSON contract.
10//!
11//! On `local` the tuner uses the deterministic in-process reference harness (no
12//! Python). On `devbox`/`cluster` the [`PythonAdapter`] takes over and calls
13//! `python -m harness_lock.tasks.worker <request.json> <response.json>`; that
14//! worker may in turn lift each body to a real Flyte task via the Python `flyte`
15//! / `flyteplugins-rs` packages, but returns the identical response shape either
16//! way — so the Rust core never changes based on where the work actually ran.
17
18use std::path::Path;
19
20use serde::{Deserialize, Serialize};
21use serde_json::json;
22
23use crate::adapters::{HarnessAdapter, RunOutcome};
24use crate::error::{Error, Result};
25use crate::search::Stack;
26use crate::tasks::Task;
27
28/// Where the tuner executes.
29#[derive(Debug, Clone, Copy, PartialEq, Eq)]
30pub enum Backend {
31    Local,
32    Devbox,
33    Cluster,
34}
35
36impl Backend {
37    pub fn parse(s: &str) -> Backend {
38        match s.trim().to_ascii_lowercase().as_str() {
39            "devbox" => Backend::Devbox,
40            "cluster" => Backend::Cluster,
41            _ => Backend::Local,
42        }
43    }
44
45    pub fn as_str(&self) -> &'static str {
46        match self {
47            Backend::Local => "local",
48            Backend::Devbox => "devbox",
49            Backend::Cluster => "cluster",
50        }
51    }
52
53    /// Whether this backend routes work through the Python worker.
54    pub fn needs_worker(&self) -> bool {
55        matches!(self, Backend::Devbox | Backend::Cluster)
56    }
57}
58
59/// Current version of the Rust <-> Python task contract.
60pub const CONTRACT_VERSION: u64 = 1;
61
62/// One unit of work sent to the Python worker.
63#[derive(Debug, Clone, Serialize, Deserialize)]
64pub struct TaskRequest {
65    pub version: u64,
66    /// Handler id: `llm.run_trial | data.load | data.synthetic | grade |
67    /// prompt.search | ft.lora`.
68    pub kind: String,
69    /// Backend the worker should run the body under (devbox/cluster -> Flyte lift).
70    pub backend: String,
71    pub stack: serde_json::Value,
72    pub task: serde_json::Value,
73    pub seed: u64,
74    pub workdir: String,
75    pub plan: serde_json::Value,
76}
77
78#[derive(Debug, Clone, Serialize, Deserialize)]
79pub struct TaskResult {
80    #[serde(default)]
81    pub success: bool,
82    #[serde(default)]
83    pub score: f64,
84    #[serde(default)]
85    pub cost_usd: f64,
86    #[serde(default)]
87    pub gpu_h: f64,
88    #[serde(default)]
89    pub latency_s: f64,
90    #[serde(default)]
91    pub steps: u64,
92    #[serde(default)]
93    pub trace: serde_json::Value,
94    #[serde(default)]
95    pub artifact_digests: Vec<String>,
96    #[serde(default)]
97    pub model_digests: serde_json::Value,
98}
99
100#[derive(Debug, Clone, Deserialize)]
101pub struct TaskResponse {
102    pub version: u64,
103    #[serde(default)]
104    pub ok: bool,
105    #[serde(default)]
106    pub result: Option<TaskResult>,
107    #[serde(default)]
108    pub error: Option<String>,
109}
110
111/// A `HarnessAdapter` that delegates each trial to the Python worker.
112///
113/// Thin wrapper over [`PythonWorker`] for the trial path (kind `llm.run_trial`).
114pub struct PythonAdapter {
115    worker: PythonWorker,
116}
117
118impl PythonAdapter {
119    pub fn new() -> Self {
120        PythonAdapter {
121            worker: PythonWorker::new(),
122        }
123    }
124}
125
126impl Default for PythonAdapter {
127    fn default() -> Self {
128        Self::new()
129    }
130}
131
132impl HarnessAdapter for PythonAdapter {
133    fn name(&self) -> &str {
134        "python-worker"
135    }
136
137    fn run(&self, stack: &Stack, task: &Task, workspace: &Path, seed: u64) -> Result<RunOutcome> {
138        let r = self
139            .worker
140            .call("llm.run_trial", "cluster", Some(stack), Some(task), seed, workspace, &json!({}))?;
141        Ok(RunOutcome {
142            success: r.success,
143            cost_usd: r.cost_usd,
144            latency_s: r.latency_s,
145            steps: r.steps,
146            trace: r.trace,
147        })
148    }
149}
150
151/// Shells any work item out to the Python worker over the versioned JSON contract.
152///
153/// Used by the tuner (as `PythonAdapter`) for trials and directly by `ops::optimize`
154/// for the data / prompt / fine-tune stages of the pipeline.
155///
156/// The interpreter is taken from `HL_PYTHON` (falling back to `python3`) so it can
157/// be pinned to the interpreter that actually has `harness_lock` installed (e.g. a
158/// project venv), independent of the interpreter the Rust wheel itself runs inside.
159pub struct PythonWorker {
160    interp: String,
161    module: String,
162}
163
164impl PythonWorker {
165    pub fn new() -> Self {
166        let interp = std::env::var("HL_PYTHON").unwrap_or_else(|_| "python3".to_string());
167        PythonWorker {
168            interp,
169            module: "harness_lock.tasks.worker".to_string(),
170        }
171    }
172
173    /// Invoke the worker for `kind`, writing a request into `workspace`, and return
174    /// the parsed result. `plan` carries the depth/data/axes so the worker is
175    /// depth-aware on devbox/cluster.
176    pub fn call(
177        &self,
178        kind: &str,
179        backend: &str,
180        stack: Option<&Stack>,
181        task: Option<&Task>,
182        seed: u64,
183        workspace: &Path,
184        plan: &serde_json::Value,
185    ) -> Result<TaskResult> {
186        let req_path = workspace.join("_al_request.json");
187        let out_path = workspace.join("_al_response.json");
188        if let Some(parent) = req_path.parent() {
189            std::fs::create_dir_all(parent)?;
190        }
191
192        let request = TaskRequest {
193            version: CONTRACT_VERSION,
194            kind: kind.to_string(),
195            backend: backend.to_string(),
196            stack: stack.map(|s| s.to_value()).unwrap_or(json!({})),
197            task: task.map(|t| {
198                json!({
199                    "id": t.id,
200                    "prompt_text": t.prompt_text,
201                    "expected": t.expected_artifacts(),
202                    "fixtures": t.fixtures.as_ref().map(|p| p.to_string_lossy().to_string()),
203                })
204            })
205            .unwrap_or(json!({})),
206            seed,
207            workdir: workspace.to_string_lossy().to_string(),
208            plan: plan.clone(),
209        };
210        std::fs::write(&req_path, serde_json::to_vec(&request)?)?;
211
212        let out = std::process::Command::new(&self.interp)
213            .arg("-m")
214            .arg(&self.module)
215            .arg(&req_path)
216            .arg(&out_path)
217            .output()
218            .map_err(|e| {
219                Error::Worker(format!(
220                    "could not start Python worker ({}) — is harness-lock installed for `{}`? ({e})",
221                    self.module, self.interp
222                ))
223            })?;
224
225        if !out.status.success() {
226            let stderr = String::from_utf8_lossy(&out.stderr);
227            return Err(Error::Worker(format!(
228                "python worker failed ({}): {}",
229                out.status,
230                stderr.trim()
231            )));
232        }
233        let text = std::fs::read_to_string(&out_path).map_err(|e| {
234            Error::Worker(format!("worker did not write a response file {}: {e}", out_path.display()))
235        })?;
236        let resp: TaskResponse = serde_json::from_str(&text)?;
237        if !resp.ok {
238            return Err(Error::Worker(
239                resp.error.unwrap_or_else(|| "worker reported failure".into()),
240            ));
241        }
242        resp.result.ok_or_else(|| Error::Worker("worker returned no result".into()))
243    }
244}
245
246impl Default for PythonWorker {
247    fn default() -> Self {
248        Self::new()
249    }
250}
251
harness_lock_core/worker.rs

harness_lock_core/
worker.rs
