Skip to content

Commit 31b5aa3

Browse files
ricardoV94ricardoV94
committed
Use asv for benchmarking
1 parent 1243d58 commit 31b5aa3

50 files changed

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

‎.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: 2 additions & 56 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
@@ -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": ["main"],
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)