Benchmarking

pylcm uses ASV (Airspeed Velocity) to track performance across commits. Benchmarks run locally on GPU hardware and results are published to a dashboard.

Machine Setup¶

On first use, register your machine with ASV:

pixi run asv-machine

This runs asv machine --yes and then records the RAM in whole gigabytes, so the kilobyte drift of the kernel’s reported memory does not start a new dashboard series. ASV copies the entry to .asv/results/<machine-name>/machine.json on the next run.

Running Benchmarks¶

There are two primary workflows depending on whether you are working on main or a PR branch:

PR branches — run benchmarks and post a comparison comment on your PR:

pixi run -e benchmarks-cuda12 asv-run-and-pr-comment

Main branch — run benchmarks and publish results to the dashboard:

pixi run -e benchmarks-cuda12 asv-run-and-publish-main

Both workflows run asv-run (which requires a clean worktree) followed by their respective post-processing step. The PR workflow passes asv-run an ASV --bench regex that skips the ACA debug-log timing (AcaBaselineDebugLog); only the main-branch run measures it. The PR comparison table lists the benchmarks HEAD ran, so a benchmark with a main result but none on HEAD gets no row.

Individual tasks are also available:

# Run all benchmarks (GPU required)
pixi run -e benchmarks-cuda12 asv-run

# Quick smoke test (not saved)
pixi run -e benchmarks-cuda12 asv-quick

# Post benchmark comment to current PR (no GPU needed)
pixi run asv-pr-comment

# Compare two commits (no GPU needed)
pixi run asv-compare HEAD~1 HEAD

# Preview dashboard locally (no GPU needed)
pixi run asv-preview

# Publish results to dashboard (no GPU needed)
pixi run asv-publish

The asv-run and asv-quick tasks set XLA_PYTHON_CLIENT_PREALLOCATE=false automatically so JAX allocates GPU memory on demand rather than grabbing it all up front. The default device-memory budget refuses an on-demand pool, so a benchmark that builds its model inside the ASV process passes _gpu_mem.default_budget_execution_kwargs(): the default budget the devices would resolve to, as an explicit device_memory_bytes. The isolated measurement subprocesses in _gpu_mem turn preallocation back on, as production runs: an on-demand pool grows in separate regions, so an explicit budget the size of the device pool can admit a buffer that no free region holds. There the ACA benchmark passes no execution policy and runs under aca-model’s default, as production does. peak_bytes_in_use counts the allocator chunks handed to buffers, not the pool. A preallocated pool carves each chunk to the requested size; an on-demand pool can hand a buffer a whole newly grown region, so it reads higher for the same programs.

Exact paired GridSearch measurements¶

Use the separate paired harness when a GridSearch execution change needs evidence against one exact base revision. Run it from a clean checkout of the head revision and provide separate, clean checkouts of both revisions:

pixi run -e benchmarks-cuda12 grid-search-pair \
    --base-checkout /path/to/base \
    --base-revision <full-base-commit> \
    --head-checkout /path/to/head \
    --head-revision <full-head-commit> \
    --harness-revision <full-harness-commit> \
    --output /path/outside/both/checkouts/grid-search-pair \
    --precision 32 \
    --backend gpu \
    --repeats 4

All three revisions must be full 40-character commit hashes. The checkout you run the task from supplies the harness and is validated against --harness-revision, so it need not be the head checkout; all three checkouts must agree on the lock file and scenario-owned sources. The output directory must not already exist and must lie outside the harness, base, and head checkouts.

The harness measures ten named scenarios. Four are GridSearch routes: singleton hard max, singleton EV1, collective GridSearch with value-dependent inputs, and distributed co-map. The other six are the full 18-regime ACA baseline at each combination of 3 or 6 asset points and 16, 64, or 256 consumption points, named aca-a<assets>-c<consumption>. Pass --scenario once per name to measure a subset; omitting it measures all ten. The distributed row always uses exactly four CPU devices; the other rows use the requested backend. Use at least two repeats so execution alternates base--head on even repeats and head--base on odd repeats; four repeats are the balanced evidence profile.

