김경종
107 lines
3.9 KiB
Python
107 lines
3.9 KiB
Python
from __future__ import annotations
|
|
|
|
from pdf2md.gpu import GpuInfo
|
|
from pdf2md.ir import WarningCode, WarningSeverity
|
|
from pdf2md.mineru_profile import resolve_mineru_profile
|
|
|
|
|
|
SAFE_ENV = {
|
|
"MINERU_PROCESSING_WINDOW_SIZE": "1",
|
|
"MINERU_API_MAX_CONCURRENT_REQUESTS": "1",
|
|
"MINERU_PDF_RENDER_THREADS": "1",
|
|
}
|
|
|
|
|
|
def test_auto_profile_uses_safe_values_without_gpu_inventory() -> None:
|
|
profile = resolve_mineru_profile("auto", selected_gpu=None, cuda_requested=True)
|
|
|
|
assert profile.applied_profile == "safe"
|
|
assert profile.environment == SAFE_ENV
|
|
assert [warning.code for warning in profile.warnings] == [WarningCode.GPU_UNAVAILABLE]
|
|
|
|
|
|
def test_auto_profile_uses_safe_values_for_gtx_1070_ti() -> None:
|
|
gpu = GpuInfo(index=0, name="NVIDIA GeForce GTX 1070 Ti", memory_total_mib=8192, driver_version="577.00")
|
|
|
|
profile = resolve_mineru_profile("auto", selected_gpu=gpu, cuda_requested=True)
|
|
|
|
assert profile.requested_profile == "auto"
|
|
assert profile.applied_profile == "safe"
|
|
assert profile.environment == SAFE_ENV
|
|
assert profile.selected_gpu_name == "NVIDIA GeForce GTX 1070 Ti"
|
|
|
|
|
|
def test_auto_profile_uses_moderate_values_for_16gb_turing_or_newer_gpu() -> None:
|
|
gpu = GpuInfo(index=0, name="NVIDIA RTX A4000", memory_total_mib=16384, driver_version="577.00")
|
|
|
|
profile = resolve_mineru_profile("auto", selected_gpu=gpu, cuda_requested=True)
|
|
|
|
assert profile.applied_profile == "auto"
|
|
assert profile.environment == {
|
|
"MINERU_PROCESSING_WINDOW_SIZE": "8",
|
|
"MINERU_API_MAX_CONCURRENT_REQUESTS": "1",
|
|
"MINERU_PDF_RENDER_THREADS": "4",
|
|
}
|
|
assert profile.warnings == ()
|
|
|
|
|
|
def test_auto_profile_uses_conservative_values_for_12gb_to_16gb_gpu() -> None:
|
|
gpu = GpuInfo(index=0, name="NVIDIA RTX 4070", memory_total_mib=12288, driver_version="577.00")
|
|
|
|
profile = resolve_mineru_profile("auto", selected_gpu=gpu, cuda_requested=True)
|
|
|
|
assert profile.applied_profile == "auto-conservative"
|
|
assert profile.environment == {
|
|
"MINERU_PROCESSING_WINDOW_SIZE": "4",
|
|
"MINERU_API_MAX_CONCURRENT_REQUESTS": "1",
|
|
"MINERU_PDF_RENDER_THREADS": "2",
|
|
}
|
|
|
|
|
|
def test_performance_profile_uses_performance_values_only_on_strong_gpu() -> None:
|
|
gpu = GpuInfo(index=1, name="NVIDIA RTX 4090", memory_total_mib=24564, driver_version="577.00")
|
|
|
|
profile = resolve_mineru_profile("performance", selected_gpu=gpu, cuda_requested=True)
|
|
|
|
assert profile.applied_profile == "performance"
|
|
assert profile.environment == {
|
|
"MINERU_PROCESSING_WINDOW_SIZE": "16",
|
|
"MINERU_API_MAX_CONCURRENT_REQUESTS": "1",
|
|
"MINERU_PDF_RENDER_THREADS": "4",
|
|
}
|
|
assert profile.selected_gpu_index == 1
|
|
assert profile.selected_gpu_vram_mib == 24564
|
|
|
|
|
|
def test_performance_profile_downgrades_to_safe_on_weak_gpu() -> None:
|
|
gpu = GpuInfo(index=0, name="NVIDIA GeForce GTX 1070 Ti", memory_total_mib=8192, driver_version="577.00")
|
|
|
|
profile = resolve_mineru_profile("performance", selected_gpu=gpu, cuda_requested=True)
|
|
|
|
assert profile.applied_profile == "safe"
|
|
assert profile.environment == SAFE_ENV
|
|
assert [warning.code for warning in profile.warnings] == [WarningCode.MINERU_PROFILE_ADJUSTED]
|
|
assert profile.warnings[0].severity == WarningSeverity.WARNING
|
|
|
|
|
|
def test_profile_details_are_json_ready() -> None:
|
|
gpu = GpuInfo(index=0, name="NVIDIA RTX A5000", memory_total_mib=24564, driver_version="577.00")
|
|
|
|
profile = resolve_mineru_profile("auto", selected_gpu=gpu, cuda_requested=True)
|
|
|
|
assert profile.to_engine_options() == {
|
|
"requested": "auto",
|
|
"applied": "auto",
|
|
"environment": {
|
|
"MINERU_API_MAX_CONCURRENT_REQUESTS": "1",
|
|
"MINERU_PDF_RENDER_THREADS": "4",
|
|
"MINERU_PROCESSING_WINDOW_SIZE": "8",
|
|
},
|
|
"selected_gpu": {
|
|
"index": 0,
|
|
"name": "NVIDIA RTX A5000",
|
|
"memory_total_mib": 24564,
|
|
"pre_turing_risk": False,
|
|
},
|
|
}
|