Add performance profiling (#923)

Author: ofek

.github/CONTRIBUTING.md

@@ -140,6 +140,16 @@ just bench encodings
 
 Omit the benchmark name to list all available benchmarks.
 
+## Profiling
+
+You can run all performance profiling tests with the `test-perf` recipe, for example:
+
+```
+just test-perf basic
+```
+
+Omit the profiling test name to list all available profiling tests.
+
 ## Python Version Management
 
 You can use the `python` variable to specify the version of Python to use for each environment. For example, to run tests with Python 3.12:

.github/workflows/ci.yml

@@ -168,6 +168,37 @@ jobs:
         use_oidc: true
         directory: coverage-files
 
+  profile:
+    name: Run profiling tests
+    needs:
+    - test
+    runs-on: ubuntu-latest
+
+    steps:
+    - name: Checkout code
+      uses: actions/checkout@v5
+
+    - name: Load environment file
+      uses: ./.github/actions/load-env
+
+    - name: Install Python
+      uses: actions/setup-python@v6
+      with:
+        python-version-file: pyproject.toml
+
+    - name: Install uv
+      uses: astral-sh/setup-uv@v7
+
+    - name: Install command runner
+      run: uv tool install rust-just
+
+    - name: Install profiling dependencies
+      run: just env-sync prof
+
+    # TODO: set up infra to commit results after each merge and then make tests compare
+    - name: Run profiling tests
+      run: just test-perf all --calibrate --rounds 10000
+
   build-wheels:
     name: Build wheels on ${{ matrix.os }} for ${{ matrix.archs }}
     needs:

.justfile

@@ -51,7 +51,7 @@ alias pre-commit := prek
 [group: "Static Analysis"]
 hooks-fix *args: (
   env-run "hooks"
-  "prek run --show-diff-on-failure --all-files"
+  "prek run --all-files"
   args
 )
 alias hooks := hooks-fix
@@ -60,7 +60,7 @@ alias hooks := hooks-fix
 [group: "Static Analysis"]
 hooks-check *args: (
   env-run "hooks"
-  "prek run --show-diff-on-failure --all-files --hook-stage manual"
+  "prek run --all-files --hook-stage manual"
   args
 )
 alias check := hooks-check
@@ -112,13 +112,38 @@ doc-serve: (
   "python -c \"import pathlib,webbrowser;webbrowser.open_new_tab(pathlib.Path('docs/build/html/index.html').absolute().as_uri())\""
 )
 
+# Run performance tests.
+[group: "Profiling"]
+test-perf name="" *args: (
+  env-run "prof"
+  (
+    if name == "" {
+      "python -c \"import os;print('Available tests:', os.linesep.join(['', *sorted(e[5:].split('.')[0] for e in os.listdir('tests/prof/perf') if e.startswith('test_')), 'all (run everything)']))\""
+    } else {
+      (
+        if name =~ "^-" {
+          error(
+            "Invalid test name: " + name
+          )
+        } else if name == "all" {
+          "pytest -o testpaths=tests/prof/perf"
+        } else {
+          "pytest tests/prof/perf/test_" + name + ".py"
+        }
+      )
+      + " --benchmark-name short --benchmark-disable-gc"
+    }
+  )
+  args
+)
+
 # Run benchmarks.
 [group: "Benchmarking"]
 bench-run name="" *args: (
   env-run "bench"
   (
     if name == "" {
-      "python -c \"import os;print(os.linesep.join(sorted(e[6:].split('.')[0] for e in os.listdir('benchmarks') if e.startswith('bench_'))))\""
+      "python -c \"import os;print('Available benchmarks:', os.linesep.join(['', *sorted(e[6:].split('.')[0] for e in os.listdir('benchmarks') if e.startswith('bench_')), 'all (run everything)']))\""
     } else {
       "python -m benchmarks.bench_" + name
     }
@@ -212,6 +237,8 @@ _env_sync_all: (
   _with_env "doc" "sync"
 ) (
   _with_env "bench" "sync"
+) (
+  _with_env "prof" "sync"
 )
 
 [private]
@@ -254,7 +281,9 @@ _with_env env action *args:
     } else if action == "sync" { \
       "sync" \
     } else { \
-      error("Unknown action: " + action) \
+      error( \
+        "Unknown action: " + action \
+      ) \
     } \
   }}{{ \
     if rebuild =~ "^(true|1)$" { \
@@ -267,13 +296,17 @@ _with_env env action *args:
       "--group test-unit --group test-typing" \
     } else if env == "doc" { \
       "--group doc" \
+    } else if env == "prof" { \
+      "--group prof" \
     } else if env == "bench" { \
       "--group bench" \
     } else if env == "hooks" { \
       "--only-group hooks" \
     } else if env == "none" { \
       "--isolated --no-sync --no-config" \
     } else { \
-      error("Unknown environment: " + env) \
+      error( \
+        "Unknown environment: " + env \
+      ) \
     } \
   }} {{ args }}

pyproject.toml

@@ -104,6 +104,12 @@ hooks = [
   "prek>=0.2.13",
   "ruff==0.14.1",
 ]
