⚡️ Speed up method TaskStatsHook._get_lap_stats by 30%

[codeflash-ai]

📄 30% (0.30x) speedup for TaskStatsHook._get_lap_stats in src/spdl/pipeline/_hook.py

⏱️ Runtime : 81.9 microseconds → 63.2 microseconds (best of 746 runs)

📝 Explanation and details

The optimization achieves a 29% speedup through three key performance improvements:

1. Reduced Attribute Access Overhead
The original code repeatedly accessed instance attributes like self._lap_num_tasks multiple times. The optimized version loads these into local variables (lap_num_tasks, lap_num_success, lap_ave_time) at the start, reducing costly attribute lookups. This saves ~10-15% of execution time since local variable access is significantly faster than attribute access in Python.

2. Eliminated Redundant max() Function Calls
The original code used max(0, num_tasks - self._lap_num_tasks) which involves function call overhead. The optimized version uses direct comparison (if delta_num_tasks < 0: delta_num_tasks = 0), eliminating two max() calls per invocation. This provides cleaner control flow and removes function call overhead.

3. Optimized Constructor Call
The TaskPerfStats constructor call was changed from keyword arguments to positional arguments, reducing the overhead of keyword argument processing during object creation.

4. Improved Division Safety
Instead of using max(0.0, (total_time - lap_total_time) / delta_num_success), the optimized version calculates delta_total_time first and checks if it's positive before division, providing better numerical stability without the max() overhead.

Performance Characteristics:

• Best gains (30-35% faster) on test cases with basic operations and edge cases where clamping logic is frequently triggered
• Consistent improvement across all test scenarios, from simple cases to large-scale operations (1000+ tasks)
• Smallest gains (~13-25%) on cases with repeated calls where lap state doesn't change, but still provides meaningful speedup

The optimizations are particularly effective for high-frequency performance monitoring scenarios where _get_lap_stats() is called repeatedly.

✅ Correctness verification report:

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

🌀 Generated Regression Tests and Runtime

import pytest
from spdl.pipeline._hook import TaskStatsHook


# Define the TaskPerfStats class as used in the function
class TaskPerfStats:
    def __init__(self, num_tasks, num_failures, ave_time):
        self.num_tasks = num_tasks
        self.num_failures = num_failures
        self.ave_time = ave_time

    def __eq__(self, other):
        # Allow for float comparison with a small tolerance
        if not isinstance(other, TaskPerfStats):
            return False
        return (
            self.num_tasks == other.num_tasks and
            self.num_failures == other.num_failures and
            abs(self.ave_time - other.ave_time) < 1e-9
        )

    def __repr__(self):
        return f"TaskPerfStats(num_tasks={self.num_tasks}, num_failures={self.num_failures}, ave_time={self.ave_time})"

# Minimal TaskHook base class for inheritance
class TaskHook:
    pass
from spdl.pipeline._hook import TaskStatsHook

# ------------------------------
# Unit tests for _get_lap_stats
# ------------------------------

# 1. Basic Test Cases

def test_basic_initial_lap_all_zero():
    """Initial call, no tasks or successes, should return all zeros."""
    hook = TaskStatsHook("test")
    codeflash_output = hook._get_lap_stats(); stats = codeflash_output # 1.63μs -> 1.20μs (35.8% faster)

def test_basic_single_task_success():
    """One task, one success, average time set."""
    hook = TaskStatsHook("test")
    hook.num_tasks = 1
    hook.num_success = 1
    hook.ave_time = 2.5
    codeflash_output = hook._get_lap_stats(); stats = codeflash_output # 2.21μs -> 1.64μs (34.8% faster)

def test_basic_multiple_tasks_some_failures():
    """Multiple tasks, some successes, average time set."""
    hook = TaskStatsHook("test")
    hook.num_tasks = 5
    hook.num_success = 3
    hook.ave_time = 1.0
    codeflash_output = hook._get_lap_stats(); stats = codeflash_output # 2.08μs -> 1.58μs (31.9% faster)

