⚡️ Speed up method CodeflashTrace.write_function_timings by 12,039% in PR #59 (codeflash-trace-decorator)

[codeflash-ai]

⚡️ This pull request contains optimizations for PR #59

If you approve this dependent PR, these changes will be merged into the original PR branch codeflash-trace-decorator.

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

---

📄 12,039% (120.39x) speedup for CodeflashTrace.write_function_timings in codeflash/benchmarking/codeflash_trace.py

⏱️ Runtime : 998 microseconds → 8.22 microseconds (best of 8 runs)

📝 Explanation and details

Here is the optimized version of the CodeflashTrace class focusing on performance improvements, particularly within the write_function_timings function.

• Reuse the same cursor for multiple insertions to minimize the overhead of repeatedly creating cursors.
• Instead of accumulating entries and writing to the database in large chunks, write entries to the database more frequently to prevent large data handling and reduce memory usage.
• Batch the arguments and keyword arguments pickling process.

Explanation.

1. Primary optimization related to Database handling.

   • Connection Initialization: The database connection is initialized in the constructor if trace_path is provided, eliminating the need to reinitialize it each time in the decorator method.
   • Cursor Reuse: The cursor is created once during initialization and reused.
   • Batch Control: Instead of waiting for a very large list to accumulate, intermediate batches (threshold set at 100) are written to minimize memory usage and eliminate any potential latency due to large insertions.

2. Pickling.

   • Batch Pickling: The arguments and keyword arguments are pickled immediately or on-call basis, minimizing the pickling overhead time.
   • Error Handling: Improved error handling within _pickle_args_kwargs function.

3. Code Organization.

   • Helper functions (_initialize_db_connection, _pickle_args_kwargs, _write_batch_and_clear) improve readability.

By adopting these optimizations, the code's performance, especially for database write operations and argument serialization, should be significantly improved.

✅ Correctness verification report:

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

🌀 Generated Regression Tests Details

import functools
import os
import pickle
import sqlite3
import sys
import time
from typing import Callable

# function to test
import dill
# imports
import pytest
from codeflash.benchmarking.codeflash_trace import CodeflashTrace

# Create a singleton instance
codeflash_trace = CodeflashTrace()

# unit tests
@pytest.fixture
def setup_database():
    # Setup an in-memory SQLite database for testing
    codeflash_trace._trace_path = ":memory:"
    connection = sqlite3.connect(codeflash_trace._trace_path)
    cursor = connection.cursor()
    cursor.execute("""
        CREATE TABLE benchmark_function_timings (
            function_name TEXT,
            class_name TEXT,
            module_name TEXT,
            file_path TEXT,
            benchmark_function_name TEXT,
            benchmark_module_path TEXT,
            benchmark_line_number TEXT,
            function_time_ns INTEGER,
            overhead_time_ns INTEGER,
            args BLOB,
            kwargs BLOB
        )
    """)
    connection.commit()
    codeflash_trace._connection = connection
    yield
    connection.close()

# Basic Functionality
def test_single_function_call(setup_database):
    @codeflash_trace
    def add(a, b):
        return a + b

    add(1, 2)
    codeflash_trace.write_function_timings()
    cursor = codeflash_trace._connection.cursor()
    cursor.execute("SELECT * FROM benchmark_function_timings")
    rows = cursor.fetchall()

def test_multiple_function_calls(setup_database):
    @codeflash_trace
    def add(a, b):
        return a + b

    @codeflash_trace
    def multiply(a, b):
        return a * b

    add(1, 2)
    multiply(2, 5)
    codeflash_trace.write_function_timings()
    cursor = codeflash_trace._connection.cursor()
    cursor.execute("SELECT * FROM benchmark_function_timings")
    rows = cursor.fetchall()

# Edge Cases
def test_no_function_calls(setup_database):
    codeflash_trace.write_function_timings()
    cursor = codeflash_trace._connection.cursor()
    cursor.execute("SELECT * FROM benchmark_function_timings")
    rows = cursor.fetchall()

def test_empty_arguments(setup_database):
    @codeflash_trace
    def no_args():
        return "no args"

    @codeflash_trace
    def empty_list(lst=[]):
        return len(lst)

    no_args()
    empty_list()
    codeflash_trace.write_function_timings()
    cursor = codeflash_trace._connection.cursor()
    cursor.execute("SELECT * FROM benchmark_function_timings")
    rows = cursor.fetchall()












def test_benchmarking_disabled(setup_database):
    os.environ["CODEFLASH_BENCHMARKING"] = "False"
    @codeflash_trace
    def add(a, b):
        return a + b

    add(1, 2)
    codeflash_trace.write_function_timings()
    cursor = codeflash_trace._connection.cursor()
    cursor.execute("SELECT * FROM benchmark_function_timings")
    rows = cursor.fetchall()

def test_missing_environment_variables(setup_database):
    if "CODEFLASH_BENCHMARK_FUNCTION_NAME" in os.environ:
        del os.environ["CODEFLASH_BENCHMARK_FUNCTION_NAME"]
    if "CODEFLASH_BENCHMARK_MODULE_PATH" in os.environ:
        del os.environ["CODEFLASH_BENCHMARK_MODULE_PATH"]
    if "CODEFLASH_BENCHMARK_LINE_NUMBER" in os.environ:
        del os.environ["CODEFLASH_BENCHMARK_LINE_NUMBER"]
    @codeflash_trace
    def add(a, b):
        return a + b

    add(1, 2)
    codeflash_trace.write_function_timings()
    cursor = codeflash_trace._connection.cursor()
    cursor.execute("SELECT * FROM benchmark_function_timings")
    rows = cursor.fetchall()

