⚡️ Speed up function _gather_error by 30%

[codeflash-ai]

📄 30% (0.30x) speedup for _gather_error in src/spdl/pipeline/_node.py

⏱️ Runtime : 1.19 milliseconds → 914 microseconds (best of 431 runs)

📝 Explanation and details

The optimization adds a conditional check if len(errs) > 1: before sorting the error list. This simple change provides a 30% speedup by avoiding unnecessary sorting operations when there are 0 or 1 errors.

Key optimization:

• Conditional sorting: Only sorts when len(errs) > 1, since lists with 0-1 elements are already "sorted"
• Line profiler evidence: Shows sorting went from 3,086 calls (every function call) to only 80 calls (when multiple errors exist)

Why this works:

• Python's list.sort() has overhead even for small lists - it still needs to allocate temporary space and perform comparisons
• In many practical scenarios, most nodes have 0-1 errors, making sorting unnecessary
• The len() check is extremely fast (O(1)) compared to sorting overhead

Performance by test case:

• Best gains (30-35%): Large-scale tests with many nodes having few/no errors (test_performance_with_large_flat_upstream, test_large_mixed_tree_with_some_errors)
• Moderate gains (15-30%): Single nodes and small hierarchies where sorting is frequently skipped
• Minimal impact (2-5% slower): Cases with multiple errors where sorting still occurs, due to added length check overhead

This optimization is particularly effective for error collection in pipeline scenarios where most nodes succeed and only occasional failures need sorting.

✅ Correctness verification report:

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

🌀 Generated Regression Tests and Runtime

from typing import Generic, List, Tuple, TypeVar

# imports
import pytest
from spdl.pipeline._node import _gather_error

T = TypeVar("T")

# Minimal _Node and Task mocks for testing
class MockTask:
    def __init__(self, name, cancelled=False, exception=None):
        self._name = name
        self._cancelled = cancelled
        self._exception = exception

    def cancelled(self):
        return self._cancelled

    def exception(self):
        return self._exception

    def get_name(self):
        return self._name

class _Node(Generic[T]):
    def __init__(self, task: MockTask, upstream: List['_Node[T]'] = None):
        self.task = task
        self.upstream = upstream if upstream is not None else []

# unit tests

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

def test_no_error_and_no_upstream():
    # Node with no error and no upstream
    node = _Node(MockTask("A"))
    codeflash_output = _gather_error(node) # 1.31μs -> 1.13μs (16.0% faster)

def test_error_and_no_upstream():
    # Node with an error and no upstream
    exc = ValueError("fail")
    node = _Node(MockTask("A", exception=exc))
    codeflash_output = _gather_error(node); result = codeflash_output # 1.51μs -> 1.26μs (19.9% faster)

def test_cancelled_task_no_error():
    # Node with cancelled task, no error
    node = _Node(MockTask("A", cancelled=True))
    codeflash_output = _gather_error(node) # 968ns -> 828ns (16.9% faster)

def test_upstream_with_error():
    # Node with upstream node that has error
    exc = RuntimeError("upstream error")
    upstream = _Node(MockTask("B", exception=exc))
    node = _Node(MockTask("A"), upstream=[upstream])
    codeflash_output = _gather_error(node); result = codeflash_output # 2.28μs -> 1.84μs (23.6% faster)

def test_multiple_upstream_errors():
    # Node with multiple upstream errors
    exc1 = Exception("err1")
    exc2 = Exception("err2")
    up1 = _Node(MockTask("B", exception=exc1))
    up2 = _Node(MockTask("C", exception=exc2))
    node = _Node(MockTask("A"), upstream=[up1, up2])
    codeflash_output = _gather_error(node); result = codeflash_output # 2.98μs -> 3.03μs (1.78% slower)

def test_error_on_node_and_upstream():
    # Node and upstream both have errors
    exc1 = Exception("err1")
    exc2 = Exception("err2")
    up = _Node(MockTask("B", exception=exc2))
    node = _Node(MockTask("A", exception=exc1), upstream=[up])
    codeflash_output = _gather_error(node); result = codeflash_output # 2.53μs -> 2.46μs (2.60% faster)

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

def test_empty_upstream_list():
    # Node with empty upstream list
    node = _Node(MockTask("A"), upstream=[])
    codeflash_output = _gather_error(node) # 1.11μs -> 950ns (17.2% faster)

