Skip to content

Commit 7852564

Browse files
porcelijnporcelijn
committed
Error handling in memory allocation of responses
1 parent b9fedc4 commit 7852564

5 files changed

Lines changed: 164 additions & 38 deletions

File tree

‎triton-rs/src/lib.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -30,7 +30,7 @@ pub use triton_sys as sys;
3030
pub type Error = Box<dyn std::error::Error>;
3131

3232
#[allow(non_snake_case)]
33-
pub fn to_TRITONSERVER_Error(err: Error) -> *const triton_sys::TRITONSERVER_Error {
33+
pub fn to_TRITONSERVER_Error(err: Error) -> *mut triton_sys::TRITONSERVER_Error {
3434
let err = std::ffi::CString::new(err.to_string()).expect("CString::new failed");
3535
unsafe {
3636
triton_sys::TRITONSERVER_ErrorNew(

‎triton-rs/src/model_executor.rs‎

Lines changed: 152 additions & 33 deletions
Original file line numberDiff line numberDiff line change
@@ -4,8 +4,9 @@ use std::ffi::CString;
44
use std::os::raw::c_void;
55
use std::ptr;
66
// use async_trait::async_trait;
7-
use crate::{check_err, Error, InferenceRequest, InferenceResponse};
8-
use crate::Server;
7+
use crate::{check_err, Error, InferenceRequest, InferenceResponse, Server};
8+
#[cfg(not(test))]
9+
use crate::to_TRITONSERVER_Error;
910
use tokio::sync::oneshot;
1011

1112
// // #[async_trait]
@@ -119,61 +120,171 @@ extern "C" fn response_alloc(
119120
unsafe {
120121
*actual_memory_type = preferred_memory_type;
121122
*actual_memory_type_id = preferred_memory_type_id;
122-
123-
if byte_size == 0 {
123+
}
124+
if byte_size == 0 {
125+
unsafe {
124126
*buffer = ptr::null_mut();
125127
*buffer_userp = ptr::null_mut();
126-
return ptr::null_mut();
127128
}
129+
return ptr::null_mut();
130+
}
128131

129-
let allocated_ptr = match preferred_memory_type {
130-
// Implement memory allocation based on type (CPU/GPU)
131-
_ => libc::malloc(byte_size),
132-
};
132+
let allocated_ptr = match preferred_memory_type {
133+
// Implement memory allocation based on type (CPU/GPU)
134+
triton_sys::TRITONSERVER_memorytype_enum_TRITONSERVER_MEMORY_CPU => unsafe {
135+
libc::malloc(byte_size)
136+
},
137+
triton_sys::TRITONSERVER_memorytype_enum_TRITONSERVER_MEMORY_CPU_PINNED =>
138+
unsafe {
139+
*actual_memory_type =
140+
triton_sys::TRITONSERVER_memorytype_enum_TRITONSERVER_MEMORY_CPU;
141+
*actual_memory_type_id = 0;
142+
libc::malloc(byte_size)
143+
},
144+
triton_sys::TRITONSERVER_memorytype_enum_TRITONSERVER_MEMORY_GPU => {
145+
return to_TRITONSERVER_Error("GPU memory not implemented".into());
146+
},
147+
_ => {
148+
return to_TRITONSERVER_Error("Invalid memory type requested".into());
149+
},
150+
};
133151

134-
if !allocated_ptr.is_null() {
135-
*buffer = allocated_ptr;
136-
*buffer_userp = Box::into_raw(Box::new(CString::from_vec_unchecked(
137-
std::ffi::CStr::from_ptr(tensor_name).to_bytes().to_vec(),
138-
))) as *mut c_void;
139-
}
152+
if allocated_ptr.is_null() {
153+
return to_TRITONSERVER_Error("Failed to allocate memory for response".into());
154+
}
140155

141-
ptr::null_mut() // Success
156+
// Store tensor name on Rust heap
157+
let tensor_name = unsafe { CString::from_vec_unchecked(
158+
std::ffi::CStr::from_ptr(tensor_name).to_bytes().to_vec()) };
159+
let tensor_name = Box::into_raw(Box::new(tensor_name));
160+
unsafe {
161+
*buffer = allocated_ptr;
162+
*buffer_userp = tensor_name as *mut c_void;
142163
}
164+
165+
ptr::null_mut() // Success
143166
}
144167

145168
extern "C" fn response_release(
146169
_allocator: *mut triton_sys::TRITONSERVER_ResponseAllocator,
147170
buffer: *mut c_void,
148171
buffer_userp: *mut c_void,
149-
_byte_size: libc::size_t,
172+
byte_size: libc::size_t,
150173
memory_type: triton_sys::TRITONSERVER_MemoryType,
151174
_memory_type_id: i64,
152175
) -> *mut triton_sys::TRITONSERVER_Error {
153-
unsafe {
154-
if !buffer.is_null() {
155-
match memory_type {
156-
triton_sys::TRITONSERVER_memorytype_enum_TRITONSERVER_MEMORY_CPU => {
157-
libc::free(buffer)
158-
}
159-
_ => libc::free(buffer), // Simplified for example
160-
};
161-
}
162-
163-
if !buffer_userp.is_null() {
164-
drop(Box::from_raw(buffer_userp as *mut CString));
176+
if !buffer.is_null() {
177+
match memory_type {
178+
triton_sys::TRITONSERVER_memorytype_enum_TRITONSERVER_MEMORY_CPU =>
179+
unsafe {
180+
// Detect use after free errors
181+
libc::free(libc::memset(buffer, 123, byte_size));
182+
// libc::free(buffer);
183+
},
184+
_ => unimplemented!("Memory type not implemented"),
165185
}
186+
}
166187

167-
ptr::null_mut() // Success
188+
if !buffer_userp.is_null() {
189+
let tensor_name = unsafe { Box::from_raw(buffer_userp as *mut CString) };
190+
drop(tensor_name);
168191
}
192+
193+
ptr::null_mut() // Success
169194
}
170195

171196
#[cfg(test)]
172197
mod tests {
173-
// use super::*;
174-
// use std::ptr;
198+
use super::*;
199+
use std::ptr;
175200
// use tokio::test;
176-
//
201+
202+
#[test]
203+
fn test_allocation_cpu() {
204+
let allocator = ptr::null_mut();
205+
let tensor_name = c"foo";
206+
let byte_size: libc::size_t = 123;
207+
let mut memory_type =
208+
triton_sys::TRITONSERVER_memorytype_enum_TRITONSERVER_MEMORY_CPU;
209+
let mut memory_type_id: i64 = 321;
210+
let mut buffer = ptr::null_mut();
211+
let mut buffer_userp = ptr::null_mut();
212+
213+
let r = response_alloc(allocator, tensor_name.as_ptr(), byte_size,
214+
memory_type, memory_type_id, ptr::null_mut(),
215+
&mut buffer, &mut buffer_userp,
216+
&mut memory_type, &mut memory_type_id);
217+
assert!(r.is_null());
218+
assert!(!buffer.is_null());
219+
assert!(!buffer_userp.is_null());
220+
assert_eq!(memory_type, 0);
221+
assert_eq!(memory_type_id, 321);
222+
223+
let slice: &mut [u8] = unsafe {
224+
std::slice::from_raw_parts_mut(buffer as *mut u8, 123)
225+
};
226+
slice[122] = 0xff; // poke in buffer
227+
let tensor_name = unsafe { Box::from_raw(buffer_userp as *mut CString) };
228+
assert_eq!(tensor_name.as_c_str(), c"foo"); // peek
229+
Box::leak(tensor_name); // avoid double free
230+
231+
let r = response_release(allocator, buffer, buffer_userp, byte_size,
232+
memory_type, memory_type_id);
233+
assert!(r.is_null());
234+
}
235+
236+
#[test]
237+
fn test_allocation_cpu_pinned() {
238+
let allocator = ptr::null_mut();
239+
let tensor_name = c"bar";
240+
let byte_size: libc::size_t = 123;
241+
let mut memory_type =
242+
triton_sys::TRITONSERVER_memorytype_enum_TRITONSERVER_MEMORY_CPU_PINNED;
243+
let mut memory_type_id: i64 = 321;
244+
let mut buffer = ptr::null_mut();
245+
let mut buffer_userp = ptr::null_mut();
246+
247+
let r = response_alloc(allocator, tensor_name.as_ptr(), byte_size,
248+
memory_type, memory_type_id, ptr::null_mut(),
249+
&mut buffer, &mut buffer_userp,
250+
&mut memory_type, &mut memory_type_id);
251+
assert!(r.is_null());
252+
assert!(!buffer.is_null());
253+
assert!(!buffer_userp.is_null());
254+
assert_eq!(memory_type, 0); // <-- changed PINNED -> Regular
255+
assert_eq!(memory_type_id, 0); // changed 321 -> 0
256+
257+
let slice: &mut [u8] = unsafe {
258+
std::slice::from_raw_parts_mut(buffer as *mut u8, 123)
259+
};
260+
slice[122] = 0xff; // poke in buffer
261+
let tensor_name = unsafe { Box::from_raw(buffer_userp as *mut CString) };
262+
assert_eq!(tensor_name.as_c_str(), c"bar"); // peek
263+
Box::leak(tensor_name); // avoid double free
264+
265+
let r = response_release(allocator, buffer, buffer_userp, byte_size,
266+
memory_type, memory_type_id);
267+
assert!(r.is_null());
268+
}
269+
270+
#[test]
271+
fn test_allocation_gpu() {
272+
let allocator = ptr::null_mut();
273+
let tensor_name = c"baz";
274+
let byte_size: libc::size_t = 123;
275+
let mut memory_type =
276+
triton_sys::TRITONSERVER_memorytype_enum_TRITONSERVER_MEMORY_GPU;
277+
let mut memory_type_id: i64 = 321;
278+
let mut buffer = ptr::null_mut();
279+
let mut buffer_userp = ptr::null_mut();
280+
281+
let r = response_alloc(allocator, tensor_name.as_ptr(), byte_size,
282+
memory_type, memory_type_id, ptr::null_mut(),
283+
&mut buffer, &mut buffer_userp,
284+
&mut memory_type, &mut memory_type_id);
285+
assert!(!r.is_null());
286+
}
287+
177288
// #[test]
178289
// async fn test_executor_creation() {
179290
// // Create with null server pointer for test
@@ -184,3 +295,11 @@ mod tests {
184295
// let executor = executor.unwrap();
185296
// }
186297
}
298+
299+
// TEST STUB
300+
#[cfg(test)]
301+
#[allow(non_snake_case)]
302+
pub fn to_TRITONSERVER_Error(_err: Error) -> *mut triton_sys::TRITONSERVER_Error {
303+
std::ptr::NonNull::<triton_sys::TRITONSERVER_Error>::dangling().as_ptr()
304+
}
305+

‎triton-rs/src/request.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -80,7 +80,7 @@ impl<'a> Input {
8080
pub fn slice<T>(&self) -> Result<&[T], Error> {
8181
let mut buffer: *const c_void = ptr::null_mut();
8282
let index = 0;
83-
let mut memory_type: triton_sys::TRITONSERVER_MemoryType = 0;
83+
let mut memory_type = triton_sys::TRITONSERVER_memorytype_enum_TRITONSERVER_MEMORY_CPU;
8484
let mut memory_type_id = 0;
8585
let mut buffer_byte_size = 0;
8686
check_err(unsafe {

‎triton-rs/src/response.rs‎

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -79,7 +79,9 @@ impl Response {
7979
let data_type = <T as SupportedTypes>::of();
8080
assert_eq!(data_type.byte_size() as usize, std::mem::size_of::<T>());
8181
let mut output = self.output(name, data_type, shape)?;
82-
output.set_data(data)?;
82+
if data.len() > 0 {
83+
output.set_data(data)?;
84+
}
8385
Ok(())
8486
}
8587

@@ -115,7 +117,7 @@ impl Output {
115117
let mut buffer: *mut c_void = ptr::null_mut();
116118
let element_len = data.len();
117119
let buffer_byte_size = std::mem::size_of_val(data) as u64;
118-
let mut memory_type: triton_sys::TRITONSERVER_MemoryType = 0;
120+
let mut memory_type = triton_sys::TRITONSERVER_memorytype_enum_TRITONSERVER_MEMORY_CPU;
119121
let mut memory_type_id = 0;
120122
check_err(unsafe {
121123
triton_sys::TRITONBACKEND_OutputBuffer(
@@ -127,6 +129,10 @@ impl Output {
127129
)
128130
})?;
129131

132+
if buffer.is_null() {
133+
return Err("Failed to allocate output buffer".into());
134+
}
135+
130136
let mem: &mut [T] = unsafe {
131137
slice::from_raw_parts_mut(buffer as *mut T, element_len)
132138
};

‎triton-sys/src/lib.rs‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,4 +4,5 @@
44

55
// include!(concat!(env!("OUT_DIR"), "/bindings.rs"));
66
pub mod bindings;
7-
pub use bindings::*;
7+
pub use bindings::*;
8+

0 commit comments

Comments
 (0)