update queue_lens on generation ends.

fastchat/serve/base_model_worker.py

@@ -190,6 +190,7 @@ def acquire_worker_semaphore():
 def create_background_tasks():
     background_tasks = BackgroundTasks()
     background_tasks.add_task(release_worker_semaphore)
+    background_tasks.add_task(worker.send_heart_beat)  # refrech queue_lens
     return background_tasks
 
 
@@ -208,6 +209,7 @@ async def api_generate(request: Request):
     await acquire_worker_semaphore()
     output = await asyncio.to_thread(worker.generate_gate, params)
     release_worker_semaphore()
+    worker.send_heart_beat()  # refrech queue_lens
     return JSONResponse(output)
 
 

fastchat/serve/controller.py

@@ -196,7 +196,10 @@ def get_worker_address(self, model_name: str):
                     worker_qlen.append(w_info.queue_length / w_info.speed)
             if len(worker_names) == 0:
                 return ""
-            min_index = np.argmin(worker_qlen)
+            # min_index = np.argmin(worker_qlen)
+            min_index = np.random.choice(
+                np.where(worker_qlen == np.min(worker_qlen))[0]
+            )  # tie break with random choice
             w_name = worker_names[min_index]
             self.worker_info[w_name].queue_length += 1
             logger.info(

fastchat/serve/vllm_worker.py

@@ -191,6 +191,7 @@ async def abort_request() -> None:
 
     background_tasks = BackgroundTasks()
     background_tasks.add_task(release_worker_semaphore)
+    background_tasks.add_task(worker.send_heart_beat)  # refrech queue_lens
     background_tasks.add_task(abort_request)
     return background_tasks
 
@@ -216,6 +217,7 @@ async def api_generate(request: Request):
     params["request"] = request
     output = await worker.generate(params)
     release_worker_semaphore()
+    worker.send_heart_beat()  # refrech queue_lens
     await engine.abort(request_id)
     return JSONResponse(output)