Skip to content

Commit 94cb476

Browse files
BF667BF667
committed
fix: faiss AVX512/AVX2 import — use FAISS_DISABLE_CPU_FEATURES to skip AVX
Instead of catching ModuleNotFoundError after the fact, set FAISS_DISABLE_CPU_FEATURES=AVX512,AVX2,AVX512_SPR,SVE before importing faiss. This tells faiss's loader.py to skip straight to the basic swigfaiss module, eliminating noisy INFO logs and avoiding the crash when AVX native modules are missing.
1 parent 58ac85f commit 94cb476

4 files changed

Lines changed: 97 additions & 54 deletions

File tree

‎arvc/engine/models/utils.py‎

Lines changed: 54 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -11,37 +11,62 @@
1111

1212

1313
# ── Safe FAISS import ────────────────────────────────────────────────────
14-
# faiss-cpu may fail with ModuleNotFoundError("No module named 'faiss.swigfaiss_avx2'")
15-
# or "No module named 'faiss.swigfaiss_avx512'" when the AVX-optimized native
16-
# module isn't available (e.g. older CPUs, some cloud VMs).
17-
# Fall back to the non-AVX version gracefully.
18-
try:
19-
import faiss
20-
except ModuleNotFoundError as _faiss_err:
21-
_faiss_msg = str(_faiss_err)
22-
if "swigfaiss_avx" in _faiss_msg:
23-
import warnings
24-
warnings.warn(
25-
f"{_faiss_msg} — falling back to non-AVX faiss. "
26-
"If you need AVX support, install a compatible faiss-cpu wheel.",
27-
stacklevel=2,
28-
)
29-
# Force faiss to use the fallback (non-AVX) implementation
30-
try:
31-
import faiss.loader as _faiss_loader
32-
if hasattr(_faiss_loader, "toggle_swigfaiss_avx2"):
33-
_faiss_loader.toggle_swigfaiss_avx2 = False
34-
except Exception:
35-
pass
36-
# Clear cached failed import so faiss can retry
37-
if "faiss" in sys.modules:
38-
del sys.modules["faiss"]
39-
# Set env var that tells faiss to skip AVX
40-
os.environ["FAISS_NO_AVX2"] = "1"
41-
import faiss
42-
else:
14+
# faiss-cpu tries AVX512 → AVX2 → basic swigfaiss at import time.
15+
# When AVX native modules are missing it logs INFO messages and falls back,
16+
# but on some systems the ModuleNotFoundError propagates instead of being
17+
# caught internally. We handle both cases:
18+
# 1. Pre-disable AVX features via FAISS_DISABLE_CPU_FEATURES so faiss
19+
# never attempts the AVX path (avoids noisy INFO logs).
20+
# 2. If import still fails with a swigfaiss_avx* error, force the basic
21+
# loader and retry.
22+
def _import_faiss():
23+
"""Import faiss with safe AVX fallback."""
24+
# If faiss is already loaded, return it
25+
if "faiss" in sys.modules:
26+
return sys.modules["faiss"]
27+
28+
# First try: disable AVX features so faiss skips straight to basic
29+
_prev = os.environ.get("FAISS_DISABLE_CPU_FEATURES", "")
30+
os.environ["FAISS_DISABLE_CPU_FEATURES"] = "AVX512, AVX2, AVX512_SPR, SVE"
31+
try:
32+
import faiss as _faiss
33+
return _faiss
34+
except ModuleNotFoundError:
35+
pass
36+
finally:
37+
if _prev:
38+
os.environ["FAISS_DISABLE_CPU_FEATURES"] = _prev
39+
else:
40+
os.environ.pop("FAISS_DISABLE_CPU_FEATURES", None)
41+
42+
# Second try: clear any partial import and retry with env set
43+
for _mod in list(sys.modules.keys()):
44+
if _mod.startswith("faiss"):
45+
del sys.modules[_mod]
46+
os.environ["FAISS_DISABLE_CPU_FEATURES"] = "AVX512, AVX2, AVX512_SPR, SVE"
47+
try:
48+
import faiss as _faiss
49+
return _faiss
50+
except ModuleNotFoundError as _faiss_err:
51+
_faiss_msg = str(_faiss_err)
52+
if "swigfaiss" in _faiss_msg:
53+
import warnings
54+
warnings.warn(
55+
f"{_faiss_msg} — faiss AVX modules unavailable. "
56+
"Install a compatible faiss-cpu wheel or set FAISS_OPT_LEVEL.",
57+
stacklevel=3,
58+
)
59+
# Last resort: try to force the basic loader
60+
for _mod in list(sys.modules.keys()):
61+
if _mod.startswith("faiss"):
62+
del sys.modules[_mod]
63+
os.environ["FAISS_OPT_LEVEL"] = ""
64+
import faiss as _faiss
65+
return _faiss
4366
raise
4467

