refactor: integrate prepare_shm function into the worker_server class.

Author: Super User

xllm/core/distributed_runtime/dist_manager.cpp

@@ -117,27 +117,6 @@ std::unique_ptr<CommChannel> create_channel(const std::string& worker_addrs,
 
   return channel;
 }
-
-void prepare_shm(
-    int dp_local_tp_size,
-    int rank,
-    std::unique_ptr<ForwardSharedMemoryManager>& input_shm_manager,
-    std::unique_ptr<ForwardSharedMemoryManager>& output_shm_manager) {
-  bool is_creator;
-  int32_t dp_group = rank / dp_local_tp_size;
-
-  string name = ForwardSharedMemoryManager::create_unique_name(
-      dp_group, FORWARD_RAW_INPUT_TYPE, rank);
-  input_shm_manager = std::make_unique<ForwardSharedMemoryManager>(
-      name, PB_INPUT_SHM_SIZE, is_creator, FORWARD_RAW_INPUT_TYPE);
-  LOG(INFO) << "Create input shared memory manager with name: " << name;
-
-  name = ForwardSharedMemoryManager::create_unique_name(
-      dp_group, FORWARD_RAW_OUTPUT_TYPE, rank);
-  output_shm_manager = std::make_unique<ForwardSharedMemoryManager>(
-      name, PB_OUTPUT_SHM_SIZE, is_creator, FORWARD_RAW_OUTPUT_TYPE);
-  LOG(INFO) << "Create output shared memory manager with name: " << name;
-}
 }  // namespace
 
 void DistManager::setup_multi_node_workers(
@@ -201,24 +180,15 @@ void DistManager::setup_multi_node_workers(
     bool use_spawn_worker = options.enable_offline_inference() && i > 0;
     ParallelArgs parallel_args(rank, world_size, dp_size, nullptr, ep_size);
 
-    std::unique_ptr<ForwardSharedMemoryManager> input_shm_manager = nullptr;
-    std::unique_ptr<ForwardSharedMemoryManager> output_shm_manager = nullptr;
-    if (options.is_local() && FLAGS_enable_shm) {
-      prepare_shm(
-          dp_local_tp_size, rank, input_shm_manager, output_shm_manager);
-    }
-    servers_.emplace_back(
-        std::make_unique<WorkerServer>(i,
-                                       master_node_addr,
-                                       // done,
-                                       dones[i],
-                                       parallel_args,
-                                       devices[i],
-                                       worker_server_options,
-                                       worker_type,
-                                       use_spawn_worker,
-                                       std::move(input_shm_manager),
-                                       std::move(output_shm_manager)));
+    servers_.emplace_back(std::make_unique<WorkerServer>(i,
+                                                         master_node_addr,
+                                                         // done,
+                                                         dones[i],
+                                                         parallel_args,
+                                                         devices[i],
+                                                         worker_server_options,
+                                                         worker_type,
+                                                         use_spawn_worker));
   }
 
   // Master node need to wait all workers done

xllm/core/distributed_runtime/worker_server.cpp

@@ -168,17 +168,38 @@ void WorkerServer::create_spawn_server(int local_rank,
   done.store(true);
 }
 
-WorkerServer::WorkerServer(
-    int local_worker_idx,
-    const std::string& master_node_addr,
-    std::atomic<bool>& done,
+void WorkerServer::prepare_shm(
     const ParallelArgs& parallel_args,
-    const torch::Device& d,
     const runtime::Options& options,
-    WorkerType worker_type,
-    bool use_spawn_worker,
-    std::unique_ptr<ForwardSharedMemoryManager> input_shm_manager,
-    std::unique_ptr<ForwardSharedMemoryManager> output_shm_manager) {
+    std::unique_ptr<ForwardSharedMemoryManager>& input_shm_manager,
+    std::unique_ptr<ForwardSharedMemoryManager>& output_shm_manager) {
+  if (options.is_local() && FLAGS_enable_shm) {
+    bool is_creator;
+    int dp_local_tp_size = parallel_args.world_size() / parallel_args.dp_size();
+    int dp_group = parallel_args.rank() / dp_local_tp_size;
+
+    string name = ForwardSharedMemoryManager::create_unique_name(
+        dp_group, FORWARD_RAW_INPUT_TYPE, parallel_args.rank());
+    input_shm_manager = std::make_unique<ForwardSharedMemoryManager>(
+        name, PB_INPUT_SHM_SIZE, is_creator, FORWARD_RAW_INPUT_TYPE);
+    LOG(INFO) << "Create input shared memory manager with name: " << name;
+
+    name = ForwardSharedMemoryManager::create_unique_name(
+        dp_group, FORWARD_RAW_OUTPUT_TYPE, parallel_args.rank());
+    output_shm_manager = std::make_unique<ForwardSharedMemoryManager>(
+        name, PB_OUTPUT_SHM_SIZE, is_creator, FORWARD_RAW_OUTPUT_TYPE);
+    LOG(INFO) << "Create output shared memory manager with name: " << name;
+  }
+}
+
+WorkerServer::WorkerServer(int local_worker_idx,
+                           const std::string& master_node_addr,
+                           std::atomic<bool>& done,
+                           const ParallelArgs& parallel_args,
+                           const torch::Device& d,
+                           const runtime::Options& options,
+                           WorkerType worker_type,
+                           bool use_spawn_worker) {
   if (worker_type == WorkerType::LLM || worker_type == WorkerType::ELM) {
     if (use_spawn_worker) {
       // start worker in a spawn process(for offline inference worker.)
@@ -187,6 +208,9 @@ WorkerServer::WorkerServer(
       return;
     }
 
+    std::unique_ptr<ForwardSharedMemoryManager> input_shm_manager = nullptr;
+    std::unique_ptr<ForwardSharedMemoryManager> output_shm_manager = nullptr;
+    prepare_shm(parallel_args, options, input_shm_manager, output_shm_manager);
     // start worker in a thread.
     worker_thread_ =
         std::make_unique<std::thread>(&WorkerServer::create_server,

xllm/core/distributed_runtime/worker_server.h

@@ -37,17 +37,14 @@ namespace xllm {
 
 class WorkerServer {
  public:
-  WorkerServer(
-      int local_worker_idx,
-      const std::string& master_node_addr,
-      std::atomic<bool>& done,
-      const ParallelArgs& parallel_args,
-      const torch::Device& d,
-      const runtime::Options& options,
-      WorkerType worker_type,
-      bool use_spawn_worker = false,
-      std::unique_ptr<ForwardSharedMemoryManager> input_shm_manager = nullptr,
-      std::unique_ptr<ForwardSharedMemoryManager> output_shm_manager = nullptr);
+  WorkerServer(int local_worker_idx,
+               const std::string& master_node_addr,
+               std::atomic<bool>& done,
+               const ParallelArgs& parallel_args,
+               const torch::Device& d,
+               const runtime::Options& options,
+               WorkerType worker_type,
+               bool use_spawn_worker = false);
 
   virtual ~WorkerServer();
 
@@ -82,6 +79,12 @@ class WorkerServer {
                         proto::AddressInfo& addr_info,
                         proto::CommUniqueIdList& uids);
 
+  void prepare_shm(
+      const ParallelArgs& parallel_args,
+      const runtime::Options& options,
+      std::unique_ptr<ForwardSharedMemoryManager>& input_shm_manager,
+      std::unique_ptr<ForwardSharedMemoryManager>& output_shm_manager);
+
  private:
   std::unique_ptr<std::thread> worker_thread_;