Construct server engines on their owner worker

crates/onnx-genai-server/src/lib.rs

@@ -8,6 +8,8 @@
 //! native-runtime, and execution-provider spans when available. OTLP span export
 //! is intentionally deferred (see issue #13).
 
+#![forbid(unsafe_code)]
+
 use std::{net::SocketAddr, time::Instant};
 
 use axum::{

crates/onnx-genai-server/src/main.rs

@@ -1,3 +1,5 @@
+#![forbid(unsafe_code)]
+
 use clap::Parser;
 use onnx_genai_server::{ServeArgs, run_serve};
 

crates/onnx-genai-server/src/multimodal.rs

@@ -255,7 +255,7 @@ pub fn load(model_dir: &Path) -> anyhow::Result<Option<PipelineSetup>> {
     {
         return Ok(None);
     }
-    let models = PipelineModels::load_with_ort_session_filter(
+    let _models = PipelineModels::load_with_ort_session_filter(
         model_dir,
         onnx_genai_ort::SessionOptions::default(),
         |_| false,
@@ -275,7 +275,7 @@ pub fn load(model_dir: &Path) -> anyhow::Result<Option<PipelineSetup>> {
             .and_then(|metadata| metadata.generation.as_ref())
             .and_then(|generation| generation.defaults.clone()),
         tokenizer_path: tokenizer_path(model_dir, &directory)?,
-        multimodal: build(&directory, &models)?,
+        multimodal: build(&directory)?,
     }))
 }
 
@@ -320,10 +320,7 @@ pub fn audio_decoder_prompt(
     Ok(tokens)
 }
 
-pub fn build(
-    directory: &PipelineModelDirectory,
-    _models: &PipelineModels,
-) -> anyhow::Result<MultimodalSpecs> {
+pub fn build(directory: &PipelineModelDirectory) -> anyhow::Result<MultimodalSpecs> {
     derive_specs(
         &directory.spec.workflow,
         directory.preprocessing.as_ref(),

crates/onnx-genai-server/src/registry.rs

@@ -91,7 +91,7 @@ pub(crate) struct ModelHandleParts {
 }
 
 impl ModelHandle {
-    pub(crate) fn new(parts: ModelHandleParts) -> anyhow::Result<Self> {
+    pub(crate) fn new(parts: ModelHandleParts) -> Self {
         let ModelHandleParts {
             id,
             model_dir,
@@ -111,7 +111,7 @@ impl ModelHandle {
         // that keys a session or a lease reads it back off the driver, which is
         // why the two can never disagree.
         engine.bind_model(&id);
-        Ok(Self {
+        Self {
             id,
             engine,
             tokenizer,
@@ -126,7 +126,7 @@ impl ModelHandle {
             last_request_at: AtomicU64::new(now_millis()),
             warmed: AtomicBool::new(false),
             warmup_lock: std::sync::Mutex::new(()),
-        })
+        }
     }
 
     /// Stop this model's engine worker and wait for it to exit.
@@ -244,11 +244,10 @@ impl std::error::Error for WarmupError {
 /// `Arc`s, so `AppState` (and therefore every Axum request) shares one registry.
 /// All mutable state lives behind `Arc<RwLock<RegistryInner>>`.
 ///
-/// **Locking discipline:** the heavy model build (`Engine::from_dir`, tokenizer
-/// and chat template) is always performed *outside* the lock via
-/// `tokio::task::spawn_blocking`; the `RwLock` is only ever held for the short,
-/// synchronous critical sections that mutate the maps. No lock is ever held
-/// across an `.await`, so the synchronous `std::sync::RwLock` is deadlock-free here.
+/// **Locking discipline:** model preparation and the wait for worker readiness
+/// happen outside the lock. Lazy loading waits from `tokio::task::spawn_blocking`,
+/// while `Engine::from_dir` itself runs on the new owner worker. The `RwLock` is
+/// held only for short synchronous map mutations and never across an `.await`.
 #[derive(Clone)]
 pub(crate) struct ModelRegistry {
     inner: Arc<RwLock<RegistryInner>>,
@@ -433,33 +432,6 @@ impl ModelRegistry {
         Ok(registry)
     }
 
-    /// Build a registry around a single, already-constructed handle.
-    ///
-    /// Used by the `AppState::new*` constructors that start from a live `Engine`
-    /// rather than a spec path.  Because there is no backing spec, the model is
-    /// not recorded in `available` and therefore cannot be lazily reloaded after
-    /// an unload.
-    pub(crate) fn from_handle(handle: Arc<ModelHandle>, config: ServerConfig) -> Self {
-        let authority_provider = Arc::new(ServerMemoryAuthorities::new(
-            config.engine_config.limits.vram_limit,
-        ));
-        let default_id = Some(handle.id.clone());
-        let mut inner = RegistryInner {
-            models: HashMap::new(),
-            order: Vec::new(),
-            default_id,
-            available: HashMap::new(),
-        };
-        inner.insert_loaded(handle);
-        Self {
-            inner: Arc::new(RwLock::new(inner)),
-            config: Arc::new(config),
-            load_locks: Arc::new(tokio::sync::Mutex::new(HashMap::new())),
-            policy_lock: Arc::new(tokio::sync::RwLock::new(())),
-            authority_provider,
-        }
-    }
-
     /// Resolve an already-loaded handle by name, updating `last_request_at`.
     ///
     /// - **Empty / whitespace** — falls back to the default model.
@@ -529,9 +501,10 @@ impl ModelRegistry {
 
     /// Load (or return the already-loaded) model for `id`.
     ///
-    /// The heavy construction runs on a blocking thread pool via
-    /// `spawn_blocking`; the registry lock is only taken for the brief insert +
-    /// eviction critical section afterwards.  A per-id async guard serialises
+    /// Model preparation and the blocking wait for the owner's ready handshake
+    /// run on a blocking thread pool; engine construction runs on the owner
+    /// worker itself. The registry lock is only taken for the brief insert +
+    /// eviction critical section afterwards. A per-id async guard serialises
     /// concurrent loads of the same id so the model is built only once.
     pub(crate) async fn load(&self, id: &str) -> anyhow::Result<Arc<ModelHandle>> {
         // Fast path: already loaded.
@@ -1019,8 +992,7 @@ mod tests {
             multimodal: None,
             speech: None,
             image_pipeline: None,
-        })
-        .expect("handle");
+        });
 
         assert!(
             handle.speech.is_none(),
@@ -1321,6 +1293,31 @@ mod tests {
         assert!(registry.contains_available("tiny").unwrap());
     }
 
+    #[test]
+    fn dropping_the_last_registry_joins_the_engine_worker() {
+        let model_dir =
+            PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("../../tests/fixtures/tiny-llm");
+        let specs = vec![ModelSpec {
+            id: "tiny".to_string(),
+            path: model_dir,
+            eager: true,
+            warmup: false,
+        }];
+        let registry = ModelRegistry::from_specs(&specs, ServerConfig::default()).unwrap();
+        let worker = {
+            let handle = registry.resolve("tiny").unwrap().unwrap();
+            Arc::clone(handle.engine.workers.primary())
+        };
+
+        drop(registry);
+
+        assert!(
+            worker.is_joined(),
+            "dropping the last registry must wait for engine teardown",
+        );
+        assert!(!worker.is_running());
+    }
+
     /// An unload racing an in-flight request does not tear the engine out from
     /// under it: the worker keeps running for whoever still holds the handle.
     #[tokio::test]

crates/onnx-genai-server/src/state.rs

@@ -504,65 +504,6 @@ impl AppState {
         })
     }
 
-    pub fn new(model_id: String, engine: Engine, tokenizer: Tokenizer) -> Self {
-        Self::new_with_template(model_id, engine, tokenizer, None)
-    }
-
-    pub fn new_with_template(
-        model_id: String,
-        engine: Engine,
-        tokenizer: Tokenizer,
-        chat_template: Option<ChatTemplate>,
-    ) -> Self {
-        Self::new_with_template_and_config(
-            model_id,
-            engine,
-            tokenizer,
-            chat_template,
-            ServerConfig::default(),
-            None,
-        )
-    }
-
-    fn new_with_template_and_config(
-        model_id: String,
-        engine: Engine,
-        tokenizer: Tokenizer,
-        chat_template: Option<ChatTemplate>,
-        config: ServerConfig,
-        model_max_context: Option<usize>,
-    ) -> Self {
-        let config = config.validate().expect("validated server config");
-        let fim_config = engine.fim_config().cloned();
-        let engine_driver = EngineDriver::start(
-            engine,
-            config.max_batch.unwrap_or(DEFAULT_MAX_BATCH),
-            config.max_queue_depth,
-        );
-        let handle = ModelHandle::new(ModelHandleParts {
-            id: model_id,
-            // Test-only constructor: the model was handed in already loaded, so
-            // there is no package directory to resolve files against.
-            model_dir: std::path::PathBuf::new(),
-            engine: engine_driver,
-            tokenizer: Arc::new(tokenizer),
-            chat_template: chat_template.map(Arc::new),
-            model_max_context,
-            generation_defaults: None,
-            fim_config,
-            multimodal: None,
-            speech: None,
-            image_pipeline: None,
-        })
-        .expect("test model handle");
-        let registry = ModelRegistry::from_handle(Arc::new(handle), config.clone());
-        Self {
-            registry,
-            sessions: SessionRegistry::new(config.max_sessions),
-            config,
-        }
-    }
-
     /// Returns the id of the first loaded model, for use in log messages and the CLI.
     pub fn model_id(&self) -> String {
         match self.registry.default_id() {
@@ -655,9 +596,10 @@ pub(crate) fn enforce_requested_max_batch(
 ///
 /// `config` must already be validated.  This is the single shared construction
 /// path used by both startup (`ModelRegistry::from_specs`) and runtime lazy
-/// loading (`ModelRegistry::load`).  It is a **blocking** function (it calls
-/// `Engine::from_dir`, which takes seconds) and must therefore be invoked from a
-/// blocking context (e.g. at startup or via `tokio::task::spawn_blocking`).
+/// loading (`ModelRegistry::load`). It blocks until the new worker reports its
+/// typed initialization result, so async callers must invoke it from a blocking
+/// context. The engine itself is constructed inside that worker, never on the
+/// caller's blocking-pool thread.
 pub(crate) fn build_handle_with_authorities(
     spec: &ModelSpec,
     config: &ServerConfig,
@@ -707,21 +649,29 @@ pub(crate) fn build_handle_with_authorities(
     let model_max_context = load_model_max_context(model_directory.metadata_path.as_deref())?;
     let tokenizer = Tokenizer::from_file(&model_directory.tokenizer_path)
         .map_err(|e| anyhow::anyhow!("Failed to load tokenizer: {e}"))?;
-    let engine = Engine::from_dir_with_memory_authority_provider(
-        model_dir,
-        config.engine_config.clone(),
-        authorities,
-    )?;
-    let fim_config = engine.fim_config().cloned();
     let generation_defaults = None;
-    // Refuse an explicit `--max-batch > 1` the decode path cannot honor, rather
-    // than accepting it and silently falling back to per-request decoding. The
-    // check is sourced from the engine's decode path, not from whether an ORT
-    // session exists (issue #750).
-    enforce_requested_max_batch(&engine, config.max_batch)?;
     let requested_max_batch = config.max_batch.unwrap_or(DEFAULT_MAX_BATCH);
-    let engine_driver = EngineDriver::start(engine, requested_max_batch, config.max_queue_depth);
-    ModelHandle::new(ModelHandleParts {
+    let engine_model_dir = model_dir.to_path_buf();
+    let engine_config = config.engine_config.clone();
+    let configured_max_batch = config.max_batch;
+    let (engine_driver, fim_config) = EngineDriver::start(
+        move || {
+            let engine = Engine::from_dir_with_memory_authority_provider(
+                &engine_model_dir,
+                engine_config,
+                authorities,
+            )?;
+            // Refuse an explicit `--max-batch > 1` the decode path cannot honor,
+            // rather than accepting it and silently falling back to per-request
+            // decoding. This runs beside the engine before readiness is reported.
+            enforce_requested_max_batch(&engine, configured_max_batch)?;
+            let fim_config = engine.fim_config().cloned();
+            Ok((engine, fim_config))
+        },
+        requested_max_batch,
+        config.max_queue_depth,
+    )?;
+    Ok(ModelHandle::new(ModelHandleParts {
         id: model_id,
         model_dir: model_dir.to_path_buf(),
         engine: engine_driver,
@@ -733,7 +683,7 @@ pub(crate) fn build_handle_with_authorities(
         multimodal: None,
         speech: None,
         image_pipeline: None,
-    })
+    }))
 }
 
 fn build_pipeline_handle(
@@ -749,13 +699,7 @@ fn build_pipeline_handle(
     let tokenizer_path = crate::multimodal::tokenizer_path(model_dir, &directory)?;
     let tokenizer = Tokenizer::from_file(&tokenizer_path)
         .map_err(|e| anyhow::anyhow!("Failed to load pipeline tokenizer: {e}"))?;
-
-    let engine = Engine::from_dir_with_memory_authority_provider(
-        model_dir,
-        config.engine_config.clone(),
-        authorities,
-    )?;
-    let multimodal = crate::multimodal::build(&directory, engine.models()?)?;
+    let multimodal = crate::multimodal::build(&directory)?;
     // Resolve one exact speech capability at load time: bind the single
     // text-assembly processor to the single compatible buffered PCM16 WAV audio
     // output that serving will encode. Fail closed on zero, ambiguous, or
@@ -789,14 +733,30 @@ fn build_pipeline_handle(
     // superseded generation metadata surfaces, so the pipeline path resolves
     // sampling from request options only — same as the non-pipeline path above.
     let generation_defaults = None;
-    // The same admission the decoder path applies: an explicit `--max-batch`
-    // the decode path cannot honor is refused rather than silently ignored.
-    enforce_requested_max_batch(&engine, config.max_batch)?;
     let requested_max_batch = config.max_batch.unwrap_or(DEFAULT_MAX_BATCH);
-    ModelHandle::new(ModelHandleParts {
+    let engine_model_dir = model_dir.to_path_buf();
+    let engine_config = config.engine_config.clone();
+    let configured_max_batch = config.max_batch;
+    let (engine_driver, ()) = EngineDriver::start(
+        move || {
+            let engine = Engine::from_dir_with_memory_authority_provider(
+                &engine_model_dir,
+                engine_config,
+                authorities,
+            )?;
+            // The same admission the decoder path applies: an explicit
+            // `--max-batch` the decode path cannot honor is refused rather than
+            // silently ignored.
+            enforce_requested_max_batch(&engine, configured_max_batch)?;
+            Ok((engine, ()))
+        },
+        requested_max_batch,
+        config.max_queue_depth,
+    )?;
+    Ok(ModelHandle::new(ModelHandleParts {
         id: model_id,
         model_dir: model_dir.to_path_buf(),
-        engine: EngineDriver::start(engine, requested_max_batch, config.max_queue_depth),
+        engine: engine_driver,
         tokenizer: Arc::new(tokenizer),
         chat_template: chat_template.map(Arc::new),
         model_max_context,
@@ -805,7 +765,7 @@ fn build_pipeline_handle(
         multimodal: Some(multimodal),
         speech,
         image_pipeline,
-    })
+    }))
 }
 
 #[cfg(test)]