fix(ai-bundles): stop CPU onnxruntime from clobbering onnxruntime-gpu (#544)

Both PyPI onnxruntime flavors unpack into the same site-packages directory, so a bundle carrying the CPU build (transcription, via faster-whisper) overwrote the GPU build's native libraries during install while the stale onnxruntime_gpu dist-info kept claiming otherwise. Every ONNX-backed tool then silently ran on CPU.

The installer now reconciles the flavor before the venv merge and the GPU build always wins, in both install orders; reinstalling any GPU bundle repairs a previously clobbered venv. gpu.py's warning now says exactly that. Build-side, build-bundle.sh gains the same reconcile and verify-bundle-compatibility.sh layers bundles through the real installer merge and asserts a single flavor.

Verified live on an RTX 4070 against the published bundles: reproduced the clobber with the stock installer, then confirmed both the prevention and repair paths with the patched one.

Fixes #490

packages/ai/python/gpu.py

@@ -168,8 +168,10 @@ def onnx_providers():
             available = _ort.get_available_providers()
             if "CUDAExecutionProvider" in available:
                 return (["CUDAExecutionProvider", "CPUExecutionProvider"], "cuda")
-            emit_info("GPU detected by torch but CUDAExecutionProvider not available in onnxruntime "
-                      "-- install onnxruntime-gpu for GPU acceleration")
+            emit_info("GPU detected by torch but CUDAExecutionProvider is missing from onnxruntime. "
+                      "A bundle install likely replaced onnxruntime-gpu with the CPU build (#490); "
+                      "reinstall the Background Removal bundle from Settings > AI Features to "
+                      "restore the GPU build")
         except ImportError:
             emit_info("onnxruntime not installed, cannot check CUDA provider")
     emit_info("No GPU detected, processing on CPU")

packages/ai/python/install_feature.py

@@ -540,6 +540,70 @@ def move_tree(src: str, dst: str) -> None:
     shutil.rmtree(src, ignore_errors=True)
 
 
+# -- ONNX Runtime flavor reconciliation (GPU wins) --
+
+def _onnx_dist_infos(sp_dir: str) -> tuple:
+    """Return (cpu, gpu) lists of onnxruntime dist-info directory names.
+
+    Wheel metadata dirs are `<normalized-name>-<version>.dist-info`, and name
+    normalization turns every other onnxruntime distribution into
+    `onnxruntime_<suffix>` (onnxruntime-gpu -> onnxruntime_gpu), so a directory
+    starting with exactly "onnxruntime-" can only be the CPU build.
+    """
+    if not os.path.isdir(sp_dir):
+        return ([], [])
+    names = os.listdir(sp_dir)
+    cpu = sorted(n for n in names if n.startswith("onnxruntime-") and n.endswith(".dist-info"))
+    gpu = sorted(n for n in names if n.startswith("onnxruntime_gpu-") and n.endswith(".dist-info"))
+    return (cpu, gpu)
+
+
+def reconcile_onnxruntime(staging_sp: str, site_packages_dir: str) -> None:
+    """Never let a bundle downgrade the venv's ONNX Runtime from GPU to CPU.
+
+    `onnxruntime` (CPU) and `onnxruntime-gpu` unpack into the SAME package
+    directory (`onnxruntime/`), so whichever bundle lands last wins
+    file-by-file. A CPU build arriving after the GPU build (e.g. transcription's
+    faster-whisper dependency, after background-removal put onnxruntime-gpu in
+    place) silently strips CUDAExecutionProvider while the surviving
+    onnxruntime_gpu dist-info keeps claiming it is installed (#490). The GPU
+    build ships CPUExecutionProvider too, so it satisfies every consumer of the
+    CPU build; the reverse is not true. Hence: GPU wins, in both install orders.
+    """
+    staging_cpu, staging_gpu = _onnx_dist_infos(staging_sp)
+    venv_cpu, venv_gpu = _onnx_dist_infos(site_packages_dir)
+
+    if staging_cpu and (venv_gpu or staging_gpu):
+        # Incoming CPU flavor while a GPU build exists: drop the CPU metadata,
+        # and when the staged package files ARE the CPU build (no GPU flavor in
+        # this same bundle), drop them too so they cannot clobber the venv's
+        # GPU libraries.
+        for name in staging_cpu:
+            shutil.rmtree(os.path.join(staging_sp, name), ignore_errors=True)
+        if not staging_gpu:
+            pkg_dir = os.path.join(staging_sp, "onnxruntime")
+            if os.path.isdir(pkg_dir):
+                shutil.rmtree(pkg_dir, ignore_errors=True)
+        sys.stderr.write(
+            "[install] kept the GPU build of onnxruntime: dropped this bundle's CPU "
+            "onnxruntime so it cannot disable CUDA for other AI tools (#490)\n"
+        )
+        sys.stderr.flush()
+
+    if staging_gpu and venv_cpu:
+        # Incoming GPU flavor over a CPU install: the GPU files win the merge;
+        # clear the venv's stale CPU metadata so pip reflects reality. This also
+        # makes reinstalling any GPU bundle repair a venv clobbered before this
+        # guard existed.
+        for name in venv_cpu:
+            shutil.rmtree(os.path.join(site_packages_dir, name), ignore_errors=True)
+        sys.stderr.write(
+            "[install] replacing the CPU build of onnxruntime with the GPU build; "
+            "removed its stale metadata\n"
+        )
+        sys.stderr.flush()
+
+
 # -- Fixups (NCCL wheel) --
 
 def apply_fixups(staging_dir: str, venv_path: str) -> None:
@@ -819,6 +883,7 @@ def _install() -> None:
                     },
                     mf,
                 )
