Misc chores

.github/workflows/tests.yml

@@ -43,6 +43,8 @@ env:
 jobs:
   pre-commit:
     runs-on: ubuntu-latest
+    env:
+      SKIP: no-commit-to-branch
     steps:
       - uses: actions/checkout@v6
         with:
@@ -84,11 +86,36 @@ jobs:
       - id: set-matrix
         run: |
           BASE_MATRIX='[
-            {"os":"ubuntu-latest","runs-on":"ubuntu-latest","target-suffix":""},
-            {"os":"ubuntu-22.04","runs-on":"ubuntu-22.04","target-suffix":"-old"}
+            {
+              "id": "latest",
+              "runs-on": "ubuntu-latest",
+              "target-suffix": "",
+              "tests-args": ""
+            },
+            {
+              "id": "old",
+              "runs-on": "ubuntu-22.04",
+              "target-suffix": "-old",
+              "tests-args": ""
+            },
+            {
+              "id": "single-cpu",
+              "runs-on": "ubuntu-latest",
+              "target-suffix": "",
+              "tests-args": "ARGS=--num-cpu-devices=1"
+            }
           ]'
           BENCHMARKS_BASE_MATRIX='[
-            {"os":"ubuntu-latest","runs-on":"ubuntu-latest","target-suffix":""}
+            {
+              "id": "latest",
+              "runs-on": "ubuntu-latest",
+              "vars": ""
+            },
+            {
+              "id": "single-cpu",
+              "runs-on": "ubuntu-latest",
+              "vars": "BARTZ_BENCHMARKS_SINGLE_CPU_DEVICE=1 "
+            }
           ]'
 
           MATRIX=$(echo "$BASE_MATRIX" | jq -c ".")
@@ -98,7 +125,7 @@ jobs:
           echo "benchmarks-matrix={\"include\":$BENCHMARKS_MATRIX}" >> $GITHUB_OUTPUT
 
   tests:
-    name: tests-${{ matrix.os }}${{ matrix.target-suffix }}
+    name: tests-${{ matrix.id }}
     needs: setup-matrix
     runs-on: ${{ matrix.runs-on }}
     continue-on-error: ${{ matrix.target-suffix == '-gpu' }}
@@ -115,7 +142,7 @@ jobs:
         uses: astral-sh/setup-uv@v7
         with:
           version: ${{ env.UV_VERSION }}
-          cache-suffix: -${{ matrix.os }}${{ matrix.target-suffix }}
+          cache-suffix: -${{ matrix.id }}
 
       - name: Install R
         uses: r-lib/actions/setup-r@v2
@@ -134,21 +161,15 @@ jobs:
           ldd "$(R RHOME)/library/methods/libs/methods.so" || true
           uv run --dev python -c "import rpy2.situation as s; print('\n'.join(map(str, s.iter_info())))"
 
-      - name: Install jax GPU support
-        if: matrix.target-suffix == '-gpu'
-        run: |
-          CUDA_VERSION=`nvidia-smi 2>/dev/null | grep -o 'CUDA Version: [0-9]*' | cut -d' ' -f3`
-          uv pip install "jax[cuda$CUDA_VERSION]"
-
       - name: Run unit tests
         timeout-minutes: 30
-        run: make COVERAGE_SUFFIX=-${{ matrix.os }}${{ matrix.target-suffix }} tests${{ matrix.target-suffix }}
+        run: make COVERAGE_SUFFIX=-${{ matrix.id }} tests${{ matrix.target-suffix }} ${{ matrix.tests-args }}
 
       - name: Save coverage information
         uses: actions/upload-artifact@v6
         with:
-          name: coverage.tests-${{ matrix.os }}${{ matrix.target-suffix }}
-          path: .coverage.tests-${{ matrix.os }}${{ matrix.target-suffix }}
+          name: coverage.tests-${{ matrix.id }}
+          path: .coverage.tests-${{ matrix.id }}
           include-hidden-files: true
           if-no-files-found: error
 
@@ -279,7 +300,7 @@ jobs:
         uses: actions/deploy-pages@v4
 
   benchmarks:
-    name: benchmarks-${{ matrix.os }}${{ matrix.target-suffix }}
+    name: benchmarks-${{ matrix.id }}
     needs: setup-matrix
     runs-on: ${{ matrix.runs-on }}
     continue-on-error: ${{ matrix.target-suffix == '-gpu' }}
