neuronx 0.1.0

Bindings for the AWS Neuron runtime (libnrt).
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

use crate::sys;

/// Main library error type.
#[derive(thiserror::Error, Debug)]
pub enum Error {
    /// A libnrt call returned a status other than `NRT_SUCCESS`.
    #[error("{op} failed: {msg} (status {status})")]
    Nrt { op: &'static str, status: sys::NRT_STATUS, msg: String },

    /// libnrt is process-wide state, only one `Runtime` can exist at a time.
    #[error("the neuron runtime is already initialized in this process")]
    AlreadyInitialized,

    /// A host buffer does not match the size of the tensor region it is copied to/from.
    #[error("size mismatch in {op}: tensor region has {expected} bytes, host buffer has {got}")]
    SizeMismatch { op: &'static str, expected: usize, got: usize },

    /// A string handed to libnrt contains an interior nul byte.
    #[error("interior nul byte in string: {0}")]
    NulError(#[from] std::ffi::NulError),
}

pub type Result<T> = std::result::Result<T, Error>;

pub(crate) fn check(status: sys::NRT_STATUS, op: &'static str) -> Result<()> {
    if status == sys::NRT_SUCCESS {
        return Ok(());
    }
    let msg = unsafe { sys::nrt_get_status_as_str(status) };
    let msg = if msg.is_null() {
        "unknown status".to_string()
    } else {
        unsafe { std::ffi::CStr::from_ptr(msg) }.to_string_lossy().into_owned()
    };
    Err(Error::Nrt { op, status, msg })
}