[infer] get_output and save_output recover input without msg_queue_id (#10574)

* check

* fix input_text dy_insert

* check code

Author: ckl117

csrc/gpu/cpp_extensions.cu

@@ -238,11 +238,9 @@ std::vector<paddle::Tensor> GetPaddingOffsetV2(const paddle::Tensor& input_ids,
 
 void SaveOutMmsg(const paddle::Tensor& x,
                  const paddle::Tensor& not_need_stop, // cpu
-                 const paddle::Tensor& msg_queue_id,      // cpu
                  int64_t rank_id);
 
 void GetOutput(const paddle::Tensor& x,
-               const paddle::Tensor& msg_queue_id, // cpu
                int64_t rank_id,
                bool wait_flag);
 

csrc/gpu/get_output.cc

@@ -20,21 +20,33 @@
 #include "paddle/extension.h"
 
 #define MAX_BSZ 512
+// #define GET_OUTPUT_DEBUG
 
 struct msgdata {
     long mtype;
     int mtext[MAX_BSZ + 2];   // stop_flag, bsz, tokens
 };
 
 void GetOutput(const paddle::Tensor& x,
-               const paddle::Tensor& msg_queue_id,
                int64_t rank_id,
                bool wait_flag) {
   if (rank_id > 0) return;
 
   static struct msgdata msg_rcv;
-  int queue_id_val = msg_queue_id.data<int>()[0];
-  static key_t key = ftok("./", queue_id_val);
+  int msg_queue_id = 1;
+  if (const char* inference_msg_queue_id_env_p =
+          std::getenv("INFERENCE_MSG_QUEUE_ID")) {
+      std::string inference_msg_queue_id_env_str(
+          inference_msg_queue_id_env_p);
+      int inference_msg_queue_id_from_env =
+          std::stoi(inference_msg_queue_id_env_str);
+#ifdef GET_OUTPUT_DEBUG
+      std::cout << "Your INFERENCE_MSG_QUEUE_ID is: "
+                << inference_msg_queue_id_from_env << std::endl;
+#endif
+      msg_queue_id = inference_msg_queue_id_from_env;
+  }
+  static key_t key = ftok("./", msg_queue_id);
 
   static int msgid = msgget(key, IPC_CREAT | 0666);
 
@@ -62,7 +74,7 @@ void GetOutput(const paddle::Tensor& x,
 }
 
 PD_BUILD_OP(get_output)
-    .Inputs({"x", "msg_queue_id"})
+    .Inputs({"x"})
     .Attrs({"rank_id: int64_t",
             "wait_flag: bool"})
     .Outputs({"x_out"})

csrc/gpu/save_with_output_msg.cc

@@ -20,6 +20,7 @@
 #include "paddle/extension.h"
 
 #define MAX_BSZ 512
+// #define SAVE_WITH_OUTPUT_DEBUG
 
 struct msgdata {
     long mtype;
@@ -28,16 +29,32 @@ struct msgdata {
 
 void SaveOutMmsg(const paddle::Tensor& x,
                  const paddle::Tensor& not_need_stop, // cpu
-                 const paddle::Tensor& msg_queue_id,      // cpu
                  int64_t rank_id) {
     if (rank_id > 0) return;
     auto x_cpu = x.copy_to(paddle::CPUPlace(), false);
     int64_t *x_data = x_cpu.data<int64_t>();
     auto not_need_stop_data = not_need_stop.data<bool>()[0];
 
     static struct msgdata msg_sed;
-    int queue_id_val = msg_queue_id.data<int>()[0];
-    static key_t key = ftok("./", queue_id_val);
+    int msg_queue_id = 1;
+    if (const char* inference_msg_queue_id_env_p =
+            std::getenv("INFERENCE_MSG_QUEUE_ID")) {
+        std::string inference_msg_queue_id_env_str(
+            inference_msg_queue_id_env_p);
+        int inference_msg_queue_id_from_env =
+            std::stoi(inference_msg_queue_id_env_str);
+        msg_queue_id = inference_msg_queue_id_from_env;
+#ifdef SAVE_WITH_OUTPUT_DEBUG
+        std::cout << "Your INFERENCE_MSG_QUEUE_ID is: "
+                  << inference_msg_queue_id_from_env << std::endl;
+#endif
+    } else {
+#ifdef SAVE_WITH_OUTPUT_DEBUG
+        std::cout << "Failed to got INFERENCE_MSG_QUEUE_ID at env, use default."
+                  << std::endl;
+#endif
+    }
+    static key_t key = ftok("./", msg_queue_id);
     static int msgid = msgget(key, IPC_CREAT | 0666);
 
     msg_sed.mtype = 1;
@@ -54,7 +71,7 @@ void SaveOutMmsg(const paddle::Tensor& x,
 }
 
 PD_BUILD_OP(save_output)
-    .Inputs({"x", "not_need_stop", "msg_queue_id"})
+    .Inputs({"x", "not_need_stop"})
     .Attrs({"rank_id: int64_t"})
     .Outputs({"x_out"})
     .SetInplaceMap({{"x", "x_out"}})

llm/docs/predict/best_practices.md

@@ -44,3 +44,6 @@ PaddleNLP 提供了多种环境变量，用于优化推理性能和资源使用
 
 **自定义算子通用 优化**
 - `DYNAMIC_INFERENCE_MODE`：动态图推理时自定义算子是否采用 pybind 调用方式，默认为 True。
+
+**其他**
+- `INFERENCE_MSG_QUEUE_ID`：多实例消息队列 id，默认为 进程 id。

llm/predict/predictor.py

@@ -848,7 +848,8 @@ def __init__(
 
         self.pre_cache_length = 0
 
-        self.msg_queue_id = os.getpid()
+        msg_queue_id_str = os.getenv("INFERENCE_MSG_QUEUE_ID", str(os.getpid()))
+        os.environ["INFERENCE_MSG_QUEUE_ID"] = msg_queue_id_str
 
         if config.export_precache:
             pre_cache_npy = np.load(config.prefix_path)
@@ -948,7 +949,6 @@ def init_model_inputs(self, config: PredictorArgument):
         )
         self.model_inputs["bad_tokens"] = paddle.to_tensor([-1], dtype="int64")
         self.model_inputs["is_block_step"] = paddle.full(shape=[config.batch_size], fill_value=False, dtype="bool")
-        self.model_inputs["msg_queue_id"] = paddle.full(shape=[1], fill_value=self.msg_queue_id, dtype="int32").cpu()
 
         # bloom model needs src_mask and tgt_mask!
         if "bloom" in self.architectures:
@@ -1185,7 +1185,7 @@ def predict_via_mq(self, input_texts: list[str], return_tokens=False):
 
         read_res_process = mp.Process(
             target=read_res_func,
-            args=[self.model_name_or_path, tensor_queue, result_queue, done_event, self.model_inputs["msg_queue_id"]],
+            args=[self.model_name_or_path, tensor_queue, result_queue, done_event],
         )
         if self.tensor_parallel_rank == 0:
             read_res_process.start()
@@ -1311,7 +1311,7 @@ def insert_task(self, pos, task_id, repeat_num):
         self.model_inputs["stop_flags"][pos] = False
         self.model_inputs["result_id"][pos][0] = task_id
         self.model_inputs["step_idx"][pos, 0] = 1
-        self.model_inputs["pre_ids"][pos][0] = self.input_ids[query_id][-1]
+        self.model_inputs["pre_ids"][pos][0] = np.array(self.input_ids[query_id][-1])
         self.model_inputs["pre_ids"][pos][1:] = -1
         self.model_inputs["not_need_stop"][0] = True
 
@@ -1477,7 +1477,6 @@ def predict_dy_insert(
                     task_queue,
                     result_queue,
                     done_event,
-                    self.model_inputs["msg_queue_id"],
                     len(self.input_ids),
                     detokenize,
                 ],
@@ -1716,7 +1715,7 @@ def predict_via_mq(self, input_texts: list[str], return_tokens=False):
 
         read_res_process = mp.Process(
             target=read_res_func,
-            args=[self.model_name_or_path, tensor_queue, result_queue, done_event, self.model_inputs["msg_queue_id"]],
+            args=[self.model_name_or_path, tensor_queue, result_queue, done_event],
         )
         if self.tensor_parallel_rank == 0:
             read_res_process.start()

llm/server/server/server/engine/engine.py

@@ -24,8 +24,7 @@
 
 import numpy as np
 from server.engine.resource_manager import ResourceManager
-from server.engine.task_queue_manager import (TaskQueueManager,
-                                              launch_queue_service)
+from server.engine.task_queue_manager import TaskQueueManager, launch_queue_service
 from server.engine.token_processor import TokenProcessor, WarmUpTokenProcessor
 from server.utils import model_server_logger
 
@@ -34,10 +33,11 @@ class Engine(object):
     """
     Engine Class
     """
+
     def __init__(self, cfg, token_processor):
         self.cfg = cfg
         # Master node only
-        if self.cfg.nnode == 1 or self.cfg.host_ip == os.getenv('POD_0_IP', '127.0.0.1'):
+        if self.cfg.nnode == 1 or self.cfg.host_ip == os.getenv("POD_0_IP", "127.0.0.1"):
             self.queue_service = self._start_tasks_queue_service()
         self.tasks_queue = TaskQueueManager(mp_num=self.cfg.mp_num, port=self.cfg.infer_port)
         self.resource_manager = ResourceManager(self.cfg)
@@ -53,8 +53,10 @@ def start(self):
         initialize engine and start sub services
         """
         assert not self.is_started, "The engine is already started.!"
+        msg_queue_id_str = os.getenv("INFERENCE_MSG_QUEUE_ID", str(os.getpid()))
+        os.environ["INFERENCE_MSG_QUEUE_ID"] = msg_queue_id_str
         start_time = time.time()
-        
+
         self.token_processor.tasks_queue = self.tasks_queue
         self.infer_proc = self._start_infer_service()
         model_server_logger.info("Waiting infer processes ready...")
@@ -80,17 +82,18 @@ def warmup(self):
         """
         # get eos_token_id
         from server.data.processor import DataProcessor
+
         eos_token_ids = DataProcessor().get_eos_tokens()
 
-       # construct test tasks
+        # construct test tasks
         res_task = []
         for j in range(2 * self.cfg.max_batch_size):
             data = {
                 "input_ids": [5],
                 "req_id": j,
                 "max_dec_len": self.cfg.dec_len_limit,
                 "min_dec_len": int(self.cfg.dec_len_limit * 0.5) + 1,
-                "eos_token_ids": eos_token_ids
+                "eos_token_ids": eos_token_ids,
             }
             res_task.append(data)
         for j in range(2 * self.cfg.max_prefill_batch):
@@ -99,7 +102,7 @@ def warmup(self):
                 "req_id": j + 2 * self.cfg.max_batch_size,
                 "max_dec_len": 1,
                 "min_dec_len": 1,
-                "eos_token_ids": eos_token_ids
+                "eos_token_ids": eos_token_ids,
             }
             res_task.append(data)
 
@@ -130,8 +133,9 @@ def insert_tasks(self, tasks):
 
         available_batch = np.sum(self.resource_manager.stop_flags)
         if len(tasks) > available_batch:
-            model_server_logger.error("Inserting batch:{} exceeds the available batch:{}.".format(
-                len(tasks), available_batch))
+            model_server_logger.error(
+                "Inserting batch:{} exceeds the available batch:{}.".format(len(tasks), available_batch)
+            )
             model_server_logger.error("The exceeded part will be ignored!")
             tasks = tasks[:available_batch]
 
@@ -140,21 +144,23 @@ def insert_tasks(self, tasks):
             input_token_num = len(tasks[i]["input_ids"])
             if input_token_num >= self.cfg.max_seq_len - 1:
                 model_server_logger.warning(f"{req_id}: Input length:{input_token_num}, exceed the limits.")
-                tasks[i]["input_ids"] = tasks[i]["input_ids"][:self.cfg.max_seq_len - 1]
+                tasks[i]["input_ids"] = tasks[i]["input_ids"][: self.cfg.max_seq_len - 1]
             if "seq_len" in tasks[i] and "max_dec_len" not in tasks[i]:
                 tasks[i]["max_dec_len"] = tasks[i]["seq_len"]
 
             # max_dec_len + input_token_num > MAX_SEQ_LEN
             if input_token_num + tasks[i]["max_dec_len"] > self.cfg.max_seq_len:
                 tasks[i]["max_dec_len"] = self.cfg.max_seq_len - input_token_num
-                model_server_logger.warning("Force max_dec_len to be {} for req_id={}.".format(
-                    tasks[i]["max_dec_len"], tasks[i]["req_id"]))
+                model_server_logger.warning(
+                    "Force max_dec_len to be {} for req_id={}.".format(tasks[i]["max_dec_len"], tasks[i]["req_id"])
+                )
 
             # min_dec_len + input_token_num > MAX_SEQ_LEN
             if input_token_num + tasks[i]["min_dec_len"] > self.cfg.max_seq_len:
                 tasks[i]["min_dec_len"] = self.cfg.max_seq_len - input_token_num
-                model_server_logger.warning("Force min_dec_len to be {} for req_id={}.".format(
-                    tasks[i]["min_dec_len"], tasks[i]["req_id"]))
+                model_server_logger.warning(
+                    "Force min_dec_len to be {} for req_id={}.".format(tasks[i]["min_dec_len"], tasks[i]["req_id"])
+                )
 
         tasks = self.resource_manager.allocate_resources_for_new_tasks(tasks)
         if not tasks:
@@ -292,9 +298,7 @@ def _init_engine_flags(self):
         self.shm_flag_ready = shared_memory.SharedMemory(
             create=True, size=flag_array.nbytes, name=self.cfg.get_unique_name("shm_flag_infer_ready")
         )
-        self.flag_ready_array = np.ndarray(
-            flag_array.shape, dtype=flag_array.dtype, buffer=self.shm_flag_ready.buf
-        )
+        self.flag_ready_array = np.ndarray(flag_array.shape, dtype=flag_array.dtype, buffer=self.shm_flag_ready.buf)
         self.flag_ready_array[:] = 0
 
         # broadcast flag for engine
@@ -324,19 +328,22 @@ def _init_engine_flags(self):
             tmp = shared_memory.SharedMemory(
                 create=False,
                 size=has_block_step_flag_array.nbytes,
-                name=self.cfg.get_unique_name("shm_flag_has_block_step"))
+                name=self.cfg.get_unique_name("shm_flag_has_block_step"),
+            )
             tmp.close()
             tmp.unlink()
         except:
             pass
         self.shm_flag_has_block_step = shared_memory.SharedMemory(
             create=True,
             size=has_block_step_flag_array.nbytes,
-            name=self.cfg.get_unique_name("shm_flag_has_block_step"))
+            name=self.cfg.get_unique_name("shm_flag_has_block_step"),
+        )
         self.flag_has_block_step_array = np.ndarray(
             has_block_step_flag_array.shape,
             dtype=has_block_step_flag_array.dtype,
-            buffer=self.shm_flag_has_block_step.buf)
+            buffer=self.shm_flag_has_block_step.buf,
+        )
         self.flag_has_block_step_array[:] = 0
 
     def _exit_sub_services(self):
@@ -362,8 +369,9 @@ def _start_tasks_queue_service(self):
         if p.is_alive():
             model_server_logger.info("start tasks queue service successfully")
         else:
-            error_msg = "Failed to start tasks queue service, please check " \
-                        "the log/task_queue_manager.log for details"
+            error_msg = (
+                "Failed to start tasks queue service, please check " "the log/task_queue_manager.log for details"
+            )
             model_server_logger.info(error_msg)
             raise Exception(error_msg)
         return p
@@ -380,14 +388,16 @@ def _start_gpu_infer_service(self):
         pd_cmd = "python3 -m paddle.distributed.launch "
         py_script = os.path.join(current_dir_path, "infer.py")
 
-        arguments = (f" --nnodes {str(self.cfg.nnode)}"
-                    f" --devices {self.cfg.device_ids} {py_script} --model_dir {self.cfg.model_dir}"
-                    f" --max_batch_size {self.cfg.max_batch_size} --max_seq_len {self.cfg.max_seq_len}"
-                    f" --max_dec_len {self.cfg.max_dec_len}"
-                    f" --max_block_num {self.cfg.total_block_num} --block_size {self.cfg.block_size}"
-                    f" --use_cache_kv_int8 {self.cfg.use_cache_kv_int8}"
-                    f" --enc_dec_block_num {self.cfg.enc_dec_block_num}"
-                    f" --block_ratio {self.cfg.block_ratio} --dtype {self.cfg.dtype}")
+        arguments = (
+            f" --nnodes {str(self.cfg.nnode)}"
+            f" --devices {self.cfg.device_ids} {py_script} --model_dir {self.cfg.model_dir}"
+            f" --max_batch_size {self.cfg.max_batch_size} --max_seq_len {self.cfg.max_seq_len}"
+            f" --max_dec_len {self.cfg.max_dec_len}"
+            f" --max_block_num {self.cfg.total_block_num} --block_size {self.cfg.block_size}"
+            f" --use_cache_kv_int8 {self.cfg.use_cache_kv_int8}"
+            f" --enc_dec_block_num {self.cfg.enc_dec_block_num}"
+            f" --block_ratio {self.cfg.block_ratio} --dtype {self.cfg.dtype}"
+        )
         if self.cfg.nnode > 1:
             pd_cmd = pd_cmd + f" --ips {self.cfg.ips}"
         pd_cmd = pd_cmd + arguments + " >log/launch_infer.log 2>&1"

llm/server/server/server/engine/infer.py

@@ -279,7 +279,6 @@ def init_inputs(self):
         self.share_inputs["input_ids"] = paddle.full(
             shape=[self.args.max_batch_size, self.args.max_seq_len], fill_value=self.pad_token_id, dtype="int64"
         )
-        self.share_inputs["msg_queue_id"] = paddle.full(shape=[1], fill_value=1, dtype="int32").cpu()
         self.share_inputs["top_p"] = paddle.full(
             shape=[self.args.max_batch_size, 1], fill_value=self.top_p, dtype="float32"
         )
@@ -743,7 +742,7 @@ def _init_predictor(self):
             config.set_xpu_device_id(device_id)
             xpu_config = paddle.inference.XpuConfig()
             xpu_config.device_id = device_id
-            xpu_config.l3_size = 0 
+            xpu_config.l3_size = 0
             xpu_config.l3_autotune_size = 0
             config.set_xpu_config(xpu_config)
             config.switch_ir_optim(True)

llm/server/server/server/engine/token_processor.py

@@ -47,7 +47,6 @@ def __init__(self, cfg):
         self.tokens_counter = Counter()
 
         self.is_speculate_decoding = self.cfg.get_speculate_config().speculate_method != "None"
-        self.msg_queue_id = paddle.full(shape=[1], fill_value=1, dtype="int32")
         if self.is_speculate_decoding:
             self.output_tokens = paddle.full(
                 shape=[SPECULATE_MAX_BSZ * MAX_DRAFT_TOKENS + SPECULATE_MAX_BSZ + 2, 1], fill_value=2, dtype="int64"
@@ -97,7 +96,7 @@ def process_sampling_results(self):
                 if self.is_speculate_decoding:
                     speculate_get_output(self.output_tokens, rank_id, is_blocking)
                 else:
-                    get_output(self.output_tokens, self.msg_queue_id, rank_id, is_blocking)
+                    get_output(self.output_tokens, rank_id, is_blocking)
 
                 if self.output_tokens[0, 0] == -2:
                     continue
@@ -281,7 +280,7 @@ def process_sampling_results(self):
                 if self.is_speculate_decoding:
                     speculate_get_output(self.output_tokens, rank_id, self._is_blocking)
                 else:
-                    get_output(self.output_tokens, self.msg_queue_id, rank_id, self._is_blocking)
+                    get_output(self.output_tokens, rank_id, self._is_blocking)
 
                 if self.output_tokens[0, 0] == -2:
                     continue