fix: improve GPU detection diagnostics and fallback for container environments

The GPU detection in gpu.py had two issues preventing GPU usage in
containers (especially rootless podman with CDI):

1. When torch was installed but torch.cuda.is_available() returned
   False, the function returned immediately without trying the
   ONNX Runtime + nvidia-smi fallback. This meant a CPU-only torch
   build (installed before GPU was available) would block all GPU
   detection, even for ONNX-based tools.

2. The failure logged a generic "torch loaded but CUDA not available"
   with no diagnostic information, making it impossible to debug
   whether the issue was a CPU-only build, missing libraries, or
   device permissions.

The fix restructures gpu_available() into three detection tiers
(torch -> ONNX Runtime -> nvidia-smi) that always fall through on
failure. When torch CUDA fails, it now checks torch.version.cuda to
distinguish CPU-only builds from CUDA builds that can't access the
GPU, and logs LD_LIBRARY_PATH, torch.cuda.init() errors, and
nvidia-smi results.

Also fixes two env var passthrough bugs in buildMinimalEnv():
- SNAPOTTER_GPU was never passed to the Python subprocess, so the
  user-facing GPU override env var had no effect
- MODELS_DIR was a dead entry (never set as env var); replaced with
  MODELS_PATH which the Dockerfile sets and Python scripts read

Closes #134
This commit is contained in:
SnapOtter
2026-05-14 23:17:21 +08:00
parent b449bd1a56
commit 19a607454a
3 changed files with 303 additions and 33 deletions
+2 -1
View File
@@ -26,11 +26,12 @@ function buildMinimalEnv(): Record<string, string> {
"LD_LIBRARY_PATH",
// Application-specific vars the sidecar scripts depend on
"DATA_DIR",
"MODELS_DIR",
"MODELS_PATH",
"U2NET_HOME",
"PROCESSING_TIMEOUT_S",
"DISPATCHER_MAX_REQUESTS",
"PYTHON_VENV_PATH",
"SNAPOTTER_GPU",
];
for (const key of passthrough) {
if (process.env[key] !== undefined) {