⚡️ Speed up method _ImportHookChainedLoader._load_module by 333%

[codeflash-ai]

📄 333% (3.33x) speedup for _ImportHookChainedLoader._load_module in wandb/sdk/lib/import_hooks.py

⏱️ Runtime : 543 microseconds → 126 microseconds (best of 101 runs)

📝 Explanation and details

The optimized code achieves a 332% speedup through three key micro-optimizations that eliminate expensive operations:

1. Eliminated Class Instantiation Overhead (84% reduction in _set_loader)

• Original: Created a new UNDEFINED class on every call (class UNDEFINED: pass)
• Optimized: Uses a single shared object() sentinel (_UNDEFINED = object())
• Impact: The line profiler shows the class definition took 642,143ns (84% of _set_loader time). Using object() reduces this to 38,253ns - a 94% reduction in sentinel creation overhead.

2. Reduced Dictionary Operations in Hook Processing

• Original: Used _post_import_hooks.pop(name, {}) which creates an empty dict when no hooks exist
• Optimized: Uses _post_import_hooks.pop(name, None) and early returns if no hooks
• Impact: Avoids unnecessary dictionary creation and iteration when no hooks are registered, which is common in many test cases.

3. Optimized Attribute Access Pattern

• Original: Combined getattr calls with complex boolean logic
• Optimized: Separates attribute lookups and uses intermediate variables to cache results
• Impact: Reduces redundant attribute access, particularly for __spec__ and its loader property.

Performance Results by Test Type:

• Basic module loading: 300-400% faster across all scenarios
• Large-scale operations: Maintains 300%+ speedup even with 500+ hooks/modules
• Edge cases: Particularly effective when modules lack __loader__ or __spec__ attributes

The optimizations are most effective for workloads with frequent module loading where hooks are rarely registered, as the code now short-circuits expensive operations early.

✅ Correctness verification report:

Test	Status
⚙️ Existing Unit Tests	🔘 None Found
🌀 Generated Regression Tests	✅ 159 Passed
⏪ Replay Tests	🔘 None Found
🔎 Concolic Coverage Tests	🔘 None Found
📊 Tests Coverage	80.0%

🌀 Generated Regression Tests and Runtime

import threading
import types
from typing import Any, Dict

# imports
import pytest
from wandb.sdk.lib.import_hooks import _ImportHookChainedLoader

_post_import_hooks: Dict = {}

# ===========================
# UNIT TESTS FOR _load_module
# ===========================

# Helper: Dummy loader with load_module
class DummyLoader:
    def __init__(self):
        self.called_with = []
        self.module_to_return = None

    def load_module(self, fullname):
        self.called_with.append(fullname)
        return self.module_to_return

# Helper: Dummy module object
def make_dummy_module(name="dummy_mod", with_loader=None, with_spec=False):
    mod = types.ModuleType(name)
    if with_loader is not None:
        mod.__loader__ = with_loader
    if with_spec:
        class Spec:
            def __init__(self, loader):
                self.loader = loader
        mod.__spec__ = Spec(None)
    return mod

# Helper: Dummy post-import hook
def make_hook():
    called = {}
    def hook(module):
        called['module'] = module
    return hook, called

# 1. BASIC TEST CASES

def test_load_module_basic_sets_loader_and_calls_hook():
    # Basic: loader returns a module, __loader__ is None, hook is called
    loader = DummyLoader()
    mod = make_dummy_module("mod1", with_loader=None)
    loader.module_to_return = mod
    chained = _ImportHookChainedLoader(loader)
    # Register a post-import hook
    hook, called = make_hook()
    _post_import_hooks[mod.__name__] = {"h1": hook}
    # Call _load_module
    codeflash_output = chained._load_module(mod.__name__); result = codeflash_output # 13.5μs -> 2.81μs (382% faster)

