Skip to content

Commit e64246b

Browse files
committed
fix(loader): detect linker-locked TF .so via /proc/self/maps; guard zombie purge
Problem: after loading TF 2.13.0 and exiting the context, a second with-block requesting TF 2.12.0 would still activate in-process, corrupting interpreter state and crashing. Two root causes fixed in loader.py: 1. /proc/self/maps fallback for cext detection The existing _tf_loaded_pids check requires omnipkg.isolation.patchers to be imported first. In bare subprocesses (test runner, user scripts) that module is never loaded, so _abi_so_mapped stayed False and the loader fell through to in-process activation of the wrong version. Fix: read /proc/self/maps directly — OS-level, survives sys.modules purging, no import dependencies. 2. Zombie purge was wiping linker-locked TF modules _aggressive_module_cleanup correctly preserved pywrap_tensorflow and friends, but the POST-EXIT zombie scan (which checks __file__ for .omnipkg_versions paths) then purged those same modules, leaving Load 2 with no Python layer for a still-mapped C backend — causing reinit failures. Fix: zombie purge now skips all tensorflow.* modules when pywrap_tensorflow is present in sys.modules, matching the same guard already used by _aggressive_module_cleanup. 3. Graceful no-daemon path When _abi_so_mapped=True but no daemon is available, execution previously fell through to Strategy 1/2/3 (in-process activation). Now returns self immediately with _abi_conflict_detected=True — caller gets the already-loaded version, no crash, no corrupt state. Also in this commit: - refactor(tests): move src/tests/ → src/omnipkg/tests/ so the test package lives inside the omnipkg namespace and does not pollute site-packages for end users - fix(pyproject.toml): update packages.find include to omnipkg.tests* - chore(.gitignore): exclude __pycache__/ and *.pyc
1 parent 981340d commit e64246b

20 files changed

Lines changed: 50 additions & 24 deletions

.gitignore

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -209,3 +209,8 @@ src/omnipkg/_vendor/wheels/
209209
.omnipkg_baseline
210210
ci-wheel-cache/
211211
build/
212+
213+
# Python cache
214+
__pycache__/
215+
*.pyc
216+
*.pyo

pyproject.toml

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -152,7 +152,7 @@ omnipkg = "omnipkg.dispatcher:main"
152152

153153
[tool.setuptools.packages.find]
154154
where = ["src"]
155-
include = ["omnipkg*", "omnipkg.*", "tests*"]
155+
include = ["omnipkg*", "omnipkg.*", "omnipkg.tests*"]
156156
exclude = [
157157
"build*",
158158
"dist*",

src/omnipkg/cli.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -32,7 +32,7 @@
3232
from .core import omnipkg as OmnipkgCore
3333

3434
project_root = Path(__file__).resolve().parent.parent
35-
TESTS_DIR = Path(__file__).parent.parent / "tests"
35+
TESTS_DIR = Path(__file__).parent / "tests"
3636
DEMO_DIR = Path(__file__).parent
3737
try:
3838
FILE_PATH = Path(__file__).resolve()

src/omnipkg/isolation/patchers.py

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -428,7 +428,6 @@ class DummyBackend:
428428

429429
# Only check for reload if THIS WORKER previously loaded TF successfully
430430
if current_pid in _tf_loaded_pids and "tensorflow" not in sys.modules:
431-
safe_print("☢️ [OMNIPKG] FATAL TENSORFLOW RELOAD DETECTED!")
432431
raise ProcessCorruptedException(
433432
"Attempted to reload TensorFlow in a process where its C++ libraries were already initialized."
434433
)

src/omnipkg/loader.py

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4085,8 +4085,13 @@ def __exit__(self, exc_type, exc_val, exc_tb):
40854085
# but they may leave modules from a different version in sys.modules.
40864086
# At depth=1 we own the final cleanup, so we do a full pass over
40874087
# the known C-extension switcher families.
4088+
_patchers_mod = sys.modules.get("omnipkg.isolation.patchers")
4089+
_tf_pids = getattr(_patchers_mod, "_tf_loaded_pids", set()) if _patchers_mod else set()
4090+
_tf_linker_locked = os.getpid() in _tf_pids
40884091
_CE_SWITCHERS = ("numpy", "torch", "scipy", "pandas", "tensorflow")
40894092
for _sw in _CE_SWITCHERS:
4093+
if _sw == "tensorflow" and _tf_linker_locked:
4094+
continue
40904095
_stale = [
40914096
k for k in list(sys.modules)
40924097
if k == _sw or k.startswith(_sw + ".")
@@ -5719,6 +5724,13 @@ def __enter__(self):
57195724
_pkg_lower in omnipkgLoader.ABI_PACKAGES
57205725
and _abi_indicator_map.get(_pkg_lower, f"{_pkg_lower}._") in sys.modules
57215726
)
5727+
# Also check _tf_loaded_pids — TF purges its modules on __exit__ but
5728+
# the .so stays mapped. sys.modules check alone misses this case.
5729+
if not _abi_so_mapped and _pkg_lower == "tensorflow":
5730+
_patchers = sys.modules.get("omnipkg.isolation.patchers")
5731+
_tf_pids = getattr(_patchers, "_tf_loaded_pids", set()) if _patchers else set()
5732+
if os.getpid() in _tf_pids:
5733+
_abi_so_mapped = True
57225734
if _abi_so_mapped and self._worker_fallback_enabled and DAEMON_AVAILABLE:
57235735
if not self.quiet:
57245736
safe_print(
File renamed without changes.

0 commit comments

Comments
 (0)