@@ -296,19 +317,14 @@ jobs:
         uses: astral-sh/setup-uv@v7
         with:
           version: ${{ env.UV_VERSION }}
-          cache-suffix: -${{ matrix.os }}${{ matrix.target-suffix }}
+          cache-suffix: -${{ matrix.id }}
 
       - name: Get machine information for asv
         run: RPY2_CFFI_MODE=ABI uv run --dev python -m asv machine --yes
 
-      - name: Install jax GPU support
-        if: matrix.target-suffix == '-gpu'
-        run: |
-          CUDA_VERSION=`nvidia-smi 2>/dev/null | grep -o 'CUDA Version: [0-9]*' | cut -d' ' -f3`
-          uv pip install "jax[cuda$CUDA_VERSION]"
-
       - name: Run asv in quick mode
-        run: make asv-quick ARGS='--attribute timeout=120' # long timeout because CI runners are slow
+        # long timeout because CI runners are slow
+        run: ${{ matrix.vars }}make asv-quick ARGS='--attribute timeout=120'
 
   dev-commands:
     runs-on: ubuntu-latest
@@ -328,6 +344,9 @@ jobs:
       - name: run IPython with old toolchain
         run: RPY2_CFFI_MODE=ABI make ipython-old ARGS=config/ipython/profile_default/startup/startup.ipy
 
+      - name: Update package version
+        run: make copy-version
+
   required-dummy:
     runs-on: ubuntu-latest
     needs:

.pre-commit-config.yaml

@@ -31,6 +31,7 @@ repos:
         exclude: tests/(rbartpackages/.+\.py|util\.py)$
       - id: trailing-whitespace
         exclude: '^renv/activate\.R$' # because renv edits it automatically, with trailing whitespace
+      - id: no-commit-to-branch
   - repo: https://github.com/sbrunner/hooks
     rev: 1.6.1
     hooks:
@@ -55,4 +56,8 @@ repos:
   - repo: https://github.com/jorisroovers/gitlint
     rev: v0.19.1
     hooks:
-      - id: gitlint # this is already defined on stage commit-msg
+      - id: gitlint # this is already defined to run in stage commit-msg
+  - repo: https://github.com/astral-sh/uv-pre-commit
+    rev: 0.10.4
+    hooks:
+      - id: uv-lock

Makefile

@@ -35,8 +35,8 @@ UV_RUN = uv run --dev $(EXTRAS)
 OLD_PYTHON = $(shell grep 'requires-python' pyproject.toml | sed 's/.*>=\([0-9.]*\).*/\1/')
 UV_RUN_OLD = $(UV_RUN) --python=$(OLD_PYTHON) --resolution=lowest-direct --exclude-newer=2025-05-15 --isolated
 
-.PHONY: all
-all:
+.PHONY: help
+help:
 	@echo "Available targets:"
 	@echo "- setup: create R and Python environments for development"
 	@echo "- tests: run unit tests on cpu, saving coverage information"
@@ -61,8 +61,10 @@ all:
 	@echo "- asv-quick: run quick benchmarks on current code, no saving"
 	@echo "- ipython: start an ipython shell with stuff pre-imported"
 	@echo "- ipython-old: start an ipython shell with oldest supported python and dependencies"
+	@echo "- lint: run the linter used in pre-commit"
 	@echo
 	@echo "Release workflow:"
+	@echo "- do a PR that re-runs benchmarks"
 	@echo "- create a new branch"
 	@echo "- $$ uv version --bump major|minor|patch"
 	@echo "- describe release in docs/changelog.md"
@@ -71,7 +73,7 @@ all:
 	@echo "- $$ make release (iterate to fix problems)"
 	@echo "- if CI does not pass, debug and go back to make release"
 	@echo "- merge PR"
-	@echo "- if CI does not pass, debug and go back to make release"
+	@echo "- if CI does not pass, debug and go back to open PR"
 	@echo "- $$ make upload"
 	@echo "- publish github release (updates zenodo automatically)"
 	@echo "- if the online docs are not up-to-date, merge another PR to trigger a new merge CI"
@@ -82,14 +84,17 @@ setup:
 	Rscript -e "renv::restore()"
 	$(UV_RUN) pre-commit install --install-hooks
 
