⚡️ Speed up method SharedObjectLoadingScope.__enter__ by 9%

[codeflash-ai]

📄 9% (0.09x) speedup for SharedObjectLoadingScope.__enter__ in keras/src/legacy/saving/serialization.py

⏱️ Runtime : 1.57 milliseconds → 1.44 milliseconds (best of 250 runs)

📝 Explanation and details

The optimization eliminates a function call overhead by inlining the getattr operation directly in the __enter__ method instead of calling _shared_object_disabled().

Key Changes:

• Replaced if _shared_object_disabled(): with disabled = getattr(SHARED_OBJECT_DISABLED, "disabled", False) followed by if disabled:
• This removes the function call overhead while maintaining identical behavior

Why This Improves Performance:
In Python, function calls have inherent overhead due to stack frame creation, argument passing, and return value handling. The line profiler shows that the original _shared_object_disabled() function was called 2,632 times, consuming 1.63ms total. By inlining this simple getattr call directly, we eliminate:

• Function call setup/teardown overhead
• Stack frame allocation
• Return value handling

The profiler results confirm this optimization: the conditional check time dropped from 7.55ms (76.7% of total time) to 1.69ms (39.2% of total time), achieving a 9% overall speedup.

Test Case Performance:
The annotated tests show consistent 6-13% improvements across various scenarios, with particularly strong gains in:

• Disabled state handling (9-11% faster)
• Large-scale operations with 500+ scope entries (9.4% faster)
• Multiple sequential calls (8-13% faster)

This optimization is especially valuable since SharedObjectLoadingScope is used in Keras model serialization/deserialization workflows, where it may be called frequently during model loading operations.

✅ Correctness verification report:

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

🌀 Generated Regression Tests and Runtime

import threading

# imports
import pytest
from keras.src.legacy.saving.serialization import SharedObjectLoadingScope

# function to test
# (copied from keras/src/legacy/saving/serialization.py)
SHARED_OBJECT_DISABLED = threading.local()
SHARED_OBJECT_LOADING = threading.local()

class NoopLoadingScope:
    """A dummy scope that does nothing, used when shared object handling is disabled."""
    pass
from keras.src.legacy.saving.serialization import SharedObjectLoadingScope

# ----------- Basic Test Cases -----------

def test_enter_returns_self_when_enabled():
    """Test that __enter__ returns self when shared object handling is enabled."""
    scope = SharedObjectLoadingScope()
    codeflash_output = scope.__enter__(); result = codeflash_output # 2.97μs -> 2.77μs (7.38% faster)

def test_enter_sets_scope_on_threadlocal():
    """Test that __enter__ sets the scope on SHARED_OBJECT_LOADING threadlocal."""
    scope = SharedObjectLoadingScope()
    scope.__enter__() # 2.05μs -> 1.93μs (6.38% faster)

def test_enter_initializes_obj_ids_to_obj_dict():
    """Test that __enter__ initializes the _obj_ids_to_obj dictionary."""
    scope = SharedObjectLoadingScope()
    scope.__enter__() # 1.94μs -> 1.82μs (6.81% faster)

def test_enter_returns_noop_when_disabled():
    """Test that __enter__ returns NoopLoadingScope when disabled."""
    SHARED_OBJECT_DISABLED.disabled = True
    scope = SharedObjectLoadingScope()
    codeflash_output = scope.__enter__(); result = codeflash_output # 1.85μs -> 1.69μs (9.28% faster)

# ----------- Edge Test Cases -----------

def test_enter_multiple_calls_resets_obj_ids_to_obj():
    """Test that multiple __enter__ calls on the same scope resets _obj_ids_to_obj each time."""
    scope = SharedObjectLoadingScope()
    scope.__enter__() # 2.66μs -> 2.37μs (12.5% faster)
    scope._obj_ids_to_obj['test'] = 123
    scope.__enter__() # 884ns -> 874ns (1.14% faster)

def test_enter_disabled_then_enabled_switch():
    """Test switching between disabled and enabled states."""
    SHARED_OBJECT_DISABLED.disabled = True
    scope = SharedObjectLoadingScope()
    codeflash_output = scope.__enter__(); result = codeflash_output # 1.83μs -> 1.65μs (11.0% faster)
    SHARED_OBJECT_DISABLED.disabled = False
    codeflash_output = scope.__enter__(); result2 = codeflash_output # 842ns -> 832ns (1.20% faster)

