⚡️ Speed up function _estimate_min_resource by 141%

[codeflash-ai]

📄 141% (1.41x) speedup for _estimate_min_resource in optuna/pruners/_successive_halving.py

⏱️ Runtime : 702 microseconds → 291 microseconds (best of 250 runs)

📝 Explanation and details

The optimized code achieves a 140% speedup by eliminating unnecessary list allocation and reducing attribute lookups. Here are the key optimizations:

1. Eliminated List Creation: The original code builds an entire list n_steps in memory before finding the maximum value. The optimized version uses a generator expression and finds the maximum incrementally during iteration, avoiding the memory allocation overhead.

2. Cached Attribute Lookup: Storing TrialState.COMPLETE in a local variable COMPLETE eliminates repeated attribute lookups during the filtering operation, reducing per-iteration overhead.

3. Manual Max Search: Instead of using Python's built-in max() function on a list, the optimization performs a manual maximum search that processes values as they're generated, further reducing memory pressure.

Performance Impact by Test Case:

• Small datasets (few trials): Shows modest slowdowns of 5-25% due to the overhead of the manual loop outweighing benefits
• Large datasets (1000+ trials): Shows dramatic speedups of 160-170%, where memory allocation costs dominate

The line profiler confirms this - the original code spends 98.8% of time in list creation, while the optimized version distributes time across the generator loop (92.4%) with much lower total runtime.

This optimization is particularly valuable for hyperparameter tuning workloads where _estimate_min_resource processes large numbers of completed trials, making the memory efficiency gains substantial.

✅ Correctness verification report:

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

🌀 Generated Regression Tests and Runtime

from __future__ import annotations

# Simulate minimal optuna.trial.FrozenTrial and TrialState for testing
from enum import Enum, auto

# imports
import pytest  # used for our unit tests
from optuna.pruners._successive_halving import _estimate_min_resource


class TrialState(Enum):
    COMPLETE = auto()
    PRUNED = auto()
    RUNNING = auto()
    WAITING = auto()
    FAIL = auto()

class FrozenTrial:
    def __init__(self, state, last_step):
        self.state = state
        self.last_step = last_step
from optuna.pruners._successive_halving import _estimate_min_resource

# unit tests

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

def test_empty_trials_returns_none():
    # No trials at all
    codeflash_output = _estimate_min_resource([]) # 663ns -> 1.18μs (43.8% slower)

def test_all_trials_incomplete_returns_none():
    # All trials are not COMPLETE
    trials = [
        FrozenTrial(TrialState.PRUNED, 10),
        FrozenTrial(TrialState.RUNNING, 20),
        FrozenTrial(TrialState.WAITING, 30),
        FrozenTrial(TrialState.FAIL, 40),
    ]
    codeflash_output = _estimate_min_resource(trials) # 1.23μs -> 1.34μs (7.93% slower)

def test_all_trials_complete_with_steps():
    # All trials COMPLETE with valid last_step values
    trials = [
        FrozenTrial(TrialState.COMPLETE, 100),
        FrozenTrial(TrialState.COMPLETE, 200),
        FrozenTrial(TrialState.COMPLETE, 150),
    ]
    # max last_step = 200, 200//100 = 2
    codeflash_output = _estimate_min_resource(trials) # 1.11μs -> 1.20μs (7.64% slower)

def test_some_trials_complete_some_not():
    # Only some trials are COMPLETE
    trials = [
        FrozenTrial(TrialState.COMPLETE, 50),
        FrozenTrial(TrialState.PRUNED, 100),
        FrozenTrial(TrialState.COMPLETE, 300),
        FrozenTrial(TrialState.RUNNING, 400),
    ]
    # max last_step among COMPLETE = 300, 300//100 = 3
    codeflash_output = _estimate_min_resource(trials) # 1.18μs -> 1.25μs (5.76% slower)

def test_some_trials_complete_with_none_last_step():
    # Some COMPLETE trials have last_step=None
    trials = [
        FrozenTrial(TrialState.COMPLETE, None),
        FrozenTrial(TrialState.COMPLETE, 120),
        FrozenTrial(TrialState.COMPLETE, None),
        FrozenTrial(TrialState.COMPLETE, 80),
    ]
    # max last_step = 120, 120//100 = 1
    codeflash_output = _estimate_min_resource(trials) # 1.14μs -> 1.22μs (7.20% slower)

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

