preview

Author: KRRT7

codeflash/benchmarking/codeflash_trace.py

@@ -102,19 +102,17 @@ def __call__(self, func: Callable) -> Callable:
 
         @functools.wraps(func)
         def wrapper(*args: tuple, **kwargs: dict) -> object:
-            # Initialize thread-local active functions set if it doesn't exist
             if not hasattr(self._thread_local, "active_functions"):
                 self._thread_local.active_functions = set()
             # If it's in a recursive function, just return the result
             if func_id in self._thread_local.active_functions:
                 return func(*args, **kwargs)
             # Track active functions so we can detect recursive functions
             self._thread_local.active_functions.add(func_id)
-            # Measure execution time
             start_time = time.thread_time_ns()
             result = func(*args, **kwargs)
             end_time = time.thread_time_ns()
-            # Calculate execution time
+            logger.info(f"CodeflashTrace: Function {func.__name__} executed in {end_time - start_time} ns")
             execution_time = end_time - start_time
             self.function_call_count += 1
 

codeflash/benchmarking/plugin/plugin.py

@@ -78,7 +78,7 @@ def get_function_benchmark_timings(trace_path: Path) -> dict[str, dict[Benchmark
             A nested dictionary where:
             - Outer keys are module_name.qualified_name (module.class.function)
             - Inner keys are of type BenchmarkKey
-            - Values are function timing in milliseconds
+            - Values are function timing in nanoseconds
 
         """
         result = {}
@@ -120,16 +120,16 @@ def get_function_benchmark_timings(trace_path: Path) -> dict[str, dict[Benchmark
         return result
 
     @staticmethod
-    def get_benchmark_timings(trace_path: Path) -> dict[BenchmarkKey, int]:
+    def get_benchmark_timings(trace_path: Path) -> dict[tuple[str, str, int], int]:
         """Extract total benchmark timings from trace files.
 
         Args:
             trace_path: Path to the trace file
 
         Returns:
             A dictionary mapping where:
-            - Keys are of type BenchmarkKey
-            - Values are total benchmark timing in milliseconds (with overhead subtracted)
+            - Keys are (module_path, function_name, line_number)
+            - Values are total benchmark timing in nanoseconds (with overhead subtracted)
 
         """
         result = {}
@@ -148,8 +148,8 @@ def get_benchmark_timings(trace_path: Path) -> dict[BenchmarkKey, int]:
 
             for row in cursor.fetchall():
                 benchmark_file, benchmark_func, benchmark_line, total_overhead_ns = row
-                benchmark_key = BenchmarkKey(module_path=benchmark_file, function_name=benchmark_func)
-                overhead_by_benchmark[benchmark_key] = total_overhead_ns or 0  # Handle NULL sum case
+                key = (benchmark_file, benchmark_func, benchmark_line)
+                overhead_by_benchmark[key] = total_overhead_ns or 0  # Handle NULL sum case
 
             # Query the benchmark_timings table for total times
             cursor.execute(
@@ -159,13 +159,9 @@ def get_benchmark_timings(trace_path: Path) -> dict[BenchmarkKey, int]:
 
             for row in cursor.fetchall():
                 benchmark_file, benchmark_func, benchmark_line, time_ns = row
-
-                benchmark_key = BenchmarkKey(
-                    module_path=benchmark_file, function_name=benchmark_func
-                )  # (file::function::line)
-                # Subtract overhead from total time
-                overhead = overhead_by_benchmark.get(benchmark_key, 0)
-                result[benchmark_key] = time_ns - overhead
+                key = (benchmark_file, benchmark_func, benchmark_line)
+                overhead = overhead_by_benchmark.get(key, 0)
+                result[key] = time_ns - overhead
 
         finally:
             connection.close()
@@ -245,7 +241,6 @@ def _run_benchmark(
             else:
                 call_identifier = f"{benchmark_function_name}::call_{self._call_count}"
 
-            os.environ["CODEFLASH_BENCHMARKING"] = "True"
             os.environ["CODEFLASH_BENCHMARK_FUNCTION_NAME"] = call_identifier
             os.environ["CODEFLASH_BENCHMARK_MODULE_PATH"] = benchmark_module_path
             os.environ["CODEFLASH_BENCHMARK_LINE_NUMBER"] = str(line_number)