diff --git a/py/src/braintrust/integrations/base.py b/py/src/braintrust/integrations/base.py index 127ed893..66e708ef 100644 --- a/py/src/braintrust/integrations/base.py +++ b/py/src/braintrust/integrations/base.py @@ -627,21 +627,77 @@ def detect_version(cls, module: Any) -> str | None: return detect_module_version(module, cls.import_names) +# Module names that raised ImportError once. Bounded by the patchers' static +# target_module attributes. +_UNIMPORTABLE_MODULES: set[str] = set() + +# CPython < 3.12 records "which lock is this thread blocked on" in a single +# dict slot per thread (``importlib._bootstrap._blocking_on[tid]``) and clears +# it unconditionally in ``_ModuleLock.acquire``'s finally. When an import +# nests on one thread -- a meta-path finder importing something while +# resolving a submodule, say -- the inner frame clears the slot and the outer +# frame's cleanup raises ``KeyError: ``. It is a bug in the +# interpreter, not in the module being imported. +_IMPORT_LOCK_BOOKKEEPING_IS_BUGGY = sys.version_info < (3, 12) + + +def _import_module_tolerating_lock_bug(name: str) -> Any: + """``importlib.import_module`` that retries the interpreter's lock bug. + + Only the spurious ``KeyError`` is handled; ``ImportError`` and every other + failure propagate untouched, so a genuinely absent module is still + reported as absent rather than papered over. + """ + try: + return importlib.import_module(name) + except KeyError: + if not _IMPORT_LOCK_BOOKKEEPING_IS_BUGGY: + raise + # The bookkeeping is per-call, so a second attempt normally gets a clean + # slot. If the interpreter trips again, report the module as unavailable: + # setup() runs inside the caller's application, and skipping one optional + # patch target beats raising an interpreter-internal KeyError at them. + try: + return importlib.import_module(name) + except KeyError: + if not _IMPORT_LOCK_BOOKKEEPING_IS_BUGGY: + raise + return None + + def _import_optional_module(name: str) -> Any | None: """Return the named module, or ``None`` when it cannot be imported. - ``sys.modules`` is consulted first so an already-imported module never - reacquires the import lock. Patcher resolution calls this for every - patcher on every ``setup()``, and on CPython 3.10 that lock traffic can - trip the interpreter's own re-entrancy bookkeeping, surfacing as - ``KeyError: `` from ``importlib._bootstrap``. + A module that is already imported is returned straight from + ``sys.modules`` so it never reacquires the import lock. Patcher resolution + calls this for every patcher on every ``setup()``, and CPython 3.10's + ``_find_and_load`` locks unconditionally, so that traffic can trip the + interpreter's own re-entrancy bookkeeping and surface as + ``KeyError: `` from ``importlib._bootstrap``. (3.11+ added this + same shortcut upstream.) + + The ``_initializing`` check is what makes the shortcut safe: a module whose + body is still executing is already in ``sys.modules`` but does not have its + attributes yet, so returning it would make patch targets look absent and + silently skip instrumentation. Falling through to ``import_module`` blocks + on the per-module lock until the other thread finishes. This mirrors the + predicate CPython uses for the same decision. """ module = sys.modules.get(name) - if module is not None: + if module is not None and not getattr(getattr(module, "__spec__", None), "_initializing", False): return module + if name in _UNIMPORTABLE_MODULES: + return None try: - return importlib.import_module(name) + return _import_module_tolerating_lock_bug(name) except ImportError: + # Patchers carry target modules for layouts that only some versions of + # a provider ship, so a miss is normal and permanent -- mistralai 2.x, + # for example, leaves seven of them unimportable and two setup() calls + # retry them ~96 times between them. Each retry takes the import lock, + # which is what CPython 3.10 trips over. The sys.modules check above + # runs first, so a module that genuinely shows up later is still found. + _UNIMPORTABLE_MODULES.add(name) return None diff --git a/py/src/braintrust/integrations/test_base.py b/py/src/braintrust/integrations/test_base.py index 00b8495f..eba6538d 100644 --- a/py/src/braintrust/integrations/test_base.py +++ b/py/src/braintrust/integrations/test_base.py @@ -1,6 +1,10 @@ +import importlib.machinery +import importlib.util import sys import types +import pytest +from braintrust.integrations import base from braintrust.integrations.base import _import_optional_module, _resolve_attr_path @@ -70,6 +74,129 @@ def explode(name): # pragma: no cover - must never be reached assert _import_optional_module("braintrust_fake_sdk") is sentinel -def test_import_optional_module_imports_when_absent(): +def test_import_optional_module_imports_when_absent(monkeypatch): + monkeypatch.setattr("braintrust.integrations.base._UNIMPORTABLE_MODULES", set()) assert _import_optional_module("json") is sys.modules["json"] assert _import_optional_module("braintrust_module_that_does_not_exist") is None + + +def test_import_optional_module_waits_for_initializing_module(monkeypatch): + """A half-built module must not short-circuit the import machinery. + + The loader puts a module in sys.modules *before* running its body, so + during a concurrent import the attributes a patcher looks for do not + exist yet. Returning it would make the target look absent and silently + skip instrumentation, so we must fall through and let import_module + block on the per-module lock. + """ + # module_from_spec wires up __spec__ the way the real loader does, so the + # flag can be flipped on the spec itself. + spec = importlib.machinery.ModuleSpec("braintrust_partial_sdk", loader=None) + spec._initializing = True + partial = importlib.util.module_from_spec(spec) + monkeypatch.setitem(sys.modules, "braintrust_partial_sdk", partial) + + finished = types.ModuleType("braintrust_partial_sdk") + monkeypatch.setattr( + "braintrust.integrations.base.importlib.import_module", + lambda name: finished, + ) + + assert _import_optional_module("braintrust_partial_sdk") is finished + + # Once initialization completes the shortcut applies again. + spec._initializing = False + assert _import_optional_module("braintrust_partial_sdk") is partial + + +def _raise_import_error(name): + raise ImportError(name) + + +def test_import_optional_module_caches_failures(monkeypatch): + """A permanently-absent target must not be retried on every setup(). + + Patchers carry target modules for provider layouts that only some + versions ship, so misses are normal -- and each retry takes the import + lock, which is what CPython 3.10 trips over. + """ + monkeypatch.setattr("braintrust.integrations.base._UNIMPORTABLE_MODULES", set()) + attempts = [] + + def failing(name): + attempts.append(name) + raise ImportError(name) + + monkeypatch.setattr("braintrust.integrations.base.importlib.import_module", failing) + + assert _import_optional_module("braintrust_absent_sdk") is None + assert _import_optional_module("braintrust_absent_sdk") is None + assert _import_optional_module("braintrust_absent_sdk") is None + assert attempts == ["braintrust_absent_sdk"] + + +def test_import_optional_module_still_sees_a_late_arrival(monkeypatch): + """A cached failure must not hide a module that shows up afterwards.""" + monkeypatch.setattr("braintrust.integrations.base._UNIMPORTABLE_MODULES", set()) + monkeypatch.setattr("braintrust.integrations.base.importlib.import_module", _raise_import_error) + + assert _import_optional_module("braintrust_late_sdk") is None + + late = types.ModuleType("braintrust_late_sdk") + monkeypatch.setitem(sys.modules, "braintrust_late_sdk", late) + assert _import_optional_module("braintrust_late_sdk") is late + + +def test_import_optional_module_survives_the_import_lock_bug(monkeypatch): + """CPython < 3.12 can raise KeyError from its own import bookkeeping. + + _blocking_on holds one slot per thread and acquire() clears it + unconditionally, so a nested import makes the outer frame's cleanup raise + KeyError: . That is an interpreter bug, not a signal about the + module, and it must not escape into the caller's setup(). + """ + monkeypatch.setattr("braintrust.integrations.base._UNIMPORTABLE_MODULES", set()) + monkeypatch.setattr("braintrust.integrations.base._IMPORT_LOCK_BOOKKEEPING_IS_BUGGY", True) + + recovered = types.ModuleType("braintrust_locky_sdk") + calls = [] + + def flaky(name): + calls.append(name) + if len(calls) == 1: + raise KeyError(140000000000000) + return recovered + + monkeypatch.setattr("braintrust.integrations.base.importlib.import_module", flaky) + + assert _import_optional_module("braintrust_locky_sdk") is recovered + assert len(calls) == 2 + + +def test_import_optional_module_gives_up_if_the_lock_bug_persists(monkeypatch): + """A second KeyError reports the module absent rather than crashing.""" + monkeypatch.setattr("braintrust.integrations.base._UNIMPORTABLE_MODULES", set()) + monkeypatch.setattr("braintrust.integrations.base._IMPORT_LOCK_BOOKKEEPING_IS_BUGGY", True) + + def always_keyerror(name): + raise KeyError(140000000000000) + + monkeypatch.setattr("braintrust.integrations.base.importlib.import_module", always_keyerror) + + assert _import_optional_module("braintrust_locky_sdk") is None + # Not cached as unimportable: the import never actually resolved. + assert "braintrust_locky_sdk" not in base._UNIMPORTABLE_MODULES + + +def test_import_optional_module_propagates_keyerror_on_fixed_interpreters(monkeypatch): + """On 3.12+ a KeyError is a real error and must not be swallowed.""" + monkeypatch.setattr("braintrust.integrations.base._UNIMPORTABLE_MODULES", set()) + monkeypatch.setattr("braintrust.integrations.base._IMPORT_LOCK_BOOKKEEPING_IS_BUGGY", False) + + def boom(name): + raise KeyError("a real bug in the module's import") + + monkeypatch.setattr("braintrust.integrations.base.importlib.import_module", boom) + + with pytest.raises(KeyError): + _import_optional_module("braintrust_locky_sdk")