def test_complete_trials_all_none_last_step_returns_none():
    # All COMPLETE trials have last_step=None
    trials = [
        FrozenTrial(TrialState.COMPLETE, None),
        FrozenTrial(TrialState.COMPLETE, None),
    ]
    codeflash_output = _estimate_min_resource(trials) # 977ns -> 1.14μs (14.3% slower)

def test_mixed_states_and_none_last_step():
    # Mixed states, some with None last_step
    trials = [
        FrozenTrial(TrialState.PRUNED, 500),
        FrozenTrial(TrialState.COMPLETE, None),
        FrozenTrial(TrialState.COMPLETE, 0),
        FrozenTrial(TrialState.COMPLETE, 100),
    ]
    # max last_step = 100, 100//100 = 1
    codeflash_output = _estimate_min_resource(trials) # 1.10μs -> 1.19μs (7.08% slower)

def test_max_last_step_less_than_100():
    # All last_steps < 100, should return 1
    trials = [
        FrozenTrial(TrialState.COMPLETE, 10),
        FrozenTrial(TrialState.COMPLETE, 99),
        FrozenTrial(TrialState.COMPLETE, 50),
    ]
    codeflash_output = _estimate_min_resource(trials) # 1.05μs -> 1.20μs (12.6% slower)

def test_max_last_step_exactly_100():
    # max last_step exactly 100
    trials = [
        FrozenTrial(TrialState.COMPLETE, 100),
        FrozenTrial(TrialState.COMPLETE, 80),
    ]
    # 100//100 = 1
    codeflash_output = _estimate_min_resource(trials) # 970ns -> 1.13μs (14.2% slower)

def test_max_last_step_just_over_100():
    # max last_step just over 100
    trials = [
        FrozenTrial(TrialState.COMPLETE, 101),
        FrozenTrial(TrialState.COMPLETE, 100),
    ]
    # 101//100 = 1
    codeflash_output = _estimate_min_resource(trials) # 941ns -> 1.10μs (14.8% slower)

def test_max_last_step_large_but_not_multiple_of_100():
    # max last_step = 999, 999//100 = 9
    trials = [
        FrozenTrial(TrialState.COMPLETE, 999),
        FrozenTrial(TrialState.COMPLETE, 800),
    ]
    codeflash_output = _estimate_min_resource(trials) # 986ns -> 1.12μs (12.2% slower)

def test_last_step_zero():
    # last_step = 0, should return 1 (since max(0//100, 1) == 1)
    trials = [
        FrozenTrial(TrialState.COMPLETE, 0),
    ]
    codeflash_output = _estimate_min_resource(trials) # 819ns -> 1.09μs (24.7% slower)

def test_negative_last_step():
    # Negative last_step, should still compute (max(-5//100, 1) == 1)
    trials = [
        FrozenTrial(TrialState.COMPLETE, -5),
        FrozenTrial(TrialState.COMPLETE, -100),
    ]
    codeflash_output = _estimate_min_resource(trials) # 985ns -> 1.13μs (13.1% slower)

def test_mixed_positive_and_negative_last_steps():
    # Should use the maximum last_step, even if negative values are present
    trials = [
        FrozenTrial(TrialState.COMPLETE, -100),
        FrozenTrial(TrialState.COMPLETE, 200),
        FrozenTrial(TrialState.COMPLETE, 0),
    ]
    # max last_step = 200, 200//100 = 2
    codeflash_output = _estimate_min_resource(trials) # 1.01μs -> 1.15μs (11.9% slower)

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

def test_large_number_of_trials_all_complete():
    # 1000 COMPLETE trials with increasing last_step
    trials = [FrozenTrial(TrialState.COMPLETE, i) for i in range(1, 1001)]
    # max last_step = 1000, 1000//100 = 10
    codeflash_output = _estimate_min_resource(trials) # 54.4μs -> 20.9μs (161% faster)