# Overhead Calculation
def test_overhead_calculation(setup_database):
    @codeflash_trace
    def overhead_test(a, b):
        return a + b

    overhead_test(1, 2)
    codeflash_trace.write_function_timings()
    cursor = codeflash_trace._connection.cursor()
    cursor.execute("SELECT * FROM benchmark_function_timings")
    rows = cursor.fetchall()

# Performance and Scalability
def test_large_input_data(setup_database):
    @codeflash_trace
    def large_input_data(lst):
        return sum(lst)

    large_input_data(list(range(1000)))
    codeflash_trace.write_function_timings()
    cursor = codeflash_trace._connection.cursor()
    cursor.execute("SELECT * FROM benchmark_function_timings")
    rows = cursor.fetchall()

def test_high_frequency_calls(setup_database):
    @codeflash_trace
    def high_frequency(a):
        return a * a

    for i in range(1000):
        high_frequency(i)
    codeflash_trace.write_function_timings()
    cursor = codeflash_trace._connection.cursor()
    cursor.execute("SELECT COUNT(*) FROM benchmark_function_timings")
    count = cursor.fetchone()[0]
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

import functools
import os
import pickle
import sqlite3
import sys
import time
from typing import Callable

# function to test
import dill
# imports
import pytest  # used for our unit tests
from codeflash.benchmarking.codeflash_trace import CodeflashTrace

# Create a singleton instance
codeflash_trace = CodeflashTrace()

# unit tests

@pytest.fixture
def setup_database(tmp_path):
    """Fixture to set up and tear down a temporary database."""
    db_path = tmp_path / "test_db.sqlite"
    codeflash_trace._trace_path = str(db_path)
    conn = sqlite3.connect(db_path)
    conn.execute("""
    CREATE TABLE benchmark_function_timings (
        function_name TEXT, class_name TEXT, module_name TEXT, file_path TEXT,
        benchmark_function_name TEXT, benchmark_module_path TEXT, benchmark_line_number TEXT,
        function_time_ns INTEGER, overhead_time_ns INTEGER, args BLOB, kwargs BLOB
    )""")
    conn.commit()
    yield conn
    conn.close()

def test_single_function_call(setup_database):
    """Test a single function call without arguments."""
    @codeflash_trace
    def simple_function():
        pass

    simple_function()
    codeflash_trace.write_function_timings()

    cur = setup_database.cursor()
    cur.execute("SELECT * FROM benchmark_function_timings")
    rows = cur.fetchall()

def test_function_with_arguments(setup_database):
    """Test a function call with basic arguments."""
    @codeflash_trace
    def add(a, b):
        return a + b

    add(1, 2)
    codeflash_trace.write_function_timings()

    cur = setup_database.cursor()
    cur.execute("SELECT * FROM benchmark_function_timings")
    rows = cur.fetchall()

def test_multiple_function_calls(setup_database):
    """Test multiple function calls."""
    @codeflash_trace
    def simple_function():
        pass

    simple_function()
    simple_function()
    codeflash_trace.write_function_timings()

    cur = setup_database.cursor()
    cur.execute("SELECT * FROM benchmark_function_timings")
    rows = cur.fetchall()

def test_function_with_complex_arguments(setup_database):
    """Test a function call with complex arguments."""
    @codeflash_trace
    def process_list(data):
        return [x * 2 for x in data]

    process_list([1, 2, 3])
    codeflash_trace.write_function_timings()

    cur = setup_database.cursor()
    cur.execute("SELECT * FROM benchmark_function_timings")
    rows = cur.fetchall()

def test_function_with_none_argument(setup_database):
    """Test a function call with None as an argument."""
    @codeflash_trace
    def handle_none(value):
        return value is None

    handle_none(None)
    codeflash_trace.write_function_timings()

    cur = setup_database.cursor()
    cur.execute("SELECT * FROM benchmark_function_timings")
    rows = cur.fetchall()

def test_function_with_empty_list(setup_database):
    """Test a function call with an empty list as an argument."""
    @codeflash_trace
    def process_empty_list(data):
        return len(data) == 0

    process_empty_list([])
    codeflash_trace.write_function_timings()

    cur = setup_database.cursor()
    cur.execute("SELECT * FROM benchmark_function_timings")
    rows = cur.fetchall()



def test_benchmarking_environment(setup_database):
    """Test function call in benchmarking environment."""
    os.environ["CODEFLASH_BENCHMARKING"] = "True"
    os.environ["CODEFLASH_BENCHMARK_FUNCTION_NAME"] = "benchmark_func"
    os.environ["CODEFLASH_BENCHMARK_MODULE_PATH"] = "benchmark_module"
    os.environ["CODEFLASH_BENCHMARK_LINE_NUMBER"] = "42"

    @codeflash_trace
    def simple_function():
        pass

    simple_function()
    codeflash_trace.write_function_timings()

    cur = setup_database.cursor()
    cur.execute("SELECT * FROM benchmark_function_timings")
    rows = cur.fetchall()

To edit these changes git checkout codeflash/optimize-pr59-2025-04-03T23.05.44 and push.