⚡️ Speed up method _CachedStorage.get_trial by 8%

[codeflash-ai]

📄 8% (0.08x) speedup for _CachedStorage.get_trial in optuna/storages/_cached_storage.py

⏱️ Runtime : 2.41 milliseconds → 2.23 milliseconds (best of 153 runs)

📝 Explanation and details

The optimized code achieves an 8% speedup through two key optimizations that reduce overhead in cache operations:

1. Eliminated double dictionary lookup in _get_cached_trial:

• Original: if trial_id not in self._trial_id_to_study_id_and_number: followed by self._trial_id_to_study_id_and_number[trial_id] performs two hash table lookups
• Optimized: Single dict.get(trial_id) lookup stores the result, eliminating redundant hash computation and lookup

2. Reduced lock contention in get_trial:

• Original: Always acquires lock before checking cache, causing unnecessary blocking when trial isn't cached
• Optimized: Pre-checks if trial exists in cache mapping before acquiring lock, avoiding lock contention for backend-only lookups

3. Minor efficiency improvement in RDBStorage:

• Eliminated intermediate variable assignment by directly returning the result of _build_frozen_trial_from_trial_model

These optimizations are particularly effective for workloads with:

• High cache miss rates (43-76% faster for backend-only lookups based on test results)
• Mixed cache/backend scenarios (52% faster for backend lookups in mixed workloads)
• Large-scale operations where lock contention becomes significant

The optimizations maintain identical behavior while reducing CPU overhead from hash table operations and thread synchronization, explaining the consistent 8% overall speedup.

✅ Correctness verification report:

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

🌀 Generated Regression Tests and Runtime

import threading
from types import SimpleNamespace

# imports
import pytest
from optuna.storages._cached_storage import _CachedStorage

# --- Minimal stubs for external classes/types used by get_trial ---

# Simulate optuna.trial.TrialState enum with is_finished()
class TrialState:
    RUNNING = 0
    WAITING = 4
    COMPLETE = 1
    PRUNED = 2
    FAIL = 3

    def __init__(self, value):
        self.value = value

    def is_finished(self):
        return self.value not in (self.RUNNING, self.WAITING)

# Simulate optuna.trial.FrozenTrial
class FrozenTrial:
    def __init__(self, trial_id, number=0, state=None, value=None):
        self.trial_id = trial_id
        self.number = number
        self.state = state or TrialState(TrialState.COMPLETE)
        self.value = value

# Simulate study info
class _StudyInfo:
    def __init__(self, trials):
        self.trials = trials

# Simulate backend storage
class DummyBackend:
    def __init__(self):
        self._trials = {}

    def add_trial(self, trial):
        self._trials[trial.trial_id] = trial

    def get_trial(self, trial_id):
        # Simulate backend retrieval
        if trial_id not in self._trials:
            raise KeyError(f"Trial {trial_id} not found in backend.")
        return self._trials[trial_id]

# --- Unit tests ---

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

def test_get_trial_returns_cached_finished_trial():
    """Test that get_trial returns a finished trial from cache."""
    backend = DummyBackend()
    storage = _CachedStorage(backend)
    trial = FrozenTrial(trial_id=1, number=0, state=TrialState(TrialState.COMPLETE))
    study_id = 10
    storage._studies[study_id] = _StudyInfo([trial])
    storage._trial_id_to_study_id_and_number[1] = (study_id, 0)
    codeflash_output = storage.get_trial(1); result = codeflash_output # 2.12μs -> 2.39μs (11.2% slower)

def test_get_trial_returns_backend_trial_when_not_in_cache():
    """Test that get_trial falls back to backend if trial not cached."""
    backend = DummyBackend()
    backend_trial = FrozenTrial(trial_id=2, number=0, state=TrialState(TrialState.COMPLETE))
    backend.add_trial(backend_trial)
    storage = _CachedStorage(backend)
    codeflash_output = storage.get_trial(2); result = codeflash_output # 1.49μs -> 1.04μs (43.8% faster)

def test_get_trial_returns_backend_trial_when_cached_trial_unfinished():
    """Test that get_trial falls back to backend if cached trial is unfinished."""
    backend = DummyBackend()
    backend_trial = FrozenTrial(trial_id=3, number=0, state=TrialState(TrialState.COMPLETE))
    backend.add_trial(backend_trial)
    storage = _CachedStorage(backend)
    unfinished_trial = FrozenTrial(trial_id=3, number=0, state=TrialState(TrialState.RUNNING))
    study_id = 11
    storage._studies[study_id] = _StudyInfo([unfinished_trial])
    storage._trial_id_to_study_id_and_number[3] = (study_id, 0)
    codeflash_output = storage.get_trial(3); result = codeflash_output # 2.22μs -> 2.44μs (9.17% slower)