def test_large_number_of_trials_mixed_states():
    # 500 COMPLETE, 500 PRUNED, last_step from 1 to 1000
    trials = [FrozenTrial(TrialState.COMPLETE, i) for i in range(1, 501)] + \
             [FrozenTrial(TrialState.PRUNED, i) for i in range(501, 1001)]
    # max last_step among COMPLETE = 500, 500//100 = 5
    codeflash_output = _estimate_min_resource(trials) # 55.7μs -> 20.8μs (167% faster)

def test_large_number_of_trials_many_none_last_step():
    # 1000 trials, half COMPLETE with None last_step, half with increasing last_step
    trials = [FrozenTrial(TrialState.COMPLETE, None) for _ in range(500)] + \
             [FrozenTrial(TrialState.COMPLETE, i) for i in range(1, 501)]
    # max last_step among COMPLETE = 500, 500//100 = 5
    codeflash_output = _estimate_min_resource(trials) # 54.8μs -> 20.7μs (166% faster)

def test_large_number_of_trials_all_incomplete():
    # 1000 trials, none COMPLETE
    trials = [FrozenTrial(TrialState.PRUNED, i) for i in range(1, 1001)]
    codeflash_output = _estimate_min_resource(trials) # 54.4μs -> 20.6μs (164% faster)

def test_large_number_of_trials_with_negative_and_positive_last_steps():
    # 500 COMPLETE with negative last_steps, 500 with positive
    trials = [FrozenTrial(TrialState.COMPLETE, -i) for i in range(1, 501)] + \
             [FrozenTrial(TrialState.COMPLETE, i) for i in range(1, 501)]
    # max last_step = 500, 500//100 = 5
    codeflash_output = _estimate_min_resource(trials) # 55.6μs -> 21.1μs (164% faster)

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

def test_determinism():
    # Running the function twice on the same input should give same result
    trials = [
        FrozenTrial(TrialState.COMPLETE, 123),
        FrozenTrial(TrialState.COMPLETE, 456),
        FrozenTrial(TrialState.PRUNED, 789),
    ]
    codeflash_output = _estimate_min_resource(trials); result1 = codeflash_output # 1.16μs -> 1.29μs (10.3% slower)
    codeflash_output = _estimate_min_resource(trials); result2 = codeflash_output # 514ns -> 525ns (2.10% slower)

# ---------------
# Unusual Types Test
# ---------------

def test_non_integer_last_step():
    # If last_step is float, should still work (Python's // with float returns float)
    trials = [
        FrozenTrial(TrialState.COMPLETE, 150.0),
        FrozenTrial(TrialState.COMPLETE, 99.9),
    ]
    # max last_step = 150.0, 150.0//100 = 1.0, max(1.0, 1) == 1.0
    codeflash_output = _estimate_min_resource(trials); result = codeflash_output # 997ns -> 1.16μs (13.9% slower)


#------------------------------------------------
from __future__ import annotations

# imports
import pytest  # used for our unit tests
from optuna.pruners._successive_halving import _estimate_min_resource


# Minimal mock of optuna.trial._state.TrialState and optuna.trial.FrozenTrial for testing
class TrialState:
    COMPLETE = "COMPLETE"
    RUNNING = "RUNNING"
    PRUNED = "PRUNED"
    FAIL = "FAIL"
    WAITING = "WAITING"

class FrozenTrial:
    def __init__(self, state, last_step):
        self.state = state
        self.last_step = last_step
from optuna.pruners._successive_halving import _estimate_min_resource

# unit tests

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

def test_single_complete_trial_with_last_step():
    # Single trial, COMPLETE, last_step=100
    trials = [FrozenTrial(TrialState.COMPLETE, 100)]
    codeflash_output = _estimate_min_resource(trials) # 1.28μs -> 1.38μs (7.02% slower)

def test_multiple_complete_trials_with_varied_last_step():
    # Multiple COMPLETE trials, last_step=50, 200, 300
    trials = [
        FrozenTrial(TrialState.COMPLETE, 50),
        FrozenTrial(TrialState.COMPLETE, 200),
        FrozenTrial(TrialState.COMPLETE, 300)
    ]
    # max(last_step)=300, 300//100=3
    codeflash_output = _estimate_min_resource(trials) # 1.28μs -> 1.39μs (8.33% slower)

