From dff38573def2a4d14d79e326ff8c7496f8ca7e2d Mon Sep 17 00:00:00 2001 From: Sai Kiran Polisetty Date: Sun, 29 Jun 2025 19:35:55 +0530 Subject: [PATCH 1/5] Add validation --- qa/L0_shared_memory/shared_memory_test.py | 18 ++++++++++++++++++ qa/L0_shared_memory/test.sh | 1 + src/common.h | 3 +++ src/shared_memory_manager.cc | 12 +++++++++++- 4 files changed, 33 insertions(+), 1 deletion(-) diff --git a/qa/L0_shared_memory/shared_memory_test.py b/qa/L0_shared_memory/shared_memory_test.py index 146025cf26..6bae7b6aa8 100755 --- a/qa/L0_shared_memory/shared_memory_test.py +++ b/qa/L0_shared_memory/shared_memory_test.py @@ -509,6 +509,24 @@ def test_python_client_leak(self): "client memory usage is increasing", ) + def test_register_reserved_names(self): + """ + Test that registration fails if attempting to use a reserved + prefix for the shm key. + """ + # This matches kTritonSharedMemoryRegionPrefix in the server code. + reserved_prefix = "triton_" + + # The shared memory key cannot start with the reserved prefix. + shm_name = "my_test_shm_name" + shm_key = f"{reserved_prefix}_my_test_shm_key" + + with self.assertRaisesRegex( + utils.InferenceServerException, + f"cannot register shared memory region '{shm_name}' with key '{shm_key}' as it uses the reserved prefix '{reserved_prefix}'", + ) as e: + self.triton_client.register_system_shared_memory(shm_name, shm_key, 10000) + def callback(user_data, result, error): if error: diff --git a/qa/L0_shared_memory/test.sh b/qa/L0_shared_memory/test.sh index 97b81397ec..55cdb80951 100755 --- a/qa/L0_shared_memory/test.sh +++ b/qa/L0_shared_memory/test.sh @@ -55,6 +55,7 @@ for i in \ test_infer_byte_size_out_of_bound \ test_infer_integer_overflow \ test_register_out_of_bound \ + test_register_reserved_names \ test_python_client_leak; do for client_type in http grpc; do SERVER_ARGS="--model-repository=`pwd`/models --log-verbose=1 ${SERVER_ARGS_EXTRA}" diff --git a/src/common.h b/src/common.h index fe7e3be4cd..7d92851387 100644 --- a/src/common.h +++ b/src/common.h @@ -45,6 +45,9 @@ constexpr char kContentEncodingHTTPHeader[] = "Content-Encoding"; constexpr char kContentTypeHeader[] = "Content-Type"; constexpr char kContentLengthHeader[] = "Content-Length"; +// This prefix is reserved for shared memory regions created internally by Triton +constexpr char kTritonSharedMemoryRegionPrefix[] = "triton_"; + constexpr int MAX_GRPC_MESSAGE_SIZE = INT32_MAX; /// The value for a dimension in a shape that indicates that that diff --git a/src/shared_memory_manager.cc b/src/shared_memory_manager.cc index 17b0e43da0..d1aab9f5e8 100644 --- a/src/shared_memory_manager.cc +++ b/src/shared_memory_manager.cc @@ -354,6 +354,17 @@ SharedMemoryManager::RegisterSystemSharedMemory( const std::string& name, const std::string& shm_key, const size_t offset, const size_t byte_size) { + // Check if the shared memory key starts with the reserved prefix + if (shm_key.rfind(kTritonSharedMemoryRegionPrefix, 0) == 0) { + return TRITONSERVER_ErrorNew( + TRITONSERVER_ERROR_INVALID_ARG, + std::string( + "cannot register shared memory region '" + name + "' with key '" + + shm_key + "' as it uses the reserved prefix '" + + kTritonSharedMemoryRegionPrefix + "'") + .c_str()); + } + std::lock_guard lock(mu_); if (shared_memory_map_.find(name) != shared_memory_map_.end()) { @@ -670,7 +681,6 @@ SharedMemoryManager::UnregisterAll(TRITONSERVER_MemoryType memory_type) ++next_it; if (it->second->kind_ == TRITONSERVER_MEMORY_GPU) { TRITONSERVER_Error* err = UnregisterHelper(it->first, memory_type); - ; if (err != nullptr) { unregister_fails.push_back(it->first); LOG_VERBOSE(1) << TRITONSERVER_ErrorMessage(err); From 92ecad0acd54f288b00e6a8909dd833d8010ac76 Mon Sep 17 00:00:00 2001 From: Sai Kiran Polisetty Date: Sun, 29 Jun 2025 19:45:36 +0530 Subject: [PATCH 2/5] Update --- src/common.h | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/common.h b/src/common.h index 7d92851387..bf10930ecb 100644 --- a/src/common.h +++ b/src/common.h @@ -45,7 +45,7 @@ constexpr char kContentEncodingHTTPHeader[] = "Content-Encoding"; constexpr char kContentTypeHeader[] = "Content-Type"; constexpr char kContentLengthHeader[] = "Content-Length"; -// This prefix is reserved for shared memory regions created internally by Triton +// This prefix is reserved for shm regions created internally by Triton constexpr char kTritonSharedMemoryRegionPrefix[] = "triton_"; constexpr int MAX_GRPC_MESSAGE_SIZE = INT32_MAX; From 62247a9b952fe9e74a58748306141d156327bac2 Mon Sep 17 00:00:00 2001 From: Sai Kiran Polisetty Date: Wed, 2 Jul 2025 10:53:25 +0530 Subject: [PATCH 3/5] Update src/shared_memory_manager.cc Co-authored-by: Yingge He <157551214+yinggeh@users.noreply.github.com> --- src/shared_memory_manager.cc | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/shared_memory_manager.cc b/src/shared_memory_manager.cc index d1aab9f5e8..80b739ee1b 100644 --- a/src/shared_memory_manager.cc +++ b/src/shared_memory_manager.cc @@ -360,7 +360,7 @@ SharedMemoryManager::RegisterSystemSharedMemory( TRITONSERVER_ERROR_INVALID_ARG, std::string( "cannot register shared memory region '" + name + "' with key '" + - shm_key + "' as it uses the reserved prefix '" + + shm_key + "' as the key contains the reserved prefix '" + kTritonSharedMemoryRegionPrefix + "'") .c_str()); } From 61045aabf029eba8f287056143078c966c01d0ef Mon Sep 17 00:00:00 2001 From: Sai Kiran Polisetty Date: Wed, 2 Jul 2025 10:53:32 +0530 Subject: [PATCH 4/5] Update qa/L0_shared_memory/shared_memory_test.py Co-authored-by: Yingge He <157551214+yinggeh@users.noreply.github.com> --- qa/L0_shared_memory/shared_memory_test.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/qa/L0_shared_memory/shared_memory_test.py b/qa/L0_shared_memory/shared_memory_test.py index 6bae7b6aa8..3d7c7ed4de 100755 --- a/qa/L0_shared_memory/shared_memory_test.py +++ b/qa/L0_shared_memory/shared_memory_test.py @@ -523,7 +523,7 @@ def test_register_reserved_names(self): with self.assertRaisesRegex( utils.InferenceServerException, - f"cannot register shared memory region '{shm_name}' with key '{shm_key}' as it uses the reserved prefix '{reserved_prefix}'", + f"cannot register shared memory region '{shm_name}' with key '{shm_key}' as the key contains the reserved prefix '{reserved_prefix}'", ) as e: self.triton_client.register_system_shared_memory(shm_name, shm_key, 10000) From 01a425eee1c6097a484b5c9315781b21dd91de74 Mon Sep 17 00:00:00 2001 From: Sai Kiran Polisetty Date: Wed, 2 Jul 2025 15:16:44 +0530 Subject: [PATCH 5/5] Update --- qa/L0_shared_memory/shared_memory_test.py | 2 +- src/common.h | 3 ++- 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/qa/L0_shared_memory/shared_memory_test.py b/qa/L0_shared_memory/shared_memory_test.py index 3d7c7ed4de..35667aacfa 100755 --- a/qa/L0_shared_memory/shared_memory_test.py +++ b/qa/L0_shared_memory/shared_memory_test.py @@ -515,7 +515,7 @@ def test_register_reserved_names(self): prefix for the shm key. """ # This matches kTritonSharedMemoryRegionPrefix in the server code. - reserved_prefix = "triton_" + reserved_prefix = "triton_python_backend_shm_region_" # The shared memory key cannot start with the reserved prefix. shm_name = "my_test_shm_name" diff --git a/src/common.h b/src/common.h index bf10930ecb..7df618e965 100644 --- a/src/common.h +++ b/src/common.h @@ -46,7 +46,8 @@ constexpr char kContentTypeHeader[] = "Content-Type"; constexpr char kContentLengthHeader[] = "Content-Length"; // This prefix is reserved for shm regions created internally by Triton -constexpr char kTritonSharedMemoryRegionPrefix[] = "triton_"; +constexpr char kTritonSharedMemoryRegionPrefix[] = + "triton_python_backend_shm_region_"; constexpr int MAX_GRPC_MESSAGE_SIZE = INT32_MAX;