def test_deeply_nested_upstream_errors():
    # Deep nesting of upstream nodes with errors
    exc1 = Exception("err1")
    exc2 = Exception("err2")
    exc3 = Exception("err3")
    n3 = _Node(MockTask("C", exception=exc3))
    n2 = _Node(MockTask("B", exception=exc2), upstream=[n3])
    n1 = _Node(MockTask("A", exception=exc1), upstream=[n2])
    codeflash_output = _gather_error(n1); result = codeflash_output # 3.25μs -> 3.34μs (2.81% slower)

def test_duplicate_task_names():
    # Upstream nodes with duplicate task names
    exc1 = Exception("err1")
    exc2 = Exception("err2")
    n1 = _Node(MockTask("A", exception=exc1))
    n2 = _Node(MockTask("A", exception=exc2))
    node = _Node(MockTask("B"), upstream=[n1, n2])
    codeflash_output = _gather_error(node); result = codeflash_output # 3.02μs -> 2.92μs (3.42% faster)

def test_upstream_with_cancelled_and_error():
    # Upstream node is cancelled but has an exception
    exc = Exception("err")
    up = _Node(MockTask("B", cancelled=True, exception=exc))
    node = _Node(MockTask("A"))
    # Cancelled tasks should not report errors
    codeflash_output = _gather_error(node) # 1.15μs -> 926ns (23.7% faster)

def test_upstream_with_none_exception():
    # Upstream node with exception() returning None
    up = _Node(MockTask("B", exception=None))
    node = _Node(MockTask("A"))
    codeflash_output = _gather_error(node) # 1.03μs -> 890ns (15.6% faster)

def test_upstream_with_mixed_errors_and_cancelled():
    exc1 = Exception("err1")
    up1 = _Node(MockTask("B", exception=exc1))
    up2 = _Node(MockTask("C", cancelled=True, exception=Exception("should not appear")))
    node = _Node(MockTask("A"), upstream=[up1, up2])
    codeflash_output = _gather_error(node); result = codeflash_output # 2.78μs -> 2.19μs (26.7% faster)

def test_sorting_of_errors_by_task_name():
    # Errors should be sorted by task name
    exc1 = Exception("err1")
    exc2 = Exception("err2")
    exc3 = Exception("err3")
    up1 = _Node(MockTask("Z", exception=exc1))
    up2 = _Node(MockTask("Y", exception=exc2))
    up3 = _Node(MockTask("X", exception=exc3))
    node = _Node(MockTask("A"), upstream=[up1, up2, up3])
    codeflash_output = _gather_error(node); result = codeflash_output # 3.69μs -> 3.44μs (7.03% faster)

def test_upstream_with_no_errors():
    # Upstream nodes without errors
    up1 = _Node(MockTask("B"))
    up2 = _Node(MockTask("C"))
    node = _Node(MockTask("A"), upstream=[up1, up2])
    codeflash_output = _gather_error(node) # 2.23μs -> 1.85μs (20.3% faster)

def test_upstream_is_none():
    # Upstream is None (should default to empty)
    node = _Node(MockTask("A"), upstream=None)
    codeflash_output = _gather_error(node) # 1.07μs -> 904ns (18.0% faster)

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

def test_many_upstream_nodes_with_errors():
    # Large number of upstream nodes with errors
    N = 500
    excs = [Exception(f"err{i}") for i in range(N)]
    upstreams = [_Node(MockTask(f"TASK{i}", exception=excs[i])) for i in range(N)]
    node = _Node(MockTask("ROOT"), upstream=upstreams)
    codeflash_output = _gather_error(node); result = codeflash_output # 227μs -> 174μs (30.6% faster)
    # Should be sorted by name, which is TASK0 ... TASK499
    expected = [(f"TASK{i}", excs[i]) for i in range(N)]

def test_large_nested_tree_of_errors():
    # Tree structure: root -> [level1 nodes] -> [level2 nodes]
    N = 10
    level2_excs = [Exception(f"L2_{i}") for i in range(N)]
    level2_nodes = [_Node(MockTask(f"L2_{i}", exception=level2_excs[i])) for i in range(N)]
    level1_excs = [Exception(f"L1_{i}") for i in range(N)]
    level1_nodes = [_Node(MockTask(f"L1_{i}", exception=level1_excs[i]), upstream=[level2_nodes[i]]) for i in range(N)]
    root = _Node(MockTask("ROOT"), upstream=level1_nodes)
    codeflash_output = _gather_error(root); result = codeflash_output # 13.3μs -> 12.5μs (5.79% faster)
    expected = []
    for i in range(N):
        expected.append((f"L1_{i}", level1_excs[i]))
        expected.append((f"L2_{i}", level2_excs[i]))
    # Should be sorted by name
    expected.sort(key=lambda i: i[0])

