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
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

use std::sync::atomic::{AtomicBool, Ordering};

use crate::error::check;
use crate::{c_chars_to_string, sys, Error, Result};

static INITIALIZED: AtomicBool = AtomicBool::new(false);

/// The neuron runtime, initialized with `nrt_init` and closed with `nrt_close` on drop.
#[derive(Debug)]
pub struct Runtime {
    _private: (),
}

impl Runtime {
    /// Initializes the runtime for framework-less execution.
    pub fn new() -> Result<Self> {
        Self::with_framework(sys::NRT_FRAMEWORK_TYPE_NO_FW, "", "")
    }

    /// Initializes the runtime, `framework` is one of the `sys::NRT_FRAMEWORK_TYPE_*` values.
    pub fn with_framework(
        framework: sys::nrt_framework_type_t,
        fw_version: &str,
        fal_version: &str,
    ) -> Result<Self> {
        if INITIALIZED.swap(true, Ordering::SeqCst) {
            return Err(Error::AlreadyInitialized);
        }
        let init = || -> Result<()> {
            let fw_version = std::ffi::CString::new(fw_version)?;
            let fal_version = std::ffi::CString::new(fal_version)?;
            let status =
                unsafe { sys::nrt_init(framework, fw_version.as_ptr(), fal_version.as_ptr()) };
            check(status, "nrt_init")
        };
        if let Err(err) = init() {
            INITIALIZED.store(false, Ordering::SeqCst);
            return Err(err);
        }
        Ok(Self { _private: () })
    }

    /// Number of NeuronCores (logical, i.e. VNCs) on the instance.
    pub fn total_vnc_count(&self) -> Result<u32> {
        let mut count = 0;
        check(unsafe { sys::nrt_get_total_vnc_count(&mut count) }, "nrt_get_total_vnc_count")?;
        Ok(count)
    }

    /// Number of NeuronCores (logical, i.e. VNCs) visible to this process.
    pub fn visible_vnc_count(&self) -> Result<u32> {
        let mut count = 0;
        check(unsafe { sys::nrt_get_visible_vnc_count(&mut count) }, "nrt_get_visible_vnc_count")?;
        Ok(count)
    }
}

impl Drop for Runtime {
    fn drop(&mut self) {
        unsafe { sys::nrt_close() };
        INITIALIZED.store(false, Ordering::SeqCst);
    }
}

/// Version of the libnrt library linked at runtime.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Version {
    pub major: u64,
    pub minor: u64,
    pub patch: u64,
    pub maintenance: u64,
    pub detail: String,
    pub git_hash: String,
}

pub fn version() -> Result<Version> {
    let mut v = sys::nrt_version_t::default();
    let status = unsafe { sys::nrt_get_version(&mut v, std::mem::size_of::<sys::nrt_version_t>()) };
    check(status, "nrt_get_version")?;
    Ok(Version {
        major: v.rt_major,
        minor: v.rt_minor,
        patch: v.rt_patch,
        maintenance: v.rt_maintenance,
        detail: c_chars_to_string(&v.rt_detail),
        git_hash: c_chars_to_string(&v.git_hash),
    })
}

/// Instance family and size, `family` is one of the `sys::NRT_INSTANCE_*` values.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct InstanceInfo {
    pub family: u32,
    pub size: u32,
    pub arch_name: String,
    pub device_revision: String,
}

pub fn instance_info() -> Result<InstanceInfo> {
    let mut info = sys::nrt_instance_info_t::default();
    let status = unsafe {
        sys::nrt_get_instance_info(&mut info, std::mem::size_of::<sys::nrt_instance_info_t>())
    };
    check(status, "nrt_get_instance_info")?;
    Ok(InstanceInfo {
        family: info.family,
        size: info.size,
        arch_name: c_chars_to_string(&info.arch_name),
        device_revision: c_chars_to_string(&info.device_revision),
    })
}