lmdeploy_server.py 3.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293
  1. import os
  2. import sys
  3. from loguru import logger
  4. from mineru.backend.vlm.utils import set_lmdeploy_backend
  5. from mineru.utils.models_download_utils import auto_download_and_get_model_root_path
  6. def main():
  7. args = sys.argv[1:]
  8. has_port_arg = False
  9. has_gpu_memory_utilization_arg = False
  10. has_log_level_arg = False
  11. device_type = ""
  12. lm_backend = ""
  13. # 检查现有参数
  14. indices_to_remove = []
  15. for i, arg in enumerate(args):
  16. if arg == "--server-port" or arg.startswith("--server-port="):
  17. has_port_arg = True
  18. if arg == "--cache-max-entry-count" or arg.startswith("--cache-max-entry-count="):
  19. has_gpu_memory_utilization_arg = True
  20. if arg == "--log-level" or arg.startswith("--log-level="):
  21. has_log_level_arg = True
  22. if arg == "--backend" or arg == "--lmdeploy-backend":
  23. if i + 1 < len(args):
  24. lm_backend = args[i + 1]
  25. indices_to_remove.extend([i, i + 1])
  26. elif arg.startswith("--backend=") or arg.startswith("--lmdeploy-backend="):
  27. lm_backend = arg.split("=", 1)[1]
  28. indices_to_remove.append(i)
  29. if arg == "--device" or arg == "--lmdeploy-device":
  30. if i + 1 < len(args):
  31. device_type = args[i + 1]
  32. indices_to_remove.extend([i, i + 1])
  33. elif arg.startswith("--device=") or arg.startswith("--lmdeploy-device="):
  34. device_type = arg.split("=", 1)[1]
  35. indices_to_remove.append(i)
  36. # 从后往前删除,避免索引错位
  37. for i in sorted(set(indices_to_remove), reverse=True):
  38. args.pop(i)
  39. # 添加默认参数
  40. if not has_port_arg:
  41. args.extend(["--server-port", "30000"])
  42. if not has_gpu_memory_utilization_arg:
  43. args.extend(["--cache-max-entry-count", "0.5"])
  44. if not has_log_level_arg:
  45. args.extend(["--log-level", "ERROR"])
  46. device_type = os.getenv("MINERU_LMDEPLOY_DEVICE", device_type)
  47. if device_type == "":
  48. device_type = "cuda"
  49. elif device_type not in ["cuda", "ascend", "maca", "camb"]:
  50. raise ValueError(f"Unsupported lmdeploy device type: {device_type}")
  51. lm_backend = os.getenv("MINERU_LMDEPLOY_BACKEND", lm_backend)
  52. if lm_backend == "":
  53. lm_backend = set_lmdeploy_backend(device_type)
  54. elif lm_backend not in ["pytorch", "turbomind"]:
  55. raise ValueError(f"Unsupported lmdeploy backend: {lm_backend}")
  56. logger.info(f"lmdeploy device is: {device_type}, lmdeploy backend is: {lm_backend}")
  57. if lm_backend == "pytorch":
  58. os.environ["TOKENIZERS_PARALLELISM"] = "false"
  59. args.extend(["--device", device_type])
  60. args.extend(["--backend", lm_backend])
  61. model_path = auto_download_and_get_model_root_path("/", "vlm")
  62. # logger.debug(args)
  63. # 重构参数,将模型路径作为位置参数
  64. sys.argv = [sys.argv[0]] + ["serve", "api_server", model_path] + args
  65. if os.getenv('OMP_NUM_THREADS') is None:
  66. os.environ["OMP_NUM_THREADS"] = "1"
  67. # 启动 lmdeploy 服务器
  68. print(f"start lmdeploy server: {sys.argv}")
  69. # 使用os.system调用启动lmdeploy服务器
  70. os.system("lmdeploy " + " ".join(sys.argv[1:]))
  71. if __name__ == "__main__":
  72. main()