⚡️ 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.