uni_xervo/
reliability.rs

1//! Reliability primitives: circuit breaker, instrumented model wrappers with
2//! timeout and retry support, and metrics emission.
3
4use crate::error::{Result, RuntimeError};
5use crate::traits::{
6    EmbeddingModel, GenerationOptions, GenerationResult, GeneratorModel, Message, RawTensorModel,
7    RerankerModel, ScoredDoc, TensorBatch, TensorSpec,
8};
9use async_trait::async_trait;
10use std::sync::{Arc, Mutex};
11use std::time::{Duration, Instant};
12
13/// Internal circuit breaker state machine.
14#[derive(Debug, Clone, Copy, PartialEq, Eq)]
15enum State {
16    Closed,
17    Open,
18    HalfOpen,
19}
20
21/// Tunable parameters for the circuit breaker.
22pub struct CircuitBreakerConfig {
23    /// Number of consecutive failures before the breaker opens.
24    pub failure_threshold: u32,
25    /// Seconds to wait in the open state before allowing a probe call.
26    pub open_wait_seconds: u64,
27}
28
29impl Default for CircuitBreakerConfig {
30    fn default() -> Self {
31        Self {
32            failure_threshold: 5,
33            open_wait_seconds: 10,
34        }
35    }
36}
37
38struct Inner {
39    state: State,
40    failures: u32,
41    last_failure: Option<Instant>,
42    config: CircuitBreakerConfig,
43    half_open_probe_in_flight: bool,
44}
45
46/// Thread-safe circuit breaker that tracks failures and short-circuits calls
47/// when a provider is unhealthy.
48///
49/// State transitions: **Closed** -> (failures >= threshold) -> **Open** ->
50/// (wait period elapsed) -> **HalfOpen** -> (probe succeeds) -> **Closed**
51/// (or probe fails -> back to **Open**).
52#[derive(Clone)]
53pub struct CircuitBreakerWrapper {
54    inner: Arc<Mutex<Inner>>,
55}
56
57impl CircuitBreakerWrapper {
58    /// Create a new circuit breaker with the given configuration.
59    pub fn new(config: CircuitBreakerConfig) -> Self {
60        Self {
61            inner: Arc::new(Mutex::new(Inner {
62                state: State::Closed,
63                failures: 0,
64                last_failure: None,
65                config,
66                half_open_probe_in_flight: false,
67            })),
68        }
69    }
70
71    /// Execute `f` through the circuit breaker.
72    ///
73    /// Returns [`RuntimeError::Unavailable`] immediately when the breaker is
74    /// open.  In the half-open state only a single probe call is allowed;
75    /// concurrent callers receive `Unavailable` until the probe completes.
76    pub async fn call<F, Fut, T>(&self, f: F) -> Result<T>
77    where
78        F: FnOnce() -> Fut,
79        Fut: std::future::Future<Output = Result<T>>,
80    {
81        let is_probe_call;
82
83        // 1. Check state
84        {
85            let mut inner = self.inner.lock().unwrap();
86            match inner.state {
87                State::Open => {
88                    if let Some(last) = inner.last_failure {
89                        if last.elapsed() >= Duration::from_secs(inner.config.open_wait_seconds) {
90                            inner.state = State::HalfOpen;
91                        } else {
92                            return Err(RuntimeError::Unavailable);
93                        }
94                    }
95                }
96                State::HalfOpen => {
97                    if inner.half_open_probe_in_flight {
98                        return Err(RuntimeError::Unavailable);
99                    }
100                }
101                State::Closed => {}
102            }
103            is_probe_call = inner.state == State::HalfOpen;
104            if is_probe_call {
105                inner.half_open_probe_in_flight = true;
106            }
107        }
108
109        // 2. Execute
110        let result = f().await;
111
112        // 3. Update state
113        let mut inner = self.inner.lock().unwrap();
114        match result {
115            Ok(val) => {
116                if is_probe_call {
117                    inner.state = State::Closed;
118                    inner.failures = 0;
119                    inner.half_open_probe_in_flight = false;
120                } else if inner.state == State::Closed {
121                    inner.failures = 0;
122                }
123                Ok(val)
124            }
125            Err(e) => {
126                if is_probe_call {
127                    inner.half_open_probe_in_flight = false;
128                }
129                inner.failures += 1;
130                inner.last_failure = Some(Instant::now());
131
132                if is_probe_call
133                    || (inner.state == State::Closed
134                        && inner.failures >= inner.config.failure_threshold)
135                {
136                    inner.state = State::Open;
137                }
138                Err(e)
139            }
140        }
141    }
142}
143
144/// Wrapper around an [`EmbeddingModel`] that adds per-call timeout enforcement,
145/// exponential-backoff retries for transient errors, and metrics emission
146/// (`model_inference.duration_seconds`, `model_inference.total`).
147pub struct InstrumentedEmbeddingModel {
148    pub inner: Arc<dyn EmbeddingModel>,
149    pub alias: String,
150    pub provider_id: String,
151    pub timeout: Option<Duration>,
152    pub retry: Option<crate::api::RetryConfig>,
153}
154
155#[async_trait]
156impl EmbeddingModel for InstrumentedEmbeddingModel {
157    async fn embed(&self, texts: &[&str]) -> Result<crate::traits::EmbedResult> {
158        let start = Instant::now();
159        let mut attempts = 0;
160        let max_attempts = self.retry.as_ref().map(|r| r.max_attempts).unwrap_or(1);
161
162        let res = loop {
163            attempts += 1;
164            let fut = self.inner.embed(texts);
165
166            let res = if let Some(timeout) = self.timeout {
167                match tokio::time::timeout(timeout, fut).await {
168                    Ok(r) => r,
169                    Err(_) => Err(RuntimeError::Timeout),
170                }
171            } else {
172                fut.await
173            };
174
175            match res {
176                Ok(val) => break Ok(val),
177                Err(e) if e.is_retryable() && attempts < max_attempts => {
178                    let backoff = self.retry.as_ref().unwrap().get_backoff(attempts);
179                    tracing::warn!(
180                        alias = %self.alias,
181                        attempt = attempts,
182                        backoff_ms = backoff.as_millis(),
183                        error = %e,
184                        "Retrying embedding call"
185                    );
186                    tokio::time::sleep(backoff).await;
187                    continue;
188                }
189                Err(e) => break Err(e),
190            }
191        };
192
193        let duration = start.elapsed();
194        let status = if res.is_ok() { "success" } else { "failure" };
195
196        metrics::histogram!(
197            "model_inference.duration_seconds",
198            "alias" => self.alias.clone(),
199            "task" => "embed",
200            "provider" => self.provider_id.clone()
201        )
202        .record(duration.as_secs_f64());
203
204        metrics::counter!(
205            "model_inference.total",
206            "alias" => self.alias.clone(),
207            "task" => "embed",
208            "provider" => self.provider_id.clone(),
209            "status" => status
210        )
211        .increment(1);
212
213        res
214    }
215
216    fn dimensions(&self) -> u32 {
217        self.inner.dimensions()
218    }
219
220    async fn warmup(&self) -> Result<()> {
221        self.inner.warmup().await
222    }
223}
224
225impl crate::traits::ModelInfo for InstrumentedEmbeddingModel {
226    fn model_id(&self) -> &str {
227        self.inner.model_id()
228    }
229    fn active_execution_providers(&self) -> Vec<String> {
230        self.inner.active_execution_providers()
231    }
232}
233
234/// Wrapper around a [`GeneratorModel`] that adds timeout, retry, and metrics.
235///
236/// See [`InstrumentedEmbeddingModel`] for details on the instrumentation behavior.
237pub struct InstrumentedGeneratorModel {
238    pub inner: Arc<dyn GeneratorModel>,
239    pub alias: String,
240    pub provider_id: String,
241    pub timeout: Option<Duration>,
242    pub retry: Option<crate::api::RetryConfig>,
243}
244
245#[async_trait]
246impl GeneratorModel for InstrumentedGeneratorModel {
247    async fn generate(
248        &self,
249        messages: &[Message],
250        options: GenerationOptions,
251    ) -> Result<GenerationResult> {
252        let start = Instant::now();
253        let mut attempts = 0;
254        let max_attempts = self.retry.as_ref().map(|r| r.max_attempts).unwrap_or(1);
255
256        let res = loop {
257            attempts += 1;
258            let fut = self.inner.generate(messages, options.clone());
259
260            let res = if let Some(timeout) = self.timeout {
261                match tokio::time::timeout(timeout, fut).await {
262                    Ok(r) => r,
263                    Err(_) => Err(RuntimeError::Timeout),
264                }
265            } else {
266                fut.await
267            };
268
269            match res {
270                Ok(val) => break Ok(val),
271                Err(e) if e.is_retryable() && attempts < max_attempts => {
272                    let backoff = self.retry.as_ref().unwrap().get_backoff(attempts);
273                    tracing::warn!(
274                        alias = %self.alias,
275                        attempt = attempts,
276                        backoff_ms = backoff.as_millis(),
277                        error = %e,
278                        "Retrying generation call"
279                    );
280                    tokio::time::sleep(backoff).await;
281                    continue;
282                }
283                Err(e) => break Err(e),
284            }
285        };
286
287        let duration = start.elapsed();
288        let status = if res.is_ok() { "success" } else { "failure" };
289
290        metrics::histogram!(
291            "model_inference.duration_seconds",
292            "alias" => self.alias.clone(),
293            "task" => "generate",
294            "provider" => self.provider_id.clone()
295        )
296        .record(duration.as_secs_f64());
297
298        metrics::counter!(
299            "model_inference.total",
300            "alias" => self.alias.clone(),
301            "task" => "generate",
302            "provider" => self.provider_id.clone(),
303            "status" => status
304        )
305        .increment(1);
306
307        res
308    }
309
310    async fn warmup(&self) -> Result<()> {
311        self.inner.warmup().await
312    }
313}
314
315impl crate::traits::ModelInfo for InstrumentedGeneratorModel {
316    fn model_id(&self) -> &str {
317        self.inner.model_id()
318    }
319    fn active_execution_providers(&self) -> Vec<String> {
320        self.inner.active_execution_providers()
321    }
322}
323
324/// Wrapper around an [`RawTensorModel`] that adds timeout, retry, and metrics.
325pub struct InstrumentedRawTensorModel {
326    pub inner: Arc<dyn RawTensorModel>,
327    pub alias: String,
328    pub provider_id: String,
329    pub timeout: Option<Duration>,
330    pub retry: Option<crate::api::RetryConfig>,
331}
332
333#[async_trait]
334impl RawTensorModel for InstrumentedRawTensorModel {
335    async fn run(&self, inputs: &TensorBatch) -> Result<TensorBatch> {
336        let start = Instant::now();
337        let mut attempts = 0;
338        let max_attempts = self.retry.as_ref().map(|r| r.max_attempts).unwrap_or(1);
339
340        let res = loop {
341            attempts += 1;
342            let fut = self.inner.run(inputs);
343
344            let res = if let Some(timeout) = self.timeout {
345                match tokio::time::timeout(timeout, fut).await {
346                    Ok(r) => r,
347                    Err(_) => Err(RuntimeError::Timeout),
348                }
349            } else {
350                fut.await
351            };
352
353            match res {
354                Ok(val) => break Ok(val),
355                Err(e) if e.is_retryable() && attempts < max_attempts => {
356                    let backoff = self.retry.as_ref().unwrap().get_backoff(attempts);
357                    tracing::warn!(
358                        alias = %self.alias,
359                        attempt = attempts,
360                        backoff_ms = backoff.as_millis(),
361                        error = %e,
362                        "Retrying ONNX run call"
363                    );
364                    tokio::time::sleep(backoff).await;
365                    continue;
366                }
367                Err(e) => break Err(rewrite_onnx_error_alias(e, &self.alias)),
368            }
369        };
370
371        record_onnx_metrics(
372            &self.alias,
373            &self.provider_id,
374            start.elapsed(),
375            if res.is_ok() { "success" } else { "failure" },
376        );
377        res
378    }
379
380    async fn run_batch(&self, inputs: &[TensorBatch]) -> Result<Vec<TensorBatch>> {
381        let start = Instant::now();
382        let mut attempts = 0;
383        let max_attempts = self.retry.as_ref().map(|r| r.max_attempts).unwrap_or(1);
384
385        let res = loop {
386            attempts += 1;
387            let fut = self.inner.run_batch(inputs);
388
389            let res = if let Some(timeout) = self.timeout {
390                match tokio::time::timeout(timeout, fut).await {
391                    Ok(r) => r,
392                    Err(_) => Err(RuntimeError::Timeout),
393                }
394            } else {
395                fut.await
396            };
397
398            match res {
399                Ok(val) => break Ok(val),
400                Err(e) if e.is_retryable() && attempts < max_attempts => {
401                    let backoff = self.retry.as_ref().unwrap().get_backoff(attempts);
402                    tracing::warn!(
403                        alias = %self.alias,
404                        attempt = attempts,
405                        backoff_ms = backoff.as_millis(),
406                        error = %e,
407                        "Retrying ONNX run_batch call"
408                    );
409                    tokio::time::sleep(backoff).await;
410                    continue;
411                }
412                Err(e) => break Err(rewrite_onnx_error_alias(e, &self.alias)),
413            }
414        };
415
416        record_onnx_metrics(
417            &self.alias,
418            &self.provider_id,
419            start.elapsed(),
420            if res.is_ok() { "success" } else { "failure" },
421        );
422        res
423    }
424
425    fn max_batch_size(&self) -> usize {
426        self.inner.max_batch_size()
427    }
428
429    fn input_signature(&self) -> &[TensorSpec] {
430        self.inner.input_signature()
431    }
432
433    fn output_signature(&self) -> &[TensorSpec] {
434        self.inner.output_signature()
435    }
436
437    async fn warmup(&self) -> Result<()> {
438        self.inner.warmup().await
439    }
440}
441
442impl crate::traits::ModelInfo for InstrumentedRawTensorModel {
443    fn model_id(&self) -> &str {
444        self.inner.model_id()
445    }
446    fn active_execution_providers(&self) -> Vec<String> {
447        self.inner.active_execution_providers()
448    }
449}
450
451fn rewrite_onnx_error_alias(error: RuntimeError, alias: &str) -> RuntimeError {
452    match error {
453        RuntimeError::OnnxModelNotFound { path, .. } => RuntimeError::OnnxModelNotFound {
454            alias: alias.to_string(),
455            path,
456        },
457        RuntimeError::OnnxArtifactSelectionFailure { cause, .. } => {
458            RuntimeError::OnnxArtifactSelectionFailure {
459                alias: alias.to_string(),
460                cause,
461            }
462        }
463        RuntimeError::OnnxDownloadFailure { cause, .. } => RuntimeError::OnnxDownloadFailure {
464            alias: alias.to_string(),
465            cause,
466        },
467        RuntimeError::OnnxLoadFailure { path, cause, .. } => RuntimeError::OnnxLoadFailure {
468            alias: alias.to_string(),
469            path,
470            cause,
471        },
472        RuntimeError::OnnxSignatureIntrospectionFailure { cause, .. } => {
473            RuntimeError::OnnxSignatureIntrospectionFailure {
474                alias: alias.to_string(),
475                cause,
476            }
477        }
478        RuntimeError::OnnxInputMissing { required_input, .. } => RuntimeError::OnnxInputMissing {
479            alias: alias.to_string(),
480            required_input,
481        },
482        RuntimeError::OnnxInputTypeMismatch {
483            input_name,
484            expected,
485            got,
486            ..
487        } => RuntimeError::OnnxInputTypeMismatch {
488            alias: alias.to_string(),
489            input_name,
490            expected,
491            got,
492        },
493        RuntimeError::OnnxInputShapeMismatch {
494            input_name,
495            expected,
496            got,
497            ..
498        } => RuntimeError::OnnxInputShapeMismatch {
499            alias: alias.to_string(),
500            input_name,
501            expected,
502            got,
503        },
504        RuntimeError::OnnxInvocationFailure { cause, .. } => RuntimeError::OnnxInvocationFailure {
505            alias: alias.to_string(),
506            cause,
507        },
508        RuntimeError::OnnxBatchStackingFailure { cause, .. } => {
509            RuntimeError::OnnxBatchStackingFailure {
510                alias: alias.to_string(),
511                cause,
512            }
513        }
514        other => other,
515    }
516}
517
518fn record_onnx_metrics(alias: &str, provider_id: &str, duration: Duration, status: &str) {
519    metrics::histogram!(
520        "model_inference.duration_seconds",
521        "alias" => alias.to_string(),
522        "task" => "raw",
523        "provider" => provider_id.to_string()
524    )
525    .record(duration.as_secs_f64());
526
527    metrics::counter!(
528        "model_inference.total",
529        "alias" => alias.to_string(),
530        "task" => "raw",
531        "provider" => provider_id.to_string(),
532        "status" => status.to_string()
533    )
534    .increment(1);
535}
536
537/// Wrapper around a [`RerankerModel`] that adds timeout, retry, and metrics.
538///
539/// See [`InstrumentedEmbeddingModel`] for details on the instrumentation behavior.
540pub struct InstrumentedRerankerModel {
541    pub inner: Arc<dyn RerankerModel>,
542    pub alias: String,
543    pub provider_id: String,
544    pub timeout: Option<Duration>,
545    pub retry: Option<crate::api::RetryConfig>,
546}
547
548#[async_trait]
549impl RerankerModel for InstrumentedRerankerModel {
550    async fn rerank(&self, query: &str, docs: &[&str]) -> Result<Vec<ScoredDoc>> {
551        let start = Instant::now();
552        let mut attempts = 0;
553        let max_attempts = self.retry.as_ref().map(|r| r.max_attempts).unwrap_or(1);
554
555        let res = loop {
556            attempts += 1;
557            let fut = self.inner.rerank(query, docs);
558
559            let res = if let Some(timeout) = self.timeout {
560                match tokio::time::timeout(timeout, fut).await {
561                    Ok(r) => r,
562                    Err(_) => Err(RuntimeError::Timeout),
563                }
564            } else {
565                fut.await
566            };
567
568            match res {
569                Ok(val) => break Ok(val),
570                Err(e) if e.is_retryable() && attempts < max_attempts => {
571                    let backoff = self.retry.as_ref().unwrap().get_backoff(attempts);
572                    tracing::warn!(
573                        alias = %self.alias,
574                        attempt = attempts,
575                        backoff_ms = backoff.as_millis(),
576                        error = %e,
577                        "Retrying rerank call"
578                    );
579                    tokio::time::sleep(backoff).await;
580                    continue;
581                }
582                Err(e) => break Err(e),
583            }
584        };
585
586        let duration = start.elapsed();
587        let status = if res.is_ok() { "success" } else { "failure" };
588
589        metrics::histogram!(
590            "model_inference.duration_seconds",
591            "alias" => self.alias.clone(),
592            "task" => "rerank",
593            "provider" => self.provider_id.clone()
594        )
595        .record(duration.as_secs_f64());
596
597        metrics::counter!(
598            "model_inference.total",
599            "alias" => self.alias.clone(),
600            "task" => "rerank",
601            "provider" => self.provider_id.clone(),
602            "status" => status
603        )
604        .increment(1);
605
606        res
607    }
608
609    async fn warmup(&self) -> Result<()> {
610        self.inner.warmup().await
611    }
612}
613
614impl crate::traits::ModelInfo for InstrumentedRerankerModel {
615    fn model_id(&self) -> &str {
616        self.inner.model_id()
617    }
618    fn active_execution_providers(&self) -> Vec<String> {
619        self.inner.active_execution_providers()
620    }
621}
622
623// ---------------------------------------------------------------------------
624// Multimodal instrumentation wrappers.
625//
626// Each wraps a future-producing call (the trait's business method) with the
627// same timeout-then-retry-then-metrics pattern used by
628// `InstrumentedEmbeddingModel` and `InstrumentedGeneratorModel`. The pattern
629// is centralized into the helper `run_instrumented` below so the seven new
630// wrappers stay short. Existing wrappers keep their inlined form (could be
631// migrated to this helper in a follow-up).
632// ---------------------------------------------------------------------------
633
634/// Run an async operation with the standard timeout / retry / metrics
635/// envelope used by every instrumented model wrapper.
636///
637/// `make_fut` is a closure invoked once per attempt — it must produce a fresh
638/// future each call so retries operate on a fresh inner call rather than
639/// re-polling a completed future. `task` is the wire-format task name used
640/// for the metric labels (`embed_image`, `nlp`, etc.).
641async fn run_instrumented<F, Fut, T>(
642    alias: &str,
643    provider_id: &str,
644    task: &'static str,
645    timeout: Option<Duration>,
646    retry: Option<&crate::api::RetryConfig>,
647    mut make_fut: F,
648) -> Result<T>
649where
650    F: FnMut() -> Fut,
651    Fut: std::future::Future<Output = Result<T>>,
652{
653    let start = Instant::now();
654    let mut attempts = 0u32;
655    let max_attempts = retry.as_ref().map(|r| r.max_attempts).unwrap_or(1);
656
657    let res = loop {
658        attempts += 1;
659        let fut = make_fut();
660        let attempt_res = if let Some(t) = timeout {
661            match tokio::time::timeout(t, fut).await {
662                Ok(r) => r,
663                Err(_) => Err(RuntimeError::Timeout),
664            }
665        } else {
666            fut.await
667        };
668
669        match attempt_res {
670            Ok(val) => break Ok(val),
671            Err(e) if e.is_retryable() && attempts < max_attempts => {
672                let backoff = retry.as_ref().unwrap().get_backoff(attempts);
673                tracing::warn!(
674                    alias = %alias,
675                    attempt = attempts,
676                    backoff_ms = backoff.as_millis(),
677                    error = %e,
678                    "Retrying {task} call",
679                );
680                tokio::time::sleep(backoff).await;
681                continue;
682            }
683            Err(e) => break Err(e),
684        }
685    };
686
687    let duration = start.elapsed();
688    let status = if res.is_ok() { "success" } else { "failure" };
689    metrics::histogram!(
690        "model_inference.duration_seconds",
691        "alias" => alias.to_string(),
692        "task" => task,
693        "provider" => provider_id.to_string()
694    )
695    .record(duration.as_secs_f64());
696    metrics::counter!(
697        "model_inference.total",
698        "alias" => alias.to_string(),
699        "task" => task,
700        "provider" => provider_id.to_string(),
701        "status" => status
702    )
703    .increment(1);
704
705    res
706}
707
708/// Wrapper around an [`ImageEmbeddingModel`](crate::traits::ImageEmbeddingModel)
709/// adding timeout / retry / metrics.
710pub struct InstrumentedImageEmbeddingModel {
711    pub inner: Arc<dyn crate::traits::ImageEmbeddingModel>,
712    pub alias: String,
713    pub provider_id: String,
714    pub timeout: Option<Duration>,
715    pub retry: Option<crate::api::RetryConfig>,
716}
717
718#[async_trait]
719impl crate::traits::ImageEmbeddingModel for InstrumentedImageEmbeddingModel {
720    async fn embed(
721        &self,
722        images: Vec<crate::traits::ImageInput>,
723    ) -> Result<crate::traits::EmbedResult> {
724        run_instrumented(
725            &self.alias,
726            &self.provider_id,
727            "embed_image",
728            self.timeout,
729            self.retry.as_ref(),
730            || self.inner.embed(images.clone()),
731        )
732        .await
733    }
734
735    fn dimensions(&self) -> u32 {
736        self.inner.dimensions()
737    }
738
739    async fn warmup(&self) -> Result<()> {
740        self.inner.warmup().await
741    }
742}
743
744impl crate::traits::ModelInfo for InstrumentedImageEmbeddingModel {
745    fn model_id(&self) -> &str {
746        self.inner.model_id()
747    }
748    fn active_execution_providers(&self) -> Vec<String> {
749        self.inner.active_execution_providers()
750    }
751}
752
753/// Wrapper around an [`AudioEmbeddingModel`](crate::traits::AudioEmbeddingModel)
754/// adding timeout / retry / metrics.
755pub struct InstrumentedAudioEmbeddingModel {
756    pub inner: Arc<dyn crate::traits::AudioEmbeddingModel>,
757    pub alias: String,
758    pub provider_id: String,
759    pub timeout: Option<Duration>,
760    pub retry: Option<crate::api::RetryConfig>,
761}
762
763#[async_trait]
764impl crate::traits::AudioEmbeddingModel for InstrumentedAudioEmbeddingModel {
765    async fn embed(
766        &self,
767        audios: Vec<crate::traits::AudioInput>,
768    ) -> Result<crate::traits::EmbedResult> {
769        run_instrumented(
770            &self.alias,
771            &self.provider_id,
772            "embed_audio",
773            self.timeout,
774            self.retry.as_ref(),
775            || self.inner.embed(audios.clone()),
776        )
777        .await
778    }
779
780    fn dimensions(&self) -> u32 {
781        self.inner.dimensions()
782    }
783
784    async fn warmup(&self) -> Result<()> {
785        self.inner.warmup().await
786    }
787}
788
789impl crate::traits::ModelInfo for InstrumentedAudioEmbeddingModel {
790    fn model_id(&self) -> &str {
791        self.inner.model_id()
792    }
793    fn active_execution_providers(&self) -> Vec<String> {
794        self.inner.active_execution_providers()
795    }
796}
797
798/// Wrapper around a [`MultimodalEmbeddingModel`](crate::traits::MultimodalEmbeddingModel)
799/// adding timeout / retry / metrics.
800pub struct InstrumentedMultimodalEmbeddingModel {
801    pub inner: Arc<dyn crate::traits::MultimodalEmbeddingModel>,
802    pub alias: String,
803    pub provider_id: String,
804    pub timeout: Option<Duration>,
805    pub retry: Option<crate::api::RetryConfig>,
806}
807
808#[async_trait]
809impl crate::traits::MultimodalEmbeddingModel for InstrumentedMultimodalEmbeddingModel {
810    async fn embed(
811        &self,
812        inputs: Vec<crate::traits::MultimodalInput>,
813    ) -> Result<crate::traits::EmbedResult> {
814        run_instrumented(
815            &self.alias,
816            &self.provider_id,
817            "embed_multimodal",
818            self.timeout,
819            self.retry.as_ref(),
820            || self.inner.embed(inputs.clone()),
821        )
822        .await
823    }
824
825    fn dimensions(&self) -> u32 {
826        self.inner.dimensions()
827    }
828
829    fn supported_modalities(&self) -> &[crate::traits::Modality] {
830        self.inner.supported_modalities()
831    }
832
833    async fn warmup(&self) -> Result<()> {
834        self.inner.warmup().await
835    }
836}
837
838impl crate::traits::ModelInfo for InstrumentedMultimodalEmbeddingModel {
839    fn model_id(&self) -> &str {
840        self.inner.model_id()
841    }
842    fn active_execution_providers(&self) -> Vec<String> {
843        self.inner.active_execution_providers()
844    }
845}
846
847/// Wrapper around a [`SparseEmbeddingModel`](crate::traits::SparseEmbeddingModel)
848/// adding timeout / retry / metrics.
849pub struct InstrumentedSparseEmbeddingModel {
850    pub inner: Arc<dyn crate::traits::SparseEmbeddingModel>,
851    pub alias: String,
852    pub provider_id: String,
853    pub timeout: Option<Duration>,
854    pub retry: Option<crate::api::RetryConfig>,
855}
856
857#[async_trait]
858impl crate::traits::SparseEmbeddingModel for InstrumentedSparseEmbeddingModel {
859    async fn embed(&self, texts: &[&str]) -> Result<crate::traits::SparseEmbedResult> {
860        run_instrumented(
861            &self.alias,
862            &self.provider_id,
863            "embed_sparse",
864            self.timeout,
865            self.retry.as_ref(),
866            || self.inner.embed(texts),
867        )
868        .await
869    }
870
871    fn vocab_size(&self) -> u32 {
872        self.inner.vocab_size()
873    }
874
875    async fn warmup(&self) -> Result<()> {
876        self.inner.warmup().await
877    }
878}
879
880impl crate::traits::ModelInfo for InstrumentedSparseEmbeddingModel {
881    fn model_id(&self) -> &str {
882        self.inner.model_id()
883    }
884    fn active_execution_providers(&self) -> Vec<String> {
885        self.inner.active_execution_providers()
886    }
887}
888
889/// Wrapper around a [`MultiVectorEmbeddingModel`](crate::traits::MultiVectorEmbeddingModel)
890/// adding timeout / retry / metrics.
891pub struct InstrumentedMultiVectorEmbeddingModel {
892    pub inner: Arc<dyn crate::traits::MultiVectorEmbeddingModel>,
893    pub alias: String,
894    pub provider_id: String,
895    pub timeout: Option<Duration>,
896    pub retry: Option<crate::api::RetryConfig>,
897}
898
899#[async_trait]
900impl crate::traits::MultiVectorEmbeddingModel for InstrumentedMultiVectorEmbeddingModel {
901    async fn embed(&self, texts: &[&str]) -> Result<crate::traits::MultiVectorEmbedResult> {
902        run_instrumented(
903            &self.alias,
904            &self.provider_id,
905            "embed_multi_vector",
906            self.timeout,
907            self.retry.as_ref(),
908            || self.inner.embed(texts),
909        )
910        .await
911    }
912
913    fn dimensions(&self) -> u32 {
914        self.inner.dimensions()
915    }
916
917    async fn warmup(&self) -> Result<()> {
918        self.inner.warmup().await
919    }
920}
921
922impl crate::traits::ModelInfo for InstrumentedMultiVectorEmbeddingModel {
923    fn model_id(&self) -> &str {
924        self.inner.model_id()
925    }
926    fn active_execution_providers(&self) -> Vec<String> {
927        self.inner.active_execution_providers()
928    }
929}
930
931/// Wrapper around a [`HybridEmbeddingModel`](crate::traits::HybridEmbeddingModel)
932/// adding timeout / retry / metrics.
933pub struct InstrumentedHybridEmbeddingModel {
934    pub inner: Arc<dyn crate::traits::HybridEmbeddingModel>,
935    pub alias: String,
936    pub provider_id: String,
937    pub timeout: Option<Duration>,
938    pub retry: Option<crate::api::RetryConfig>,
939}
940
941#[async_trait]
942impl crate::traits::HybridEmbeddingModel for InstrumentedHybridEmbeddingModel {
943    async fn embed(
944        &self,
945        texts: &[&str],
946        heads: crate::traits::HeadSet,
947    ) -> Result<crate::traits::HybridEmbedResult> {
948        run_instrumented(
949            &self.alias,
950            &self.provider_id,
951            "embed_hybrid",
952            self.timeout,
953            self.retry.as_ref(),
954            || self.inner.embed(texts, heads),
955        )
956        .await
957    }
958
959    fn available_heads(&self) -> crate::traits::HeadSet {
960        self.inner.available_heads()
961    }
962
963    async fn warmup(&self) -> Result<()> {
964        self.inner.warmup().await
965    }
966}
967
968impl crate::traits::ModelInfo for InstrumentedHybridEmbeddingModel {
969    fn model_id(&self) -> &str {
970        self.inner.model_id()
971    }
972    fn active_execution_providers(&self) -> Vec<String> {
973        self.inner.active_execution_providers()
974    }
975}
976
977/// Wrapper around an [`NlpModel`](crate::traits::NlpModel) adding timeout / retry / metrics.
978pub struct InstrumentedNlpModel {
979    pub inner: Arc<dyn crate::traits::NlpModel>,
980    pub alias: String,
981    pub provider_id: String,
982    pub timeout: Option<Duration>,
983    pub retry: Option<crate::api::RetryConfig>,
984}
985
986#[async_trait]
987impl crate::traits::NlpModel for InstrumentedNlpModel {
988    async fn analyze(
989        &self,
990        requests: Vec<crate::traits::NlpRequest<'_>>,
991    ) -> Result<Vec<crate::traits::NlpResult>> {
992        run_instrumented(
993            &self.alias,
994            &self.provider_id,
995            "nlp",
996            self.timeout,
997            self.retry.as_ref(),
998            || self.inner.analyze(requests.clone()),
999        )
1000        .await
1001    }
1002
1003    fn supported_tasks(&self) -> crate::traits::NlpTasks {
1004        self.inner.supported_tasks()
1005    }
1006
1007    fn label_maps(&self) -> Option<&crate::traits::NlpLabelMaps> {
1008        self.inner.label_maps()
1009    }
1010
1011    async fn warmup(&self) -> Result<()> {
1012        self.inner.warmup().await
1013    }
1014}
1015
1016impl crate::traits::ModelInfo for InstrumentedNlpModel {
1017    fn model_id(&self) -> &str {
1018        self.inner.model_id()
1019    }
1020    fn active_execution_providers(&self) -> Vec<String> {
1021        self.inner.active_execution_providers()
1022    }
1023}
1024
1025/// Wrapper around a [`DocumentExtractionModel`](crate::traits::DocumentExtractionModel)
1026/// adding timeout / retry / metrics.
1027pub struct InstrumentedDocumentExtractionModel {
1028    pub inner: Arc<dyn crate::traits::DocumentExtractionModel>,
1029    pub alias: String,
1030    pub provider_id: String,
1031    pub timeout: Option<Duration>,
1032    pub retry: Option<crate::api::RetryConfig>,
1033}
1034
1035#[async_trait]
1036impl crate::traits::DocumentExtractionModel for InstrumentedDocumentExtractionModel {
1037    async fn extract(
1038        &self,
1039        pages: Vec<crate::traits::ImageInput>,
1040        options: crate::traits::DocExtractOptions,
1041    ) -> Result<Vec<crate::traits::DocExtractResult>> {
1042        run_instrumented(
1043            &self.alias,
1044            &self.provider_id,
1045            "document_extract",
1046            self.timeout,
1047            self.retry.as_ref(),
1048            || self.inner.extract(pages.clone(), options.clone()),
1049        )
1050        .await
1051    }
1052
1053    async fn warmup(&self) -> Result<()> {
1054        self.inner.warmup().await
1055    }
1056}
1057
1058impl crate::traits::ModelInfo for InstrumentedDocumentExtractionModel {
1059    fn model_id(&self) -> &str {
1060        self.inner.model_id()
1061    }
1062    fn active_execution_providers(&self) -> Vec<String> {
1063        self.inner.active_execution_providers()
1064    }
1065}
1066
1067/// Wrapper around a [`TranscriptionModel`](crate::traits::TranscriptionModel)
1068/// adding timeout / retry / metrics to the batch `transcribe` call.
1069///
1070/// The timeout applies batch-wide — operators tuning `timeout` for batched
1071/// ingest need to budget for the full batch, not a per-item slice.
1072pub struct InstrumentedTranscriptionModel {
1073    pub inner: Arc<dyn crate::traits::TranscriptionModel>,
1074    pub alias: String,
1075    pub provider_id: String,
1076    pub timeout: Option<Duration>,
1077    pub retry: Option<crate::api::RetryConfig>,
1078}
1079
1080#[async_trait]
1081impl crate::traits::TranscriptionModel for InstrumentedTranscriptionModel {
1082    async fn transcribe(
1083        &self,
1084        audios: Vec<crate::traits::AudioInput>,
1085        options: crate::traits::TranscribeOptions,
1086    ) -> Result<Vec<crate::traits::TranscribeResult>> {
1087        run_instrumented(
1088            &self.alias,
1089            &self.provider_id,
1090            "transcribe",
1091            self.timeout,
1092            self.retry.as_ref(),
1093            || self.inner.transcribe(audios.clone(), options.clone()),
1094        )
1095        .await
1096    }
1097
1098    fn supported_languages(&self) -> &[String] {
1099        self.inner.supported_languages()
1100    }
1101
1102    async fn warmup(&self) -> Result<()> {
1103        self.inner.warmup().await
1104    }
1105}
1106
1107impl crate::traits::ModelInfo for InstrumentedTranscriptionModel {
1108    fn model_id(&self) -> &str {
1109        self.inner.model_id()
1110    }
1111    fn active_execution_providers(&self) -> Vec<String> {
1112        self.inner.active_execution_providers()
1113    }
1114}
1115
1116/// Wrapper around an [`OcrModel`](crate::traits::OcrModel) adding timeout / retry / metrics.
1117pub struct InstrumentedOcrModel {
1118    pub inner: Arc<dyn crate::traits::OcrModel>,
1119    pub alias: String,
1120    pub provider_id: String,
1121    pub timeout: Option<Duration>,
1122    pub retry: Option<crate::api::RetryConfig>,
1123}
1124
1125#[async_trait]
1126impl crate::traits::OcrModel for InstrumentedOcrModel {
1127    async fn recognize(
1128        &self,
1129        images: Vec<crate::traits::ImageInput>,
1130    ) -> Result<Vec<crate::traits::OcrResult>> {
1131        run_instrumented(
1132            &self.alias,
1133            &self.provider_id,
1134            "ocr",
1135            self.timeout,
1136            self.retry.as_ref(),
1137            || self.inner.recognize(images.clone()),
1138        )
1139        .await
1140    }
1141
1142    async fn warmup(&self) -> Result<()> {
1143        self.inner.warmup().await
1144    }
1145}
1146
1147impl crate::traits::ModelInfo for InstrumentedOcrModel {
1148    fn model_id(&self) -> &str {
1149        self.inner.model_id()
1150    }
1151    fn active_execution_providers(&self) -> Vec<String> {
1152        self.inner.active_execution_providers()
1153    }
1154}
1155
1156#[cfg(test)]
1157mod tests {
1158    use super::*;
1159    use crate::traits::{ModelInfo, NlpModel};
1160    use std::sync::atomic::{AtomicU32, Ordering};
1161
1162    #[tokio::test]
1163    async fn test_circuit_breaker_transitions() {
1164        let config = CircuitBreakerConfig {
1165            failure_threshold: 2,
1166            open_wait_seconds: 1,
1167        };
1168        let cb = CircuitBreakerWrapper::new(config);
1169        let counter = Arc::new(AtomicU32::new(0));
1170
1171        // 1. Success calls - state remains Closed
1172        let res = cb.call(|| async { Ok::<_, RuntimeError>(()) }).await;
1173        assert!(res.is_ok());
1174
1175        // 2. Failures - state transitions to Open
1176        let res = cb
1177            .call(|| async { Err::<(), _>(RuntimeError::InferenceError("fail".into())) })
1178            .await;
1179        assert!(res.is_err()); // Fail 1
1180
1181        let res = cb
1182            .call(|| async { Err::<(), _>(RuntimeError::InferenceError("fail".into())) })
1183            .await;
1184        assert!(res.is_err()); // Fail 2 -> Open
1185
1186        // 3. Open state - calls rejected immediately
1187        let res = cb
1188            .call(|| async {
1189                counter.fetch_add(1, Ordering::SeqCst);
1190                Ok(())
1191            })
1192            .await;
1193        assert!(res.is_err());
1194        assert_eq!(res.err().unwrap().to_string(), "Unavailable");
1195        assert_eq!(counter.load(Ordering::SeqCst), 0); // Should not have run
1196
1197        // 4. Wait for HalfOpen
1198        tokio::time::sleep(Duration::from_millis(1100)).await;
1199
1200        // 5. HalfOpen - allow one call
1201        // If it fails, go back to Open
1202        let res = cb
1203            .call(|| async { Err::<(), _>(RuntimeError::InferenceError("fail".into())) })
1204            .await;
1205        assert!(res.is_err());
1206
1207        // Should be Open again
1208        let res = cb.call(|| async { Ok(()) }).await;
1209        assert!(res.is_err());
1210        assert_eq!(res.err().unwrap().to_string(), "Unavailable");
1211
1212        // 6. Wait again for HalfOpen
1213        tokio::time::sleep(Duration::from_millis(1100)).await;
1214
1215        // 7. Success - transition to Closed
1216        let res = cb.call(|| async { Ok(()) }).await;
1217        assert!(res.is_ok());
1218
1219        // Should be closed now, next call works
1220        let res = cb.call(|| async { Ok(()) }).await;
1221        assert!(res.is_ok());
1222    }
1223
1224    #[tokio::test]
1225    async fn test_half_open_allows_single_probe() {
1226        let config = CircuitBreakerConfig {
1227            failure_threshold: 1,
1228            open_wait_seconds: 1,
1229        };
1230        let cb = CircuitBreakerWrapper::new(config);
1231
1232        // Open breaker.
1233        let _ = cb
1234            .call(|| async { Err::<(), _>(RuntimeError::InferenceError("fail".into())) })
1235            .await;
1236
1237        tokio::time::sleep(Duration::from_millis(1100)).await;
1238
1239        let started = Arc::new(std::sync::atomic::AtomicU32::new(0));
1240        let finished = Arc::new(std::sync::atomic::AtomicU32::new(0));
1241
1242        let cb_probe = cb.clone();
1243        let started_probe = started.clone();
1244        let finished_probe = finished.clone();
1245        let probe = tokio::spawn(async move {
1246            cb_probe
1247                .call(|| async move {
1248                    started_probe.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
1249                    tokio::time::sleep(Duration::from_millis(150)).await;
1250                    finished_probe.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
1251                    Ok::<_, RuntimeError>(())
1252                })
1253                .await
1254        });
1255
1256        // Allow the first probe to enter.
1257        tokio::time::sleep(Duration::from_millis(20)).await;
1258
1259        // A concurrent call during half-open probe should fail fast.
1260        let second = cb.call(|| async { Ok::<_, RuntimeError>(()) }).await;
1261        assert!(matches!(second, Err(RuntimeError::Unavailable)));
1262
1263        let probe_result = probe.await.unwrap();
1264        assert!(probe_result.is_ok());
1265        assert_eq!(started.load(std::sync::atomic::Ordering::SeqCst), 1);
1266        assert_eq!(finished.load(std::sync::atomic::Ordering::SeqCst), 1);
1267
1268        // Closed again.
1269        let res = cb.call(|| async { Ok::<_, RuntimeError>(()) }).await;
1270        assert!(res.is_ok());
1271    }
1272
1273    /// A minimal `NlpModel` whose `label_maps` returns a marker vocabulary, used
1274    /// to verify the instrumentation wrapper forwards introspection methods.
1275    struct LabelMapNlpModel {
1276        labels: crate::traits::NlpLabelMaps,
1277    }
1278
1279    #[async_trait]
1280    impl crate::traits::NlpModel for LabelMapNlpModel {
1281        async fn analyze(
1282            &self,
1283            _requests: Vec<crate::traits::NlpRequest<'_>>,
1284        ) -> Result<Vec<crate::traits::NlpResult>> {
1285            Ok(Vec::new())
1286        }
1287
1288        fn supported_tasks(&self) -> crate::traits::NlpTasks {
1289            crate::traits::NlpTasks::CLS
1290        }
1291
1292        fn label_maps(&self) -> Option<&crate::traits::NlpLabelMaps> {
1293            Some(&self.labels)
1294        }
1295    }
1296
1297    impl crate::traits::ModelInfo for LabelMapNlpModel {
1298        fn model_id(&self) -> &str {
1299            "mock/labelmaps"
1300        }
1301    }
1302
1303    /// The instrumentation wrapper forwards `label_maps()` to the inner model
1304    /// rather than falling back to the trait default `None` (R5 regression: the
1305    /// runtime wraps every model, so a non-forwarded method is invisible to
1306    /// callers).
1307    #[tokio::test]
1308    async fn instrumented_nlp_forwards_label_maps() {
1309        let labels = crate::traits::NlpLabelMaps {
1310            cls: vec!["statement".to_string(), "question".to_string()],
1311            ..Default::default()
1312        };
1313        let inner: Arc<dyn crate::traits::NlpModel> = Arc::new(LabelMapNlpModel { labels });
1314        let wrapped = InstrumentedNlpModel {
1315            inner,
1316            alias: "nlp/x".to_string(),
1317            provider_id: "test".to_string(),
1318            timeout: None,
1319            retry: None,
1320        };
1321
1322        let maps = wrapped.label_maps().expect("wrapper forwards label_maps");
1323        assert_eq!(maps.cls, ["statement", "question"]);
1324        assert_eq!(wrapped.model_id(), "mock/labelmaps");
1325        assert_eq!(wrapped.supported_tasks(), crate::traits::NlpTasks::CLS);
1326    }
1327
1328    /// A model without label maps reports `None` through the wrapper, confirming
1329    /// the default path is preserved (remote/mock providers stay `None`).
1330    #[tokio::test]
1331    async fn instrumented_nlp_label_maps_none_passes_through() {
1332        let inner: Arc<dyn crate::traits::NlpModel> = Arc::new(crate::mock::MockNlpModel::new());
1333        let wrapped = InstrumentedNlpModel {
1334            inner,
1335            alias: "nlp/y".to_string(),
1336            provider_id: "test".to_string(),
1337            timeout: None,
1338            retry: None,
1339        };
1340        assert!(wrapped.label_maps().is_none());
1341    }
1342}
uni_xervo/reliability.rs

uni_xervo/
reliability.rs
