Skip to main content

mold_server/
queue.rs

1use std::collections::VecDeque;
2use std::sync::Arc;
3
4use base64::Engine as _;
5use mold_core::{
6    ImageData, OutputFormat, OutputMetadata, SseCompleteEvent, SseErrorEvent, SseProgressEvent,
7};
8use mold_db::{GenerationRecord, MetadataDb, RecordSource};
9use sha2::{Digest, Sha256};
10use std::sync::atomic::Ordering;
11use std::time::{Instant, SystemTime, UNIX_EPOCH};
12
13use crate::gpu_pool::GpuJob;
14use crate::model_manager;
15use crate::state::{
16    ActiveGenerationSnapshot, AppState, GenerationJob, GenerationJobResult, SseMessage,
17};
18
19/// Convert an inference-crate progress event to an SSE wire event.
20fn progress_to_sse(event: mold_inference::ProgressEvent) -> SseProgressEvent {
21    event.into()
22}
23
24/// Strips backtrace frames from candle error messages.
25///
26/// Renders the full anyhow cause chain (`{:#}`) so wrappers like
27/// `with_context("mmap single-file checkpoint at …")` carry their root cause
28/// through to the wire — otherwise users see the outer wrapper only.
29pub(crate) fn clean_error_message(e: &anyhow::Error) -> String {
30    let full = format!("{e:#}");
31    let mut lines: Vec<&str> = Vec::new();
32    for line in full.lines() {
33        let trimmed = line.trim_start();
34        if (trimmed.starts_with("0:") || trimmed.starts_with("1:"))
35            && trimmed.len() > 3
36            && trimmed
37                .as_bytes()
38                .first()
39                .is_some_and(|b| b.is_ascii_digit())
40        {
41            break;
42        }
43        if trimmed.len() > 2
44            && trimmed.as_bytes()[0].is_ascii_digit()
45            && trimmed.contains("::")
46            && trimmed.contains("at ")
47        {
48            break;
49        }
50        lines.push(line);
51    }
52    let msg = lines.join("\n").trim().to_string();
53    if msg.is_empty() {
54        format!("{}", e.root_cause())
55    } else {
56        msg
57    }
58}
59
60fn set_active_generation(state: &AppState, model: &str, prompt: &str) {
61    let prompt_sha256 = format!("{:x}", Sha256::digest(prompt.as_bytes()));
62    let started_at_unix_ms = SystemTime::now()
63        .duration_since(UNIX_EPOCH)
64        .unwrap_or_default()
65        .as_millis() as u64;
66
67    let mut active = state
68        .active_generation
69        .write()
70        .unwrap_or_else(|e| e.into_inner());
71    *active = Some(ActiveGenerationSnapshot {
72        model: model.to_string(),
73        prompt_sha256,
74        started_at_unix_ms,
75        started_at: Instant::now(),
76    });
77}
78
79fn clear_active_generation(state: &AppState) {
80    let mut active = state
81        .active_generation
82        .write()
83        .unwrap_or_else(|e| e.into_inner());
84    *active = None;
85}
86
87/// Save an image to disk and (best-effort) record a row in the metadata DB.
88///
89/// Errors writing to disk are logged and skipped. DB errors are also logged
90/// but do not fail the save — the file is the source of truth.
91///
92/// Shared between the legacy single-GPU `process_job` (this file) and the
93/// per-GPU worker (`gpu_worker.rs`). Keep these on one helper so the DB
94/// upsert can never silently regress on one path while the other keeps
95/// working.
96pub(crate) fn save_image_to_dir(
97    dir: &std::path::Path,
98    img: &mold_core::ImageData,
99    model: &str,
100    batch_size: u32,
101    metadata: Option<&OutputMetadata>,
102    generation_time_ms: Option<i64>,
103    db: Option<&MetadataDb>,
104) {
105    if let Err(e) = std::fs::create_dir_all(dir) {
106        tracing::warn!("failed to create output dir {}: {e}", dir.display());
107        return;
108    }
109    let now = SystemTime::now()
110        .duration_since(UNIX_EPOCH)
111        .unwrap_or_default();
112    let timestamp_ms = now.as_millis() as u64;
113    let ext = img.format.to_string();
114    let filename =
115        mold_core::default_output_filename(model, timestamp_ms, &ext, batch_size, img.index);
116    let path = dir.join(&filename);
117    match std::fs::write(&path, &img.data) {
118        Ok(()) => tracing::info!("saved image to {}", path.display()),
119        Err(e) => {
120            tracing::warn!("failed to save image to {}: {e}", path.display());
121            return;
122        }
123    }
124    if let (Some(db), Some(meta)) = (db, metadata) {
125        let mut rec = GenerationRecord::from_save(
126            dir,
127            filename,
128            img.format,
129            meta.clone(),
130            RecordSource::Server,
131            now.as_millis() as i64,
132        );
133        rec.stat_from_disk(&path);
134        rec.generation_time_ms = generation_time_ms;
135        rec.hostname = hostname_string();
136        rec.backend = current_backend_label();
137        if let Err(e) = db.upsert(&rec) {
138            tracing::warn!("metadata DB upsert failed for {}: {e:#}", rec.filename);
139        }
140    }
141}
142
143/// Save a video file to disk and (best-effort) record its metadata row.
144/// Mirrors `save_image_to_dir` for the video-output path. See that helper
145/// for the multi-path-callers note.
146///
147/// When `gif_preview` is non-empty, also persists
148/// `$MOLD_HOME/cache/previews/<filename>.preview.gif`. The gallery preview
149/// endpoint (`GET /api/gallery/preview/:filename`) streams from that path
150/// so remote TUI clients can animate the detail pane without re-fetching
151/// the full MP4.
152#[allow(clippy::too_many_arguments)]
153pub(crate) fn save_video_to_dir(
154    dir: &std::path::Path,
155    bytes: &[u8],
156    gif_preview: &[u8],
157    format: OutputFormat,
158    model: &str,
159    metadata: &OutputMetadata,
160    generation_time_ms: Option<i64>,
161    db: Option<&MetadataDb>,
162) {
163    if let Err(e) = std::fs::create_dir_all(dir) {
164        tracing::warn!("failed to create output dir {}: {e}", dir.display());
165        return;
166    }
167    let now = SystemTime::now()
168        .duration_since(UNIX_EPOCH)
169        .unwrap_or_default();
170    let ts = now.as_millis() as u64;
171    let ext = format.extension();
172    let filename = mold_core::default_output_filename(model, ts, ext, 1, 0);
173    let path = dir.join(&filename);
174    if let Err(e) = std::fs::write(&path, bytes) {
175        tracing::error!("failed to save video to {}: {e}", path.display());
176        return;
177    }
178    if !gif_preview.is_empty() {
179        save_video_preview_gif(&filename, gif_preview);
180    }
181    if let Some(db) = db {
182        let mut rec = GenerationRecord::from_save(
183            dir,
184            filename,
185            format,
186            metadata.clone(),
187            RecordSource::Server,
188            now.as_millis() as i64,
189        );
190        rec.stat_from_disk(&path);
191        rec.generation_time_ms = generation_time_ms;
192        rec.hostname = hostname_string();
193        rec.backend = current_backend_label();
194        if let Err(e) = db.upsert(&rec) {
195            tracing::warn!("metadata DB upsert failed for {}: {e:#}", rec.filename);
196        }
197    }
198}
199
200/// Persist a video's `.preview.gif` sidecar to the server's preview cache
201/// (`$MOLD_HOME/cache/previews/<filename>.preview.gif`). Best-effort —
202/// warnings log and return so a failure here never fails the save path.
203///
204/// Shared with the multi-GPU worker path (`gpu_worker::process_job`) so
205/// video outputs land a preview regardless of which save flow wrote the
206/// MP4; otherwise `/api/gallery/preview/:filename` would 404 whenever the
207/// server is running with GPU workers enabled.
208pub(crate) fn save_video_preview_gif(filename: &str, gif_bytes: &[u8]) {
209    let preview_dir = mold_core::Config::mold_dir()
210        .unwrap_or_else(|| std::path::PathBuf::from(".mold"))
211        .join("cache")
212        .join("previews");
213    save_video_preview_gif_to(&preview_dir, filename, gif_bytes);
214}
215
216/// Testable inner of [`save_video_preview_gif`] that accepts an explicit
217/// preview directory (lets unit tests exercise the write path without
218/// racing on the `MOLD_HOME` env var).
219fn save_video_preview_gif_to(preview_dir: &std::path::Path, filename: &str, gif_bytes: &[u8]) {
220    if let Err(e) = std::fs::create_dir_all(preview_dir) {
221        tracing::warn!(
222            "failed to create preview cache dir {}: {e}",
223            preview_dir.display()
224        );
225        return;
226    }
227    let preview_path = preview_dir.join(format!("{filename}.preview.gif"));
228    if let Err(e) = std::fs::write(&preview_path, gif_bytes) {
229        tracing::warn!(
230            "failed to write preview gif {}: {e}",
231            preview_path.display()
232        );
233    }
234}
235
236/// Best-effort hostname for the `hostname` DB column. Falls back to `None`.
237fn hostname_string() -> Option<String> {
238    hostname::get().ok().and_then(|s| s.into_string().ok())
239}
240
241/// Compile-time backend label so DB rows say where the work happened.
242fn current_backend_label() -> Option<String> {
243    if cfg!(feature = "cuda") {
244        Some("cuda".into())
245    } else if cfg!(feature = "metal") {
246        Some("metal".into())
247    } else {
248        Some("cpu".into())
249    }
250}
251
252/// Build the SSE `complete` wire event from a finished generation response.
253///
254/// Video responses encode the actual video bytes (MP4/GIF/APNG/WebP) as the
255/// payload and populate every `video_*` metadata field; image responses
256/// encode the image bytes with the video fields cleared. `img` is the
257/// `ImageData` chosen by the caller — either the first generated image or an
258/// `ImageData` synthesized from the video thumbnail (the single-primary-image
259/// shape that the internal `GenerationJobResult` still expects).
260///
261/// Shared between the single-GPU path (`process_job` in this file) and the
262/// multi-GPU path (`gpu_worker::process_job`) so the two can never drift on
263/// which `video_*` fields are populated. Before this helper existed the
264/// multi-GPU worker always encoded the thumbnail PNG as the payload and
265/// hard-coded every `video_*` field to `None`, which silently degraded every
266/// LTX-Video / LTX-2 generation into an image response on hosts with at
267/// least one GPU worker.
268pub(crate) fn build_sse_complete_event(
269    response: &mold_core::GenerateResponse,
270    img: &mold_core::ImageData,
271) -> SseCompleteEvent {
272    let b64 = base64::engine::general_purpose::STANDARD;
273    if let Some(ref video) = response.video {
274        SseCompleteEvent {
275            image: b64.encode(&video.data),
276            format: video.format,
277            width: video.width,
278            height: video.height,
279            seed_used: response.seed_used,
280            generation_time_ms: response.generation_time_ms,
281            model: response.model.clone(),
282            video_frames: Some(video.frames),
283            video_fps: Some(video.fps),
284            video_thumbnail: Some(b64.encode(&video.thumbnail)),
285            video_gif_preview: if video.gif_preview.is_empty() {
286                None
287            } else {
288                Some(b64.encode(&video.gif_preview))
289            },
290            video_has_audio: video.has_audio,
291            video_duration_ms: video.duration_ms,
292            video_audio_sample_rate: video.audio_sample_rate,
293            video_audio_channels: video.audio_channels,
294            gpu: response.gpu,
295        }
296    } else {
297        SseCompleteEvent {
298            image: b64.encode(&img.data),
299            format: img.format,
300            width: img.width,
301            height: img.height,
302            seed_used: response.seed_used,
303            generation_time_ms: response.generation_time_ms,
304            model: response.model.clone(),
305            video_frames: None,
306            video_fps: None,
307            video_thumbnail: None,
308            video_gif_preview: None,
309            video_has_audio: false,
310            video_duration_ms: None,
311            video_audio_sample_rate: None,
312            video_audio_channels: None,
313            gpu: response.gpu,
314        }
315    }
316}
317
318/// Runs the generation queue worker loop. Processes one job at a time (FIFO),
319/// but uses a small bounded lookahead buffer to prefer jobs whose model is
320/// already loaded — minimizing model swaps when the queue interleaves models.
321/// Exits when the sender half of the channel is dropped (server shutdown).
322pub async fn run_queue_worker(
323    mut job_rx: tokio::sync::mpsc::Receiver<GenerationJob>,
324    state: AppState,
325) {
326    tracing::debug!("generation queue worker started");
327    let buffer_size = resolve_lookahead_buffer();
328    let max_deferrals = resolve_max_deferrals();
329    let mut buffer: VecDeque<BufferedJob> = VecDeque::with_capacity(buffer_size);
330
331    loop {
332        if buffer.is_empty() {
333            match job_rx.recv().await {
334                Some(j) => buffer.push_back(BufferedJob::new(j)),
335                None => break,
336            }
337        }
338        // Top up the buffer without blocking — drain the channel up to capacity.
339        top_up_buffer(&mut buffer, &mut job_rx, buffer_size);
340
341        let loaded = single_gpu_loaded_models(&state).await;
342        let job = pick_next_job(&mut buffer, &loaded, max_deferrals);
343        let job_id = job.id.clone();
344
345        #[cfg(feature = "metrics")]
346        crate::metrics::record_queue_depth(state.queue.pending());
347        process_job(&state, job).await;
348        state.queue.decrement();
349        // Drop the registry entry on every terminal path — the worker
350        // here doesn't own a drop guard, so we do it inline alongside
351        // the queue counter decrement.
352        state.job_registry.remove(&job_id);
353        #[cfg(feature = "metrics")]
354        crate::metrics::record_queue_depth(state.queue.pending());
355    }
356    tracing::info!("generation queue worker shutting down");
357}
358
359async fn single_gpu_loaded_models(state: &AppState) -> std::collections::HashSet<String> {
360    let mut set = std::collections::HashSet::new();
361    let cache = state.model_cache.lock().await;
362    if let Some(name) = cache.active_model() {
363        set.insert(name.to_string());
364    }
365    set
366}
367
368/// Build the set of "currently loaded somewhere" model names across every
369/// worker in the multi-GPU pool. A worker counts the model as loaded if
370/// either it's in the worker's cache as Gpu-resident OR it's the worker's
371/// `active_generation` (covering the take-and-restore window where the
372/// cache entry briefly disappears).
373fn multi_gpu_loaded_models(state: &AppState) -> std::collections::HashSet<String> {
374    let mut set = std::collections::HashSet::new();
375    for worker in &state.gpu_pool.workers {
376        if let Ok(active_gen) = worker.active_generation.read() {
377            if let Some(g) = active_gen.as_ref() {
378                set.insert(g.model.clone());
379            }
380        }
381        if let Ok(cache) = worker.model_cache.lock() {
382            if let Some(name) = cache.active_model() {
383                set.insert(name.to_string());
384            }
385        }
386    }
387    set
388}
389
390/// In-flight wrapper that tracks how many times the picker has skipped this
391/// job. Once the count exceeds `max_deferrals`, the picker force-dispatches
392/// it to bound starvation.
393pub(crate) struct BufferedJob {
394    pub(crate) job: GenerationJob,
395    pub(crate) deferred: usize,
396}
397
398impl BufferedJob {
399    fn new(job: GenerationJob) -> Self {
400        Self { job, deferred: 0 }
401    }
402}
403
404/// Drain the receive channel into the lookahead buffer, capped at
405/// `buffer_size`. Returns when the buffer is full or the channel has no
406/// immediately-available jobs (the receiver is unchanged on `Empty`). Pure
407/// helper extracted so tests can lock in the cap as a load-bearing invariant
408/// without spinning up the full async dispatcher.
409pub(crate) fn top_up_buffer(
410    buffer: &mut VecDeque<BufferedJob>,
411    job_rx: &mut tokio::sync::mpsc::Receiver<GenerationJob>,
412    buffer_size: usize,
413) {
414    while buffer.len() < buffer_size {
415        match job_rx.try_recv() {
416            Ok(j) => buffer.push_back(BufferedJob::new(j)),
417            Err(_) => break,
418        }
419    }
420}
421
422/// Pure picker for the lookahead buffer. Selects the buffered job whose
423/// model is already loaded somewhere in `loaded`; ties broken by arrival
424/// order (front of the deque wins). The head's `deferred` count bounds
425/// starvation: if the head has been skipped `max_deferrals` times, it wins
426/// regardless of `loaded` membership.
427///
428/// The returned job is removed from the buffer; remaining buffered jobs that
429/// were skipped have their `deferred` count incremented. Increments
430/// `mold_queue_reorders_total` whenever a non-head job is picked.
431pub(crate) fn pick_next_job(
432    buffer: &mut VecDeque<BufferedJob>,
433    loaded: &std::collections::HashSet<String>,
434    max_deferrals: usize,
435) -> GenerationJob {
436    debug_assert!(
437        !buffer.is_empty(),
438        "pick_next_job requires non-empty buffer"
439    );
440
441    // Force-dispatch the head if it's hit the starvation budget.
442    if let Some(head) = buffer.front() {
443        if head.deferred >= max_deferrals {
444            return buffer.pop_front().expect("checked non-empty").job;
445        }
446    }
447
448    // Find the front-most buffered job whose model is already loaded.
449    let pick_idx = buffer
450        .iter()
451        .position(|b| loaded.contains(&b.job.request.model))
452        .unwrap_or(0);
453
454    if pick_idx > 0 {
455        for (i, b) in buffer.iter_mut().enumerate() {
456            if i < pick_idx {
457                b.deferred += 1;
458            }
459        }
460        let model = buffer[pick_idx].job.request.model.clone();
461        tracing::debug!(
462            picked_model = %model,
463            head_model = %buffer.front().map(|b| b.job.request.model.as_str()).unwrap_or(""),
464            picked_index = pick_idx,
465            "queue reorder picked non-head job"
466        );
467        #[cfg(feature = "metrics")]
468        crate::metrics::record_queue_reorder();
469    }
470
471    buffer.remove(pick_idx).expect("pick_idx in range").job
472}
473
474pub(crate) const DEFAULT_LOOKAHEAD_BUFFER: usize = 8;
475pub(crate) const DEFAULT_MAX_DEFERRALS: usize = 3;
476pub(crate) const LOOKAHEAD_BUFFER_ENV: &str = "MOLD_QUEUE_LOOKAHEAD_BUFFER";
477pub(crate) const MAX_DEFERRALS_ENV: &str = "MOLD_QUEUE_MAX_DEFERRALS";
478const LOOKAHEAD_BUFFER_LOWER: usize = 1;
479const LOOKAHEAD_BUFFER_UPPER: usize = 64;
480const MAX_DEFERRALS_UPPER: usize = 32;
481
482/// Resolve the lookahead buffer size from env, falling back to the default.
483/// Out-of-range or unparseable values log a warning and use the default —
484/// matching the warn-then-default pattern of `resolve_max_cached_models`.
485pub(crate) fn resolve_lookahead_buffer() -> usize {
486    match std::env::var(LOOKAHEAD_BUFFER_ENV) {
487        Ok(raw) => match raw.trim().parse::<usize>() {
488            Ok(n) if (LOOKAHEAD_BUFFER_LOWER..=LOOKAHEAD_BUFFER_UPPER).contains(&n) => n,
489            Ok(n) => {
490                tracing::warn!(
491                    env = LOOKAHEAD_BUFFER_ENV,
492                    value = n,
493                    lower = LOOKAHEAD_BUFFER_LOWER,
494                    upper = LOOKAHEAD_BUFFER_UPPER,
495                    "ignoring out-of-range queue lookahead buffer; using default"
496                );
497                DEFAULT_LOOKAHEAD_BUFFER
498            }
499            Err(e) => {
500                tracing::warn!(
501                    env = LOOKAHEAD_BUFFER_ENV,
502                    raw = %raw,
503                    error = %e,
504                    "ignoring unparseable queue lookahead buffer; using default"
505                );
506                DEFAULT_LOOKAHEAD_BUFFER
507            }
508        },
509        Err(_) => DEFAULT_LOOKAHEAD_BUFFER,
510    }
511}
512
513/// Resolve the max-deferrals starvation budget from env. Out-of-range or
514/// unparseable values log a warning and use the default.
515pub(crate) fn resolve_max_deferrals() -> usize {
516    match std::env::var(MAX_DEFERRALS_ENV) {
517        Ok(raw) => match raw.trim().parse::<usize>() {
518            Ok(n) if n <= MAX_DEFERRALS_UPPER => n,
519            Ok(n) => {
520                tracing::warn!(
521                    env = MAX_DEFERRALS_ENV,
522                    value = n,
523                    upper = MAX_DEFERRALS_UPPER,
524                    "ignoring out-of-range queue max-deferrals; using default"
525                );
526                DEFAULT_MAX_DEFERRALS
527            }
528            Err(e) => {
529                tracing::warn!(
530                    env = MAX_DEFERRALS_ENV,
531                    raw = %raw,
532                    error = %e,
533                    "ignoring unparseable queue max-deferrals; using default"
534                );
535                DEFAULT_MAX_DEFERRALS
536            }
537        },
538        Err(_) => DEFAULT_MAX_DEFERRALS,
539    }
540}
541
542async fn process_job(state: &AppState, job: GenerationJob) {
543    // Check if client already disconnected before doing any work
544    if job.result_tx.is_closed() {
545        tracing::debug!("skipping queued job — client disconnected");
546        return;
547    }
548
549    // Single-GPU path: there's only one slot. `gpu=None` keeps the wire
550    // shape consistent with multi-GPU even when we don't know the ordinal.
551    state.job_registry.mark_running(&job.id, None);
552
553    // Send "now processing" event (position 0). `id` echoes the
554    // server-assigned UUID so reconnecting clients can match progress
555    // updates to their persisted card.
556    if let Some(ref tx) = job.progress_tx {
557        let _ = tx.send(SseMessage::Progress(SseProgressEvent::Queued {
558            position: 0,
559            id: job.id.clone(),
560        }));
561    }
562
563    // 1. Ensure model is ready (with progress forwarding)
564    let progress_callback = job.progress_tx.as_ref().map(|tx| {
565        let tx = tx.clone();
566        Arc::new(move |event: mold_inference::ProgressEvent| {
567            let _ = tx.send(SseMessage::Progress(progress_to_sse(event)));
568        }) as model_manager::EngineProgressCallback
569    });
570
571    let activation_hint = model_manager::activation_hint_for_request(state, &job.request).await;
572    if let Err(api_err) = model_manager::ensure_model_ready(
573        state,
574        &job.request.model,
575        progress_callback,
576        activation_hint,
577    )
578    .await
579    {
580        let err_msg = api_err.error.clone();
581        if let Some(ref tx) = job.progress_tx {
582            let _ = tx.send(SseMessage::Error(SseErrorEvent {
583                message: err_msg.clone(),
584            }));
585        }
586        let _ = job.result_tx.send(Err(err_msg));
587        return;
588    }
589
590    // 2. Low-memory warning (MPS/unified memory only — observability aid)
591    #[cfg(target_os = "macos")]
592    if let Some(available) = mold_inference::device::available_system_memory_bytes() {
593        if available < 1_000_000_000 {
594            tracing::warn!(
595                available_mb = available / 1_000_000,
596                "low memory before inference — system may become unstable"
597            );
598        }
599    }
600
601    // 3. Take the engine out of the cache so the cache mutex stays free during
602    //    generation. Mirrors the multi-GPU `gpu_worker::process_job` pattern —
603    //    holding the cache lock through inference would block /api/models,
604    //    /api/cache, and any concurrent gallery/admin reads.
605    let taken = {
606        let mut cache = state.model_cache.lock().await;
607        cache.take(&job.request.model)
608    };
609    let Some(mut cached_engine) = taken else {
610        let err_msg = "no engine available after model readiness check".to_string();
611        if let Some(ref tx) = job.progress_tx {
612            let _ = tx.send(SseMessage::Error(SseErrorEvent {
613                message: err_msg.clone(),
614            }));
615        }
616        let _ = job.result_tx.send(Err(err_msg));
617        return;
618    };
619
620    let active_gen = state.active_generation.clone();
621    let gen_req = job.request.clone();
622    let progress_tx = job.progress_tx.clone();
623
624    set_active_generation(state, &job.request.model, &job.request.prompt);
625
626    // Install progress callback before crossing into spawn_blocking — keeps
627    // the callback installation off the blocking thread. Mirrors the pre-
628    // refactor behavior: when streaming, set the callback; when not, clear
629    // it (and only clear after generate when streaming).
630    let was_streaming = progress_tx.is_some();
631    if let Some(ref ptx) = progress_tx {
632        let ptx = ptx.clone();
633        cached_engine.engine.set_on_progress(Box::new(move |event| {
634            let _ = ptx.send(SseMessage::Progress(progress_to_sse(event)));
635        }));
636    } else {
637        cached_engine.engine.clear_on_progress();
638    }
639
640    #[cfg(feature = "metrics")]
641    let inference_start = Instant::now();
642    // RSS sample taken just before inference; the post-inference sample below
643    // logs the per-job delta so RAM growth can be attributed to a specific
644    // generation rather than tracked at process granularity.
645    let rss_before = crate::resources::ram_snapshot().used_by_mold;
646    // Run generation on the blocking pool. Move the engine in, return it back
647    // out (alongside the result + any panic payload) so we can restore it to
648    // the cache in async context regardless of outcome.
649    let join_result = tokio::task::spawn_blocking(move || {
650        let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
651            cached_engine.engine.generate(&gen_req)
652        }));
653        if was_streaming {
654            cached_engine.engine.clear_on_progress();
655        }
656        (cached_engine, result)
657    })
658    .await;
659
660    let rss_after = crate::resources::ram_snapshot().used_by_mold;
661    let rss_delta = rss_after as i64 - rss_before as i64;
662    tracing::info!(
663        model = %job.request.model,
664        rss_before_mb = rss_before / 1_000_000,
665        rss_after_mb = rss_after / 1_000_000,
666        rss_delta_mb = rss_delta / 1_000_000,
667        "generation memory delta"
668    );
669
670    #[cfg(feature = "metrics")]
671    let inference_duration = inference_start.elapsed().as_secs_f64();
672
673    // Restore the engine to the cache as soon as the blocking task joins —
674    // even panics must restore so the cache isn't left with a hole. If the
675    // tokio task itself failed (JoinError), the engine is gone — restoration
676    // is impossible. Without `clear_in_flight` the model name would leak
677    // forever in `in_flight`, so `ensure_model_ready` keeps fast-pathing
678    // through `contains()` while every subsequent `take()` returns `None`,
679    // permanently jamming this model. Clear the marker so the cache will
680    // legitimately re-load the engine on the next request.
681    let result = match join_result {
682        Ok((cached_engine, panic_or_result)) => {
683            {
684                let mut cache = state.model_cache.lock().await;
685                cache.restore(cached_engine);
686            }
687            clear_active_generation(state);
688            Ok(panic_or_result)
689        }
690        Err(join_err) => {
691            {
692                let mut cache = state.model_cache.lock().await;
693                cache.clear_in_flight(&job.request.model);
694            }
695            clear_active_generation(state);
696            Err(join_err)
697        }
698    };
699
700    match result {
701        Ok(Ok(Ok(mut response))) => {
702            #[cfg(feature = "metrics")]
703            crate::metrics::record_generation(&job.request.model, inference_duration);
704
705            if response.images.is_empty() && response.video.is_none() {
706                let err_msg = "generation error: engine returned no images or video".to_string();
707                if let Some(ref tx) = job.progress_tx {
708                    let _ = tx.send(SseMessage::Error(SseErrorEvent {
709                        message: err_msg.clone(),
710                    }));
711                }
712                let _ = job.result_tx.send(Err(err_msg));
713                return;
714            }
715            // For video-only responses, synthesize an ImageData from the thumbnail
716            // so the existing queue/SSE pipeline can handle it.
717            let img = if !response.images.is_empty() {
718                response.images.remove(0)
719            } else if let Some(ref video) = response.video {
720                ImageData {
721                    data: video.thumbnail.clone(),
722                    format: OutputFormat::Png,
723                    width: video.width,
724                    height: video.height,
725                    index: 0,
726                }
727            } else {
728                unreachable!("checked above");
729            };
730
731            // Save to output directory if configured.
732            // Builds OutputMetadata from the request + the engine's actual
733            // seed_used so the DB and embedded chunks agree.
734            if let Some(ref dir) = job.output_dir {
735                let dir = dir.clone();
736                let model = job.request.model.clone();
737                let batch_size = job.request.batch_size;
738                let generation_time_ms = response.generation_time_ms as i64;
739                let metadata = OutputMetadata::from_generate_request(
740                    &job.request,
741                    response.seed_used,
742                    None,
743                    mold_core::build_info::version_string(),
744                );
745                let db = state.metadata_db.clone();
746                if let Some(ref video) = response.video {
747                    let video_data = video.data.clone();
748                    let video_gif_preview = video.gif_preview.clone();
749                    let video_format = video.format;
750                    let video_metadata = metadata.clone();
751                    tokio::task::spawn_blocking(move || {
752                        save_video_to_dir(
753                            &dir,
754                            &video_data,
755                            &video_gif_preview,
756                            video_format,
757                            &model,
758                            &video_metadata,
759                            Some(generation_time_ms),
760                            db.as_ref().as_ref(),
761                        );
762                    });
763                } else {
764                    let img_clone = img.clone();
765                    let metadata_clone = metadata.clone();
766                    tokio::task::spawn_blocking(move || {
767                        save_image_to_dir(
768                            &dir,
769                            &img_clone,
770                            &model,
771                            batch_size,
772                            Some(&metadata_clone),
773                            Some(generation_time_ms),
774                            db.as_ref().as_ref(),
775                        );
776                    });
777                }
778            }
779
780            // Send SSE complete event
781            if let Some(ref tx) = job.progress_tx {
782                let event = build_sse_complete_event(&response, &img);
783                let _ = tx.send(SseMessage::Complete(event));
784            }
785
786            // Send result through oneshot
787            let _ = job.result_tx.send(Ok(GenerationJobResult {
788                image: img,
789                response,
790            }));
791        }
792        Ok(Ok(Err(e))) => {
793            #[cfg(feature = "metrics")]
794            crate::metrics::record_generation_error(&job.request.model);
795
796            *active_gen.write().unwrap_or_else(|e| e.into_inner()) = None;
797            tracing::error!("generation error: {e:#}");
798            let err_msg = format!("generation error: {}", clean_error_message(&e));
799            if let Some(ref tx) = job.progress_tx {
800                let _ = tx.send(SseMessage::Error(SseErrorEvent {
801                    message: err_msg.clone(),
802                }));
803            }
804            let _ = job.result_tx.send(Err(err_msg));
805        }
806        Ok(Err(panic_payload)) => {
807            #[cfg(feature = "metrics")]
808            crate::metrics::record_generation_error(&job.request.model);
809
810            *active_gen.write().unwrap_or_else(|e| e.into_inner()) = None;
811            let msg = panic_payload
812                .downcast_ref::<String>()
813                .map(|s| s.as_str())
814                .or_else(|| panic_payload.downcast_ref::<&str>().copied())
815                .unwrap_or("unknown panic");
816            tracing::error!("inference panicked: {msg}");
817            let err_msg = format!("inference panicked: {msg}");
818            if let Some(ref tx) = job.progress_tx {
819                let _ = tx.send(SseMessage::Error(SseErrorEvent {
820                    message: err_msg.clone(),
821                }));
822            }
823            let _ = job.result_tx.send(Err(err_msg));
824        }
825        Err(join_err) => {
826            #[cfg(feature = "metrics")]
827            crate::metrics::record_generation_error(&job.request.model);
828
829            *active_gen.write().unwrap_or_else(|e| e.into_inner()) = None;
830            tracing::error!("inference task join error: {join_err:?}");
831            let err_msg = "inference task failed".to_string();
832            if let Some(ref tx) = job.progress_tx {
833                let _ = tx.send(SseMessage::Error(SseErrorEvent {
834                    message: err_msg.clone(),
835                }));
836            }
837            let _ = job.result_tx.send(Err(err_msg));
838        }
839    }
840}
841
842// ── Multi-GPU queue dispatcher ──────────────────────────────────────────────
843
844/// Runs the multi-GPU dispatch loop. Routes each generation job to the best
845/// GPU worker based on the placement strategy (model-loaded > idle > evict LRU).
846/// Uses a small lookahead buffer so an interleaved queue (`[A, B, A, B]`)
847/// doesn't force a sibling worker to swap models when one already has the
848/// right one warm.
849///
850/// Exits when the sender half of the channel is dropped (server shutdown).
851pub async fn run_queue_dispatcher(
852    mut job_rx: tokio::sync::mpsc::Receiver<GenerationJob>,
853    state: AppState,
854) {
855    tracing::debug!("multi-GPU queue dispatcher started");
856    let buffer_size = resolve_lookahead_buffer();
857    let max_deferrals = resolve_max_deferrals();
858    let mut buffer: VecDeque<BufferedJob> = VecDeque::with_capacity(buffer_size);
859
860    loop {
861        if buffer.is_empty() {
862            match job_rx.recv().await {
863                Some(j) => buffer.push_back(BufferedJob::new(j)),
864                None => break,
865            }
866        }
867        top_up_buffer(&mut buffer, &mut job_rx, buffer_size);
868
869        let loaded = multi_gpu_loaded_models(&state);
870        let job = pick_next_job(&mut buffer, &loaded, max_deferrals);
871
872        #[cfg(feature = "metrics")]
873        crate::metrics::record_queue_depth(state.queue.pending());
874
875        let job_id = job.id.clone();
876        let model_name = job.request.model.clone();
877        let estimated_vram = estimate_model_vram(&model_name);
878
879        if let Some(err_msg) = crate::gpu_pool::model_unschedulable_message(&model_name) {
880            tracing::warn!(model = %model_name, "{err_msg}");
881            if let Some(tx) = job.progress_tx {
882                let _ = tx.send(SseMessage::Error(SseErrorEvent {
883                    message: err_msg.clone(),
884                }));
885            }
886            let _ = job.result_tx.send(Err(err_msg));
887            state.queue.decrement();
888            state.job_registry.remove(&job_id);
889            #[cfg(feature = "metrics")]
890            crate::metrics::record_queue_depth(state.queue.pending());
891            continue;
892        }
893
894        let preferred_gpu = match state
895            .gpu_pool
896            .resolve_explicit_placement_gpu(job.request.placement.as_ref())
897        {
898            Ok(ordinal) => ordinal,
899            Err(err_msg) => {
900                tracing::warn!(model = %model_name, "{err_msg}");
901                if let Some(tx) = job.progress_tx {
902                    let _ = tx.send(SseMessage::Error(SseErrorEvent {
903                        message: err_msg.clone(),
904                    }));
905                }
906                let _ = job.result_tx.send(Err(err_msg));
907                state.queue.decrement();
908                state.job_registry.remove(&job_id);
909                #[cfg(feature = "metrics")]
910                crate::metrics::record_queue_depth(state.queue.pending());
911                continue;
912            }
913        };
914
915        if job.result_tx.is_closed() {
916            tracing::debug!(model = %model_name, "skipping queued multi-GPU job — client disconnected");
917            state.queue.decrement();
918            state.job_registry.remove(&job_id);
919            #[cfg(feature = "metrics")]
920            crate::metrics::record_queue_depth(state.queue.pending());
921            continue;
922        }
923
924        // Build the GpuJob once; the retry loop moves it between attempts.
925        let mut gpu_job = Some(GpuJob {
926            id: job.id.clone(),
927            model: model_name.clone(),
928            request: job.request,
929            progress_tx: job.progress_tx,
930            result_tx: job.result_tx,
931            output_dir: job.output_dir,
932            config: state.config.clone(),
933            metadata_db: state.metadata_db.clone(),
934            queue: state.queue.clone(),
935            registry: state.job_registry.clone(),
936        });
937
938        let mut skip: Vec<usize> = if preferred_gpu.is_none() {
939            let failed = crate::gpu_pool::failed_ordinals_for_model(&model_name);
940            if failed.len() < state.gpu_pool.worker_count() {
941                failed
942            } else {
943                Vec::new()
944            }
945        } else {
946            Vec::new()
947        };
948        let mut dispatched = false;
949
950        while !dispatched {
951            if gpu_job
952                .as_ref()
953                .is_some_and(|pending| pending.result_tx.is_closed())
954            {
955                tracing::debug!(
956                    model = %model_name,
957                    "dropping queued multi-GPU job before dispatch — client disconnected"
958                );
959                state.queue.decrement();
960                state.job_registry.remove(&job_id);
961                break;
962            }
963
964            let worker = if let Some(ordinal) = preferred_gpu {
965                state.gpu_pool.worker_by_ordinal(ordinal)
966            } else {
967                state
968                    .gpu_pool
969                    .select_worker_excluding(&model_name, estimated_vram, &skip)
970            };
971
972            let Some(worker) = worker else {
973                let rejected = gpu_job
974                    .take()
975                    .expect("gpu_job retained after failed dispatch");
976                let err_msg = if state.gpu_pool.worker_count() == 0 {
977                    format!("no GPU available for model {model_name}")
978                } else if let Some(ordinal) = preferred_gpu {
979                    format!("gpu:{ordinal} is not available for model {model_name}")
980                } else {
981                    format!("no GPU worker available for model {model_name}")
982                };
983                tracing::error!(model = %model_name, "{err_msg}");
984                if let Some(tx) = rejected.progress_tx {
985                    let _ = tx.send(SseMessage::Error(SseErrorEvent {
986                        message: err_msg.clone(),
987                    }));
988                }
989                let _ = rejected.result_tx.send(Err(err_msg));
990                state.queue.decrement();
991                state.job_registry.remove(&job_id);
992                break;
993            };
994
995            // Increment in-flight BEFORE sending to reserve the slot.
996            worker.in_flight.fetch_add(1, Ordering::SeqCst);
997            let pending = gpu_job.take().expect("gpu_job present in retry loop");
998            match worker.job_tx.try_send(pending) {
999                Ok(()) => {
1000                    dispatched = true;
1001                }
1002                Err(std::sync::mpsc::TrySendError::Full(j)) => {
1003                    worker.in_flight.fetch_sub(1, Ordering::SeqCst);
1004                    gpu_job = Some(j);
1005                    if preferred_gpu.is_none() {
1006                        skip.push(worker.gpu.ordinal);
1007                        if skip.len() >= state.gpu_pool.worker_count().max(1) {
1008                            skip.clear();
1009                            tokio::time::sleep(std::time::Duration::from_millis(10)).await;
1010                        }
1011                    } else {
1012                        tokio::time::sleep(std::time::Duration::from_millis(10)).await;
1013                    }
1014                }
1015                Err(std::sync::mpsc::TrySendError::Disconnected(j)) => {
1016                    worker.in_flight.fetch_sub(1, Ordering::SeqCst);
1017                    tracing::warn!(
1018                        gpu = worker.gpu.ordinal,
1019                        "GPU worker disconnected — retrying dispatch"
1020                    );
1021                    gpu_job = Some(j);
1022                    if preferred_gpu.is_none() {
1023                        skip.push(worker.gpu.ordinal);
1024                    } else {
1025                        let rejected = gpu_job.take().expect("gpu_job retained after disconnect");
1026                        let err_msg = format!(
1027                            "gpu:{} disconnected while dispatching model {model_name}",
1028                            worker.gpu.ordinal
1029                        );
1030                        if let Some(tx) = rejected.progress_tx {
1031                            let _ = tx.send(SseMessage::Error(SseErrorEvent {
1032                                message: err_msg.clone(),
1033                            }));
1034                        }
1035                        let _ = rejected.result_tx.send(Err(err_msg));
1036                        state.queue.decrement();
1037                        state.job_registry.remove(&job_id);
1038                        break;
1039                    }
1040                }
1041            }
1042        }
1043        #[cfg(feature = "metrics")]
1044        crate::metrics::record_queue_depth(state.queue.pending());
1045    }
1046    tracing::info!("multi-GPU queue dispatcher shutting down");
1047}
1048
1049/// Rough VRAM estimate for a model (used for placement decisions).
1050pub fn estimate_model_vram(model_name: &str) -> u64 {
1051    // Use a simple heuristic based on model name patterns.
1052    // Quantized models are smaller; BF16/FP16 are larger.
1053    let lower = model_name.to_lowercase();
1054    if lower.contains("flux2")
1055        && lower.contains("9b")
1056        && (lower.contains(":bf16") || lower.contains(":fp16"))
1057    {
1058        32_000_000_000 // Klein-9B BF16 needs a 32GB-class card in practice.
1059    } else if lower.contains(":q4") {
1060        6_000_000_000 // ~6GB
1061    } else if lower.contains(":q8") || lower.contains(":fp8") {
1062        12_000_000_000 // ~12GB
1063    } else if lower.contains(":bf16") || lower.contains(":fp16") {
1064        24_000_000_000 // ~24GB
1065    } else if lower.contains("sd15") || lower.contains("sd1.5") {
1066        4_000_000_000 // ~4GB
1067    } else {
1068        // SDXL (~8GB) and other models default to 8GB.
1069        8_000_000_000
1070    }
1071}
1072
1073#[cfg(test)]
1074mod tests {
1075    use super::*;
1076    use crate::gpu_pool::{GpuPool, GpuWorker};
1077    use crate::model_cache::ModelCache;
1078    use crate::state::QueueHandle;
1079    use mold_core::{GenerateRequest, ImageData, OutputFormat};
1080    use mold_db::MetadataDb;
1081    use mold_inference::device::DiscoveredGpu;
1082    use mold_inference::shared_pool::SharedPool;
1083    use std::sync::atomic::AtomicUsize;
1084    use std::sync::{Arc, Mutex, RwLock};
1085    use tempfile::TempDir;
1086
1087    /// A `GenerateRequest` with the bare minimum fields populated — enough to
1088    /// hand to `OutputMetadata::from_generate_request` in tests.
1089    fn fake_request(model: &str) -> GenerateRequest {
1090        GenerateRequest {
1091            prompt: "a cat".to_string(),
1092            negative_prompt: None,
1093            model: model.to_string(),
1094            width: 512,
1095            height: 512,
1096            steps: 4,
1097            guidance: 3.5,
1098            seed: Some(7),
1099            batch_size: 1,
1100            output_format: Some(OutputFormat::Png),
1101            embed_metadata: None,
1102            scheduler: None,
1103            cfg_plus: None,
1104            source_image: None,
1105            edit_images: None,
1106            strength: 0.75,
1107            mask_image: None,
1108            control_image: None,
1109            control_model: None,
1110            control_scale: 1.0,
1111            expand: None,
1112            original_prompt: None,
1113            lora: None,
1114            frames: None,
1115            fps: None,
1116            upscale_model: None,
1117            gif_preview: false,
1118            enable_audio: None,
1119            audio_file: None,
1120            source_video: None,
1121            keyframes: None,
1122            pipeline: None,
1123            loras: None,
1124            retake_range: None,
1125            spatial_upscale: None,
1126            temporal_upscale: None,
1127            placement: None,
1128        }
1129    }
1130
1131    fn fake_image() -> ImageData {
1132        ImageData {
1133            // PNG magic bytes — the helpers don't validate, but this keeps
1134            // the on-disk file from being trivially mistaken for empty.
1135            data: vec![0x89, 0x50, 0x4E, 0x47, 0x0D, 0x0A, 0x1A, 0x0A],
1136            format: OutputFormat::Png,
1137            width: 512,
1138            height: 512,
1139            index: 0,
1140        }
1141    }
1142
1143    fn test_worker(
1144        ordinal: usize,
1145        channel_size: usize,
1146    ) -> (
1147        Arc<GpuWorker>,
1148        std::sync::mpsc::Receiver<crate::gpu_pool::GpuJob>,
1149    ) {
1150        let (job_tx, job_rx) = std::sync::mpsc::sync_channel(channel_size);
1151        let worker = Arc::new(GpuWorker {
1152            gpu: DiscoveredGpu {
1153                ordinal,
1154                name: format!("gpu{ordinal}"),
1155                total_vram_bytes: 24_000_000_000,
1156                free_vram_bytes: 24_000_000_000,
1157            },
1158            model_cache: Arc::new(Mutex::new(ModelCache::new(3))),
1159            active_generation: Arc::new(RwLock::new(None)),
1160            model_load_lock: Arc::new(Mutex::new(())),
1161            shared_pool: Arc::new(Mutex::new(SharedPool::new())),
1162            in_flight: AtomicUsize::new(0),
1163            consecutive_failures: AtomicUsize::new(0),
1164            degraded_until: RwLock::new(None),
1165            job_tx,
1166        });
1167        (worker, job_rx)
1168    }
1169
1170    #[test]
1171    fn save_image_to_dir_writes_file_and_creates_missing_dir() {
1172        let tmp = TempDir::new().unwrap();
1173        let nested = tmp.path().join("sub/output");
1174        assert!(!nested.exists());
1175
1176        save_image_to_dir(&nested, &fake_image(), "flux-dev:q4", 1, None, None, None);
1177
1178        assert!(nested.exists(), "save should mkdir -p");
1179        let entries: Vec<_> = std::fs::read_dir(&nested).unwrap().collect();
1180        assert_eq!(entries.len(), 1);
1181        let name = entries[0].as_ref().unwrap().file_name();
1182        let name_str = name.to_string_lossy();
1183        // Filename uses model-with-colon-replaced-by-dash + ms timestamp + .png.
1184        assert!(name_str.starts_with("mold-flux-dev-q4-"), "{name_str}");
1185        assert!(name_str.ends_with(".png"), "{name_str}");
1186    }
1187
1188    #[test]
1189    fn save_image_to_dir_includes_batch_index_when_batch_size_gt_1() {
1190        let tmp = TempDir::new().unwrap();
1191        let mut img = fake_image();
1192        img.index = 3;
1193        img.format = OutputFormat::Jpeg;
1194        img.data = vec![0xFF, 0xD8, 0xFF, 0xE0]; // JPEG magic
1195
1196        save_image_to_dir(tmp.path(), &img, "sdxl", 4, None, None, None);
1197
1198        let entries: Vec<_> = std::fs::read_dir(tmp.path()).unwrap().collect();
1199        let name = entries[0]
1200            .as_ref()
1201            .unwrap()
1202            .file_name()
1203            .to_string_lossy()
1204            .to_string();
1205        assert!(
1206            name.contains("-3.jpeg"),
1207            "expected batch index suffix: {name}"
1208        );
1209    }
1210
1211    #[test]
1212    fn save_image_to_dir_upserts_metadata_row_when_db_provided() {
1213        let tmp = TempDir::new().unwrap();
1214        let db = MetadataDb::open_in_memory().unwrap();
1215        let req = fake_request("flux-dev:q4");
1216        let meta = OutputMetadata::from_generate_request(&req, 42, None, "test-version");
1217
1218        save_image_to_dir(
1219            tmp.path(),
1220            &fake_image(),
1221            "flux-dev:q4",
1222            1,
1223            Some(&meta),
1224            Some(1234),
1225            Some(&db),
1226        );
1227
1228        let rows = db.list(Some(tmp.path())).unwrap();
1229        assert_eq!(rows.len(), 1, "exactly one DB row for the saved file");
1230        let rec = &rows[0];
1231        assert_eq!(rec.metadata.prompt, "a cat");
1232        assert_eq!(rec.metadata.seed, 42);
1233        assert_eq!(rec.metadata.version, "test-version");
1234        assert_eq!(rec.format, OutputFormat::Png);
1235        assert_eq!(rec.generation_time_ms, Some(1234));
1236        // stat_from_disk should have populated the size from the actual file.
1237        assert!(rec.file_size_bytes.unwrap_or(0) > 0);
1238    }
1239
1240    #[test]
1241    fn save_image_to_dir_skips_db_when_metadata_is_none() {
1242        let tmp = TempDir::new().unwrap();
1243        let db = MetadataDb::open_in_memory().unwrap();
1244
1245        save_image_to_dir(
1246            tmp.path(),
1247            &fake_image(),
1248            "flux-dev:q4",
1249            1,
1250            None, // ← metadata absent
1251            Some(1234),
1252            Some(&db),
1253        );
1254
1255        // File still on disk, but no DB row recorded — both gates must hold
1256        // for the upsert to fire.
1257        assert_eq!(std::fs::read_dir(tmp.path()).unwrap().count(), 1);
1258        assert_eq!(db.list(None).unwrap().len(), 0);
1259    }
1260
1261    #[test]
1262    fn save_image_to_dir_invalid_path_does_not_panic() {
1263        // /dev/null is a file, not a directory — create_dir_all should fail
1264        // and the helper must log + return cleanly rather than panic.
1265        save_image_to_dir(
1266            std::path::Path::new("/dev/null/cant-mkdir-here"),
1267            &fake_image(),
1268            "test",
1269            1,
1270            None,
1271            None,
1272            None,
1273        );
1274    }
1275
1276    #[test]
1277    fn save_video_to_dir_writes_mp4_and_records_metadata() {
1278        let tmp = TempDir::new().unwrap();
1279        let db = MetadataDb::open_in_memory().unwrap();
1280        let mut req = fake_request("ltx-video:fp16");
1281        req.frames = Some(25);
1282        req.fps = Some(24);
1283        let meta = OutputMetadata::from_generate_request(&req, 99, None, "test-version");
1284
1285        // Minimal MP4-ish bytes: an `ftyp` box header. The helper writes
1286        // bytes verbatim — content validation happens at gallery scan time.
1287        let bytes = b"\x00\x00\x00\x18ftypmp42\x00\x00\x00\x00mp42isom".to_vec();
1288
1289        save_video_to_dir(
1290            tmp.path(),
1291            &bytes,
1292            b"",
1293            OutputFormat::Mp4,
1294            "ltx-video:fp16",
1295            &meta,
1296            Some(5000),
1297            Some(&db),
1298        );
1299
1300        let entries: Vec<_> = std::fs::read_dir(tmp.path()).unwrap().collect();
1301        assert_eq!(entries.len(), 1);
1302        let name = entries[0]
1303            .as_ref()
1304            .unwrap()
1305            .file_name()
1306            .to_string_lossy()
1307            .to_string();
1308        assert!(name.starts_with("mold-ltx-video-fp16-"), "{name}");
1309        assert!(name.ends_with(".mp4"), "{name}");
1310
1311        let rows = db.list(Some(tmp.path())).unwrap();
1312        assert_eq!(rows.len(), 1);
1313        assert_eq!(rows[0].format, OutputFormat::Mp4);
1314        assert_eq!(rows[0].metadata.frames, Some(25));
1315        assert_eq!(rows[0].metadata.fps, Some(24));
1316        assert_eq!(rows[0].generation_time_ms, Some(5000));
1317    }
1318
1319    #[test]
1320    fn save_video_to_dir_without_db_still_writes_file() {
1321        let tmp = TempDir::new().unwrap();
1322        let req = fake_request("ltx-video:fp16");
1323        let meta = OutputMetadata::from_generate_request(&req, 1, None, "v");
1324
1325        save_video_to_dir(
1326            tmp.path(),
1327            b"fake gif bytes",
1328            b"",
1329            OutputFormat::Gif,
1330            "ltx-video:fp16",
1331            &meta,
1332            None,
1333            None,
1334        );
1335
1336        let entries: Vec<_> = std::fs::read_dir(tmp.path()).unwrap().collect();
1337        assert_eq!(entries.len(), 1);
1338        let name = entries[0]
1339            .as_ref()
1340            .unwrap()
1341            .file_name()
1342            .to_string_lossy()
1343            .to_string();
1344        assert!(name.ends_with(".gif"), "{name}");
1345    }
1346
1347    #[test]
1348    fn save_video_to_dir_invalid_path_does_not_panic() {
1349        let req = fake_request("ltx-video:fp16");
1350        let meta = OutputMetadata::from_generate_request(&req, 1, None, "v");
1351        save_video_to_dir(
1352            std::path::Path::new("/dev/null/nope"),
1353            b"x",
1354            b"",
1355            OutputFormat::Mp4,
1356            "test",
1357            &meta,
1358            None,
1359            None,
1360        );
1361    }
1362
1363    /// `save_video_preview_gif_to` must write to
1364    /// `<preview_dir>/<filename>.preview.gif` — the exact location
1365    /// `GET /api/gallery/preview/:filename` streams from. Without this
1366    /// sidecar the preview endpoint would 404 on every real generation
1367    /// and the TUI detail pane would only ever see the PNG thumbnail
1368    /// fallback.
1369    #[test]
1370    fn save_video_preview_gif_writes_to_preview_cache() {
1371        let td = tempfile::tempdir().unwrap();
1372        let preview_dir = td.path().join("cache").join("previews");
1373
1374        const GIF: &[u8] = b"GIF89a\x01\x00\x01\x00\x00\x00\x00\x3b";
1375        save_video_preview_gif_to(&preview_dir, "ltx2-42.mp4", GIF);
1376
1377        let expected = preview_dir.join("ltx2-42.mp4.preview.gif");
1378        assert!(
1379            expected.is_file(),
1380            "preview gif should land at {}",
1381            expected.display()
1382        );
1383        assert_eq!(std::fs::read(&expected).unwrap(), GIF);
1384    }
1385
1386    #[test]
1387    fn build_sse_complete_event_video_carries_mp4_payload_and_metadata() {
1388        // Regression guard for the multi-GPU bug: if `response.video` is set,
1389        // the SSE complete event must encode the actual video bytes and
1390        // populate every `video_*` field so the client can reconstruct a
1391        // `VideoData`. Before the shared helper, `gpu_worker.rs` encoded the
1392        // thumbnail PNG and hard-coded every `video_*` field to `None`,
1393        // silently degrading every LTX-Video / LTX-2 response to an image.
1394        let video = mold_core::VideoData {
1395            data: vec![0x00, 0x00, 0x00, 0x18, b'f', b't', b'y', b'p'],
1396            format: OutputFormat::Mp4,
1397            width: 768,
1398            height: 512,
1399            frames: 25,
1400            fps: 24,
1401            thumbnail: vec![0x89, 0x50, 0x4E, 0x47],
1402            gif_preview: vec![b'G', b'I', b'F', b'8'],
1403            has_audio: true,
1404            duration_ms: Some(1040),
1405            audio_sample_rate: Some(44100),
1406            audio_channels: Some(2),
1407        };
1408        let resp = mold_core::GenerateResponse {
1409            images: vec![],
1410            video: Some(video.clone()),
1411            generation_time_ms: 1234,
1412            model: "ltx-2-19b-distilled:fp8".to_string(),
1413            seed_used: 7,
1414            gpu: Some(0),
1415        };
1416        // The `img` the caller synthesizes from the video thumbnail — must be
1417        // ignored for the video branch.
1418        let thumb_img = ImageData {
1419            data: video.thumbnail.clone(),
1420            format: OutputFormat::Png,
1421            width: video.width,
1422            height: video.height,
1423            index: 0,
1424        };
1425
1426        let event = build_sse_complete_event(&resp, &thumb_img);
1427
1428        let b64 = base64::engine::general_purpose::STANDARD;
1429        assert_eq!(event.image, b64.encode(&video.data));
1430        assert_eq!(event.format, OutputFormat::Mp4);
1431        assert_eq!(event.video_frames, Some(25));
1432        assert_eq!(event.video_fps, Some(24));
1433        assert_eq!(event.video_thumbnail, Some(b64.encode(&video.thumbnail)));
1434        assert_eq!(
1435            event.video_gif_preview,
1436            Some(b64.encode(&video.gif_preview))
1437        );
1438        assert!(event.video_has_audio);
1439        assert_eq!(event.video_duration_ms, Some(1040));
1440        assert_eq!(event.gpu, Some(0));
1441    }
1442
1443    #[test]
1444    fn build_sse_complete_event_video_empty_gif_preview_omits_field() {
1445        let video = mold_core::VideoData {
1446            data: vec![0x00, 0x00, 0x00, 0x18],
1447            format: OutputFormat::Mp4,
1448            width: 256,
1449            height: 256,
1450            frames: 17,
1451            fps: 12,
1452            thumbnail: vec![0x89, 0x50],
1453            gif_preview: Vec::new(),
1454            has_audio: false,
1455            duration_ms: None,
1456            audio_sample_rate: None,
1457            audio_channels: None,
1458        };
1459        let resp = mold_core::GenerateResponse {
1460            images: vec![],
1461            video: Some(video),
1462            generation_time_ms: 0,
1463            model: "m".to_string(),
1464            seed_used: 0,
1465            gpu: None,
1466        };
1467        let event = build_sse_complete_event(&resp, &fake_image());
1468        assert!(event.video_gif_preview.is_none());
1469        assert!(!event.video_has_audio);
1470    }
1471
1472    #[test]
1473    fn build_sse_complete_event_image_clears_all_video_fields() {
1474        let resp = mold_core::GenerateResponse {
1475            images: vec![fake_image()],
1476            video: None,
1477            generation_time_ms: 100,
1478            model: "flux-schnell:q8".to_string(),
1479            seed_used: 5,
1480            gpu: None,
1481        };
1482        let event = build_sse_complete_event(&resp, &fake_image());
1483        assert_eq!(event.format, OutputFormat::Png);
1484        assert!(event.video_frames.is_none());
1485        assert!(event.video_fps.is_none());
1486        assert!(event.video_thumbnail.is_none());
1487        assert!(event.video_gif_preview.is_none());
1488        assert!(!event.video_has_audio);
1489        assert!(event.video_duration_ms.is_none());
1490    }
1491
1492    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
1493    async fn queue_dispatcher_waits_for_worker_capacity_instead_of_rejecting() {
1494        let (worker, worker_rx) = test_worker(0, 1);
1495        let (job_tx, job_rx) = tokio::sync::mpsc::channel(4);
1496        let queue = QueueHandle::new(job_tx.clone());
1497        let state = crate::state::AppState::empty(
1498            mold_core::Config::default(),
1499            queue.clone(),
1500            Arc::new(GpuPool {
1501                workers: vec![worker.clone()],
1502            }),
1503            8,
1504        );
1505
1506        let (filler_result_tx, _filler_result_rx) = tokio::sync::oneshot::channel();
1507        let filler_job = crate::gpu_pool::GpuJob {
1508            id: String::new(),
1509            model: "busy-model".to_string(),
1510            request: fake_request("busy-model"),
1511            progress_tx: None,
1512            result_tx: filler_result_tx,
1513            output_dir: None,
1514            config: state.config.clone(),
1515            metadata_db: state.metadata_db.clone(),
1516            queue: state.queue.clone(),
1517            registry: state.job_registry.clone(),
1518        };
1519        worker.job_tx.send(filler_job).unwrap();
1520
1521        let dispatcher = tokio::spawn(run_queue_dispatcher(job_rx, state.clone()));
1522
1523        let (result_tx, mut result_rx) = tokio::sync::oneshot::channel();
1524        let job = crate::state::GenerationJob {
1525            id: String::new(),
1526            request: fake_request("flux-dev:q4"),
1527            progress_tx: None,
1528            result_tx,
1529            output_dir: None,
1530        };
1531        let _position = queue.submit(job, 8).await.unwrap();
1532
1533        tokio::time::sleep(std::time::Duration::from_millis(25)).await;
1534        assert!(
1535            result_rx.try_recv().is_err(),
1536            "dispatcher should keep the job pending while all worker channels are full"
1537        );
1538
1539        let _filler = worker_rx
1540            .recv()
1541            .expect("filler job should occupy the local channel");
1542        let dispatched = worker_rx
1543            .recv_timeout(std::time::Duration::from_secs(1))
1544            .expect("queued job should dispatch once capacity is available");
1545        assert_eq!(dispatched.model, "flux-dev:q4");
1546
1547        drop(job_tx);
1548        dispatcher.abort();
1549    }
1550
1551    /// Regression for the take-and-restore refactor in `process_job`: when
1552    /// the engine vanishes from the cache between `ensure_model_ready` and
1553    /// `cache.take()`, the take path must produce `None` (handled with a
1554    /// clean error in `process_job`) rather than panicking. The pure cache
1555    /// invariant — `take()` on an absent model returns `None` — is what
1556    /// keeps the take-and-restore safe.
1557    #[tokio::test]
1558    async fn cache_take_on_vanished_engine_returns_none_not_panic() {
1559        use crate::model_cache::ModelCache;
1560        use mold_core::GenerateResponse;
1561        use mold_inference::InferenceEngine;
1562
1563        struct StubEngine(&'static str);
1564        impl InferenceEngine for StubEngine {
1565            fn generate(&mut self, _r: &GenerateRequest) -> anyhow::Result<GenerateResponse> {
1566                unimplemented!()
1567            }
1568            fn model_name(&self) -> &str {
1569                self.0
1570            }
1571            fn is_loaded(&self) -> bool {
1572                true
1573            }
1574            fn load(&mut self) -> anyhow::Result<()> {
1575                Ok(())
1576            }
1577        }
1578
1579        let mut cache = ModelCache::new(3);
1580        // Cache empty (engine never inserted, or evicted/removed by a
1581        // concurrent admin call between `ensure_model_ready` and `take`).
1582        assert!(cache.take("vanished-model").is_none());
1583
1584        // After a take of a present engine, a subsequent take of the same
1585        // name must also return None — guards against double-take in the
1586        // restore path.
1587        cache.insert(Box::new(StubEngine("present-model")), 0);
1588        let first = cache.take("present-model");
1589        assert!(first.is_some());
1590        assert!(
1591            cache.take("present-model").is_none(),
1592            "double-take must return None"
1593        );
1594    }
1595
1596    fn buf_job(model: &str) -> BufferedJob {
1597        let (tx, _rx) = tokio::sync::oneshot::channel();
1598        BufferedJob::new(crate::state::GenerationJob {
1599            id: String::new(),
1600            request: fake_request(model),
1601            progress_tx: None,
1602            result_tx: tx,
1603            output_dir: None,
1604        })
1605    }
1606
1607    #[test]
1608    fn pick_next_job_picks_head_when_head_model_loaded() {
1609        use std::collections::{HashSet, VecDeque};
1610        let mut buffer: VecDeque<BufferedJob> = VecDeque::new();
1611        buffer.push_back(buf_job("a"));
1612        buffer.push_back(buf_job("b"));
1613        buffer.push_back(buf_job("a"));
1614        let loaded: HashSet<String> = ["a".to_string()].into_iter().collect();
1615        let picked = pick_next_job(&mut buffer, &loaded, 3);
1616        assert_eq!(picked.request.model, "a");
1617        assert_eq!(buffer.len(), 2);
1618        assert_eq!(buffer.front().unwrap().job.request.model, "b");
1619        assert_eq!(
1620            buffer.front().unwrap().deferred,
1621            0,
1622            "head shouldn't be deferred when picker chose the head itself"
1623        );
1624    }
1625
1626    #[test]
1627    fn pick_next_job_picks_non_head_when_only_non_head_model_loaded() {
1628        use std::collections::{HashSet, VecDeque};
1629        let mut buffer: VecDeque<BufferedJob> = VecDeque::new();
1630        buffer.push_back(buf_job("a"));
1631        buffer.push_back(buf_job("b"));
1632        buffer.push_back(buf_job("a"));
1633        let loaded: HashSet<String> = ["b".to_string()].into_iter().collect();
1634        let picked = pick_next_job(&mut buffer, &loaded, 3);
1635        assert_eq!(picked.request.model, "b");
1636        assert_eq!(buffer.len(), 2);
1637        // The head ("a") was skipped once and now sits at deferral=1.
1638        assert_eq!(buffer.front().unwrap().job.request.model, "a");
1639        assert_eq!(buffer.front().unwrap().deferred, 1);
1640    }
1641
1642    #[test]
1643    fn pick_next_job_force_dispatches_head_after_max_deferrals() {
1644        use std::collections::{HashSet, VecDeque};
1645        let mut buffer: VecDeque<BufferedJob> = VecDeque::new();
1646        let mut head = buf_job("a");
1647        head.deferred = 3;
1648        buffer.push_back(head);
1649        buffer.push_back(buf_job("b"));
1650        // Even though only `b` is loaded, head ("a") has hit the budget and wins.
1651        let loaded: HashSet<String> = ["b".to_string()].into_iter().collect();
1652        let picked = pick_next_job(&mut buffer, &loaded, 3);
1653        assert_eq!(picked.request.model, "a");
1654        assert_eq!(buffer.len(), 1);
1655        assert_eq!(buffer.front().unwrap().job.request.model, "b");
1656    }
1657
1658    #[test]
1659    fn pick_next_job_falls_back_to_head_when_nothing_loaded() {
1660        use std::collections::{HashSet, VecDeque};
1661        let mut buffer: VecDeque<BufferedJob> = VecDeque::new();
1662        buffer.push_back(buf_job("a"));
1663        buffer.push_back(buf_job("b"));
1664        let loaded: HashSet<String> = HashSet::new();
1665        let picked = pick_next_job(&mut buffer, &loaded, 3);
1666        assert_eq!(picked.request.model, "a");
1667    }
1668
1669    /// Fix D: with `max_deferrals = 0`, every reorder would exceed the
1670    /// budget on the very first skip, so the picker degenerates to FIFO —
1671    /// the head wins regardless of which model is loaded.
1672    #[test]
1673    fn pick_next_job_max_deferrals_zero_picks_head_even_when_non_head_loaded() {
1674        use std::collections::{HashSet, VecDeque};
1675        let mut buffer: VecDeque<BufferedJob> = VecDeque::new();
1676        buffer.push_back(buf_job("b")); // head
1677        buffer.push_back(buf_job("a")); // non-head
1678        let loaded: HashSet<String> = ["a".to_string()].into_iter().collect();
1679        let picked = pick_next_job(&mut buffer, &loaded, 0);
1680        assert_eq!(
1681            picked.request.model, "b",
1682            "max_deferrals=0 must force FIFO — head must win even when only the non-head model is loaded"
1683        );
1684        assert_eq!(buffer.len(), 1);
1685        assert_eq!(buffer.front().unwrap().job.request.model, "a");
1686    }
1687
1688    /// Fix D: with `max_deferrals = 0` and an empty `loaded` set, the head
1689    /// is the only candidate anyway. Locks in the FIFO behaviour when
1690    /// nothing is warm.
1691    #[test]
1692    fn pick_next_job_max_deferrals_zero_with_empty_loaded_picks_head() {
1693        use std::collections::{HashSet, VecDeque};
1694        let mut buffer: VecDeque<BufferedJob> = VecDeque::new();
1695        buffer.push_back(buf_job("a")); // head
1696        buffer.push_back(buf_job("b"));
1697        let loaded: HashSet<String> = HashSet::new();
1698        let picked = pick_next_job(&mut buffer, &loaded, 0);
1699        assert_eq!(picked.request.model, "a");
1700        assert_eq!(buffer.len(), 1);
1701        assert_eq!(buffer.front().unwrap().job.request.model, "b");
1702    }
1703
1704    /// Fix E: when both head and a non-head match `loaded`, the picker must
1705    /// pick the front-most match — i.e. the first `A` in `[A, B, A, B]`
1706    /// when both `A` and `B` are loaded. Locks in arrival-order stability
1707    /// across multiple matching jobs.
1708    #[test]
1709    fn pick_next_job_picks_front_most_match_when_multiple_loaded() {
1710        use std::collections::{HashSet, VecDeque};
1711        let mut buffer: VecDeque<BufferedJob> = VecDeque::new();
1712        buffer.push_back(buf_job("a"));
1713        buffer.push_back(buf_job("b"));
1714        buffer.push_back(buf_job("a"));
1715        buffer.push_back(buf_job("b"));
1716        let loaded: HashSet<String> = ["a".to_string(), "b".to_string()].into_iter().collect();
1717        let picked = pick_next_job(&mut buffer, &loaded, 3);
1718        assert_eq!(
1719            picked.request.model, "a",
1720            "front-most match wins (the first `a`), not the loaded model with the most copies later in the buffer"
1721        );
1722        // Three jobs remain: [b, a, b]; head was the picked first `a` so the
1723        // new head is the original-index-1 `b`. Nothing was deferred because
1724        // the picker chose the head itself.
1725        assert_eq!(buffer.len(), 3);
1726        let remaining: Vec<&str> = buffer
1727            .iter()
1728            .map(|b| b.job.request.model.as_str())
1729            .collect();
1730        assert_eq!(remaining, vec!["b", "a", "b"]);
1731        assert_eq!(buffer.front().unwrap().deferred, 0);
1732    }
1733
1734    /// Integration: an interleaved `[A, B, A, B]` queue dispatched against a
1735    /// single worker that has model `A` warm should reorder so both `A` jobs
1736    /// run first, then both `B` jobs — minimizing model swaps from 4 → 1.
1737    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
1738    async fn queue_dispatcher_reorders_interleaved_jobs_to_minimize_swaps() {
1739        let (worker, worker_rx) = test_worker(0, 8);
1740        // Pre-mark the worker as having model "a" loaded so the picker
1741        // recognises it as warm.
1742        {
1743            let mut cache = worker.model_cache.lock().unwrap();
1744            struct Engine(&'static str);
1745            impl mold_inference::InferenceEngine for Engine {
1746                fn generate(
1747                    &mut self,
1748                    _r: &GenerateRequest,
1749                ) -> anyhow::Result<mold_core::GenerateResponse> {
1750                    unimplemented!()
1751                }
1752                fn model_name(&self) -> &str {
1753                    self.0
1754                }
1755                fn is_loaded(&self) -> bool {
1756                    true
1757                }
1758                fn load(&mut self) -> anyhow::Result<()> {
1759                    Ok(())
1760                }
1761            }
1762            cache.insert(Box::new(Engine("a")), 0);
1763        }
1764
1765        let (job_tx, job_rx) = tokio::sync::mpsc::channel(8);
1766        let queue = QueueHandle::new(job_tx.clone());
1767        let state = crate::state::AppState::empty(
1768            mold_core::Config::default(),
1769            queue.clone(),
1770            Arc::new(GpuPool {
1771                workers: vec![worker.clone()],
1772            }),
1773            8,
1774        );
1775
1776        // Submit [a, b, a, b] BEFORE the dispatcher spins up so the buffer
1777        // top-up sees all four at once.
1778        let mut result_rxs = Vec::new();
1779        for model in ["a", "b", "a", "b"] {
1780            let (tx, rx) = tokio::sync::oneshot::channel();
1781            let job = crate::state::GenerationJob {
1782                id: String::new(),
1783                request: fake_request(model),
1784                progress_tx: None,
1785                result_tx: tx,
1786                output_dir: None,
1787            };
1788            queue.submit(job, 8).await.unwrap();
1789            result_rxs.push(rx);
1790        }
1791
1792        let dispatcher = tokio::spawn(run_queue_dispatcher(job_rx, state.clone()));
1793
1794        let mut order = Vec::new();
1795        for _ in 0..4 {
1796            let dispatched = worker_rx
1797                .recv_timeout(std::time::Duration::from_secs(2))
1798                .expect("worker should receive the dispatched job");
1799            order.push(dispatched.model);
1800        }
1801        drop(job_tx);
1802        dispatcher.abort();
1803
1804        assert_eq!(
1805            order,
1806            vec![
1807                "a".to_string(),
1808                "a".to_string(),
1809                "b".to_string(),
1810                "b".to_string(),
1811            ],
1812            "lookahead reorder should batch all `a` jobs together before swapping to `b`"
1813        );
1814    }
1815
1816    /// Fix F: the `top_up_buffer` helper must never grow the buffer past
1817    /// `buffer_size`, no matter how many jobs are sitting in the channel.
1818    /// This is the load-bearing invariant that bounds the working set the
1819    /// picker considers — without it a burst submission could let the
1820    /// dispatcher reorder across the entire pending queue, defeating the
1821    /// fairness guarantees the `deferred` counter is built around.
1822    #[tokio::test]
1823    async fn top_up_buffer_never_exceeds_capacity() {
1824        use std::collections::VecDeque;
1825        let (job_tx, mut job_rx) = tokio::sync::mpsc::channel::<GenerationJob>(32);
1826
1827        // Submit 10 jobs into the channel synchronously so the buffer's top-up
1828        // call sees them all immediately available via try_recv.
1829        for i in 0..10 {
1830            let (tx, _rx) = tokio::sync::oneshot::channel();
1831            let job = GenerationJob {
1832                id: String::new(),
1833                request: fake_request(&format!("model-{i}")),
1834                progress_tx: None,
1835                result_tx: tx,
1836                output_dir: None,
1837            };
1838            job_tx.send(job).await.unwrap();
1839        }
1840
1841        // buffer_size = 4 — top_up must stop at 4 even with 10 in the channel.
1842        let mut buffer: VecDeque<BufferedJob> = VecDeque::with_capacity(4);
1843        top_up_buffer(&mut buffer, &mut job_rx, 4);
1844        assert_eq!(
1845            buffer.len(),
1846            4,
1847            "top_up_buffer must cap at buffer_size, leaving the rest in the channel"
1848        );
1849
1850        // Drain the four buffered jobs, then top up again; the next call must
1851        // pull only the next four from the channel (FIFO order preserved).
1852        while buffer.pop_front().is_some() {}
1853        top_up_buffer(&mut buffer, &mut job_rx, 4);
1854        assert_eq!(buffer.len(), 4);
1855        let names: Vec<&str> = buffer
1856            .iter()
1857            .map(|b| b.job.request.model.as_str())
1858            .collect();
1859        assert_eq!(
1860            names,
1861            vec!["model-4", "model-5", "model-6", "model-7"],
1862            "second top-up must drain the next FIFO window from the channel"
1863        );
1864
1865        // Drop sender so the channel reports closed; remaining 2 jobs still
1866        // arrive via try_recv before the channel goes dry.
1867        drop(job_tx);
1868        while buffer.pop_front().is_some() {}
1869        top_up_buffer(&mut buffer, &mut job_rx, 4);
1870        assert_eq!(
1871            buffer.len(),
1872            2,
1873            "top_up_buffer drains the channel tail when fewer jobs than capacity remain"
1874        );
1875        let names: Vec<&str> = buffer
1876            .iter()
1877            .map(|b| b.job.request.model.as_str())
1878            .collect();
1879        assert_eq!(names, vec!["model-8", "model-9"]);
1880    }
1881
1882    /// Same invariant, but reached via the dispatcher loop (integration). A
1883    /// burst of N > buffer_size jobs must still dispatch in FIFO order with
1884    /// no jobs lost — the buffer cap can't drop traffic, only delay it. We
1885    /// drain the worker channel as fast as the dispatcher fills it, so the
1886    /// test exercises buffer rotation rather than worker-channel back-pressure.
1887    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
1888    async fn queue_dispatcher_dispatches_all_jobs_when_submission_exceeds_buffer() {
1889        let (worker, worker_rx) = test_worker(0, 4);
1890        let (job_tx, job_rx) = tokio::sync::mpsc::channel(32);
1891        let queue = QueueHandle::new(job_tx.clone());
1892        let state = crate::state::AppState::empty(
1893            mold_core::Config::default(),
1894            queue.clone(),
1895            Arc::new(GpuPool {
1896                workers: vec![worker.clone()],
1897            }),
1898            32,
1899        );
1900
1901        // Drain the worker channel concurrently and decrement in_flight as
1902        // a real worker would, so the dispatcher's worker-selection sees the
1903        // worker as idle for each subsequent send (otherwise `in_flight`
1904        // grows unbounded and the worker never re-classifies as eligible
1905        // when the sync-channel fills).
1906        let drain_worker = worker.clone();
1907        let drainer = std::thread::spawn(move || {
1908            let mut order = Vec::new();
1909            while order.len() < 10 {
1910                match worker_rx.recv_timeout(std::time::Duration::from_secs(5)) {
1911                    Ok(j) => {
1912                        drain_worker.in_flight.fetch_sub(1, Ordering::SeqCst);
1913                        order.push(j.model);
1914                    }
1915                    Err(e) => panic!("drain stalled at {:?}: {e:?}", order),
1916                }
1917            }
1918            order
1919        });
1920
1921        let dispatcher = tokio::spawn(run_queue_dispatcher(job_rx, state.clone()));
1922
1923        // Submit AFTER the dispatcher and drainer are running so we exercise
1924        // the live top-up loop rather than a one-shot drain of a pre-filled
1925        // channel. Hold result_rx values past the dispatch — the dispatcher
1926        // skips jobs whose result_tx is closed, which would otherwise drop
1927        // every job before it reaches the worker channel.
1928        let mut held_rxs = Vec::new();
1929        for i in 0..10 {
1930            let (tx, rx) = tokio::sync::oneshot::channel();
1931            held_rxs.push(rx);
1932            let job = crate::state::GenerationJob {
1933                id: String::new(),
1934                request: fake_request(&format!("model-{i}")),
1935                progress_tx: None,
1936                result_tx: tx,
1937                output_dir: None,
1938            };
1939            queue.submit(job, 32).await.unwrap();
1940        }
1941
1942        let order = drainer.join().expect("drainer thread panic");
1943        drop(job_tx);
1944        dispatcher.abort();
1945
1946        let expected: Vec<String> = (0..10).map(|i| format!("model-{i}")).collect();
1947        assert_eq!(
1948            order, expected,
1949            "10 distinct jobs must come out in FIFO across buffer rotations"
1950        );
1951    }
1952
1953    /// Serializes every test that mutates queue env vars (process-global).
1954    static QUEUE_ENV_LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(());
1955
1956    fn with_queue_env<R>(name: &str, value: Option<&str>, f: impl FnOnce() -> R) -> R {
1957        let _g = QUEUE_ENV_LOCK.lock().unwrap_or_else(|e| e.into_inner());
1958        let prev = std::env::var(name).ok();
1959        match value {
1960            Some(v) => std::env::set_var(name, v),
1961            None => std::env::remove_var(name),
1962        }
1963        let out = f();
1964        match prev {
1965            Some(v) => std::env::set_var(name, v),
1966            None => std::env::remove_var(name),
1967        }
1968        out
1969    }
1970
1971    #[test]
1972    fn resolve_lookahead_buffer_uses_default_when_env_missing() {
1973        let n = with_queue_env(LOOKAHEAD_BUFFER_ENV, None, resolve_lookahead_buffer);
1974        assert_eq!(n, DEFAULT_LOOKAHEAD_BUFFER);
1975    }
1976
1977    #[test]
1978    fn resolve_lookahead_buffer_honors_env_within_range() {
1979        let n = with_queue_env(LOOKAHEAD_BUFFER_ENV, Some("4"), resolve_lookahead_buffer);
1980        assert_eq!(n, 4);
1981    }
1982
1983    #[test]
1984    fn resolve_lookahead_buffer_falls_back_when_out_of_range() {
1985        // 0 is below the 1 lower bound; 999 is above the 64 upper bound.
1986        let n = with_queue_env(LOOKAHEAD_BUFFER_ENV, Some("0"), resolve_lookahead_buffer);
1987        assert_eq!(n, DEFAULT_LOOKAHEAD_BUFFER);
1988        let n = with_queue_env(LOOKAHEAD_BUFFER_ENV, Some("999"), resolve_lookahead_buffer);
1989        assert_eq!(n, DEFAULT_LOOKAHEAD_BUFFER);
1990    }
1991
1992    #[test]
1993    fn resolve_lookahead_buffer_falls_back_when_unparseable() {
1994        let n = with_queue_env(
1995            LOOKAHEAD_BUFFER_ENV,
1996            Some("not-a-number"),
1997            resolve_lookahead_buffer,
1998        );
1999        assert_eq!(n, DEFAULT_LOOKAHEAD_BUFFER);
2000    }
2001
2002    #[test]
2003    fn resolve_max_deferrals_uses_default_when_env_missing() {
2004        let n = with_queue_env(MAX_DEFERRALS_ENV, None, resolve_max_deferrals);
2005        assert_eq!(n, DEFAULT_MAX_DEFERRALS);
2006    }
2007
2008    #[test]
2009    fn resolve_max_deferrals_honors_env_within_range() {
2010        // 0 is the in-range "FIFO" sentinel, 32 is the upper edge.
2011        let n = with_queue_env(MAX_DEFERRALS_ENV, Some("0"), resolve_max_deferrals);
2012        assert_eq!(n, 0);
2013        let n = with_queue_env(MAX_DEFERRALS_ENV, Some("32"), resolve_max_deferrals);
2014        assert_eq!(n, 32);
2015        let n = with_queue_env(MAX_DEFERRALS_ENV, Some("5"), resolve_max_deferrals);
2016        assert_eq!(n, 5);
2017    }
2018
2019    #[test]
2020    fn resolve_max_deferrals_falls_back_when_out_of_range() {
2021        let n = with_queue_env(MAX_DEFERRALS_ENV, Some("999"), resolve_max_deferrals);
2022        assert_eq!(n, DEFAULT_MAX_DEFERRALS);
2023    }
2024
2025    #[test]
2026    fn resolve_max_deferrals_falls_back_when_unparseable() {
2027        let n = with_queue_env(
2028            MAX_DEFERRALS_ENV,
2029            Some("not-a-number"),
2030            resolve_max_deferrals,
2031        );
2032        assert_eq!(n, DEFAULT_MAX_DEFERRALS);
2033    }
2034
2035    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
2036    async fn queue_dispatcher_honors_explicit_placement_gpu() {
2037        let (worker0, rx0) = test_worker(0, 1);
2038        let (worker1, rx1) = test_worker(1, 1);
2039        let (job_tx, job_rx) = tokio::sync::mpsc::channel(4);
2040        let queue = QueueHandle::new(job_tx.clone());
2041        let state = crate::state::AppState::empty(
2042            mold_core::Config::default(),
2043            queue.clone(),
2044            Arc::new(GpuPool {
2045                workers: vec![worker0, worker1],
2046            }),
2047            8,
2048        );
2049
2050        let dispatcher = tokio::spawn(run_queue_dispatcher(job_rx, state));
2051
2052        let mut request = fake_request("flux-dev:q4");
2053        request.placement = Some(mold_core::types::DevicePlacement {
2054            text_encoders: mold_core::types::DeviceRef::Auto,
2055            advanced: Some(mold_core::types::AdvancedPlacement {
2056                transformer: mold_core::types::DeviceRef::gpu(1),
2057                ..mold_core::types::AdvancedPlacement::default()
2058            }),
2059        });
2060
2061        let (result_tx, _result_rx) = tokio::sync::oneshot::channel();
2062        let job = crate::state::GenerationJob {
2063            id: String::new(),
2064            request,
2065            progress_tx: None,
2066            result_tx,
2067            output_dir: None,
2068        };
2069        let _position = queue.submit(job, 8).await.unwrap();
2070
2071        let dispatched = rx1
2072            .recv_timeout(std::time::Duration::from_secs(1))
2073            .expect("explicit placement should route to gpu 1");
2074        assert_eq!(dispatched.model, "flux-dev:q4");
2075        assert!(rx0.try_recv().is_err(), "gpu 0 should not receive the job");
2076
2077        drop(job_tx);
2078        dispatcher.abort();
2079    }
2080}