+            reconcile_onnxruntime(staging_sp, site_packages_dir)
             move_tree(staging_sp, site_packages_dir)
             if os.path.exists(venv_writing_marker):
                 os.unlink(venv_writing_marker)

packages/ai/python/tests/test_install_feature_onnx_reconcile.py

@@ -0,0 +1,176 @@
+"""ONNX Runtime flavor reconciliation tests (GPU wins).
+
+The invariant under test: a bundle install must never downgrade the shared
+venv's ONNX Runtime from the GPU build to the CPU build. Both PyPI packages
+(`onnxruntime` and `onnxruntime-gpu`) install into the SAME site-packages
+directory (`onnxruntime/`), so a bundle that carries the CPU build (e.g.
+transcription, via faster-whisper's transitive dependency) overwrites the GPU
+build's native libraries file-by-file during move_tree, while the stale
+`onnxruntime_gpu-*.dist-info` metadata survives. pip then claims the GPU build
+is installed but `CUDAExecutionProvider` is gone and every ONNX tool silently
+runs on CPU (snapotter-hq/SnapOtter#490).
+"""
+
+import importlib.util
+import os
+
+CUDA_PROVIDER_LIB = "libonnxruntime_providers_cuda.so"
+
+
+def load_installer():
+    script_path = os.path.join(os.path.dirname(__file__), "..", "install_feature.py")
+    spec = importlib.util.spec_from_file_location("install_feature_onnx_under_test", script_path)
+    module = importlib.util.module_from_spec(spec)
+    assert spec.loader is not None
+    spec.loader.exec_module(module)
+    return module
+
+
+def make_onnxruntime(sp, flavor, version="1.20.1"):
+    """Lay out a fake onnxruntime install of the given flavor under sp."""
+    capi = sp / "onnxruntime" / "capi"
+    capi.mkdir(parents=True, exist_ok=True)
+    (sp / "onnxruntime" / "__init__.py").write_text(flavor)
+    (capi / "onnxruntime_pybind11_state.so").write_text(flavor)
+    if flavor == "gpu":
+        (capi / CUDA_PROVIDER_LIB).write_text("cuda")
+        dist_info = sp / f"onnxruntime_gpu-{version}.dist-info"
+        dist_info.mkdir()
+        (dist_info / "METADATA").write_text("Name: onnxruntime-gpu")
+    else:
+        dist_info = sp / f"onnxruntime-{version}.dist-info"
+        dist_info.mkdir()
+        (dist_info / "METADATA").write_text("Name: onnxruntime")
+
+
+def dist_infos(sp, prefix):
+    return sorted(
+        n for n in os.listdir(sp)
+        if n.startswith(prefix) and n.endswith(".dist-info")
+    )
+
+
+def cpu_dist_infos(sp):
+    return [n for n in dist_infos(sp, "onnxruntime-")]
+
+
+def gpu_dist_infos(sp):
+    return [n for n in dist_infos(sp, "onnxruntime_gpu-")]
+
+
+def test_incoming_cpu_bundle_cannot_clobber_gpu_install(tmp_path):
+    """The #490 scenario: GPU bundle installed first, transcription second."""
+    installer = load_installer()
+    staging = tmp_path / "staging"
+    venv_sp = tmp_path / "venv-sp"
+    staging.mkdir()
+    venv_sp.mkdir()
+    make_onnxruntime(venv_sp, "gpu")
+    make_onnxruntime(staging, "cpu")
+    # The rest of the incoming bundle must survive untouched.
+    (staging / "faster_whisper").mkdir()
+    (staging / "faster_whisper" / "__init__.py").write_text("fw")
+
+    installer.reconcile_onnxruntime(str(staging), str(venv_sp))
+    installer.move_tree(str(staging), str(venv_sp))
+
+    # The GPU build survives: CUDA provider lib intact, core lib not downgraded.
+    assert (venv_sp / "onnxruntime" / "capi" / CUDA_PROVIDER_LIB).exists()
+    assert (venv_sp / "onnxruntime" / "capi" / "onnxruntime_pybind11_state.so").read_text() == "gpu"
+    # Metadata stays truthful: GPU dist-info only, no CPU dist-info shipped in.
+    assert gpu_dist_infos(venv_sp) == ["onnxruntime_gpu-1.20.1.dist-info"]
+    assert cpu_dist_infos(venv_sp) == []
+    # The bundle's own packages still landed.
+    assert (venv_sp / "faster_whisper" / "__init__.py").read_text() == "fw"
+
+
+def test_incoming_gpu_bundle_replaces_cpu_and_clears_stale_metadata(tmp_path):
+    """Reverse order: transcription first, then a GPU bundle. GPU wins."""
+    installer = load_installer()
+    staging = tmp_path / "staging"
+    venv_sp = tmp_path / "venv-sp"
+    staging.mkdir()
+    venv_sp.mkdir()
+    make_onnxruntime(venv_sp, "cpu")
+    make_onnxruntime(staging, "gpu")
+
+    installer.reconcile_onnxruntime(str(staging), str(venv_sp))
+    installer.move_tree(str(staging), str(venv_sp))
+
+    assert (venv_sp / "onnxruntime" / "capi" / CUDA_PROVIDER_LIB).exists()
+    assert (venv_sp / "onnxruntime" / "capi" / "onnxruntime_pybind11_state.so").read_text() == "gpu"
+    # The stale CPU dist-info is gone, so pip metadata matches reality.
+    assert cpu_dist_infos(venv_sp) == []
+    assert gpu_dist_infos(venv_sp) == ["onnxruntime_gpu-1.20.1.dist-info"]
+
+
+def test_cpu_bundle_into_empty_venv_installs_normally(tmp_path):
+    """No GPU build anywhere (e.g. transcription alone): nothing to protect."""
+    installer = load_installer()
+    staging = tmp_path / "staging"
+    venv_sp = tmp_path / "venv-sp"
+    staging.mkdir()
+    venv_sp.mkdir()
+    make_onnxruntime(staging, "cpu")
+
+    installer.reconcile_onnxruntime(str(staging), str(venv_sp))
+    installer.move_tree(str(staging), str(venv_sp))
+
+    assert (venv_sp / "onnxruntime" / "capi" / "onnxruntime_pybind11_state.so").read_text() == "cpu"
+    assert cpu_dist_infos(venv_sp) == ["onnxruntime-1.20.1.dist-info"]
+
+
+def test_cpu_reinstall_over_cpu_is_untouched(tmp_path):
+    """arm64 and CPU-only stacks: CPU-over-CPU merges are business as usual."""
+    installer = load_installer()
+    staging = tmp_path / "staging"
+    venv_sp = tmp_path / "venv-sp"
+    staging.mkdir()
+    venv_sp.mkdir()
+    make_onnxruntime(venv_sp, "cpu")
+    make_onnxruntime(staging, "cpu")
+
+    installer.reconcile_onnxruntime(str(staging), str(venv_sp))
+    installer.move_tree(str(staging), str(venv_sp))
+
+    assert (venv_sp / "onnxruntime" / "capi" / "onnxruntime_pybind11_state.so").read_text() == "cpu"
+    assert cpu_dist_infos(venv_sp) == ["onnxruntime-1.20.1.dist-info"]
+
+
+def test_gpu_reinstall_over_gpu_is_untouched(tmp_path):
+    installer = load_installer()
+    staging = tmp_path / "staging"
+    venv_sp = tmp_path / "venv-sp"
+    staging.mkdir()
+    venv_sp.mkdir()
+    make_onnxruntime(venv_sp, "gpu")
+    make_onnxruntime(staging, "gpu")
+
+    installer.reconcile_onnxruntime(str(staging), str(venv_sp))
+    installer.move_tree(str(staging), str(venv_sp))
+
+    assert (venv_sp / "onnxruntime" / "capi" / CUDA_PROVIDER_LIB).exists()
+    assert gpu_dist_infos(venv_sp) == ["onnxruntime_gpu-1.20.1.dist-info"]
+
+
+def test_staging_with_both_flavors_ships_gpu_metadata_only(tmp_path):
+    """Defensive: a bundle built with both flavors (GPU installed last, so the
+    package files are the GPU build) must not ship the stale CPU metadata."""
+    installer = load_installer()
+    staging = tmp_path / "staging"
+    venv_sp = tmp_path / "venv-sp"
+    staging.mkdir()
+    venv_sp.mkdir()
+    make_onnxruntime(staging, "gpu")
+    # Simulate the leftover CPU dist-info next to the GPU files.
+    cpu_meta = staging / "onnxruntime-1.20.1.dist-info"
+    cpu_meta.mkdir()
+    (cpu_meta / "METADATA").write_text("Name: onnxruntime")
+
+    installer.reconcile_onnxruntime(str(staging), str(venv_sp))
+    installer.move_tree(str(staging), str(venv_sp))
+
+    # The GPU package files ship; the stale CPU metadata does not.
+    assert (venv_sp / "onnxruntime" / "capi" / CUDA_PROVIDER_LIB).exists()
+    assert cpu_dist_infos(venv_sp) == []
+    assert gpu_dist_infos(venv_sp) == ["onnxruntime_gpu-1.20.1.dist-info"]