⚡️ Speed up method CodeFlashBenchmarkPlugin.write_benchmark_timings by 111% in PR #217 (proper-cleanup)

[codeflash-ai]

⚡️ This pull request contains optimizations for PR #217

If you approve this dependent PR, these changes will be merged into the original PR branch proper-cleanup.

This PR will be automatically closed if the original PR is merged.

---

📄 111% (1.11x) speedup for CodeFlashBenchmarkPlugin.write_benchmark_timings in codeflash/benchmarking/plugin/plugin.py

⏱️ Runtime : 26.9 milliseconds → 12.8 milliseconds (best of 121 runs)

📝 Explanation and details

Here's a rewritten, optimized version of your program, focusing on what the line profile indicates are bottlenecks.

• Reuse cursor: Opening a new cursor repeatedly is slow. Maintain a persistent cursor.
• Batching commits: Commit after many inserts if possible. However, since you clear the buffer after each write, one commit per call is necessary.
• Pragma optimizations: Set SQLite pragmas (synchronous = OFF, journal_mode = MEMORY) for faster inserts if durability isn't paramount.
• Avoid excessive object recreation: Only connect if needed, and clear but do not reallocate the benchmark list.
• Reduce exception handling cost: Trap and re-raise only actual DB exceptions.

Note: For highest speed, executemany and single-transaction-batch inserts are already optimal for SQLite. If even faster, use bulk insert with INSERT INTO ... VALUES (...), (...), ..., but this requires constructing SQL dynamically.

Here’s the optimized version.

Key points:

• self._ensure_connection() ensures both persistent connection and cursor.
• Pragmas are set only once for connection.
• Use self.benchmark_timings.clear() to avoid list reallocation.
• The cursor is reused for the lifetime of the object.

If your stability requirements are stricter (durability required), remove or tune the PRAGMA statements. If you want even higher throughput and can collect many queries per transaction, consider accepting a "bulk flush" mode to reduce commit frequency, but this requires API change.

This code preserves your public API and all comments, while running considerably faster especially on large inserts.

✅ Correctness verification report:

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

🌀 Generated Regression Tests Details

from __future__ import annotations

import os
import sqlite3
import tempfile

# imports
import pytest  # used for our unit tests
from codeflash.benchmarking.plugin.plugin import CodeFlashBenchmarkPlugin

# ---- Helper functions for setting up the test DB ----

def create_benchmark_db(path):
    """Create a SQLite database at path with the correct schema."""
    conn = sqlite3.connect(path)
    cur = conn.cursor()
    cur.execute("""
        CREATE TABLE benchmark_timings (
            benchmark_module_path TEXT,
            benchmark_function_name TEXT,
            benchmark_line_number INTEGER,
            benchmark_time_ns INTEGER
        )
    """)
    conn.commit()
    conn.close()

def fetch_benchmark_timings(path):
    """Fetch all rows from the benchmark_timings table."""
    conn = sqlite3.connect(path)
    cur = conn.cursor()
    cur.execute("SELECT benchmark_module_path, benchmark_function_name, benchmark_line_number, benchmark_time_ns FROM benchmark_timings")
    rows = cur.fetchall()
    conn.close()
    return rows

# ---- Unit Tests ----

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

def test_write_single_benchmark_timing(tmp_path):
    """Test writing a single benchmark timing entry."""
    db_path = tmp_path / "bench.db"
    create_benchmark_db(str(db_path))
    plugin = CodeFlashBenchmarkPlugin()
    plugin._trace_path = str(db_path)
    plugin.benchmark_timings = [("mod.py", "func", 42, 1000)]
    plugin.write_benchmark_timings()
    # Check DB contents
    rows = fetch_benchmark_timings(str(db_path))

