pich

Author: shell-nlp

gpt_server/model_worker/base/model_worker_base.py

@@ -203,7 +203,7 @@ def get_worker(
         controller_addr: str = "http://localhost:21001",
         worker_id: str = str(uuid.uuid4())[:8],
         model_names: List[str] = [""],
-        limit_worker_concurrency: int = 512,
+        limit_worker_concurrency: int = 1024,
         conv_template: str = None,  # type: ignore
     ):
         worker = cls(
@@ -251,6 +251,8 @@ def run(cls):
         parser.add_argument("--log_level", type=str, default="WARNING")
         # task_type
         parser.add_argument("--task_type", type=str, default="auto")
+        # limit_worker_concurrency
+        parser.add_argument("--limit_worker_concurrency", type=int, default=1024)
         args = parser.parse_args()
         os.environ["num_gpus"] = str(args.num_gpus)
         if args.backend == "vllm":
@@ -279,6 +281,7 @@ def run(cls):
         os.environ["dtype"] = args.dtype
         os.environ["log_level"] = args.log_level
         os.environ["task_type"] = args.task_type
+        limit_worker_concurrency = int(args.limit_worker_concurrency)
         logger.remove(0)
         log_level = os.getenv("log_level", "WARNING")
         logger.add(sys.stderr, level=log_level)
@@ -291,6 +294,8 @@ def run(cls):
         os.environ["WORKER_PORT"] = str(port)
         os.environ["WORKER_HOST"] = str(local_ip)
         worker_addr = f"http://{host}:{port}"
+        model_names = args.model_names
+        logger.info(f"{model_names[0]} args: \n{args}")
 
         @app.on_event("startup")
         async def startup():
@@ -299,9 +304,10 @@ async def startup():
             worker = cls.get_worker(
                 worker_addr=worker_addr,
                 model_path=args.model_name_or_path,
-                model_names=args.model_names,
+                model_names=model_names,
                 conv_template="chatglm3",  # TODO 默认是chatglm3用于统一处理
                 controller_addr=controller_address,
+                limit_worker_concurrency=limit_worker_concurrency,
             )
 
         uvicorn.run(app, host=host, port=port)

gpt_server/script/config_example.yaml

@@ -22,6 +22,7 @@ model_worker_args:
   host: 0.0.0.0
   controller_address: http://localhost:21001 # # 将模型注册到 控制器的 地址
   log_level: WARNING # DEBUG INFO WARNING ERROR
+  limit_worker_concurrency: 1024 # worker的最大并发数,默认为 1024
 
 models:
 - qwen:

gpt_server/serving/openai_api_server.py

@@ -686,7 +686,7 @@ async def generate_completion_stream(payload: Dict[str, Any], worker_addr: str):
             worker_addr + "/worker_generate_stream",
             headers=headers,
             json=payload,
-            timeout=30,
+            timeout=60,
         ) as response:
             # content = await response.aread()
             buffer = b""

gpt_server/utils.py

@@ -136,6 +136,9 @@ def start_model_worker(config: dict):
         host = config["model_worker_args"]["host"]
         controller_address = config["model_worker_args"]["controller_address"]
         log_level = config["model_worker_args"].get("log_level", "WARNING")
+        limit_worker_concurrency = config["model_worker_args"].get(
+            "limit_worker_concurrency", 1024
+        )
     except KeyError as e:
         error_msg = f"请参照 https://github.com/shell-nlp/gpt_server/blob/main/gpt_server/script/config.yaml 设置正确的 model_worker_args"
         logger.error(error_msg)
@@ -254,6 +257,7 @@ def start_model_worker(config: dict):
                         + f" --kv_cache_quant_policy {kv_cache_quant_policy}"  # kv cache 量化策略
                         + f" --log_level {log_level}"  # 日志水平
                         + f" --task_type {task_type}"  # 日志水平
+                        + f" --limit_worker_concurrency {limit_worker_concurrency}"  # 限制worker并发数
                     )
                     # 处理为 None的情况
                     if lora: