From f31ffe94a875a3e77044763a57de0c0d24822c2b Mon Sep 17 00:00:00 2001 From: Katherine Yang Date: Mon, 8 Apr 2024 19:53:45 -0700 Subject: [PATCH 1/5] update according to server changes --- .../client_backend/triton_c_api/triton_loader.cc | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/src/c++/perf_analyzer/client_backend/triton_c_api/triton_loader.cc b/src/c++/perf_analyzer/client_backend/triton_c_api/triton_loader.cc index 12c767c3c..35f7657f3 100644 --- a/src/c++/perf_analyzer/client_backend/triton_c_api/triton_loader.cc +++ b/src/c++/perf_analyzer/client_backend/triton_c_api/triton_loader.cc @@ -945,7 +945,8 @@ TritonLoader::Infer( TRITONSERVER_MemoryType memory_type; int64_t memory_type_id; RETURN_IF_ERROR(shm_manager_->GetMemoryInfo( - shm_name, offset, &buf, &memory_type, &memory_type_id)); + shm_name, offset, shm_byte_size, &buf, &memory_type, + &memory_type_id)); alloc_payload.output_map_.emplace( std::piecewise_construct, std::forward_as_tuple(output->Name()), @@ -1149,7 +1150,8 @@ TritonLoader::AddInputs( TRITONSERVER_MemoryType memory_type; int64_t memory_type_id; RETURN_IF_ERROR(shm_manager_->GetMemoryInfo( - shm_name, offset, &buf, &memory_type, &memory_type_id)); + shm_name, offset, shm_byte_size, &buf, &memory_type, + &memory_type_id)); RETURN_IF_TRITONSERVER_ERROR( inference_request_append_input_data_fn_( irequest, input_name, buf, byte_size, From 9f29920933f4925d276b6d85d593973290826b7a Mon Sep 17 00:00:00 2001 From: Katherine Yang Date: Tue, 9 Apr 2024 16:35:09 -0700 Subject: [PATCH 2/5] update get memory info --- .../triton_c_api/shared_memory_manager.cc | 5 +++-- .../triton_c_api/shared_memory_manager.h | 13 ++++++++----- 2 files changed, 11 insertions(+), 7 deletions(-) diff --git a/src/c++/perf_analyzer/client_backend/triton_c_api/shared_memory_manager.cc b/src/c++/perf_analyzer/client_backend/triton_c_api/shared_memory_manager.cc index 25dee2d74..2e6770a47 100644 --- a/src/c++/perf_analyzer/client_backend/triton_c_api/shared_memory_manager.cc +++ b/src/c++/perf_analyzer/client_backend/triton_c_api/shared_memory_manager.cc @@ -89,8 +89,9 @@ SharedMemoryManager::RegisterSystemMemory( Error SharedMemoryManager::GetMemoryInfo( - const std::string& name, size_t offset, void** shm_mapped_addr, - TRITONSERVER_MemoryType* memory_type, int64_t* device_id) + const std::string& name, size_t offset, size_t byte_size, + void** shm_mapped_addr, TRITONSERVER_MemoryType* memory_type, + int64_t* device_id) { // protect shared_memory_map_ from concurrent access std::lock_guard lock(mu_); diff --git a/src/c++/perf_analyzer/client_backend/triton_c_api/shared_memory_manager.h b/src/c++/perf_analyzer/client_backend/triton_c_api/shared_memory_manager.h index 1ffeb292c..78a7ccfaf 100644 --- a/src/c++/perf_analyzer/client_backend/triton_c_api/shared_memory_manager.h +++ b/src/c++/perf_analyzer/client_backend/triton_c_api/shared_memory_manager.h @@ -70,19 +70,22 @@ class SharedMemoryManager { Error RegisterSystemMemory( const std::string& name, void* ptr, const size_t byte_size); - /// Get the access information for the shared memory block with the specified - /// name. Return an Error if named block doesn't exist. + /// Get the access information for the shared memory block + /// with the specified name. Return TRITONSERVER_ERROR_NOT_FOUND + /// if named block doesn't exist. /// \param name The name of the shared memory block to get. /// \param offset The offset in the block + /// \param byte_size The byte size to request for the shm region /// \param shm_mapped_addr Returns the pointer to the shared /// memory block with the specified name and offset /// \param memory_type Returns the type of the memory /// \param device_id Returns the device id associated with the /// memory block - /// \return an Error indicating success or failure. + /// \return a TRITONSERVER_Error indicating success or failure. Error GetMemoryInfo( - const std::string& name, size_t offset, void** shm_mapped_addr, - TRITONSERVER_MemoryType* memory_type, int64_t* device_id); + const std::string& name, size_t offset, size_t byte_size, + void** shm_mapped_addr, TRITONSERVER_MemoryType* memory_type, + int64_t* device_id); /// Removes the named shared memory block of the specified type from /// the manager. Any future attempt to get the details of this block From 1c9d3a544657c1e16c91cd5f8ef401d72ddfb7b2 Mon Sep 17 00:00:00 2001 From: Katherine Yang <80359429+jbkyang-nvi@users.noreply.github.com> Date: Wed, 10 Apr 2024 10:24:27 -0700 Subject: [PATCH 3/5] Update documentation suggestion Co-authored-by: Iman Tabrizian --- .../client_backend/triton_c_api/shared_memory_manager.h | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/c++/perf_analyzer/client_backend/triton_c_api/shared_memory_manager.h b/src/c++/perf_analyzer/client_backend/triton_c_api/shared_memory_manager.h index 78a7ccfaf..face37991 100644 --- a/src/c++/perf_analyzer/client_backend/triton_c_api/shared_memory_manager.h +++ b/src/c++/perf_analyzer/client_backend/triton_c_api/shared_memory_manager.h @@ -71,7 +71,7 @@ class SharedMemoryManager { const std::string& name, void* ptr, const size_t byte_size); /// Get the access information for the shared memory block - /// with the specified name. Return TRITONSERVER_ERROR_NOT_FOUND + /// with the specified name. Return an Error /// if named block doesn't exist. /// \param name The name of the shared memory block to get. /// \param offset The offset in the block From 70ffb7659e55142382377c0228682614eb68ad8d Mon Sep 17 00:00:00 2001 From: Katherine Yang <80359429+jbkyang-nvi@users.noreply.github.com> Date: Wed, 10 Apr 2024 10:24:35 -0700 Subject: [PATCH 4/5] Update documentation suggestion Co-authored-by: Iman Tabrizian --- .../client_backend/triton_c_api/shared_memory_manager.h | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/c++/perf_analyzer/client_backend/triton_c_api/shared_memory_manager.h b/src/c++/perf_analyzer/client_backend/triton_c_api/shared_memory_manager.h index face37991..6b2082c44 100644 --- a/src/c++/perf_analyzer/client_backend/triton_c_api/shared_memory_manager.h +++ b/src/c++/perf_analyzer/client_backend/triton_c_api/shared_memory_manager.h @@ -81,7 +81,7 @@ class SharedMemoryManager { /// \param memory_type Returns the type of the memory /// \param device_id Returns the device id associated with the /// memory block - /// \return a TRITONSERVER_Error indicating success or failure. + /// \return an Error indicating success or failure. Error GetMemoryInfo( const std::string& name, size_t offset, size_t byte_size, void** shm_mapped_addr, TRITONSERVER_MemoryType* memory_type, From 92983fbad8d6baf0150429d83961773b61e3122a Mon Sep 17 00:00:00 2001 From: Katherine Yang Date: Wed, 10 Apr 2024 11:05:33 -0700 Subject: [PATCH 5/5] Include actual server changes --- .../triton_c_api/shared_memory_manager.cc | 23 +++++++++++++++++++ 1 file changed, 23 insertions(+) diff --git a/src/c++/perf_analyzer/client_backend/triton_c_api/shared_memory_manager.cc b/src/c++/perf_analyzer/client_backend/triton_c_api/shared_memory_manager.cc index 2e6770a47..0658daedd 100644 --- a/src/c++/perf_analyzer/client_backend/triton_c_api/shared_memory_manager.cc +++ b/src/c++/perf_analyzer/client_backend/triton_c_api/shared_memory_manager.cc @@ -101,6 +101,29 @@ SharedMemoryManager::GetMemoryInfo( return Error( std::string("Unable to find shared memory region: '" + name + "'")); } + + // validate offset + size_t shm_region_end = 0; + if (it->second->kind_ == TRITONSERVER_MEMORY_CPU) { + shm_region_end = it->second->offset_; + } + if (it->second->byte_size_ > 0) { + shm_region_end += it->second->byte_size_ - 1; + } + if (offset > shm_region_end) { + return Error( + std::string("Invalid offset for shared memory region: '" + name + "'") + .c_str()); + } + // validate byte_size + offset is within memory bounds + size_t total_req_shm = offset + byte_size - 1; + if (total_req_shm > shm_region_end) { + return Error(std::string( + "Invalid offset + byte size for shared memory region: '" + + name + "'") + .c_str()); + } + if (it->second->kind_ == TRITONSERVER_MEMORY_CPU) { *shm_mapped_addr = (void*)((uint8_t*)it->second->mapped_addr_ + it->second->offset_ + offset);