def test_write_multiple_benchmark_timings(tmp_path):
    """Test writing multiple benchmark timing entries."""
    db_path = tmp_path / "bench.db"
    create_benchmark_db(str(db_path))
    plugin = CodeFlashBenchmarkPlugin()
    plugin._trace_path = str(db_path)
    plugin.benchmark_timings = [
        ("mod1.py", "func1", 10, 100),
        ("mod2.py", "func2", 20, 200),
        ("mod3.py", "func3", 30, 300),
    ]
    plugin.write_benchmark_timings()
    rows = fetch_benchmark_timings(str(db_path))

def test_write_benchmark_timings_twice(tmp_path):
    """Test writing benchmark timings in two separate calls."""
    db_path = tmp_path / "bench.db"
    create_benchmark_db(str(db_path))
    plugin = CodeFlashBenchmarkPlugin()
    plugin._trace_path = str(db_path)
    # First call
    plugin.benchmark_timings = [("mod.py", "func", 1, 111)]
    plugin.write_benchmark_timings()
    # Second call
    plugin.benchmark_timings = [("mod2.py", "func2", 2, 222)]
    plugin.write_benchmark_timings()
    rows = fetch_benchmark_timings(str(db_path))

def test_no_benchmark_timings_no_write(tmp_path):
    """Test that nothing is written if benchmark_timings is empty."""
    db_path = tmp_path / "bench.db"
    create_benchmark_db(str(db_path))
    plugin = CodeFlashBenchmarkPlugin()
    plugin._trace_path = str(db_path)
    plugin.benchmark_timings = []
    plugin.write_benchmark_timings()
    rows = fetch_benchmark_timings(str(db_path))

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

def test_write_benchmark_with_null_like_values(tmp_path):
    """Test writing entries with empty strings and zero/negative numbers."""
    db_path = tmp_path / "bench.db"
    create_benchmark_db(str(db_path))
    plugin = CodeFlashBenchmarkPlugin()
    plugin._trace_path = str(db_path)
    plugin.benchmark_timings = [
        ("", "", 0, 0),
        ("mod.py", "func", -1, -100),
    ]
    plugin.write_benchmark_timings()
    rows = fetch_benchmark_timings(str(db_path))

def test_write_benchmark_timings_with_long_strings(tmp_path):
    """Test writing entries with very long strings."""
    db_path = tmp_path / "bench.db"
    create_benchmark_db(str(db_path))
    plugin = CodeFlashBenchmarkPlugin()
    plugin._trace_path = str(db_path)
    long_str = "x" * 1000
    plugin.benchmark_timings = [
        (long_str, long_str, 123, 456789),
    ]
    plugin.write_benchmark_timings()
    rows = fetch_benchmark_timings(str(db_path))

def test_write_benchmark_timings_invalid_db_path(tmp_path):
    """Test error handling when the DB path is invalid."""
    plugin = CodeFlashBenchmarkPlugin()
    # Set to a path in a non-existent directory
    plugin._trace_path = str(tmp_path / "nonexistent" / "bench.db")
    plugin.benchmark_timings = [("mod.py", "func", 1, 2)]
    with pytest.raises(sqlite3.OperationalError):
        plugin.write_benchmark_timings()

def test_write_benchmark_timings_with_existing_connection(tmp_path):
    """Test that an existing _connection is used and not replaced."""
    db_path = tmp_path / "bench.db"
    create_benchmark_db(str(db_path))
    plugin = CodeFlashBenchmarkPlugin()
    plugin._trace_path = str(db_path)
    plugin._connection = sqlite3.connect(str(db_path))
    plugin.benchmark_timings = [("mod.py", "func", 1, 2)]
    plugin.write_benchmark_timings()
    # Check that the connection is still open (should not be closed)
    try:
        plugin._connection.execute("SELECT 1")
    except sqlite3.ProgrammingError:
        pass
    rows = fetch_benchmark_timings(str(db_path))

def test_write_benchmark_timings_rollback_on_failure(tmp_path):
    """Test that rollback is called and no data is written if insert fails."""
    db_path = tmp_path / "bench.db"
    # Create a DB with the wrong schema (missing a column)
    conn = sqlite3.connect(str(db_path))
    cur = conn.cursor()
    cur.execute("""
        CREATE TABLE benchmark_timings (
            benchmark_module_path TEXT,
            benchmark_function_name TEXT
        )
    """)
    conn.commit()
    conn.close()
    plugin = CodeFlashBenchmarkPlugin()
    plugin._trace_path = str(db_path)
    plugin.benchmark_timings = [("mod.py", "func", 1, 2)]
    with pytest.raises(sqlite3.OperationalError):
        plugin.write_benchmark_timings()
    # Table should still be empty
    conn = sqlite3.connect(str(db_path))
    cur = conn.cursor()
    cur.execute("SELECT * FROM benchmark_timings")
    rows = cur.fetchall()
    conn.close()

# ----------- LARGE SCALE TEST CASES -----------

def test_write_large_number_of_benchmark_timings(tmp_path):
    """Test writing a large number (1000) of benchmark timings."""
    db_path = tmp_path / "bench.db"
    create_benchmark_db(str(db_path))
    plugin = CodeFlashBenchmarkPlugin()
    plugin._trace_path = str(db_path)
    N = 1000
    data = [(f"mod{i}.py", f"func{i}", i, i*100) for i in range(N)]
    plugin.benchmark_timings = list(data)
    plugin.write_benchmark_timings()
    rows = fetch_benchmark_timings(str(db_path))

def test_write_large_benchmark_timings_multiple_batches(tmp_path):
    """Test multiple large writes to the DB."""
    db_path = tmp_path / "bench.db"
    create_benchmark_db(str(db_path))
    plugin = CodeFlashBenchmarkPlugin()
    plugin._trace_path = str(db_path)
    N = 500
    # First batch
    data1 = [(f"modA{i}.py", f"funcA{i}", i, i*10) for i in range(N)]
    plugin.benchmark_timings = list(data1)
    plugin.write_benchmark_timings()
    # Second batch
    data2 = [(f"modB{i}.py", f"funcB{i}", i, i*20) for i in range(N)]
    plugin.benchmark_timings = list(data2)
    plugin.write_benchmark_timings()
    rows = fetch_benchmark_timings(str(db_path))

def test_write_benchmark_timings_performance(tmp_path):
    """Test that writing 1000 entries does not take excessive time."""
    import time
    db_path = tmp_path / "bench.db"
    create_benchmark_db(str(db_path))
    plugin = CodeFlashBenchmarkPlugin()
    plugin._trace_path = str(db_path)
    N = 1000
    plugin.benchmark_timings = [(f"mod{i}.py", f"func{i}", i, i*123) for i in range(N)]
    start = time.time()
    plugin.write_benchmark_timings()
    duration = time.time() - start
    rows = fetch_benchmark_timings(str(db_path))
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

from __future__ import annotations

import os
import sqlite3
import tempfile

# imports
import pytest  # used for our unit tests
from codeflash.benchmarking.plugin.plugin import CodeFlashBenchmarkPlugin


# Helper function to create a temporary SQLite DB with the required table
def create_temp_db_with_benchmark_table():
    temp_db = tempfile.NamedTemporaryFile(delete=False)
    conn = sqlite3.connect(temp_db.name)
    cur = conn.cursor()
    # Create the table with the expected schema
    cur.execute(
        """
        CREATE TABLE benchmark_timings (
            benchmark_module_path TEXT,
            benchmark_function_name TEXT,
            benchmark_line_number INTEGER,
            benchmark_time_ns INTEGER
        )
        """
    )
    conn.commit()
    conn.close()
    return temp_db

# Helper function to read all rows from the benchmark_timings table
def read_benchmark_timings_from_db(db_path):
    conn = sqlite3.connect(db_path)
    cur = conn.cursor()
    cur.execute("SELECT benchmark_module_path, benchmark_function_name, benchmark_line_number, benchmark_time_ns FROM benchmark_timings")
    rows = cur.fetchall()
    conn.close()
    return rows

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

def test_write_single_benchmark_timing_basic():
    """Test writing a single benchmark timing entry to the database."""
    temp_db = create_temp_db_with_benchmark_table()
    plugin = CodeFlashBenchmarkPlugin()
    plugin._trace_path = temp_db.name
    plugin.benchmark_timings = [
        ("module.py", "func", 42, 123456789)
    ]
    plugin.write_benchmark_timings()
    # Check the DB contents
    rows = read_benchmark_timings_from_db(temp_db.name)
    os.unlink(temp_db.name)

def test_write_multiple_benchmark_timings_basic():
    """Test writing multiple benchmark timings in one call."""
    temp_db = create_temp_db_with_benchmark_table()
    plugin = CodeFlashBenchmarkPlugin()
    plugin._trace_path = temp_db.name
    plugin.benchmark_timings = [
        ("a.py", "f1", 1, 100),
        ("b.py", "f2", 2, 200),
        ("c.py", "f3", 3, 300)
    ]
    plugin.write_benchmark_timings()
    rows = read_benchmark_timings_from_db(temp_db.name)
    os.unlink(temp_db.name)

def test_write_benchmark_timings_twice_accumulates():
    """Test that calling write_benchmark_timings twice appends new rows."""
    temp_db = create_temp_db_with_benchmark_table()
    plugin = CodeFlashBenchmarkPlugin()
    plugin._trace_path = temp_db.name

    # First write
    plugin.benchmark_timings = [("x.py", "foo", 10, 999)]
    plugin.write_benchmark_timings()
    # Second write
    plugin.benchmark_timings = [("y.py", "bar", 20, 888)]
    plugin.write_benchmark_timings()

    rows = read_benchmark_timings_from_db(temp_db.name)
    os.unlink(temp_db.name)

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

def test_write_empty_benchmark_timings_list():
    """Test that nothing is written if benchmark_timings is empty."""
    temp_db = create_temp_db_with_benchmark_table()
    plugin = CodeFlashBenchmarkPlugin()
    plugin._trace_path = temp_db.name
    plugin.benchmark_timings = []
    plugin.write_benchmark_timings()
    rows = read_benchmark_timings_from_db(temp_db.name)
    os.unlink(temp_db.name)

def test_write_benchmark_timings_none_connection():
    """Test that the function creates a connection if _connection is None."""
    temp_db = create_temp_db_with_benchmark_table()
    plugin = CodeFlashBenchmarkPlugin()
    plugin._trace_path = temp_db.name
    plugin._connection = None
    plugin.benchmark_timings = [("mod.py", "func", 5, 1000)]
    plugin.write_benchmark_timings()
    rows = read_benchmark_timings_from_db(temp_db.name)
    os.unlink(temp_db.name)

def test_write_benchmark_timings_existing_connection():
    """Test that the function uses an existing connection if present."""
    temp_db = create_temp_db_with_benchmark_table()
    plugin = CodeFlashBenchmarkPlugin()
    plugin._trace_path = temp_db.name
    plugin._connection = sqlite3.connect(temp_db.name)
    plugin.benchmark_timings = [("mod2.py", "func2", 6, 2000)]
    plugin.write_benchmark_timings()
    rows = read_benchmark_timings_from_db(temp_db.name)
    plugin._connection.close()
    os.unlink(temp_db.name)

def test_write_benchmark_timings_unicode_and_long_strings():
    """Test writing entries with Unicode and very long strings."""
    temp_db = create_temp_db_with_benchmark_table()
    plugin = CodeFlashBenchmarkPlugin()
    plugin._trace_path = temp_db.name
    long_str = "a" * 500
    unicode_str = "模块.py"
    plugin.benchmark_timings = [
        (long_str, "f", 1, 1),
        (unicode_str, "函数", 2, 2)
    ]
    plugin.write_benchmark_timings()
    rows = read_benchmark_timings_from_db(temp_db.name)
    os.unlink(temp_db.name)

def test_write_benchmark_timings_negative_and_zero_values():
    """Test writing entries with zero and negative numbers."""
    temp_db = create_temp_db_with_benchmark_table()
    plugin = CodeFlashBenchmarkPlugin()
    plugin._trace_path = temp_db.name
    plugin.benchmark_timings = [
        ("mod.py", "zero", 0, 0),
        ("mod.py", "neg", -1, -100)
    ]
    plugin.write_benchmark_timings()
    rows = read_benchmark_timings_from_db(temp_db.name)
    os.unlink(temp_db.name)

def test_write_benchmark_timings_rollback_on_failure():
    """Test that if an error occurs, no partial data is committed."""
    temp_db = create_temp_db_with_benchmark_table()
    plugin = CodeFlashBenchmarkPlugin()
    plugin._trace_path = temp_db.name
    # Insert a valid row first
    plugin.benchmark_timings = [("ok.py", "ok", 1, 1)]
    plugin.write_benchmark_timings()
    # Now, insert an invalid row (too few columns)
    plugin.benchmark_timings = [("bad.py", "bad", 2)]  # Missing one value
    with pytest.raises(sqlite3.ProgrammingError):
        plugin.write_benchmark_timings()
    # The DB should still only have the first row
    rows = read_benchmark_timings_from_db(temp_db.name)
    os.unlink(temp_db.name)

def test_write_benchmark_timings_no_table():
    """Test that writing to a DB without the table raises an error."""
    temp_db = tempfile.NamedTemporaryFile(delete=False)
    plugin = CodeFlashBenchmarkPlugin()
    plugin._trace_path = temp_db.name
    plugin.benchmark_timings = [("no_table.py", "func", 1, 1)]
    with pytest.raises(sqlite3.OperationalError):
        plugin.write_benchmark_timings()
    os.unlink(temp_db.name)

def test_write_benchmark_timings_non_string_module_and_func():
    """Test that non-string types for module/function raise an error."""
    temp_db = create_temp_db_with_benchmark_table()
    plugin = CodeFlashBenchmarkPlugin()
    plugin._trace_path = temp_db.name
    plugin.benchmark_timings = [
        (123, 456, 7, 8)
    ]
    # SQLite will coerce types, so this should succeed, but let's check the DB
    plugin.write_benchmark_timings()
    rows = read_benchmark_timings_from_db(temp_db.name)
    os.unlink(temp_db.name)

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

def test_write_benchmark_timings_large_scale():
    """Test writing a large number of benchmark timings (scalability)."""
    temp_db = create_temp_db_with_benchmark_table()
    plugin = CodeFlashBenchmarkPlugin()
    plugin._trace_path = temp_db.name
    N = 1000  # Upper limit as per instructions
    plugin.benchmark_timings = [
        (f"mod{i}.py", f"func{i}", i, i * 100) for i in range(N)
    ]
    plugin.write_benchmark_timings()
    rows = read_benchmark_timings_from_db(temp_db.name)
    os.unlink(temp_db.name)

def test_write_benchmark_timings_large_scale_multiple_batches():
    """Test writing large batches in multiple calls to ensure accumulation."""
    temp_db = create_temp_db_with_benchmark_table()
    plugin = CodeFlashBenchmarkPlugin()
    plugin._trace_path = temp_db.name
    N = 500
    # First batch
    plugin.benchmark_timings = [
        (f"A{i}.py", f"F{i}", i, i) for i in range(N)
    ]
    plugin.write_benchmark_timings()
    # Second batch
    plugin.benchmark_timings = [
        (f"B{i}.py", f"G{i}", i, i*2) for i in range(N)
    ]
    plugin.write_benchmark_timings()
    rows = read_benchmark_timings_from_db(temp_db.name)
    os.unlink(temp_db.name)
# 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-pr217-2025-05-19T03.55.14 and push.

[misrasaurabh1]

i like this actually...
optimizations here actually matter, @KRRT7 can you review this?