mullama 0.3.0

Comprehensive Rust bindings for llama.cpp with memory-safe API and advanced features
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49

use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::Arc;

use tokio::sync::{RwLock as TokioRwLock, RwLockReadGuard, RwLockWriteGuard};

use crate::{Context, ContextParams, Model, MullamaError};

/// Round-robin context pool for a loaded model.
pub struct ContextPool {
    contexts: Vec<TokioRwLock<Context>>,
    next_context: AtomicUsize,
}

impl ContextPool {
    pub fn new(
        model: Arc<Model>,
        context: Context,
        ctx_params: ContextParams,
        context_pool_size: usize,
    ) -> Result<Self, MullamaError> {
        let context_pool_size = context_pool_size.max(1);
        let mut contexts = Vec::with_capacity(context_pool_size);
        contexts.push(TokioRwLock::new(context));

        for _ in 1..context_pool_size {
            let ctx = Context::new(model.clone(), ctx_params.clone())?;
            contexts.push(TokioRwLock::new(ctx));
        }

        Ok(Self {
            contexts,
            next_context: AtomicUsize::new(0),
        })
    }

    pub async fn acquire(&self) -> RwLockWriteGuard<'_, Context> {
        let idx = self.next_context.fetch_add(1, Ordering::Relaxed) % self.contexts.len();
        self.contexts[idx].write().await
    }

    pub async fn read(&self) -> RwLockReadGuard<'_, Context> {
        let idx = self.next_context.load(Ordering::Relaxed) % self.contexts.len();
        self.contexts[idx].read().await
    }

    pub fn size(&self) -> usize {
        self.contexts.len()
    }
}