+prof = [
+  "memray; sys_platform != 'win32'",
+  "pytest",
+  "pytest-benchmark",
+  "pytest-memray; sys_platform != 'win32'",
+]
 test-typing = [
   "mypy",
   "pyright",

tests/conftest.py

@@ -0,0 +1,94 @@
+import os
+import re
+import sys
+from pathlib import Path
+
+
+def pytest_addoption(parser):
+    # Declaring additional command-line options within test suite-specific `conftest.py`
+    # files is documented as unsupported by pytest. See:
+    # https://docs.pytest.org/en/stable/reference/reference.html#pytest.hookspec.pytest_addoption.
+    #
+    # However, it seems to work in practice with the following limitations:
+    #   - The options only appear in the help text for suites that are configured in
+    #     the default `testpaths` configuration. This is undesirable because we don't
+    #     want a missing argument to trigger heavy tests like performance profiling.
+    #   - The options are parsed differently such that an equals sign is required for
+    #     options with a value rather than the more natural space-separated format.
+    #
+    # As a workaround, we parse the command line for options that match the patterns of
+    # the configurable test suites and add their options to the parser.
+    configurable_test_suites = {
+        "tests/prof/perf": add_performance_profiling_options,
+    }
+
+    cwd = Path.cwd()
+    project_root = Path(__file__).parent.parent
+    for test_suite_relpath, add_options_func in list(configurable_test_suites.items()):
+        test_suite_path = project_root.joinpath(test_suite_relpath)
+        if test_suite_path == cwd or test_suite_path in cwd.parents:
+            add_options_func(parser)
+            configurable_test_suites.pop(test_suite_relpath)
+            break
+
+    # Always accept forward slashes and the system separator, escaping the latter in case of backslashes (Windows).
+    allowed_separators = {"/", re.escape(os.sep)}
+    separator_pattern = f"[{''.join(sorted(allowed_separators))}]"
+    configurable_test_suites_patterns = {
+        re.compile(
+            rf"\b{test_suite_relpath.replace('/', separator_pattern)}\b"
+        ): add_options_func
+        for test_suite_relpath, add_options_func in configurable_test_suites.items()
+    }
+
+    for arg in sys.argv:
+        if not configurable_test_suites_patterns:
+            break
+
+        for pattern, add_options_func in list(
+            configurable_test_suites_patterns.items()
+        ):
+            if arg in {"-h", "--help"} or pattern.search(arg):
+                add_options_func(parser)
+                configurable_test_suites_patterns.pop(pattern)
+                break
+
+
+def add_performance_profiling_options(parser):
+    group = parser.getgroup("performance profiling")
+    group.addoption(
+        "--calibrate",
+        action="store_true",
+        help=(
+            "Override automatic calibration of benchmarks and enable the following options: ",
+            "--rounds, --warmup-rounds, --iterations",
+        ),
+    )
+    group.addoption(
+        "--rounds",
+        type=int,
+        default=1000,
+        metavar="ROUNDS",
+        help="Number of rounds for benchmarks (default: 1000)",
+    )
+    group.addoption(
+        "--warmup-rounds",
+        type=int,
+        default=50,
+        metavar="WARMUP_ROUNDS",
+        help="Number of warmup rounds for benchmarks (default: 50)",
+    )
+    group.addoption(
+        "--iterations",
+        type=int,
+        default=10,
+        metavar="ITERATIONS",
+        help="Number of iterations in each round of benchmarks (default: 10)",
+    )
+    group.addoption(
+        "--size",
+        type=int,
+        default=1000,
+        metavar="SIZE",
+        help="Size of the benchmark data (default: 1000)",
+    )

tests/prof/__init__.py

@@ -0,0 +1,42 @@
+import pytest
+
+from benchmarks.generate_data import make_filesystem_data
+
+
+@pytest.fixture
+def bench(benchmark, bench_config):
+    if not bench_config["calibrate"]:
+        return lambda target, *args, **kwargs: benchmark(target, *args, **kwargs)
+
+    return lambda target, *args, **kwargs: benchmark.pedantic(
+        target,
+        args=args,
+        kwargs=kwargs,
+        rounds=bench_config["rounds"],
+        iterations=bench_config["iterations"],
+        warmup_rounds=bench_config["warmup_rounds"],
+    )
+
+
+@pytest.fixture
+def filesystem_data(bench_config):
+    return lambda size=bench_config["size"]: make_filesystem_data(size)
+
+
+@pytest.fixture(scope="class")
+def shared_data():
+    """
+    This is used to share data between serially-executed tests within a class.
+    """
+    return {}
+
+
+@pytest.fixture
+def bench_config(request):
+    return {
+        "calibrate": request.config.getoption("--calibrate"),
+        "rounds": request.config.getoption("--rounds"),
+        "warmup_rounds": request.config.getoption("--warmup-rounds"),
+        "iterations": request.config.getoption("--iterations"),
+        "size": request.config.getoption("--size"),
+    }

tests/prof/perf/__init__.py

@@ -0,0 +1,19 @@
+import pytest
+
+from benchmarks.bench_validation import bench_msgspec
+
+
+@pytest.mark.benchmark(group="roundtrip")
+class TestRoundtrip:
+    def test_encode(self, bench, filesystem_data, shared_data):
+        data = filesystem_data()
+        encoded = bench(bench_msgspec.encode, data)
+        shared_data["encoded"] = encoded
+
+    def test_decode(self, bench, shared_data):
+        if shared_data["encoded"] is None:
+            raise ValueError(
+                "test_encode must run before test_decode to provide encoded data"
+            )
+
+        bench(bench_msgspec.decode, shared_data["encoded"])