1pub mod auth;
2pub mod catalog_api;
3pub mod chain_job_runner;
4pub mod chain_limits;
5pub mod test_support;
6pub mod downloads;
8pub mod gpu_pool;
9pub mod gpu_worker;
10pub mod job_registry;
11pub mod logging;
12mod memory_preflight;
13#[cfg(feature = "metrics")]
14pub mod metrics;
15pub mod model_cache;
16pub mod model_manager;
17pub mod queue;
18pub mod rate_limit;
19pub mod request_id;
20pub mod resources;
21pub mod routes;
22pub mod routes_chain;
23pub mod routes_chain_jobs;
24mod signals;
25pub mod state;
26pub mod web_ui;
27
28#[cfg(all(test, feature = "metrics"))]
29mod metrics_test;
30#[cfg(test)]
31mod resources_test;
32#[cfg(test)]
33mod routes_test;
34
35use anyhow::Result;
36use axum::{extract::DefaultBodyLimit, middleware};
37use mold_core::types::GpuSelection;
38use mold_core::{Config, ModelPaths};
39use std::net::SocketAddr;
40use std::path::PathBuf;
41use std::sync::atomic::AtomicUsize;
42use tokio::net::TcpListener;
43use tower_http::cors::CorsLayer;
44use tower_http::trace::TraceLayer;
45use tracing::info;
46
47use state::QueueHandle;
48
49const MAX_REQUEST_BODY_BYTES: usize = 64 * 1024 * 1024;
50
51pub async fn run_server(
52 bind: &str,
53 port: u16,
54 models_dir: PathBuf,
55 gpu_selection: GpuSelection,
56 queue_size: usize,
57) -> Result<()> {
58 signals::ignore_sigpipe();
64
65 Config::install_runtime_models_dir_override(models_dir.clone());
66
67 let mut config = Config::load_or_default();
68 config.models_dir = models_dir.to_string_lossy().into_owned();
69 let model_name = config.resolved_default_model();
70
71 let shared_pool = std::sync::Arc::new(std::sync::Mutex::new(
73 mold_inference::shared_pool::SharedPool::new(),
74 ));
75
76 let discovered = mold_inference::device::discover_gpus();
77 let selected = mold_inference::device::filter_gpus(&discovered, &gpu_selection);
78
79 if selected.is_empty() && !discovered.is_empty() {
80 anyhow::bail!(
81 "No GPUs matched selection {:?} (discovered: {:?})",
82 gpu_selection,
83 discovered.iter().map(|g| g.ordinal).collect::<Vec<_>>()
84 );
85 }
86
87 let mut workers = Vec::new();
88 let mut _gpu_thread_handles = Vec::new();
89
90 const PER_WORKER_CHANNEL_SIZE: usize = 2;
95
96 let max_cached = state::resolve_max_cached_models();
97 for gpu in &selected {
98 let (job_tx, job_rx) = std::sync::mpsc::sync_channel(PER_WORKER_CHANNEL_SIZE);
99 let worker = std::sync::Arc::new(gpu_pool::GpuWorker {
100 gpu: gpu.clone(),
101 model_cache: std::sync::Arc::new(std::sync::Mutex::new(model_cache::ModelCache::new(
102 max_cached,
103 ))),
104 active_generation: std::sync::Arc::new(std::sync::RwLock::new(None)),
105 model_load_lock: std::sync::Arc::new(std::sync::Mutex::new(())),
106 shared_pool: shared_pool.clone(),
107 in_flight: AtomicUsize::new(0),
108 consecutive_failures: AtomicUsize::new(0),
109 degraded_until: std::sync::RwLock::new(None),
110 job_tx,
111 });
112
113 let handle = gpu_worker::spawn_gpu_thread(worker.clone(), job_rx);
114 _gpu_thread_handles.push(handle);
115 workers.push(worker);
116 }
117
118 let gpu_pool = std::sync::Arc::new(gpu_pool::GpuPool { workers });
119
120 for status in gpu_pool.gpu_status() {
122 info!(
123 gpu = status.ordinal,
124 name = %status.name,
125 vram_mb = status.vram_total_bytes / 1_000_000,
126 "GPU worker ready"
127 );
128 }
129
130 if selected.is_empty() {
131 info!("no GPUs discovered — server will operate in CPU/pull-only mode");
132 }
133
134 let (job_tx, job_rx) = tokio::sync::mpsc::channel(queue_size.max(1));
136 let queue_handle = QueueHandle::new(job_tx);
137
138 let mut state = if gpu_pool.worker_count() > 0 {
140 if let Some(paths) = ModelPaths::resolve(&model_name, &config) {
141 info!(model = %model_name, "configured model");
142 info!(transformer = %paths.transformer.display());
143 info!(vae = %paths.vae.display());
144 if let Some(spatial_upscaler) = &paths.spatial_upscaler {
145 info!(spatial_upscaler = %spatial_upscaler.display());
146 }
147 if let Some(t5) = &paths.t5_encoder {
148 info!(t5 = %t5.display());
149 }
150 if let Some(clip) = &paths.clip_encoder {
151 info!(clip = %clip.display());
152 }
153 if let Some(t5_tok) = &paths.t5_tokenizer {
154 info!(t5_tok = %t5_tok.display());
155 }
156 if let Some(clip_tok) = &paths.clip_tokenizer {
157 info!(clip_tok = %clip_tok.display());
158 }
159 if let Some(clip2) = &paths.clip_encoder_2 {
160 info!(clip2 = %clip2.display());
161 }
162 if let Some(clip2_tok) = &paths.clip_tokenizer_2 {
163 info!(clip2_tok = %clip2_tok.display());
164 }
165 for (i, te) in paths.text_encoder_files.iter().enumerate() {
166 info!(text_encoder_shard = i, path = %te.display());
167 }
168 if let Some(text_tok) = &paths.text_tokenizer {
169 info!(text_tok = %text_tok.display());
170 }
171 info!("multi-GPU mode defers model loading to per-GPU workers");
172 } else {
173 info!("no default model configured — models will be pulled on first request");
174 }
175 let mut state = state::AppState::empty(config, queue_handle, gpu_pool.clone(), queue_size);
176 state.shared_pool = shared_pool;
177 state
178 } else {
179 match ModelPaths::resolve(&model_name, &config) {
180 Some(paths) => {
181 info!(model = %model_name, "configured model");
182 info!(transformer = %paths.transformer.display());
183 info!(vae = %paths.vae.display());
184 if let Some(spatial_upscaler) = &paths.spatial_upscaler {
185 info!(spatial_upscaler = %spatial_upscaler.display());
186 }
187 if let Some(t5) = &paths.t5_encoder {
188 info!(t5 = %t5.display());
189 }
190 if let Some(clip) = &paths.clip_encoder {
191 info!(clip = %clip.display());
192 }
193 if let Some(t5_tok) = &paths.t5_tokenizer {
194 info!(t5_tok = %t5_tok.display());
195 }
196 if let Some(clip_tok) = &paths.clip_tokenizer {
197 info!(clip_tok = %clip_tok.display());
198 }
199 if let Some(clip2) = &paths.clip_encoder_2 {
200 info!(clip2 = %clip2.display());
201 }
202 if let Some(clip2_tok) = &paths.clip_tokenizer_2 {
203 info!(clip2_tok = %clip2_tok.display());
204 }
205 for (i, te) in paths.text_encoder_files.iter().enumerate() {
206 info!(text_encoder_shard = i, path = %te.display());
207 }
208 if let Some(text_tok) = &paths.text_tokenizer {
209 info!(text_tok = %text_tok.display());
210 }
211
212 let offload = std::env::var("MOLD_OFFLOAD").is_ok_and(|v| v == "1");
213 let engine = mold_inference::create_engine_with_pool(
214 model_name,
215 paths,
216 &config,
217 mold_inference::LoadStrategy::Eager,
218 0,
219 offload,
220 Some(shared_pool.clone()),
221 )?;
222 let mut state = state::AppState::new(
223 engine,
224 config,
225 queue_handle,
226 gpu_pool.clone(),
227 queue_size,
228 );
229 state.shared_pool = shared_pool;
230 state
231 }
232 None => {
233 info!("no default model configured — models will be pulled on first request");
234 state::AppState::empty(config, queue_handle, gpu_pool.clone(), queue_size)
235 }
236 }
237 };
238
239 match mold_db::open_default() {
241 Ok(Some(db)) => {
242 info!(db = %db.path().display(), "metadata DB opened");
243 state.metadata_db = std::sync::Arc::new(Some(db));
244 }
245 Ok(None) => {
246 tracing::info!("metadata DB disabled (MOLD_DB_DISABLE set or MOLD_HOME unresolved)");
247 }
248 Err(e) => {
249 tracing::warn!(
250 "failed to open metadata DB: {e:#} — gallery falls back to filesystem scan"
251 );
252 }
253 }
254
255 if state.metadata_db.is_some() {
256 let Some(jobs_root) = Config::mold_dir().map(|dir| dir.join("jobs")) else {
257 anyhow::bail!("metadata DB opened but MOLD_HOME could not be resolved for chain jobs");
258 };
259 std::fs::create_dir_all(&jobs_root)?;
260 let db_arc = state.metadata_db.clone();
261 let reconcile_root = jobs_root.clone();
262 let (flipped, repaired) = tokio::task::spawn_blocking(move || {
263 let Some(db) = db_arc.as_ref().as_ref() else {
264 anyhow::bail!("metadata DB disappeared before chain reconcile");
265 };
266 chain_job_runner::startup_reconcile(db, &reconcile_root)
267 })
268 .await??;
269 tracing::info!(
270 flipped,
271 repaired,
272 jobs_root = %jobs_root.display(),
273 "chain job startup reconcile complete"
274 );
275
276 let gc_db = state.metadata_db.clone();
277 let gc_root = jobs_root.clone();
278 let startup_gc = tokio::task::spawn_blocking(move || {
279 let Some(db) = gc_db.as_ref().as_ref() else {
280 anyhow::bail!("metadata DB disappeared before chain startup GC");
281 };
282 chain_job_runner::startup_gc_sweep(db, &gc_root)
283 })
284 .await??;
285 tracing::info!(
286 swept_ephemeral_jobs = startup_gc.swept_ephemeral_jobs,
287 pruned_artifact_dirs = startup_gc.pruned_artifact_dirs,
288 jobs_root = %jobs_root.display(),
289 "chain job startup GC complete"
290 );
291
292 let config_snapshot = state.config.read().await.clone();
293 let output_dir = if config_snapshot.is_output_disabled() {
294 None
295 } else {
296 Some(config_snapshot.effective_output_dir())
297 };
298 let deps = chain_job_runner::RunnerDeps {
299 db: state.metadata_db.clone(),
300 jobs_root,
301 executor: std::sync::Arc::new(chain_job_runner::ProductionStageExecutor::new(
302 state.gpu_pool.clone(),
303 config_snapshot,
304 )),
305 queue_probe: std::sync::Arc::new(chain_job_runner::ProductionQueueProbe::new(
306 state.queue.clone(),
307 state.gpu_pool.clone(),
308 )),
309 events: std::sync::Arc::new(chain_job_runner::JobEventBus::new()),
310 cancel: std::sync::Arc::new(chain_job_runner::CancelRegistry::new()),
311 job_locks: std::sync::Arc::new(chain_job_runner::JobMutationLocks::new()),
312 claims: std::sync::Arc::new(chain_job_runner::EphemeralClaims::default()),
313 output_dir,
314 };
315 state.chain_jobs = Some(std::sync::Arc::new(chain_job_runner::spawn_runner(deps)));
316 }
317
318 let worker_state = state.clone();
322 if gpu_pool.worker_count() > 0 {
323 tokio::spawn(queue::run_queue_dispatcher(job_rx, worker_state));
324 } else {
325 tokio::spawn(queue::run_queue_worker(job_rx, worker_state));
326 }
327
328 let idle_evict_handle = spawn_cache_idle_evictor(
332 state.model_cache.clone(),
333 state.model_load_lock.clone(),
334 gpu_pool.clone(),
335 std::time::Duration::from_secs(state::resolve_cache_idle_ttl_secs()),
336 );
337
338 let downloads_shutdown = tokio_util::sync::CancellationToken::new();
345 let downloads_models_dir = state.config.read().await.resolved_models_dir();
346 let downloads_driver = crate::downloads::spawn_driver(
347 state.downloads.clone(),
348 std::sync::Arc::new(crate::downloads::HfPullDriver),
349 std::sync::Arc::new(crate::downloads::CivitaiRecipeDriver),
350 downloads_models_dir,
351 downloads_shutdown.clone(),
352 );
353
354 {
356 let config = state.config.read().await;
357 if config.is_output_disabled() {
358 tracing::warn!(
359 "image output is disabled (output_dir is empty) — \
360 generated images will not be saved and the TUI gallery will be empty"
361 );
362 } else {
363 let output_dir = config.effective_output_dir();
364 let _ = std::fs::create_dir_all(&output_dir);
365 info!(output_dir = %output_dir.display(), "gallery output directory");
366 routes::spawn_thumbnail_warmup(&config);
367
368 if state.metadata_db.is_some() {
372 let db_arc = state.metadata_db.clone();
373 let dir = output_dir.clone();
374 tokio::spawn(async move {
375 let join = tokio::task::spawn_blocking(move || {
376 if let Some(db) = db_arc.as_ref() {
377 db.reconcile(&dir)
378 } else {
379 Ok(mold_db::ReconcileStats::default())
380 }
381 })
382 .await;
383 match join {
384 Ok(Ok(stats)) => tracing::info!(
385 imported = stats.imported,
386 updated = stats.updated,
387 removed = stats.removed,
388 kept = stats.kept,
389 "metadata DB reconciled with gallery directory"
390 ),
391 Ok(Err(e)) => tracing::warn!("metadata DB reconcile failed: {e:#}"),
392 Err(e) => tracing::warn!("metadata DB reconcile task join error: {e}"),
393 }
394 });
395 }
396 }
397 }
398
399 let auth_state = auth::load_api_keys()?;
401 let rl_config = rate_limit::load_rate_limit_config()?;
402
403 let cors = build_cors_layer()?;
404
405 #[cfg(feature = "metrics")]
408 let prometheus_handle = metrics::install_recorder();
409
410 let (shutdown_tx, shutdown_rx) = tokio::sync::oneshot::channel::<()>();
416 *state.shutdown_tx.lock().await = Some(shutdown_tx);
417
418 #[cfg(unix)]
419 {
420 let sigterm_state = state.clone();
421 tokio::spawn(async move {
422 if let Ok(mut sig) =
423 tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate())
424 {
425 sig.recv().await;
426 tracing::info!("received SIGTERM, initiating graceful shutdown");
427 if let Some(tx) = sigterm_state.shutdown_tx.lock().await.take() {
428 let _ = tx.send(());
429 }
430 }
431 });
432 }
433
434 let resources_aggregator = resources::spawn_aggregator(state.resources.clone());
438
439 #[cfg(feature = "metrics")]
441 let server_start_time = state.start_time;
442
443 #[allow(unused_mut)]
446 let mut app = routes::create_router(state)
447 .merge(web_ui::router())
448 .layer(DefaultBodyLimit::max(MAX_REQUEST_BODY_BYTES))
449 .layer(middleware::from_fn(rate_limit::rate_limit_middleware))
450 .layer(middleware::from_fn_with_state(
451 rl_config,
452 rate_limit::inject_rate_limit_state,
453 ))
454 .layer(middleware::from_fn(auth::require_api_key))
455 .layer(middleware::from_fn_with_state(
456 auth_state,
457 auth::inject_auth_state,
458 ));
459
460 #[cfg(feature = "metrics")]
463 {
464 app = app.layer(middleware::from_fn(metrics::http_metrics_middleware));
465 }
466
467 #[cfg(feature = "metrics")]
468 {
469 let metrics_state = metrics::MetricsState {
470 handle: prometheus_handle,
471 start_time: server_start_time,
472 };
473 app = app.route(
474 "/metrics",
475 axum::routing::get(metrics::metrics_endpoint).with_state(metrics_state),
476 );
477 }
478
479 let app = app
480 .layer(middleware::from_fn(request_id::request_id_middleware))
481 .layer(TraceLayer::new_for_http())
482 .layer(cors);
483
484 let addr: SocketAddr = format!("{bind}:{port}").parse()?;
485 let version = mold_core::build_info::version_string();
486 info!(%addr, %version, "starting mold server");
487
488 let listener = TcpListener::bind(addr).await?;
489 axum::serve(
490 listener,
491 app.into_make_service_with_connect_info::<SocketAddr>(),
492 )
493 .with_graceful_shutdown(async {
494 let _ = shutdown_rx.await;
495 tracing::info!("shutting down");
496 })
497 .await?;
498
499 downloads_shutdown.cancel();
504 downloads_driver.abort();
505 idle_evict_handle.abort();
506 resources_aggregator.abort();
509
510 Ok(())
511}
512
513fn spawn_cache_idle_evictor(
526 legacy_cache: std::sync::Arc<tokio::sync::Mutex<model_cache::ModelCache>>,
527 legacy_load_lock: std::sync::Arc<tokio::sync::Mutex<()>>,
528 gpu_pool: std::sync::Arc<gpu_pool::GpuPool>,
529 ttl: std::time::Duration,
530) -> tokio::task::JoinHandle<()> {
531 use tokio::time::{interval, MissedTickBehavior};
532 tokio::spawn(async move {
533 let mut tick = interval(std::time::Duration::from_secs(60));
534 tick.set_missed_tick_behavior(MissedTickBehavior::Skip);
535 tick.tick().await;
538 loop {
539 tick.tick().await;
540
541 {
548 let _load_guard = legacy_load_lock.lock().await;
549 let evicted = {
550 let mut cache = legacy_cache.lock().await;
551 cache.evict_idle(ttl)
552 };
553 let evicted_count = evicted.len();
554 drop(evicted);
557
558 let legacy_active = legacy_cache.lock().await.active_model().is_some();
564 if evicted_count > 0 && !legacy_active {
565 tokio::task::spawn_blocking(|| {
566 mold_inference::reclaim_gpu_memory(0);
567 })
568 .await
569 .ok();
570 }
571 }
572
573 for worker in &gpu_pool.workers {
579 let worker = worker.clone();
580 let ttl = ttl;
581 tokio::task::spawn_blocking(move || {
582 let _load_guard = match worker.model_load_lock.lock() {
583 Ok(g) => g,
584 Err(poisoned) => poisoned.into_inner(),
585 };
586 let evicted = {
587 let mut cache =
588 worker.model_cache.lock().unwrap_or_else(|e| e.into_inner());
589 cache.evict_idle(ttl)
590 };
591 let evicted_count = evicted.len();
592 drop(evicted);
593
594 let active = worker
595 .model_cache
596 .lock()
597 .map(|c| c.active_model().is_some())
598 .unwrap_or(true);
599 if evicted_count > 0 && !active {
600 mold_inference::device::init_thread_gpu_ordinal(worker.gpu.ordinal);
606 mold_inference::reclaim_gpu_memory(worker.gpu.ordinal);
607 mold_inference::device::clear_thread_gpu_ordinal();
608 }
609 })
610 .await
611 .ok();
612 }
613 }
614 })
615}
616
617fn build_cors_layer() -> Result<CorsLayer> {
618 let cors = match std::env::var("MOLD_CORS_ORIGIN") {
619 Ok(origin) if !origin.is_empty() => {
620 let origin = origin
621 .parse::<axum::http::HeaderValue>()
622 .map_err(|_| anyhow::anyhow!("invalid MOLD_CORS_ORIGIN value: {origin}"))?;
623 CorsLayer::new()
624 .allow_origin(origin)
625 .allow_methods([
626 axum::http::Method::GET,
627 axum::http::Method::POST,
628 axum::http::Method::DELETE,
629 ])
630 .allow_headers(tower_http::cors::Any)
631 .expose_headers([
632 axum::http::header::HeaderName::from_static("x-mold-seed-used"),
633 axum::http::header::HeaderName::from_static("x-request-id"),
634 axum::http::header::HeaderName::from_static("retry-after"),
635 axum::http::header::HeaderName::from_static("x-mold-video-frames"),
636 axum::http::header::HeaderName::from_static("x-mold-video-fps"),
637 axum::http::header::HeaderName::from_static("x-mold-video-width"),
638 axum::http::header::HeaderName::from_static("x-mold-video-height"),
639 axum::http::header::HeaderName::from_static("x-mold-video-has-audio"),
640 axum::http::header::HeaderName::from_static("x-mold-video-duration-ms"),
641 axum::http::header::HeaderName::from_static("x-mold-video-audio-sample-rate"),
642 axum::http::header::HeaderName::from_static("x-mold-video-audio-channels"),
643 axum::http::header::HeaderName::from_static("x-mold-dimension-warning"),
644 ])
645 }
646 _ => CorsLayer::permissive(),
647 };
648 Ok(cors)
649}
650
651#[cfg(test)]
652mod tests {
653 use super::build_cors_layer;
654 use std::sync::Mutex;
655
656 static ENV_LOCK: Mutex<()> = Mutex::new(());
657
658 #[test]
659 fn invalid_cors_origin_returns_error() {
660 let _lock = ENV_LOCK.lock().unwrap();
661 std::env::set_var("MOLD_CORS_ORIGIN", "\nnot-a-header");
662 let result = build_cors_layer();
663 std::env::remove_var("MOLD_CORS_ORIGIN");
664 assert!(result.is_err());
665 }
666
667 #[test]
668 fn valid_cors_origin_builds_layer() {
669 let _lock = ENV_LOCK.lock().unwrap();
670 std::env::set_var("MOLD_CORS_ORIGIN", "https://example.com");
671 let result = build_cors_layer();
672 std::env::remove_var("MOLD_CORS_ORIGIN");
673 assert!(result.is_ok());
674 }
675}