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

packages/ai/python/gpu.py

@@ -11,44 +11,101 @@ def emit_info(msg):
     print(json.dumps({"info": msg}), file=sys.stderr, flush=True)
 
 
-@functools.lru_cache(maxsize=1)
+def _nvidia_smi_gpu_name():
-def gpu_available():
+    """Return GPU name from nvidia-smi, or None if unavailable."""
-    """Return True if a usable CUDA GPU is present at runtime."""
-    # Allow explicit disable via env var (set to "false" or "0")
-    override = os.environ.get("SNAPOTTER_GPU")
-    if override is not None and override.lower() in ("0", "false", "no"):
-        return False
-
-    # Use torch.cuda as the source of truth when available. It actually
-    # probes the hardware. Fall back to onnxruntime provider detection
-    # when torch is not installed (e.g. CPU-only images without PyTorch).
     try:
-        import torch
-        avail = torch.cuda.is_available()
-        if avail:
-            name = torch.cuda.get_device_name(0)
-            print(f"[gpu] CUDA available via torch: {name}", file=sys.stderr, flush=True)
-        else:
-            print("[gpu] torch loaded but CUDA not available", file=sys.stderr, flush=True)
-        return avail
-    except ImportError as e:
-        print(f"[gpu] torch not importable: {e}", file=sys.stderr, flush=True)
-
-    # Fallback: check if onnxruntime-gpu is installed and CUDA EP is available,
-    # then verify an actual NVIDIA GPU is present via nvidia-smi.
-    try:
-        import onnxruntime as _ort
-        providers = _ort.get_available_providers()
-        if "CUDAExecutionProvider" not in providers:
-            return False
-        # CUDA EP is compiled in — verify hardware is actually present.
-        # nvidia-smi is the most reliable cross-platform check.
         result = subprocess.run(
             ["nvidia-smi", "--query-gpu=name", "--format=csv,noheader"],
             capture_output=True, text=True, timeout=5,
         )
         if result.returncode == 0 and result.stdout.strip():