Every pair must pass numerical value parity and exact output shape, dtype, and sharding parity. The harness also requires non-empty HLO and compiler-memory status for every compiled core, host high-water-mark evidence, and measured peak device memory on GPU (with an explicit not-applicable record on CPU). It verifies the native execution declaration for every named target: singleton hard max and distributed co-map stream action_product; EV1 and collective GridSearch keep canonical dense action reducers while the planner tiles their state cells. The reported streamed flag means action reduction, so a planned core can have streamed=False. It also rejects communication collectives in the distributed co-map head. Raw timings, compiler and device memory, HLO files and digests, layout manifests, and the base/head ratios are retained under the output directory; summary.json is the entry point for review.

Benchmark Scenarios¶

The reduced ACA benchmark preserves the 18-regime structure while using its benchmark profile. Production ACA runs use the production grid configuration and 3 preference types.

The ACA report labels its timing measurements as follows:

• Warm solve + simulate — reuses compiled code

• Cold solve + simulate — first run, including compilation

Publishing Results¶

After running benchmarks on the main branch, publish them to the dashboard:

pixi run asv-publish

This generates the ASV HTML dashboard and pushes results to the OpenSourceEconomics.github.io repo under pylcm-benchmarks/. A persistent clone is kept in .benchmark-site/ (gitignored) to avoid re-cloning on every publish.

Before asv publish, the merged history is normalised so each machine draws one continuous line per benchmark:

• every stored result takes the machine’s current machine.json params (with the RAM in whole gigabytes), so a kernel update or a RAM drift does not split the series;

• every stored version stamp is nulled, so a version bump shows as a step on the same line instead of hiding all earlier results. This gives up ASV’s guard against mixing measurement semantics on one line; read a step at a version bump as a change of workload, not a regression;

• a benchmark that changed from an ASV-native time_* method to a track_* method carries its history over under the new name.

• the stored results of benchmark classes retired from the suite (the ACA GPU-memory profiles) are dropped.

The normalised results are pushed back with the dashboard, so the pass is applied once and repeating it changes nothing.

The dashboard takes its titles and order from the PR comparison table’s labels in benchmarks/pr_comment.py: each benchmark is titled “ — ”, the front-page grid groups by benchmark, and each group starts with its execution time.

CI Check¶

The benchmark-check workflow runs on every pull request. It looks for a PR comment with the <!-- benchmark-check --> marker and checks the embedded commit hash against the PR’s HEAD:

• Passes if the comment’s commit hash matches the PR HEAD

• Passes with warning if a benchmark comment exists but for an older commit

• Fails if no benchmark comment is found

To satisfy the check:

pixi run -e benchmarks-cuda12 asv-run-and-pr-comment

Adding New Benchmarks¶

Create a new bench_*.py file in the benchmarks/ directory. Benchmarks use ASV’s class-based API:

import gc
import statistics
import time

from . import _gpu_mem


class TimeMyModel:
    timeout = 600

    def setup(self):
        # Lazy imports — JAX must not be imported at module level
        import jax.numpy as jnp
        from lcm_examples import my_model

        self.model = my_model.get_model()
        self.model_params = my_model.get_params()
        self.initial_conditions = {
            "wealth": jnp.full(1_000, 5.0),
            "regime_id": jnp.zeros(1_000, dtype=jnp.int32),
        }

        # JIT warmup (timed separately)
        start = time.perf_counter()
        self.model.simulate(
            params=self.model_params,
            initial_conditions=self.initial_conditions,
            log_level="off",
        )
        self._warmup_time = time.perf_counter() - start

    def execute_for_measurement(self):
        self.model.solve(params=self.model_params, log_level="off")

    def track_execution_time(self):
        return statistics.median(
            _gpu_mem.warm_call_seconds(
                execute=self.execute_for_measurement, warm_samples=3
            )
        )

    track_execution_time.unit = "seconds"

    def teardown(self):
        import jax

        jax.clear_caches()
        gc.collect()

    def track_warmup(self):
        return self._warmup_time

    track_warmup.unit = "seconds"

Key points:

• Lazy imports: All imports of JAX and lcm must happen inside setup(), not at module level, to avoid os.fork() conflicts with JAX’s GPU runtime.

• teardown(): Clear JAX caches and run garbage collection between benchmarks.

• track_warmup: Measure JIT compilation time separately from steady-state execution time.

• track_execution_time: Report execution time under this name, as the median of warm calls, rather than as an ASV-native time_* method, so the dashboard and the PR table show every benchmark’s execution time the same way.

• Use the params and param_names class attributes to vary grid sizes or other parameters.