utils.py 3.7 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102
  1. import os
  2. from loguru import logger
  3. from packaging import version
  4. from mineru.utils.check_sys_env import is_windows_environment, is_linux_environment
  5. from mineru.utils.config_reader import get_device
  6. from mineru.utils.model_utils import get_vram
  7. def enable_custom_logits_processors() -> bool:
  8. import torch
  9. from vllm import __version__ as vllm_version
  10. if not torch.cuda.is_available():
  11. logger.info("CUDA not available, disabling custom_logits_processors")
  12. return False
  13. major, minor = torch.cuda.get_device_capability()
  14. # 正确计算Compute Capability
  15. compute_capability = f"{major}.{minor}"
  16. # 安全地处理环境变量
  17. vllm_use_v1_str = os.getenv('VLLM_USE_V1', "1")
  18. if vllm_use_v1_str.isdigit():
  19. vllm_use_v1 = int(vllm_use_v1_str)
  20. else:
  21. vllm_use_v1 = 1
  22. if vllm_use_v1 == 0:
  23. logger.info("VLLM_USE_V1 is set to 0, disabling custom_logits_processors")
  24. return False
  25. elif version.parse(vllm_version) < version.parse("0.10.1"):
  26. logger.info(f"vllm version: {vllm_version} < 0.10.1, disable custom_logits_processors")
  27. return False
  28. elif version.parse(compute_capability) < version.parse("8.0"):
  29. if version.parse(vllm_version) >= version.parse("0.10.2"):
  30. logger.info(f"compute_capability: {compute_capability} < 8.0, but vllm version: {vllm_version} >= 0.10.2, enable custom_logits_processors")
  31. return True
  32. else:
  33. logger.info(f"compute_capability: {compute_capability} < 8.0 and vllm version: {vllm_version} < 0.10.2, disable custom_logits_processors")
  34. return False
  35. else:
  36. logger.info(f"compute_capability: {compute_capability} >= 8.0 and vllm version: {vllm_version} >= 0.10.1, enable custom_logits_processors")
  37. return True
  38. def set_lmdeploy_backend(device_type:str) -> str:
  39. lmdeploy_backend = ""
  40. if device_type.lower() in ["ascend", "maca", "camb"]:
  41. lmdeploy_backend = "pytorch"
  42. elif device_type.lower() in ["cuda"]:
  43. import torch
  44. if not torch.cuda.is_available():
  45. raise ValueError("CUDA is not available.")
  46. if is_windows_environment():
  47. lmdeploy_backend = "turbomind"
  48. elif is_linux_environment():
  49. major, minor = torch.cuda.get_device_capability()
  50. compute_capability = f"{major}.{minor}"
  51. if version.parse(compute_capability) >= version.parse("8.0"):
  52. lmdeploy_backend = "pytorch"
  53. else:
  54. lmdeploy_backend = "turbomind"
  55. else:
  56. raise ValueError("Unsupported operating system.")
  57. else:
  58. raise ValueError(f"Unsupported device type: {device_type}")
  59. return lmdeploy_backend
  60. def set_default_gpu_memory_utilization() -> float:
  61. from vllm import __version__ as vllm_version
  62. if version.parse(vllm_version) >= version.parse("0.11.0"):
  63. return 0.7
  64. else:
  65. return 0.5
  66. def set_default_batch_size() -> int:
  67. try:
  68. device = get_device()
  69. vram = get_vram(device)
  70. if vram is not None:
  71. gpu_memory = int(os.getenv('MINERU_VIRTUAL_VRAM_SIZE', round(vram)))
  72. if gpu_memory >= 16:
  73. batch_size = 8
  74. elif gpu_memory >= 8:
  75. batch_size = 4
  76. else:
  77. batch_size = 1
  78. logger.info(f'gpu_memory: {gpu_memory} GB, batch_size: {batch_size}')
  79. else:
  80. # Default batch_ratio when VRAM can't be determined
  81. batch_size = 1
  82. logger.info(f'Could not determine GPU memory, using default batch_ratio: {batch_size}')
  83. except Exception as e:
  84. logger.warning(f'Error determining VRAM: {e}, using default batch_ratio: 1')
  85. batch_size = 1
  86. return batch_size