def test_complete_and_noncomplete_trials():
    # Only COMPLETE trials with last_step count; others ignored
    trials = [
        FrozenTrial(TrialState.COMPLETE, 100),
        FrozenTrial(TrialState.RUNNING, 1000),
        FrozenTrial(TrialState.PRUNED, 1000),
        FrozenTrial(TrialState.FAIL, 1000),
        FrozenTrial(TrialState.COMPLETE, 200)
    ]
    # Only two COMPLETE trials: 100, 200 -> max=200 //100=2
    codeflash_output = _estimate_min_resource(trials) # 1.34μs -> 1.37μs (1.97% slower)

def test_complete_trials_with_none_last_step():
    # Some COMPLETE trials with None last_step
    trials = [
        FrozenTrial(TrialState.COMPLETE, None),
        FrozenTrial(TrialState.COMPLETE, 150),
        FrozenTrial(TrialState.COMPLETE, None),
        FrozenTrial(TrialState.COMPLETE, 50)
    ]
    # Only 150 and 50 are counted
    codeflash_output = _estimate_min_resource(trials) # 1.14μs -> 1.30μs (12.3% slower)

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

def test_no_trials():
    # Empty list
    trials = []
    codeflash_output = _estimate_min_resource(trials) # 645ns -> 1.03μs (37.3% slower)

def test_no_complete_trials():
    # No COMPLETE trials
    trials = [
        FrozenTrial(TrialState.RUNNING, 100),
        FrozenTrial(TrialState.PRUNED, 200)
    ]
    codeflash_output = _estimate_min_resource(trials) # 1.08μs -> 1.19μs (8.52% slower)

def test_complete_trials_all_none_last_step():
    # All COMPLETE trials have None last_step
    trials = [
        FrozenTrial(TrialState.COMPLETE, None),
        FrozenTrial(TrialState.COMPLETE, None)
    ]
    codeflash_output = _estimate_min_resource(trials) # 1.06μs -> 1.22μs (13.4% slower)

def test_last_step_zero():
    # COMPLETE trial with last_step=0
    trials = [FrozenTrial(TrialState.COMPLETE, 0)]
    # max(0)//100=0, but min resource is at least 1
    codeflash_output = _estimate_min_resource(trials) # 910ns -> 1.16μs (21.7% slower)

def test_last_step_less_than_100():
    # COMPLETE trial with last_step=99
    trials = [FrozenTrial(TrialState.COMPLETE, 99)]
    # 99//100=0, but min resource is at least 1
    codeflash_output = _estimate_min_resource(trials) # 957ns -> 1.17μs (18.2% slower)

def test_last_step_exactly_100():
    # COMPLETE trial with last_step=100
    trials = [FrozenTrial(TrialState.COMPLETE, 100)]
    codeflash_output = _estimate_min_resource(trials) # 890ns -> 1.10μs (18.9% slower)

def test_last_step_just_over_100():
    # COMPLETE trial with last_step=101
    trials = [FrozenTrial(TrialState.COMPLETE, 101)]
    codeflash_output = _estimate_min_resource(trials) # 887ns -> 1.09μs (18.8% slower)

def test_last_step_199():
    # COMPLETE trial with last_step=199
    trials = [FrozenTrial(TrialState.COMPLETE, 199)]
    codeflash_output = _estimate_min_resource(trials) # 942ns -> 1.12μs (16.0% slower)

def test_last_step_200():
    # COMPLETE trial with last_step=200
    trials = [FrozenTrial(TrialState.COMPLETE, 200)]
    codeflash_output = _estimate_min_resource(trials) # 882ns -> 1.12μs (21.2% slower)

def test_non_integer_last_step():
    # Should handle float last_step by integer division
    trials = [FrozenTrial(TrialState.COMPLETE, 250.9)]
    # 250.9//100 = 2.0, max(2.0,1)=2.0
    codeflash_output = _estimate_min_resource(trials) # 930ns -> 1.07μs (12.8% slower)

