modify pdftomd

src/pdf2md/doctor.py

@@ -13,8 +13,10 @@ from dataclasses import dataclass
 from pathlib import Path
 from typing import Any, Literal, Protocol
 
+from pdf2md.gpu import NVIDIA_SMI_QUERY, GpuInfo, parse_nvidia_smi_gpus, select_gpu
 from pdf2md.math_render import default_mathjax_helper_path
 from pdf2md.mineru_adapter import CommandResult, MinerUAdapter, MinerUVersionResult
+from pdf2md.mineru_profile import resolve_mineru_profile
 
 
 DoctorStatus = Literal["pass", "warn", "fail"]
@@ -178,13 +180,7 @@ def _check_gpu(which: Which, run_command: CommandRunner) -> DoctorCheck:
     if nvidia_smi_path is None:
         return DoctorCheck("gpu", "warn", "nvidia-smi was not found; NVIDIA GPU visibility could not be confirmed.")
 
-    result = run_command(
-        (
-            "nvidia-smi",
-            "--query-gpu=name,memory.total,driver_version",
-            "--format=csv,noheader",
-        )
-    )
+    result = run_command(NVIDIA_SMI_QUERY)
     if result.exit_code != 0:
         return DoctorCheck(
             "gpu",
@@ -193,20 +189,30 @@ def _check_gpu(which: Which, run_command: CommandRunner) -> DoctorCheck:
             (f"path: {nvidia_smi_path}", f"exit code: {result.exit_code}", _trim_detail(result.stderr)),
         )
 
-    gpu_lines = tuple(line.strip() for line in result.stdout.splitlines() if line.strip())
-    if not gpu_lines:
+    try:
+        gpus = parse_nvidia_smi_gpus(result.stdout)
+    except ValueError as error:
+        return DoctorCheck(
+            "gpu",
+            "warn",
+            "nvidia-smi output could not be parsed.",
+            (f"path: {nvidia_smi_path}", str(error), _trim_detail(result.stdout)),
+        )
+
+    if not gpus:
         return DoctorCheck("gpu", "warn", "nvidia-smi reported no visible NVIDIA GPU.", (f"path: {nvidia_smi_path}",))
 
-    risky_names = tuple(line for line in gpu_lines if _is_pascal_or_pre_turing(line))
-    if risky_names:
+    details = [f"path: {nvidia_smi_path}", *_gpu_detail_lines(gpus), *_gpu_recommendation_details(gpus)]
+    risky_gpus = tuple(gpu for gpu in gpus if gpu.pre_turing_risk)
+    if risky_gpus:
         return DoctorCheck(
             "gpu",
             "warn",
             "NVIDIA GPU is visible, but Pascal/pre-Turing compatibility risk was detected.",
-            (f"path: {nvidia_smi_path}", *risky_names),
+            tuple(details),
         )
 
-    return DoctorCheck("gpu", "pass", "NVIDIA GPU is visible.", (f"path: {nvidia_smi_path}", *gpu_lines))
+    return DoctorCheck("gpu", "pass", "NVIDIA GPU is visible.", tuple(details))
 
 
 def _check_pytorch(import_module: ImportModule) -> DoctorCheck:
@@ -269,6 +275,25 @@ def _check_pytorch(import_module: ImportModule) -> DoctorCheck:
     return DoctorCheck("pytorch", "pass", f"PyTorch {version} reports CUDA available.", tuple(details))
 
 
+def _gpu_detail_lines(gpus: tuple[GpuInfo, ...]) -> tuple[str, ...]:
+    return tuple(
+        f"gpu {gpu.index}: {gpu.name}, {gpu.memory_total_mib} MiB, driver {gpu.driver_version}"
+        for gpu in gpus
+    )
+
+
+def _gpu_recommendation_details(gpus: tuple[GpuInfo, ...]) -> tuple[str, ...]:
+    try:
+        selection = select_gpu(gpus, "auto")
+    except ValueError:
+        return ()
+    profile = resolve_mineru_profile("auto", selected_gpu=selection.gpu, cuda_requested=True)
+    return (
+        f"auto gpu: {selection.cuda_device} ({selection.gpu.name}, {selection.gpu.memory_total_mib} MiB)",
+        f"recommended MinerU profile: {profile.applied_profile}",
+    )
+
+
 def _check_model_cache(env: Mapping[str, str], path_exists: PathExists, home: Path) -> DoctorCheck:
     configured_values: list[str] = []
     existing_paths: list[str] = []