diff --git a/docker/build-bundle.sh b/docker/build-bundle.sh index 44529099..b71828cf 100755 --- a/docker/build-bundle.sh +++ b/docker/build-bundle.sh @@ -252,6 +252,38 @@ for patch in patches: print(" Source patches applied") PYPATCH +# ── Step 4.6: ONNX Runtime flavor hygiene ──────────────────────────────── +# A pinned onnxruntime-gpu plus a transitive CPU onnxruntime (rembg and +# faster-whisper both depend on plain `onnxruntime`) leaves BOTH dist-infos in +# site-packages, with the package files belonging to whichever installed last. +# Force the GPU build's files back and drop the CPU metadata, so a bundle can +# never ship CPU files under GPU metadata, or two flavors at once (#490). +echo "=== Reconciling ONNX Runtime flavor ===" +python3 << 'PYONNX' +import os, shlex, shutil, site, subprocess, sys + +site_packages = site.getsitepackages()[0] +names = os.listdir(site_packages) +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")) + +if not (cpu and gpu): + print(" Single ONNX Runtime flavor (or none); nothing to reconcile") + sys.exit(0) + +version = gpu[-1][len("onnxruntime_gpu-"):-len(".dist-info")] +cmd = (f"{sys.executable} -m pip install --no-cache-dir --force-reinstall " + f"--no-deps onnxruntime-gpu=={version}") +print(f" > {cmd}", flush=True) +if subprocess.run(shlex.split(cmd)).returncode != 0: + print("ERROR: onnxruntime-gpu reinstall failed", file=sys.stderr) + sys.exit(1) + +for name in cpu: + shutil.rmtree(os.path.join(site_packages, name), ignore_errors=True) +print(f" Kept onnxruntime-gpu=={version}; removed CPU dist-info: {', '.join(cpu)}") +PYONNX + # ── Step 5: Download models ────────────────────────────────────────────── echo "=== Downloading models ===" python3 << 'PYMODELS' diff --git a/docker/verify-bundle-compatibility.sh b/docker/verify-bundle-compatibility.sh index 624e434a..a4ca8279 100755 --- a/docker/verify-bundle-compatibility.sh +++ b/docker/verify-bundle-compatibility.sh @@ -77,7 +77,28 @@ for bundle_id in ${BUNDLE_IDS}; do mkdir -p "${staging}" tar -xzf "${archive}" -C "${staging}" if [[ -d "${staging}/site-packages" ]]; then - cp -a "${staging}/site-packages/." "${SITE_PACKAGES}/" + # Layer through the REAL installer merge (reconcile + move_tree), not a raw + # cp: install_feature.py reconciles the ONNX Runtime flavor between bundles + # (#490), and this script must verify the state users actually end up with. + # INSTALL_FEATURE_PY can point at a repo checkout mounted over an older + # image (same pattern as the install_runtime.py mounts in ai-bundles.yml). + "${VENV}/bin/python3" - "${staging}/site-packages" "${SITE_PACKAGES}" <<'PYMERGE' +import importlib.util, os, sys + +script = os.environ.get("INSTALL_FEATURE_PY", "/app/packages/ai/python/install_feature.py") +spec = importlib.util.spec_from_file_location("installer_under_verify", script) +mod = importlib.util.module_from_spec(spec) +spec.loader.exec_module(mod) +if not hasattr(mod, "reconcile_onnxruntime"): + print(f"FAIL: {script} has no reconcile_onnxruntime; mount the current " + "packages/ai/python/install_feature.py into the container " + "(-v $PWD/packages/ai/python/install_feature.py:/app/packages/ai/python/install_feature.py:ro)", + file=sys.stderr) + sys.exit(1) +staging_sp, site_packages = sys.argv[1], sys.argv[2] +mod.reconcile_onnxruntime(staging_sp, site_packages) +mod.move_tree(staging_sp, site_packages) +PYMERGE fi rm -rf "${staging}" echo " Installed ${bundle_id}" @@ -166,6 +187,37 @@ if [[ $? -ne 0 ]]; then fail "Constrained-package consistency check failed -- see violations above" 2 fi +log "Checking ONNX Runtime flavor consistency" + +"${VENV}/bin/python3" - "${SITE_PACKAGES}" <<'PYFLAVOR' || fail "ONNX Runtime flavor check failed" 2 +import os, sys + +sp = sys.argv[1] +names = os.listdir(sp) +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")) + +if cpu and gpu: + print(f"FAIL: both ONNX Runtime flavors present after layering: {cpu + gpu}. " + "A bundle is shipping the CPU build alongside another bundle's GPU build; " + "the merged venv's flavor then depends on install order (#490).", + file=sys.stderr) + sys.exit(1) + +if gpu: + cuda_lib = os.path.join(sp, "onnxruntime", "capi", "libonnxruntime_providers_cuda.so") + if not os.path.exists(cuda_lib): + print("FAIL: onnxruntime_gpu metadata is present but the CUDA provider library " + "is missing: a CPU build clobbered the GPU build's files (#490).", + file=sys.stderr) + sys.exit(1) + print(f" ONNX Runtime flavor: GPU ({gpu[-1]}), CUDA provider library present") +elif cpu: + print(f" ONNX Runtime flavor: CPU ({cpu[-1]})") +else: + print(" ONNX Runtime not present in any bundle") +PYFLAVOR + pass "All bundles for ${ARCH} are mutually compatible" log "Done" exit 0 diff --git a/packages/ai/python/gpu.py b/packages/ai/python/gpu.py index c51e8620..a79647d6 100644 --- a/packages/ai/python/gpu.py +++ b/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") diff --git a/packages/ai/python/install_feature.py b/packages/ai/python/install_feature.py index 489c81d2..6f88ea0d 100644 --- a/packages/ai/python/install_feature.py +++ b/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 `-.dist-info`, and name + normalization turns every other onnxruntime distribution into + `onnxruntime_` (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) diff --git a/packages/ai/python/tests/test_install_feature_onnx_reconcile.py b/packages/ai/python/tests/test_install_feature_onnx_reconcile.py new file mode 100644 index 00000000..db6bde73 --- /dev/null +++ b/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"] diff --git a/tests/unit/features/install-feature-prebuilt.test.ts b/tests/unit/features/install-feature-prebuilt.test.ts index cc75b9a7..d4d271a1 100644 --- a/tests/unit/features/install-feature-prebuilt.test.ts +++ b/tests/unit/features/install-feature-prebuilt.test.ts @@ -241,3 +241,103 @@ describe("install_feature.py prebuilt mode", () => { expect(installed.bundles["face-detection"]).toBeDefined(); }); }); + +/** Build a tar whose site-packages carries a fake onnxruntime of the given flavor. */ +function createOnnxTar( + bundleId: string, + flavor: "cpu" | "gpu", +): { tarPath: string; sha256: string } { + const buildDir = join(tempDir, `build-${bundleId}`); + const capi = join(buildDir, "site-packages", "onnxruntime", "capi"); + mkdirSync(capi, { recursive: true }); + writeFileSync(join(buildDir, "site-packages", "onnxruntime", "__init__.py"), flavor); + writeFileSync(join(capi, "onnxruntime_pybind11_state.so"), flavor); + if (flavor === "gpu") { + writeFileSync(join(capi, "libonnxruntime_providers_cuda.so"), "cuda"); + } + const distInfo = + flavor === "gpu" ? "onnxruntime_gpu-1.20.1.dist-info" : "onnxruntime-1.20.1.dist-info"; + mkdirSync(join(buildDir, "site-packages", distInfo), { recursive: true }); + writeFileSync( + join(buildDir, "site-packages", distInfo, "METADATA"), + flavor === "gpu" ? "Name: onnxruntime-gpu" : "Name: onnxruntime", + ); + writeFileSync( + join(buildDir, "bundle.json"), + JSON.stringify({ + bundleId, + version: "1.0.0-test", + arch: "amd64-gpu", + imageVersion: "2.0.0", + pythonVersion: "3.12", + models: [], + }), + ); + + const tarPath = join(tempDir, `${bundleId}-test.tar.gz`); + execFileSync("tar", ["czf", tarPath, "-C", buildDir, "."]); + rmSync(buildDir, { recursive: true }); + + const hash = createHash("sha256").update(readFileSync(tarPath)).digest("hex"); + return { tarPath, sha256: hash }; +} + +function installBundle(bundleId: string, tarPath: string) { + return spawnSync("python3", [scriptPath, bundleId, manifestPath, modelsDir], { + env: { + ...process.env, + DATA_DIR: tempDir, + PYTHON_VENV_PATH: venvDir, + SNAPOTTER_BUNDLE_LOCAL_PATH: tarPath, + }, + timeout: 30_000, + }); +} + +describe("onnxruntime flavor reconciliation (#490)", () => { + const cudaLib = () => + join(sitePackagesDir, "onnxruntime", "capi", "libonnxruntime_providers_cuda.so"); + const coreLib = () => + join(sitePackagesDir, "onnxruntime", "capi", "onnxruntime_pybind11_state.so"); + + it("a bundle carrying CPU onnxruntime cannot clobber the venv's GPU build", () => { + const gpu = createOnnxTar("gpu-bundle", "gpu"); + writeManifest("gpu-bundle", gpu.tarPath, gpu.sha256, { models: [] }); + let result = installBundle("gpu-bundle", gpu.tarPath); + expect(result.status, `stderr: ${result.stderr?.toString()}`).toBe(0); + + const cpu = createOnnxTar("cpu-bundle", "cpu"); + writeManifest("cpu-bundle", cpu.tarPath, cpu.sha256, { models: [] }); + result = installBundle("cpu-bundle", cpu.tarPath); + expect(result.status, `stderr: ${result.stderr?.toString()}`).toBe(0); + + // The GPU build survives: CUDA provider intact, core lib not downgraded. + expect(existsSync(cudaLib())).toBe(true); + expect(readFileSync(coreLib(), "utf-8")).toBe("gpu"); + // Metadata stays truthful: no CPU dist-info shipped in. + expect(existsSync(join(sitePackagesDir, "onnxruntime-1.20.1.dist-info"))).toBe(false); + expect(existsSync(join(sitePackagesDir, "onnxruntime_gpu-1.20.1.dist-info"))).toBe(true); + // Both bundles still recorded as installed. + const installed = JSON.parse(readFileSync(join(aiDir, "installed.json"), "utf-8")); + expect(installed.bundles["gpu-bundle"]).toBeDefined(); + expect(installed.bundles["cpu-bundle"]).toBeDefined(); + }); + + it("installing a GPU bundle repairs a venv previously downgraded to the CPU build", () => { + const cpu = createOnnxTar("cpu-bundle", "cpu"); + writeManifest("cpu-bundle", cpu.tarPath, cpu.sha256, { models: [] }); + let result = installBundle("cpu-bundle", cpu.tarPath); + expect(result.status, `stderr: ${result.stderr?.toString()}`).toBe(0); + + const gpu = createOnnxTar("gpu-bundle", "gpu"); + writeManifest("gpu-bundle", gpu.tarPath, gpu.sha256, { models: [] }); + result = installBundle("gpu-bundle", gpu.tarPath); + expect(result.status, `stderr: ${result.stderr?.toString()}`).toBe(0); + + expect(existsSync(cudaLib())).toBe(true); + expect(readFileSync(coreLib(), "utf-8")).toBe("gpu"); + // The stale CPU metadata is cleared so pip metadata matches reality. + expect(existsSync(join(sitePackagesDir, "onnxruntime-1.20.1.dist-info"))).toBe(false); + expect(existsSync(join(sitePackagesDir, "onnxruntime_gpu-1.20.1.dist-info"))).toBe(true); + }); +});