def test_get_trial_returns_backend_trial_when_trial_id_not_in_cache():
    """Test that get_trial falls back to backend if trial_id not in cache mapping."""
    backend = DummyBackend()
    backend_trial = FrozenTrial(trial_id=4, number=0, state=TrialState(TrialState.COMPLETE))
    backend.add_trial(backend_trial)
    storage = _CachedStorage(backend)
    # No mapping for trial_id=4
    codeflash_output = storage.get_trial(4); result = codeflash_output # 1.42μs -> 854ns (66.6% faster)

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

def test_get_trial_raises_keyerror_when_not_found_anywhere():
    """Test that get_trial raises KeyError if trial not in cache or backend."""
    backend = DummyBackend()
    storage = _CachedStorage(backend)
    with pytest.raises(KeyError):
        storage.get_trial(999) # 2.12μs -> 1.56μs (36.1% faster)

def test_get_trial_with_multiple_studies_and_trials():
    """Test get_trial returns correct trial when multiple studies/trials exist."""
    backend = DummyBackend()
    storage = _CachedStorage(backend)
    # Study 1 with 2 trials
    trial_a = FrozenTrial(trial_id=5, number=0, state=TrialState(TrialState.COMPLETE))
    trial_b = FrozenTrial(trial_id=6, number=1, state=TrialState(TrialState.COMPLETE))
    study1_id = 100
    storage._studies[study1_id] = _StudyInfo([trial_a, trial_b])
    storage._trial_id_to_study_id_and_number[5] = (study1_id, 0)
    storage._trial_id_to_study_id_and_number[6] = (study1_id, 1)
    # Study 2 with 1 trial
    trial_c = FrozenTrial(trial_id=7, number=0, state=TrialState(TrialState.COMPLETE))
    study2_id = 200
    storage._studies[study2_id] = _StudyInfo([trial_c])
    storage._trial_id_to_study_id_and_number[7] = (study2_id, 0)
    # All should be returned from cache
    codeflash_output = storage.get_trial(5) # 1.96μs -> 2.39μs (17.8% slower)
    codeflash_output = storage.get_trial(6) # 829ns -> 895ns (7.37% slower)
    codeflash_output = storage.get_trial(7) # 608ns -> 669ns (9.12% slower)

def test_get_trial_with_waiting_state_trial():
    """Test that a trial in WAITING state is not returned from cache."""
    backend = DummyBackend()
    backend_trial = FrozenTrial(trial_id=8, number=0, state=TrialState(TrialState.COMPLETE))
    backend.add_trial(backend_trial)
    storage = _CachedStorage(backend)
    waiting_trial = FrozenTrial(trial_id=8, number=0, state=TrialState(TrialState.WAITING))
    study_id = 300
    storage._studies[study_id] = _StudyInfo([waiting_trial])
    storage._trial_id_to_study_id_and_number[8] = (study_id, 0)
    codeflash_output = storage.get_trial(8); result = codeflash_output # 2.09μs -> 2.45μs (14.5% slower)

def test_get_trial_with_pruned_and_fail_states():
    """Test that trials in PRUNED and FAIL states are returned from cache."""
    backend = DummyBackend()
    storage = _CachedStorage(backend)
    pruned_trial = FrozenTrial(trial_id=9, number=0, state=TrialState(TrialState.PRUNED))
    fail_trial = FrozenTrial(trial_id=10, number=0, state=TrialState(TrialState.FAIL))
    study_id = 400
    storage._studies[study_id] = _StudyInfo([pruned_trial, fail_trial])
    storage._trial_id_to_study_id_and_number[9] = (study_id, 0)
    storage._trial_id_to_study_id_and_number[10] = (study_id, 1)
    codeflash_output = storage.get_trial(9) # 1.69μs -> 1.99μs (15.2% slower)
    codeflash_output = storage.get_trial(10) # 755ns -> 867ns (12.9% slower)