def test_load_module_no_post_import_hook():
    # Basic: no post-import hook registered, should not fail
    loader = DummyLoader()
    mod = make_dummy_module("mod_nohook")
    loader.module_to_return = mod
    chained = _ImportHookChainedLoader(loader)
    # No hook registered
    if mod.__name__ in _post_import_hooks:
        del _post_import_hooks[mod.__name__]
    codeflash_output = chained._load_module(mod.__name__); result = codeflash_output # 13.5μs -> 2.81μs (382% faster)

def test_load_module_preserves_existing_loader():
    # Basic: module already has __loader__ set to something else, should not overwrite
    loader = DummyLoader()
    sentinel_loader = object()
    mod = make_dummy_module("mod2", with_loader=sentinel_loader)
    loader.module_to_return = mod
    chained = _ImportHookChainedLoader(loader)
    codeflash_output = chained._load_module(mod.__name__); result = codeflash_output # 13.4μs -> 2.52μs (433% faster)

def test_load_module_sets_spec_loader_if_self():
    # Basic: __spec__.loader is self, should set to underlying loader
    loader = DummyLoader()
    mod = make_dummy_module("mod3", with_loader=None, with_spec=True)
    chained = _ImportHookChainedLoader(loader)
    mod.__spec__.loader = chained
    loader.module_to_return = mod
    chained._load_module(mod.__name__) # 10.6μs -> 2.95μs (261% faster)

def test_load_module_skips_setting_loader_on_attribute_error():
    # Basic: module does not allow setting __loader__ (simulate built-in)
    class DummyModule:
        __name__ = "mod4"
        def __setattr__(self, name, value):
            if name == "__loader__":
                raise AttributeError("cannot set")
            object.__setattr__(self, name, value)
    loader = DummyLoader()
    mod = DummyModule()
    loader.module_to_return = mod
    chained = _ImportHookChainedLoader(loader)
    # Should not raise
    codeflash_output = chained._load_module(mod.__name__); result = codeflash_output # 10.2μs -> 2.68μs (281% faster)

# 2. EDGE TEST CASES



def test_load_module_with_multiple_hooks():
    # Edge: multiple hooks, all called
    loader = DummyLoader()
    mod = make_dummy_module("mod7")
    loader.module_to_return = mod
    chained = _ImportHookChainedLoader(loader)
    hook1, called1 = make_hook()
    hook2, called2 = make_hook()
    _post_import_hooks[mod.__name__] = {"h1": hook1, "h2": hook2}
    chained._load_module(mod.__name__) # 13.8μs -> 3.02μs (357% faster)

def test_load_module_with_none_hook_skips():
    # Edge: hook value is None, should skip
    loader = DummyLoader()
    mod = make_dummy_module("mod8")
    loader.module_to_return = mod
    chained = _ImportHookChainedLoader(loader)
    # One hook is None
    hook, called = make_hook()
    _post_import_hooks[mod.__name__] = {"h1": None, "h2": hook}
    chained._load_module(mod.__name__) # 13.5μs -> 2.72μs (397% faster)

def test_load_module_with_no_loader_attribute():
    # Edge: module does not have __loader__ attribute at all
    loader = DummyLoader()
    mod = make_dummy_module("mod9")
    delattr(mod, "__loader__")
    loader.module_to_return = mod
    chained = _ImportHookChainedLoader(loader)
    chained._load_module(mod.__name__) # 14.9μs -> 4.01μs (270% faster)

def test_load_module_with_spec_but_no_loader():
    # Edge: __spec__ exists but no loader attribute
    loader = DummyLoader()
    mod = make_dummy_module("mod10", with_loader=None, with_spec=True)
    delattr(mod.__spec__, "loader")
    loader.module_to_return = mod
    chained = _ImportHookChainedLoader(loader)
    # Should not raise
    chained._load_module(mod.__name__) # 10.4μs -> 2.88μs (262% faster)