def test_negative_last_step():
    # Negative last_step (should still apply //100, but negative //100 is -1)
    trials = [FrozenTrial(TrialState.COMPLETE, -50)]
    # max(-50//100, 1) == max(-1, 1) == 1
    codeflash_output = _estimate_min_resource(trials) # 820ns -> 1.12μs (27.1% slower)

def test_mixed_states_and_last_steps():
    # Mixed states and last_step values
    trials = [
        FrozenTrial(TrialState.COMPLETE, 10),
        FrozenTrial(TrialState.PRUNED, 1000),
        FrozenTrial(TrialState.COMPLETE, None),
        FrozenTrial(TrialState.RUNNING, 100),
        FrozenTrial(TrialState.COMPLETE, 500)
    ]
    # Only 10 and 500, max=500//100=5
    codeflash_output = _estimate_min_resource(trials) # 1.25μs -> 1.31μs (4.28% slower)

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

def test_large_number_of_trials_all_complete():
    # 1000 COMPLETE trials, last_step from 1 to 1000
    trials = [FrozenTrial(TrialState.COMPLETE, i) for i in range(1, 1001)]
    # max(last_step)=1000, 1000//100=10
    codeflash_output = _estimate_min_resource(trials) # 55.3μs -> 21.1μs (163% faster)

def test_large_number_of_trials_mixed_states():
    # 500 COMPLETE (last_step=0..499), 500 PRUNED (last_step=500..999)
    trials = (
        [FrozenTrial(TrialState.COMPLETE, i) for i in range(500)]
        + [FrozenTrial(TrialState.PRUNED, i) for i in range(500, 1000)]
    )
    # max(last_step for COMPLETE)=499, 499//100=4
    codeflash_output = _estimate_min_resource(trials) # 56.0μs -> 21.3μs (162% faster)

def test_large_number_of_trials_some_none_last_step():
    # 1000 COMPLETE, half with None, half with last_step=1000
    trials = (
        [FrozenTrial(TrialState.COMPLETE, None) for _ in range(500)]
        + [FrozenTrial(TrialState.COMPLETE, 1000) for _ in range(500)]
    )
    # max(last_step)=1000, 1000//100=10
    codeflash_output = _estimate_min_resource(trials) # 56.8μs -> 20.9μs (172% faster)

def test_large_number_of_trials_all_noncomplete():
    # 1000 PRUNED trials, should return None
    trials = [FrozenTrial(TrialState.PRUNED, i) for i in range(1000)]
    codeflash_output = _estimate_min_resource(trials) # 55.7μs -> 21.1μs (164% faster)

def test_large_number_of_trials_all_none_last_step():
    # 1000 COMPLETE trials, all None last_step
    trials = [FrozenTrial(TrialState.COMPLETE, None) for _ in range(1000)]
    codeflash_output = _estimate_min_resource(trials) # 56.8μs -> 20.9μs (172% faster)

def test_large_number_of_trials_high_last_step():
    # 1000 COMPLETE trials, last_step=1_000_000
    trials = [FrozenTrial(TrialState.COMPLETE, 1_000_000) for _ in range(1000)]
    # max(last_step)=1_000_000, //100=10_000
    codeflash_output = _estimate_min_resource(trials) # 55.6μs -> 20.8μs (168% faster)

def test_large_number_of_trials_mixed_extremes():
    # 500 COMPLETE with last_step=1, 500 COMPLETE with last_step=999
    trials = (
        [FrozenTrial(TrialState.COMPLETE, 1) for _ in range(500)]
        + [FrozenTrial(TrialState.COMPLETE, 999) for _ in range(500)]
    )
    # max(last_step)=999, 999//100=9
    codeflash_output = _estimate_min_resource(trials) # 56.4μs -> 21.1μs (168% faster)
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.
#------------------------------------------------
from optuna.pruners._successive_halving import _estimate_min_resource

def test__estimate_min_resource():
    _estimate_min_resource([])

🔎 Concolic Coverage Tests and Runtime

Test File::Test Function	Original ⏱️	Optimized ⏱️	Speedup
codeflash_concolic_bg1jh046/tmp9fbjf2ge/test_concolic_coverage.py::test__estimate_min_resource	796ns	1.45μs	-45.1%⚠️

To edit these changes git checkout codeflash/optimize-_estimate_min_resource-mho97akg and push.