def test_enter_disabled_scope_not_initialized():
    """Test that _obj_ids_to_obj is not initialized when disabled."""
    SHARED_OBJECT_DISABLED.disabled = True
    scope = SharedObjectLoadingScope()
    scope.__enter__() # 1.77μs -> 1.61μs (9.67% faster)

def test_enter_noop_scope_is_new_instance():
    """Test that NoopLoadingScope returned is a new instance each time."""
    SHARED_OBJECT_DISABLED.disabled = True
    scope = SharedObjectLoadingScope()
    codeflash_output = scope.__enter__(); result1 = codeflash_output # 1.73μs -> 1.58μs (9.40% faster)
    codeflash_output = scope.__enter__(); result2 = codeflash_output # 806ns -> 776ns (3.87% faster)

# ----------- Large Scale Test Cases -----------

def test_enter_large_number_of_scopes():
    """Test entering a large number of scopes to ensure scalability."""
    scopes = [SharedObjectLoadingScope() for _ in range(500)]
    # Enter each scope and check threadlocal is set correctly
    for i, scope in enumerate(scopes):
        scope.__enter__() # 303μs -> 276μs (9.56% faster)
        # Add something to the dict and check it is reset
        scope._obj_ids_to_obj['key'] = i
        scope.__enter__() # 303μs -> 279μs (8.36% faster)

def test_enter_large_obj_ids_to_obj_dict():
    """Test that the _obj_ids_to_obj dict can handle a large number of entries."""
    scope = SharedObjectLoadingScope()
    scope.__enter__() # 2.21μs -> 2.05μs (7.45% faster)
    for i in range(1000):
        scope._obj_ids_to_obj[i] = f"obj_{i}"
    # Re-enter resets dict
    scope.__enter__() # 7.20μs -> 7.00μs (2.77% faster)

def test_enter_performance_under_load():
    """Test performance does not degrade with many enter calls (sanity check)."""
    import time
    scope = SharedObjectLoadingScope()
    start = time.time()
    for _ in range(1000):
        scope.__enter__() # 569μs -> 520μs (9.42% faster)
    duration = time.time() - start

# ----------- Determinism Test -----------

def test_enter_deterministic_behavior():
    """Test that repeated calls produce deterministic results."""
    scope = SharedObjectLoadingScope()
    results = []
    for _ in range(10):
        results.append(scope.__enter__()) # 7.59μs -> 7.07μs (7.49% faster)

# ----------- Clean up -----------

def test_exit_resets_scope_to_noop():
    """Test that __exit__ resets the scope to NoopLoadingScope."""
    scope = SharedObjectLoadingScope()
    scope.__enter__() # 1.64μs -> 1.53μs (7.45% faster)
    scope.__exit__()
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

import sys
import threading

# imports
import pytest
from keras.src.legacy.saving.serialization import SharedObjectLoadingScope

# Simulate keras/src/legacy/saving/serialization.py
SHARED_OBJECT_DISABLED = threading.local()
SHARED_OBJECT_LOADING = threading.local()

class NoopLoadingScope:
    """A no-operation context manager to use when shared object handling is disabled."""
    pass
from keras.src.legacy.saving.serialization import SharedObjectLoadingScope

# unit tests

# ---- Basic Test Cases ----

def test_enter_returns_self_when_not_disabled():
    """Test that __enter__ returns self when shared object handling is enabled."""
    # Ensure the disabled flag is not set
    if hasattr(SHARED_OBJECT_DISABLED, "disabled"):
        del SHARED_OBJECT_DISABLED.disabled

    scope = SharedObjectLoadingScope()
    codeflash_output = scope.__enter__(); result = codeflash_output # 1.36μs -> 1.22μs (11.5% faster)
    # Clean up
    if hasattr(SHARED_OBJECT_LOADING, "scope"):
        del SHARED_OBJECT_LOADING.scope

