mirror of
https://github.com/snapotter-hq/SnapOtter.git
synced 2026-08-03 07:46:42 +02:00
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
This commit is contained in:
@@ -252,6 +252,38 @@ for patch in patches:
|
|||||||
print(" Source patches applied")
|
print(" Source patches applied")
|
||||||
PYPATCH
|
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 ──────────────────────────────────────────────
|
# ── Step 5: Download models ──────────────────────────────────────────────
|
||||||
echo "=== Downloading models ==="
|
echo "=== Downloading models ==="
|
||||||
python3 << 'PYMODELS'
|
python3 << 'PYMODELS'
|
||||||
|
|||||||
@@ -77,7 +77,28 @@ for bundle_id in ${BUNDLE_IDS}; do
|
|||||||
mkdir -p "${staging}"
|
mkdir -p "${staging}"
|
||||||
tar -xzf "${archive}" -C "${staging}"
|
tar -xzf "${archive}" -C "${staging}"
|
||||||
if [[ -d "${staging}/site-packages" ]]; then
|
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
|
fi
|
||||||
rm -rf "${staging}"
|
rm -rf "${staging}"
|
||||||
echo " Installed ${bundle_id}"
|
echo " Installed ${bundle_id}"
|
||||||
@@ -166,6 +187,37 @@ if [[ $? -ne 0 ]]; then
|
|||||||
fail "Constrained-package consistency check failed -- see violations above" 2
|
fail "Constrained-package consistency check failed -- see violations above" 2
|
||||||
fi
|
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"
|
pass "All bundles for ${ARCH} are mutually compatible"
|
||||||
log "Done"
|
log "Done"
|
||||||
exit 0
|
exit 0
|
||||||
|
|||||||
@@ -168,8 +168,10 @@ def onnx_providers():
|
|||||||
available = _ort.get_available_providers()
|
available = _ort.get_available_providers()
|
||||||
if "CUDAExecutionProvider" in available:
|
if "CUDAExecutionProvider" in available:
|
||||||
return (["CUDAExecutionProvider", "CPUExecutionProvider"], "cuda")
|
return (["CUDAExecutionProvider", "CPUExecutionProvider"], "cuda")
|
||||||
emit_info("GPU detected by torch but CUDAExecutionProvider not available in onnxruntime "
|
emit_info("GPU detected by torch but CUDAExecutionProvider is missing from onnxruntime. "
|
||||||
"-- install onnxruntime-gpu for GPU acceleration")
|
"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:
|
except ImportError:
|
||||||
emit_info("onnxruntime not installed, cannot check CUDA provider")
|
emit_info("onnxruntime not installed, cannot check CUDA provider")
|
||||||
emit_info("No GPU detected, processing on CPU")
|
emit_info("No GPU detected, processing on CPU")
|
||||||
|
|||||||
@@ -540,6 +540,70 @@ def move_tree(src: str, dst: str) -> None:
|
|||||||
shutil.rmtree(src, ignore_errors=True)
|
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) --
|
# -- Fixups (NCCL wheel) --
|
||||||
|
|
||||||
def apply_fixups(staging_dir: str, venv_path: str) -> None:
|
def apply_fixups(staging_dir: str, venv_path: str) -> None:
|
||||||
@@ -819,6 +883,7 @@ def _install() -> None:
|
|||||||
},
|
},
|
||||||
mf,
|
mf,
|
||||||
)
|
)
|
||||||
|
reconcile_onnxruntime(staging_sp, site_packages_dir)
|
||||||
move_tree(staging_sp, site_packages_dir)
|
move_tree(staging_sp, site_packages_dir)
|
||||||
if os.path.exists(venv_writing_marker):
|
if os.path.exists(venv_writing_marker):
|
||||||
os.unlink(venv_writing_marker)
|
os.unlink(venv_writing_marker)
|
||||||
|
|||||||
@@ -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"]
|
||||||
@@ -241,3 +241,103 @@ describe("install_feature.py prebuilt mode", () => {
|
|||||||
expect(installed.bundles["face-detection"]).toBeDefined();
|
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);
|
||||||
|
});
|
||||||
|
});
|
||||||
|
|||||||
Reference in New Issue
Block a user