-            print(f"[gpu] CUDA available via ONNX Runtime + nvidia-smi: {result.stdout.strip()}",
+            return result.stdout.strip()
+    except (FileNotFoundError, subprocess.TimeoutExpired):
+        pass
+    return None
+
+
+@functools.lru_cache(maxsize=1)
+def gpu_available():
+    """Return True if a usable CUDA GPU is present at runtime."""
+    override = os.environ.get("SNAPOTTER_GPU")
+    if override is not None and override.lower() in ("0", "false", "no"):
+        return False
+
+    # Try torch first -- it probes the hardware directly.
+    torch_available = _try_torch_cuda()
+    if torch_available:
+        return True
+
+    # torch either isn't installed or can't use CUDA. Fall through to
+    # ONNX Runtime + nvidia-smi so ONNX-based tools can still use GPU.
+    onnx_available = _try_onnx_cuda()
+    if onnx_available:
+        return True
+
+    # Last resort: check nvidia-smi alone. The GPU is present even if
+    # neither torch nor ONNX Runtime can use it (e.g. CPU-only packages).
+    gpu_name = _nvidia_smi_gpu_name()
+    if gpu_name:
+        print(f"[gpu] nvidia-smi found GPU ({gpu_name}) but neither torch "
+              "nor ONNX Runtime can use it -- reinstall AI features for GPU support",
+              file=sys.stderr, flush=True)
+    return False
+
+
+def _try_torch_cuda():
+    """Check GPU via torch.cuda. Returns True if CUDA is usable."""
+    try:
+        import torch
+    except ImportError as e:
+        print(f"[gpu] torch not importable: {e}", file=sys.stderr, flush=True)
+        return False
+
+    if torch.cuda.is_available():
+        name = torch.cuda.get_device_name(0)
+        print(f"[gpu] CUDA available via torch: {name}", file=sys.stderr, flush=True)
+        return True
+
+    # CUDA not available -- diagnose why.
+    cuda_version = getattr(torch.version, "cuda", None)
+    if not cuda_version:
+        gpu_name = _nvidia_smi_gpu_name()
+        if gpu_name:
+            print(f"[gpu] torch is a CPU-only build but GPU is present ({gpu_name}) "
+                  "-- reinstall AI features to get CUDA support",
+                  file=sys.stderr, flush=True)
+        else:
+            print("[gpu] torch is a CPU-only build and no GPU detected",
+                  file=sys.stderr, flush=True)
+        return False
+
+    # torch has CUDA compiled in but can't access the GPU.
+    diag = [f"torch has CUDA {cuda_version} but cannot access GPU"]
+    ld_path = os.environ.get("LD_LIBRARY_PATH", "")
+    diag.append(f"LD_LIBRARY_PATH={'<empty>' if not ld_path else ld_path}")
+    try:
+        torch.cuda.init()
+    except RuntimeError as e:
+        diag.append(f"torch.cuda.init(): {e}")
+    gpu_name = _nvidia_smi_gpu_name()
+    if gpu_name:
+        diag.append(f"nvidia-smi sees GPU ({gpu_name}) but torch cannot use it")
+    else:
+        diag.append("nvidia-smi also cannot find a GPU")
+    print(f"[gpu] {'; '.join(diag)}", file=sys.stderr, flush=True)
+    return False
+
+
+def _try_onnx_cuda():
+    """Check GPU via ONNX Runtime CUDAExecutionProvider + nvidia-smi."""
+    try:
+        import onnxruntime as _ort
+        providers = _ort.get_available_providers()
+        if "CUDAExecutionProvider" not in providers:
+            return False
+        gpu_name = _nvidia_smi_gpu_name()
+        if gpu_name:
+            print(f"[gpu] CUDA available via ONNX Runtime + nvidia-smi: {gpu_name}",
                   file=sys.stderr, flush=True)
             return True
         return False

packages/ai/src/bridge.ts

@@ -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) {

tests/unit/ai/gpu-detection.test.ts

@@ -0,0 +1,212 @@
+import { type ChildProcess, spawn } from "node:child_process";
+import { EventEmitter } from "node:events";
+import { Writable } from "node:stream";
+import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
+
+vi.mock("node:child_process", () => ({
+  spawn: vi.fn(),
+}));
+
+vi.mock("sharp", () => ({
+  default: vi.fn(),
+}));
+
+function createMockProcess(): {
+  process: ChildProcess;
+  stdin: Writable;
+  stdout: EventEmitter;
+  stderr: EventEmitter;
+  emitEvent: (event: string, ...args: unknown[]) => void;
+  stdinWrites: string[];
+} {
+  const stdinWrites: string[] = [];
+  const stdin = new Writable({
+    write(chunk, _encoding, callback) {
+      stdinWrites.push(chunk.toString());
+      callback();
+    },
+  });
+  const stdout = new EventEmitter();
+  const stderr = new EventEmitter();
+
+  const proc = new EventEmitter() as unknown as ChildProcess;
+  Object.assign(proc, {
+    stdin,
+    stdout,
+    stderr,
+    pid: 12345,
+    killed: false,
+    kill: vi.fn(() => {
+      (proc as { killed: boolean }).killed = true;
+      return true;
+    }),
+  });
+
+  return {
+    process: proc,
+    stdin,
+    stdout,
+    stderr,
+    emitEvent: (event: string, ...args: unknown[]) => proc.emit(event, ...args),
+    stdinWrites,
+  };
+}
+
+describe("buildMinimalEnv - env passthrough", () => {
+  let runPythonWithProgress: (
+    scriptName: string,
+    args: string[],
+    options?: { onProgress?: (p: number, s: string) => void; timeout?: number },
+  ) => Promise<{ stdout: string; stderr: string }>;
+
+  beforeEach(async () => {
+    vi.resetModules();
+    const mod = await import("../../../packages/ai/src/bridge.js");
+    runPythonWithProgress = mod.runPythonWithProgress;
+  });
+
+  afterEach(() => {
+    vi.restoreAllMocks();
+    delete process.env.SNAPOTTER_GPU;
+    delete process.env.MODELS_PATH;
+    delete process.env.MODELS_DIR;
+  });
+
+  function getSpawnEnv(): Record<string, string> | undefined {
+    const allCalls = vi.mocked(spawn).mock.calls;
+    const lastCall = allCalls[allCalls.length - 1];
+    return (lastCall?.[2] as { env?: Record<string, string> })?.env;
+  }
+
+  it("passes SNAPOTTER_GPU to Python subprocess", async () => {
+    process.env.SNAPOTTER_GPU = "1";
+
+    const mock = createMockProcess();
+    vi.mocked(spawn).mockReturnValue(mock.process);
+
+    const promise = runPythonWithProgress("test.py", []);
+    mock.stdout.emit("data", Buffer.from('{"ok": true}\n'));
+    mock.emitEvent("close", 0, null);
+    await promise;
+
+    const env = getSpawnEnv();
+    expect(env).toBeDefined();
+    expect(env!.SNAPOTTER_GPU).toBe("1");
+  });
+
+  it("passes MODELS_PATH to Python subprocess", async () => {
+    process.env.MODELS_PATH = "/data/ai/models";
+
+    const mock = createMockProcess();
+    vi.mocked(spawn).mockReturnValue(mock.process);
+
+    const promise = runPythonWithProgress("test.py", []);
+    mock.stdout.emit("data", Buffer.from('{"ok": true}\n'));
+    mock.emitEvent("close", 0, null);
+    await promise;
+
+    const env = getSpawnEnv();
+    expect(env).toBeDefined();
+    expect(env!.MODELS_PATH).toBe("/data/ai/models");
+  });
+
+  it("does not pass MODELS_DIR (removed dead entry)", async () => {
+    process.env.MODELS_DIR = "/some/path";
+
+    const mock = createMockProcess();
+    vi.mocked(spawn).mockReturnValue(mock.process);
+
+    const promise = runPythonWithProgress("test.py", []);
+    mock.stdout.emit("data", Buffer.from('{"ok": true}\n'));
+    mock.emitEvent("close", 0, null);
+    await promise;
+
+    const env = getSpawnEnv();
+    expect(env).toBeDefined();
+    expect(env!.MODELS_DIR).toBeUndefined();
+  });
+
+  it("does not include SNAPOTTER_GPU when not set in parent env", async () => {
+    delete process.env.SNAPOTTER_GPU;
+
+    const mock = createMockProcess();
+    vi.mocked(spawn).mockReturnValue(mock.process);
+
+    const promise = runPythonWithProgress("test.py", []);
+    mock.stdout.emit("data", Buffer.from('{"ok": true}\n'));
+    mock.emitEvent("close", 0, null);
+    await promise;
+
+    const env = getSpawnEnv();
+    expect(env).toBeDefined();
+    expect(env!.SNAPOTTER_GPU).toBeUndefined();
+  });
+
+  it("always includes PYTHONUNBUFFERED=1", async () => {
+    const mock = createMockProcess();
+    vi.mocked(spawn).mockReturnValue(mock.process);
+
+    const promise = runPythonWithProgress("test.py", []);
+    mock.stdout.emit("data", Buffer.from('{"ok": true}\n'));
+    mock.emitEvent("close", 0, null);
+    await promise;
+
+    const env = getSpawnEnv();
+    expect(env).toBeDefined();
+    expect(env!.PYTHONUNBUFFERED).toBe("1");
+  });
+
+  it("passes LD_LIBRARY_PATH when set", async () => {
+    process.env.LD_LIBRARY_PATH = "/usr/local/nvidia/lib64";
+
+    const mock = createMockProcess();
+    vi.mocked(spawn).mockReturnValue(mock.process);
+
+    const promise = runPythonWithProgress("test.py", []);
+    mock.stdout.emit("data", Buffer.from('{"ok": true}\n'));
+    mock.emitEvent("close", 0, null);
+    await promise;
+
+    const env = getSpawnEnv();
+    expect(env!.LD_LIBRARY_PATH).toBe("/usr/local/nvidia/lib64");
+
+    delete process.env.LD_LIBRARY_PATH;
+  });
+
+  it("passes CUDA_VISIBLE_DEVICES when set", async () => {
+    process.env.CUDA_VISIBLE_DEVICES = "0,1";
+
+    const mock = createMockProcess();
+    vi.mocked(spawn).mockReturnValue(mock.process);
+
+    const promise = runPythonWithProgress("test.py", []);
+    mock.stdout.emit("data", Buffer.from('{"ok": true}\n'));
+    mock.emitEvent("close", 0, null);
+    await promise;
+
+    const env = getSpawnEnv();
+    expect(env!.CUDA_VISIBLE_DEVICES).toBe("0,1");
+
+    delete process.env.CUDA_VISIBLE_DEVICES;
+  });
+
+  it("does not leak unrelated env vars to Python subprocess", async () => {
+    process.env.SECRET_API_KEY = "should-not-leak";
+    process.env.DATABASE_URL = "sqlite://secret.db";
+
+    const mock = createMockProcess();
+    vi.mocked(spawn).mockReturnValue(mock.process);
+
+    const promise = runPythonWithProgress("test.py", []);
+    mock.stdout.emit("data", Buffer.from('{"ok": true}\n'));
+    mock.emitEvent("close", 0, null);
+    await promise;
+
+    const env = getSpawnEnv();
+    expect(env!.SECRET_API_KEY).toBeUndefined();
+    expect(env!.DATABASE_URL).toBeUndefined();
+
+    delete process.env.SECRET_API_KEY;
+    delete process.env.DATABASE_URL;
+  });
+});