fix(integrations): don't shortcut a module that is still initializing

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: <thread id>``. 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: <thread id>`` 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: <thread id>`` 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
 
 

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: <thread id>. 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")