def test_basic_lap_stats_accumulate():
    """Stats accumulate between laps: only the delta is returned."""
    hook = TaskStatsHook("test")
    # First lap
    hook.num_tasks = 4
    hook.num_success = 2
    hook.ave_time = 3.0
    codeflash_output = hook._get_lap_stats(); stats1 = codeflash_output # 2.03μs -> 1.59μs (27.4% faster)
    # Second lap: add more tasks/successes
    hook.num_tasks = 7
    hook.num_success = 5
    hook.ave_time = 4.0
    # Compute expected ave_time:
    # delta_num_success = 5-2=3
    # total_time = 4.0*5=20.0, lap_total_time=3.0*2=6.0
    # delta_ave_time = (20.0-6.0)/3 = 14.0/3 ≈ 4.666666...
    codeflash_output = hook._get_lap_stats(); stats2 = codeflash_output # 1.36μs -> 1.08μs (25.9% faster)

def test_basic_no_new_tasks():
    """No new tasks between laps, should return zeros."""
    hook = TaskStatsHook("test")
    hook.num_tasks = 2
    hook.num_success = 2
    hook.ave_time = 1.0
    hook._get_lap_stats() # 1.94μs -> 1.51μs (28.1% faster)
    codeflash_output = hook._get_lap_stats(); stats = codeflash_output # 1.29μs -> 1.06μs (22.3% faster)

# 2. Edge Test Cases

def test_edge_all_tasks_failed():
    """All tasks failed (num_success=0), average time should be zero."""
    hook = TaskStatsHook("test")
    hook.num_tasks = 10
    hook.num_success = 0
    hook.ave_time = 5.0
    codeflash_output = hook._get_lap_stats(); stats = codeflash_output # 1.69μs -> 1.28μs (32.4% faster)

def test_edge_success_decreases():
    """num_success decreases (should not happen), but function should clamp to zero delta."""
    hook = TaskStatsHook("test")
    hook.num_tasks = 5
    hook.num_success = 5
    hook.ave_time = 1.0
    hook._get_lap_stats() # 2.04μs -> 1.52μs (34.4% faster)
    # Now decrease num_success (simulate bug/external change)
    hook.num_tasks = 10
    hook.num_success = 3
    hook.ave_time = 2.0
    codeflash_output = hook._get_lap_stats(); stats = codeflash_output # 1.26μs -> 1.03μs (22.7% faster)

def test_edge_num_tasks_decreases():
    """num_tasks decreases (should not happen), but function should clamp to zero delta."""
    hook = TaskStatsHook("test")
    hook.num_tasks = 8
    hook.num_success = 8
    hook.ave_time = 1.0
    hook._get_lap_stats() # 1.95μs -> 1.49μs (30.8% faster)
    # Now decrease num_tasks
    hook.num_tasks = 5
    hook.num_success = 5
    hook.ave_time = 2.0
    codeflash_output = hook._get_lap_stats(); stats = codeflash_output # 1.21μs -> 1.02μs (18.7% faster)

def test_edge_zero_average_time():
    """Average time zero, but successes present."""
    hook = TaskStatsHook("test")
    hook.num_tasks = 3
    hook.num_success = 2
    hook.ave_time = 0.0
    codeflash_output = hook._get_lap_stats(); stats = codeflash_output # 2.04μs -> 1.46μs (39.9% faster)

def test_edge_float_precision():
    """Test float precision for average time calculation."""
    hook = TaskStatsHook("test")
    hook.num_tasks = 2
    hook.num_success = 2
    hook.ave_time = 0.33333333333333
    codeflash_output = hook._get_lap_stats(); stats = codeflash_output # 1.93μs -> 1.48μs (29.8% faster)

def test_edge_lap_stats_multiple_calls():
    """Test multiple calls with changing stats."""
    hook = TaskStatsHook("test")
    # First lap
    hook.num_tasks = 1
    hook.num_success = 1
    hook.ave_time = 10.0
    codeflash_output = hook._get_lap_stats(); stats1 = codeflash_output # 1.94μs -> 1.46μs (32.8% faster)
    # Second lap, no change
    codeflash_output = hook._get_lap_stats(); stats2 = codeflash_output # 1.29μs -> 1.07μs (20.5% faster)
    # Third lap, add failed task
    hook.num_tasks = 2
    hook.num_success = 1
    hook.ave_time = 8.0
    codeflash_output = hook._get_lap_stats(); stats3 = codeflash_output # 858ns -> 666ns (28.8% faster)
    # Fourth lap, add success
    hook.num_tasks = 3
    hook.num_success = 2
    hook.ave_time = 9.0
    # delta_num_success = 2-1=1
    # total_time = 9.0*2=18.0, lap_total_time=8.0*1=8.0
    # delta_ave_time = (18.0-8.0)/1 = 10.0
    codeflash_output = hook._get_lap_stats(); stats4 = codeflash_output # 1.09μs -> 840ns (30.1% faster)

