Add pytest-benchmark benchmarks against a fixed synthetic dataset

.github/workflows/benchmarks.yml

@@ -0,0 +1,79 @@
+name: benchmarks
+on:
+  push:
+    branches:
+      - master
+permissions:
+  # Each run pushes its results to the gh-pages branch.
+  contents: write
+# Each run adds to the history on gh-pages, so runs can't overlap.
+concurrency: benchmarks
+jobs:
+  benchmark:
+    runs-on: ubuntu-24.04
+    services:
+      postgres:
+        image: pgvector/pgvector:0.8.6-pg18
+        env:
+          POSTGRES_PASSWORD: postgres
+        options: >-
+          --health-cmd "pg_isready --username postgres"
+          --health-start-period 30s
+          --health-start-interval 2s
+        ports:
+          - 5432:5432
+      elasticsearch:
+        image: elasticsearch:9.5.2
+        env:
+          ES_JAVA_OPTS: "-Xms250m -Xmx750m"
+          discovery.type: single-node
+          xpack.security.enabled: "true"
+          ELASTIC_PASSWORD: elastic
+        options: >-
+          --health-cmd "curl --fail --user elastic:elastic http://localhost:9200/"
+          --health-start-period 30s
+          --health-start-interval 2s
+        ports:
+          - 9200:9200
+    env:
+      DJANGO_DATABASE_URL: postgres://postgres:postgres@localhost:5432/django
+      DJANGO_ISIC_ELASTICSEARCH_URL: http://elastic:elastic@localhost:9200
+      # The same as ci.yml. No benchmark connects to them, so this job has no services for them.
+      DJANGO_CELERY_BROKER_URL: amqp://localhost:5672/
+      DJANGO_CACHE_URL: redis://localhost:6379/0
+    steps:
+      - name: Checkout repository
+        uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
+      - name: Install uv
+        uses: astral-sh/setup-uv@11f9893b081a58869d3b5fccaea48c9e9e46f990 # v8.3.2
+      - name: Increase maintenance_work_mem for pgvector
+        # The dataset rebuilds the imageembedding IVFFlat index; see ci.yml.
+        run: psql -c "ALTER SYSTEM SET maintenance_work_mem = '512MB'" -c "SELECT pg_reload_conf();"
+        env:
+          PGHOST: localhost
+          PGUSER: postgres
+          PGPASSWORD: postgres
+      - name: Run benchmarks
+        run: uv run --locked tox -e benchmark -- --benchmark-json=.benchmarks/ci.json
+      - name: Checkout gh-pages
+        uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
+        with:
+          ref: gh-pages
+          path: gh-pages
+      - name: Record results
+        run: |
+          uv run --locked python benchmarks/history.py record .benchmarks/ci.json gh-pages/benchmarks/data.js
+          cp benchmarks/index.html gh-pages/benchmarks/
+      - name: Push results
+        # A regression is still recorded, so the next run compares against it.
+        working-directory: gh-pages
+        run: |
+          # github-actions[bot]'s noreply address (41898282 is its user ID), as in actions/checkout's
+          # README, so GitHub attributes the commit to that bot.
+          git config user.name "github-actions[bot]"
+          git config user.email "41898282+github-actions[bot]@users.noreply.github.com"
+          git add benchmarks
+          git commit --message "Record benchmarks for $GITHUB_SHA"
+          git push
+      - name: Check for regressions
+        run: uv run --locked python benchmarks/history.py check gh-pages/benchmarks/data.js

.gitignore

@@ -68,6 +68,7 @@ pip-delete-this-directory.txt
 htmlcov/
 .tox/
 .nox/
+.benchmarks/
 .coverage
 .coverage.*
 .cache

benchmarks/README.md