def test_get_trial_large_number_of_trials():
    """Test get_trial with a large number of cached trials."""
    backend = DummyBackend()
    storage = _CachedStorage(backend)
    study_id = 1000
    num_trials = 500  # Keep under 1000 as per instructions
    trials = []
    for i in range(num_trials):
        trial = FrozenTrial(trial_id=10000 + i, number=i, state=TrialState(TrialState.COMPLETE))
        trials.append(trial)
        storage._trial_id_to_study_id_and_number[10000 + i] = (study_id, i)
    storage._studies[study_id] = _StudyInfo(trials)
    # All should be returned from cache
    for i in range(num_trials):
        codeflash_output = storage.get_trial(10000 + i) # 288μs -> 305μs (5.41% slower)

def test_get_trial_large_number_of_backend_trials():
    """Test get_trial with a large number of backend trials, none cached."""
    backend = DummyBackend()
    num_trials = 500
    for i in range(num_trials):
        trial = FrozenTrial(trial_id=20000 + i, number=i, state=TrialState(TrialState.COMPLETE))
        backend.add_trial(trial)
    storage = _CachedStorage(backend)
    for i in range(num_trials):
        codeflash_output = storage.get_trial(20000 + i) # 217μs -> 144μs (50.0% faster)

def test_get_trial_large_number_mixed_cache_and_backend():
    """Test get_trial with a mix of cached and backend trials."""
    backend = DummyBackend()
    storage = _CachedStorage(backend)
    study_id = 3000
    num_trials = 250
    # Half cached, half only in backend
    cached_trials = []
    for i in range(num_trials):
        trial = FrozenTrial(trial_id=30000 + i, number=i, state=TrialState(TrialState.COMPLETE))
        cached_trials.append(trial)
        storage._trial_id_to_study_id_and_number[30000 + i] = (study_id, i)
    storage._studies[study_id] = _StudyInfo(cached_trials)
    for i in range(num_trials, 2 * num_trials):
        trial = FrozenTrial(trial_id=30000 + i, number=i, state=TrialState(TrialState.COMPLETE))
        backend.add_trial(trial)
    # Check cached
    for i in range(num_trials):
        codeflash_output = storage.get_trial(30000 + i) # 146μs -> 154μs (5.35% slower)
    # Check backend
    for i in range(num_trials, 2 * num_trials):
        codeflash_output = storage.get_trial(30000 + i) # 107μs -> 70.5μs (52.2% faster)
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.
#------------------------------------------------
import threading
from datetime import datetime, timedelta
from types import SimpleNamespace

# imports
import pytest
from optuna.storages._cached_storage import _CachedStorage

# --- Minimal stubs for optuna.trial.FrozenTrial and TrialState ---

class TrialState:
    RUNNING = 0
    COMPLETE = 1
    PRUNED = 2
    FAIL = 3
    WAITING = 4

    def __init__(self, state):
        self._state = state

    def is_finished(self):
        return self._state not in (TrialState.RUNNING, TrialState.WAITING)

    def __eq__(self, other):
        if isinstance(other, int):
            return self._state == other
        if isinstance(other, TrialState):
            return self._state == other._state
        return False

    def __repr__(self):
        return f"TrialState({self._state})"

class FrozenTrial:
    def __init__(self, number, state, trial_id):
        self.number = number
        self.state = TrialState(state)
        self.trial_id = trial_id

    def __eq__(self, other):
        return (
            isinstance(other, FrozenTrial)
            and self.number == other.number
            and self.state == other.state
            and self.trial_id == other.trial_id
        )

    def __repr__(self):
        return f"FrozenTrial(number={self.number}, state={self.state}, trial_id={self.trial_id})"

# --- Minimal stub for RDBStorage ---

class DummyBackend:
    """A dummy backend to simulate persistent storage."""
    def __init__(self):
        self.trials = {}
        self.lock = threading.Lock()
        self.call_log = []

    def get_trial(self, trial_id):
        self.call_log.append(trial_id)
        if trial_id not in self.trials:
            raise KeyError(f"Trial {trial_id} not found in backend")
        return self.trials[trial_id]

# --- Minimal stub for _StudyInfo ---

class _StudyInfo:
    def __init__(self, trials):
        self.trials = trials

# --- Unit Tests ---

# ----------- BASIC TEST CASES -----------

def test_get_trial_returns_cached_finished_trial():
    """Test that get_trial returns a finished trial from cache."""
    backend = DummyBackend()
    cache = _CachedStorage(backend)
    # Setup: One study with one trial, finished
    study_id = 1
    trial_id = 100
    trial = FrozenTrial(number=0, state=TrialState.COMPLETE, trial_id=trial_id)
    cache._studies[study_id] = _StudyInfo([trial])
    cache._trial_id_to_study_id_and_number[trial_id] = (study_id, 0)
    cache._study_id_and_number_to_trial_id[(study_id, 0)] = trial_id
    # Backend should not be called
    codeflash_output = cache.get_trial(trial_id); result = codeflash_output # 2.46μs -> 2.56μs (4.03% slower)

