utils.py 3.8 KB

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