freeze the class

Author: KRRT7

codeflash/benchmarking/plugin/plugin.py

@@ -12,10 +12,8 @@
 from codeflash.benchmarking.codeflash_trace import codeflash_trace
 from codeflash.code_utils.code_utils import module_name_from_file_path
 
-# from codeflash.models.models import BenchmarkKey
 
-
-@dataclass
+@dataclass(frozen=True)
 class BenchmarkKey:
     module_path: str
     function_name: str
@@ -192,6 +190,7 @@ def get_benchmark_timings(trace_path: Path) -> dict[BenchmarkKey, int]:
 
                 # Create the benchmark key (file::function::line)
                 benchmark_key = BenchmarkKey(module_path=benchmark_file, function_name=benchmark_func)
+                print(f"XXX Processing benchmark: {benchmark_key}")
                 # Subtract overhead from total time
                 overhead = overhead_by_benchmark.get(benchmark_key, 0)
                 result[benchmark_key] = time_ns - overhead
@@ -232,26 +231,16 @@ def pytest_configure(config: pytest.Config) -> None:
 
     @staticmethod
     def pytest_collection_modifyitems(config: pytest.Config, items: list[pytest.Item]) -> None:
-        # Skip tests that don't have the benchmark fixture
+        # Skip tests that don't have the benchmark marker
         if not config.getoption("--codeflash-trace"):
             return
 
-        skip_no_benchmark = pytest.mark.skip(reason="Test requires benchmark fixture")
         for item in items:
-            # Check for direct benchmark fixture usage
-            has_fixture = hasattr(item, "fixturenames") and "benchmark" in item.fixturenames
-
             # Check for @pytest.mark.benchmark marker
-            has_marker = False
             if hasattr(item, "get_closest_marker"):
                 marker = item.get_closest_marker("benchmark")
-                if marker is not None:
-                    has_marker = True
-                    print("XXX FOUND THE BENCHMARK MARKER")
-
-            # Skip if neither fixture nor marker is present
-            if not (has_fixture or has_marker):
-                item.add_marker(skip_no_benchmark)
+                if marker is None:
+                    item.add_marker(pytest.mark.skip(reason="Test requires benchmark marker"))
 
     # Benchmark fixture
     class Benchmark:  # noqa: D106
@@ -307,11 +296,10 @@ def _run_benchmark(self, func, *args, **kwargs):  # noqa: ANN001, ANN002, ANN003
     @staticmethod
     @pytest.fixture
     def benchmark(request: pytest.FixtureRequest) -> object:
+        """Fixture to provide the benchmark functionality."""
         if not request.config.getoption("--codeflash-trace"):
-            print("XXX NOOOO PLS")
             return None
         print("XXX BENCHMARK PLUGIN INITIATED")
-
         return CodeFlashBenchmarkPlugin.Benchmark(request)
 
 

codeflash/benchmarking/pytest_new_process_trace_benchmarks.py

@@ -28,12 +28,12 @@
                 "-p",
                 "no:profiling",
                 "-s",
+                "-vv",
                 "-o",
                 "addopts=",
             ],
             plugins=[codeflash_benchmark_plugin],
         )  # Errors will be printed to stdout, not stderr
-
     except Exception as e:
         print(f"Failed to collect tests: {e!s}", file=sys.stderr)
         exitcode = -1

codeflash/optimization/optimizer.py

@@ -89,13 +89,13 @@ def run_benchmarks(
                     logger.info(
                         f"No valid benchmarks found in {self.args.benchmarks_root} for functions to optimize, continuing optimization"
                     )
-                else:
-                    function_benchmark_timings = CodeFlashBenchmarkPlugin.get_function_benchmark_timings(trace_file)
-                    total_benchmark_timings = CodeFlashBenchmarkPlugin.get_benchmark_timings(trace_file)
-                    function_to_results = validate_and_format_benchmark_table(
-                        function_benchmark_timings, total_benchmark_timings
-                    )
-                    print_benchmark_table(function_to_results)
+                    raise SystemExit  # noqa: TRY301
+                function_benchmark_timings = CodeFlashBenchmarkPlugin.get_function_benchmark_timings(trace_file)
+                total_benchmark_timings = CodeFlashBenchmarkPlugin.get_benchmark_timings(trace_file)
+                function_to_results = validate_and_format_benchmark_table(
+                    function_benchmark_timings, total_benchmark_timings
+                )
+                print_benchmark_table(function_to_results)
             except Exception as e:
                 logger.info(f"Error while tracing existing benchmarks: {e}")
                 logger.info("Information on existing benchmarks will not be available for this run.")