feat: add memory profiling support

.github/workflows/codspeed.yml

@@ -10,6 +10,10 @@ env:
   PYTHON_VERSION: "3.14"
   SHARDS: 4
 
+permissions:
+  contents: read # required for actions/checkout
+  id-token: write # required for OIDC authentication with CodSpeed
+
 jobs:
   benchmarks:
     strategy:
@@ -33,12 +37,14 @@ jobs:
           sudo apt-get install valgrind -y
           uv sync --dev
           sudo apt-get remove valgrind -y
+      - uses: dtolnay/rust-toolchain@stable
       - name: Run benchmarks
         uses: CodSpeedHQ/action@main
         with:
+          # runner-version: branch:cod-2105-support-memory-profiling-for-python
           mode: ${{ matrix.mode }}
           run: uv run pytest tests/benchmarks/ --codspeed --test-group=${{ matrix.shard }} --test-group-count=${{ env.SHARDS }}
-          token: ${{ secrets.CODSPEED_TOKEN }}
+          # upload-url: https://api.staging.preview.codspeed.io/upload
 
   all-checks:
     runs-on: ubuntu-latest

src/pytest_codspeed/instruments/__init__.py

@@ -21,7 +21,7 @@ class Instrument(metaclass=ABCMeta):
     instrument: ClassVar[str]
 
     @abstractmethod
-    def __init__(self, config: CodSpeedConfig): ...
+    def __init__(self, config: CodSpeedConfig, mode: MeasurementMode): ...
 
     @abstractmethod
     def get_instrument_config_str_and_warns(self) -> tuple[str, list[str]]: ...
