triton_rs/
model.rs

1use crate::{check_err, Error};
2use libc::c_char;
3use std::ffi::CStr;
4use std::fs::File;
5use std::io::prelude::*;
6use std::path::PathBuf;
7use std::ptr;
8
9pub struct Model {
10    ptr: *mut triton_sys::TRITONBACKEND_Model,
11}
12
13impl Model {
14    pub fn from_ptr(ptr: *mut triton_sys::TRITONBACKEND_Model) -> Self {
15        Self { ptr }
16    }
17
18    pub fn name(&self) -> Result<String, Error> {
19        let mut model_name: *const c_char = ptr::null_mut();
20        check_err(unsafe { triton_sys::TRITONBACKEND_ModelName(self.ptr, &mut model_name) })?;
21
22        let c_str = unsafe { CStr::from_ptr(model_name) };
23        Ok(c_str.to_string_lossy().to_string())
24    }
25
26    pub fn version(&self) -> Result<u64, Error> {
27        let mut version = 0u64;
28        check_err(unsafe { triton_sys::TRITONBACKEND_ModelVersion(self.ptr, &mut version) })?;
29        Ok(version)
30    }
31
32    pub fn location(&self) -> Result<String, Error> {
33        let mut artifact_type: triton_sys::TRITONBACKEND_ArtifactType = 0u32;
34        let mut location: *const c_char = ptr::null_mut();
35        check_err(unsafe {
36            triton_sys::TRITONBACKEND_ModelRepository(self.ptr, &mut artifact_type, &mut location)
37        })?;
38
39        let c_str = unsafe { CStr::from_ptr(location) };
40        Ok(c_str.to_string_lossy().to_string())
41    }
42
43    pub fn path(&self, filename: &str) -> Result<PathBuf, Error> {
44        Ok(PathBuf::from(format!(
45            "{}/{}/{}",
46            self.location()?,
47            self.version()?,
48            filename
49        )))
50    }
51
52    pub fn load_file(&self, filename: &str) -> Result<Vec<u8>, Error> {
53        let path = self.path(filename)?;
54        let mut f = File::open(path)?;
55
56        let mut buffer = Vec::new();
57        f.read_to_end(&mut buffer)?;
58
59        Ok(buffer)
60    }
61}
triton_rs/model.rs

triton_rs/
model.rs