@@ -0,0 +1,26 @@
+# Benchmarks
+
+[pytest-benchmark](https://pytest-benchmark.readthedocs.io) benchmarks against a fixed synthetic dataset (`dataset.py`), so results are comparable between runs and commits.
+
+They need the Postgres and Elasticsearch from `docker compose`. They use their own test database (`test_isic_benchmark`) and `benchmark-*` indices, so the development and test databases and indices are never touched.
+
+```sh
+uv run tox -e benchmark                             # everything
+uv run tox -e benchmark -- -k image_search          # a subset
+uv run tox -e benchmark -- --benchmark-save=before  # save the run to .benchmarks/
+uv run tox -e benchmark -- --benchmark-compare      # compare with the last saved run
+uv run tox -e benchmark -- --benchmark-disable      # run each benchmark once, without timing
+```
+
+- Each session starts from a fresh copy of the dataset. The first run for a schema generates about 100k images, which takes about 2 minutes, and keeps them in a `test_isic_benchmark_<hash>` database that later runs copy. Drop old ones with `DROP DATABASE test_isic_benchmark_<hash>`.
+- Caching is disabled (`DummyCache`, cachalot off), so every request does its full work.
+- Benchmarks that write data roll back each round.
+- Benchmarks that use the `benchmark_with_memory` fixture also record the peak memory Python allocates, measured in a separate call. Memory allocated outside Python, such as by psycopg, isn't counted.
+- CI (`.github/workflows/benchmarks.yml`) runs on each push to master. `history.py record` adds each benchmark's fastest round and peak memory to `benchmarks/data.js` on the `gh-pages` branch, and the job copies `index.html` next to it. Then `history.py check` fails the job when either is more than 50% higher than in the run before it.
+- `index.html` charts `data.js`, grouped by module and test, at https://imagemarkup.github.io/isic/benchmarks/. To chart local runs:
+
+  ```sh
+  uv run tox -e benchmark -- --benchmark-json=.benchmarks/run.json
+  uv run python benchmarks/history.py record .benchmarks/run.json .benchmarks/site/data.js
+  cp benchmarks/index.html .benchmarks/site/ && python -m http.server -d .benchmarks/site
+  ```

benchmarks/conftest.py

@@ -0,0 +1,156 @@
+from collections.abc import Callable
+import tracemalloc
+
+from django.conf import settings
+from django.contrib.auth.models import User
+from django.db import connection, transaction
+from django.test import Client
+from django.test.utils import setup_databases, teardown_databases
+import pytest
+
+from benchmarks import dataset
+from benchmarks.dataset import Viewer
+from isic.core.dsl import parse_query
+from isic.core.search import IMAGE_INDEX_MAPPINGS, LESION_INDEX_MAPPINGS, get_elasticsearch_client
+from isic.core.tasks import sync_elasticsearch_indices_task
+
+
+@pytest.fixture(scope="session")
+def django_db_setup(django_test_environment, django_db_blocker):
+    # pytest-django only creates the test database for tests marked django_db, and the marker
+    # wraps each test in a transaction, which some views can't run in (e.g. setting the isolation
+    # level). so the database is created the same way, but for every benchmark, and copied from
+    # the generated dataset when there is one.
+    # migrations create the indices, and refuse to when they exist with other mappings, e.g. from
+    # a run of another commit.
+    client = get_elasticsearch_client()
+    for index in [
+        settings.ISIC_ELASTICSEARCH_IMAGES_INDEX,
+        settings.ISIC_ELASTICSEARCH_LESIONS_INDEX,
+    ]:
+        client.indices.delete(index=index, ignore_unavailable=True)
+    with django_db_blocker.unblock():
+        generated = dataset.template_exists()
+        if generated:
+            connection.settings_dict["TEST"]["TEMPLATE"] = dataset.template_name()
+        config = setup_databases(
+            verbosity=0, interactive=False, aliases={"default"}, serialized_aliases=set()
+        )
+        if not generated:
+            dataset.generate()
+    yield
+    with django_db_blocker.unblock():
+        teardown_databases(config, verbosity=0)
+
+
+@pytest.fixture(scope="session")
+def manifest(django_db_setup, django_db_blocker) -> dict:
+    with django_db_blocker.unblock():
+        return dataset.manifest()
+
+
+@pytest.fixture(autouse=True)
+def _database(django_db_setup, django_db_blocker):
+    with django_db_blocker.unblock():
+        yield
+
+
+INDEX_SETTINGS = {
+    "index.number_of_replicas": 0,
+    # nothing writes during a run. an explicit interval also stops idle shards from making the
+    # next search wait for a refresh.
+    "index.refresh_interval": "-1",
+    # every request should do its full work
+    "index.requests.cache.enable": False,
+    "index.queries.cache.enabled": False,
+}
+
+
+@pytest.fixture(scope="session")
+def search_index(django_db_setup, django_db_blocker):
+    """Index the dataset into the benchmark indices."""
+    client = get_elasticsearch_client()
+    indices = {
+        settings.ISIC_ELASTICSEARCH_IMAGES_INDEX: IMAGE_INDEX_MAPPINGS,
+        settings.ISIC_ELASTICSEARCH_LESIONS_INDEX: LESION_INDEX_MAPPINGS,
+    }
+    for index, mappings in indices.items():
+        client.indices.delete(index=index, ignore_unavailable=True)
+        client.indices.create(index=index, mappings=mappings, settings=INDEX_SETTINGS)
+    with django_db_blocker.unblock():
+        sync_elasticsearch_indices_task()
+    client.indices.refresh(index=list(indices))
+    # bulk indexing leaves segments that elasticsearch keeps merging in the background, which
+    # slows whichever benchmarks run first. merged up front, every run searches the same index.
+    client.indices.forcemerge(index=list(indices), max_num_segments=1)
+    client.indices.refresh(index=list(indices))
+
+
+def _fetch(client: Client, path: str) -> bytes:
+    parse_query.cache_clear()
+    response = client.get(path)
+    if response.status_code != 200:
+        raise AssertionError(f"{path}: {response.status_code}")
+    if response.streaming:
+        return b"".join(response.streaming_content)
+    return response.content
+
+
+@pytest.fixture
+def endpoint(manifest) -> Callable[..., Callable[[], bytes]]:
+    """
+    Log in as a user and return a callable that requests a path.
+
+    The path can be a function of the logged in client. It's requested once before returning.
+    """
+
+    def make(user: Viewer, path: str | Callable[[Client], str]) -> Callable[[], bytes]:
+        client = Client()
+        if user != "anonymous":
+            client.force_login(User.objects.get(pk=manifest["users"][user]))
+        if callable(path):
+            path = path(client)
+        _fetch(client, path)
+        return lambda: _fetch(client, path)
+
+    return make
+
+
+@pytest.fixture
+def rolled_back() -> Callable[[Callable[[], object]], Callable[[], None]]:
+    """
+    Wrap a function so each call's writes are rolled back.
+
+    A benchmark calls the function many times, and each call has to start from the same data.
+    pytest-django's rollback only happens once, at the end of a test.
+    """
+
+    def wrap(fn):
+        def run():
+            with transaction.atomic():
+                fn()
+                transaction.set_rollback(True)
+
+        return run
+
+    return wrap
+
+
+@pytest.fixture
+def benchmark_with_memory(benchmark) -> Callable[[Callable[[], object]], None]:
+    """
+    Benchmark a function, and record the peak memory Python allocates while calling it.
+
+    Memory is measured in a separate call, since tracing allocations slows the function down.
+    """
+
+    def run(fn):
+        tracemalloc.start()
+        try:
+            fn()
+            benchmark.extra_info["peak_memory_mib"] = tracemalloc.get_traced_memory()[1] / 2**20
+        finally:
+            tracemalloc.stop()
+        benchmark(fn)
+
+    return run