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
19fn progress_to_sse(event: mold_inference::ProgressEvent) -> SseProgressEvent {
21 event.into()
22}
23
24pub(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
87pub(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#[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
200pub(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
216fn 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
236fn hostname_string() -> Option<String> {
238 hostname::get().ok().and_then(|s| s.into_string().ok())
239}
240
241fn 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
252pub(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
318pub 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_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 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
368fn 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
390pub(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
404pub(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
422pub(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 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 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
482pub(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
513pub(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 if job.result_tx.is_closed() {
545 tracing::debug!("skipping queued job — client disconnected");
546 return;
547 }
548
549 state.job_registry.mark_running(&job.id, None);
552
553 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 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 #[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 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 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 let rss_before = crate::resources::ram_snapshot().used_by_mold;
646 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 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 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 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 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 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
842pub 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 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 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
1049pub fn estimate_model_vram(model_name: &str) -> u64 {
1051 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 } else if lower.contains(":q4") {
1060 6_000_000_000 } else if lower.contains(":q8") || lower.contains(":fp8") {
1062 12_000_000_000 } else if lower.contains(":bf16") || lower.contains(":fp16") {
1064 24_000_000_000 } else if lower.contains("sd15") || lower.contains("sd1.5") {
1066 4_000_000_000 } else {
1068 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 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 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 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]; 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 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, Some(1234),
1252 Some(&db),
1253 );
1254
1255 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 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 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 #[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 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 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 #[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 assert!(cache.take("vanished-model").is_none());
1583
1584 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 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 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 #[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")); buffer.push_back(buf_job("a")); 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 #[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")); 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 #[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 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 #[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 {
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 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 #[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 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 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 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(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 #[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 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 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 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 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 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}