ort_parallel/
async_pool.rs

1use std::{
2    path::{Path, PathBuf},
3    sync::Arc,
4    time::{Duration, SystemTime, SystemTimeError},
5};
6
7use ort::session::{
8    Input, RunOptions, SelectedOutputMarker, Session, SessionInputs, SessionOutputs,
9    builder::{PrepackedWeights, SessionBuilder},
10};
11
12use tokio::sync::{Mutex, Semaphore};
13
14use crate::SessionBuilderFactory;
15
16pub struct AsyncSessionPool {
17    sessions: Arc<Mutex<Vec<Arc<Mutex<Session>>>>>,
18    available_sessions: Arc<Mutex<Vec<usize>>>,
19    sem: Arc<Semaphore>,
20    max: usize,
21    builder: SessionBuilderFactory,
22    file: PathBuf,
23    pub inputs: Vec<Input>,
24}
25
26fn clone_inputs(inp: &Input) -> Input {
27    Input {
28        name: inp.name.clone(),
29        input_type: inp.input_type.clone(),
30    }
31}
32
33impl AsyncSessionPool {
34    pub fn commit_from_file(
35        builder: SessionBuilder,
36        path: &Path,
37        max_sessions: usize,
38    ) -> ort::Result<Self> {
39        assert!(max_sessions > 0);
40        let prepacked_weights = PrepackedWeights::new();
41        let builder = builder.with_prepacked_weights(&prepacked_weights)?;
42        let fs = builder.clone().commit_from_file(path)?;
43        let inputs = fs.inputs.iter().map(clone_inputs).collect::<Vec<Input>>();
44        Ok(Self {
45            inputs,
46            sessions: Arc::new(Mutex::new(vec![Arc::new(Mutex::new(fs))])),
47            sem: Arc::new(Semaphore::new(max_sessions)),
48            available_sessions: Arc::new(Mutex::new(vec![0])),
49            max: max_sessions,
50            builder: SessionBuilderFactory(builder),
51            file: path.to_path_buf(),
52        })
53    }
54
55    pub async fn load_all(&self) -> ort::Result<()> {
56        let count = self.max - self.sessions.lock().await.len();
57        let mut sessions = self.sessions.lock().await;
58        let mut avail = self.available_sessions.lock().await;
59        if count != 0 {
60            for _ in 0..count {
61                let session = self.create_new()?;
62                sessions.push(session.clone());
63                avail.push(sessions.len() - 1);
64            }
65        }
66        Ok(())
67    }
68
69    fn create_new(&self) -> Result<Arc<Mutex<Session>>, ort::Error> {
70        Ok(Arc::new(Mutex::new(
71            self.builder.generate().commit_from_file(&self.file)?,
72        )))
73    }
74
75    async fn release_session(&self, idx: usize) {
76        self.available_sessions.lock().await.push(idx);
77        self.sem.add_permits(1);
78    }
79
80    async fn get_session(&self) -> Result<(Arc<Mutex<Session>>, usize), ort::Error> {
81        let _permit = self.sem.acquire().await.unwrap();
82
83        if let Some(idx) = self.available_sessions.lock().await.pop() {
84            let sessions = self.sessions.lock().await;
85            return Ok((sessions[idx].clone(), idx));
86        }
87
88        if self.sessions.lock().await.len() < self.max {
89            let session = match self.create_new() {
90                Ok(v) => v,
91                Err(e) => {
92                    self.sem.add_permits(1);
93                    return Err(e);
94                }
95            };
96            let mut sessions = self.sessions.lock().await;
97            sessions.push(session.clone());
98            return Ok((session, sessions.len() - 1));
99        }
100
101        unreachable!()
102    }
103
104    pub async fn run_async<'r, 's: 'r, 'i, 'v: 'i + 'r, O: SelectedOutputMarker, const N: usize>(
105        &'s self,
106        input_values: impl Into<SessionInputs<'i, 'v, N>>,
107        run_options: &'r RunOptions<O>,
108    ) -> Result<SessionOutputs<'r>, ort::Error> {
109        let (ses, idx) = self.get_session().await?;
110        let ses: &'s mut Session = unsafe { &mut *(&mut *ses.lock().await as *mut Session) };
111        let out = match ses.run_async(input_values, run_options) {
112            Ok(v) => v.await,
113            Err(e) => Err(e),
114        };
115        self.release_session(idx).await;
116        out
117    }
118
119    pub async fn run_async_profiled<
120        'r,
121        's: 'r,
122        'i,
123        'v: 'i + 'r,
124        O: SelectedOutputMarker,
125        const N: usize,
126    >(
127        &'s self,
128        input_values: impl Into<SessionInputs<'i, 'v, N>>,
129        run_options: &'r RunOptions<O>,
130    ) -> Result<
131        (
132            SessionOutputs<'r>,
133            Result<Duration, SystemTimeError>,
134            String,
135        ),
136        ort::Error,
137    > {
138        let (ses, idx) = self.get_session().await?;
139        let ses_ptr = &mut *ses.lock().await as *mut Session;
140        let ses1: &'s mut Session = unsafe { &mut *(ses_ptr) };
141        let ses2: &'s mut Session = unsafe { &mut *(ses_ptr) };
142
143        let earlier = SystemTime::now();
144        let _ = ses1.profiling_start_ns()?;
145        let out = match ses1.run_async(input_values, run_options) {
146            Ok(v) => v.await,
147            Err(e) => Err(e),
148        };
149        let end = ses2.end_profiling()?;
150        let now = SystemTime::now().duration_since(earlier);
151        self.release_session(idx).await;
152        out.map(|v| (v, now, end))
153    }
154}
ort_parallel/async_pool.rs

ort_parallel/
async_pool.rs