def test_load_module_with_strange_module_object():
    # Edge: module is not a real module, but an object
    class StrangeModule:
        __name__ = "mod11"
    loader = DummyLoader()
    mod = StrangeModule()
    loader.module_to_return = mod
    chained = _ImportHookChainedLoader(loader)
    chained._load_module(mod.__name__) # 10.2μs -> 2.73μs (272% faster)

# 3. LARGE SCALE TEST CASES

def test_load_module_many_hooks():
    # Large scale: many hooks registered for a module
    loader = DummyLoader()
    mod = make_dummy_module("mod_large")
    loader.module_to_return = mod
    chained = _ImportHookChainedLoader(loader)
    called = []
    def make_large_hook(i):
        def hook(module):
            called.append(i)
        return hook
    hooks = {f"h{i}": make_large_hook(i) for i in range(100)}
    _post_import_hooks[mod.__name__] = hooks
    chained._load_module(mod.__name__) # 13.9μs -> 2.74μs (409% faster)

def test_load_module_many_modules():
    # Large scale: load many modules in sequence
    for i in range(50):
        loader = DummyLoader()
        name = f"mod_seq_{i}"
        mod = make_dummy_module(name)
        loader.module_to_return = mod
        chained = _ImportHookChainedLoader(loader)
        hook, called = make_hook()
        _post_import_hooks[name] = {"h": hook}
        codeflash_output = chained._load_module(name); result = codeflash_output # 170μs -> 36.4μs (368% faster)

def test_load_module_performance_with_large_hooks(monkeypatch):
    # Large scale: test that _load_module does not degrade with many hooks
    # (not a strict performance test, but ensures it works with many)
    loader = DummyLoader()
    mod = make_dummy_module("mod_perf")
    loader.module_to_return = mod
    chained = _ImportHookChainedLoader(loader)
    call_count = [0]
    def hook(module):
        call_count[0] += 1
    # Register 500 hooks
    _post_import_hooks[mod.__name__] = {f"h{i}": hook for i in range(500)}
    chained._load_module(mod.__name__) # 13.7μs -> 2.87μs (378% faster)

def test_load_module_does_not_deadlock_on_reentrant_hook():
    # Large scale/edge: hook registers another hook for a different module
    loader = DummyLoader()
    mod1 = make_dummy_module("mod_reentrant1")
    mod2 = make_dummy_module("mod_reentrant2")
    loader.module_to_return = mod1
    chained = _ImportHookChainedLoader(loader)
    def hook(module):
        # Register a hook for another module
        _post_import_hooks[mod2.__name__] = {"h2": lambda m: None}
    _post_import_hooks[mod1.__name__] = {"h1": hook}
    chained._load_module(mod1.__name__) # 13.8μs -> 2.82μs (388% faster)
    # Now load the second module, should not deadlock
    loader2 = DummyLoader()
    loader2.module_to_return = mod2
    chained2 = _ImportHookChainedLoader(loader2)
    chained2._load_module(mod2.__name__) # 6.23μs -> 1.09μs (471% faster)

def test_load_module_thread_safety():
    # Large scale: simulate concurrent loads (not a true concurrency test, but exercises lock)
    import threading
    loader = DummyLoader()
    mod = make_dummy_module("mod_thread")
    loader.module_to_return = mod
    chained = _ImportHookChainedLoader(loader)
    called = []
    def hook(module):
        called.append(1)
    _post_import_hooks[mod.__name__] = {"h": hook}
    # Run _load_module in multiple threads
    threads = [threading.Thread(target=chained._load_module, args=(mod.__name__,)) for _ in range(10)]
    for t in threads:
        t.start()
    for t in threads:
        t.join()
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.
#------------------------------------------------
import threading
import types
from typing import Any, Dict

# imports
import pytest
from wandb.sdk.lib.import_hooks import _ImportHookChainedLoader

# --- Function to test (from wandb/sdk/lib/import_hooks.py) ---


_post_import_hooks: Dict = {}

# --- Unit Tests ---