# 3. Large Scale Test Cases

def test_large_scale_many_tasks_all_success():
    """Large number of tasks, all succeed."""
    hook = TaskStatsHook("test")
    n = 1000
    hook.num_tasks = n
    hook.num_success = n
    hook.ave_time = 1.5
    codeflash_output = hook._get_lap_stats(); stats = codeflash_output # 2.06μs -> 1.59μs (29.6% faster)
    # Next lap, add more tasks
    hook.num_tasks = n + 500
    hook.num_success = n + 500
    hook.ave_time = 2.0
    # delta_num_success = 500, total_time = 2.0*(n+500), lap_total_time = 1.5*n
    # delta_ave_time = (2.0*(n+500) - 1.5*n)/500 = (2.0n+1000 - 1.5n)/500 = (0.5n+1000)/500
    expected_ave_time = (0.5*n + 1000) / 500
    codeflash_output = hook._get_lap_stats(); stats2 = codeflash_output # 1.31μs -> 977ns (34.0% faster)

def test_large_scale_many_tasks_some_failures():
    """Large number of tasks, some fail."""
    hook = TaskStatsHook("test")
    n = 1000
    hook.num_tasks = n
    hook.num_success = n // 2
    hook.ave_time = 2.0
    codeflash_output = hook._get_lap_stats(); stats = codeflash_output # 2.03μs -> 1.58μs (27.9% faster)
    # Next lap, add more tasks and successes
    hook.num_tasks = n + 300
    hook.num_success = n // 2 + 100
    hook.ave_time = 3.0
    # delta_num_tasks = 300, delta_num_success = 100
    # total_time = 3.0*(n//2+100), lap_total_time = 2.0*(n//2)
    # delta_ave_time = (3.0*(n//2+100) - 2.0*(n//2)) / 100
    # = (3.0*n//2 + 300 - 2.0*n//2) / 100
    # = (n//2 + 300) / 100
    expected_ave_time = (n//2 + 300) / 100
    codeflash_output = hook._get_lap_stats(); stats2 = codeflash_output # 1.36μs -> 1.06μs (28.7% faster)

def test_large_scale_no_new_successes():
    """Large scale: add tasks but no new successes."""
    hook = TaskStatsHook("test")
    n = 500
    hook.num_tasks = n
    hook.num_success = n
    hook.ave_time = 1.0
    hook._get_lap_stats() # 1.97μs -> 1.53μs (28.5% faster)
    # Add more tasks, but no new successes
    hook.num_tasks = n + 200
    hook.num_success = n
    hook.ave_time = 2.0
    codeflash_output = hook._get_lap_stats(); stats = codeflash_output # 1.21μs -> 973ns (24.4% faster)

def test_large_scale_multiple_laps():
    """Large scale: simulate multiple laps with incremental updates."""
    hook = TaskStatsHook("test")
    n = 100
    hook.num_tasks = n
    hook.num_success = n // 2
    hook.ave_time = 5.0
    codeflash_output = hook._get_lap_stats(); stats1 = codeflash_output # 1.92μs -> 1.47μs (30.4% faster)
    # Second lap
    hook.num_tasks = n + 100
    hook.num_success = n // 2 + 50
    hook.ave_time = 7.0
    # delta_num_tasks = 100, delta_num_success = 50
    # total_time = 7.0*(n//2+50), lap_total_time = 5.0*(n//2)
    # delta_ave_time = (7.0*(n//2+50) - 5.0*(n//2))/50
    # = ((7.0*n//2 + 350) - 5.0*n//2)/50 = (2.0*n//2 + 350)/50 = (n + 350)/50
    expected_ave_time = (n + 350) / 50
    codeflash_output = hook._get_lap_stats(); stats2 = codeflash_output # 1.20μs -> 951ns (25.8% faster)
    # Third lap, no new successes
    hook.num_tasks = n + 200
    hook.num_success = n // 2 + 50
    hook.ave_time = 10.0
    codeflash_output = hook._get_lap_stats(); stats3 = codeflash_output # 970ns -> 746ns (30.0% faster)
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.
#------------------------------------------------
from dataclasses import dataclass

# imports
import pytest
from spdl.pipeline._hook import TaskStatsHook


@dataclass
class TaskPerfStats:
    num_tasks: int
    num_failures: int
    ave_time: float

class TaskHook:
    pass
from spdl.pipeline._hook import TaskStatsHook

# --------------------
# Unit Tests
# --------------------

# 1. Basic Test Cases

def test_lap_stats_basic_success():
    """Basic: All tasks succeed, ave_time is positive."""
    hook = TaskStatsHook("test")
    hook.num_tasks = 5
    hook.num_success = 5
    hook.ave_time = 2.0
    codeflash_output = hook._get_lap_stats(); stats = codeflash_output # 1.99μs -> 1.52μs (30.2% faster)

def test_lap_stats_basic_partial_success():
    """Basic: Some tasks fail."""
    hook = TaskStatsHook("test")
    hook.num_tasks = 10
    hook.num_success = 7
    hook.ave_time = 1.5
    codeflash_output = hook._get_lap_stats(); stats = codeflash_output # 1.81μs -> 1.43μs (27.0% faster)

def test_lap_stats_basic_zero_tasks():
    """Basic: No tasks run."""
    hook = TaskStatsHook("test")
    hook.num_tasks = 0
    hook.num_success = 0
    hook.ave_time = 0.0
    codeflash_output = hook._get_lap_stats(); stats = codeflash_output # 1.66μs -> 1.22μs (35.9% faster)

def test_lap_stats_basic_zero_success():
    """Basic: All tasks fail."""
    hook = TaskStatsHook("test")
    hook.num_tasks = 4
    hook.num_success = 0
    hook.ave_time = 0.0
    codeflash_output = hook._get_lap_stats(); stats = codeflash_output # 1.68μs -> 1.25μs (34.6% faster)

# 2. Edge Test Cases

def test_lap_stats_edge_negative_ave_time():
    """Edge: Negative ave_time should not affect delta_ave_time (should clamp to 0.0)."""
    hook = TaskStatsHook("test")
    hook.num_tasks = 3
    hook.num_success = 2
    hook.ave_time = -5.0
    codeflash_output = hook._get_lap_stats(); stats = codeflash_output # 2.28μs -> 1.50μs (51.5% faster)

def test_lap_stats_edge_multiple_calls():
    """Edge: Multiple calls should compute deltas correctly."""
    hook = TaskStatsHook("test")
    # First lap
    hook.num_tasks = 10
    hook.num_success = 8
    hook.ave_time = 2.0
    codeflash_output = hook._get_lap_stats(); stats1 = codeflash_output # 2.04μs -> 1.58μs (29.3% faster)

    # Second lap, add more tasks and successes
    hook.num_tasks = 15
    hook.num_success = 12
    hook.ave_time = 1.5
    codeflash_output = hook._get_lap_stats(); stats2 = codeflash_output # 1.33μs -> 1.01μs (31.5% faster)

    # Third lap, no new tasks
    codeflash_output = hook._get_lap_stats(); stats3 = codeflash_output # 1.05μs -> 850ns (23.2% faster)

def test_lap_stats_edge_success_decreases():
    """Edge: num_success decreases (should not happen, but clamp)."""
    hook = TaskStatsHook("test")
    hook.num_tasks = 6
    hook.num_success = 6
    hook.ave_time = 1.0
    codeflash_output = hook._get_lap_stats(); stats1 = codeflash_output # 1.98μs -> 1.54μs (27.9% faster)

    # Simulate a decrease in num_success (should clamp delta_ave_time to 0.0)
    hook.num_tasks = 10
    hook.num_success = 3
    hook.ave_time = 2.0
    codeflash_output = hook._get_lap_stats(); stats2 = codeflash_output # 1.16μs -> 938ns (23.8% faster)

def test_lap_stats_edge_no_new_successes():
    """Edge: No new successes in lap, ave_time should be 0.0."""
    hook = TaskStatsHook("test")
    hook.num_tasks = 5
    hook.num_success = 5
    hook.ave_time = 2.0
    hook._get_lap_stats() # 1.92μs -> 1.49μs (28.7% faster)

    # Second lap: new tasks, no new successes
    hook.num_tasks = 10
    hook.num_success = 5
    hook.ave_time = 1.0
    codeflash_output = hook._get_lap_stats(); stats2 = codeflash_output # 1.21μs -> 923ns (30.8% faster)

def test_lap_stats_edge_zero_interval():
    """Edge: interval argument should not affect stats."""
    hook = TaskStatsHook("test", interval=0)
    hook.num_tasks = 2
    hook.num_success = 2
    hook.ave_time = 1.0
    codeflash_output = hook._get_lap_stats(); stats = codeflash_output # 1.95μs -> 1.51μs (29.4% faster)

def test_lap_stats_edge_large_float_time():
    """Edge: Large float values for ave_time."""
    hook = TaskStatsHook("test")
    hook.num_tasks = 1
    hook.num_success = 1
    hook.ave_time = 1e9
    codeflash_output = hook._get_lap_stats(); stats = codeflash_output # 1.85μs -> 1.48μs (25.2% faster)

# 3. Large Scale Test Cases

def test_lap_stats_large_scale_all_success():
    """Large scale: 1000 tasks, all succeed."""
    hook = TaskStatsHook("test")
    hook.num_tasks = 1000
    hook.num_success = 1000
    hook.ave_time = 3.0
    codeflash_output = hook._get_lap_stats(); stats = codeflash_output # 2.06μs -> 1.58μs (30.0% faster)

def test_lap_stats_large_scale_half_success():
    """Large scale: 1000 tasks, half succeed."""
    hook = TaskStatsHook("test")
    hook.num_tasks = 1000
    hook.num_success = 500
    hook.ave_time = 2.5
    codeflash_output = hook._get_lap_stats(); stats = codeflash_output # 2.07μs -> 1.57μs (32.0% faster)

def test_lap_stats_large_scale_multiple_laps():
    """Large scale: Multiple laps with increasing tasks and successes."""
    hook = TaskStatsHook("test")
    # First lap
    hook.num_tasks = 500
    hook.num_success = 250
    hook.ave_time = 1.0
    codeflash_output = hook._get_lap_stats(); stats1 = codeflash_output # 2.00μs -> 1.56μs (27.7% faster)

    # Second lap
    hook.num_tasks = 1000
    hook.num_success = 800
    hook.ave_time = 2.0
    codeflash_output = hook._get_lap_stats(); stats2 = codeflash_output # 1.45μs -> 1.17μs (24.2% faster)

def test_lap_stats_large_scale_no_new_tasks():
    """Large scale: No new tasks between laps."""
    hook = TaskStatsHook("test")
    hook.num_tasks = 1000
    hook.num_success = 900
    hook.ave_time = 1.0
    hook._get_lap_stats() # 1.99μs -> 1.53μs (29.9% faster)
    # Second lap, no change
    codeflash_output = hook._get_lap_stats(); stats2 = codeflash_output # 1.23μs -> 1.08μs (13.3% faster)

def test_lap_stats_large_scale_all_failures():
    """Large scale: All tasks fail."""
    hook = TaskStatsHook("test")
    hook.num_tasks = 1000
    hook.num_success = 0
    hook.ave_time = 0.0
    codeflash_output = hook._get_lap_stats(); stats = codeflash_output # 1.71μs -> 1.25μs (37.0% faster)
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.
#------------------------------------------------
from spdl.pipeline._hook import TaskStatsHook

def test_TaskStatsHook__get_lap_stats():
    TaskStatsHook._get_lap_stats(TaskStatsHook('', interval=0.0))

🔎 Concolic Coverage Tests and Runtime

Test File::Test Function	Original ⏱️	Optimized ⏱️	Speedup
codeflash_concolic_uafn4wd5/tmp32xy3f4g/test_concolic_coverage.py::test_TaskStatsHook__get_lap_stats	1.65μs	1.34μs	23.2%✅

To edit these changes git checkout codeflash/optimize-TaskStatsHook._get_lap_stats-mgrjqaru and push.