mold_server/
lib.rs

1pub mod model_cache;
2pub mod model_manager;
3pub mod queue;
4pub mod routes;
5pub mod state;
6
7#[cfg(test)]
8mod routes_test;
9
10use anyhow::Result;
11use mold_core::{Config, ModelPaths};
12use std::net::SocketAddr;
13use std::path::PathBuf;
14use tokio::net::TcpListener;
15use tower_http::cors::CorsLayer;
16use tower_http::trace::TraceLayer;
17use tracing::info;
18
19use state::QueueHandle;
20
21pub async fn run_server(bind: &str, port: u16, models_dir: PathBuf) -> Result<()> {
22    Config::install_runtime_models_dir_override(models_dir.clone());
23
24    let mut config = Config::load_or_default();
25    config.models_dir = models_dir.to_string_lossy().into_owned();
26    let model_name = config.resolved_default_model();
27
28    // Create the generation queue channel (bounded, 16 slots).
29    let (job_tx, job_rx) = tokio::sync::mpsc::channel(16);
30    let queue_handle = QueueHandle::new(job_tx);
31
32    let state = match ModelPaths::resolve(&model_name, &config) {
33        Some(paths) => {
34            info!(model = %model_name, "configured model");
35            info!(transformer = %paths.transformer.display());
36            info!(vae = %paths.vae.display());
37            if let Some(t5) = &paths.t5_encoder {
38                info!(t5 = %t5.display());
39            }
40            if let Some(clip) = &paths.clip_encoder {
41                info!(clip = %clip.display());
42            }
43            if let Some(t5_tok) = &paths.t5_tokenizer {
44                info!(t5_tok = %t5_tok.display());
45            }
46            if let Some(clip_tok) = &paths.clip_tokenizer {
47                info!(clip_tok = %clip_tok.display());
48            }
49            if let Some(clip2) = &paths.clip_encoder_2 {
50                info!(clip2 = %clip2.display());
51            }
52            if let Some(clip2_tok) = &paths.clip_tokenizer_2 {
53                info!(clip2_tok = %clip2_tok.display());
54            }
55            for (i, te) in paths.text_encoder_files.iter().enumerate() {
56                info!(text_encoder_shard = i, path = %te.display());
57            }
58            if let Some(text_tok) = &paths.text_tokenizer {
59                info!(text_tok = %text_tok.display());
60            }
61
62            let offload = std::env::var("MOLD_OFFLOAD").is_ok_and(|v| v == "1");
63            let engine = mold_inference::create_engine(
64                model_name,
65                paths,
66                &config,
67                mold_inference::LoadStrategy::Eager,
68                offload,
69            )?;
70            state::AppState::new(engine, config, queue_handle)
71        }
72        None => {
73            info!("no default model configured — models will be pulled on first request");
74            state::AppState::empty(config, queue_handle)
75        }
76    };
77
78    // Spawn the generation queue worker — processes jobs sequentially (single GPU).
79    let worker_state = state.clone();
80    tokio::spawn(queue::run_queue_worker(job_rx, worker_state));
81
82    let cors = build_cors_layer()?;
83
84    let app = routes::create_router(state)
85        .layer(TraceLayer::new_for_http())
86        .layer(cors);
87
88    let addr: SocketAddr = format!("{bind}:{port}").parse()?;
89    info!(%addr, "starting mold server");
90
91    let listener = TcpListener::bind(addr).await?;
92    axum::serve(listener, app).await?;
93
94    Ok(())
95}
96
97fn build_cors_layer() -> Result<CorsLayer> {
98    let cors = match std::env::var("MOLD_CORS_ORIGIN") {
99        Ok(origin) if !origin.is_empty() => {
100            let origin = origin
101                .parse::<axum::http::HeaderValue>()
102                .map_err(|_| anyhow::anyhow!("invalid MOLD_CORS_ORIGIN value: {origin}"))?;
103            CorsLayer::new()
104                .allow_origin(origin)
105                .allow_methods([
106                    axum::http::Method::GET,
107                    axum::http::Method::POST,
108                    axum::http::Method::DELETE,
109                ])
110                .allow_headers(tower_http::cors::Any)
111                .expose_headers([axum::http::header::HeaderName::from_static(
112                    "x-mold-seed-used",
113                )])
114        }
115        _ => CorsLayer::permissive(),
116    };
117    Ok(cors)
118}
119
120#[cfg(test)]
121mod tests {
122    use super::build_cors_layer;
123    use std::sync::Mutex;
124
125    static ENV_LOCK: Mutex<()> = Mutex::new(());
126
127    #[test]
128    fn invalid_cors_origin_returns_error() {
129        let _lock = ENV_LOCK.lock().unwrap();
130        std::env::set_var("MOLD_CORS_ORIGIN", "\nnot-a-header");
131        let result = build_cors_layer();
132        std::env::remove_var("MOLD_CORS_ORIGIN");
133        assert!(result.is_err());
134    }
135
136    #[test]
137    fn valid_cors_origin_builds_layer() {
138        let _lock = ENV_LOCK.lock().unwrap();
139        std::env::set_var("MOLD_CORS_ORIGIN", "https://example.com");
140        let result = build_cors_layer();
141        std::env::remove_var("MOLD_CORS_ORIGIN");
142        assert!(result.is_ok());
143    }
144}
mold_server/lib.rs

mold_server/
lib.rs
