From e16603d201a9b1ec5de45565723befc027b3bfed Mon Sep 17 00:00:00 2001 From: Ricardo Vieira Date: Fri, 10 Jul 2026 15:33:05 +0200 Subject: [PATCH] Allow building in free-threaded python --- .github/workflows/test.yml | 32 ++++++++++++++++++++++++++++---- conftest.py | 16 ++++++++++++++++ pyproject.toml | 14 +++++++++++--- 3 files changed, 55 insertions(+), 7 deletions(-) diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index af49b76cef..2b64f31c8f 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -93,6 +93,7 @@ jobs: install-torch: [0] install-mlx: [0] install-xarray: [0] + free-threading: [0] part: - [ "*rest", "tests --ignore=tests/scan --ignore=tests/tensor --ignore=tests/xtensor --ignore=tests/link/numba" ] - [ "scan", "tests/scan" ] @@ -142,6 +143,12 @@ jobs: default-mode: "CVM" python-version: "3.14" os: "macos-15" + # conftest's pytest_sessionfinish asserts the GIL was never re-enabled. + - part: ["free-threading smoke test", "tests/tensor/test_math_scipy.py"] + default-mode: "FAST_RUN" + python-version: "3.14" + os: "ubuntu-latest" + free-threading: 1 steps: - uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # v6.0.3 @@ -155,7 +162,12 @@ jobs: micromamba-version: "1.5.10-0" # until https://github.com/mamba-org/setup-micromamba/issues/225 is resolved init-shell: bash post-cleanup: "all" - create-args: python=${{ matrix.python-version }} + # Create the env with python-freethreading up front so it resolves to the + # cp*t build from the start; adding it later conflicts with an already-solved + # regular CPython. + create-args: >- + python=${{ matrix.python-version }} + ${{ matrix.free-threading == 1 && 'python-freethreading' || '' }} - name: Create matrix id id: matrix-id @@ -171,11 +183,17 @@ jobs: shell: micromamba-shell {0} run: | + # python-freethreading constrains the solve to the free-threaded CPython build + if [[ $FREE_THREADING == "1" ]]; then FT_PKG="python-freethreading"; else FT_PKG=""; fi if [[ $OS == "macos-15" ]]; then - 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 rich; + micromamba install --yes -q "python~=${PYTHON_VERSION}" $FT_PKG numpy "scipy<1.17.0" "numba>=0.63" pip graphviz cython pytest coverage pytest-cov pytest-benchmark pytest-mock pytest-sphinx libblas=*=*accelerate rich; + elif [[ $FREE_THREADING == "1" ]]; then + # mkl/mkl-service have no free-threaded (cp*t) builds, so rely on the default (openblas) BLAS. + micromamba install --yes -q "python~=${PYTHON_VERSION}" $FT_PKG numpy "scipy<1.17.0" "numba>=0.63" pip graphviz cython pytest coverage pytest-cov pytest-benchmark pytest-mock pytest-sphinx rich; else - 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 rich; + micromamba install --yes -q "python~=${PYTHON_VERSION}" $FT_PKG numpy "scipy<1.17.0" "numba>=0.63" pip graphviz cython pytest coverage pytest-cov pytest-benchmark pytest-mock pytest-sphinx mkl mkl-service rich; fi + if [[ $FREE_THREADING == "1" ]]; then python -c 'import sysconfig; assert sysconfig.get_config_var("Py_GIL_DISABLED") == 1, "Expected a free-threaded interpreter"'; fi 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 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 if [[ $INSTALL_MLX == "1" ]]; then micromamba install --yes -q -c conda-forge "python~=${PYTHON_VERSION}" "mlx<0.29.4"; fi @@ -186,7 +204,8 @@ jobs: python -c 'import pytensor; print(pytensor.config.__str__(print_doc=False))' if [[ $OS == "macos-15" ]]; then python -c 'import pytensor; assert pytensor.config.blas__ldflags.startswith("-framework Accelerate"), "Blas flags are not set to MacOS Accelerate"'; - else + elif [[ $FREE_THREADING != "1" ]]; then + # Blas flags are only used by the C backend; the free-threading job runs the numba backend. python -c 'import pytensor; assert pytensor.config.blas__ldflags != "", "Blas flags are empty"'; fi env: @@ -195,6 +214,7 @@ jobs: INSTALL_TORCH: ${{ matrix.install-torch }} INSTALL_XARRAY: ${{ matrix.install-xarray }} INSTALL_MLX: ${{ matrix.install-mlx }} + FREE_THREADING: ${{ matrix.free-threading }} OS: ${{ matrix.os}} - name: Run tests @@ -202,6 +222,9 @@ jobs: run: | if [[ $DEFAULT_MODE == "FAST_COMPILE" ]]; then export PYTENSOR_FLAGS=$PYTENSOR_FLAGS,mode=FAST_COMPILE; fi if [[ $DEFAULT_MODE == "CVM" ]]; then export PYTENSOR_FLAGS=$PYTENSOR_FLAGS,linker=cvm; fi + # Disable the C backend on the free-threaded build so no test imports a C + # extension, which would re-enable the GIL (checked in conftest). + if [[ $FREE_THREADING == "1" ]]; then export PYTENSOR_FLAGS=$PYTENSOR_FLAGS,cxx=; fi export PYTENSOR_FLAGS=$PYTENSOR_FLAGS,on_opt_error=raise,on_shape_error=raise,gcc__cxxflags=-pipe 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 env: @@ -211,6 +234,7 @@ jobs: OMP_NUM_THREADS: 1 PART: ${{ matrix.part[1] }} DEFAULT_MODE: ${{ matrix.default-mode }} + FREE_THREADING: ${{ matrix.free-threading }} - name: Upload coverage file uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 diff --git a/conftest.py b/conftest.py index 14ceca72fd..d02c472ba0 100644 --- a/conftest.py +++ b/conftest.py @@ -1,4 +1,6 @@ import os +import sys +import sysconfig import pytest @@ -17,6 +19,20 @@ def pytest_sessionstart(session): os.environ["NUMBA_BOUNDSCHECK"] = "1" +def pytest_sessionfinish(session, exitstatus): + # On a free-threaded build, importing any single-phase-init C extension + # re-enables the GIL. The default (numba) backend must not, so if the GIL is + # back on, some test pulled in PyTensor's own C extensions. Checked after the + # whole session so it is independent of collection order. + if sysconfig.get_config_var("Py_GIL_DISABLED") == 1 and sys._is_gil_enabled(): + print( # noqa: T201 + "\nERROR: the GIL was re-enabled during the test session. " + "A C-extension was imported on a free-threaded build.", + file=sys.stderr, + ) + session.exitstatus = pytest.ExitCode.TESTS_FAILED + + def pytest_addoption(parser): parser.addoption( "--runslow", action="store_true", default=False, help="run slow tests" diff --git a/pyproject.toml b/pyproject.toml index ea9d0728f8..d1a4890f25 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -33,6 +33,7 @@ classifiers = [ "Programming Language :: Python :: 3.12", "Programming Language :: Python :: 3.13", "Programming Language :: Python :: 3.14", + "Programming Language :: Python :: Free Threading :: 1 - Unstable", ] keywords = [ @@ -182,11 +183,18 @@ files = ["pytensor", "tests"] build = "*" # Uncomment to skip builds that compile but fail when trying to test (maybe due to incompatibility with runner) # archs = ["auto64"] -# Disable any-platform (pp*), 32-bit builds, and free-threaded builds (cp*t-*) +# Disable any-platform (pp*) and 32-bit builds # Additional options to consider: "*musllinux*" -skip = ["pp*", "*-win32", "*-manylinux_i686", "cp*t-*"] +skip = ["pp*", "*-win32", "*-manylinux_i686"] +# Free-threaded (cp*t-*) wheels are built by default. The compiled extensions +# are not marked free-threading-safe, so importing them re-enables the GIL +# (with a warning). That's fine: they're no longer loaded by default. build-frontend = "build" -test-command = 'python -c "import pytensor; print(pytensor.__version__); from pytensor.scan import scan_perform; print(scan_perform.get_version())"' +# Import the compiled extensions and evaluate a function to check the wheel loads and +# runs everywhere, including free-threaded builds. cxx="" disables the C backend (the +# minimal wheel-test env has no libpython to link a runtime-compiled thunk against) and +# FAST_COMPILE avoids numba, which crashes on free-threaded Windows. +test-command = '''python -c "import pytensor; pytensor.config.cxx = ''; import pytensor.tensor as pt; print(pytensor.__version__); from pytensor.scan import scan_perform; print(scan_perform.get_version()); x = pt.vector(); f = pytensor.function([x], (x * 2).sum(), mode='FAST_COMPILE'); assert float(f([1.0, 2.0, 3.0])) == 12.0"''' test-skip = ["*musllinux*", "*i686*"] # Testing seems to be running into issues locating libs where expected