libmir-cuda 0.1.0

CUDA inference backend for libmir
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
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116

mod batch;
mod execution;
mod model;
mod runner;
mod runtime;
mod trace;

use std::{
    collections::HashMap,
    fmt,
    sync::{Arc, Mutex, MutexGuard},
};

use ::runtime::{kv::CacheConfig, scheduler::SchedulerConfig};

use self::model::LoadedModel;
use crate::{CudaBackend, CudaConfig, Error, Result};

#[derive(Clone)]
pub struct CudaEngine {
    backend: CudaBackend,
    cache: CacheConfig,
    session_config: crate::CudaModelSessionConfig,
    scheduler: SchedulerConfig,
    models: Arc<Mutex<HashMap<String, Arc<LoadedModel>>>>,
}

#[derive(Debug, Clone)]
pub struct CudaMemoryStats {
    pub total: u64,
    pub available: u64,
    pub reserved: u64,
    pub used: u64,
    pub device: String,
    pub integrated: bool,
}

impl CudaEngine {
    pub fn new(config: CudaConfig, cache: CacheConfig) -> Result<Self> {
        Self::new_with_scheduler(config, cache, SchedulerConfig::default())
    }

    pub fn new_with_scheduler(
        config: CudaConfig,
        cache: CacheConfig,
        scheduler: SchedulerConfig,
    ) -> Result<Self> {
        let session_config = config.model_session;
        Ok(Self {
            backend: CudaBackend::new(config)?,
            cache,
            session_config,
            scheduler,
            models: Arc::new(Mutex::new(HashMap::new())),
        })
    }

    fn model(&self, id: &str) -> Result<Arc<LoadedModel>> {
        self.models()?
            .get(id)
            .cloned()
            .ok_or_else(|| Error::State(format!("model is not loaded: {id}")))
    }

    fn models(&self) -> Result<MutexGuard<'_, HashMap<String, Arc<LoadedModel>>>> {
        let Ok(models) = self.models.lock() else {
            return Err(Error::State("model registry lock is poisoned".into()));
        };
        Ok(models)
    }

    pub fn clear_model_sessions(&self, model_id: &str) -> Result<()> {
        self.model(model_id)?.clear_sessions()
    }

    pub fn release_session(&self, model_id: &str, session: uuid::Uuid) -> Result<()> {
        self.model(model_id)?.release_session(session)
    }

    pub fn unload_model(&self, model_id: &str) -> Result<bool> {
        let removed = self.models()?.remove(model_id);
        let unloaded = removed.is_some();
        drop(removed);
        Ok(unloaded)
    }

    pub fn clear_memory_cache(&self) -> Result<()> {
        self.backend.trim_memory_pool(0)
    }

    pub fn memory_stats(&self) -> Result<CudaMemoryStats> {
        let pool = self.backend.memory_pool_stats()?;
        let (available, total) = self.backend.memory_info()?;
        let device = self.backend.device_info();
        Ok(CudaMemoryStats {
            total: u64::try_from(total)?,
            available: u64::try_from(available)?,
            reserved: pool.reserved,
            used: pool.used,
            device: device.name.clone(),
            integrated: device.integrated,
        })
    }
}

impl fmt::Debug for CudaEngine {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        formatter
            .debug_struct("CudaEngine")
            .field("device", &self.backend.device_info().name)
            .field("cache", &self.cache)
            .field("session_config", &self.session_config)
            .field("scheduler", &self.scheduler)
            .finish_non_exhaustive()
    }
}