lmdeploy_server.py 1.8 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859
  1. import os
  2. import sys
  3. from mineru.utils.models_download_utils import auto_download_and_get_model_root_path
  4. def main():
  5. args = sys.argv[1:]
  6. has_port_arg = False
  7. has_gpu_memory_utilization_arg = False
  8. has_log_level_arg = False
  9. has_device_arg = False
  10. device_type = "cuda"
  11. # 检查现有参数
  12. for i, arg in enumerate(args):
  13. if arg == "--server-port" or arg.startswith("--server-port="):
  14. has_port_arg = True
  15. if arg == "--cache-max-entry-count" or arg.startswith("--cache-max-entry-count="):
  16. has_gpu_memory_utilization_arg = True
  17. if arg == "--log-level" or arg.startswith("--log-level="):
  18. has_log_level_arg = True
  19. if arg == "--device":
  20. has_device_arg = True
  21. if i + 1 < len(args):
  22. device_type = args[i + 1]
  23. elif arg.startswith("--device="):
  24. has_device_arg = True
  25. device_type = arg.split("=", 1)[1]
  26. # 添加默认参数
  27. if not has_port_arg:
  28. args.extend(["--server-port", "30000"])
  29. if not has_gpu_memory_utilization_arg:
  30. args.extend(["--cache-max-entry-count", "0.5"])
  31. if not has_log_level_arg:
  32. args.extend(["--log-level", "ERROR"])
  33. if has_device_arg:
  34. if device_type.lower() in ["ascend", "maca", "camb"]:
  35. args.extend(["--backend", "pytorch"])
  36. model_path = auto_download_and_get_model_root_path("/", "vlm")
  37. # 重构参数,将模型路径作为位置参数
  38. sys.argv = [sys.argv[0]] + ["serve", "api_server", model_path] + args
  39. if os.getenv('OMP_NUM_THREADS') is None:
  40. os.environ["OMP_NUM_THREADS"] = "1"
  41. # 启动 lmdeploy 服务器
  42. print(f"start lmdeploy server: {sys.argv}")
  43. # 使用os.system调用启动lmdeploy服务器
  44. os.system("lmdeploy " + " ".join(sys.argv[1:]))
  45. if __name__ == "__main__":
  46. main()