Skip importlib for already-loaded modules in the workflow sandbox

CHANGELOG.md

@@ -35,6 +35,13 @@ to include examples, links to docs, or any other relevant information.
 
 ### Deprecated
 
+### Fixed
+
+- Sandboxed workflow imports of already-loaded modules no longer go through importlib's module
+  locks, fixing intermittent `Failed validating workflow` errors on Python 3.10 caused by a
+  `KeyError` in `importlib._bootstrap._ModuleLock.acquire` when a garbage-collection finalizer
+  imported `warnings` during a workflow load ([#585](https://github.com/temporalio/sdk-python/issues/585)).
+
 ### :boom: Breaking Changes
 
 - Experimental external storage: `ExternalStorage.driver_selector` is now called with a

temporalio/worker/workflow_sandbox/_importer.py

@@ -258,7 +258,9 @@ def _import(
                 sys.modules[full_name] = new_mod
                 new_spec.loader.exec_module(new_mod)
 
-        mod = importlib.__import__(name, globals, locals, fromlist, level)
+        mod = _already_imported(name, full_name, fromlist, level)
+        if mod is None:
+            mod = importlib.__import__(name, globals, locals, fromlist, level)
         # Check for restrictions if necessary and apply
         if mod.__name__ not in self.modules_checked_for_restrictions:
             self.modules_checked_for_restrictions.add(mod.__name__)
@@ -539,6 +541,36 @@ def _get_thread_local_builtin(name: str) -> _ThreadLocalCallable:
     return ret
 
 
+def _already_imported(
+    name: str, full_name: str, fromlist: Sequence[str], level: int
+) -> types.ModuleType | None:
+    # Mirrors importlib.__import__ for loaded modules without taking module locks
+    mod = _fully_imported(full_name)
+    if mod is None:
+        return None
+    if fromlist:
+        # Only statically stored attributes count; module __getattr__ stays with importlib
+        mod_dict = getattr(mod, "__dict__", None)
+        if not isinstance(mod_dict, dict):
+            return None
+        if "__path__" in mod_dict and any(
+            not isinstance(x, str) or x == "*" or x not in mod_dict for x in fromlist
+        ):
+            return None
+        return mod
+    if level != 0:
+        return None
+    top = name.partition(".")[0]
+    return mod if top == full_name else _fully_imported(top)
+
+
+def _fully_imported(name: str) -> types.ModuleType | None:
+    mod = sys.modules.get(name)
+    if mod is None or getattr(getattr(mod, "__spec__", None), "_initializing", False):
+        return None
+    return mod
+
+
 def _resolve_module_name(
     name: str, globals: Mapping[str, object] | None, level: int
 ) -> str:

tests/worker/workflow_sandbox/test_importer.py

@@ -1,5 +1,7 @@
 import dataclasses
+import importlib
 import sys
+from typing import Any
 
 import pytest
 
@@ -27,6 +29,65 @@ def test_workflow_sandbox_importer_invalid_module():
     )
 
 
+def test_workflow_sandbox_importer_repeat_import_skips_import_machinery(
+    monkeypatch: pytest.MonkeyPatch,
+):
+    imported: list[str] = []
+    orig_import = importlib.__import__
+
+    def recording_import(
+        name: str,
+        globals: Any = None,
+        locals: Any = None,
+        fromlist: Any = (),
+        level: int = 0,
+    ) -> Any:
+        imported.append(name)
+        return orig_import(name, globals, locals, fromlist, level)
+
+    monkeypatch.setattr(importlib, "__import__", recording_import)
+    with Importer(restrictions, RestrictionContext()).applied():
+        import tests.worker.workflow_sandbox.testmodules.passthrough_module as passthrough
+        import tests.worker.workflow_sandbox.testmodules.stateful_module as stateful
+
+        assert imported
+        imported.clear()
+
+        # Loaded modules are served from sys.modules without re-entering importlib
+        import typing
+
+        import tests.worker.workflow_sandbox.testmodules.passthrough_module as passthrough_again
+        import tests.worker.workflow_sandbox.testmodules.stateful_module as stateful_again
+        from tests.worker.workflow_sandbox import testmodules
+        from tests.worker.workflow_sandbox.testmodules import stateful_module
+
+        assert passthrough_again is passthrough
+        assert stateful_again is stateful is stateful_module
+        assert getattr(testmodules, "stateful_module") is stateful
+        assert typing is sys.modules["typing"]
+    assert imported == []
+
+
+def test_workflow_sandbox_importer_repeat_import_leaves_module_getattr_to_importlib():
+    pkg_name = "tests.worker.workflow_sandbox.testmodules.dynamic_attr_package"
+    with Importer(restrictions, RestrictionContext()).applied():
+        dyn_pkg = importlib.import_module(pkg_name)
+        assert dyn_pkg.dynamic_value == 42
+        before = len(dyn_pkg.getattr_calls)
+
+        # importlib's fromlist hasattr plus the attribute read, same as without the sandbox
+        pkg = __import__(pkg_name, fromlist=["dynamic_value"])
+        assert pkg.dynamic_value == 42
+        assert dyn_pkg.getattr_calls[before:] == ["dynamic_value", "dynamic_value"]
+
+        # A missing name is probed once by importlib and once by the read, not more
+        before = len(dyn_pkg.getattr_calls)
+        pkg = __import__(pkg_name, fromlist=["missing_value"])
+        with pytest.raises(AttributeError):
+            getattr(pkg, "missing_value")
+        assert dyn_pkg.getattr_calls[before:] == ["missing_value", "missing_value"]
+
+
 def test_workflow_sandbox_importer_passthrough_module():
     # Import outside of importer
     import tests.worker.workflow_sandbox.testmodules.passthrough_module as outside1

tests/worker/workflow_sandbox/testmodules/dynamic_attr_package/__init__.py

@@ -0,0 +1,8 @@
+getattr_calls: list[str] = []
+
+
+def __getattr__(name: str) -> int:
+    getattr_calls.append(name)
+    if name == "dynamic_value":
+        return 42
+    raise AttributeError(name)