Hello friends! How are you all?
- Background Introduction
- Design of Shared Memory and Message Queue
- From the Perspective of the Parent Process
- Creation of Shared Memory and Message Queue
- Sending Messages
- Receiving Messages
- From the Perspective of the Stub Process
- Receiving Messages
- Calling Python Logic with pybind11
- Summary
This article introduces how to use shared memory and subprocess communication in TritonServer, focusing on the encapsulation of shared memory management and the use of the C++ Boost Inprocess module. This functionality is reusable and can be applied to other scenarios.
Background Introduction
TritonServer, as an AI inference service, supports running models or custom logic through a Python backend using subprocesses.
In simple terms, it forks a stub subprocess that runs user-defined Python code via pybind11. When the C++ implemented TritonServer receives a request, it calls the Python script, processes it, and returns the result.
- Fork a stub subprocess to execute user-defined Python code
- TritonServer, as the parent process, needs to communicate with the stub subprocess, such as controlling the stub to stop, sending requests, and controlling commands.
This designs inter-process communication; if done over the network, performance would degrade, so shared memory (shm) is used here.
Design of Shared Memory and Message Queue
For inter-process communication, a message queue is designed, with shared memory serving as the carrier to store the message queue.
Since it is a message queue for transferring data between processes, both processes must obtain the same message queue handle. The producer creates (create) and sends (push) messages, while the consumer loads the message queue and consumes (pop) messages.
From the Perspective of the Parent Process
The parent process creates the message queue, pushes messages to the message queue stub_message_queue, and then waits for the child process to consume the messages. Then, the parent message_queue receives messages returned from the child process.
Creation of Shared Memory and Message Queue
For example, below are several different message queues created by the MessageQueue class for communication between the parent process and the stub process, with the message queue type being bi::managed_external_buffer::handle_t.
// Source code address: https://github.com/triton-inference-server/python_backend/blob/main/src/stub_launcher.cc#L151
// stub_message_queue_ is the queue for the parent to send control information to the stub process, such as sending initialization, stopping, etc.
RETURN_IF_EXCEPTION(
stub_message_queue_ =
MessageQueue<bi::managed_external_buffer::handle_t>::Create(
shm_pool_, shm_message_queue_size_));
// parent_message_queue_ is the queue for the stub process to send information to the parent process
RETURN_IF_EXCEPTION(
parent_message_queue_ =
MessageQueue<bi::managed_external_buffer::handle_t>::Create(
shm_pool_, shm_message_queue_size_));
RETURN_IF_EXCEPTION(
stub_to_parent_mq_ =
MessageQueue<bi::managed_external_buffer::handle_t>::Create(
shm_pool_, shm_message_queue_size_));
RETURN_IF_EXCEPTION(
parent_to_stub_mq_ =
MessageQueue<bi::managed_external_buffer::handle_t>::Create(
shm_pool_, shm_message_queue_size_));
The SharedMemoryManager class is a reusable class that encapsulates the creation and destruction of shared memory, as well as the construction and destruction of objects within shared memory. Here, parameters such as the name of the shared memory, size, growth size, and whether to create it are specified, creating a segment of shm shared memory for process communication.
Additionally, an IPCControlShm object is created in shared memory, which contains the message queue handles and some health status flags for the processes.
std::unique_ptr<SharedMemoryManager> shm_pool_ = std::make_unique<SharedMemoryManager>(
shm_region_name_, shm_default_byte_size_, shm_growth_byte_size_,
true /* create */);
AllocatedSharedMemory<IPCControlShm> current_ipc_control =
shm_pool_->Construct<IPCControlShm>();
The IPCControlShm structure contains all the information needed for process communication, including the health status of both processes and the message queue handles for both processes.
// Control data structure for the communication between the Python stub and the
// main stub.
struct IPCControlShm {
bool stub_health;
bool parent_health;
bool uses_env;
bool decoupled;
bi::interprocess_mutex parent_health_mutex;
bi::interprocess_mutex stub_health_mutex;
bi::managed_external_buffer::handle_t stub_message_queue;
bi::managed_external_buffer::handle_t parent_message_queue;
bi::managed_external_buffer::handle_t stub_to_parent_mq;
bi::managed_external_buffer::handle_t parent_to_stub_mq;
bi::managed_external_buffer::handle_t memory_manager_message_queue;
};
Sending Messages
The following code describes how to send an initialize message from the parent process to the stub process. An IPCMessage type message is constructed, assigned values, and then pushed to the stub_message_queue.
std::unordered_map<std::string, std::string> initialize_map = {
{"model_config", model_config_buffer_.MutableContents()},
{"model_instance_kind", kind_},
{"model_instance_name", model_instance_name_},
{"model_instance_device_id", std::to_string(device_id_)},
{"model_repository", model_repository_path_},
{"model_version", std::to_string(model_version_)},
{"model_name", model_name_}};
std::unique_ptr<IPCMessage> initialize_message =
IPCMessage::Create(shm_pool_, false/* inline_response */);
initialize_message->Command() = PYTHONSTUB_InitializeRequest;
std::unique_ptr<PbMap> pb_map = PbMap::Create(shm_pool_, initialize_map);
bi::managed_external_buffer::handle_t initialize_map_handle =
pb_map->ShmHandle();
initialize_message->Args() = initialize_map_handle;
stub_message_queue_->Push(initialize_message->ShmHandle());
Receiving Messages
How is the initialize_message mentioned in the previous section received and processed?
See the code below; the ReceiveMessageFromStub function receives the message and retrieves the specific content through the shm_pool shared memory method.
bi::managed_external_buffer::handle_t message;
RETURN_IF_ERROR(ReceiveMessageFromStub(message, initialization_timeout_ms));
std::unique_ptr<IPCMessage> initialize_response_message =
IPCMessage::LoadFromSharedMemory(shm_pool_, message);
if (initialize_response_message->Command() != PYTHONSTUB_InitializeResponse) {
return TRITONSERVER_ErrorNew(
TRITONSERVER_ERROR_INTERNAL,
(std::string(
"Received unexpected response from Python backend stub: ") +
model_instance_name_)
.c_str());
}
auto initialize_response =
std::move((shm_pool_->Load<InitializeResponseShm>(
initialize_response_message->Args())))
.data_;
ReceiveMessageFromStub actually retrieves the message from the message queue parent message_queue.
message = parent_message_queue_->Pop(
timeout_miliseconds /* duration ms */, success);
From the Perspective of the Stub Process
From the perspective of the stub process, which is the Python subprocess, it does not create shared memory. Instead, it uses the shared memory created by the parent process and loads the created message queue using the LoadFromSharedMemory method for direct use.
// https://github.com/triton-inference-server/python_backend/blob/main/src/pb_stub.cc#L164
shm_pool_ = std::make_unique<SharedMemoryManager>(
shm_region_name, shm_default_size, shm_growth_size, false /* create */);
AllocatedSharedMemory<IPCControlShm> ipc_control =
shm_pool_->Load<IPCControlShm>(ipc_control_handle);
ipc_control_ = ipc_control.data_.get();
stub_message_queue_ = MessageQueue<bi::managed_external_buffer::handle_t>::
LoadFromSharedMemory(shm_pool_, ipc_control_->stub_message_queue);
As mentioned above, the parent sends an initialize message to the stub, and after the stub process receives it, it will reply with an initialize_response. The child process receives messages from the stub_message_queue using Pop.
Receiving Messages
How does the stub process receive messages? The parent process sends messages by enqueueing them into the stub message_queue, and in the PopMessage function called by the stub process, it retrieves messages from the queue using Pop.
std::unique_ptr<IPCMessage>
Stub::PopMessage()
{
bool success = false;
std::unique_ptr<IPCMessage> ipc_message;
bi::managed_external_buffer::handle_t message;
while (!success) {
message = stub_message_queue_->Pop(1000, success);
}
ipc_message = IPCMessage::LoadFromSharedMemory(shm_pool_, message);
return ipc_message;
}
The core logic of the stub process is to call the RunCommand method, which retrieves messages from the stub message_queue and executes different logic based on the message type.
std::unique_ptr<IPCMessage> ipc_message;
{
// Release the GIL lock when waiting for new message. Without this line, the
// other threads in the user's Python model cannot make progress if they
// give up GIL.
py::gil_scoped_release release;
ipc_message = this->PopMessage();
}
switch (ipc_message->Command()) {
case PYTHONSTUB_CommandType::PYTHONSTUB_InitializeRequest:
//...
}
Calling Python Logic with pybind11
The stub is a C++ implemented process that, upon receiving the initialize_message, will use the pybind11 library.
The pybind11 library can call user-implemented Python logic, such as obtaining Python classes and invoking class functions. Here is a simple code snippet for better understanding:
#include <pybind11/embed.h>
#include <pybind11/numpy.h>
#include <pybind11/stl.h>
namespace py = pybind11;
void
Stub::Initialize(bi::managed_external_buffer::handle_t map_handle)
{
py::module sys = StubSetup();
py::module python_backend_utils =
py::module_::import("triton_python_backend_utils");
py::module c_python_backend_utils =
py::module_::import("c_python_backend_utils");
// ......
c_python_backend_utils.attr("shared_memory") = py::cast(shm_pool_.get());
async_event_loop_ = py::none();
background_futures_ = py::set();
// Create a TritonPythonModel object, which is a user-implemented Python class that needs to implement initialize, execute, finalizer, etc.
py::object TritonPythonModel = sys.attr("TritonPythonModel");
model_instance_ = TritonPythonModel();
std::unordered_map<std::string, std::string> map;
std::unique_ptr<PbMap> pb_map_shm =
PbMap::LoadFromSharedMemory(shm_pool_, map_handle);
// Get the unordered_map representation of the map in shared memory.
map = pb_map_shm->UnorderedMap();
py::dict model_config_params;
for (const auto& pair : map) {
model_config_params[pair.first.c_str()] = pair.second;
}
// Call initialize if exists.
// Call the user-implemented initialize function
if (py::hasattr(model_instance_, "initialize")) {
model_instance_.attr("initialize")(model_config_params);
}
}
Summary
This article, due to space limitations, introduced the IPC communication mechanism of Triton and the initialization process of the stub process, as well as how the stub process executes user-implemented Python logic. The inter-process communication mechanism and pybind11 calling Python logic can be reused in your own projects. This kind of cross-language calling is common, with other examples including C++ calling Lua, JavaScript, etc.
The message queue utilizes the SharedMemoryManager class, which is implemented using the Boost.Interprocess library. For more information on the implementation of the message queue[1] and the use of the Boost.Interprocess library for shared memory[2], you can refer to the source code implementation. Future articles may expand on this topic.
References
[1]
Implementation of the message queue: https://github.com/triton-inference-server/python_backend/blob/main/src/message_queue.h#L67
[2]
Usage of the Boost.Interprocess library for shared memory: https://github.com/triton-inference-server/python_backend/blob/main/src/shm_manager.h#L44