def test_get_trial_returns_from_backend_if_not_in_cache():
    """Test that get_trial calls backend if trial is not in cache."""
    backend = DummyBackend()
    cache = _CachedStorage(backend)
    trial_id = 200
    trial = FrozenTrial(number=0, state=TrialState.COMPLETE, trial_id=trial_id)
    backend.trials[trial_id] = trial
    # Not in cache
    codeflash_output = cache.get_trial(trial_id); result = codeflash_output # 1.69μs -> 995ns (69.4% faster)

def test_get_trial_returns_from_backend_if_trial_not_finished():
    """Test that get_trial calls backend if cached trial is not finished."""
    backend = DummyBackend()
    cache = _CachedStorage(backend)
    study_id = 2
    trial_id = 201
    unfinished_trial = FrozenTrial(number=0, state=TrialState.RUNNING, trial_id=trial_id)
    cache._studies[study_id] = _StudyInfo([unfinished_trial])
    cache._trial_id_to_study_id_and_number[trial_id] = (study_id, 0)
    cache._study_id_and_number_to_trial_id[(study_id, 0)] = trial_id
    backend.trials[trial_id] = FrozenTrial(number=0, state=TrialState.COMPLETE, trial_id=trial_id)
    codeflash_output = cache.get_trial(trial_id); result = codeflash_output # 2.47μs -> 2.49μs (0.843% slower)

def test_get_trial_returns_from_backend_if_trial_id_not_in_cache():
    """Test that get_trial calls backend if trial_id not mapped in cache."""
    backend = DummyBackend()
    cache = _CachedStorage(backend)
    trial_id = 300
    backend.trials[trial_id] = FrozenTrial(number=0, state=TrialState.COMPLETE, trial_id=trial_id)
    codeflash_output = cache.get_trial(trial_id); result = codeflash_output # 1.57μs -> 888ns (76.4% faster)

# ----------- EDGE TEST CASES -----------

def test_get_trial_raises_if_backend_raises():
    """Test that get_trial propagates backend exception if trial not found anywhere."""
    backend = DummyBackend()
    cache = _CachedStorage(backend)
    trial_id = 999
    with pytest.raises(KeyError):
        cache.get_trial(trial_id) # 2.20μs -> 1.68μs (31.3% faster)

def test_get_trial_with_multiple_studies_and_trials():
    """Test correct mapping with multiple studies and trials."""
    backend = DummyBackend()
    cache = _CachedStorage(backend)
    # Study 1: 2 trials, Study 2: 1 trial
    study1_id = 10
    study2_id = 20
    t1 = FrozenTrial(number=0, state=TrialState.COMPLETE, trial_id=1000)
    t2 = FrozenTrial(number=1, state=TrialState.PRUNED, trial_id=1001)
    t3 = FrozenTrial(number=0, state=TrialState.COMPLETE, trial_id=2000)
    cache._studies[study1_id] = _StudyInfo([t1, t2])
    cache._studies[study2_id] = _StudyInfo([t3])
    cache._trial_id_to_study_id_and_number[1000] = (study1_id, 0)
    cache._trial_id_to_study_id_and_number[1001] = (study1_id, 1)
    cache._trial_id_to_study_id_and_number[2000] = (study2_id, 0)
    cache._study_id_and_number_to_trial_id[(study1_id, 0)] = 1000
    cache._study_id_and_number_to_trial_id[(study1_id, 1)] = 1001
    cache._study_id_and_number_to_trial_id[(study2_id, 0)] = 2000
    # Should get from cache
    codeflash_output = cache.get_trial(1000) # 2.09μs -> 2.36μs (11.6% slower)
    codeflash_output = cache.get_trial(1001) # 875ns -> 853ns (2.58% faster)
    codeflash_output = cache.get_trial(2000) # 731ns -> 712ns (2.67% faster)