# Helper class to simulate a loader for testing
class DummyLoader:
    """Simulates a loader with load_module, create_module, exec_module."""
    def __init__(self):
        self.loaded_names = []
        self.created_specs = []
        self.executed_modules = []
        self.load_module_return = None
        self.create_module_return = None

    def load_module(self, fullname):
        self.loaded_names.append(fullname)
        if self.load_module_return is not None:
            return self.load_module_return
        # Create a dummy module object
        mod = types.ModuleType(fullname)
        mod.__name__ = fullname
        mod.__loader__ = self
        return mod

    def create_module(self, spec):
        self.created_specs.append(spec)
        if self.create_module_return is not None:
            return self.create_module_return
        mod = types.ModuleType(spec.name)
        mod.__name__ = spec.name
        return mod

    def exec_module(self, module):
        self.executed_modules.append(module)

# Helper for post-import hook
def make_hook(recorder):
    def hook(module):
        recorder.append(module.__name__)
    return hook

# --- Basic Test Cases ---

def test_load_module_basic_sets_loader_and_returns_module():
    """Basic: _load_module should call loader.load_module and return its result."""
    loader = DummyLoader()
    chained = _ImportHookChainedLoader(loader)
    codeflash_output = chained._load_module("testmod"); module = codeflash_output # 14.4μs -> 3.61μs (300% faster)

def test_load_module_invokes_post_import_hooks():
    """Basic: _load_module should call post-import hooks for the module."""
    loader = DummyLoader()
    chained = _ImportHookChainedLoader(loader)
    recorder = []
    _post_import_hooks["testmod"] = {"hook1": make_hook(recorder)}
    codeflash_output = chained._load_module("testmod"); module = codeflash_output # 13.9μs -> 3.43μs (306% faster)

def test_load_module_truncates_post_import_hooks():
    """Basic: After loading, hooks for the module should be truncated (but key remains)."""
    loader = DummyLoader()
    chained = _ImportHookChainedLoader(loader)
    recorder = []
    _post_import_hooks["testmod"] = {"hook1": make_hook(recorder)}
    codeflash_output = chained._load_module("testmod"); module = codeflash_output # 14.4μs -> 3.41μs (322% faster)

def test_load_module_no_hooks_registered():
    """Basic: _load_module should not fail if no hooks are registered."""
    loader = DummyLoader()
    chained = _ImportHookChainedLoader(loader)
    # Ensure no exception is raised
    codeflash_output = chained._load_module("testmod_no_hooks"); module = codeflash_output # 14.2μs -> 3.56μs (298% faster)


def test_load_module_module_has_no_loader_attribute():
    """Edge: Module does not support __loader__ assignment (simulate by using a custom object)."""
    loader = DummyLoader()
    chained = _ImportHookChainedLoader(loader)
    # Simulate a module with __slots__ and no __loader__ attribute
    class NoLoaderModule:
        __slots__ = ['__name__']
        def __init__(self, name):
            self.__name__ = name
    loader.load_module_return = NoLoaderModule("testmod_noloader")
    # Should not raise
    codeflash_output = chained._load_module("testmod_noloader"); module = codeflash_output # 11.3μs -> 3.19μs (255% faster)

def test_load_module_module_loader_is_none():
    """Edge: Module __loader__ is None, should be set to loader."""
    loader = DummyLoader()
    chained = _ImportHookChainedLoader(loader)
    mod = types.ModuleType("testmod_none_loader")
    mod.__loader__ = None
    loader.load_module_return = mod
    codeflash_output = chained._load_module("testmod_none_loader"); module = codeflash_output # 13.3μs -> 2.91μs (359% faster)

def test_load_module_module_loader_is_self():
    """Edge: Module __loader__ is self, should be set to loader."""
    loader = DummyLoader()
    chained = _ImportHookChainedLoader(loader)
    mod = types.ModuleType("testmod_self_loader")
    mod.__loader__ = chained
    loader.load_module_return = mod
    codeflash_output = chained._load_module("testmod_self_loader"); module = codeflash_output # 13.7μs -> 2.84μs (382% faster)

