⚡️ Speed up method CodeFlashBenchmarkPlugin.write_benchmark_timings by 167% in PR #59 (codeflash-trace-decorator)

Here is the optimized version of the given Python program.



Changes made.
1. Added PRAGMA statements (`PRAGMA synchronous = OFF` and `PRAGMA journal_mode = MEMORY`) after establishing the database connection. These settings help in improving the write performance by reducing the disk I/O operations.
  
However, note that setting `PRAGMA synchronous = OFF` and `PRAGMA journal_mode = MEMORY` makes the database more vulnerable to corruption in the event of an application crash or power failure. Use these optimizations only when performance is critical, and the database is not mission-critical.

Author: codeflash-ai[bot]

codeflash/benchmarking/plugin/plugin.py

@@ -1,10 +1,13 @@
 from __future__ import annotations
+
 import os
 import sqlite3
 import sys
 import time
 from pathlib import Path
+
 import pytest
+
 from codeflash.benchmarking.codeflash_trace import codeflash_trace
 from codeflash.models.models import BenchmarkKey
 
@@ -15,7 +18,7 @@ def __init__(self) -> None:
         self._connection = None
         self.benchmark_timings = []
 
-    def setup(self, trace_path:str) -> None:
+    def setup(self, trace_path: str) -> None:
         try:
             # Open connection
             self._trace_path = trace_path
@@ -28,7 +31,7 @@ def setup(self, trace_path:str) -> None:
                 "benchmark_time_ns INTEGER)"
             )
             self._connection.commit()
-            self.close() # Reopen only at the end of pytest session
+            self.close()  # Reopen only at the end of pytest session
         except Exception as e:
             print(f"Database setup error: {e}")
             if self._connection:
@@ -42,20 +45,23 @@ def write_benchmark_timings(self) -> None:
 
         if self._connection is None:
             self._connection = sqlite3.connect(self._trace_path)
+            self._connection.execute("PRAGMA synchronous = OFF")
+            self._connection.execute("PRAGMA journal_mode = MEMORY")
 
         try:
             cur = self._connection.cursor()
             # Insert data into the benchmark_timings table
             cur.executemany(
                 "INSERT INTO benchmark_timings (benchmark_file_path, benchmark_function_name, benchmark_line_number, benchmark_time_ns) VALUES (?, ?, ?, ?)",
-                self.benchmark_timings
+                self.benchmark_timings,
             )
             self._connection.commit()
-            self.benchmark_timings = [] # Clear the benchmark timings list
+            self.benchmark_timings = []  # Clear the benchmark timings list
         except Exception as e:
             print(f"Error writing to benchmark timings database: {e}")
             self._connection.rollback()
             raise
+
     def close(self) -> None:
         if self._connection:
             self._connection.close()
@@ -189,12 +195,7 @@ def pytest_sessionfinish(self, session, exitstatus):
 
     @staticmethod
     def pytest_addoption(parser):
-        parser.addoption(
-            "--codeflash-trace",
-            action="store_true",
-            default=False,
-            help="Enable CodeFlash tracing"
-        )
+        parser.addoption("--codeflash-trace", action="store_true", default=False, help="Enable CodeFlash tracing")
 
     @staticmethod
     def pytest_plugin_registered(plugin, manager):
@@ -246,7 +247,7 @@ def test_something(benchmark):
             os.environ["CODEFLASH_BENCHMARK_LINE_NUMBER"] = str(line_number)
             os.environ["CODEFLASH_BENCHMARKING"] = "True"
 
-        # Run the function
+            # Run the function
             start = time.perf_counter_ns()
             result = func(*args, **kwargs)
             end = time.perf_counter_ns()
@@ -260,7 +261,8 @@ def test_something(benchmark):
             codeflash_trace.function_call_count = 0
             # Add to the benchmark timings buffer
             codeflash_benchmark_plugin.benchmark_timings.append(
-                (benchmark_file_path, benchmark_function_name, line_number, end - start))
+                (benchmark_file_path, benchmark_function_name, line_number, end - start)
+            )
 
             return result
 
@@ -272,4 +274,5 @@ def benchmark(request):
 
         return CodeFlashBenchmarkPlugin.Benchmark(request)
 
-codeflash_benchmark_plugin = CodeFlashBenchmarkPlugin()
+
+codeflash_benchmark_plugin = CodeFlashBenchmarkPlugin()