def test_large_mixed_tree_with_some_errors():
    # Large tree, only some nodes have errors
    N = 100
    upstreams = []
    expected = []
    for i in range(N):
        if i % 10 == 0:
            exc = Exception(f"err{i}")
            upstreams.append(_Node(MockTask(f"TASK{i}", exception=exc)))
            expected.append((f"TASK{i}", exc))
        else:
            upstreams.append(_Node(MockTask(f"TASK{i}")))
    node = _Node(MockTask("ROOT"), upstream=upstreams)
    expected.sort(key=lambda i: i[0])
    codeflash_output = _gather_error(node) # 33.1μs -> 25.0μs (32.5% faster)

def test_performance_with_large_flat_upstream():
    # Flat structure with many upstreams, no errors
    N = 999
    upstreams = [_Node(MockTask(f"TASK{i}")) for i in range(N)]
    node = _Node(MockTask("ROOT"), upstream=upstreams)
    codeflash_output = _gather_error(node) # 272μs -> 201μs (35.3% faster)

def test_performance_with_large_deep_chain():
    # Deep chain of nodes, only last has error
    N = 100
    exc = Exception("deep error")
    chain = _Node(MockTask(f"NODE{N}", exception=exc))
    for i in reversed(range(N)):
        chain = _Node(MockTask(f"NODE{i}"), upstream=[chain])
    codeflash_output = _gather_error(chain); result = codeflash_output # 34.1μs -> 25.7μs (32.7% faster)
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.
#------------------------------------------------
from typing import Generic, List, Tuple, TypeVar

# imports
import pytest
from spdl.pipeline._node import _gather_error

T = TypeVar("T")

# Minimal _Node and Task mock implementations for testing
class DummyTask:
    def __init__(self, name, exception=None, cancelled=False):
        self._name = name
        self._exception = exception
        self._cancelled = cancelled

    def get_name(self):
        return self._name

    def exception(self):
        return self._exception

    def cancelled(self):
        return self._cancelled

class _Node(Generic[T]):
    def __init__(self, task, upstream=None):
        self.task = task
        self.upstream = upstream if upstream is not None else []

# unit tests

# 1. Basic Test Cases

def test_no_error_in_single_node():
    # Single node, no error
    node = _Node(DummyTask("node1"))
    codeflash_output = _gather_error(node); result = codeflash_output # 1.29μs -> 1.18μs (9.51% faster)

def test_single_node_with_error():
    # Single node with error
    exc = ValueError("fail")
    node = _Node(DummyTask("node1", exception=exc))
    codeflash_output = _gather_error(node); result = codeflash_output # 1.58μs -> 1.23μs (29.2% faster)

def test_single_node_cancelled_with_error():
    # Single node, error present but cancelled
    exc = RuntimeError("cancelled error")
    node = _Node(DummyTask("node1", exception=exc, cancelled=True))
    codeflash_output = _gather_error(node); result = codeflash_output # 1.00μs -> 809ns (24.2% faster)

def test_two_nodes_one_error():
    # Two nodes, one with error
    exc = KeyError("upstream error")
    upstream = _Node(DummyTask("up", exception=exc))
    node = _Node(DummyTask("main"), upstream=[upstream])
    codeflash_output = _gather_error(node); result = codeflash_output # 2.25μs -> 1.76μs (27.6% faster)

def test_two_nodes_both_with_error():
    # Two nodes, both with errors
    exc1 = KeyError("upstream error")
    exc2 = ZeroDivisionError("main error")
    upstream = _Node(DummyTask("up", exception=exc1))
    node = _Node(DummyTask("main", exception=exc2), upstream=[upstream])
    codeflash_output = _gather_error(node); result = codeflash_output # 2.50μs -> 2.63μs (4.90% slower)

# 2. Edge Test Cases

def test_node_with_multiple_upstreams_errors():
    # Node with multiple upstreams, some with errors
    exc1 = KeyError("upstream1 error")
    exc2 = IndexError("upstream2 error")
    upstream1 = _Node(DummyTask("up1", exception=exc1))
    upstream2 = _Node(DummyTask("up2", exception=exc2))
    node = _Node(DummyTask("main"), upstream=[upstream1, upstream2])
    codeflash_output = _gather_error(node); result = codeflash_output # 3.01μs -> 2.90μs (3.73% faster)

