async_tensorrt/ffi/pre/
logger.rs

1use cpp::cpp;
2
3cpp! {{
4    #ifndef ODDITY_FFI_LOGGER
5    #define ODDITY_FFI_LOGGER
6
7    #include <mutex>
8
9    // The custom logger is required for TensorRT. We can use it to intercept error messages and
10    // diagnostics. TensorRT is very verbose when verbose.
11    class Logger : public ILogger
12    {
13    public:
14        // Implements the `log` function for the logger. This will be invoked
15        // by TensorRT for every log message.
16        void log(Severity severity, const char* msg) noexcept override {
17            // If severity is `ERROR` or worse, then store the error message in
18            // `m_lastError`.
19            if (severity == Severity::kERROR || severity == Severity::kINTERNAL_ERROR) {
20                std::lock_guard<std::mutex> _lastErrorGuard(m_lastErrorMutex);
21                m_lastError = std::string(msg);
22            }
23            // Pass message on to Rust handler.
24            std::int32_t severity_val = static_cast<std::int32_t>(severity);
25            rust!(Logger_handleLogMessage [
26                severity_val : i32 as "std::int32_t",
27                msg : *const std::os::raw::c_char as "const char*"
28            ] {
29                handle_log_message_raw(severity_val, msg);
30            });
31        }
32
33        // Get last logged error message.
34        const std::string getLastError() {
35            std::lock_guard<std::mutex> _lastErrorGuard(m_lastErrorMutex);
36            return m_lastError;
37        }
38    private:
39        std::mutex m_lastErrorMutex {};
40        std::string m_lastError = "";
41    }
42    GLOBAL_LOGGER;
43
44    #endif // ODDITY_FFI_LOGGER
45}}
46
47/// TensorRT logging message severity.
48#[derive(Debug, Clone, Copy, PartialEq)]
49pub enum Severity {
50    /// An internal error has occurred. Execution is unrecoverable.
51    InternalError,
52    /// An application error has occurred.
53    Error,
54    /// An application error has been discovered, but TensorRT has recovered or fallen back to a default.
55    Warning,
56    /// Informational messages with instructional information.
57    Info,
58    /// Verbose messages with debugging information.
59    Verbose,
60    /// A severity code was provied by TensorRT that was not recognized.
61    Unknown,
62}
63
64impl From<i32> for Severity {
65    /// Convert from raw log level integer to [`Severity`].
66    fn from(value: i32) -> Self {
67        match value {
68            0 => Severity::InternalError,
69            1 => Severity::Error,
70            2 => Severity::Warning,
71            3 => Severity::Info,
72            4 => Severity::Verbose,
73            _ => Severity::Unknown,
74        }
75    }
76}
77
78/// Raw handler for log messages.
79///
80/// This function redirects logging to `tracing`, with the following rules:
81/// * `InternalError` and `Error` become `error`.
82/// * `Warning` becomes `warn`.
83/// * `Info` becomes `trace`.
84/// * All other logging is ignored.
85///
86/// # Arguments
87///
88/// * `severity` - Integer severity value of log message.
89/// * `msg` - Raw C string log message.
90///
91/// # Safety
92///
93/// The caller must ensure that the message in `msg` is a valid pointer to a C string.
94unsafe fn handle_log_message_raw(severity: i32, msg: *const std::os::raw::c_char) {
95    let msg_c_str: &std::ffi::CStr = std::ffi::CStr::from_ptr(msg);
96    let msg = msg_c_str.to_str().unwrap_or("");
97    if !msg.is_empty() {
98        match severity.into() {
99            Severity::InternalError | Severity::Error => {
100                tracing::error!(target: "tensorrt", "{msg}");
101            }
102            Severity::Warning => {
103                tracing::warn!(target: "tensorrt", "{msg}");
104            }
105            Severity::Info => {
106                tracing::trace!(target: "tensorrt", "{msg}");
107            }
108            _ => {}
109        }
110    }
111}
async_tensorrt/ffi/pre/logger.rs

async_tensorrt/ffi/pre/
logger.rs
