add the plugin as a pytest entry point

Author: KRRT7

codeflash/benchmarking/plugin/plugin.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import importlib.util
 import os
 import sqlite3
 import sys
@@ -12,6 +13,8 @@
 from codeflash.benchmarking.codeflash_trace import codeflash_trace
 from codeflash.code_utils.code_utils import module_name_from_file_path
 
+IS_PYTEST_BENCHMARK_INSTALLED = importlib.util.find_spec("pytest_benchmark") is not None
+
 
 @dataclass(frozen=True)
 class BenchmarkKey:
@@ -212,19 +215,18 @@ def pytest_sessionfinish(self, session, exitstatus) -> None:  # noqa: ANN001, AR
     def pytest_addoption(parser: pytest.Parser) -> None:
         parser.addoption("--codeflash-trace", action="store_true", default=False, help="Enable CodeFlash tracing")
 
-    @staticmethod
-    def pytest_plugin_registered(plugin, manager) -> None:  # noqa: ANN001
-        # Not necessary since run with -p no:benchmark, but just in case
-        if hasattr(plugin, "name") and plugin.name == "pytest-benchmark":
-            manager.unregister(plugin)
-
     @staticmethod
     def pytest_configure(config: pytest.Config) -> None:
-        """Register the benchmark marker."""
+        """Register the benchmark marker and disable conflicting plugins."""
         config.addinivalue_line(
             "markers", "benchmark: mark test as a benchmark that should be run with codeflash tracing"
         )
 
+        if config.getoption("--codeflash-trace") and IS_PYTEST_BENCHMARK_INSTALLED:
+            object.__setattr__(config.option, "benchmark_disable", True)
+            config.pluginmanager.set_blocked("pytest_benchmark")
+            config.pluginmanager.set_blocked("pytest-benchmark")
+
     @staticmethod
     def pytest_collection_modifyitems(config: pytest.Config, items: list[pytest.Item]) -> None:
         # Skip tests that don't have the benchmark marker
@@ -297,3 +299,20 @@ def benchmark(request: pytest.FixtureRequest) -> object:
 
 
 codeflash_benchmark_plugin = CodeFlashBenchmarkPlugin()
+
+
+def pytest_configure(config: pytest.Config) -> None:
+    config.addinivalue_line("markers", "benchmark: mark test as a benchmark that should be run with codeflash tracing")
+
+
+def pytest_addoption(parser: pytest.Parser) -> None:
+    parser.addoption(
+        "--codeflash-trace", action="store_true", default=False, help="Enable CodeFlash tracing for benchmarks"
+    )
+
+
+@pytest.fixture
+def benchmark(request: pytest.FixtureRequest) -> object:
+    if not request.config.getoption("--codeflash-trace"):
+        return lambda func, *args, **kwargs: func(*args, **kwargs)
+    return codeflash_benchmark_plugin.Benchmark(request)

codeflash/optimization/optimizer.py

@@ -90,13 +90,13 @@ def run_benchmarks(
                     logger.info(
                         f"No valid benchmarks found in {self.args.benchmarks_root} for functions to optimize, continuing optimization"
                     )
-                    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)
+                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)
             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.")

pyproject.toml

@@ -67,7 +67,6 @@ dev = [
     "types-openpyxl>=3.1.5.20241020",
     "types-regex>=2024.9.11.20240912",
     "types-python-dateutil>=2.9.0.20241003",
-    "pytest-benchmark>=5.1.0",
     "types-gevent>=24.11.0.20241230,<25",
     "types-greenlet>=3.1.0.20241221,<4",
     "types-pexpect>=4.9.0.20241208,<5",
@@ -300,3 +299,6 @@ markers = [
 [build-system]
 requires = ["hatchling", "uv-dynamic-versioning"]
 build-backend = "hatchling.build"
+
+[project.entry-points.pytest11]
+codeflash = "codeflash.benchmarking.plugin.plugin"