68+
faiss = _import_faiss()
69+
4570

4671
from arvc.utils import huggingface
4772
from arvc.utils.variables import translations, configs, config, logger, embedders_model, spin_model, whisper_model

‎arvc/engine/training/create_index.py‎

Lines changed: 43 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -8,33 +8,51 @@
88
from sklearn.cluster import MiniBatchKMeans
99

1010
# ── Safe FAISS import ────────────────────────────────────────────────────
11-
# faiss-cpu may fail with ModuleNotFoundError for swigfaiss_avx2 or
12-
# swigfaiss_avx512 when the AVX-optimized native module isn't available.
13-
# Fall back to the non-AVX version gracefully.
14-
try:
15-
import faiss
16-
except ModuleNotFoundError as _faiss_err:
17-
_faiss_msg = str(_faiss_err)
18-
if "swigfaiss_avx" in _faiss_msg:
19-
import warnings
20-
warnings.warn(
21-
f"{_faiss_msg} — falling back to non-AVX faiss. "
22-
"If you need AVX support, install a compatible faiss-cpu wheel.",
23-
stacklevel=2,
24-
)
25-
try:
26-
import faiss.loader as _faiss_loader
27-
if hasattr(_faiss_loader, "toggle_swigfaiss_avx2"):
28-
_faiss_loader.toggle_swigfaiss_avx2 = False
29-
except Exception:
30-
pass
31-
if "faiss" in sys.modules:
32-
del sys.modules["faiss"]
33-
os.environ["FAISS_NO_AVX2"] = "1"
34-
import faiss
35-
else:
11+
# faiss-cpu tries AVX512 → AVX2 → basic swigfaiss at import time.
12+
# Pre-disable AVX features via FAISS_DISABLE_CPU_FEATURES so faiss
13+
# skips straight to basic and avoids noisy INFO/warning logs.
14+
def _import_faiss():
15+
"""Import faiss with safe AVX fallback."""
16+
if "faiss" in sys.modules:
17+
return sys.modules["faiss"]
18+
_prev = os.environ.get("FAISS_DISABLE_CPU_FEATURES", "")
19+
os.environ["FAISS_DISABLE_CPU_FEATURES"] = "AVX512, AVX2, AVX512_SPR, SVE"
20+
try:
21+
import faiss as _faiss
22+
return _faiss
23+
except ModuleNotFoundError:
24+
pass
25+
finally:
26+
if _prev:
27+
os.environ["FAISS_DISABLE_CPU_FEATURES"] = _prev
28+
else:
29+
os.environ.pop("FAISS_DISABLE_CPU_FEATURES", None)
30+
# Retry: clear partial import and force basic
31+
for _mod in list(sys.modules.keys()):
32+
if _mod.startswith("faiss"):
33+
del sys.modules[_mod]
34+
os.environ["FAISS_DISABLE_CPU_FEATURES"] = "AVX512, AVX2, AVX512_SPR, SVE"
35+
try:
36+
import faiss as _faiss
37+
return _faiss
38+
except ModuleNotFoundError as _faiss_err:
39+
if "swigfaiss" in str(_faiss_err):
40+
import warnings
41+
warnings.warn(
42+
f"{_faiss_err} — faiss AVX modules unavailable. "
43+
"Install a compatible faiss-cpu wheel.",
44+
stacklevel=3,
45+
)
46+
for _mod in list(sys.modules.keys()):
47+
if _mod.startswith("faiss"):
48+
del sys.modules[_mod]
49+
os.environ["FAISS_OPT_LEVEL"] = ""
50+
import faiss as _faiss
51+
return _faiss
3652
raise
3753

54+
faiss = _import_faiss()
55+
3856
# ── FIX: Ensure project root is in sys.path BEFORE any arvc imports ──
3957
_project_root = os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))))
4058
if _project_root not in sys.path:

‎docs/Advanced-RVC-Documentation.html‎

100644100755
File mode changed.

‎docs/Advanced-RVC-Documentation.pdf‎

100644100755
File mode changed.

0 commit comments

Comments
 (0)