uni_xervo/
runtime.rs

1//! The core runtime that manages providers, catalogs, and loaded model instances.
2
3use crate::api::{ModelAliasSpec, ModelRuntimeKey};
4use crate::error::{Result, RuntimeError};
5use crate::options_validation::validate_provider_options;
6use crate::reliability::{
7    InstrumentedAudioEmbeddingModel, InstrumentedDocumentExtractionModel,
8    InstrumentedEmbeddingModel, InstrumentedGeneratorModel, InstrumentedHybridEmbeddingModel,
9    InstrumentedImageEmbeddingModel, InstrumentedMultiVectorEmbeddingModel,
10    InstrumentedMultimodalEmbeddingModel, InstrumentedNlpModel, InstrumentedOcrModel,
11    InstrumentedRawTensorModel, InstrumentedRerankerModel, InstrumentedSparseEmbeddingModel,
12    InstrumentedTranscriptionModel,
13};
14use crate::traits::{
15    AudioEmbeddingModel, DocumentExtractionModel, EmbeddingModel, GeneratorModel,
16    HybridEmbeddingModel, ImageEmbeddingModel, LoadedModelHandle, ModelProvider,
17    MultiVectorEmbeddingModel, MultimodalEmbeddingModel, NlpModel, OcrModel, RawTensorModel,
18    RerankerModel, SparseEmbeddingModel, TranscriptionModel,
19};
20use dashmap::DashMap;
21use std::any::Any;
22use std::collections::HashMap;
23use std::sync::Arc;
24use tokio::sync::{Mutex, RwLock};
25
26/// Per-alias cache of fully-instrumented typed handles.
27///
28/// Eliminates repeated spec lookups, key hashing, registry reads, and
29/// wrapper allocations on hot paths where the same alias is resolved
30/// many times (e.g. per-turn inference in a pipeline).
31///
32/// Uses [`DashMap`] for lock-free concurrent reads on the hot path.
33/// Entries are populated on first successful resolve via
34/// `entry().or_insert_with(...)`, so concurrent first-callers all observe
35/// the same `Arc` and the wrapper is allocated exactly once per alias.
36/// Once populated, an entry is never mutated.
37#[derive(Default)]
38struct HandleCache {
39    embeddings: DashMap<String, Arc<dyn EmbeddingModel>>,
40    rerankers: DashMap<String, Arc<dyn RerankerModel>>,
41    generators: DashMap<String, Arc<dyn GeneratorModel>>,
42    raw_tensor_models: DashMap<String, Arc<dyn RawTensorModel>>,
43    // Multimodal extension surface — added in Phase 6, instrumentation
44    // filled in by Phase 7.
45    image_embedders: DashMap<String, Arc<dyn ImageEmbeddingModel>>,
46    audio_embedders: DashMap<String, Arc<dyn AudioEmbeddingModel>>,
47    multimodal_embedders: DashMap<String, Arc<dyn MultimodalEmbeddingModel>>,
48    sparse_embedders: DashMap<String, Arc<dyn SparseEmbeddingModel>>,
49    multi_vector_embedders: DashMap<String, Arc<dyn MultiVectorEmbeddingModel>>,
50    hybrid_embedders: DashMap<String, Arc<dyn HybridEmbeddingModel>>,
51    nlp_models: DashMap<String, Arc<dyn NlpModel>>,
52    document_extractors: DashMap<String, Arc<dyn DocumentExtractionModel>>,
53    transcribers: DashMap<String, Arc<dyn TranscriptionModel>>,
54    ocr_models: DashMap<String, Arc<dyn OcrModel>>,
55}
56
57/// Default load timeout applied when [`ModelAliasSpec::load_timeout`] is `None`.
58const DEFAULT_LOAD_TIMEOUT_SECS: u64 = 600;
59
60/// The central runtime that owns registered providers and a catalog of model
61/// aliases.
62///
63/// Obtain an instance via [`ModelRuntime::builder()`] and the
64/// [`ModelRuntimeBuilder`].  Once built, use [`embedding`](Self::embedding),
65/// [`reranker`](Self::reranker), or [`generator`](Self::generator) to obtain
66/// typed, instrumented model handles.
67///
68/// Models are loaded lazily on first access (unless configured for eager or
69/// background warmup) and cached in an internal registry so that subsequent
70/// requests for the same model are served instantly.
71pub struct ModelRuntime {
72    providers: HashMap<String, Box<dyn ModelProvider>>,
73    registry: Arc<ModelRegistry>,
74    catalog: RwLock<HashMap<String, ModelAliasSpec>>,
75    handle_cache: HandleCache,
76}
77
78/// Internal registry that caches loaded model instances and coordinates
79/// concurrent load requests to prevent duplicate work.
80#[derive(Default)]
81pub struct ModelRegistry {
82    instances: RwLock<HashMap<ModelRuntimeKey, LoadedModelHandle>>,
83    /// Per-key mutexes to prevent concurrent loads of the same model.
84    loader_locks: Mutex<HashMap<ModelRuntimeKey, Arc<Mutex<()>>>>,
85}
86
87impl ModelRuntime {
88    /// Create a new [`ModelRuntimeBuilder`] for configuring and constructing a
89    /// runtime.
90    pub fn builder() -> ModelRuntimeBuilder {
91        ModelRuntimeBuilder::default()
92    }
93
94    /// Register a new model alias at runtime.
95    pub async fn register(&self, spec: ModelAliasSpec) -> Result<()> {
96        spec.validate()?;
97        if !self.providers.contains_key(&spec.provider_id) {
98            return Err(RuntimeError::Config(format!(
99                "Unknown provider '{}' for alias '{}'",
100                spec.provider_id, spec.alias
101            )));
102        }
103        validate_provider_options(&spec.provider_id, spec.task, &spec.options)?;
104        let mut catalog = self.catalog.write().await;
105        if catalog.contains_key(&spec.alias) {
106            return Err(RuntimeError::Config(format!(
107                "Alias '{}' already exists",
108                spec.alias
109            )));
110        }
111        catalog.insert(spec.alias.clone(), spec);
112        Ok(())
113    }
114
115    /// Check if an alias exists in the catalog.
116    pub async fn contains_alias(&self, alias: &str) -> bool {
117        let catalog = self.catalog.read().await;
118        catalog.contains_key(alias)
119    }
120
121    /// Look up a spec by alias, returning an error if not found.
122    async fn lookup_spec(&self, alias: &str) -> Result<ModelAliasSpec> {
123        let catalog = self.catalog.read().await;
124        catalog
125            .get(alias)
126            .cloned()
127            .ok_or_else(|| RuntimeError::AliasNotFound {
128                alias: alias.to_string(),
129            })
130    }
131
132    /// Pre-load and cache every model in the catalog.
133    ///
134    /// Models already loaded are skipped. Fails fast on the first error.
135    /// Call this during application startup to avoid cold-start latency on
136    /// first inference.
137    pub async fn prefetch_all(&self) -> Result<()> {
138        let specs: Vec<ModelAliasSpec> = {
139            let catalog = self.catalog.read().await;
140            catalog.values().cloned().collect()
141        };
142        for spec in specs {
143            tracing::info!(alias = %spec.alias, "Prefetching model");
144            self.resolve_and_load_internal(&spec).await?;
145        }
146        Ok(())
147    }
148
149    /// Pre-load and cache specific aliases.
150    ///
151    /// Returns an error immediately if an alias is not found in the catalog
152    /// or if any model fails to load. Models already loaded are skipped.
153    pub async fn prefetch(&self, aliases: &[&str]) -> Result<()> {
154        for alias in aliases {
155            let spec = self.lookup_spec(alias).await?;
156            tracing::info!(alias = %alias, "Prefetching model");
157            self.resolve_and_load_internal(&spec).await?;
158        }
159        Ok(())
160    }
161
162    /// Resolve, load (if necessary), and return an instrumented [`EmbeddingModel`]
163    /// handle for the given alias.
164    ///
165    /// The returned handle is cached per alias so that repeated calls skip
166    /// spec lookup, key hashing, and wrapper allocation.
167    pub async fn embedding(&self, alias: &str) -> Result<Arc<dyn EmbeddingModel>> {
168        if let Some(cached) = self.handle_cache.embeddings.get(alias) {
169            return Ok(cached.clone());
170        }
171
172        let spec = self.lookup_spec(alias).await?;
173        let handle = self.resolve_and_load_internal(&spec).await?;
174        if let Some(model) = handle.downcast_ref::<Arc<dyn EmbeddingModel>>() {
175            let cached = self
176                .handle_cache
177                .embeddings
178                .entry(alias.to_string())
179                .or_insert_with(|| {
180                    let wrapper: Arc<dyn EmbeddingModel> = Arc::new(InstrumentedEmbeddingModel {
181                        inner: model.clone(),
182                        alias: alias.to_string(),
183                        provider_id: spec.provider_id.clone(),
184                        timeout: spec.timeout.map(std::time::Duration::from_secs),
185                        retry: spec.retry.clone(),
186                    });
187                    wrapper
188                })
189                .clone();
190            return Ok(cached);
191        }
192
193        Err(RuntimeError::CapabilityMismatch(format!(
194            "Model for alias '{}' does not implement EmbeddingModel",
195            alias
196        )))
197    }
198
199    /// Resolve a dense text [`EmbeddingModel`] by alias.
200    ///
201    /// Agent-noun alias for [`embedding`](Self::embedding), matching the
202    /// `image_embedder` / `sparse_embedder` / `multi_vector_embedder` naming.
203    pub async fn embedder(&self, alias: &str) -> Result<Arc<dyn EmbeddingModel>> {
204        self.embedding(alias).await
205    }
206
207    /// Resolve, load (if necessary), and return an instrumented [`RerankerModel`]
208    /// handle for the given alias.
209    ///
210    /// The returned handle is cached per alias so that repeated calls skip
211    /// spec lookup, key hashing, and wrapper allocation.
212    pub async fn reranker(&self, alias: &str) -> Result<Arc<dyn RerankerModel>> {
213        if let Some(cached) = self.handle_cache.rerankers.get(alias) {
214            return Ok(cached.clone());
215        }
216
217        let spec = self.lookup_spec(alias).await?;
218        let handle = self.resolve_and_load_internal(&spec).await?;
219        if let Some(model) = handle.downcast_ref::<Arc<dyn RerankerModel>>() {
220            let cached = self
221                .handle_cache
222                .rerankers
223                .entry(alias.to_string())
224                .or_insert_with(|| {
225                    let wrapper: Arc<dyn RerankerModel> = Arc::new(InstrumentedRerankerModel {
226                        inner: model.clone(),
227                        alias: alias.to_string(),
228                        provider_id: spec.provider_id.clone(),
229                        timeout: spec.timeout.map(std::time::Duration::from_secs),
230                        retry: spec.retry.clone(),
231                    });
232                    wrapper
233                })
234                .clone();
235            return Ok(cached);
236        }
237        Err(RuntimeError::CapabilityMismatch(format!(
238            "Model for alias '{}' does not implement RerankerModel",
239            alias
240        )))
241    }
242
243    /// Resolve, load (if necessary), and return an instrumented [`GeneratorModel`]
244    /// handle for the given alias.
245    ///
246    /// The returned handle is cached per alias so that repeated calls skip
247    /// spec lookup, key hashing, and wrapper allocation.
248    pub async fn generator(&self, alias: &str) -> Result<Arc<dyn GeneratorModel>> {
249        if let Some(cached) = self.handle_cache.generators.get(alias) {
250            return Ok(cached.clone());
251        }
252
253        let spec = self.lookup_spec(alias).await?;
254        let handle = self.resolve_and_load_internal(&spec).await?;
255        if let Some(model) = handle.downcast_ref::<Arc<dyn GeneratorModel>>() {
256            let cached = self
257                .handle_cache
258                .generators
259                .entry(alias.to_string())
260                .or_insert_with(|| {
261                    let wrapper: Arc<dyn GeneratorModel> = Arc::new(InstrumentedGeneratorModel {
262                        inner: model.clone(),
263                        alias: alias.to_string(),
264                        provider_id: spec.provider_id.clone(),
265                        timeout: spec.timeout.map(std::time::Duration::from_secs),
266                        retry: spec.retry.clone(),
267                    });
268                    wrapper
269                })
270                .clone();
271            return Ok(cached);
272        }
273        Err(RuntimeError::CapabilityMismatch(format!(
274            "Model for alias '{}' does not implement GeneratorModel",
275            alias
276        )))
277    }
278
279    /// Resolve, load (if necessary), and return an instrumented [`RawTensorModel`]
280    /// handle for the given alias.
281    ///
282    /// The returned handle is cached per alias so that repeated calls skip
283    /// spec lookup, key hashing, and wrapper allocation.
284    pub async fn raw_tensor_model(&self, alias: &str) -> Result<Arc<dyn RawTensorModel>> {
285        if let Some(cached) = self.handle_cache.raw_tensor_models.get(alias) {
286            return Ok(cached.clone());
287        }
288
289        let spec = self.lookup_spec(alias).await?;
290        let handle = self.resolve_and_load_internal(&spec).await?;
291        if let Some(model) = handle.downcast_ref::<Arc<dyn RawTensorModel>>() {
292            let cached = self
293                .handle_cache
294                .raw_tensor_models
295                .entry(alias.to_string())
296                .or_insert_with(|| {
297                    let wrapper: Arc<dyn RawTensorModel> = Arc::new(InstrumentedRawTensorModel {
298                        inner: model.clone(),
299                        alias: alias.to_string(),
300                        provider_id: spec.provider_id.clone(),
301                        timeout: spec.timeout.map(std::time::Duration::from_secs),
302                        retry: spec.retry.clone(),
303                    });
304                    wrapper
305                })
306                .clone();
307            return Ok(cached);
308        }
309
310        Err(RuntimeError::ProviderCapabilityMissing {
311            alias: alias.to_string(),
312            provider_id: spec.provider_id,
313            capability: "RawTensorModel".to_string(),
314        })
315    }
316
317    /// Resolve, load (if necessary), and return an instrumented
318    /// [`ImageEmbeddingModel`] handle for the given alias.
319    ///
320    /// # Examples
321    /// ```no_run
322    /// # use std::sync::Arc;
323    /// # use uni_xervo::runtime::ModelRuntime;
324    /// # use uni_xervo::traits::ImageInput;
325    /// # async fn run(runtime: Arc<ModelRuntime>) -> uni_xervo::error::Result<()> {
326    /// let embedder = runtime.image_embedder("embed/siglip").await?;
327    /// let image = ImageInput::Bytes {
328    ///     data: std::fs::read("photo.png").unwrap(),
329    ///     media_type: "image/png".to_string(),
330    /// };
331    /// let result = embedder.embed(vec![image]).await?;
332    /// println!("dimension: {}", result.vectors[0].len());
333    /// # Ok(())
334    /// # }
335    /// ```
336    ///
337    /// # Errors
338    /// Returns an error if the alias is unknown, the model fails to load, or
339    /// the provider does not implement image embedding.
340    pub async fn image_embedder(&self, alias: &str) -> Result<Arc<dyn ImageEmbeddingModel>> {
341        if let Some(cached) = self.handle_cache.image_embedders.get(alias) {
342            return Ok(cached.clone());
343        }
344        let spec = self.lookup_spec(alias).await?;
345        let handle = self.resolve_and_load_internal(&spec).await?;
346        if let Some(model) = handle.downcast_ref::<Arc<dyn ImageEmbeddingModel>>() {
347            let cached = self
348                .handle_cache
349                .image_embedders
350                .entry(alias.to_string())
351                .or_insert_with(|| {
352                    let wrapper: Arc<dyn ImageEmbeddingModel> =
353                        Arc::new(InstrumentedImageEmbeddingModel {
354                            inner: model.clone(),
355                            alias: alias.to_string(),
356                            provider_id: spec.provider_id.clone(),
357                            timeout: spec.timeout.map(std::time::Duration::from_secs),
358                            retry: spec.retry.clone(),
359                        });
360                    wrapper
361                })
362                .clone();
363            return Ok(cached);
364        }
365        Err(RuntimeError::ProviderCapabilityMissing {
366            alias: alias.to_string(),
367            provider_id: spec.provider_id,
368            capability: "ImageEmbeddingModel".to_string(),
369        })
370    }
371
372    /// Resolve, load (if necessary), and return an instrumented
373    /// [`AudioEmbeddingModel`] handle for the given alias.
374    pub async fn audio_embedder(&self, alias: &str) -> Result<Arc<dyn AudioEmbeddingModel>> {
375        if let Some(cached) = self.handle_cache.audio_embedders.get(alias) {
376            return Ok(cached.clone());
377        }
378        let spec = self.lookup_spec(alias).await?;
379        let handle = self.resolve_and_load_internal(&spec).await?;
380        if let Some(model) = handle.downcast_ref::<Arc<dyn AudioEmbeddingModel>>() {
381            let cached = self
382                .handle_cache
383                .audio_embedders
384                .entry(alias.to_string())
385                .or_insert_with(|| {
386                    let wrapper: Arc<dyn AudioEmbeddingModel> =
387                        Arc::new(InstrumentedAudioEmbeddingModel {
388                            inner: model.clone(),
389                            alias: alias.to_string(),
390                            provider_id: spec.provider_id.clone(),
391                            timeout: spec.timeout.map(std::time::Duration::from_secs),
392                            retry: spec.retry.clone(),
393                        });
394                    wrapper
395                })
396                .clone();
397            return Ok(cached);
398        }
399        Err(RuntimeError::ProviderCapabilityMissing {
400            alias: alias.to_string(),
401            provider_id: spec.provider_id,
402            capability: "AudioEmbeddingModel".to_string(),
403        })
404    }
405
406    /// Resolve, load (if necessary), and return an instrumented
407    /// [`MultimodalEmbeddingModel`] handle for the given alias.
408    pub async fn multimodal_embedder(
409        &self,
410        alias: &str,
411    ) -> Result<Arc<dyn MultimodalEmbeddingModel>> {
412        if let Some(cached) = self.handle_cache.multimodal_embedders.get(alias) {
413            return Ok(cached.clone());
414        }
415        let spec = self.lookup_spec(alias).await?;
416        let handle = self.resolve_and_load_internal(&spec).await?;
417        if let Some(model) = handle.downcast_ref::<Arc<dyn MultimodalEmbeddingModel>>() {
418            let cached = self
419                .handle_cache
420                .multimodal_embedders
421                .entry(alias.to_string())
422                .or_insert_with(|| {
423                    let wrapper: Arc<dyn MultimodalEmbeddingModel> =
424                        Arc::new(InstrumentedMultimodalEmbeddingModel {
425                            inner: model.clone(),
426                            alias: alias.to_string(),
427                            provider_id: spec.provider_id.clone(),
428                            timeout: spec.timeout.map(std::time::Duration::from_secs),
429                            retry: spec.retry.clone(),
430                        });
431                    wrapper
432                })
433                .clone();
434            return Ok(cached);
435        }
436        Err(RuntimeError::ProviderCapabilityMissing {
437            alias: alias.to_string(),
438            provider_id: spec.provider_id,
439            capability: "MultimodalEmbeddingModel".to_string(),
440        })
441    }
442
443    /// Resolve, load (if necessary), and return an instrumented
444    /// [`SparseEmbeddingModel`] handle for the given alias.
445    ///
446    /// # Errors
447    /// Returns an error if the alias is unknown, the model fails to load, or
448    /// the provider does not implement sparse embedding.
449    pub async fn sparse_embedder(&self, alias: &str) -> Result<Arc<dyn SparseEmbeddingModel>> {
450        if let Some(cached) = self.handle_cache.sparse_embedders.get(alias) {
451            return Ok(cached.clone());
452        }
453        let spec = self.lookup_spec(alias).await?;
454        let handle = self.resolve_and_load_internal(&spec).await?;
455        if let Some(model) = handle.downcast_ref::<Arc<dyn SparseEmbeddingModel>>() {
456            let cached = self
457                .handle_cache
458                .sparse_embedders
459                .entry(alias.to_string())
460                .or_insert_with(|| {
461                    let wrapper: Arc<dyn SparseEmbeddingModel> =
462                        Arc::new(InstrumentedSparseEmbeddingModel {
463                            inner: model.clone(),
464                            alias: alias.to_string(),
465                            provider_id: spec.provider_id.clone(),
466                            timeout: spec.timeout.map(std::time::Duration::from_secs),
467                            retry: spec.retry.clone(),
468                        });
469                    wrapper
470                })
471                .clone();
472            return Ok(cached);
473        }
474        Err(RuntimeError::ProviderCapabilityMissing {
475            alias: alias.to_string(),
476            provider_id: spec.provider_id,
477            capability: "SparseEmbeddingModel".to_string(),
478        })
479    }
480
481    /// Resolve, load (if necessary), and return an instrumented
482    /// [`MultiVectorEmbeddingModel`] handle for the given alias.
483    ///
484    /// # Errors
485    /// Returns an error if the alias is unknown, the model fails to load, or
486    /// the provider does not implement multi-vector embedding.
487    pub async fn multi_vector_embedder(
488        &self,
489        alias: &str,
490    ) -> Result<Arc<dyn MultiVectorEmbeddingModel>> {
491        if let Some(cached) = self.handle_cache.multi_vector_embedders.get(alias) {
492            return Ok(cached.clone());
493        }
494        let spec = self.lookup_spec(alias).await?;
495        let handle = self.resolve_and_load_internal(&spec).await?;
496        if let Some(model) = handle.downcast_ref::<Arc<dyn MultiVectorEmbeddingModel>>() {
497            let cached = self
498                .handle_cache
499                .multi_vector_embedders
500                .entry(alias.to_string())
501                .or_insert_with(|| {
502                    let wrapper: Arc<dyn MultiVectorEmbeddingModel> =
503                        Arc::new(InstrumentedMultiVectorEmbeddingModel {
504                            inner: model.clone(),
505                            alias: alias.to_string(),
506                            provider_id: spec.provider_id.clone(),
507                            timeout: spec.timeout.map(std::time::Duration::from_secs),
508                            retry: spec.retry.clone(),
509                        });
510                    wrapper
511                })
512                .clone();
513            return Ok(cached);
514        }
515        Err(RuntimeError::ProviderCapabilityMissing {
516            alias: alias.to_string(),
517            provider_id: spec.provider_id,
518            capability: "MultiVectorEmbeddingModel".to_string(),
519        })
520    }
521
522    /// Resolve, load (if necessary), and return an instrumented
523    /// [`HybridEmbeddingModel`] handle for the given alias.
524    ///
525    /// The handle serves dense, sparse, and multi-vector heads from a single
526    /// forward pass; select which to materialize with a
527    /// [`HeadSet`](crate::traits::HeadSet). Only multi-output graphs with a
528    /// hybrid preset (e.g. `BGEM3Hybrid`) resolve here — single-head models use
529    /// the per-task resolvers.
530    ///
531    /// # Errors
532    /// Returns an error if the alias is unknown, the model has no hybrid preset,
533    /// the model fails to load, or the loaded handle lacks the hybrid capability.
534    pub async fn hybrid_embedder(&self, alias: &str) -> Result<Arc<dyn HybridEmbeddingModel>> {
535        if let Some(cached) = self.handle_cache.hybrid_embedders.get(alias) {
536            return Ok(cached.clone());
537        }
538        let spec = self.lookup_spec(alias).await?;
539        let handle = self.resolve_and_load_internal(&spec).await?;
540        if let Some(model) = handle.downcast_ref::<Arc<dyn HybridEmbeddingModel>>() {
541            let cached = self
542                .handle_cache
543                .hybrid_embedders
544                .entry(alias.to_string())
545                .or_insert_with(|| {
546                    let wrapper: Arc<dyn HybridEmbeddingModel> =
547                        Arc::new(InstrumentedHybridEmbeddingModel {
548                            inner: model.clone(),
549                            alias: alias.to_string(),
550                            provider_id: spec.provider_id.clone(),
551                            timeout: spec.timeout.map(std::time::Duration::from_secs),
552                            retry: spec.retry.clone(),
553                        });
554                    wrapper
555                })
556                .clone();
557            return Ok(cached);
558        }
559        Err(RuntimeError::ProviderCapabilityMissing {
560            alias: alias.to_string(),
561            provider_id: spec.provider_id,
562            capability: "HybridEmbeddingModel".to_string(),
563        })
564    }
565
566    /// Resolve, load (if necessary), and return an instrumented [`NlpModel`]
567    /// handle for the given alias.
568    pub async fn nlp_model(&self, alias: &str) -> Result<Arc<dyn NlpModel>> {
569        if let Some(cached) = self.handle_cache.nlp_models.get(alias) {
570            return Ok(cached.clone());
571        }
572        let spec = self.lookup_spec(alias).await?;
573        let handle = self.resolve_and_load_internal(&spec).await?;
574        if let Some(model) = handle.downcast_ref::<Arc<dyn NlpModel>>() {
575            let cached = self
576                .handle_cache
577                .nlp_models
578                .entry(alias.to_string())
579                .or_insert_with(|| {
580                    let wrapper: Arc<dyn NlpModel> = Arc::new(InstrumentedNlpModel {
581                        inner: model.clone(),
582                        alias: alias.to_string(),
583                        provider_id: spec.provider_id.clone(),
584                        timeout: spec.timeout.map(std::time::Duration::from_secs),
585                        retry: spec.retry.clone(),
586                    });
587                    wrapper
588                })
589                .clone();
590            return Ok(cached);
591        }
592        Err(RuntimeError::ProviderCapabilityMissing {
593            alias: alias.to_string(),
594            provider_id: spec.provider_id,
595            capability: "NlpModel".to_string(),
596        })
597    }
598
599    /// Resolve, load (if necessary), and return an instrumented
600    /// [`DocumentExtractionModel`] handle for the given alias.
601    ///
602    /// # Examples
603    /// ```no_run
604    /// # use std::sync::Arc;
605    /// # use uni_xervo::runtime::ModelRuntime;
606    /// # use uni_xervo::traits::{DocExtractOptions, DocOutputFormat, ImageInput};
607    /// # async fn run(runtime: Arc<ModelRuntime>) -> uni_xervo::error::Result<()> {
608    /// let extractor = runtime.document_extractor("docext/olmocr").await?;
609    /// let page = ImageInput::Bytes {
610    ///     data: std::fs::read("page.png").unwrap(),
611    ///     media_type: "image/png".to_string(),
612    /// };
613    /// let options = DocExtractOptions {
614    ///     output: DocOutputFormat::Markdown,
615    ///     include_tables: true,
616    ///     include_formulas: true,
617    ///     include_bboxes: false,
618    /// };
619    /// let pages = extractor.extract(vec![page], options).await?;
620    /// println!("{}", pages[0].plain_markdown);
621    /// # Ok(())
622    /// # }
623    /// ```
624    ///
625    /// # Errors
626    /// Returns an error if the alias is unknown, the model fails to load, or
627    /// the provider does not implement document extraction.
628    pub async fn document_extractor(
629        &self,
630        alias: &str,
631    ) -> Result<Arc<dyn DocumentExtractionModel>> {
632        if let Some(cached) = self.handle_cache.document_extractors.get(alias) {
633            return Ok(cached.clone());
634        }
635        let spec = self.lookup_spec(alias).await?;
636        let handle = self.resolve_and_load_internal(&spec).await?;
637        if let Some(model) = handle.downcast_ref::<Arc<dyn DocumentExtractionModel>>() {
638            let cached = self
639                .handle_cache
640                .document_extractors
641                .entry(alias.to_string())
642                .or_insert_with(|| {
643                    let wrapper: Arc<dyn DocumentExtractionModel> =
644                        Arc::new(InstrumentedDocumentExtractionModel {
645                            inner: model.clone(),
646                            alias: alias.to_string(),
647                            provider_id: spec.provider_id.clone(),
648                            timeout: spec.timeout.map(std::time::Duration::from_secs),
649                            retry: spec.retry.clone(),
650                        });
651                    wrapper
652                })
653                .clone();
654            return Ok(cached);
655        }
656        Err(RuntimeError::ProviderCapabilityMissing {
657            alias: alias.to_string(),
658            provider_id: spec.provider_id,
659            capability: "DocumentExtractionModel".to_string(),
660        })
661    }
662
663    /// Resolve, load (if necessary), and return an instrumented
664    /// [`TranscriptionModel`] handle for the given alias.
665    pub async fn transcriber(&self, alias: &str) -> Result<Arc<dyn TranscriptionModel>> {
666        if let Some(cached) = self.handle_cache.transcribers.get(alias) {
667            return Ok(cached.clone());
668        }
669        let spec = self.lookup_spec(alias).await?;
670        let handle = self.resolve_and_load_internal(&spec).await?;
671        if let Some(model) = handle.downcast_ref::<Arc<dyn TranscriptionModel>>() {
672            let cached = self
673                .handle_cache
674                .transcribers
675                .entry(alias.to_string())
676                .or_insert_with(|| {
677                    let wrapper: Arc<dyn TranscriptionModel> =
678                        Arc::new(InstrumentedTranscriptionModel {
679                            inner: model.clone(),
680                            alias: alias.to_string(),
681                            provider_id: spec.provider_id.clone(),
682                            timeout: spec.timeout.map(std::time::Duration::from_secs),
683                            retry: spec.retry.clone(),
684                        });
685                    wrapper
686                })
687                .clone();
688            return Ok(cached);
689        }
690        Err(RuntimeError::ProviderCapabilityMissing {
691            alias: alias.to_string(),
692            provider_id: spec.provider_id,
693            capability: "TranscriptionModel".to_string(),
694        })
695    }
696
697    /// Resolve, load (if necessary), and return an instrumented [`OcrModel`]
698    /// handle for the given alias.
699    ///
700    /// # Examples
701    /// ```no_run
702    /// # use std::sync::Arc;
703    /// # use uni_xervo::runtime::ModelRuntime;
704    /// # use uni_xervo::traits::ImageInput;
705    /// # async fn run(runtime: Arc<ModelRuntime>) -> uni_xervo::error::Result<()> {
706    /// let ocr = runtime.ocr_model("ocr/ppocr-en").await?;
707    /// let image = ImageInput::Bytes {
708    ///     data: std::fs::read("scan.png").unwrap(),
709    ///     media_type: "image/png".to_string(),
710    /// };
711    /// let results = ocr.recognize(vec![image]).await?;
712    /// println!("{}", results[0].plain_text);
713    /// # Ok(())
714    /// # }
715    /// ```
716    ///
717    /// # Errors
718    /// Returns an error if the alias is unknown, the model fails to load, or
719    /// the provider does not implement OCR.
720    pub async fn ocr_model(&self, alias: &str) -> Result<Arc<dyn OcrModel>> {
721        if let Some(cached) = self.handle_cache.ocr_models.get(alias) {
722            return Ok(cached.clone());
723        }
724        let spec = self.lookup_spec(alias).await?;
725        let handle = self.resolve_and_load_internal(&spec).await?;
726        if let Some(model) = handle.downcast_ref::<Arc<dyn OcrModel>>() {
727            let cached = self
728                .handle_cache
729                .ocr_models
730                .entry(alias.to_string())
731                .or_insert_with(|| {
732                    let wrapper: Arc<dyn OcrModel> = Arc::new(InstrumentedOcrModel {
733                        inner: model.clone(),
734                        alias: alias.to_string(),
735                        provider_id: spec.provider_id.clone(),
736                        timeout: spec.timeout.map(std::time::Duration::from_secs),
737                        retry: spec.retry.clone(),
738                    });
739                    wrapper
740                })
741                .clone();
742            return Ok(cached);
743        }
744        Err(RuntimeError::ProviderCapabilityMissing {
745            alias: alias.to_string(),
746            provider_id: spec.provider_id,
747            capability: "OcrModel".to_string(),
748        })
749    }
750
751    #[tracing::instrument(skip(self, spec), fields(provider, model))]
752    async fn resolve_and_load_internal(
753        &self,
754        spec: &ModelAliasSpec,
755    ) -> Result<Arc<dyn Any + Send + Sync>> {
756        let key = ModelRuntimeKey::new(spec);
757
758        // Fast path: already loaded
759        {
760            let registry = self.registry.instances.read().await;
761            if let Some(handle) = registry.get(&key) {
762                return Ok(handle.clone());
763            }
764        }
765
766        // Slow path: coordinate loading
767        let lock = {
768            let mut locks = self.registry.loader_locks.lock().await;
769            locks
770                .entry(key.clone())
771                .or_insert_with(|| Arc::new(Mutex::new(())))
772                .clone()
773        };
774
775        // Acquire loader lock for this key
776        let _guard = lock.lock().await;
777
778        // Double-check after acquiring the loader lock
779        {
780            let registry = self.registry.instances.read().await;
781            if let Some(handle) = registry.get(&key) {
782                let result = Ok(handle.clone());
783                let mut locks = self.registry.loader_locks.lock().await;
784                locks.remove(&key);
785                return result;
786            }
787        }
788
789        let load_timeout =
790            std::time::Duration::from_secs(spec.load_timeout.unwrap_or(DEFAULT_LOAD_TIMEOUT_SECS));
791
792        let result = match tokio::time::timeout(load_timeout, async {
793            let provider = self.providers.get(&spec.provider_id).ok_or_else(|| {
794                RuntimeError::ProviderNotFound(format!("Provider '{}' not found", spec.provider_id))
795            })?;
796
797            tracing::info!(alias = %spec.alias, provider = %spec.provider_id, "Loading model instance");
798            let start = std::time::Instant::now();
799            let handle_result = provider.load(spec).await;
800            let duration = start.elapsed().as_secs_f64();
801
802            metrics::histogram!("model_load.duration_seconds").record(duration);
803
804            let handle = match handle_result {
805                Ok(h) => {
806                    metrics::counter!("model_load.total", "status" => "success").increment(1);
807                    h
808                }
809                Err(e) => {
810                    metrics::counter!("model_load.total", "status" => "failure").increment(1);
811                    tracing::error!(alias = %spec.alias, error = %e, "Model load failed");
812                    return Err(e);
813                }
814            };
815
816            // Model warmup. Dispatch covers every task trait — including the
817            // multimodal surface, which previously fell through to no warmup at
818            // all. The trait default is a no-op, so adding arms only enables
819            // eager weight loading for providers that override `warmup`.
820            if let Some(model) = handle.downcast_ref::<Arc<dyn EmbeddingModel>>() {
821                model.warmup().await?;
822            } else if let Some(model) = handle.downcast_ref::<Arc<dyn RerankerModel>>() {
823                model.warmup().await?;
824            } else if let Some(model) = handle.downcast_ref::<Arc<dyn GeneratorModel>>() {
825                model.warmup().await?;
826            } else if let Some(model) = handle.downcast_ref::<Arc<dyn RawTensorModel>>() {
827                model.warmup().await?;
828            } else if let Some(model) = handle.downcast_ref::<Arc<dyn ImageEmbeddingModel>>() {
829                model.warmup().await?;
830            } else if let Some(model) = handle.downcast_ref::<Arc<dyn AudioEmbeddingModel>>() {
831                model.warmup().await?;
832            } else if let Some(model) = handle.downcast_ref::<Arc<dyn MultimodalEmbeddingModel>>() {
833                model.warmup().await?;
834            } else if let Some(model) = handle.downcast_ref::<Arc<dyn SparseEmbeddingModel>>() {
835                model.warmup().await?;
836            } else if let Some(model) = handle.downcast_ref::<Arc<dyn MultiVectorEmbeddingModel>>() {
837                model.warmup().await?;
838            } else if let Some(model) = handle.downcast_ref::<Arc<dyn HybridEmbeddingModel>>() {
839                model.warmup().await?;
840            } else if let Some(model) = handle.downcast_ref::<Arc<dyn NlpModel>>() {
841                model.warmup().await?;
842            } else if let Some(model) = handle.downcast_ref::<Arc<dyn DocumentExtractionModel>>() {
843                model.warmup().await?;
844            } else if let Some(model) = handle.downcast_ref::<Arc<dyn TranscriptionModel>>() {
845                model.warmup().await?;
846            } else if let Some(model) = handle.downcast_ref::<Arc<dyn OcrModel>>() {
847                model.warmup().await?;
848            }
849
850            {
851                let mut registry = self.registry.instances.write().await;
852                registry.insert(key.clone(), handle.clone());
853            }
854
855            Ok(handle)
856        })
857        .await
858        {
859            Ok(res) => res,
860            Err(_) => {
861                metrics::counter!("model_load.total", "status" => "failure").increment(1);
862                tracing::error!(
863                    alias = %spec.alias,
864                    provider = %spec.provider_id,
865                    timeout_secs = load_timeout.as_secs(),
866                    "Model load timed out"
867                );
868                Err(RuntimeError::Timeout)
869            }
870        };
871
872        // Bound loader lock map growth by removing this key once the load path completes.
873        // Existing waiters hold cloned lock Arcs, so this is safe.
874        {
875            let mut locks = self.registry.loader_locks.lock().await;
876            locks.remove(&key);
877        }
878
879        result
880    }
881}
882
883/// Builder for constructing a [`ModelRuntime`] with registered providers,
884/// a model catalog, and a warmup policy.
885///
886/// ```rust,no_run
887/// # use uni_xervo::runtime::ModelRuntime;
888/// # async fn example() -> Result<(), Box<dyn std::error::Error>> {
889/// let runtime = ModelRuntime::builder()
890///     // .register_provider(...)
891///     // .catalog(...)
892///     .build()
893///     .await?;
894/// # Ok(())
895/// # }
896/// ```
897#[derive(Default)]
898pub struct ModelRuntimeBuilder {
899    providers: HashMap<String, Box<dyn ModelProvider>>,
900    catalog: Vec<ModelAliasSpec>,
901    warmup_policy: crate::api::WarmupPolicy,
902}
903
904impl ModelRuntimeBuilder {
905    /// Register a provider. The provider's
906    /// [`provider_id`](crate::traits::ModelProvider::provider_id) is used as
907    /// the lookup key; registering a second provider with the same ID
908    /// replaces the first.
909    pub fn register_provider<P: ModelProvider + 'static>(mut self, provider: P) -> Self {
910        self.providers
911            .insert(provider.provider_id().to_string(), Box::new(provider));
912        self
913    }
914
915    /// Set the model catalog from a pre-built vector of specs.
916    pub fn catalog(mut self, catalog: Vec<ModelAliasSpec>) -> Self {
917        self.catalog = catalog;
918        self
919    }
920
921    /// Load catalog from a JSON string (array of model alias specs).
922    pub fn catalog_from_str(mut self, s: &str) -> Result<Self> {
923        self.catalog = crate::api::catalog_from_str(s)?;
924        Ok(self)
925    }
926
927    /// Load catalog from a JSON file (array of model alias specs).
928    pub fn catalog_from_file(mut self, path: impl AsRef<std::path::Path>) -> Result<Self> {
929        self.catalog = crate::api::catalog_from_file(path)?;
930        Ok(self)
931    }
932
933    /// Set the global warmup policy applied to providers during
934    /// [`build`](Self::build).
935    pub fn warmup_policy(mut self, policy: crate::api::WarmupPolicy) -> Self {
936        self.warmup_policy = policy;
937        self
938    }
939
940    /// Validate the catalog, execute the warmup policy, and return the
941    /// constructed [`ModelRuntime`].
942    ///
943    /// Returns an error if any spec references an unknown provider, contains
944    /// invalid options, or if a required eager warmup fails.
945    pub async fn build(self) -> Result<Arc<ModelRuntime>> {
946        let mut catalog_map = HashMap::new();
947        for spec in self.catalog {
948            spec.validate()?;
949            if !self.providers.contains_key(&spec.provider_id) {
950                return Err(RuntimeError::Config(format!(
951                    "Unknown provider '{}' for alias '{}'",
952                    spec.provider_id, spec.alias
953                )));
954            }
955            validate_provider_options(&spec.provider_id, spec.task, &spec.options)?;
956            if catalog_map.insert(spec.alias.clone(), spec).is_some() {
957                return Err(RuntimeError::Config(
958                    "Duplicate alias in catalog".to_string(),
959                ));
960            }
961        }
962
963        let runtime = Arc::new(ModelRuntime {
964            providers: self.providers,
965            registry: Arc::new(ModelRegistry::default()),
966            catalog: RwLock::new(catalog_map),
967            handle_cache: HandleCache::default(),
968        });
969
970        // Provider Warmup Phase
971        match self.warmup_policy {
972            crate::api::WarmupPolicy::Eager => {
973                for (id, provider) in &runtime.providers {
974                    tracing::info!(provider = %id, "Eagerly warming up provider");
975                    provider.warmup().await.map_err(|e| {
976                        RuntimeError::Load(format!("Failed to warmup provider {}: {}", id, e))
977                    })?;
978                }
979            }
980            crate::api::WarmupPolicy::Background => {
981                for id in runtime.providers.keys() {
982                    tracing::info!(provider = %id, "Scheduling background provider warmup");
983                    // We have the Arc<ModelRuntime> already.
984                    let rt = runtime.clone();
985                    let provider_id = id.clone();
986                    tokio::spawn(async move {
987                        if let Some(provider) = rt.providers.get(&provider_id)
988                            && let Err(e) = provider.warmup().await
989                        {
990                            tracing::error!(provider = %provider_id, error = %e, "Background provider warmup failed");
991                        }
992                    });
993                }
994            }
995            crate::api::WarmupPolicy::Lazy => {
996                tracing::debug!("Lazy provider warmup (no-op)");
997            }
998        }
999
1000        // Model Warmup Phase
1001        let mut warmup_tasks = Vec::new();
1002
1003        let specs: Vec<ModelAliasSpec> = {
1004            let catalog = runtime.catalog.read().await;
1005            catalog.values().cloned().collect()
1006        };
1007
1008        for spec in specs {
1009            match spec.warmup {
1010                crate::api::WarmupPolicy::Eager => {
1011                    tracing::info!(alias = %spec.alias, "Eagerly warming up model");
1012                    if let Err(e) = runtime.resolve_and_load_internal(&spec).await {
1013                        if spec.required {
1014                            return Err(e);
1015                        }
1016                        tracing::error!(
1017                            alias = %spec.alias,
1018                            provider = %spec.provider_id,
1019                            error = %e,
1020                            "Optional eager model warmup failed; continuing startup"
1021                        );
1022                    }
1023                }
1024                crate::api::WarmupPolicy::Background => {
1025                    tracing::info!(alias = %spec.alias, "Scheduling background warmup");
1026                    let rt = runtime.clone();
1027                    let spec_clone = spec.clone();
1028                    // Spawn background task
1029                    warmup_tasks.push(tokio::spawn(async move {
1030                        if let Err(e) = rt.resolve_and_load_internal(&spec_clone).await {
1031                            tracing::error!(alias = %spec_clone.alias, error = %e, "Background warmup failed");
1032                        }
1033                    }));
1034                }
1035                crate::api::WarmupPolicy::Lazy => {
1036                    tracing::debug!(alias = %spec.alias, "Lazy warmup (no-op)");
1037                }
1038            }
1039        }
1040
1041        // We don't await background tasks here, they run detached.
1042        // Eager tasks are already awaited.
1043
1044        Ok(runtime)
1045    }
1046}
1047
1048#[cfg(test)]
1049mod tests {
1050    use super::*;
1051    use crate::api::ModelTask;
1052    use crate::mock::{MockProvider, make_spec};
1053
1054    #[tokio::test]
1055    async fn loader_lock_entries_cleaned_after_successful_load() {
1056        let spec = make_spec("embed/test", ModelTask::Embed, "mock/embed", "test-model");
1057        let runtime = ModelRuntime::builder()
1058            .register_provider(MockProvider::embed_only())
1059            .catalog(vec![spec])
1060            .build()
1061            .await
1062            .unwrap();
1063
1064        let _ = runtime.embedding("embed/test").await.unwrap();
1065
1066        let locks = runtime.registry.loader_locks.lock().await;
1067        assert!(
1068            locks.is_empty(),
1069            "loader lock map should be empty after load"
1070        );
1071    }
1072
1073    #[tokio::test]
1074    async fn loader_lock_entries_cleaned_after_failed_load() {
1075        let mut spec = make_spec("embed/test", ModelTask::Embed, "mock/failing", "test-model");
1076        spec.warmup = crate::api::WarmupPolicy::Lazy;
1077        let runtime = ModelRuntime::builder()
1078            .register_provider(MockProvider::failing())
1079            .catalog(vec![spec])
1080            .build()
1081            .await
1082            .unwrap();
1083
1084        let err = runtime.embedding("embed/test").await;
1085        assert!(err.is_err());
1086
1087        let locks = runtime.registry.loader_locks.lock().await;
1088        assert!(
1089            locks.is_empty(),
1090            "loader lock map should be empty after failure"
1091        );
1092    }
1093
1094    #[tokio::test]
1095    async fn loader_lock_entries_cleaned_after_load_timeout() {
1096        let mut spec = make_spec("embed/test", ModelTask::Embed, "mock/embed", "test-model");
1097        spec.warmup = crate::api::WarmupPolicy::Lazy;
1098        spec.load_timeout = Some(1);
1099
1100        let runtime = ModelRuntime::builder()
1101            .register_provider(MockProvider::embed_only().with_load_delay(2_000))
1102            .catalog(vec![spec])
1103            .build()
1104            .await
1105            .unwrap();
1106
1107        let err = runtime.embedding("embed/test").await;
1108        assert!(matches!(err, Err(RuntimeError::Timeout)));
1109
1110        let locks = runtime.registry.loader_locks.lock().await;
1111        assert!(
1112            locks.is_empty(),
1113            "loader lock map should be empty after load timeout"
1114        );
1115    }
1116}
uni_xervo/runtime.rs

uni_xervo/
runtime.rs
