Skip to content

Commit c0c3614

Browse files
ricardoV94ricardoV94
committed
Use asv for benchmarking
1 parent 0bd33bf commit c0c3614

50 files changed

Lines changed: 2220 additions & 189 deletions

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

‎.github/workflows/benchmarks.yml‎

Lines changed: 196 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,196 @@
1+
name: Benchmarks
2+
3+
on:
4+
push:
5+
branches:
6+
- main
7+
- v3
8+
pull_request:
9+
branches:
10+
- main
11+
- v3
12+
13+
concurrency:
14+
group: benchmarks-${{ github.event_name == 'pull_request' && github.head_ref || github.sha }}
15+
cancel-in-progress: true
16+
17+
jobs:
18+
benchmarks:
19+
name: "Run benchmarks"
20+
if: github.event_name == 'push'
21+
runs-on: ubuntu-latest
22+
permissions:
23+
contents: write
24+
pages: write
25+
id-token: write
26+
steps:
27+
- uses: actions/checkout@08c6903cd8c0fde910a37f88322edcfb5dd907a8 # v5.0.0
28+
with:
29+
fetch-depth: 0
30+
persist-credentials: false
31+
32+
- name: Set up Python
33+
uses: mamba-org/setup-micromamba@add3a49764cedee8ee24e82dfde87f5bc2914462 # v2.0.7
34+
with:
35+
environment-name: pytensor-bench
36+
micromamba-version: "1.5.10-0"
37+
init-shell: bash
38+
post-cleanup: "all"
39+
40+
- name: Install dependencies
41+
shell: micromamba-shell {0}
42+
run: |
43+
micromamba install --yes -q -c conda-forge "python=3.11" mkl numpy scipy pip mkl-service cython "numba>=0.57" jax jaxlib asv
44+
pip install -e ./
45+
python -c 'import pytensor; print(pytensor.config.__str__(print_doc=False))'
46+
python -c 'import pytensor; assert pytensor.config.blas__ldflags != "", "Blas flags are empty"'
47+
48+
- name: Fetch previous results from asv-results branch
49+
shell: bash
50+
run: |
51+
git fetch origin asv-results:asv-results 2>/dev/null || true
52+
if git rev-parse --verify asv-results 2>/dev/null; then
53+
git worktree add /tmp/asv-results asv-results
54+
if [ -d /tmp/asv-results/results ]; then
55+
mkdir -p .asv/results
56+
cp -r /tmp/asv-results/results/* .asv/results/
57+
fi
58+
git worktree remove /tmp/asv-results --force
59+
fi
60+
61+
- name: Configure ASV machine
62+
shell: micromamba-shell {0}
63+
run: asv machine --yes
64+
65+
- name: Run benchmarks
66+
shell: micromamba-shell {0}
67+
run: |
68+
export PYTENSOR_FLAGS=warn__ignore_bug_before=all,on_opt_error=raise,on_shape_error=raise,gcc__cxxflags=-pipe
69+
asv run --python=same --set-commit-hash=$(git rev-parse HEAD) --show-stderr
70+
71+
- name: Push results to asv-results branch
72+
shell: bash
73+
run: |
74+
git config user.name "github-actions[bot]"
75+
git config user.email "github-actions[bot]@users.noreply.github.com"
76+
77+
# Create or update the asv-results branch
78+
if git rev-parse --verify asv-results 2>/dev/null; then
79+
git worktree add /tmp/asv-results asv-results
80+
else
81+
git worktree add --orphan -b asv-results /tmp/asv-results
82+
cd /tmp/asv-results
83+
git rm -rf . 2>/dev/null || true
84+
cd -
85+
fi
86+
87+
mkdir -p /tmp/asv-results/results
88+
cp -r .asv/results/* /tmp/asv-results/results/
89+
cd /tmp/asv-results
90+
git add results/
91+
git commit -m "Update benchmark results for ${{ github.sha }}" || true
92+
git push origin asv-results
93+
cd -
94+
git worktree remove /tmp/asv-results --force
95+
96+
- name: Generate HTML
97+
shell: micromamba-shell {0}
98+
run: asv publish
99+
100+
- name: Upload Pages artifact
101+
uses: actions/upload-pages-artifact@56afc609e74202658d3ffba0e8f6dda462b719fa # v3.0.1
102+
with:
103+
path: .asv/html
104+
105+
deploy-pages:
106+
name: "Deploy benchmark dashboard"
107+
if: github.event_name == 'push'
108+
needs: benchmarks
109+
runs-on: ubuntu-latest
110+
permissions:
111+
pages: write
112+
id-token: write
113+
environment:
114+
name: github-pages
115+
url: ${{ steps.deployment.outputs.page_url }}
116+
steps:
117+
- name: Deploy to GitHub Pages
118+
id: deployment
119+
uses: actions/deploy-pages@d6db90164ac5ed86f2b6aed7e0febac5b3c0c03e # v4.0.5
120+
121+
benchmarks-pr:
122+
name: "Benchmark comparison"
123+
if: github.event_name == 'pull_request'
124+
runs-on: ubuntu-latest
125+
permissions:
126+
pull-requests: write
127+
steps:
128+
- uses: actions/checkout@08c6903cd8c0fde910a37f88322edcfb5dd907a8 # v5.0.0
129+
with:
130+
fetch-depth: 0
131+
persist-credentials: false
132+
133+
- name: Set up Python
134+
uses: mamba-org/setup-micromamba@add3a49764cedee8ee24e82dfde87f5bc2914462 # v2.0.7
135+
with:
136+
environment-name: pytensor-bench
137+
micromamba-version: "1.5.10-0"
138+
init-shell: bash
139+
post-cleanup: "all"
140+
141+
- name: Install dependencies
142+
shell: micromamba-shell {0}
143+
run: |
144+
micromamba install --yes -q -c conda-forge "python=3.11" mkl numpy scipy pip mkl-service cython "numba>=0.57" jax jaxlib asv
145+
pip install -e ./
146+
147+
- name: Configure ASV machine
148+
shell: micromamba-shell {0}
149+
run: asv machine --yes
150+
151+
- name: Run benchmark comparison
152+
id: bench
153+
shell: micromamba-shell {0}
154+
env:
155+
BASE_REF: ${{ github.base_ref }}
156+
run: |
157+
export PYTENSOR_FLAGS=warn__ignore_bug_before=all,on_opt_error=raise,on_shape_error=raise,gcc__cxxflags=-pipe
158+
159+
HEAD_SHA=$(git rev-parse HEAD)
160+
BASE_SHA=$(git merge-base "origin/$BASE_REF" HEAD)
161+
162+
# Benchmark the PR head (already installed)
163+
asv run --python=same --set-commit-hash="$HEAD_SHA" --show-stderr
164+
165+
# Checkout base, reinstall, and benchmark
166+
git checkout "$BASE_SHA"
167+
pip install -e ./
168+
asv run --python=same --set-commit-hash="$BASE_SHA" --show-stderr
169+
170+
# Return to PR head so asv.conf.json is available for compare
171+
git checkout "$HEAD_SHA"
172+
173+
# Compare results
174+
asv compare "$BASE_SHA" "$HEAD_SHA" --factor 1.1 | tee bench_output.txt
175+
176+
- name: Post benchmark results as PR comment
177+
if: always()
178+
uses: actions/github-script@f28e40c7f34bde8b3046d885e986cb6290c5673b # v7.1.0
179+
with:
180+
script: |
181+
const fs = require('fs');
182+
let output = '';
183+
try {
184+
output = fs.readFileSync('bench_output.txt', 'utf8');
185+
} catch (e) {
186+
output = 'Benchmark comparison failed to produce results.';
187+
}
188+
if (output.trim()) {
189+
const body = `## Benchmark comparison (main vs PR)\n\n\`\`\`\n${output}\n\`\`\`\n\nBenchmarks that changed by more than 10% are flagged.`;
190+
await github.rest.issues.createComment({
191+
issue_number: context.issue.number,
192+
owner: context.repo.owner,
193+
repo: context.repo.repo,
194+
body: body
195+
});
196+
}

‎.github/workflows/copilot-setup-steps.yml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -40,7 +40,7 @@ jobs:
4040
- name: "Install dependencies"
4141
shell: micromamba-shell {0}
4242
run: |
43-
micromamba install --yes -q -c conda-forge python=3.13 mkl "numpy>=2.0" scipy pip mkl-service graphviz cython pytest coverage pytest-cov pytest-benchmark pytest-mock pytest-sphinx
43+
micromamba install --yes -q -c conda-forge python=3.13 mkl "numpy>=2.0" scipy pip mkl-service graphviz cython pytest coverage pytest-cov pytest-mock pytest-sphinx
4444
micromamba install --yes -q -c conda-forge "numba>=0.57"
4545
micromamba install --yes -q -c conda-forge jax jaxlib numpyro equinox
4646
micromamba install --yes -q -c conda-forge mypy types-setuptools scipy-stubs pandas pre-commit

‎.github/workflows/test.yml‎

Lines changed: 3 additions & 57 deletions
Original file line numberDiff line numberDiff line change
@@ -163,9 +163,9 @@ jobs:
163163
run: |
164164
165165
if [[ $OS == "macos-15" ]]; then
166-
micromamba install --yes -q "python~=${PYTHON_VERSION}" numpy "scipy<1.17.0" "numba>=0.63" pip graphviz cython pytest coverage pytest-cov pytest-benchmark pytest-mock pytest-sphinx libblas=*=*accelerate;
166+
micromamba install --yes -q "python~=${PYTHON_VERSION}" numpy "scipy<1.17.0" "numba>=0.63" pip graphviz cython pytest coverage pytest-cov pytest-mock pytest-sphinx libblas=*=*accelerate;
167167
else
168-
micromamba install --yes -q "python~=${PYTHON_VERSION}" numpy "scipy<1.17.0" "numba>=0.63" pip graphviz cython pytest coverage pytest-cov pytest-benchmark pytest-mock pytest-sphinx mkl mkl-service;
168+
micromamba install --yes -q "python~=${PYTHON_VERSION}" numpy "scipy<1.17.0" "numba>=0.63" pip graphviz cython pytest coverage pytest-cov pytest-mock pytest-sphinx mkl mkl-service;
169169
fi
170170
if [[ $INSTALL_JAX == "1" ]]; then micromamba install --yes -q -c conda-forge "python~=${PYTHON_VERSION}" && pip install "jax>=0.8,<0.9.1" jaxlib numpyro equinox tfp-nightly; fi
171171
if [[ $INSTALL_TORCH == "1" ]]; then micromamba install --yes -q -c conda-forge "python~=${PYTHON_VERSION}" pytorch pytorch-cuda=12.1 "mkl<=2024.0" -c pytorch -c nvidia; fi
@@ -194,7 +194,7 @@ jobs:
194194
if [[ $DEFAULT_MODE == "FAST_COMPILE" ]]; then export PYTENSOR_FLAGS=$PYTENSOR_FLAGS,mode=FAST_COMPILE; fi
195195
if [[ $DEFAULT_MODE == "CVM" ]]; then export PYTENSOR_FLAGS=$PYTENSOR_FLAGS,linker=cvm; fi
196196
export PYTENSOR_FLAGS=$PYTENSOR_FLAGS,warn__ignore_bug_before=all,on_opt_error=raise,on_shape_error=raise,gcc__cxxflags=-pipe
197-
python -m pytest -r A --verbose --runslow --durations=50 --cov=pytensor/ --cov-report=xml:coverage/coverage-${MATRIX_ID}.xml --no-cov-on-fail $PART --benchmark-skip
197+
python -m pytest -r A --verbose --runslow --durations=50 --cov=pytensor/ --cov-report=xml:coverage/coverage-${MATRIX_ID}.xml --no-cov-on-fail $PART
198198
env:
199199
MATRIX_ID: ${{ steps.matrix-id.outputs.id }}
200200
MKL_THREADING_LAYER: GNU
@@ -209,60 +209,6 @@ jobs:
209209
name: coverage-${{ steps.matrix-id.outputs.id }}
210210
path: coverage/coverage-${{ steps.matrix-id.outputs.id }}.xml
211211

212-
benchmarks:
213-
name: "Benchmarks"
214-
needs:
215-
- changes
216-
- style
217-
runs-on: ubuntu-latest
218-
if: ${{ needs.changes.outputs.changes == 'true' && needs.style.result == 'success' }}
219-
strategy:
220-
fail-fast: false
221-
steps:
222-
- uses: actions/checkout@08c6903cd8c0fde910a37f88322edcfb5dd907a8 # v5.0.0
223-
with:
224-
fetch-depth: 0
225-
persist-credentials: false
226-
- name: Set up Python 3.11
227-
uses: mamba-org/setup-micromamba@add3a49764cedee8ee24e82dfde87f5bc2914462 # v2.0.7
228-
with:
229-
environment-name: pytensor-test
230-
micromamba-version: "1.5.10-0" # until https://github.com/mamba-org/setup-micromamba/issues/225 is resolved
231-
init-shell: bash
232-
post-cleanup: "all"
233-
- name: Install dependencies
234-
shell: micromamba-shell {0}
235-
run: |
236-
micromamba install --yes -q -c conda-forge "python~=${PYTHON_VERSION}" mkl numpy scipy pip mkl-service cython pytest "numba>=0.57" jax jaxlib pytest-benchmark
237-
pip install -e ./
238-
micromamba list && pip freeze
239-
python -c 'import pytensor; print(pytensor.config.__str__(print_doc=False))'
240-
python -c 'import pytensor; assert pytensor.config.blas__ldflags != "", "Blas flags are empty"'
241-
env:
242-
PYTHON_VERSION: 3.11
243-
- name: Download previous benchmark data
244-
uses: actions/cache@0057852bfaa89a56745cba8c7296529d2fc39830 # v4.3.0
245-
with:
246-
path: ./cache
247-
key: ${{ runner.os }}-benchmark
248-
- name: Run benchmarks
249-
shell: micromamba-shell {0}
250-
run: |
251-
export PYTENSOR_FLAGS=mode=FAST_COMPILE,warn__ignore_bug_before=all,on_opt_error=raise,on_shape_error=raise,gcc__cxxflags=-pipe
252-
python -m pytest --runslow --benchmark-only --benchmark-json output.json
253-
- name: Store benchmark result
254-
uses: benchmark-action/github-action-benchmark@4bdcce38c94cec68da58d012ac24b7b1155efe8b # v1.20.7
255-
with:
256-
name: Python Benchmark with pytest-benchmark
257-
tool: "pytest"
258-
output-file-path: output.json
259-
external-data-json-path: ./cache/benchmark-data.json
260-
alert-threshold: "200%"
261-
github-token: ${{ secrets.GITHUB_TOKEN }}
262-
comment-on-alert: false
263-
fail-on-alert: false
264-
auto-push: false
265-
266212
all-checks:
267213
if: ${{ always() }}
268214
runs-on: ubuntu-latest

‎.gitignore‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -49,9 +49,11 @@ core
4949
.mypy_cache/
5050
/htmlcov/
5151

52+
.venv/
5253
pytensor-venv/
5354
/notebooks/Sandbox*
5455
.vscode/
5556
testing-report.html
5657
coverage.xml
5758
.coverage.*
59+
.asv/

‎asv.conf.json‎

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,13 @@
1+
{
2+
"version": 1,
3+
"project": "pytensor",
4+
"project_url": "https://github.com/pymc-devs/pytensor",
5+
"repo": ".",
6+
"branches": ["HEAD"],
7+
"dvcs": "git",
8+
"environment_type": "existing",
9+
"benchmark_dir": "benchmarks",
10+
"env_dir": ".asv/env",
11+
"results_dir": ".asv/results",
12+
"html_dir": ".asv/html"
13+
}

‎benchmarks/__init__.py‎

Whitespace-only changes.

‎benchmarks/bench_blockwise.py‎

Lines changed: 68 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,68 @@
1+
import numpy as np
2+
3+
import pytensor
4+
from pytensor import grad
5+
from pytensor.tensor.math import log
6+
from pytensor.tensor.nlinalg import diagonal
7+
from pytensor.tensor.signal.conv import convolve1d
8+
from pytensor.tensor.slinalg import cholesky, solve_triangular
9+
from pytensor.tensor.type import dmatrix, tensor
10+
11+
12+
class BatchedMVNormalLogpAndDlogp:
13+
"""Benchmark batched multivariate normal log-probability and its gradient."""
14+
15+
params = [
16+
[(), (1000,), (4, 1000)],
17+
[(), (1000,), (4, 1000)],
18+
]
19+
param_names = ["mu_batch_shape", "cov_batch_shape"]
20+
21+
def setup(self, mu_batch_shape, cov_batch_shape):
22+
rng = np.random.default_rng(sum(map(ord, "batched_mvnormal")))
23+
24+
value_batch_shape = mu_batch_shape
25+
if len(cov_batch_shape) > len(mu_batch_shape):
26+
value_batch_shape = cov_batch_shape
27+
28+
value = tensor("value", shape=(*value_batch_shape, 10))
29+
mu = tensor("mu", shape=(*mu_batch_shape, 10))
30+
cov = tensor("cov", shape=(*cov_batch_shape, 10, 10))
31+
32+
self.test_values = [
33+
rng.normal(size=value.type.shape),
34+
rng.normal(size=mu.type.shape),
35+
np.eye(cov.type.shape[-1]) * np.abs(rng.normal(size=cov.type.shape)),
36+
]
37+
38+
chol_cov = cholesky(cov, lower=True, on_error="raise")
39+
delta_trans = solve_triangular(chol_cov, value - mu, b_ndim=1)
40+
quaddist = (delta_trans**2).sum(axis=-1)
41+
diag = diagonal(chol_cov, axis1=-2, axis2=-1)
42+
logdet = log(diag).sum(axis=-1)
43+
k = value.shape[-1]
44+
norm = -0.5 * k * (np.log(2 * np.pi))
45+
logp = norm - 0.5 * quaddist - logdet
46+
dlogp = grad(logp.sum(), wrt=[value, mu, cov])
47+
48+
self.fn = pytensor.function([value, mu, cov], [logp, *dlogp])
49+
50+
def time_batched_mvnormal_logp_and_dlogp(self, mu_batch_shape, cov_batch_shape):
51+
self.fn(*self.test_values)
52+
53+
54+
class SmallBlockwisePerformance:
55+
"""Benchmark small blockwise convolution."""
56+
57+
def setup(self):
58+
a = dmatrix(shape=(7, 128))
59+
b = dmatrix(shape=(7, 20))
60+
out = convolve1d(a, b, mode="valid")
61+
self.fn = pytensor.function([a, b], out, trust_input=True)
62+
63+
rng = np.random.default_rng(495)
64+
self.a_test = rng.normal(size=a.type.shape)
65+
self.b_test = rng.normal(size=b.type.shape)
66+
67+
def time_small_blockwise(self):
68+
self.fn(self.a_test, self.b_test)

0 commit comments

Comments
 (0)