diff --git a/src/clients/python/__init__.py b/src/clients/python/__init__.py index b691637e676..646a5d5438d 100644 --- a/src/clients/python/__init__.py +++ b/src/clients/python/__init__.py @@ -116,7 +116,7 @@ def from_param(cls, value): _crequest_shm_control_ctx_del.argtypes = [c_void_p] _crequest_shm_control_ctx_register = _crequest.SharedMemoryControlContextRegister _crequest_shm_control_ctx_register.restype = c_void_p -_crequest_shm_control_ctx_register.argtypes = [c_void_p, _utf8] +_crequest_shm_control_ctx_register.argtypes = [c_void_p, _utf8, c_void_p, c_uint64, c_uint64] _crequest_shm_control_ctx_unregister = _crequest.SharedMemoryControlContextUnregister _crequest_shm_control_ctx_unregister.restype = c_void_p _crequest_shm_control_ctx_unregister.argtypes = [c_void_p, _utf8] @@ -214,7 +214,7 @@ def from_param(cls, value): _crequest_infer_ctx_result_next_class.restype = c_void_p _crequest_infer_ctx_result_next_class.argtypes = [c_void_p, c_uint64, POINTER(c_uint64), POINTER(c_float), POINTER(c_char_p)] -_crequest_get_shared_memory_handle_info = _crequest.GetSharedMemoryHandleInfo +_crequest_get_shared_memory_handle_info = _crequest.SharedMemoryControlContextGetSharedMemoryHandle _crequest_get_shared_memory_handle_info.restype = c_void_p _crequest_get_shared_memory_handle_info.argtypes = [c_void_p, POINTER(c_void_p), POINTER(c_char_p), POINTER(c_int)] @@ -727,15 +727,15 @@ def close(self): _crequest_shm_control_ctx_del(self._ctx) self._ctx = None - def register(self, name, shm_key, offset, byte_size): + def register(self, name, shm_handle, offset, byte_size): """Request the inference server to register specified shared memory region. Parameters ---------- name : str The name of the shared memory region to be registered. - shm_key : str - The unique key of the shared memory object. + shm_handle : c_void_p + The handle for the shared memory region. offset : int The offset from the start of the shared shared memory region. byte_size : int @@ -753,7 +753,7 @@ def register(self, name, shm_key, offset, byte_size): _raise_error("SharedMemoryControlContext is closed") self._last_request_id = _raise_if_error( - c_void_p(_crequest_shm_control_ctx_register(self._ctx, name, shm_key, offset, byte_size))) + c_void_p(_crequest_shm_control_ctx_register(self._ctx, name, shm_handle, offset, byte_size))) return def unregister(self, name): @@ -963,7 +963,7 @@ def _prepare_request(self, inputs, outputs, or (not isinstance(output_format[1], (list, tuple))) or (type(output_format[2]) != int) \ or (type(output_format[3]) != int): _raise_error("shared memory requires tuple of size 4" \ - + " - output_format(RAW), [shm_key (string), base address(c_void_p)], offset (int), size (int)") + + " - output_format(RAW), [shm_key (string), shm_handle(c_void_p)], offset (int), size (int)") _raise_if_error( c_void_p( _crequest_infer_ctx_options_add_shared_memory( diff --git a/src/clients/python/crequest.cc b/src/clients/python/crequest.cc index 9e3f0c2cebc..73f655bf00a 100644 --- a/src/clients/python/crequest.cc +++ b/src/clients/python/crequest.cc @@ -373,12 +373,13 @@ SharedMemoryControlContextDelete(SharedMemoryControlContextCtx* ctx) nic::Error* SharedMemoryControlContextRegister( - SharedMemoryControlContextCtx* ctx, const char* name, const char* shm_key, + SharedMemoryControlContextCtx* ctx, const char* name, void* shm_handle, const int offset, const int byte_size) { nic::Error err = ctx->ctx->RegisterSharedMemory( - std::string(name), std::string(shm_key), (size_t)offset, - size_t(byte_size)); + std::string(name), + reinterpret_cast(shm_handle)->shm_key_, + (size_t)offset, size_t(byte_size)); if (err.IsOk()) { return nullptr; } @@ -399,7 +400,7 @@ SharedMemoryControlContextUnregister( } nic::Error* -GetSharedMemoryHandleInfo( +SharedMemoryControlContextGetSharedMemoryHandle( void* shm_handle, void** shm_addr, const char** shm_key, int* shm_fd) { SharedMemoryHandle* handle = diff --git a/src/clients/python/crequest.h b/src/clients/python/crequest.h index ecbbe75c96a..018e9c0b83f 100644 --- a/src/clients/python/crequest.h +++ b/src/clients/python/crequest.h @@ -88,11 +88,11 @@ nic::Error* SharedMemoryControlContextNew( const char** headers, int num_headers, bool verbose); void SharedMemoryControlContextDelete(SharedMemoryControlContextCtx* ctx); nic::Error* SharedMemoryControlContextRegister( - SharedMemoryControlContextCtx* ctx, const char* name, const char* shm_key, + SharedMemoryControlContextCtx* ctx, const char* name, void* shm_handle, const int offset, const int byte_size); nic::Error* SharedMemoryControlContextUnregister( SharedMemoryControlContextCtx* ctx, const char* model_name); -nic::Error* GetSharedMemoryHandleInfo( +nic::Error* SharedMemoryControlContextGetSharedMemoryHandle( void* shm_handle, void** shm_addr, const char** shm_key, int* shm_fd); //============================================================================== // InferContext diff --git a/src/clients/python/shared_memory/__init__.py b/src/clients/python/shared_memory/__init__.py index 7db0d0f8f56..2395676e24c 100644 --- a/src/clients/python/shared_memory/__init__.py +++ b/src/clients/python/shared_memory/__init__.py @@ -97,110 +97,85 @@ def _raise_error(msg): _cshm_error_del(err) raise ex +def create_shared_memory_region(shm_key, byte_size): + """Creates a shared memory region with the specified name and size. -class SharedMemoryHelper: - """Client helper functions for using shared memory in python. + Parameters + ---------- + shm_key : str + The unique key of the shared memory object. + byte_size : int + The size in bytes of the shared memory region to be created. + + Returns + ------- + shm_handle : c_void_p + The handle for the shared memory region. + + Raises + ------ + InferenceServerException + If unable to create the shared memory region. """ - def create_shared_memory_region(self, shm_key, byte_size): - """Creates a shared memory region with the specified name and size. + shm_handle = c_void_p() + _raise_if_error( + c_void_p(_cshm_shared_memory_region_create(shm_key, byte_size, byref(shm_handle)))) - Parameters - ---------- - shm_key : str - The unique key of the shared memory object. - byte_size : int - The size in bytes of the shared memory region to be created. + return shm_handle - Returns - ------- - shm_handle : c_void_p - The handle for the shared memory region. +def set_shared_memory_region(shm_handle, offset, input_values): + """Copy the contents of the numpy array into a shared memory region with + the specified identifier, offset and size. - Raises - ------ - InferenceServerException - If unable to create the shared memory region. - """ + Parameters + ---------- + shm_handle : c_void_p + The handle for the shared memory region. + offset : int + The offset from the start of the shared shared memory region. + input_values : np.array + The list of numpy arrays to be copied into the shared memory region. + + Raises + ------ + InferenceServerException + If unable to mmap or set values in the shared memory region. + """ + + if not isinstance(input_values, (list,tuple)): + _raise_error("input_values must be specified as a numpy array") + for input_value in input_values: + if not isinstance(input_value, (np.ndarray,)): + _raise_error("input_values must be specified as a list/tuple of numpy arrays") - shm_handle = c_void_p() + offset_current = offset + for input_value in input_values: + input_value = np.ascontiguousarray(input_value).flatten() + byte_size = input_value.size * input_value.itemsize _raise_if_error( - c_void_p(_cshm_shared_memory_region_create(shm_key, byte_size, byref(shm_handle)))) - - return shm_handle - - def set_shared_memory_region(self, shm_handle, offset, input_values): - """Copy the contents of the numpy array into a shared memory region with - the specified identifier, offset and size. - - Parameters - ---------- - shm_handle : c_void_p - The handle for the shared memory region. - offset : int - The offset from the start of the shared shared memory region. - input_values : np.array - The list of numpy arrays to be copied into the shared memory region. - - Raises - ------ - InferenceServerException - If unable to mmap or set values in the shared memory region. - """ + c_void_p(_cshm_shared_memory_region_set(shm_handle, c_uint64(offset_current), \ + c_uint64(byte_size), input_value.ctypes.data_as(c_void_p)))) + offset_current += byte_size + return - if not isinstance(input_values, (list,tuple)): - _raise_error("input_values must be specified as a numpy array") - for input_value in input_values: - if not isinstance(input_value, (np.ndarray,)): - _raise_error("input_values must be specified as a list/tuple of numpy arrays") - - offset_current = offset - for input_value in input_values: - input_value = np.ascontiguousarray(input_value).flatten() - byte_size = input_value.size * input_value.itemsize - _raise_if_error( - c_void_p(_cshm_shared_memory_region_set(shm_handle, c_uint64(offset_current), \ - c_uint64(byte_size), input_value.ctypes.data_as(c_void_p)))) - offset_current += byte_size - return - - def destroy_shared_memory_region(self, shm_key): - """Unlink a shared memory region with the specified name. - - Parameters - ---------- - shm_key : str - The unique key of the shared memory object. - - Raises - ------ - InferenceServerException - If unable to unlink the shared memory region. - """ +def destroy_shared_memory_region(shm_key): + """Unlink a shared memory region with the specified name. - _raise_if_error( - c_void_p(_cshm_shared_memory_region_destroy(shm_key))) - return - - def unmap_shared_memory_region(self, shm_addr, byte_size): - """Unmap a shared memory region with the specified name and size. - - Parameters - ---------- - shm_addr : void* - The base address of the shared memory region. - byte_size : int - The size in bytes of the data in the shared memory region. - - Raises - ------ - InferenceServerException - If unable to munmap the shared memory region. - """ + Parameters + ---------- + shm_key : str + The unique key of the shared memory object. - _raise_if_error( - c_void_p(_cshm_unmap_shared_memory_region(shm_addr, byte_size))) - return + Raises + ------ + InferenceServerException + If unable to unlink the shared memory region. + """ + + _raise_if_error( + c_void_p(_cshm_shared_memory_region_destroy(shm_key))) + return class InferenceServerException(Exception): diff --git a/src/clients/python/shared_memory/shared_memory.cc b/src/clients/python/shared_memory/shared_memory.cc index f63b25e4f9f..6e045d87982 100644 --- a/src/clients/python/shared_memory/shared_memory.cc +++ b/src/clients/python/shared_memory/shared_memory.cc @@ -139,10 +139,10 @@ SharedMemoryRegionSet( void* shm_handle, size_t offset, size_t byte_size, const void* data) { std::cout << shm_handle << '\n'; - // void* shm_addr = - // reinterpret_cast(shm_handle)->base_addr_; char* - // shm_addr_offset = reinterpret_cast(shm_addr); memcpy(shm_addr_offset - // + offset, data, byte_size); + void* shm_addr = + reinterpret_cast(shm_handle)->base_addr_; + char* shm_addr_offset = reinterpret_cast(shm_addr); + memcpy(shm_addr_offset + offset, data, byte_size); return nullptr; } diff --git a/src/clients/python/simple_shm_client.py b/src/clients/python/simple_shm_client.py index 76e680abd77..b824f6fe7d4 100644 --- a/src/clients/python/simple_shm_client.py +++ b/src/clients/python/simple_shm_client.py @@ -31,7 +31,7 @@ import os from builtins import range from tensorrtserver.api import * -from tensorrtserver.shared_memory import * +import tensorrtserver.shared_memory as shm from ctypes import * FLAGS = None @@ -90,20 +90,19 @@ output_byte_size = input_byte_size # Create Output0 and Output1 in Shared Memory and store shared memory handle - shm_helper = SharedMemoryHelper() shm_key = "/output_simple" - shm_op_handle = shm_helper.create_shared_memory_region(shm_key, output_byte_size * 2) + shm_op_handle = shm.create_shared_memory_region(shm_key, output_byte_size * 2) # Register Output shared memory with TRTIS - shared_memory_ctx.register("output_data", "/output_simple", 0, output_byte_size * 2) + shared_memory_ctx.register("output_data", shm_op_handle, 0, output_byte_size * 2) shm_key = "/input_simple" - shm_ip_handle = shm_helper.create_shared_memory_region(shm_key, input_byte_size * 2) + shm_ip_handle = shm.create_shared_memory_region(shm_key, input_byte_size * 2) # Put input data values into shared memory - shm_helper.set_shared_memory_region(shm_ip_handle, 0, [input0_data, input1_data]) + shm.set_shared_memory_region(shm_ip_handle, 0, [input0_data, input1_data]) # Register Input shared memory with TRTIS - shared_memory_ctx.register("input_data", "/input_simple", 0, input_byte_size * 2) + shared_memory_ctx.register("input_data", shm_ip_handle, 0, input_byte_size * 2) # Send inference request to the inference server. Get results for # both output tensors. @@ -134,6 +133,6 @@ del results shared_memory_ctx.unregister("input_data") - shm_helper.destroy_shared_memory_region("/input_simple") + shm.destroy_shared_memory_region("/input_simple") shared_memory_ctx.unregister("output_data") - shm_helper.destroy_shared_memory_region("/output_simple") + shm.destroy_shared_memory_region("/output_simple")