def test_get_trial_with_waiting_trial_state():
    """Test that WAITING state is not finished and triggers backend."""
    backend = DummyBackend()
    cache = _CachedStorage(backend)
    study_id = 30
    trial_id = 3000
    waiting_trial = FrozenTrial(number=0, state=TrialState.WAITING, trial_id=trial_id)
    cache._studies[study_id] = _StudyInfo([waiting_trial])
    cache._trial_id_to_study_id_and_number[trial_id] = (study_id, 0)
    cache._study_id_and_number_to_trial_id[(study_id, 0)] = trial_id
    backend.trials[trial_id] = FrozenTrial(number=0, state=TrialState.COMPLETE, trial_id=trial_id)
    codeflash_output = cache.get_trial(trial_id); result = codeflash_output # 2.34μs -> 2.38μs (1.56% slower)

def test_get_trial_with_pruned_and_failed_states():
    """Test that PRUNED and FAIL states are considered finished and returned from cache."""
    backend = DummyBackend()
    cache = _CachedStorage(backend)
    study_id = 40
    pruned_id, fail_id = 4000, 4001
    pruned_trial = FrozenTrial(number=0, state=TrialState.PRUNED, trial_id=pruned_id)
    fail_trial = FrozenTrial(number=1, state=TrialState.FAIL, trial_id=fail_id)
    cache._studies[study_id] = _StudyInfo([pruned_trial, fail_trial])
    cache._trial_id_to_study_id_and_number[pruned_id] = (study_id, 0)
    cache._trial_id_to_study_id_and_number[fail_id] = (study_id, 1)
    cache._study_id_and_number_to_trial_id[(study_id, 0)] = pruned_id
    cache._study_id_and_number_to_trial_id[(study_id, 1)] = fail_id
    codeflash_output = cache.get_trial(pruned_id) # 1.86μs -> 1.95μs (4.82% slower)
    codeflash_output = cache.get_trial(fail_id) # 829ns -> 822ns (0.852% faster)


def test_get_trial_large_cache():
    """Test get_trial performance and correctness with a large cache (1000 trials)."""
    backend = DummyBackend()
    cache = _CachedStorage(backend)
    study_id = 60
    num_trials = 1000
    trials = []
    for i in range(num_trials):
        trial = FrozenTrial(number=i, state=TrialState.COMPLETE, trial_id=6000 + i)
        trials.append(trial)
        cache._trial_id_to_study_id_and_number[6000 + i] = (study_id, i)
        cache._study_id_and_number_to_trial_id[(study_id, i)] = 6000 + i
    cache._studies[study_id] = _StudyInfo(trials)
    # All should be returned from cache, backend not called
    for i in range(num_trials):
        codeflash_output = cache.get_trial(6000 + i); result = codeflash_output # 567μs -> 597μs (5.07% slower)

def test_get_trial_large_backend():
    """Test get_trial performance and correctness with a large backend (1000 trials)."""
    backend = DummyBackend()
    cache = _CachedStorage(backend)
    num_trials = 1000
    for i in range(num_trials):
        trial_id = 7000 + i
        backend.trials[trial_id] = FrozenTrial(number=i, state=TrialState.COMPLETE, trial_id=trial_id)
    # None in cache, all must be fetched from backend
    for i in range(num_trials):
        trial_id = 7000 + i
        codeflash_output = cache.get_trial(trial_id); result = codeflash_output # 461μs -> 302μs (52.3% faster)


def test_get_trial_large_with_unfinished_trials():
    """Test large cache with some unfinished trials (should fall back to backend)."""
    backend = DummyBackend()
    cache = _CachedStorage(backend)
    study_id = 90
    num_trials = 1000
    cache_trials = []
    for i in range(num_trials):
        trial_id = 9000 + i
        if i % 10 == 0:
            # Unfinished
            trial = FrozenTrial(number=i, state=TrialState.RUNNING, trial_id=trial_id)
            backend.trials[trial_id] = FrozenTrial(number=i, state=TrialState.COMPLETE, trial_id=trial_id)
        else:
            # Finished
            trial = FrozenTrial(number=i, state=TrialState.COMPLETE, trial_id=trial_id)
        cache_trials.append(trial)
        cache._trial_id_to_study_id_and_number[trial_id] = (study_id, i)
        cache._study_id_and_number_to_trial_id[(study_id, i)] = trial_id
    cache._studies[study_id] = _StudyInfo(cache_trials)
    for i in range(num_trials):
        trial_id = 9000 + i
        codeflash_output = cache.get_trial(trial_id); result = codeflash_output # 587μs -> 618μs (4.92% slower)
        if i % 10 == 0:
            pass
        else:
            pass
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

To edit these changes git checkout codeflash/optimize-_CachedStorage.get_trial-mhl31e3f and push.