def test_enter_sets_obj_ids_to_obj_dict():
    """Test that __enter__ initializes _obj_ids_to_obj as an empty dict."""
    if hasattr(SHARED_OBJECT_DISABLED, "disabled"):
        del SHARED_OBJECT_DISABLED.disabled
    scope = SharedObjectLoadingScope()
    scope.__enter__() # 1.45μs -> 1.34μs (7.88% faster)
    # Clean up
    if hasattr(SHARED_OBJECT_LOADING, "scope"):
        del SHARED_OBJECT_LOADING.scope

def test_enter_sets_global_scope():
    """Test that __enter__ sets SHARED_OBJECT_LOADING.scope to self."""
    if hasattr(SHARED_OBJECT_DISABLED, "disabled"):
        del SHARED_OBJECT_DISABLED.disabled
    scope = SharedObjectLoadingScope()
    scope.__enter__() # 1.42μs -> 1.36μs (4.40% faster)
    # Clean up
    if hasattr(SHARED_OBJECT_LOADING, "scope"):
        del SHARED_OBJECT_LOADING.scope

# ---- Edge Test Cases ----

def test_enter_returns_noop_when_disabled():
    """Test that __enter__ returns NoopLoadingScope when disabled."""
    SHARED_OBJECT_DISABLED.disabled = True
    scope = SharedObjectLoadingScope()
    codeflash_output = scope.__enter__(); result = codeflash_output # 1.63μs -> 1.68μs (2.74% slower)
    # Clean up
    del SHARED_OBJECT_DISABLED.disabled
    if hasattr(SHARED_OBJECT_LOADING, "scope"):
        del SHARED_OBJECT_LOADING.scope

def test_enter_multiple_calls():
    """Test that multiple __enter__ calls reset _obj_ids_to_obj each time."""
    if hasattr(SHARED_OBJECT_DISABLED, "disabled"):
        del SHARED_OBJECT_DISABLED.disabled
    scope = SharedObjectLoadingScope()
    scope.__enter__() # 2.50μs -> 2.21μs (13.4% faster)
    scope._obj_ids_to_obj['foo'] = 123
    scope.__enter__() # 911ns -> 861ns (5.81% faster)
    # Clean up
    if hasattr(SHARED_OBJECT_LOADING, "scope"):
        del SHARED_OBJECT_LOADING.scope

def test_enter_does_not_affect_other_attributes():
    """Test that __enter__ does not overwrite unrelated attributes."""
    if hasattr(SHARED_OBJECT_DISABLED, "disabled"):
        del SHARED_OBJECT_DISABLED.disabled
    scope = SharedObjectLoadingScope()
    scope.some_other_attr = 42
    scope.__enter__() # 1.63μs -> 1.44μs (13.6% faster)
    # Clean up
    if hasattr(SHARED_OBJECT_LOADING, "scope"):
        del SHARED_OBJECT_LOADING.scope

# ---- Large Scale Test Cases ----

def test_enter_large_number_of_scopes():
    """Test that many scopes can be entered without interfering with each other."""
    if hasattr(SHARED_OBJECT_DISABLED, "disabled"):
        del SHARED_OBJECT_DISABLED.disabled

    scopes = [SharedObjectLoadingScope() for _ in range(100)]
    for i, scope in enumerate(scopes):
        codeflash_output = scope.__enter__(); result = codeflash_output # 61.7μs -> 56.4μs (9.38% faster)
    # Clean up
    if hasattr(SHARED_OBJECT_LOADING, "scope"):
        del SHARED_OBJECT_LOADING.scope

def test_enter_performance_under_load():
    """Test __enter__ performance and correctness under repeated use."""
    if hasattr(SHARED_OBJECT_DISABLED, "disabled"):
        del SHARED_OBJECT_DISABLED.disabled

    scope = SharedObjectLoadingScope()
    for _ in range(500):
        codeflash_output = scope.__enter__(); result = codeflash_output # 285μs -> 261μs (9.46% faster)
    # Clean up
    if hasattr(SHARED_OBJECT_LOADING, "scope"):
        del SHARED_OBJECT_LOADING.scope

To edit these changes git checkout codeflash/optimize-SharedObjectLoadingScope.__enter__-mjado33j and push.