@@ -57,6 +57,7 @@ def get_result_dict(
 
 class MeasurementMode(str, Enum):
     Simulation = "simulation"
+    Memory = "memory"
     WallTime = "walltime"
 
     @classmethod
@@ -68,12 +69,12 @@ def _missing_(cls, value: object):
 
 
 def get_instrument_from_mode(mode: MeasurementMode) -> type[Instrument]:
-    from pytest_codspeed.instruments.valgrind import (
-        ValgrindInstrument,
+    from pytest_codspeed.instruments.analysis import (
+        AnalysisInstrument,
     )
     from pytest_codspeed.instruments.walltime import WallTimeInstrument
 
-    if mode == MeasurementMode.Simulation:
-        return ValgrindInstrument
+    if mode in (MeasurementMode.Simulation, MeasurementMode.Memory):
+        return AnalysisInstrument
     else:
         return WallTimeInstrument

src/pytest_codspeed/instruments/valgrind.py → src/pytest_codspeed/instruments/analysis.py

@@ -6,7 +6,10 @@
 
 from pytest_codspeed import __semver_version__
 from pytest_codspeed.instruments import Instrument
-from pytest_codspeed.instruments.hooks import InstrumentHooks
+from pytest_codspeed.instruments.hooks import (
+    FEATURE_DISABLE_CALLGRIND_MARKERS,
+    InstrumentHooks,
+)
 from pytest_codspeed.utils import SUPPORTS_PERF_TRAMPOLINE
 
 if TYPE_CHECKING:
@@ -15,31 +18,33 @@
     from pytest import Session
 
     from pytest_codspeed.config import PedanticOptions
-    from pytest_codspeed.instruments import P, T
+    from pytest_codspeed.instruments import MeasurementMode, P, T
     from pytest_codspeed.plugin import BenchmarkMarkerOptions, CodSpeedConfig
 
 
-class ValgrindInstrument(Instrument):
-    instrument = "valgrind"
+class AnalysisInstrument(Instrument):
+    instrument = "analysis"
     instrument_hooks: InstrumentHooks | None
+    mode: MeasurementMode
 
-    def __init__(self, config: CodSpeedConfig) -> None:
+    def __init__(self, config: CodSpeedConfig, mode: MeasurementMode) -> None:
+        self.mode = mode
         self.benchmark_count = 0
         try:
             self.instrument_hooks = InstrumentHooks()
             self.instrument_hooks.set_integration("pytest-codspeed", __semver_version__)
         except RuntimeError as e:
             if os.environ.get("CODSPEED_ENV") is not None:
                 raise Exception(
-                    "Failed to initialize CPU simulation instrument hooks"
+                    f"Failed to initialize {self.mode.value} instrument hooks"
                 ) from e
             self.instrument_hooks = None
 
         self.should_measure = self.instrument_hooks is not None
 
     def get_instrument_config_str_and_warns(self) -> tuple[str, list[str]]:
         config = (
-            f"mode: simulation, "
+            f"mode: {self.mode.value}, "
             f"callgraph: {'enabled' if SUPPORTS_PERF_TRAMPOLINE else 'not supported'}"
         )
         warnings = []
@@ -73,6 +78,9 @@ def __codspeed_root_frame__() -> T:
             # Warmup CPython performance map cache
             __codspeed_root_frame__()
 
+        self.instrument_hooks.set_feature(FEATURE_DISABLE_CALLGRIND_MARKERS, True)
+        self.instrument_hooks.start_benchmark()
+
         # Manually call the library function to avoid an extra stack frame. Also
         # call the callgrind markers directly to avoid extra overhead.
         self.instrument_hooks.lib.callgrind_start_instrumentation()
@@ -81,6 +89,7 @@ def __codspeed_root_frame__() -> T:
         finally:
             # Ensure instrumentation is stopped even if the test failed
             self.instrument_hooks.lib.callgrind_stop_instrumentation()
+            self.instrument_hooks.stop_benchmark()
             self.instrument_hooks.set_executed_benchmark(uri)
 
     def measure_pedantic(
@@ -92,8 +101,8 @@ def measure_pedantic(
     ) -> T:
         if pedantic_options.rounds != 1 or pedantic_options.iterations != 1:
             warnings.warn(
-                "Valgrind instrument ignores rounds and iterations settings "
-                "in pedantic mode"
+                f"{self.mode.value.capitalize()} instrument ignores rounds and "
+                "iterations settings in pedantic mode"
             )
         if not self.instrument_hooks:
             args, kwargs = pedantic_options.setup_and_get_args_kwargs()
@@ -117,11 +126,18 @@ def __codspeed_root_frame__(*args, **kwargs) -> T:
 
         # Compute the actual result of the function
         args, kwargs = pedantic_options.setup_and_get_args_kwargs()
+
+        self.instrument_hooks.set_feature(FEATURE_DISABLE_CALLGRIND_MARKERS, True)
+        self.instrument_hooks.start_benchmark()
+
+        # Manually call the library function to avoid an extra stack frame. Also
+        # call the callgrind markers directly to avoid extra overhead.
         self.instrument_hooks.lib.callgrind_start_instrumentation()
         try:
             out = __codspeed_root_frame__(*args, **kwargs)
         finally:
             self.instrument_hooks.lib.callgrind_stop_instrumentation()
+            self.instrument_hooks.stop_benchmark()
             self.instrument_hooks.set_executed_benchmark(uri)
             if pedantic_options.teardown is not None:
                 pedantic_options.teardown(*args, **kwargs)
@@ -140,5 +156,5 @@ def report(self, session: Session) -> None:
     def get_result_dict(self) -> dict[str, Any]:
         return {
             "instrument": {"type": self.instrument},
-            # bench results will be dumped by valgrind
+            # bench results will be dumped by the runner
         }

src/pytest_codspeed/instruments/hooks/__init__.py

@@ -10,6 +10,9 @@
 if TYPE_CHECKING:
     from .dist_instrument_hooks import InstrumentHooksPointer, LibType
 
+# Feature flags for instrument hooks
+FEATURE_DISABLE_CALLGRIND_MARKERS = 0
+
 
 class InstrumentHooks:
     """Zig library wrapper class providing benchmark measurement functionality."""
@@ -80,3 +83,12 @@ def set_integration(self, name: str, version: str) -> None:
     def is_instrumented(self) -> bool:
         """Check if simulation is active."""
         return self.lib.instrument_hooks_is_instrumented(self.instance)
+
+    def set_feature(self, feature: int, enabled: bool) -> None:
+        """Set a feature flag in the instrument hooks library.
+
+        Args:
+            feature: The feature flag to set
+            enabled: Whether to enable or disable the feature
+        """
+        self.lib.instrument_hooks_set_feature(feature, enabled)

src/pytest_codspeed/instruments/hooks/build.py

@@ -34,6 +34,8 @@
 
 void callgrind_start_instrumentation();
 void callgrind_stop_instrumentation();
+
+void instrument_hooks_set_feature(uint64_t feature, bool enabled);
 """)
 
 ffibuilder.set_source(

src/pytest_codspeed/instruments/hooks/dist_instrument_hooks.pyi

@@ -29,5 +29,7 @@ class lib:
     def callgrind_start_instrumentation() -> int: ...
     @staticmethod
     def callgrind_stop_instrumentation() -> int: ...
+    @staticmethod
+    def instrument_hooks_set_feature(feature: int, enabled: bool) -> None: ...
 
 LibType = type[lib]

src/pytest_codspeed/instruments/walltime.py

@@ -24,7 +24,7 @@
     from pytest import Session
 
     from pytest_codspeed.config import PedanticOptions
-    from pytest_codspeed.instruments import P, T
+    from pytest_codspeed.instruments import MeasurementMode, P, T
     from pytest_codspeed.plugin import BenchmarkMarkerOptions, CodSpeedConfig
 
 DEFAULT_WARMUP_TIME_NS = 1_000_000_000
@@ -159,7 +159,7 @@ class WallTimeInstrument(Instrument):
     instrument = "walltime"
     instrument_hooks: InstrumentHooks | None
 
-    def __init__(self, config: CodSpeedConfig) -> None:
+    def __init__(self, config: CodSpeedConfig, _mode: MeasurementMode) -> None:
         try:
             self.instrument_hooks = InstrumentHooks()
             self.instrument_hooks.set_integration("pytest-codspeed", __semver_version__)

src/pytest_codspeed/plugin.py

@@ -112,8 +112,11 @@ def pytest_configure(config: pytest.Config):
     )
 
     if os.environ.get("CODSPEED_ENV") is not None:
-        if os.environ.get("CODSPEED_RUNNER_MODE") == "walltime":
+        runner_mode = os.environ.get("CODSPEED_RUNNER_MODE")
+        if runner_mode == "walltime":
             default_mode = MeasurementMode.WallTime.value
+        elif runner_mode == "memory":
+            default_mode = MeasurementMode.Memory.value
         else:
             default_mode = MeasurementMode.Simulation.value
     else:
@@ -142,7 +145,7 @@ def pytest_configure(config: pytest.Config):
         disabled_plugins=tuple(disabled_plugins),
         is_codspeed_enabled=is_codspeed_enabled,
         mode=mode,
-        instrument=instrument(codspeed_config),
+        instrument=instrument(codspeed_config, mode),
         config=codspeed_config,
         profile_folder=Path(profile_folder) if profile_folder else None,
     )

tests/test_pytest_plugin_cpu_instrumentation.py

@@ -76,7 +76,7 @@ def fixtured_child():
     with open(perf_filepath) as perf_file:
         lines = perf_file.readlines()
         assert any(
-            "py::ValgrindInstrument.measure.<locals>.__codspeed_root_frame__" in line
+            "py::AnalysisInstrument.measure.<locals>.__codspeed_root_frame__" in line
             for line in lines
         ), "No root frame found in perf map"
         assert any("py::test_some_addition_marked" in line for line in lines), (
@@ -135,8 +135,8 @@ def foo():
     result = run_pytest_codspeed_with_mode(pytester, MeasurementMode.Simulation)
     result.stdout.fnmatch_lines(
         [
-            "*UserWarning: Valgrind instrument ignores rounds and iterations settings "
-            "in pedantic mode*"
+            "*UserWarning: Simulation instrument ignores rounds and iterations settings"
+            " in pedantic mode*"
         ]
     )
     result.assert_outcomes(passed=1)