From ba8e8bec0b121d9f9f5e538693eff5b5f4265b19 Mon Sep 17 00:00:00 2001 From: Katherine Yang Date: Fri, 5 Apr 2024 16:51:23 -0700 Subject: [PATCH 1/7] add byte size check to check memory info --- src/grpc/infer_handler.cc | 2 +- src/grpc/infer_handler.h | 3 ++- src/http_server.cc | 6 ++++-- src/shared_memory_manager.cc | 22 +++++++++++++++------- src/shared_memory_manager.h | 6 ++++-- 5 files changed, 26 insertions(+), 13 deletions(-) diff --git a/src/grpc/infer_handler.cc b/src/grpc/infer_handler.cc index 5c60fb46ed..4cd16cee16 100644 --- a/src/grpc/infer_handler.cc +++ b/src/grpc/infer_handler.cc @@ -433,7 +433,7 @@ InferGRPCToInput( } void* tmp; RETURN_IF_ERR(shm_manager->GetMemoryInfo( - region_name, offset, &tmp, &memory_type, &memory_type_id)); + region_name, offset, byte_size, &tmp, &memory_type, &memory_type_id)); base = tmp; if (memory_type == TRITONSERVER_MEMORY_GPU) { #ifdef TRITON_ENABLE_GPU diff --git a/src/grpc/infer_handler.h b/src/grpc/infer_handler.h index 84095a188b..38029902aa 100644 --- a/src/grpc/infer_handler.h +++ b/src/grpc/infer_handler.h @@ -336,7 +336,8 @@ InferAllocatorPayload( TRITONSERVER_MemoryType memory_type; int64_t memory_type_id; RETURN_IF_ERR(shm_manager->GetMemoryInfo( - region_name, offset, &base, &memory_type, &memory_type_id)); + region_name, offset, byte_size, &base, &memory_type, + &memory_type_id)); if (memory_type == TRITONSERVER_MEMORY_GPU) { #ifdef TRITON_ENABLE_GPU diff --git a/src/http_server.cc b/src/http_server.cc index f5b791c06e..29fe7cbb4f 100644 --- a/src/http_server.cc +++ b/src/http_server.cc @@ -2627,7 +2627,8 @@ HTTPAPIServer::ParseJsonTritonIO( TRITONSERVER_MemoryType memory_type; int64_t memory_type_id; RETURN_IF_ERR(shm_manager_->GetMemoryInfo( - shm_region, shm_offset, &base, &memory_type, &memory_type_id)); + shm_region, shm_offset, byte_size, &base, &memory_type, + &memory_type_id)); if (memory_type == TRITONSERVER_MEMORY_GPU) { #ifdef TRITON_ENABLE_GPU cudaIpcMemHandle_t* cuda_handle; @@ -2741,7 +2742,8 @@ HTTPAPIServer::ParseJsonTritonIO( TRITONSERVER_MemoryType memory_type; int64_t memory_type_id; RETURN_IF_ERR(shm_manager_->GetMemoryInfo( - shm_region, offset, &base, &memory_type, &memory_type_id)); + shm_region, offset, byte_size, &base, &memory_type, + &memory_type_id)); if (memory_type == TRITONSERVER_MEMORY_GPU) { #ifdef TRITON_ENABLE_GPU diff --git a/src/shared_memory_manager.cc b/src/shared_memory_manager.cc index 268c6451d1..664bc10bed 100644 --- a/src/shared_memory_manager.cc +++ b/src/shared_memory_manager.cc @@ -333,8 +333,9 @@ SharedMemoryManager::RegisterCUDASharedMemory( TRITONSERVER_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_); @@ -348,20 +349,27 @@ SharedMemoryManager::GetMemoryInfo( } // validate offset - size_t max_offset = 0; + size_t shm_reigion_end = 0; if (it->second->kind_ == TRITONSERVER_MEMORY_CPU) { - max_offset = it->second->offset_; + shm_reigion_end = it->second->offset_; } if (it->second->byte_size_ > 0) { - max_offset += it->second->byte_size_ - 1; + shm_reigion_end += it->second->byte_size_ - 1; } - if (offset > max_offset) { + if (offset > shm_reigion_end) { return TRITONSERVER_ErrorNew( TRITONSERVER_ERROR_INVALID_ARG, std::string("Invalid offset for shared memory region: '" + name + "'") .c_str()); } - // TODO: should also validate byte_size from caller + size_t total_req_shm = offset + byte_size; + if (total_req_shm > shm_reigion_end) { + return TRITONSERVER_ErrorNew( + TRITONSERVER_ERROR_INVALID_ARG, + std::string( + "Invalid 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_ + diff --git a/src/shared_memory_manager.h b/src/shared_memory_manager.h index b282f77bc7..f079308bd5 100644 --- a/src/shared_memory_manager.h +++ b/src/shared_memory_manager.h @@ -83,6 +83,7 @@ class SharedMemoryManager { /// 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 @@ -90,8 +91,9 @@ class SharedMemoryManager { /// memory block /// \return a TRITONSERVER_Error indicating success or failure. TRITONSERVER_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); #ifdef TRITON_ENABLE_GPU /// Get the CUDA memory handle associated with the block name. From 2f8bf9178c08a4e2ef0d056ff4cdb3b73c8c9978 Mon Sep 17 00:00:00 2001 From: Katherine Yang Date: Fri, 5 Apr 2024 18:19:49 -0700 Subject: [PATCH 2/7] add test --- qa/L0_shared_memory/shared_memory_test.py | 29 +++++++++++++++++++++-- qa/L0_shared_memory/test.sh | 3 ++- 2 files changed, 29 insertions(+), 3 deletions(-) diff --git a/qa/L0_shared_memory/shared_memory_test.py b/qa/L0_shared_memory/shared_memory_test.py index bfc48d6261..b8b24234f8 100755 --- a/qa/L0_shared_memory/shared_memory_test.py +++ b/qa/L0_shared_memory/shared_memory_test.py @@ -163,6 +163,7 @@ def _basic_inference( big_shm_name="", big_shm_size=64, shm_output_offset=0, + shm_output_byte_size=0, ): input0_data = np.arange(start=0, stop=16, dtype=np.int32) input1_data = np.ones(shape=16, dtype=np.int32) @@ -192,8 +193,12 @@ def _basic_inference( else: inputs[1].set_shared_memory("input1_data", 64) - outputs[0].set_shared_memory("output0_data", 64, offset=shm_output_offset) - outputs[1].set_shared_memory("output1_data", 64, offset=shm_output_offset) + outputs[0].set_shared_memory( + "output0_data", 64, offset=shm_output_offset, byte_size=shm_output_byte_size + ) + outputs[1].set_shared_memory( + "output1_data", 64, offset=shm_output_offset, byte_size=shm_output_byte_size + ) try: results = triton_client.infer( @@ -348,6 +353,26 @@ def test_infer_offset_out_of_bound(self): self.assertIn("Invalid offset for shared memory region", error_msg[0]) self._cleanup_server(shm_handles) + def test_infer_byte_size_out_of_bound(self): + # Shared memory byte_size outside output region - Throws error + error_msg = [] + shm_handles = self._configure_sever() + offset = 2**64 - 1 + byte_size = 2**32 + + self._basic_inference( + shm_handles[0], + shm_handles[1], + shm_handles[2], + shm_handles[3], + error_msg, + shm_output_offset=offset, + shm_output_byte_size=byte_size, + ) + self.assertEqual(len(error_msg), 1) + self.assertIn("Invalid byte size for shared memory region", error_msg[0]) + self._cleanup_server(shm_handles) + if __name__ == "__main__": _protocol = os.environ.get("CLIENT_TYPE", "http") diff --git a/qa/L0_shared_memory/test.sh b/qa/L0_shared_memory/test.sh index 30abfca545..eb2feedb40 100755 --- a/qa/L0_shared_memory/test.sh +++ b/qa/L0_shared_memory/test.sh @@ -50,7 +50,8 @@ for i in \ test_too_big_shm \ test_mixed_raw_shm \ test_unregisterall \ - test_infer_offset_out_of_bound; do + test_infer_offset_out_of_bound \ + test_infer_byte_size_out_of_bound; do for client_type in http grpc; do SERVER_ARGS="--model-repository=`pwd`/models --log-verbose=1 ${SERVER_ARGS_EXTRA}" SERVER_LOG="./$i.$client_type.server.log" From 0bbf10df4acb45264dce4079526050096e292c39 Mon Sep 17 00:00:00 2001 From: Katherine Yang Date: Mon, 8 Apr 2024 17:12:37 -0700 Subject: [PATCH 3/7] update test issues --- qa/L0_shared_memory/shared_memory_test.py | 5 +++-- qa/L0_shared_memory/test.sh | 10 ---------- 2 files changed, 3 insertions(+), 12 deletions(-) diff --git a/qa/L0_shared_memory/shared_memory_test.py b/qa/L0_shared_memory/shared_memory_test.py index b8b24234f8..97e41dd6d1 100755 --- a/qa/L0_shared_memory/shared_memory_test.py +++ b/qa/L0_shared_memory/shared_memory_test.py @@ -194,10 +194,10 @@ def _basic_inference( inputs[1].set_shared_memory("input1_data", 64) outputs[0].set_shared_memory( - "output0_data", 64, offset=shm_output_offset, byte_size=shm_output_byte_size + "output0_data", shm_output_byte_size, offset=shm_output_offset ) outputs[1].set_shared_memory( - "output1_data", 64, offset=shm_output_offset, byte_size=shm_output_byte_size + "output1_data", shm_output_byte_size, offset=shm_output_offset ) try: @@ -333,6 +333,7 @@ def test_infer_offset_out_of_bound(self): # Shared memory offset outside output region - Throws error error_msg = [] shm_handles = self._configure_sever() + print("testing") if _protocol == "http": # -32 when placed in an int64 signed type, to get a negative offset # by overflowing diff --git a/qa/L0_shared_memory/test.sh b/qa/L0_shared_memory/test.sh index eb2feedb40..4d55d48b06 100755 --- a/qa/L0_shared_memory/test.sh +++ b/qa/L0_shared_memory/test.sh @@ -40,16 +40,6 @@ RET=0 rm -fr *.log for i in \ - test_invalid_create_shm \ - test_valid_create_set_register \ - test_unregister_before_register \ - test_unregister_after_register \ - test_reregister_after_register \ - test_unregister_after_inference \ - test_register_after_inference \ - test_too_big_shm \ - test_mixed_raw_shm \ - test_unregisterall \ test_infer_offset_out_of_bound \ test_infer_byte_size_out_of_bound; do for client_type in http grpc; do From 53f49462fe348666d2046dbffb687315b0564fb0 Mon Sep 17 00:00:00 2001 From: Katherine Yang Date: Mon, 8 Apr 2024 19:52:55 -0700 Subject: [PATCH 4/7] unified global variables, cleaned up cpp naming and updated test to work --- qa/L0_shared_memory/shared_memory_test.py | 50 +++++++++++++++-------- qa/L0_shared_memory/test.sh | 3 +- src/shared_memory_manager.cc | 14 ++++--- 3 files changed, 43 insertions(+), 24 deletions(-) diff --git a/qa/L0_shared_memory/shared_memory_test.py b/qa/L0_shared_memory/shared_memory_test.py index 97e41dd6d1..409e050d3b 100755 --- a/qa/L0_shared_memory/shared_memory_test.py +++ b/qa/L0_shared_memory/shared_memory_test.py @@ -40,6 +40,8 @@ import tritonclient.utils.shared_memory as shm from tritonclient import utils +DEFAULT_SHM_BYTE_SIZE = 64 + class SharedMemoryTest(tu.TestResultCollector): def test_invalid_create_shm(self): @@ -124,16 +126,16 @@ def test_reregister_after_register(self): def _configure_sever(self): shm_ip0_handle = shm.create_shared_memory_region( - "input0_data", "/input0_data", 64 + "input0_data", "/input0_data", DEFAULT_SHM_BYTE_SIZE ) shm_ip1_handle = shm.create_shared_memory_region( - "input1_data", "/input1_data", 64 + "input1_data", "/input1_data", DEFAULT_SHM_BYTE_SIZE ) shm_op0_handle = shm.create_shared_memory_region( - "output0_data", "/output0_data", 64 + "output0_data", "/output0_data", DEFAULT_SHM_BYTE_SIZE ) shm_op1_handle = shm.create_shared_memory_region( - "output1_data", "/output1_data", 64 + "output1_data", "/output1_data", DEFAULT_SHM_BYTE_SIZE ) input0_data = np.arange(start=0, stop=16, dtype=np.int32) input1_data = np.ones(shape=16, dtype=np.int32) @@ -143,10 +145,18 @@ def _configure_sever(self): triton_client = httpclient.InferenceServerClient(_url, verbose=True) else: triton_client = grpcclient.InferenceServerClient(_url, verbose=True) - triton_client.register_system_shared_memory("input0_data", "/input0_data", 64) - triton_client.register_system_shared_memory("input1_data", "/input1_data", 64) - triton_client.register_system_shared_memory("output0_data", "/output0_data", 64) - triton_client.register_system_shared_memory("output1_data", "/output1_data", 64) + triton_client.register_system_shared_memory( + "input0_data", "/input0_data", DEFAULT_SHM_BYTE_SIZE + ) + triton_client.register_system_shared_memory( + "input1_data", "/input1_data", DEFAULT_SHM_BYTE_SIZE + ) + triton_client.register_system_shared_memory( + "output0_data", "/output0_data", DEFAULT_SHM_BYTE_SIZE + ) + triton_client.register_system_shared_memory( + "output1_data", "/output1_data", DEFAULT_SHM_BYTE_SIZE + ) return [shm_ip0_handle, shm_ip1_handle, shm_op0_handle, shm_op1_handle] def _cleanup_server(self, shm_handles): @@ -161,9 +171,10 @@ def _basic_inference( shm_op1_handle, error_msg, big_shm_name="", - big_shm_size=64, + big_shm_size=DEFAULT_SHM_BYTE_SIZE, shm_output_offset=0, - shm_output_byte_size=0, + shm_output_byte_size=DEFAULT_SHM_BYTE_SIZE, + default_shm_byte_size=DEFAULT_SHM_BYTE_SIZE, ): input0_data = np.arange(start=0, stop=16, dtype=np.int32) input1_data = np.ones(shape=16, dtype=np.int32) @@ -184,14 +195,14 @@ def _basic_inference( outputs.append(grpcclient.InferRequestedOutput("OUTPUT0")) outputs.append(grpcclient.InferRequestedOutput("OUTPUT1")) - inputs[0].set_shared_memory("input0_data", 64) + inputs[0].set_shared_memory("input0_data", default_shm_byte_size) if type(shm_ip1_handle) == np.array: inputs[1].set_data_from_numpy(input0_data, binary_data=True) elif big_shm_name != "": inputs[1].set_shared_memory(big_shm_name, big_shm_size) else: - inputs[1].set_shared_memory("input1_data", 64) + inputs[1].set_shared_memory("input1_data", default_shm_byte_size) outputs[0].set_shared_memory( "output0_data", shm_output_byte_size, offset=shm_output_offset @@ -257,9 +268,11 @@ def test_register_after_inference(self): if len(error_msg) > 0: raise Exception(str(error_msg)) shm_ip2_handle = shm.create_shared_memory_region( - "input2_data", "/input2_data", 64 + "input2_data", "/input2_data", DEFAULT_SHM_BYTE_SIZE + ) + triton_client.register_system_shared_memory( + "input2_data", "/input2_data", DEFAULT_SHM_BYTE_SIZE ) - triton_client.register_system_shared_memory("input2_data", "/input2_data", 64) shm_status = triton_client.get_system_shared_memory_status() if _protocol == "http": self.assertTrue(len(shm_status) == 5) @@ -349,6 +362,7 @@ def test_infer_offset_out_of_bound(self): shm_handles[3], error_msg, shm_output_offset=offset, + default_shm_byte_size=0, ) self.assertEqual(len(error_msg), 1) self.assertIn("Invalid offset for shared memory region", error_msg[0]) @@ -358,8 +372,8 @@ def test_infer_byte_size_out_of_bound(self): # Shared memory byte_size outside output region - Throws error error_msg = [] shm_handles = self._configure_sever() - offset = 2**64 - 1 - byte_size = 2**32 + offset = 60 + byte_size = 64 self._basic_inference( shm_handles[0], @@ -371,7 +385,9 @@ def test_infer_byte_size_out_of_bound(self): shm_output_byte_size=byte_size, ) self.assertEqual(len(error_msg), 1) - self.assertIn("Invalid byte size for shared memory region", error_msg[0]) + self.assertIn( + "Invalid offset + byte size for shared memory region", error_msg[0] + ) self._cleanup_server(shm_handles) diff --git a/qa/L0_shared_memory/test.sh b/qa/L0_shared_memory/test.sh index 4d55d48b06..af4f40282d 100755 --- a/qa/L0_shared_memory/test.sh +++ b/qa/L0_shared_memory/test.sh @@ -58,12 +58,13 @@ for i in \ set +e python $SHM_TEST SharedMemoryTest.$i >>$CLIENT_LOG 2>&1 if [ $? -ne 0 ]; then + cat $CLIENT_LOG echo -e "\n***\n*** Test Failed\n***" RET=1 else check_test_results $TEST_RESULT_FILE 1 if [ $? -ne 0 ]; then - cat $CLIENT_LOG + cat $TEST_RESULT_FILE echo -e "\n***\n*** Test Result Verification Failed\n***" RET=1 fi diff --git a/src/shared_memory_manager.cc b/src/shared_memory_manager.cc index 664bc10bed..79d83350ea 100644 --- a/src/shared_memory_manager.cc +++ b/src/shared_memory_manager.cc @@ -349,25 +349,27 @@ SharedMemoryManager::GetMemoryInfo( } // validate offset - size_t shm_reigion_end = 0; + size_t shm_region_end = 0; if (it->second->kind_ == TRITONSERVER_MEMORY_CPU) { - shm_reigion_end = it->second->offset_; + shm_region_end = it->second->offset_; } if (it->second->byte_size_ > 0) { - shm_reigion_end += it->second->byte_size_ - 1; + shm_region_end += it->second->byte_size_ - 1; } - if (offset > shm_reigion_end) { + if (offset > shm_region_end) { return TRITONSERVER_ErrorNew( TRITONSERVER_ERROR_INVALID_ARG, 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; - if (total_req_shm > shm_reigion_end) { + if (total_req_shm > shm_region_end) { return TRITONSERVER_ErrorNew( TRITONSERVER_ERROR_INVALID_ARG, std::string( - "Invalid byte size for shared memory region: '" + name + "'") + "Invalid offset + byte size for shared memory region: '" + name + + "'") .c_str()); } From d38cda63c03ab5d7cd29869c807eda319e68b6fa Mon Sep 17 00:00:00 2001 From: Katherine Yang Date: Mon, 8 Apr 2024 19:55:49 -0700 Subject: [PATCH 5/7] add removed test back --- qa/L0_shared_memory/test.sh | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/qa/L0_shared_memory/test.sh b/qa/L0_shared_memory/test.sh index af4f40282d..62065006a3 100755 --- a/qa/L0_shared_memory/test.sh +++ b/qa/L0_shared_memory/test.sh @@ -40,6 +40,16 @@ RET=0 rm -fr *.log for i in \ + test_invalid_create_shm \ + test_valid_create_set_register \ + test_unregister_before_register \ + test_unregister_after_register \ + test_reregister_after_register \ + test_unregister_after_inference \ + test_register_after_inference \ + test_too_big_shm \ + test_mixed_raw_shm \ + test_unregisterall \ test_infer_offset_out_of_bound \ test_infer_byte_size_out_of_bound; do for client_type in http grpc; do From cb6e4a6d321ea65d19b5d811cf1ee0d70e69acb5 Mon Sep 17 00:00:00 2001 From: Katherine Yang Date: Tue, 9 Apr 2024 16:30:23 -0700 Subject: [PATCH 6/7] fix off by one error and updated test --- qa/L0_shared_memory/shared_memory_test.py | 36 ++++++++++++----------- qa/L0_shared_memory/test.sh | 13 ++++---- src/shared_memory_manager.cc | 2 +- 3 files changed, 27 insertions(+), 24 deletions(-) diff --git a/qa/L0_shared_memory/shared_memory_test.py b/qa/L0_shared_memory/shared_memory_test.py index 409e050d3b..a5a1e8fa16 100755 --- a/qa/L0_shared_memory/shared_memory_test.py +++ b/qa/L0_shared_memory/shared_memory_test.py @@ -40,10 +40,10 @@ import tritonclient.utils.shared_memory as shm from tritonclient import utils -DEFAULT_SHM_BYTE_SIZE = 64 - class SharedMemoryTest(tu.TestResultCollector): + DEFAULT_SHM_BYTE_SIZE = 64 + def test_invalid_create_shm(self): # Raises error since tried to create invalid system shared memory region try: @@ -124,18 +124,18 @@ def test_reregister_after_register(self): self.assertTrue(len(shm_status.regions) == 1) shm.destroy_shared_memory_region(shm_op0_handle) - def _configure_sever(self): + def _configure_sever(self, shm_byte_size=DEFAULT_SHM_BYTE_SIZE): shm_ip0_handle = shm.create_shared_memory_region( - "input0_data", "/input0_data", DEFAULT_SHM_BYTE_SIZE + "input0_data", "/input0_data", shm_byte_size ) shm_ip1_handle = shm.create_shared_memory_region( - "input1_data", "/input1_data", DEFAULT_SHM_BYTE_SIZE + "input1_data", "/input1_data", shm_byte_size ) shm_op0_handle = shm.create_shared_memory_region( - "output0_data", "/output0_data", DEFAULT_SHM_BYTE_SIZE + "output0_data", "/output0_data", shm_byte_size ) shm_op1_handle = shm.create_shared_memory_region( - "output1_data", "/output1_data", DEFAULT_SHM_BYTE_SIZE + "output1_data", "/output1_data", shm_byte_size ) input0_data = np.arange(start=0, stop=16, dtype=np.int32) input1_data = np.ones(shape=16, dtype=np.int32) @@ -146,16 +146,16 @@ def _configure_sever(self): else: triton_client = grpcclient.InferenceServerClient(_url, verbose=True) triton_client.register_system_shared_memory( - "input0_data", "/input0_data", DEFAULT_SHM_BYTE_SIZE + "input0_data", "/input0_data", shm_byte_size ) triton_client.register_system_shared_memory( - "input1_data", "/input1_data", DEFAULT_SHM_BYTE_SIZE + "input1_data", "/input1_data", shm_byte_size ) triton_client.register_system_shared_memory( - "output0_data", "/output0_data", DEFAULT_SHM_BYTE_SIZE + "output0_data", "/output0_data", shm_byte_size ) triton_client.register_system_shared_memory( - "output1_data", "/output1_data", DEFAULT_SHM_BYTE_SIZE + "output1_data", "/output1_data", shm_byte_size ) return [shm_ip0_handle, shm_ip1_handle, shm_op0_handle, shm_op1_handle] @@ -238,7 +238,11 @@ def test_unregister_after_inference(self): error_msg = [] shm_handles = self._configure_sever() self._basic_inference( - shm_handles[0], shm_handles[1], shm_handles[2], shm_handles[3], error_msg + shm_handles[0], + shm_handles[1], + shm_handles[2], + shm_handles[3], + error_msg, ) if len(error_msg) > 0: raise Exception(str(error_msg)) @@ -268,10 +272,10 @@ def test_register_after_inference(self): if len(error_msg) > 0: raise Exception(str(error_msg)) shm_ip2_handle = shm.create_shared_memory_region( - "input2_data", "/input2_data", DEFAULT_SHM_BYTE_SIZE + "input2_data", "/input2_data", self.DEFAULT_SHM_BYTE_SIZE ) triton_client.register_system_shared_memory( - "input2_data", "/input2_data", DEFAULT_SHM_BYTE_SIZE + "input2_data", "/input2_data", self.DEFAULT_SHM_BYTE_SIZE ) shm_status = triton_client.get_system_shared_memory_status() if _protocol == "http": @@ -313,7 +317,7 @@ def test_too_big_shm(self): def test_mixed_raw_shm(self): # Mix of shared memory and RAW inputs error_msg = [] - shm_handles = self._configure_sever() + shm_handles = self._configure_sever(shm_byte_size=256) input1_data = np.ones(shape=16, dtype=np.int32) self._basic_inference( shm_handles[0], [input1_data], shm_handles[2], shm_handles[3], error_msg @@ -346,7 +350,6 @@ def test_infer_offset_out_of_bound(self): # Shared memory offset outside output region - Throws error error_msg = [] shm_handles = self._configure_sever() - print("testing") if _protocol == "http": # -32 when placed in an int64 signed type, to get a negative offset # by overflowing @@ -362,7 +365,6 @@ def test_infer_offset_out_of_bound(self): shm_handles[3], error_msg, shm_output_offset=offset, - default_shm_byte_size=0, ) self.assertEqual(len(error_msg), 1) self.assertIn("Invalid offset for shared memory region", error_msg[0]) diff --git a/qa/L0_shared_memory/test.sh b/qa/L0_shared_memory/test.sh index 62065006a3..ea90b53163 100755 --- a/qa/L0_shared_memory/test.sh +++ b/qa/L0_shared_memory/test.sh @@ -63,12 +63,13 @@ for i in \ fi export CLIENT_TYPE=$client_type - echo "Test: $i, client type: $client_type" >>$CLIENT_LOG + TMP_CLIENT_LOG="./tmp_client.log" + echo "Test: $i, client type: $client_type" >>$TMP_CLIENT_LOG set +e - python $SHM_TEST SharedMemoryTest.$i >>$CLIENT_LOG 2>&1 + python $SHM_TEST SharedMemoryTest.$i >>$TMP_CLIENT_LOG 2>&1 if [ $? -ne 0 ]; then - cat $CLIENT_LOG + cat $TMP_CLIENT_LOG echo -e "\n***\n*** Test Failed\n***" RET=1 else @@ -79,17 +80,17 @@ for i in \ RET=1 fi fi - set -e - + cat $TMP_CLIENT_LOG >>$CLIENT_LOG + rm $TMP_CLIENT_LOG kill $SERVER_PID wait $SERVER_PID + set -e done done if [ $RET -eq 0 ]; then echo -e "\n***\n*** Test Passed\n***" else - cat $CLIENT_LOG echo -e "\n***\n*** Test Failed\n***" fi diff --git a/src/shared_memory_manager.cc b/src/shared_memory_manager.cc index 79d83350ea..0e0e9ea400 100644 --- a/src/shared_memory_manager.cc +++ b/src/shared_memory_manager.cc @@ -363,7 +363,7 @@ SharedMemoryManager::GetMemoryInfo( .c_str()); } // validate byte_size + offset is within memory bounds - size_t total_req_shm = offset + byte_size; + size_t total_req_shm = offset + byte_size - 1; if (total_req_shm > shm_region_end) { return TRITONSERVER_ErrorNew( TRITONSERVER_ERROR_INVALID_ARG, From b880c59ea958814b34ad311cd24b98d3dff860ab Mon Sep 17 00:00:00 2001 From: Katherine Yang Date: Wed, 10 Apr 2024 16:06:09 -0700 Subject: [PATCH 7/7] add check for server kill health --- qa/L0_shared_memory/test.sh | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/qa/L0_shared_memory/test.sh b/qa/L0_shared_memory/test.sh index ea90b53163..7329f28998 100755 --- a/qa/L0_shared_memory/test.sh +++ b/qa/L0_shared_memory/test.sh @@ -84,6 +84,10 @@ for i in \ rm $TMP_CLIENT_LOG kill $SERVER_PID wait $SERVER_PID + if [ $? -ne 0 ]; then + echo -e "\n***\n*** Test Server shut down non-gracefully\n***" + RET=1 + fi set -e done done