Эх сурвалжийг харах

fix: add default arguments for port and GPU memory utilization in server.py

myhloli 2 сар өмнө
parent
commit
b35c3345c0

+ 16 - 7
mineru/model/vlm_vllm_model/server.py

@@ -5,19 +5,28 @@ from vllm.entrypoints.cli.main import main as vllm_main
 
 
 def main():
-    # 检查命令行参数中是否包含--model-path
     args = sys.argv[1:]
-    has_model_path_arg = False
+
+    has_port_arg = False
+    has_gpu_memory_utilization_arg = False
+    has_model_arg = False
+    has_model_tag_arg = False
 
     for i, arg in enumerate(args):
+        if arg == "--port" or arg.startswith("--port="):
+            has_port_arg = True
+        if arg == "--gpu-memory-utilization" or arg.startswith("--gpu-memory-utilization="):
+            has_gpu_memory_utilization_arg = True
         if arg == "--model" or arg.startswith("--model="):
-            has_model_path_arg = True
-            break
+            has_model_arg = True
 
-    # 如果没有--model-path参数,在参数列表中添加它
-    if not has_model_path_arg:
+    if not has_port_arg:
+        args.extend(["--port", "30000"])
+    if not has_gpu_memory_utilization_arg:
+        args.extend(["--gpu-memory-utilization", "0.5"])
+    if not has_model_arg:
         default_path = auto_download_and_get_model_root_path("/", "vlm")
-        args.extend(["--model", default_path])
+        args.extend([default_path])
 
     # 重新构造sys.argv,以便透传所有参数给vllm
     sys.argv = [sys.argv[0]] + args