fix

Author: hiworldwzj

lightllm/server/router/batch.py

@@ -69,23 +69,23 @@ def is_clear(self):
     def merge(self, mini_batch: "Batch"):
         if mini_batch is None:
             return
-        
+
         for _req in mini_batch.reqs:
             self.reqs.append(_req)
         self.id_to_reqs = {req.request_id: req for req in self.reqs}
         return
-    
+
     @staticmethod
-    def merge_two_batch(batch1: "Batch", batch2: "Batch"):
+    def merge_two_batch(batch1: "Batch", batch2: "Batch") -> "Batch":
         if batch1 is None and batch2 is None:
             return None
-        
+
         not_none_batch = batch1 if batch1 is not None else batch2
 
         merge_batch = Batch(-1, [], not_none_batch.dp_size_in_node)
         merge_batch.merge(batch1)
         merge_batch.merge(batch2)
-        return
+        return merge_batch
 
     def __repr__(self):
         return f"batch_id={self.batch_id}, " f"reqs={self.reqs}, "

lightllm/server/router/manager.py

@@ -295,9 +295,9 @@ def generate_new_batch(self):
 
         # 调度的时候需要考虑当前运行的batch，和调度了但是暂时还没有推理的部分请求。
         new_batch = self.req_queue.generate_new_batch(
-            Batch.merge(self.running_batch, self.schedule_new_batch), limit_router_queue_length
+            Batch.merge_two_batch(self.running_batch, self.schedule_new_batch), limit_router_queue_length
         )
-        self.schedule_new_batch = Batch.merge(self.schedule_new_batch, new_batch)
+        self.schedule_new_batch = Batch.merge_two_batch(self.schedule_new_batch, new_batch)
         return
 
     async def _step(self):

lightllm/server/router/stats.py

@@ -4,8 +4,8 @@
 
 logger = init_logger(__name__)
 
-class Stats:
 
+class Stats:
     def __init__(self, log_status, log_stats_interval) -> None:
         self.log_stats = log_status
         self.log_stats_interval = log_stats_interval
@@ -14,16 +14,16 @@ def __init__(self, log_status, log_stats_interval) -> None:
         self.output_tokens = 0
         self.prompt_tokens = 0
         return
-    
+
     def count_prompt_tokens(self, run_batch: Batch):
-        if self.log_stats:
+        if self.log_stats and run_batch is not None:
             tokens = run_batch.input_tokens()
             self.prompt_tokens += tokens
             self.all_tokens += tokens
         return
-    
+
     def count_output_tokens(self, run_batch: Batch):
-        if self.log_stats:
+        if self.log_stats and run_batch is not None:
             tokens = len(run_batch.reqs)
             self.output_tokens += tokens
             self.all_tokens += tokens
@@ -35,13 +35,13 @@ def print_stats(self):
 
         now = time.time()
         if now - self.last_log_time > self.log_stats_interval:
-            logger.debug(f"Avg tokens(prompt+generate) throughput: {self.all_tokens/(now-self.last_log_time):8.3f} tokens/s\n"
-                         f"Avg prompt tokens throughput:           {self.prompt_tokens/(now-self.last_log_time):8.3f} tokens/s\n"
-                         f"Avg generate tokens throughput:         {self.output_tokens/(now-self.last_log_time):8.3f} tokens/s")
+            logger.debug(
+                f"Avg tokens(prompt+generate) throughput: {self.all_tokens/(now-self.last_log_time):8.3f} tokens/s\n"
+                f"Avg prompt tokens throughput:           {self.prompt_tokens/(now-self.last_log_time):8.3f} tokens/s\n"
+                f"Avg generate tokens throughput:         {self.output_tokens/(now-self.last_log_time):8.3f} tokens/s"
+            )
             self.all_tokens = 0
             self.output_tokens = 0
             self.prompt_tokens = 0
             self.last_log_time = now
         return
-
-