def test_deeply_nested_upstream():
    # Deep nesting: main -> up1 -> up2 (error in up2)
    exc = OSError("deep error")
    up2 = _Node(DummyTask("up2", exception=exc))
    up1 = _Node(DummyTask("up1"), upstream=[up2])
    node = _Node(DummyTask("main"), upstream=[up1])
    codeflash_output = _gather_error(node); result = codeflash_output # 2.59μs -> 2.01μs (28.6% faster)

def test_upstream_with_cancelled_error():
    # Upstream error but cancelled, should not report
    exc = RuntimeError("cancelled error")
    upstream = _Node(DummyTask("up", exception=exc, cancelled=True))
    node = _Node(DummyTask("main"), upstream=[upstream])
    codeflash_output = _gather_error(node); result = codeflash_output # 1.79μs -> 1.50μs (19.1% faster)

def test_duplicate_task_names():
    # Nodes with duplicate names, both errors
    exc1 = ValueError("error1")
    exc2 = ValueError("error2")
    up1 = _Node(DummyTask("dup", exception=exc1))
    up2 = _Node(DummyTask("dup", exception=exc2))
    node = _Node(DummyTask("main"), upstream=[up1, up2])
    codeflash_output = _gather_error(node); result = codeflash_output # 3.06μs -> 2.91μs (4.98% faster)

def test_node_with_no_upstream_and_no_error():
    # Node with no upstream, no error
    node = _Node(DummyTask("solo"))
    codeflash_output = _gather_error(node); result = codeflash_output # 1.12μs -> 932ns (20.7% faster)

def test_node_with_upstream_none():
    # Node with upstream explicitly set to None
    node = _Node(DummyTask("solo"), upstream=None)
    codeflash_output = _gather_error(node); result = codeflash_output # 1.08μs -> 880ns (22.4% faster)

# 3. Large Scale Test Cases

def test_many_nodes_some_with_errors():
    # 100 nodes, every 10th node has an error
    nodes = []
    errors = []
    for i in range(100):
        name = f"node{i}"
        if i % 10 == 0:
            exc = Exception(f"error{i}")
            errors.append((name, exc))
            task = DummyTask(name, exception=exc)
        else:
            task = DummyTask(name)
        nodes.append(_Node(task))
    # Create a main node with all as upstream
    main = _Node(DummyTask("main"), upstream=nodes)
    codeflash_output = _gather_error(main); result = codeflash_output # 32.5μs -> 25.8μs (26.0% faster)
    # Should include all error nodes sorted by name
    expected = sorted(errors, key=lambda x: x[0])

def test_large_tree_of_nodes():
    # Tree of nodes: main -> 10 upstreams, each with 10 upstreams (total 111 nodes)
    errors = []
    level2_nodes = []
    for i in range(10):
        subnodes = []
        for j in range(10):
            name = f"leaf{i}_{j}"
            exc = None
            if (i + j) % 7 == 0:
                exc = Exception(f"error_{name}")
                errors.append((name, exc))
            subnodes.append(_Node(DummyTask(name, exception=exc)))
        level2_nodes.append(_Node(DummyTask(f"mid{i}"), upstream=subnodes))
    main = _Node(DummyTask("main"), upstream=level2_nodes)
    codeflash_output = _gather_error(main); result = codeflash_output # 37.9μs -> 30.4μs (24.6% faster)
    expected = sorted([e for e in errors], key=lambda x: x[0])

def test_performance_with_max_nodes():
    # 999 nodes, every node has error
    nodes = []
    errors = []
    for i in range(999):
        name = f"node{i}"
        exc = Exception(f"error{i}")
        errors.append((name, exc))
        nodes.append(_Node(DummyTask(name, exception=exc)))
    main = _Node(DummyTask("main"), upstream=nodes)
    codeflash_output = _gather_error(main); result = codeflash_output # 448μs -> 336μs (33.4% faster)
    expected = sorted(errors, key=lambda x: x[0])

def test_large_chain_of_nodes():
    # Chain of 100 nodes, each upstream of the next, error in every 25th node
    prev = None
    errors = []
    for i in reversed(range(100)):
        name = f"node{i}"
        exc = Exception(f"error{i}") if i % 25 == 0 else None
        if exc:
            errors.append((name, exc))
        node = _Node(DummyTask(name, exception=exc), upstream=[prev] if prev else [])
        prev = node
    codeflash_output = _gather_error(prev); result = codeflash_output # 39.0μs -> 36.4μs (7.02% faster)
    expected = sorted(errors, key=lambda x: x[0])
# 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-_gather_error-mgqntodt and push.