studio_worker/
local.rs

1//! Local image generation — submit a prompt straight to the engine, no studio.
2//!
3//! Mirrors the studio job path (`ws::session::run_offered_job`): resolve the
4//! model's [`ModelSource`] (here from the local [`Catalog`] instead of a studio
5//! offer), build a [`Task::Image`], dispatch, and record the finished job — into
6//! the dedicated local-queue ring so it shows up in the app.
7
8use std::sync::atomic::{AtomicU64, Ordering};
9
10use chrono::Utc;
11
12use crate::catalog::Catalog;
13use crate::engine::Engine;
14use crate::job_run::JobRun;
15use crate::runtime::{truncate_prompt, CurrentJob, JobOutcome, JobSource, WorkerObservers};
16use crate::types::{ImageParams, Task, TaskKind, TaskResult};
17
18/// A local image-generation request. Optional fields fall back to the model's
19/// CLI defaults from the catalog.
20#[derive(Debug, Clone, Default)]
21pub struct LocalImageRequest {
22    pub prompt: String,
23    /// Model id; `None` uses the catalog's default image model.
24    pub model: Option<String>,
25    pub negative_prompt: Option<String>,
26    pub width: Option<u32>,
27    pub height: Option<u32>,
28    pub steps: Option<u32>,
29    pub seed: Option<u64>,
30    pub ext: Option<String>,
31}
32
33/// Why a local generation could not run.
34#[derive(Debug, thiserror::Error)]
35pub enum LocalError {
36    #[error("unknown model '{0}' (not in the local catalog)")]
37    UnknownModel(String),
38    #[error("no {0} model configured in the local catalog")]
39    NoModelForKind(TaskKind),
40    #[error("no image model configured in the local catalog")]
41    NoDefaultModel,
42    #[error("model '{0}' is not an image model")]
43    NotImageModel(String),
44    #[error("model '{id}' is a {got} model, not {want}")]
45    WrongKind {
46        id: String,
47        want: TaskKind,
48        got: TaskKind,
49    },
50    #[error("engine error: {0}")]
51    Engine(String),
52}
53
54pub(crate) fn next_job_id() -> String {
55    static COUNTER: AtomicU64 = AtomicU64::new(0);
56    let n = COUNTER.fetch_add(1, Ordering::Relaxed);
57    format!("local-{}-{n}", Utc::now().timestamp_millis())
58}
59
60/// Run one local image job: resolve the model from `catalog`, dispatch it on
61/// `engine`, record it in the local-queue ring, and return the image bytes.
62pub fn run_image(
63    engine: &dyn Engine,
64    catalog: &Catalog,
65    observers: &WorkerObservers,
66    req: &LocalImageRequest,
67) -> Result<TaskResult, LocalError> {
68    let model = match &req.model {
69        Some(id) => catalog
70            .get(id)
71            .ok_or_else(|| LocalError::UnknownModel(id.clone()))?,
72        None => catalog
73            .default_image_model()
74            .ok_or(LocalError::NoDefaultModel)?,
75    };
76    if model.kind != TaskKind::Image {
77        return Err(LocalError::NotImageModel(model.id.clone()));
78    }
79
80    let defaults = &model.source.cli_defaults;
81    let params = ImageParams {
82        prompt: req.prompt.clone(),
83        negative_prompt: req.negative_prompt.clone(),
84        width: req.width.unwrap_or(defaults.width).max(1),
85        height: req.height.unwrap_or(defaults.height).max(1),
86        steps: req.steps.unwrap_or(defaults.steps).max(1),
87        seed: req.seed,
88        cfg_scale: Some(defaults.cfg_scale),
89        sampling_method: defaults.sampling_method.clone(),
90        ext: req.ext.clone().unwrap_or_else(|| "webp".to_string()),
91        ..Default::default()
92    };
93
94    dispatch_and_record(engine, model, observers, &req.prompt, Task::Image(params))
95}
96
97/// Resolve a model of `kind` (an explicit id, else the catalog's
98/// default for that kind), dispatch `task`, record the local job, and
99/// return the result.  The generic core behind every non-image local
100/// endpoint (chat / tts / stt / video) so each stays a thin adapter.
101pub fn run_kind(
102    engine: &dyn Engine,
103    catalog: &Catalog,
104    observers: &WorkerObservers,
105    kind: TaskKind,
106    model_id: Option<&str>,
107    prompt_preview: &str,
108    task: Task,
109) -> Result<TaskResult, LocalError> {
110    let model = match model_id {
111        Some(id) => catalog
112            .get(id)
113            .ok_or_else(|| LocalError::UnknownModel(id.to_string()))?,
114        None => catalog
115            .default_model_for(kind)
116            .ok_or(LocalError::NoModelForKind(kind))?,
117    };
118    if model.kind != kind {
119        return Err(LocalError::WrongKind {
120            id: model.id.clone(),
121            want: kind,
122            got: model.kind,
123        });
124    }
125    dispatch_and_record(engine, model, observers, prompt_preview, task)
126}
127
128/// Serve a chat on the lane of a loaded model, recording it like any
129/// local job.  `None` when the resolved model is not loaded (or is not
130/// a chat model): the caller runs it as a transient job instead.
131pub fn chat_on_lane(
132    host: &crate::host::ModelHost,
133    catalog: &Catalog,
134    observers: &WorkerObservers,
135    model_id: Option<&str>,
136    prompt_preview: &str,
137    params: crate::types::LlmParams,
138) -> Option<Result<TaskResult, LocalError>> {
139    let model = match model_id {
140        Some(id) => catalog.get(id)?,
141        None => catalog.default_model_for(TaskKind::Llm)?,
142    };
143    if model.kind != TaskKind::Llm {
144        return None;
145    }
146    let served = host.with_lane(&model.id, |loaded, lane| {
147        let chat = loaded.as_chat()?;
148        let run = JobRun::begin(
149            observers,
150            CurrentJob {
151                job_id: next_job_id(),
152                kind: TaskKind::Llm,
153                model: model.id.clone(),
154                prompt: truncate_prompt(prompt_preview),
155                started_at: Utc::now(),
156                source: JobSource::Lane,
157            },
158        );
159        let result = run
160            .span()
161            .in_scope(|| chat.chat(params, &|| lane.cancelled()));
162        Some((run, result))
163    });
164    let (run, result) = match served {
165        Ok(Some(served)) => served,
166        // Not loaded, or loaded but not a chat model: the transient path decides.
167        Ok(None) | Err(_) => return None,
168    };
169    let result = result.map(|json| TaskResult::Llm { json });
170    run.finish(outcome_of(&result));
171    Some(result.map_err(|err| LocalError::Engine(err.to_string())))
172}
173
174/// The outcome a finished job is recorded with.
175fn outcome_of<E: std::fmt::Display>(result: &Result<TaskResult, E>) -> JobOutcome {
176    match result {
177        Ok(_) => JobOutcome::Completed,
178        Err(err) => JobOutcome::Failed {
179            reason: err.to_string(),
180        },
181    }
182}
183
184/// Dispatch `task` on `model` and record the finished job in the
185/// local-queue ring.  Shared by [`run_image`] and [`run_kind`] so the
186/// dispatch + bookkeeping lives in one place.
187fn dispatch_and_record(
188    engine: &dyn Engine,
189    model: &crate::catalog::CatalogModel,
190    observers: &WorkerObservers,
191    prompt_preview: &str,
192    task: Task,
193) -> Result<TaskResult, LocalError> {
194    let run = JobRun::begin(
195        observers,
196        CurrentJob {
197            job_id: next_job_id(),
198            kind: task.kind(),
199            model: model.id.clone(),
200            prompt: truncate_prompt(prompt_preview),
201            started_at: Utc::now(),
202            source: JobSource::Local,
203        },
204    );
205    let result = run
206        .span()
207        .in_scope(|| engine.dispatch_with_source(&model.id, task, &model.source));
208    if let Ok(result) = &result {
209        run.keep_thumbnail(result);
210    }
211    run.finish(outcome_of(&result));
212
213    result.map_err(|err| LocalError::Engine(err.to_string()))
214}
215
216#[cfg(test)]
217mod tests {
218    use super::*;
219    use crate::catalog::CatalogModel;
220    use crate::engine::SyntheticEngine;
221    use crate::types::{ModelCliDefaults, ModelEngine, ModelSource};
222
223    fn synthetic_model(id: &str, kind: TaskKind) -> CatalogModel {
224        CatalogModel {
225            id: id.into(),
226            display_name: id.into(),
227            kind,
228            vram_gb_estimate: 0.0,
229            description: None,
230            source: ModelSource {
231                engine: ModelEngine::Synthetic,
232                files: vec![],
233                cli_defaults: ModelCliDefaults {
234                    cfg_scale: 1.0,
235                    steps: 4,
236                    width: 64,
237                    height: 64,
238                    ..Default::default()
239                },
240            },
241            enabled: true,
242            origin: "local".into(),
243            exclusive_group: None,
244        }
245    }
246
247    fn catalog_with(models: Vec<CatalogModel>) -> Catalog {
248        Catalog {
249            models,
250            ..Default::default()
251        }
252    }
253
254    #[test]
255    fn generates_image_and_records_local_job() {
256        let engine = SyntheticEngine::new();
257        let catalog = catalog_with(vec![synthetic_model("synthetic-img", TaskKind::Image)]);
258        let observers = WorkerObservers::default();
259        let req = LocalImageRequest {
260            prompt: "a red fox".into(),
261            ..Default::default()
262        };
263
264        let result = run_image(&engine, &catalog, &observers, &req).unwrap();
265        match result {
266            TaskResult::Image { bytes, ext } => {
267                assert!(!bytes.is_empty());
268                assert_eq!(ext, "webp");
269            }
270            other => panic!("expected image, got {other:?}"),
271        }
272
273        let ring = observers.local_jobs.lock();
274        assert_eq!(ring.len(), 1);
275        let job = &ring[0];
276        assert_eq!(job.model, "synthetic-img");
277        assert_eq!(job.outcome, JobOutcome::Completed);
278        assert_eq!(job.prompt, "a red fox");
279        // The studio ring stays empty — local jobs are their own queue.
280        assert!(observers.recent_jobs.lock().is_empty());
281    }
282
283    #[test]
284    fn an_image_job_keeps_a_thumbnail_and_a_job_log() {
285        crate::test_support::install_job_log_capture();
286        let engine = SyntheticEngine::new();
287        let catalog = catalog_with(vec![synthetic_model("img", TaskKind::Image)]);
288        let observers = WorkerObservers::default();
289        let req = LocalImageRequest {
290            prompt: "a lighthouse".into(),
291            ..Default::default()
292        };
293        run_image(&engine, &catalog, &observers, &req).unwrap();
294
295        let job = observers.local_jobs.lock()[0].clone();
296        assert_eq!(job.source, JobSource::Local);
297        assert!(observers.thumbnails.contains(&job.job_id));
298        assert!(observers.active_jobs.lock().is_empty());
299        let log = crate::job_log::global().get(&job.job_id).expect("job log");
300        assert!(log.lines[0].message.starts_with("job started"));
301        assert!(log
302            .lines
303            .last()
304            .is_some_and(|l| l.message.starts_with("job finished")));
305    }
306
307    #[test]
308    fn a_chat_on_a_loaded_model_is_recorded_as_a_lane_job() {
309        let mut model = synthetic_model("chat", TaskKind::Llm);
310        model.source.engine = ModelEngine::LlamaCpp;
311        let catalog = catalog_with(vec![model]);
312        let shared = std::sync::Arc::new(parking_lot::Mutex::new(catalog.clone()));
313        let host = crate::host::ModelHost::new(
314            shared,
315            std::sync::Arc::new(crate::test_support::InstantRuntime),
316            std::sync::Arc::new(crate::test_support::FixedProbe(20.0)),
317            crate::residency::Residency::load_for_serving(None),
318        );
319        host.load("chat").unwrap();
320        host.wait_for(
321            "chat",
322            crate::lifecycle::ModelState::serves,
323            std::time::Duration::from_secs(5),
324        )
325        .expect("loaded");
326        let observers = WorkerObservers::default();
327        let params = crate::types::LlmParams {
328            messages: vec![crate::types::ChatMessage {
329                role: "user".into(),
330                content: "hello".into(),
331            }],
332            ..Default::default()
333        };
334
335        let answer = chat_on_lane(&host, &catalog, &observers, None, "hello", params)
336            .expect("served on the lane")
337            .unwrap();
338
339        assert!(matches!(answer, TaskResult::Llm { .. }));
340        let job = observers.local_jobs.lock()[0].clone();
341        assert_eq!(job.source, JobSource::Lane);
342        assert_eq!(job.outcome, JobOutcome::Completed);
343        assert!(observers.active_jobs.lock().is_empty());
344    }
345
346    #[test]
347    fn defaults_to_the_only_image_model_when_unspecified() {
348        let engine = SyntheticEngine::new();
349        let catalog = catalog_with(vec![synthetic_model("only-img", TaskKind::Image)]);
350        let observers = WorkerObservers::default();
351        let req = LocalImageRequest {
352            prompt: "x".into(),
353            model: None,
354            ..Default::default()
355        };
356        let out = run_image(&engine, &catalog, &observers, &req).unwrap();
357        assert!(matches!(out, TaskResult::Image { .. }));
358        assert_eq!(observers.local_jobs.lock()[0].model, "only-img");
359    }
360
361    #[test]
362    fn unknown_model_is_rejected() {
363        let engine = SyntheticEngine::new();
364        let catalog = catalog_with(vec![synthetic_model("a", TaskKind::Image)]);
365        let observers = WorkerObservers::default();
366        let req = LocalImageRequest {
367            prompt: "x".into(),
368            model: Some("missing".into()),
369            ..Default::default()
370        };
371        let err = run_image(&engine, &catalog, &observers, &req).unwrap_err();
372        assert!(matches!(err, LocalError::UnknownModel(m) if m == "missing"));
373        assert!(observers.local_jobs.lock().is_empty());
374    }
375
376    #[test]
377    fn no_image_model_yields_no_default() {
378        let engine = SyntheticEngine::new();
379        let catalog = catalog_with(vec![]);
380        let observers = WorkerObservers::default();
381        let req = LocalImageRequest {
382            prompt: "x".into(),
383            ..Default::default()
384        };
385        let err = run_image(&engine, &catalog, &observers, &req).unwrap_err();
386        assert!(matches!(err, LocalError::NoDefaultModel));
387    }
388
389    #[test]
390    fn non_image_model_is_rejected() {
391        let engine = SyntheticEngine::new();
392        let catalog = catalog_with(vec![synthetic_model("chat", TaskKind::Llm)]);
393        let observers = WorkerObservers::default();
394        let req = LocalImageRequest {
395            prompt: "x".into(),
396            model: Some("chat".into()),
397            ..Default::default()
398        };
399        let err = run_image(&engine, &catalog, &observers, &req).unwrap_err();
400        assert!(matches!(err, LocalError::NotImageModel(m) if m == "chat"));
401    }
402
403    #[test]
404    fn run_kind_dispatches_llm_and_records_the_job() {
405        let engine = SyntheticEngine::new();
406        let catalog = catalog_with(vec![synthetic_model("chat", TaskKind::Llm)]);
407        let observers = WorkerObservers::default();
408        let task = Task::Llm(crate::types::LlmParams {
409            messages: vec![crate::types::ChatMessage {
410                role: "user".into(),
411                content: "hi".into(),
412            }],
413            ..Default::default()
414        });
415        let out = run_kind(
416            &engine,
417            &catalog,
418            &observers,
419            TaskKind::Llm,
420            None,
421            "hi",
422            task,
423        )
424        .unwrap();
425        assert!(matches!(out, TaskResult::Llm { .. }));
426        assert_eq!(observers.local_jobs.lock()[0].kind, TaskKind::Llm);
427    }
428
429    #[test]
430    fn run_kind_rejects_a_wrong_kind_model() {
431        let engine = SyntheticEngine::new();
432        // An image model explicitly requested for an LLM job.
433        let catalog = catalog_with(vec![synthetic_model("img", TaskKind::Image)]);
434        let observers = WorkerObservers::default();
435        let task = Task::Llm(crate::types::LlmParams::default());
436        let err = run_kind(
437            &engine,
438            &catalog,
439            &observers,
440            TaskKind::Llm,
441            Some("img"),
442            "",
443            task,
444        )
445        .unwrap_err();
446        assert!(
447            matches!(err, LocalError::WrongKind { ref id, want, got }
448                if id == "img" && want == TaskKind::Llm && got == TaskKind::Image),
449            "got {err:?}"
450        );
451        assert!(
452            observers.local_jobs.lock().is_empty(),
453            "no job recorded on reject"
454        );
455    }
456
457    #[test]
458    fn run_kind_reports_no_model_for_kind() {
459        let engine = SyntheticEngine::new();
460        let catalog = catalog_with(vec![]);
461        let observers = WorkerObservers::default();
462        let err = run_kind(
463            &engine,
464            &catalog,
465            &observers,
466            TaskKind::AudioTts,
467            None,
468            "",
469            Task::AudioTts(crate::types::AudioTtsParams::default()),
470        )
471        .unwrap_err();
472        assert!(matches!(
473            err,
474            LocalError::NoModelForKind(TaskKind::AudioTts)
475        ));
476    }
477}
studio_worker/local.rs

studio_worker/
local.rs
