diff --git a/qa/L0_shared_memory/shared_memory_test.py b/qa/L0_shared_memory/shared_memory_test.py index 146025cf26..35667aacfa 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_python_backend_shm_region_" + + # 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 the key contains 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..7df618e965 100644 --- a/src/common.h +++ b/src/common.h @@ -45,6 +45,10 @@ constexpr char kContentEncodingHTTPHeader[] = "Content-Encoding"; 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_python_backend_shm_region_"; + 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..80b739ee1b 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 the key contains 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);