+.PHONY: lint
+lint:
+	$(UV_RUN) pre-commit run --all-files ruff-check
 
 ################# TESTS #################
 
 TESTS_VARS = COVERAGE_FILE=.coverage.tests$(COVERAGE_SUFFIX)
 TESTS_COMMAND = python -m pytest --cov --cov-context=test --dist=worksteal --durations=1000
 TESTS_CPU_VARS = $(TESTS_VARS) JAX_PLATFORMS=cpu
 TESTS_CPU_COMMAND = $(TESTS_COMMAND) --platform=cpu --numprocesses=2
-TESTS_GPU_VARS = $(TESTS_VARS) XLA_PYTHON_CLIENT_MEM_FRACTION=.20
+TESTS_GPU_VARS = $(TESTS_VARS) XLA_PYTHON_CLIENT_PREALLOCATE=false
 TESTS_GPU_COMMAND = $(TESTS_COMMAND) --platform=gpu --numprocesses=3
 
 .PHONY: tests
@@ -129,7 +134,6 @@ docs-latest:
 	WORKTREE_DIR=$$(mktemp -d) && \
 	trap "git worktree remove --force '$$WORKTREE_DIR' 2>/dev/null || rm -rf '$$WORKTREE_DIR'" EXIT && \
 	git worktree add --detach "$$WORKTREE_DIR" "$$LATEST_TAG" && \
-	uv sync --all-groups --directory "$$WORKTREE_DIR" && \
 	$(MAKE) -C "$$WORKTREE_DIR" docs && \
 	test ! -d _site/docs || rm -r _site/docs && \
 	mv "$$WORKTREE_DIR/_site/docs-dev" _site/docs
@@ -172,6 +176,7 @@ clean:
 	rm -fr .venv
 	rm -fr dist
 	rm -fr config/jax_cache
+	rm -fr docs/_build
 
 .PHONY: release
 release: clean update-deps copy-version check-committed tests tests-old docs
@@ -180,11 +185,17 @@ release: clean update-deps copy-version check-committed tests tests-old docs
 .PHONY: version-tag
 version-tag: copy-version check-committed
 	git fetch --tags
-	git tag v$(shell uv run python -c 'import bartz; print(bartz.__version__)')
-	git push --tags
+	$(eval VERSION_TAG := v$(shell uv run python -c 'import bartz; print(bartz.__version__)'))
+	git tag $(VERSION_TAG)
+	git push origin $(VERSION_TAG)
+
+.PHONY: smoke-test
+smoke-test:
+	uv run --isolated --no-project --with dist/*.whl python -c 'import bartz'
+	uv run --isolated --no-project --with dist/*.tar.gz python -c 'import bartz'
 
 .PHONY: upload
-upload: version-tag
+upload: smoke-test version-tag
 	@echo "Enter PyPI token:"
 	@read -s UV_PUBLISH_TOKEN && \
 	export UV_PUBLISH_TOKEN="$$UV_PUBLISH_TOKEN" && \
@@ -194,7 +205,7 @@ upload: version-tag
 	uv tool run --with="bartz==$$VERSION" python -c 'import bartz; print(bartz.__version__)'
 
 .PHONY: upload-test
-upload-test: check-committed
+upload-test: smoke-test check-committed
 	@echo "Enter TestPyPI token:"
 	@read -s UV_PUBLISH_TOKEN && \
 	export UV_PUBLISH_TOKEN="$$UV_PUBLISH_TOKEN" && \

benchmarks/__init__.py

@@ -24,10 +24,13 @@
 
 """Benchmarking code run by asv."""
 
+from os import getenv
+
 from jax import config
 
 from benchmarks._vendor_latest_bartz import ensure_vendored
 
 ensure_vendored()
 
-config.update('jax_num_cpu_devices', 16)
+if getenv('BARTZ_BENCHMARKS_SINGLE_CPU_DEVICE') is None:
+    config.update('jax_num_cpu_devices', 16)

benchmarks/rmse.py

@@ -99,7 +99,7 @@ def run_sim_impl(
     # adapt for older versions
     sig = signature(mc_gbart)
 
-    def drop_if_missing(arg: str):
+    def drop_if_missing(arg: str) -> None:
         if arg not in sig.parameters:
             kw.pop(arg)