def test_load_module_module_spec_loader_is_self():
    """Edge: Module __spec__.loader is self, should be set to loader."""
    loader = DummyLoader()
    chained = _ImportHookChainedLoader(loader)
    mod = types.ModuleType("testmod_spec_loader")
    class Spec:
        def __init__(self):
            self.loader = chained
    mod.__spec__ = Spec()
    loader.load_module_return = mod
    codeflash_output = chained._load_module("testmod_spec_loader"); module = codeflash_output # 10.9μs -> 3.09μs (252% faster)

def test_load_module_module_has_no_spec():
    """Edge: Module has no __spec__ attribute, should not fail."""
    loader = DummyLoader()
    chained = _ImportHookChainedLoader(loader)
    mod = types.ModuleType("testmod_no_spec")
    if hasattr(mod, "__spec__"):
        delattr(mod, "__spec__")
    loader.load_module_return = mod
    codeflash_output = chained._load_module("testmod_no_spec"); module = codeflash_output # 14.6μs -> 3.99μs (266% faster)

def test_load_module_module_has_no_name():
    """Edge: Module has no __name__ attribute, should not fail."""
    loader = DummyLoader()
    chained = _ImportHookChainedLoader(loader)
    class NoNameModule:
        pass
    mod = NoNameModule()
    loader.load_module_return = mod
    # Should not raise
    codeflash_output = chained._load_module("testmod_noname"); module = codeflash_output # 8.52μs -> 3.07μs (178% faster)

def test_load_module_hook_is_none():
    """Edge: Hook value is None, should not fail or call."""
    loader = DummyLoader()
    chained = _ImportHookChainedLoader(loader)
    _post_import_hooks["testmod"] = {"hook1": None}
    codeflash_output = chained._load_module("testmod"); module = codeflash_output # 14.5μs -> 3.51μs (312% faster)
    # Should not raise

# --- Large Scale Test Cases ---

def test_load_module_many_hooks():
    """Large Scale: _load_module should handle many hooks efficiently."""
    loader = DummyLoader()
    chained = _ImportHookChainedLoader(loader)
    recorder = []
    num_hooks = 500
    _post_import_hooks["testmod"] = {f"hook{i}": make_hook(recorder) for i in range(num_hooks)}
    codeflash_output = chained._load_module("testmod"); module = codeflash_output # 14.6μs -> 3.70μs (295% faster)

def test_load_module_many_modules():
    """Large Scale: _load_module should handle many modules loaded in sequence."""
    loader = DummyLoader()
    chained = _ImportHookChainedLoader(loader)
    num_modules = 500
    names = [f"mod{i}" for i in range(num_modules)]
    for name in names:
        _post_import_hooks[name] = {}
    modules = [chained._load_module(name) for name in names] # 14.4μs -> 3.45μs (317% faster)

def test_load_module_large_module_object():
    """Large Scale: _load_module should handle a module with many attributes."""
    loader = DummyLoader()
    chained = _ImportHookChainedLoader(loader)
    mod = types.ModuleType("largemod")
    # Add many attributes
    for i in range(999):
        setattr(mod, f"attr{i}", i)
    loader.load_module_return = mod
    codeflash_output = chained._load_module("largemod"); module = codeflash_output # 14.0μs -> 3.23μs (335% faster)

def test_load_module_large_hook_payload():
    """Large Scale: _load_module should handle hooks that process large data."""
    loader = DummyLoader()
    chained = _ImportHookChainedLoader(loader)
    large_data = []
    def big_hook(module):
        # Simulate processing of large data
        for i in range(999):
            large_data.append(i)
    _post_import_hooks["testmod"] = {"big_hook": big_hook}
    codeflash_output = chained._load_module("testmod"); module = codeflash_output # 14.5μs -> 3.56μs (307% faster)

To edit these changes git checkout codeflash/optimize-_ImportHookChainedLoader._load_module-mhdsc6lm and push.