diff --git a/.azure_pipeline.yml b/.azure_pipeline.yml deleted file mode 100644 index 15daa590..00000000 --- a/.azure_pipeline.yml +++ /dev/null @@ -1,46 +0,0 @@ -# Adapted from https://github.com/pandas-dev/pandas/blob/master/azure-pipelines.yml -jobs: - -- template: continuous_integration/windows.yml - parameters: - name: Windows - vmImage: vs2017-win2016 - matrix: - py37_conda: - VERSION_PYTHON: '3.7' - PACKAGER: 'conda' - py36_conda: - VERSION_PYTHON: '3.6' - PACKAGER: 'conda' - py35_pip: - VERSION_PYTHON: '3.5' - PACKAGER: 'pip' - - -- template: continuous_integration/posix.yml - parameters: - name: Linux - vmImage: ubuntu-16.04 - matrix: - # Linux environment to test that packages that comes with Ubuntu Xenial - # 16.04 are correctly handled. - py35_np_atlas: - PACKAGER: 'ubuntu' - VERSION_PYTHON: '3.5' - # Linux + Python 3.6 - py36_conda_openblas: - PACKAGER: 'conda' - VERSION_PYTHON: '3.6' - # Linux environment to test the latest available dependencies and MKL. - pylatest_conda: - PACKAGER: 'conda' - VERSION_PYTHON: '*' - -- template: continuous_integration/posix.yml - parameters: - name: macOS - vmImage: xcode9-macos10.13 - matrix: - pylatest_conda: - PACKAGER: 'conda' - VERSION_PYTHON: '*' \ No newline at end of file diff --git a/.codecov.yml b/.codecov.yml new file mode 100644 index 00000000..50979600 --- /dev/null +++ b/.codecov.yml @@ -0,0 +1,2 @@ +comment: off + diff --git a/.coveragerc b/.coveragerc new file mode 100644 index 00000000..f25962cc --- /dev/null +++ b/.coveragerc @@ -0,0 +1,3 @@ +[run] +source=threadpoolctl + diff --git a/.github/workflows/publish-to-pypi.yml b/.github/workflows/publish-to-pypi.yml new file mode 100644 index 00000000..b38ed7d7 --- /dev/null +++ b/.github/workflows/publish-to-pypi.yml @@ -0,0 +1,128 @@ +name: Publish threadpoolctl 🎮 distribution 📦 to PyPI and TestPyPI +# Taken from: +# https://packaging.python.org/en/latest/guides/publishing-package-distribution-releases-using-github-actions-ci-cd-workflows/ + +on: push + +jobs: + build: + name: Build distribution 📦 + # Don't run on forked repositories + if: github.event.repository.fork != true + runs-on: ubuntu-latest + + steps: + - uses: actions/checkout@v4 + with: + persist-credentials: false + + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: "3.x" + + - name: Install pypa/build + run: python -m pip install build --user + + - name: Build a binary wheel and a source tarball + run: python -m build + + - name: Store the distribution packages + uses: actions/upload-artifact@v4 + with: + name: python-package-distributions + path: dist/ + retention-days: 1 + + publish-to-pypi: + name: >- + Publish threadpoolctl 🎮 distribution 📦 to PyPI + if: startsWith(github.ref, 'refs/tags/') # only publish to PyPI on tag pushes + needs: + - build + runs-on: ubuntu-latest + environment: + name: pypi + url: https://pypi.org/project/threadpoolctl/ + permissions: + id-token: write # IMPORTANT: mandatory for trusted publishing + + steps: + - name: Download all the dists + uses: actions/download-artifact@v4 + with: + name: python-package-distributions + path: dist/ + - name: Publish distribution 📦 to PyPI + uses: pypa/gh-action-pypi-publish@release/v1 + + github-release: + name: >- + Sign the threadpoolctl 🎮 distribution 📦 with Sigstore + and upload them to GitHub Release + needs: + - publish-to-pypi + runs-on: ubuntu-latest + + permissions: + contents: write # IMPORTANT: mandatory for making GitHub Releases + id-token: write # IMPORTANT: mandatory for sigstore + + steps: + - name: Download all the dists + uses: actions/download-artifact@v4 + with: + name: python-package-distributions + path: dist/ + - name: Sign the dists with Sigstore + uses: sigstore/gh-action-sigstore-python@v3.0.0 + with: + inputs: >- + ./dist/*.tar.gz + ./dist/*.whl + - name: Create GitHub Release + env: + GITHUB_TOKEN: ${{ github.token }} + run: >- + gh release create + "$GITHUB_REF_NAME" + --repo "$GITHUB_REPOSITORY" + --notes "" + - name: Upload artifact signatures to GitHub Release + env: + GITHUB_TOKEN: ${{ github.token }} + # Upload to GitHub Release using the `gh` CLI. + # `dist/` contains the built packages, and the + # sigstore-produced signatures and certificates. + run: >- + gh release upload + "$GITHUB_REF_NAME" dist/** + --repo "$GITHUB_REPOSITORY" + + publish-to-testpypi: + name: Publish threadpoolctl 🎮 distribution 📦 to TestPyPI + needs: + - build + runs-on: ubuntu-latest + + environment: + name: testpypi + url: https://test.pypi.org/project/threadpoolctl/ + + permissions: + id-token: write # IMPORTANT: mandatory for trusted publishing + + steps: + - name: Download all the dists + uses: actions/download-artifact@v4 + with: + name: python-package-distributions + path: dist/ + - name: Publish distribution 📦 to TestPyPI + uses: pypa/gh-action-pypi-publish@release/v1 + with: + verbose: true + # skip-existing is required for .dev0 versions with fixed names but + # different contents. See also: https://github.com/pypa/flit/issues/257 + skip-existing: true + repository-url: https://test.pypi.org/legacy/ diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml new file mode 100644 index 00000000..39fc19e2 --- /dev/null +++ b/.github/workflows/test.yml @@ -0,0 +1,274 @@ +name: CI +permissions: + contents: read + +on: + push: + branches: + - master + pull_request: + branches: + - master + schedule: + # Daily build at 1:00 AM UTC + - cron: "0 1 * * *" + +# Cancel in-progress workflows when pushing +# a new commit on the same branch +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: true + +jobs: + + linting: + runs-on: ubuntu-latest + + steps: + - uses: actions/checkout@v3 + - name: Set up Python + uses: actions/setup-python@v4 + with: + python-version: "3.12" + - name: Install black + run: | + pip install black==25.1.0 + - name: Run black + run: | + black --check --diff . + + testing: + needs: linting + timeout-minutes: 30 + strategy: + fail-fast: false + matrix: + include: + + # Windows env with numpy, scipy,MKL installed through conda-forge + - name: pylatest_conda_forge_mkl + os: windows-latest + PYTHON_VERSION: "*" + PACKAGER: "conda-forge" + BLAS: "mkl" + # Windows env with numpy, scipy, OpenBLAS installed through conda-forge + - name: py311_conda_forge_openblas + os: windows-latest + PYTHON_VERSION: "3.11" + PACKAGER: "conda-forge" + BLAS: "openblas" + # Windows env with numpy, scipy installed through conda + - name: py310_conda + os: windows-latest + PYTHON_VERSION: "3.10" + PACKAGER: "conda" + # Windows env with numpy, scipy installed through pip + - name: py39_pip + os: windows-latest + PYTHON_VERSION: "3.9" + PACKAGER: "pip" + + # MacOS env with OpenMP installed through homebrew + - name: py39_conda_homebrew_libomp + os: macos-latest + PYTHON_VERSION: "3.9" + PACKAGER: "conda" + BLAS: "openblas" + CC_OUTER_LOOP: "clang" + CC_INNER_LOOP: "clang" + INSTALL_LIBOMP: "homebrew" + # MacOS env with OpenBLAS and OpenMP installed through conda-forge compilers + - name: pylatest_conda_forge_clang_openblas + os: macos-latest + PYTHON_VERSION: "*" + PACKAGER: "conda-forge" + BLAS: "openblas" + INSTALL_LIBOMP: "conda-forge" + # MacOS env with FlexiBLAS + - name: pylatest_flexiblas + os: macos-latest + PYTHON_VERSION: "*" + INSTALL_BLAS: "flexiblas" + PLATFORM_SPECIFIC_PACKAGES: "llvm-openmp" + + # Linux environments to test that packages that comes with Ubuntu 22.04 + # are correctly handled. + - name: py39_ubuntu_atlas_gcc_gcc + os: ubuntu-22.04 + PYTHON_VERSION: "3.9" + PACKAGER: "ubuntu" + APT_BLAS: "libatlas3-base libatlas-base-dev" + CC_OUTER_LOOP: "gcc" + CC_INNER_LOOP: "gcc" + - name: py39_ubuntu_openblas_gcc_gcc + os: ubuntu-22.04 + PYTHON_VERSION: "3.9" + PACKAGER: "ubuntu" + APT_BLAS: "libopenblas-base libopenblas-dev" + CC_OUTER_LOOP: "gcc" + CC_INNER_LOOP: "gcc" + + # Linux environment with development versions of numpy and scipy + - name: pylatest_pip_dev + os : ubuntu-latest + PACKAGER: "pip-dev" + PYTHON_VERSION: "*" + CC_OUTER_LOOP: "gcc" + CC_INNER_LOOP: "gcc" + # Linux + Python 3.9 and homogeneous runtime nesting. + - name: py39_conda_openblas_clang_clang + os: ubuntu-latest + PACKAGER: "conda" + PYTHON_VERSION: "3.9" + BLAS: "openblas" + CC_OUTER_LOOP: "clang-18" + CC_INNER_LOOP: "clang-18" + # Linux environment with MKL and Clang (known to be unsafe for + # threadpoolctl) to only test the warning from multiple OpenMP. + - name: pylatest_conda_mkl_clang_gcc + os: ubuntu-latest + PYTHON_VERSION: "*" + PACKAGER: "conda" + BLAS: "mkl" + CC_OUTER_LOOP: "clang-18" + CC_INNER_LOOP: "gcc" + TESTS: "libomp_libiomp_warning" + # Linux environment with MKL, safe for threadpoolctl. + - name: pylatest_conda_mkl_gcc_gcc + os: ubuntu-latest + PYTHON_VERSION: "*" + PACKAGER: "conda" + BLAS: "mkl" + CC_OUTER_LOOP: "gcc" + CC_INNER_LOOP: "gcc" + MKL_THREADING_LAYER: "INTEL" + # Linux + Python 3.11 with numpy / scipy installed with pip from PyPI + # and heterogeneous OpenMP runtimes. + - name: py311_pip_openblas_gcc_clang + os: ubuntu-latest + PACKAGER: "pip" + PYTHON_VERSION: "3.11" + CC_OUTER_LOOP: "gcc" + CC_INNER_LOOP: "clang-18" + # Linux environment with numpy from conda-forge channel and openblas-openmp + - name: pylatest_conda_forge + os: ubuntu-latest + PACKAGER: "conda-forge" + PYTHON_VERSION: "*" + BLAS: "openblas" + OPENBLAS_THREADING_LAYER: "openmp" + CC_OUTER_LOOP: "gcc" + CC_INNER_LOOP: "gcc" + # Linux environment with no numpy and heterogeneous OpenMP runtimes. + - name: pylatest_conda_nonumpy_gcc_clang + os: ubuntu-latest + PACKAGER: "conda" + PYTHON_VERSION: "*" + NO_NUMPY: "true" + CC_OUTER_LOOP: "gcc" + CC_INNER_LOOP: "clang-18" + + # Linux environments with numpy linked to BLIS + - name: pylatest_blis_gcc_clang_openmp + os: ubuntu-latest + PYTHON_VERSION: "*" + INSTALL_BLAS: "blis" + BLIS_NUM_THREAEDS: "4" + CC_OUTER_LOOP: "gcc" + CC_INNER_LOOP: "gcc" + BLIS_CC: "clang-18" + BLIS_ENABLE_THREADING: "openmp" + - name: pylatest_blis_clang_gcc_pthreads + os: ubuntu-latest + PYTHON_VERSION: "*" + INSTALL_BLAS: "blis" + BLIS_NUM_THREADS: "4" + CC_OUTER_LOOP: "clang-18" + CC_INNER_LOOP: "clang-18" + BLIS_CC: "gcc-12" + BLIS_ENABLE_THREADING: "pthreads" + - name: pylatest_blis_no_threading + os: ubuntu-latest + PYTHON_VERSION: "*" + INSTALL_BLAS: "blis" + BLIS_NUM_THREADS: "1" + CC_OUTER_LOOP: "gcc" + CC_INNER_LOOP: "gcc" + BLIS_CC: "gcc-12" + BLIS_ENABLE_THREADING: "no" + + # Linux env with FlexiBLAS + - name: pylatest_flexiblas + os: ubuntu-latest + PYTHON_VERSION: "*" + INSTALL_BLAS: "flexiblas" + PLATFORM_SPECIFIC_PACKAGES: "mkl" + CC_OUTER_LOOP: "gcc" + CC_INNER_LOOP: "gcc" + + env: ${{ matrix }} + + runs-on: ${{ matrix.os }} + + defaults: + run: + # Need to use this shell to get conda working properly. + # See https://github.com/marketplace/actions/setup-miniconda#important + shell: ${{ matrix.os == 'windows-latest' && 'cmd /C CALL {0}' || 'bash -el {0}' }} + + + steps: + - name: Checkout code + uses: actions/checkout@v3 + + - name: Setup conda + uses: conda-incubator/setup-miniconda@v3 + with: + auto-activate-base: true + auto-update-conda: true + miniforge-version: latest + + - name: Install dependencies + run: | + bash -el continuous_integration/install.sh + + - name: Test library + run: | + bash -el continuous_integration/run_tests.sh + + - name: Upload test results + uses: actions/upload-artifact@v4 + with: + # Requires a unique name for each job in the matrix of this run + name: test_result_${{github.run_id}}_${{ matrix.os }}_${{ matrix.name }} + path: test_result.xml + retention-days: 1 + + - name: Upload to Codecov + uses: codecov/codecov-action@v5 + with: + files: coverage.xml + + # Meta-test to ensure that at least one of the above CI configurations had + # the necessary platform settings to execute each test without raising + # skipping. + meta_test: + needs: testing + runs-on: ubuntu-latest + steps: + - name: Set up Python + uses: actions/setup-python@v4 + with: + python-version: "3.12" + + - name: Checkout code + uses: actions/checkout@v3 + + - name: Download tests results + uses: actions/download-artifact@v4 + with: + path: test_results + + - name: Check no test always skipped + run: | + python continuous_integration/check_no_test_skipped.py test_results diff --git a/.gitignore b/.gitignore new file mode 100644 index 00000000..387f3b4c --- /dev/null +++ b/.gitignore @@ -0,0 +1,26 @@ +# Python generated files +*.pyc +__pycache__ +.cache +.pytest_cache + +# Cython/C generated files +*.o +*.so +*.dylib +tests/_openmp_test_helper/*.c + +# Python install files, build and release artifacts +*.egg-info/ +build +dist + +# Coverage data +.coverage +/htmlcov + +# Developer tools +.vscode + +# pytest +.pytest_cache diff --git a/CHANGES.md b/CHANGES.md new file mode 100644 index 00000000..7040d695 --- /dev/null +++ b/CHANGES.md @@ -0,0 +1,158 @@ +3.6.0 (2025-03-13) +================== + +- Added support for libraries with a path longer than 260 on Windows. The supported path + length is now 10 times higher but not unlimited for security reasons. + https://github.com/joblib/threadpoolctl/pull/189 + +- Dropped official support for Python 3.8. + https://github.com/joblib/threadpoolctl/pull/186 + https://github.com/joblib/threadpoolctl/pull/191 + +3.5.0 (2024-04-29) +================== + +- Added support for the Scientific Python version of OpenBLAS + (https://github.com/MacPython/openblas-libs), which exposes symbols with different + names than the ones of the original OpenBLAS library. + https://github.com/joblib/threadpoolctl/pull/175 + +3.4.0 (2024-03-20) +================== + +- Added support for Python interpreters statically linked against libc or linked against + alternative implementations of libc like musl (on Alpine Linux for instance). + https://github.com/joblib/threadpoolctl/pull/171 + +- Added support for Pyodide + https://github.com/joblib/threadpoolctl/pull/169 + +3.3.0 (2024-02-14) +================== + +- Extended FlexiBLAS support to be able to switch backend at runtime. + https://github.com/joblib/threadpoolctl/pull/163 + +- Added support for FlexiBLAS + https://github.com/joblib/threadpoolctl/pull/156 + +- Fixed a bug where an unsupported library would be detected because it shares a common + prefix with one of the supported libraries. Now the symbols are also checked to + identify the supported libraries. + https://github.com/joblib/threadpoolctl/pull/151 + +3.2.0 (2023-07-13) +================== + +- Dropped support for Python 3.6 and 3.7. + +- Added support for custom library controllers. Custom controllers must inherit from + the `threadpoolctl.LibController` class and be registered to threadpoolctl using the + `threadpoolctl.register` function. + https://github.com/joblib/threadpoolctl/pull/138 + +- A warning is raised on macOS when threadpoolctl finds both Intel OpenMP and LLVM + OpenMP runtimes loaded simultaneously by the same Python program. See details and + workarounds at https://github.com/joblib/threadpoolctl/blob/master/multiple_openmp.md. + https://github.com/joblib/threadpoolctl/pull/142 + +3.1.0 (2022-01-31) +================== + +- Fixed a detection issue of the BLAS libraires packaged by conda-forge on Windows. + https://github.com/joblib/threadpoolctl/pull/112 + +- `threadpool_limits` and `ThreadpoolController.limit` now accept the string + "sequential_blas_under_openmp" for the `limits` parameter. It should only be used for + the specific case when one wants to have sequential BLAS calls within an OpenMP + parallel region. It takes into account the unexpected behavior of OpenBLAS with the + OpenMP threading layer. + https://github.com/joblib/threadpoolctl/pull/114 + +3.0.0 (2021-10-01) +================== + +- New object `threadpooctl.ThreadpoolController` which holds controllers for all the + supported native libraries. The states of these libraries is accessible through the + `info` method (equivalent to `threadpoolctl.threadpool_info()`) and their number of + threads can be limited with the `limit` method which can be used as a context + manager (equivalent to `threadpoolctl.threadpool_limits()`). This is especially useful + to avoid searching through all loaded shared libraries each time. + https://github.com/joblib/threadpoolctl/pull/95 + +- Added support for OpenBLAS built for 64bit integers in Fortran. + https://github.com/joblib/threadpoolctl/pull/101 + +- Added the possibility to use `threadpoolctl.threadpool_limits` and + `threadpooctl.ThreadpoolController` as decorators through their `wrap` method. + https://github.com/joblib/threadpoolctl/pull/102 + +- Fixed an attribute error when using old versions of OpenBLAS or BLIS that are + missing version query functions. + https://github.com/joblib/threadpoolctl/pull/88 + https://github.com/joblib/threadpoolctl/pull/91 + +- Fixed an attribute error when python is run with -OO. + https://github.com/joblib/threadpoolctl/pull/87 + +2.2.0 (2021-07-09) +================== + +- `threadpoolctl.threadpool_info()` now reports the architecture of the CPU + cores detected by OpenBLAS (via `openblas_get_corename`) and BLIS (via + `bli_arch_query_id` and `bli_arch_string`). + +- Fixed a bug when the version of MKL was not found. The + "version" field is now set to None in that case. + https://github.com/joblib/threadpoolctl/pull/82 + +2.1.0 (2020-05-29) +================== + +- New commandline interface: + + python -m threadpoolctl -i numpy + + will try to import the `numpy` package and then return the output of + `threadpoolctl.threadpool_info()` on STDOUT formatted using the JSON + syntax. This makes it easier to quickly introspect a Python environment. + + +2.0.0 (2019-12-05) +================== + +- Expose MKL, BLIS and OpenBLAS threading layer in information displayed by + `threadpool_info`. This information is referenced in the `threading_layer` + field. + https://github.com/joblib/threadpoolctl/pull/48 + https://github.com/joblib/threadpoolctl/pull/60 + +- When threadpoolctl finds libomp (LLVM OpenMP) and libiomp (Intel OpenMP) + both loaded, a warning is raised to recall that using threadpoolctl with + this mix of OpenMP libraries may cause crashes or deadlocks. + https://github.com/joblib/threadpoolctl/pull/49 + +1.1.0 (2019-09-12) +================== + +- Detect libraries referenced by symlinks (e.g. BLAS libraries from + conda-forge). + https://github.com/joblib/threadpoolctl/pull/34 + +- Add support for BLIS. + https://github.com/joblib/threadpoolctl/pull/23 + +- Breaking change: method `get_original_num_threads` on the `threadpool_limits` + context manager to cheaply access the initial state of the runtime: + - drop the `user_api` parameter; + - instead return a dict `{user_api: num_threads}`; + - fixed a bug when the limit parameter of `threadpool_limits` was set to + `None`. + + https://github.com/joblib/threadpoolctl/pull/32 + + +1.0.0 (2019-06-03) +================== + +Initial release. diff --git a/LICENSE b/LICENSE new file mode 100644 index 00000000..f2927f5f --- /dev/null +++ b/LICENSE @@ -0,0 +1,24 @@ +Copyright (c) 2019, threadpoolctl contributors + +Redistribution and use in source and binary forms, with or without +modification, are permitted provided that the following conditions are met: + + * Redistributions of source code must retain the above copyright notice, + this list of conditions and the following disclaimer. + * Redistributions in binary form must reproduce the above copyright + notice, this list of conditions and the following disclaimer in the + documentation and/or other materials provided with the distribution. + * Neither the name of copyright holder nor the names of its contributors + may be used to endorse or promote products derived from this software + without specific prior written permission. + +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE +FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. \ No newline at end of file diff --git a/README.md b/README.md index e54e3755..65fba02c 100644 --- a/README.md +++ b/README.md @@ -1,25 +1,404 @@ -# Thread-pool Controls [![Build Status](https://dev.azure.com/karementmoi/loky/_apis/build/status/joblib.threadpoolctl?branchName=master)](https://dev.azure.com/karementmoi/loky/_build/latest?definitionId=2&branchName=master) [![codecov](https://codecov.io/gh/joblib/threadpoolctl/branch/master/graph/badge.svg)](https://codecov.io/gh/joblib/threadpoolctl) +# Thread-pool Controls [![Build Status](https://github.com/joblib/threadpoolctl/actions/workflows/test.yml/badge.svg?branch=master)](https://github.com/joblib/threadpoolctl/actions?query=branch%3Amaster) [![codecov](https://codecov.io/gh/joblib/threadpoolctl/branch/master/graph/badge.svg)](https://codecov.io/gh/joblib/threadpoolctl) +Python helpers to limit the number of threads used in the +threadpool-backed of common native libraries used for scientific +computing and data science (e.g. BLAS and OpenMP). +Fine control of the underlying thread-pool size can be useful in +workloads that involve nested parallelism so as to mitigate +oversubscription issues. -Python helpers to limit the number of threads used in threadpool-backed parallelism for C-libraries. +> **Important:** In its current state, `threadpoolctl` is only designed for +> situations where BLAS and OpenMP are only called from the main Python thread. +> Or, to be more accurate, `threadpoolctl` and BLAS/OpenMP APIs should only ever +> called from the same, single Python thread. For example: +> +> * When you're using it to configure a worker in a process pool, which then calls BLAS or OpenMP APIs directly in the main thread. +> * A Jupyter notebook, where the BLAS or OpenMP APIs are being called from code running in the cell's main thread. +> +> However, once you start calling BLAS or OpenMP APIs and `threadpoolctl` from +> multiple different Python threads, the impact of the `threadpoolctl` limiting +> APIs will be very inconsistent. For more details and a plan to fix this, see +> https://github.com/joblib/threadpoolctl/issues/208 +## Installation -# Installation +- For users, install the last published version from PyPI: + ```bash + pip install threadpoolctl + ``` + +- For contributors, install from the source repository in developer + mode: + + ```bash + pip install -r dev-requirements.txt + flit install --symlink + ``` + + then you run the tests with pytest: + + ```bash + pytest + ``` + +## Usage + +### Command Line Interface + +Get a JSON description of thread-pools initialized when importing python +packages such as numpy or scipy for instance: + +``` +python -m threadpoolctl -i numpy scipy.linalg +[ + { + "filepath": "/home/ogrisel/miniconda3/envs/tmp/lib/libmkl_rt.so", + "prefix": "libmkl_rt", + "user_api": "blas", + "internal_api": "mkl", + "version": "2019.0.4", + "num_threads": 2, + "threading_layer": "intel" + }, + { + "filepath": "/home/ogrisel/miniconda3/envs/tmp/lib/libiomp5.so", + "prefix": "libiomp", + "user_api": "openmp", + "internal_api": "openmp", + "version": null, + "num_threads": 4 + } +] +``` + +The JSON information is written on STDOUT. If some of the packages are missing, +a warning message is displayed on STDERR. + +### Python Runtime Programmatic Introspection + +Introspect the current state of the threadpool-enabled runtime libraries +that are loaded when importing Python packages: + +```python +>>> from threadpoolctl import threadpool_info +>>> from pprint import pprint +>>> pprint(threadpool_info()) +[] + +>>> import numpy +>>> pprint(threadpool_info()) +[{'filepath': '/home/ogrisel/miniconda3/envs/tmp/lib/libmkl_rt.so', + 'internal_api': 'mkl', + 'num_threads': 2, + 'prefix': 'libmkl_rt', + 'threading_layer': 'intel', + 'user_api': 'blas', + 'version': '2019.0.4'}, + {'filepath': '/home/ogrisel/miniconda3/envs/tmp/lib/libiomp5.so', + 'internal_api': 'openmp', + 'num_threads': 4, + 'prefix': 'libiomp', + 'user_api': 'openmp', + 'version': None}] + +>>> import xgboost +>>> pprint(threadpool_info()) +[{'filepath': '/home/ogrisel/miniconda3/envs/tmp/lib/libmkl_rt.so', + 'internal_api': 'mkl', + 'num_threads': 2, + 'prefix': 'libmkl_rt', + 'threading_layer': 'intel', + 'user_api': 'blas', + 'version': '2019.0.4'}, + {'filepath': '/home/ogrisel/miniconda3/envs/tmp/lib/libiomp5.so', + 'internal_api': 'openmp', + 'num_threads': 4, + 'prefix': 'libiomp', + 'user_api': 'openmp', + 'version': None}, + {'filepath': '/home/ogrisel/miniconda3/envs/tmp/lib/libgomp.so.1.0.0', + 'internal_api': 'openmp', + 'num_threads': 4, + 'prefix': 'libgomp', + 'user_api': 'openmp', + 'version': None}] +``` + +In the above example, `numpy` was installed from the default anaconda channel and comes +with MKL and its Intel OpenMP (`libiomp5`) implementation while `xgboost` was installed +from pypi.org and links against GNU OpenMP (`libgomp`) so both OpenMP runtimes are +loaded in the same Python program. + +The state of these libraries is also accessible through the object oriented API: + +```python +>>> from threadpoolctl import ThreadpoolController, threadpool_info +>>> from pprint import pprint +>>> import numpy +>>> controller = ThreadpoolController() +>>> pprint(controller.info()) +[{'architecture': 'Haswell', + 'filepath': '/home/jeremie/miniconda/envs/dev/lib/libopenblasp-r0.3.17.so', + 'internal_api': 'openblas', + 'num_threads': 4, + 'prefix': 'libopenblas', + 'threading_layer': 'pthreads', + 'user_api': 'blas', + 'version': '0.3.17'}] + +>>> controller.info() == threadpool_info() +True +``` + +### Setting the Maximum Size of Thread-Pools + +Control the number of threads used by the underlying runtime libraries +in specific sections of your Python program: + +```python +>>> from threadpoolctl import threadpool_limits +>>> import numpy as np + +>>> with threadpool_limits(limits=1, user_api='blas'): +... # In this block, calls to blas implementation (like openblas or MKL) +... # will be limited to use only one thread. They can thus be used jointly +... # with thread-parallelism. +... a = np.random.randn(1000, 1000) +... a_squared = a @ a ``` -pip install threadpoolctl + +The threadpools can also be controlled via the object oriented API, which is especially +useful to avoid searching through all the loaded shared libraries each time. It will +however not act on libraries loaded after the instantiation of the +`ThreadpoolController`: + +```python +>>> from threadpoolctl import ThreadpoolController +>>> import numpy as np +>>> controller = ThreadpoolController() + +>>> with controller.limit(limits=1, user_api='blas'): +... a = np.random.randn(1000, 1000) +... a_squared = a @ a +``` + +### Restricting the limits to the scope of a function + +`threadpool_limits` and `ThreadpoolController` can also be used as decorators to set +the maximum number of threads used by the supported libraries at a function level. The +decorators are accessible through their `wrap` method: + +```python +>>> from threadpoolctl import ThreadpoolController, threadpool_limits +>>> import numpy as np +>>> controller = ThreadpoolController() + +>>> @controller.wrap(limits=1, user_api='blas') +... # or @threadpool_limits.wrap(limits=1, user_api='blas') +... def my_func(): +... # Inside this function, calls to blas implementation (like openblas or MKL) +... # will be limited to use only one thread. +... a = np.random.randn(1000, 1000) +... a_squared = a @ a +... ``` -# Usage +### Switching the FlexiBLAS backend + +`FlexiBLAS` is a BLAS wrapper for which the BLAS backend can be switched at runtime. +`threadpoolctl` exposes python bindings for this feature. Here's an example but note +that this part of the API is experimental and subject to change without deprecation: ```python -from threadpoolctl import threadpool_limits +>>> from threadpoolctl import ThreadpoolController +>>> import numpy as np +>>> controller = ThreadpoolController() + +>>> controller.info() +[{'user_api': 'blas', + 'internal_api': 'flexiblas', + 'num_threads': 1, + 'prefix': 'libflexiblas', + 'filepath': '/usr/local/lib/libflexiblas.so.3.3', + 'version': '3.3.1', + 'available_backends': ['NETLIB', 'OPENBLASPTHREAD', 'ATLAS'], + 'loaded_backends': ['NETLIB'], + 'current_backend': 'NETLIB'}] +# Retrieve the flexiblas controller +>>> flexiblas_ct = controller.select(internal_api="flexiblas").lib_controllers[0] -with threadpool_limits(limits=1, user_api='blas'): - # In this block, calls to blas implementation (like openblas or MKL) - # will be limited to use only one thread. They can thus be used jointly - # with thread-parallelism or openmp calls. - ... +# Switch the backend with one predefined at build time (listed in "available_backends") +>>> flexiblas_ct.switch_backend("OPENBLASPTHREAD") +>>> controller.info() +[{'user_api': 'blas', + 'internal_api': 'flexiblas', + 'num_threads': 4, + 'prefix': 'libflexiblas', + 'filepath': '/usr/local/lib/libflexiblas.so.3.3', + 'version': '3.3.1', + 'available_backends': ['NETLIB', 'OPENBLASPTHREAD', 'ATLAS'], + 'loaded_backends': ['NETLIB', 'OPENBLASPTHREAD'], + 'current_backend': 'OPENBLASPTHREAD'}, + {'user_api': 'blas', + 'internal_api': 'openblas', + 'num_threads': 4, + 'prefix': 'libopenblas', + 'filepath': '/usr/lib/x86_64-linux-gnu/openblas-pthread/libopenblasp-r0.3.8.so', + 'version': '0.3.8', + 'threading_layer': 'pthreads', + 'architecture': 'Haswell'}] + +# It's also possible to directly give the path to a shared library +>>> flexiblas_controller.switch_backend("/home/jeremie/miniforge/envs/flexiblas_threadpoolctl/lib/libmkl_rt.so") +>>> controller.info() +[{'user_api': 'blas', + 'internal_api': 'flexiblas', + 'num_threads': 2, + 'prefix': 'libflexiblas', + 'filepath': '/usr/local/lib/libflexiblas.so.3.3', + 'version': '3.3.1', + 'available_backends': ['NETLIB', 'OPENBLASPTHREAD', 'ATLAS'], + 'loaded_backends': ['NETLIB', + 'OPENBLASPTHREAD', + '/home/jeremie/miniforge/envs/flexiblas_threadpoolctl/lib/libmkl_rt.so'], + 'current_backend': '/home/jeremie/miniforge/envs/flexiblas_threadpoolctl/lib/libmkl_rt.so'}, + {'user_api': 'openmp', + 'internal_api': 'openmp', + 'num_threads': 4, + 'prefix': 'libomp', + 'filepath': '/home/jeremie/miniforge/envs/flexiblas_threadpoolctl/lib/libomp.so', + 'version': None}, + {'user_api': 'blas', + 'internal_api': 'openblas', + 'num_threads': 4, + 'prefix': 'libopenblas', + 'filepath': '/usr/lib/x86_64-linux-gnu/openblas-pthread/libopenblasp-r0.3.8.so', + 'version': '0.3.8', + 'threading_layer': 'pthreads', + 'architecture': 'Haswell'}, + {'user_api': 'blas', + 'internal_api': 'mkl', + 'num_threads': 2, + 'prefix': 'libmkl_rt', + 'filepath': '/home/jeremie/miniforge/envs/flexiblas_threadpoolctl/lib/libmkl_rt.so.2', + 'version': '2024.0-Product', + 'threading_layer': 'gnu'}] ``` + +You can observe that the previously linked OpenBLAS shared object stays loaded by +the Python program indefinitely, but FlexiBLAS itself no longer delegates BLAS calls +to OpenBLAS as indicated by the `current_backend` attribute. +### Writing a custom library controller + +Currently, `threadpoolctl` has support for `OpenMP` and the main `BLAS` libraries. +However it can also be used to control the threadpool of other native libraries, +provided that they expose an API to get and set the limit on the number of threads. +For that, one must implement a controller for this library and register it to +`threadpoolctl`. + +A custom controller must be a subclass of the `LibController` class and implement +the attributes and methods described in the docstring of `LibController`. Then this +new controller class must be registered using the `threadpoolctl.register` function. +An complete example can be found [here]( + https://github.com/joblib/threadpoolctl/blob/master/tests/_pyMylib/__init__.py). + +### Sequential BLAS within OpenMP parallel region + +When one wants to have sequential BLAS calls within an OpenMP parallel region, it's +safer to set `limits="sequential_blas_under_openmp"` since setting `limits=1` and +`user_api="blas"` might not lead to the expected behavior in some configurations +(e.g. OpenBLAS with the OpenMP threading layer +https://github.com/xianyi/OpenBLAS/issues/2985). + +### Known Limitations + +- `threadpool_limits` can fail to limit the number of inner threads when nesting + parallel loops managed by distinct OpenMP runtime implementations (for instance + libgomp from GCC and libomp from clang/llvm or libiomp from ICC). + + See the `test_openmp_nesting` function in [tests/test_threadpoolctl.py]( + https://github.com/joblib/threadpoolctl/blob/master/tests/test_threadpoolctl.py) + for an example. More information can be found at: + https://github.com/jeremiedbb/Nested_OpenMP + + Note however that this problem does not happen when `threadpool_limits` is + used to limit the number of threads used internally by BLAS calls that are + themselves nested under OpenMP parallel loops. `threadpool_limits` works as + expected, even if the inner BLAS implementation relies on a distinct OpenMP + implementation. + +- Using Intel OpenMP (ICC) and LLVM OpenMP (clang) in the same Python program + under Linux is known to cause problems. See the following guide for more details + and workarounds: + https://github.com/joblib/threadpoolctl/blob/master/multiple_openmp.md + +- Setting the maximum number of threads of the OpenMP and BLAS libraries has + inconsistent scope and semantics (thread-local vs process-wide) depending on + the underlying library. For more details see + https://github.com/joblib/threadpoolctl/issues/208 + + For example, if you're using OpenMP with libgomp (gcc) or libomp (clang), the + setting is thread-local and sets how many OpenMP threads will be started in + the current thread. On the other hand, with OpenBLAS with pthreads backend or + on Windows, the setting is process-wide and impacts the size of a process-wide + thread pool shared across all threads in the process. + +## Maintainers + +To make a release: + +- Create a PR to bump the version number (`__version__`) in `threadpoolctl.py` and + update the release date in `CHANGES.md`. + +- Merge the PR and check that the `Publish threadpoolctl distribution to TestPyPI` job + of the `publish-to-pypi.yml` workflow successfully uploaded the wheel and source + distribution to Test PyPI. + +- If everything is fine create a tag for the release and push it to github: + + ```bash + git tag -a X.Y.Z + git push git@github.com:joblib/threadpoolctl.git X.Y.Z + ``` + +- Check that the `Publish threadpoolctl distribution to PyPI` job of the + `publish-to-pypi.yml` workflow successfully uploaded the wheel and source distribution + to PyPI this time. + +- Create a PR for the release on the [conda-forge feedstock](https://github.com/conda-forge/threadpoolctl-feedstock) (or wait for the bot to make it). + +- Publish the release on github. + +If for some reason the steps above can't be achieved and a munual upload of the wheel +and source distribution is needed: + +- Build the distribution archives: + + ```bash + pip install flit + flit build + ``` + +- Upload the wheels and source distribution to PyPI using flit. Since PyPI doesn't + allow password authentication anymore, the username needs to be changed to the + generic name `__token__`: + + ```bash + FLIT_USERNAME=__token__ flit publish + ``` + + and a PyPI token has to be passed in place of the password. + +### Credits + +The initial dynamic library introspection code was written by @anton-malakhov +for the smp package available at https://github.com/IntelPython/smp . + +threadpoolctl extends this for other operating systems. Contrary to smp, +threadpoolctl does not attempt to limit the size of Python multiprocessing +pools (threads or processes) or set operating system-level CPU affinity +constraints: threadpoolctl only interacts with native libraries via their +public runtime APIs. diff --git a/benchmarks/bench_context_manager_overhead.py b/benchmarks/bench_context_manager_overhead.py new file mode 100644 index 00000000..d3b69c15 --- /dev/null +++ b/benchmarks/bench_context_manager_overhead.py @@ -0,0 +1,30 @@ +import time +from argparse import ArgumentParser +from pprint import pprint +from statistics import mean, stdev +from threadpoolctl import threadpool_info, threadpool_limits + +parser = ArgumentParser(description="Measure threadpool_limits call overhead.") +parser.add_argument( + "--import", + dest="packages", + default=[], + nargs="+", + help="Python packages to import to load threadpool enabled libraries.", +) +parser.add_argument("--n-calls", type=int, default=100, help="Number of iterations") + +args = parser.parse_args() +for package_name in args.packages: + __import__(package_name) + +pprint(threadpool_info()) + +timings = [] +for _ in range(args.n_calls): + t = time.time() + with threadpool_limits(limits=1): + pass + timings.append(time.time() - t) + +print(f"Overhead per call: {mean(timings) * 1e3:.3f} +/-{stdev(timings) * 1e3:.3f} ms") diff --git a/conftest.py b/conftest.py index 6f3e2e54..bf303839 100644 --- a/conftest.py +++ b/conftest.py @@ -1,35 +1 @@ -import pytest -import warnings - - -def pytest_addoption(parser): - parser.addoption("--openblas-present", action='store_true', - help="Fail test_limit_openblas_threads if BLAS is not " - "found") - parser.addoption("--mkl-present", action='store_true', - help="Fail test_limit_mkl_threads if MKL is not " - "found") - - -@pytest.fixture -def openblas_present(request): - """Fail the test if OpenBLAS is not found""" - return request.config.getoption("--openblas-present") - - -@pytest.fixture -def mkl_present(request): - """Fail the test if MKL is not found""" - return request.config.getoption("--mkl-present") - - -def pytest_configure(config): - """Verify the environment for testing threadpoolctl""" - warnings.simplefilter('always') - - # When using this option, make sure numpy is accessible - if config.getoption("--mkl-present"): - try: - import numpy as np # noqa: F401 - except ImportError: - raise ImportError("Need 'numpy' with option --mkl-present") +collect_ignore = ["tests/_openmp_test_helper"] diff --git a/continuous_integration/build_test_ext.sh b/continuous_integration/build_test_ext.sh old mode 100644 new mode 100755 index 0160a9ad..4bc27696 --- a/continuous_integration/build_test_ext.sh +++ b/continuous_integration/build_test_ext.sh @@ -1,4 +1,22 @@ +#!/bin/bash -cd threadpoolctl/tests/_openmp_test_helper -python setup.py build_ext -i || echo 'No openmp' -cd ../.. \ No newline at end of file +set -xe + +if [[ "$OSTYPE" == "linux-gnu"* ]]; then + pushd tests/_pyMylib + rm -rf *.so *.o + gcc -c -Wall -Werror -fpic -o my_threaded_lib.o my_threaded_lib.c + gcc -shared -o my_threaded_lib.so my_threaded_lib.o + popd +fi + +pushd tests/_openmp_test_helper +rm -rf *.c *.so *.dylib build/ +python setup_inner.py build_ext -i +python setup_outer.py build_ext -i + +# skip scipy required extension if no numpy +if [[ "$NO_NUMPY" != "true" ]]; then + python setup_nested_prange_blas.py build_ext -i +fi +popd diff --git a/continuous_integration/check_no_test_skipped.py b/continuous_integration/check_no_test_skipped.py new file mode 100644 index 00000000..18fd6bc2 --- /dev/null +++ b/continuous_integration/check_no_test_skipped.py @@ -0,0 +1,51 @@ +"""Check tests are not skipped in every ci job""" + +from __future__ import print_function + +import os +import sys +import xml.etree.ElementTree as ET + +base_dir = sys.argv[1] + +# dict {test: result} where result is False if the test was skipped in every +# job and True otherwise. +always_skipped = {} + +for name in os.listdir(base_dir): + # all test result files are in base_dir/test_result_*/ dirs + if name.startswith("test_result_"): + print(f"> processing test result from job {name.replace('test_result_', '')}") + print(" > tests skipped:") + result_file = os.path.join(base_dir, name, "test_result.xml") + root = ET.parse(result_file).getroot() + + # All tests are identified by the xml tag testcase. + for test in root.iter("testcase"): + test_name = test.attrib["name"] + skipped = any(child.tag == "skipped" for child in test) + if skipped: + print(" -", test_name) + if test_name in always_skipped: + always_skipped[test_name] &= skipped + else: + always_skipped[test_name] = skipped + +print("\n------------------------------------------------------------------\n") + +# List of tests that we don't want to fail the CI if they are skipped in +# every job. This is useful for tests that depend on specific versions of +# numpy or scipy and we don't want to pin old versions of these libraries. +SAFE_SKIPPED_TESTS = ["test_multiple_shipped_openblas"] + +fail = False +for test, skipped in always_skipped.items(): + if skipped: + if test in SAFE_SKIPPED_TESTS: + print(test, "was skipped in every job but it's fine to skip it") + else: + fail = True + print(test, "was skipped in every job") + +if fail: + sys.exit(1) diff --git a/continuous_integration/install.cmd b/continuous_integration/install.cmd deleted file mode 100644 index 38e896f1..00000000 --- a/continuous_integration/install.cmd +++ /dev/null @@ -1,33 +0,0 @@ -@rem https://github.com/numba/numba/blob/master/buildscripts/incremental/setup_conda_environment.cmd -@rem The cmd /C hack circumvents a regression where conda installs a conda.bat -@rem script in non-root environments. -set CONDA_INSTALL=cmd /C conda install -q -y -set PIP_INSTALL=pip install -q - -@echo on - -@rem Deactivate any environment -call deactivate -@rem Clean up any left-over from a previous build and install version of python -conda remove --all -q -y -n %VIRTUALENV% -conda create -n %VIRTUALENV% -q -y python=%VERSION_PYTHON% - -call activate %VIRTUALENV% -python -m pip install -U pip -python --version -pip --version - -@rem Install dependencies with either conda or pip. -if "%PACKAGER%" == "conda" (%CONDA_INSTALL% numpy=1.15 pytest cython) -if "%PACKAGER%" == "pip" (%PIP_INSTALL% numpy scipy pytest cython) - -@rem Install extra dependency -pip install -q coverage pytest-cov - -@rem Install package -pip install -e . - -@rem Build the cython test helper for openmp -bash ./continuous_integration/build_test_ext.sh - -if %errorlevel% neq 0 exit /b %errorlevel% diff --git a/continuous_integration/install.sh b/continuous_integration/install.sh old mode 100755 new mode 100644 index 4f552be0..7a958980 --- a/continuous_integration/install.sh +++ b/continuous_integration/install.sh @@ -1,49 +1,127 @@ #!/bin/bash -set -e +# License: BSD 3-Clause + +set -xe UNAMESTR=`uname` -if [[ "$UNAMESTR" == "Darwin" ]]; then - # install OpenMP not present by default on osx - HOMEBREW_NO_AUTO_UPDATE=1 brew install libomp - - # enable OpenMP support for Apple-clang - export CC=/usr/bin/clang - export CXX=/usr/bin/clang++ - export CPPFLAGS="$CPPFLAGS -Xpreprocessor -fopenmp" - export CFLAGS="$CFLAGS -I/usr/local/opt/libomp/include" - export CXXFLAGS="$CXXFLAGS -I/usr/local/opt/libomp/include" - export LDFLAGS="$LDFLAGS -L/usr/local/opt/libomp/lib -lomp" - export DYLD_LIBRARY_PATH=/usr/local/opt/libomp/lib + +# Install a recent version of clang and libomp if needed +# Only applicable to linux jobs +if [[ "$CC_OUTER_LOOP" == "clang-18" ]] || \ + [[ "$CC_INNER_LOOP" == "clang-18" ]] || \ + [[ "$BLIS_CC" == "clang-18" ]] +then + wget https://apt.llvm.org/llvm.sh + chmod +x llvm.sh + sudo ./llvm.sh 18 + sudo apt-get install libomp-dev fi + make_conda() { - TO_INSTALL="$@" - conda create -n $VIRTUALENV -q --yes $TO_INSTALL - source activate $VIRTUALENV + CHANNEL="$1" + TO_INSTALL="$2" + if [[ "$UNAMESTR" == "Darwin" ]]; then + if [[ "$INSTALL_LIBOMP" == "conda-forge" ]]; then + # Install an OpenMP-enabled clang/llvm from conda-forge + # assumes conda-forge is set on priority channel + TO_INSTALL="$TO_INSTALL compilers llvm-openmp" + + elif [[ "$INSTALL_LIBOMP" == "homebrew" ]]; then + # Install a compiler with a working openmp + HOMEBREW_NO_AUTO_UPDATE=1 brew install libomp + + # enable OpenMP support for Apple-clang + export CC=/usr/bin/clang + export CPPFLAGS="$CPPFLAGS -Xpreprocessor -fopenmp" + export CFLAGS="$CFLAGS -I/opt/homebrew/opt/libomp/include" + export LDFLAGS="$LDFLAGS -Wl,-rpath,/opt/homebrew/opt/libomp/lib -L/opt/homebrew/opt/libomp/lib -lomp" + fi + fi + + if [[ "$PYTHON_VERSION" == "*" ]]; then + # Avoid installing free-threaded python + TO_INSTALL="$TO_INSTALL python-gil" + fi + + # prevent mixing conda channels + conda config --set channel_priority strict + conda config --add channels $CHANNEL + + conda update -n base conda conda-libmamba-solver -q --yes + conda config --set solver libmamba + + conda create -n testenv -q --yes python=$PYTHON_VERSION $TO_INSTALL + conda activate testenv } + if [[ "$PACKAGER" == "conda" ]]; then - TO_INSTALL="python=$VERSION_PYTHON pip pytest pytest-cov \ - numpy cython" + TO_INSTALL="" + if [[ "$NO_NUMPY" != "true" ]]; then + TO_INSTALL="$TO_INSTALL numpy scipy" + if [[ -n "$BLAS" ]]; then + TO_INSTALL="$TO_INSTALL blas=*=$BLAS" + fi + fi + make_conda "defaults" "$TO_INSTALL" + +elif [[ "$PACKAGER" == "conda-forge" ]]; then + TO_INSTALL="numpy scipy blas=*=$BLAS" + if [[ "$BLAS" == "openblas" && "$OPENBLAS_THREADING_LAYER" == "openmp" ]]; then + TO_INSTALL="$TO_INSTALL libopenblas=*=*openmp*" + fi + make_conda "conda-forge" "$TO_INSTALL" + +elif [[ "$PACKAGER" == "pip" ]]; then + # Use conda to build an empty python env and then use pip to install + # numpy and scipy + make_conda "conda-forge" "" + if [[ "$NO_NUMPY" != "true" ]]; then + pip install numpy scipy + fi - make_conda $TO_INSTALL +elif [[ "$PACKAGER" == "pip-dev" ]]; then + # Use conda to build an empty python env and then use pip to install + # numpy and scipy dev versions + make_conda "conda-forge" "" + + dev_anaconda_url=https://pypi.anaconda.org/scientific-python-nightly-wheels/simple + pip install --pre --upgrade --timeout=60 --extra-index $dev_anaconda_url numpy scipy elif [[ "$PACKAGER" == "ubuntu" ]]; then - sudo apt-get install python3-scipy libatlas3-base libatlas-base-dev libatlas-dev libopenblas-base python3-virtualenv - python3 -m virtualenv --system-site-packages --python=python3 $VIRTUALENV - source $VIRTUALENV/bin/activate - python -m pip install pytest pytest-cov cython -fi + # Remove the ubuntu toolchain PPA that seems to be invalid: + # https://github.com/scikit-learn/scikit-learn/pull/13934 + sudo add-apt-repository --remove ppa:ubuntu-toolchain-r/test + sudo apt-get update + sudo apt-get install python3-scipy python3-virtualenv $APT_BLAS + python3 -m virtualenv --system-site-packages --python=python3 testenv + source testenv/bin/activate +elif [[ "$INSTALL_BLAS" == "blis" ]]; then + TO_INSTALL="cython meson-python pkg-config" + make_conda "conda-forge" "$TO_INSTALL" + source ./continuous_integration/install_blis.sh -python -m pip install coverage +elif [[ "$INSTALL_BLAS" == "flexiblas" ]]; then + TO_INSTALL="cython openblas $PLATFORM_SPECIFIC_PACKAGES meson-python pkg-config compilers" + make_conda "conda-forge" "$TO_INSTALL" + source ./continuous_integration/install_flexiblas.sh +fi + +python -m pip install -v -q -r dev-requirements.txt bash ./continuous_integration/build_test_ext.sh +# Check which BLAS is linked (only available on linux) +if [[ "$UNAMESTR" == "Linux" && "$NO_NUMPY" != "true" ]]; then + ldd tests/_openmp_test_helper/nested_prange_blas.cpython*.so +fi + python --version -python -c "import numpy; print('numpy %s' % numpy.__version__)" || echo "no numpy" -python -c "import scipy; print('scipy %s' % scipy.__version__)" || echo "no scipy" +python -c "import numpy; print(f'numpy {numpy.__version__}')" || echo "no numpy" +python -c "import scipy; print(f'scipy {scipy.__version__}')" || echo "no scipy" -pip install -e . +python -m flit install --symlink diff --git a/continuous_integration/install_blis.sh b/continuous_integration/install_blis.sh new file mode 100644 index 00000000..e8fa8074 --- /dev/null +++ b/continuous_integration/install_blis.sh @@ -0,0 +1,49 @@ +#!/bin/bash + +set -xe + + +# Install gcc 12 to build BLIS +if [[ "$BLIS_CC" == "gcc-12" ]]; then + sudo apt install gcc-12 +fi + +# step outside of threadpoolctl directory +pushd .. +ABS_PATH=$(pwd) + +# build & install blis +mkdir BLIS_install +git clone https://github.com/flame/blis.git +pushd blis + +./configure --prefix=$ABS_PATH/BLIS_install --enable-cblas --enable-threading=$BLIS_ENABLE_THREADING CC=$BLIS_CC auto +make -j4 +make install +popd + +# build & install numpy +git clone https://github.com/numpy/numpy.git +pushd numpy +git submodule update --init + +echo "libdir=$ABS_PATH/BLIS_install/lib/ +includedir=$ABS_PATH/BLIS_install/include/blis/ +version=latest +extralib=-lm -lpthread -lgfortran +Name: blis +Description: BLIS +Version: \${version} +Libs: -L\${libdir} -lblis +Libs.private: \${extralib} +Cflags: -I\${includedir}" > blis.pc + +PKG_CONFIG_PATH=$ABS_PATH/numpy/ pip install . -v --no-build-isolation -Csetup-args=-Dblas=blis + +export CFLAGS=-I$ABS_PATH/BLIS_install/include/blis +export LDFLAGS="-L$ABS_PATH/BLIS_install/lib -Wl,-rpath,$ABS_PATH/BLIS_install/lib" + +popd + +# back to threadpoolctl directory +popd diff --git a/continuous_integration/install_flexiblas.sh b/continuous_integration/install_flexiblas.sh new file mode 100644 index 00000000..1f50801b --- /dev/null +++ b/continuous_integration/install_flexiblas.sh @@ -0,0 +1,68 @@ +#!/bin/bash + +set -xe + +# step outside of threadpoolctl directory +pushd .. +ABS_PATH=$(pwd) + +# build & install FlexiBLAS +mkdir flexiblas_install +git clone https://github.com/mpimd-csc/flexiblas.git +pushd flexiblas + +# Temporary ping Flexiblas commit to avoid openmp symbols not found at link time +git checkout v3.4.2 + +mkdir build +pushd build + +EXTENSION=".so" +if [[ $(uname) == "Darwin" ]]; then + EXTENSION=".dylib" +fi + +# We intentionally restrict the list of backends to make it easier to +# write platform agnostic tests. In particular, we do not detect OS +# provided backends such as macOS' Apple/Accelerate/vecLib nor plaftorm +# specific BLAS implementations such as MKL that cannot be installed on +# arm64 hardware. +cmake ../ -DCMAKE_INSTALL_PREFIX=$ABS_PATH"/flexiblas_install" \ + -DCMAKE_POLICY_VERSION_MINIMUM=3.5 \ + -DBLAS_AUTO_DETECT="OFF" \ + -DEXTRA="OPENBLAS_CONDA" \ + -DFLEXIBLAS_DEFAULT="OPENBLAS_CONDA" \ + -DOPENBLAS_CONDA_LIBRARY=$CONDA_PREFIX"/lib/libopenblas"$EXTENSION \ +make +make install + +# Check that all 3 BLAS are listed in FlexiBLAS configuration +$ABS_PATH/flexiblas_install/bin/flexiblas list +popd +popd + +# build & install numpy +git clone https://github.com/numpy/numpy.git +pushd numpy +git submodule update --init + +echo "libdir=$ABS_PATH/flexiblas_install/lib/ +includedir=$ABS_PATH/flexiblas_install/include/flexiblas/ +version=3.3.1 +extralib=-lm -lpthread -lgfortran +Name: flexiblas +Description: FlexiBLAS - a BLAS wrapper +Version: \${version} +Libs: -L\${libdir} -lflexiblas +Libs.private: \${extralib} +Cflags: -I\${includedir}" > flexiblas.pc + +PKG_CONFIG_PATH=$ABS_PATH/numpy/ pip install . -v --no-build-isolation -Csetup-args=-Dblas=flexiblas -Csetup-args=-Dlapack=flexiblas + +export CFLAGS=-I$ABS_PATH/flexiblas_install/include/flexiblas \ +export LDFLAGS="-L$ABS_PATH/flexiblas_install/lib -Wl,-rpath,$ABS_PATH/flexiblas_install/lib" \ + +popd + +# back to threadpoolctl directory +popd diff --git a/continuous_integration/posix.yml b/continuous_integration/posix.yml deleted file mode 100644 index 8e486799..00000000 --- a/continuous_integration/posix.yml +++ /dev/null @@ -1,37 +0,0 @@ -parameters: - name: '' - vmImage: '' - matrix: [] - -jobs: -- job: ${{ parameters.name }} - pool: - vmImage: ${{ parameters.vmImage }} - variables: - TEST_DIR: '$(Agent.WorkFolder)/tmp_folder' - VIRTUALENV: 'testvenv' - JUNITXML: 'test-data.xml' - strategy: - matrix: - ${{ insert }}: ${{ parameters.matrix }} - - steps: - - bash: echo "##vso[task.prependpath]$CONDA/bin" - displayName: Add conda to PATH - condition: eq(variables['PACKAGER'], 'conda') - - script: | - continuous_integration/install.sh - displayName: 'Install' - - script: | - continuous_integration/test_script.sh - displayName: 'Test Library' - - task: PublishTestResults@2 - inputs: - testResultsFiles: '$(JUNITXML)' - testRunTitle: ${{ format('{0}-$(Agent.JobName)', parameters.name) }} - displayName: 'Publish Test Results' - condition: succeededOrFailed() - - script: | - bash continuous_integration/upload_codecov.sh - displayName: 'Upload to codecov' - condition: succeeded() diff --git a/continuous_integration/run_tests.sh b/continuous_integration/run_tests.sh new file mode 100755 index 00000000..f26f0fdc --- /dev/null +++ b/continuous_integration/run_tests.sh @@ -0,0 +1,21 @@ +#!/bin/bash + +set -xe + +if [[ "$PACKAGER" == conda* ]] || [[ -z "$PACKAGER" ]]; then + conda activate testenv + conda list +elif [[ "$PACKAGER" == pip* ]]; then + # we actually use conda to install the base environment: + conda activate testenv + pip list +elif [[ "$PACKAGER" == "ubuntu" ]]; then + source testenv/bin/activate + pip list +fi + +# Use the CLI to display the effective runtime environment prior to +# launching the tests: +python -m threadpoolctl -i numpy scipy.linalg tests._openmp_test_helper.openmp_helpers_inner + +pytest -vlrxXs -W error -k "$TESTS" --junitxml=test_result.xml --cov=threadpoolctl --cov-report xml diff --git a/continuous_integration/test_script.cmd b/continuous_integration/test_script.cmd deleted file mode 100644 index 29d718c0..00000000 --- a/continuous_integration/test_script.cmd +++ /dev/null @@ -1,5 +0,0 @@ -set DEFAULT_PYTEST_ARGS=-vlx --cov=threadpoolctl - -call activate %VIRTUALENV% - -pytest --junitxml=%JUNITXML% %DEFAULT_PYTEST_ARGS% diff --git a/continuous_integration/test_script.sh b/continuous_integration/test_script.sh deleted file mode 100755 index 6372b09b..00000000 --- a/continuous_integration/test_script.sh +++ /dev/null @@ -1,13 +0,0 @@ -#!/bin/bash - -set -e - -if [[ "$PACKAGER" == "conda" ]]; then - source activate $VIRTUALENV -elif [[ "$PACKAGER" == "ubuntu" ]]; then - source $VIRTUALENV/bin/activate -fi - -set -x -pytest -vl --junitxml=$JUNITXML --cov=threadpoolctl -set +x diff --git a/continuous_integration/upload_codecov.sh b/continuous_integration/upload_codecov.sh deleted file mode 100755 index f188ba15..00000000 --- a/continuous_integration/upload_codecov.sh +++ /dev/null @@ -1,7 +0,0 @@ -#!/bin/bash - -set -e - -python -m pip install --user codecov - -python -m codecov || echo "codecov upload failed" diff --git a/continuous_integration/windows.yml b/continuous_integration/windows.yml deleted file mode 100644 index d05ba3dc..00000000 --- a/continuous_integration/windows.yml +++ /dev/null @@ -1,38 +0,0 @@ - -parameters: - name: '' - vmImage: '' - matrix: [] - -jobs: -- job: ${{ parameters.name }} - pool: - vmImage: ${{ parameters.vmImage }} - variables: - VIRTUALENV: 'testvenv' - JUNITXML: 'test-data.xml' - CODECOV_TOKEN: 'cee0e505-c12e-4139-aa43-621fb16a2347' - strategy: - matrix: - ${{ insert }}: ${{ parameters.matrix }} - - steps: - - powershell: Write-Host "##vso[task.prependpath]$env:CONDA\Scripts" - displayName: Add conda to PATH - condition: eq(variables['PACKAGER'], 'conda') - - script: | - continuous_integration\\install.cmd - displayName: 'Install' - - script: | - continuous_integration\\test_script.cmd - displayName: 'Test threadpoolctl' - - task: PublishTestResults@2 - inputs: - testResultsFiles: '$(JUNITXML)' - testRunTitle: ${{ format('{0}-$(Agent.JobName)', parameters.name) }} - displayName: 'Publish Test Results' - condition: succeededOrFailed() - - script: | - bash continuous_integration\\upload_codecov.sh - displayName: 'Upload to codecov' - condition: succeeded() diff --git a/dev-requirements.txt b/dev-requirements.txt new file mode 100644 index 00000000..04901b66 --- /dev/null +++ b/dev-requirements.txt @@ -0,0 +1,6 @@ +flit +coverage +pytest +pytest-cov +cython +setuptools diff --git a/multiple_openmp.md b/multiple_openmp.md new file mode 100644 index 00000000..8da2cb91 --- /dev/null +++ b/multiple_openmp.md @@ -0,0 +1,87 @@ +# Multiple OpenMP Runtimes + +## Context + +OpenMP is an API specification for parallel programming. There are many +implementations of it, tied to a compiler most of the time: + +- `libgomp` for GCC (GNU C/C++ Compiler), +- `libomp` for Clang (LLVM C/C++ Compiler), +- `libiomp` for ICC (Intel C/C++ Compiler), +- `vcomp` for MSVC (Microsoft Visual Studio C/C++ Compiler). + +In general, it is not advised to have different OpenMP runtime libraries (or +even different copies of the same library) loaded at the same time in a +program. It's considered an undefined behavior. Fortunately it is not as bad as +it sounds in most situations. + +However this situation is frequent in the Python ecosystem since you can +install packages compiled with different compilers (hence linked to different +OpenMP implementations) and import them together in a Python program. + +A typical example is installing NumPy from Anaconda which is linked against MKL +(Intel's math library) and another package that uses multi-threading with OpenMP +directly in a compiled extension, as is the case in Scikit-learn (via Cython +`prange`), LightGBM and XGBoost (via pragmas in the C++ source code). + +From our experience, **most OpenMP libraries can seamlessly coexist in a same +program**. For instance, on Linux, we never observed any issue between +`libgomp` and `libiomp`, which is the most common mix (NumPy with MKL + a +package compiled with GCC, the most widely used C compiler on that platform). + +## Incompatibility between Intel OpenMP and LLVM OpenMP + +The only unrecoverable incompatibility we encountered happens when loading a +mix of compiled extensions linked with **`libomp` (LLVM/Clang) and `libiomp` +(ICC), on Linux and macOS**, manifested by crashes or deadlocks. It can happen +even with the simplest OpenMP calls like getting the maximum number of threads +that will be used in a subsequent parallel region. A possible explanation is that +`libomp` is actually a fork of `libiomp` causing name colliding for instance. +Using `threadpoolctl` may crash your program in such a setting. + +**Fortunately this problem is very rare**: at the time of writing, all major +binary distributions of Python packages for Linux use either GCC or ICC to +build the Python scientific packages. Therefore this problem would only happen +if some packagers decide to start shipping Python packages built with +LLVM/Clang instead of GCC (this is the case for instance with conda's default channel). + +## Workarounds for Intel OpenMP and LLVM OpenMP case + +As far as we know, the only workaround consists in making sure only of one of +the two incompatible OpenMP libraries is loaded. For example: + +- Tell MKL (used by NumPy) to use another threading implementation instead of the Intel + OpenMP runtime. It can be the GNU OpenMP runtime on Linux or TBB on Linux and MacOS + for instance. This is done by setting the following environment variable: + + export MKL_THREADING_LAYER=GNU + + or, if TBB is installed: + + export MKL_THREADING_LAYER=TBB + +- Install a build of NumPy and SciPy linked against OpenBLAS instead of MKL. + This can be done for instance by installing NumPy and SciPy from PyPI: + + pip install numpy scipy + + from the conda-forge conda channel: + + conda install -c conda-forge numpy scipy + + or from the default conda channel: + + conda install numpy scipy blas[build=openblas] + +- Re-build your OpenMP-enabled extensions from source with GCC (or ICC) instead + of Clang if you want to keep on using NumPy/SciPy linked against MKL with the + default `libiomp`-based threading layer. + +## References + +The above incompatibility has been reported upstream to the LLVM and Intel +developers on the following public issue trackers/forums along with a minimal +reproducer written in C: + +- https://bugs.llvm.org/show_bug.cgi?id=43565 +- https://community.intel.com/t5/Intel-C-Compiler/Cannot-call-OpenMP-functions-from-libiomp-after-calling-from/m-p/1176406 diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 00000000..c0ea1696 --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,30 @@ +[build-system] +requires = ["flit_core >=2,<4"] +build-backend = "flit_core.buildapi" + +# TODO: replace the following section by the standard [project] section to be able +# to use flit v4. +[tool.flit.metadata] +module = "threadpoolctl" +author = "Thomas Moreau" +author-email = "thomas.moreau.2010@gmail.com" +home-page = "https://github.com/joblib/threadpoolctl" +description-file = "README.md" +requires-python = ">=3.9" +license = "BSD-3-Clause" +classifiers = [ + "Intended Audience :: Developers", + "License :: OSI Approved :: BSD License", + "Programming Language :: Python :: 3", + "Programming Language :: Python :: 3.9", + "Programming Language :: Python :: 3.10", + "Programming Language :: Python :: 3.11", + "Programming Language :: Python :: 3.12", + "Programming Language :: Python :: 3.13", + "Topic :: Software Development :: Libraries :: Python Modules", +] + +[tool.black] +line-length = 88 +target_version = ['py39', 'py310', 'py311', 'py312', 'py313'] +preview = true diff --git a/setup.py b/setup.py deleted file mode 100644 index b89146ad..00000000 --- a/setup.py +++ /dev/null @@ -1,88 +0,0 @@ -import os -import re -import shutil -from setuptools import setup, find_packages -from distutils.command.clean import clean as Clean - -packages = find_packages( - exclude=['threadpoolctl.tests', 'threadpoolctl.tests._openmp_test_helper']) - - -# Function to parse __version__ in `threadpoolctl` -def find_version(): - here = os.path.abspath(os.path.dirname(__file__)) - with open(os.path.join(here, 'threadpoolctl', '__init__.py'), 'r') as fp: - version_file = fp.read() - version_match = re.search(r"^__version__ = ['\"]([^'\"]*)['\"]", - version_file, re.M) - if version_match: - return version_match.group(1) - raise RuntimeError("Unable to find version string.") - - -# Custom clean command to remove build artifacts -class CleanCommand(Clean): - description = "Remove build artifacts from the source tree" - - def run(self): - Clean.run(self) - # Remove c files if we are not within a sdist package - cwd = os.path.abspath(os.path.dirname(__file__)) - remove_c_files = not os.path.exists(os.path.join(cwd, 'PKG-INFO')) - if remove_c_files: - print('Will remove generated .c files') - if os.path.exists('build'): - shutil.rmtree('build') - for dirpath, dirnames, filenames in os.walk('.'): - for filename in filenames: - if any(filename.endswith(suffix) for suffix in - (".so", ".pyd", ".dll", ".pyc")): - os.unlink(os.path.join(dirpath, filename)) - continue - extension = os.path.splitext(filename)[1] - if remove_c_files and extension in ['.c', '.cpp']: - pyx_file = str.replace(filename, extension, '.pyx') - if os.path.exists(os.path.join(dirpath, pyx_file)): - os.unlink(os.path.join(dirpath, filename)) - for dirname in dirnames: - if dirname == '__pycache__': - shutil.rmtree(os.path.join(dirpath, dirname)) - - -cmdclass = {'clean': CleanCommand} - -setup( - name='threadpoolctl', - version=find_version(), - description=("Python helpers to limit the number of thread used in " - "thread-pool backed parallelism for C-libraries"), - long_description=open('README.md', 'rb').read().decode('utf-8'), - long_description_content_type='text/markdown', - url='https://github.com/joblib/threadpoolctl/', - author='Thomas Moreau', - author_email='thomas.moreau.2010@gmail.com', - packages=packages, - zip_safe=False, - license='BSD', - classifiers=[ - 'Development Status :: 5 - Production/Stable', - 'Environment :: Console', - 'Intended Audience :: Developers', - 'Intended Audience :: Science/Research', - 'License :: OSI Approved :: BSD License', - 'Operating System :: OS Independent', - 'Programming Language :: Python :: 2.7', - 'Programming Language :: Python :: 3', - 'Programming Language :: Python :: 3.4', - 'Programming Language :: Python :: 3.5', - 'Programming Language :: Python :: 3.6', - 'Programming Language :: Python :: 3.7', - 'Programming Language :: Python :: 3.8', - 'Topic :: Scientific/Engineering', - 'Topic :: Utilities', - 'Topic :: Software Development :: Libraries', - ], - cmdclass=cmdclass, - platforms='any', - tests_require=['pytest'], -) diff --git a/threadpoolctl/tests/__init__.py b/tests/__init__.py similarity index 100% rename from threadpoolctl/tests/__init__.py rename to tests/__init__.py diff --git a/tests/_openmp_test_helper/__init__.py b/tests/_openmp_test_helper/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/_openmp_test_helper/build_utils.py b/tests/_openmp_test_helper/build_utils.py new file mode 100644 index 00000000..90356b33 --- /dev/null +++ b/tests/_openmp_test_helper/build_utils.py @@ -0,0 +1,22 @@ +import os +import sys + + +def set_cc_variables(var_name="CC"): + cc_var = os.environ.get(var_name) + if cc_var is not None: + os.environ["CC"] = cc_var + if sys.platform == "darwin": + os.environ["LDSHARED"] = cc_var + " -bundle -undefined dynamic_lookup" + else: + os.environ["LDSHARED"] = cc_var + " -shared" + + return cc_var + + +def get_openmp_flag(): + if sys.platform == "win32": + return ["/openmp"] + elif sys.platform == "darwin" and "openmp" in os.getenv("CPPFLAGS", ""): + return [] + return ["-fopenmp"] diff --git a/tests/_openmp_test_helper/nested_prange_blas.pyx b/tests/_openmp_test_helper/nested_prange_blas.pyx new file mode 100644 index 00000000..af2a6693 --- /dev/null +++ b/tests/_openmp_test_helper/nested_prange_blas.pyx @@ -0,0 +1,44 @@ +cimport openmp +from cython.parallel import parallel, prange + +import numpy as np +from scipy.linalg.cython_blas cimport dgemm + +from threadpoolctl import ThreadpoolController + + +def check_nested_prange_blas(double[:, ::1] A, double[:, ::1] B, int nthreads): + """Run multithreaded BLAS calls within OpenMP parallel loop""" + cdef: + int m = A.shape[0] + int n = B.shape[0] + int k = A.shape[1] + + double[:, ::1] C = np.empty((m, n)) + int n_chunks = 100 + int chunk_size = A.shape[0] // n_chunks + + char* trans = 't' + char* no_trans = 'n' + double alpha = 1.0 + double beta = 0.0 + + int i + int prange_num_threads + int *prange_num_threads_ptr = &prange_num_threads + + inner_info = [None] + + with nogil, parallel(num_threads=nthreads): + if openmp.omp_get_thread_num() == 0: + with gil: + inner_info[0] = ThreadpoolController().info() + + prange_num_threads_ptr[0] = openmp.omp_get_num_threads() + + for i in prange(n_chunks): + dgemm(trans, no_trans, &n, &chunk_size, &k, + &alpha, &B[0, 0], &k, &A[i * chunk_size, 0], &k, + &beta, &C[i * chunk_size, 0], &n) + + return np.asarray(C), prange_num_threads, inner_info[0] diff --git a/tests/_openmp_test_helper/nested_prange_blas_custom.pyx b/tests/_openmp_test_helper/nested_prange_blas_custom.pyx new file mode 100644 index 00000000..6c0a9771 --- /dev/null +++ b/tests/_openmp_test_helper/nested_prange_blas_custom.pyx @@ -0,0 +1,55 @@ +cimport openmp +from cython.parallel import parallel, prange + +import numpy as np + +cdef extern from 'cblas.h' nogil: + ctypedef enum CBLAS_ORDER: + CblasRowMajor=101 + CblasColMajor=102 + ctypedef enum CBLAS_TRANSPOSE: + CblasNoTrans=111 + CblasTrans=112 + CblasConjTrans=113 + void dgemm 'cblas_dgemm' ( + CBLAS_ORDER Order, CBLAS_TRANSPOSE TransA, + CBLAS_TRANSPOSE TransB, int M, int N, + int K, double alpha, double *A, int lda, + double *B, int ldb, double beta, double *C, int ldc) + +from threadpoolctl import ThreadpoolController + + +def check_nested_prange_blas(double[:, ::1] A, double[:, ::1] B, int nthreads): + """Run multithreaded BLAS calls within OpenMP parallel loop""" + cdef: + int m = A.shape[0] + int n = B.shape[0] + int k = A.shape[1] + + double[:, ::1] C = np.empty((m, n)) + int n_chunks = 100 + int chunk_size = A.shape[0] // n_chunks + + double alpha = 1.0 + double beta = 0.0 + + int i + int prange_num_threads + int *prange_num_threads_ptr = &prange_num_threads + + inner_info = [None] + + with nogil, parallel(num_threads=nthreads): + if openmp.omp_get_thread_num() == 0: + with gil: + inner_info[0] = ThreadpoolController().info() + + prange_num_threads_ptr[0] = openmp.omp_get_num_threads() + + for i in prange(n_chunks): + dgemm(CblasRowMajor, CblasNoTrans, CblasTrans, + chunk_size, n, k, alpha, &A[i * chunk_size, 0], k, + &B[0, 0], k, beta, &C[i * chunk_size, 0], n) + + return np.asarray(C), prange_num_threads, inner_info[0] diff --git a/tests/_openmp_test_helper/openmp_helpers_inner.pxd b/tests/_openmp_test_helper/openmp_helpers_inner.pxd new file mode 100644 index 00000000..f4d0852c --- /dev/null +++ b/tests/_openmp_test_helper/openmp_helpers_inner.pxd @@ -0,0 +1 @@ +cdef int inner_openmp_loop(int) noexcept nogil diff --git a/tests/_openmp_test_helper/openmp_helpers_inner.pyx b/tests/_openmp_test_helper/openmp_helpers_inner.pyx new file mode 100644 index 00000000..e7928d2f --- /dev/null +++ b/tests/_openmp_test_helper/openmp_helpers_inner.pyx @@ -0,0 +1,38 @@ +cimport openmp +from cython.parallel import prange + + +def check_openmp_num_threads(int n): + """Run a short parallel section with OpenMP + + Return the number of threads that where effectively used by the + OpenMP runtime. + """ + cdef int num_threads = -1 + + with nogil: + num_threads = inner_openmp_loop(n) + return num_threads + + +cdef int inner_openmp_loop(int n) noexcept nogil: + """Run a short parallel section with OpenMP + + Return the number of threads that where effectively used by the + OpenMP runtime. + + This function is expected to run without the GIL and can be called + by an outer OpenMP / prange loop written in Cython in another file. + """ + cdef long n_sum = 0 + cdef int i, num_threads + + for i in prange(n): + num_threads = openmp.omp_get_num_threads() + n_sum += i + + if n_sum != (n - 1) * n / 2: + # error + return -1 + + return num_threads diff --git a/tests/_openmp_test_helper/openmp_helpers_outer.pyx b/tests/_openmp_test_helper/openmp_helpers_outer.pyx new file mode 100644 index 00000000..2c8a383c --- /dev/null +++ b/tests/_openmp_test_helper/openmp_helpers_outer.pyx @@ -0,0 +1,22 @@ +cimport openmp +from cython.parallel import prange +from openmp_helpers_inner cimport inner_openmp_loop + + +def check_nested_openmp_loops(int n, nthreads=None): + """Run a short parallel section with OpenMP with nested calls + + The inner OpenMP loop has not necessarily been built/linked with the + same runtime OpenMP runtime. + """ + cdef: + int outer_num_threads = -1 + int inner_num_threads = -1 + int num_threads = nthreads or openmp.omp_get_max_threads() + int i + + for i in prange(n, num_threads=num_threads, nogil=True): + inner_num_threads = inner_openmp_loop(n) + outer_num_threads = openmp.omp_get_num_threads() + + return outer_num_threads, inner_num_threads diff --git a/tests/_openmp_test_helper/setup_inner.py b/tests/_openmp_test_helper/setup_inner.py new file mode 100644 index 00000000..547e257a --- /dev/null +++ b/tests/_openmp_test_helper/setup_inner.py @@ -0,0 +1,38 @@ +import os +from setuptools import Extension, setup +from Cython.Build import cythonize + +from build_utils import set_cc_variables +from build_utils import get_openmp_flag + +original_environ = os.environ.copy() +try: + # Make it possible to compile the 2 OpenMP enabled Cython extensions + # with different compilers and therefore different OpenMP runtimes. + inner_loop_cc_var = set_cc_variables("CC_INNER_LOOP") + openmp_flag = get_openmp_flag() + + ext_modules = [ + Extension( + "openmp_helpers_inner", + ["openmp_helpers_inner.pyx"], + extra_compile_args=openmp_flag, + extra_link_args=openmp_flag, + ) + ] + + setup( + name="_openmp_test_helper_inner", + ext_modules=cythonize( + ext_modules, + compiler_directives={ + "language_level": 3, + "boundscheck": False, + "wraparound": False, + }, + compile_time_env={"CC_INNER_LOOP": inner_loop_cc_var or "unknown"}, + ), + ) + +finally: + os.environ.update(original_environ) diff --git a/tests/_openmp_test_helper/setup_nested_prange_blas.py b/tests/_openmp_test_helper/setup_nested_prange_blas.py new file mode 100644 index 00000000..dcb39ca6 --- /dev/null +++ b/tests/_openmp_test_helper/setup_nested_prange_blas.py @@ -0,0 +1,41 @@ +import os +from setuptools import Extension, setup +from Cython.Build import cythonize + +from build_utils import set_cc_variables +from build_utils import get_openmp_flag + +original_environ = os.environ.copy() +try: + set_cc_variables("CC_OUTER_LOOP") + openmp_flag = get_openmp_flag() + + use_custom_blas = os.getenv("INSTALL_BLAS", False) + libraries = [use_custom_blas] if use_custom_blas else [] + custom_suffix = "_custom" if use_custom_blas else "" + filename = f"nested_prange_blas{custom_suffix}.pyx" + + ext_modules = [ + Extension( + "nested_prange_blas", + [filename], + extra_compile_args=openmp_flag, + extra_link_args=openmp_flag, + libraries=libraries, + ) + ] + + setup( + name="_openmp_test_helper_nested_prange_blas", + ext_modules=cythonize( + ext_modules, + compiler_directives={ + "language_level": 3, + "boundscheck": False, + "wraparound": False, + }, + ), + ) + +finally: + os.environ.update(original_environ) diff --git a/tests/_openmp_test_helper/setup_outer.py b/tests/_openmp_test_helper/setup_outer.py new file mode 100644 index 00000000..dee39b1a --- /dev/null +++ b/tests/_openmp_test_helper/setup_outer.py @@ -0,0 +1,38 @@ +import os +from setuptools import Extension, setup +from Cython.Build import cythonize + +from build_utils import set_cc_variables +from build_utils import get_openmp_flag + +original_environ = os.environ.copy() +try: + # Make it possible to compile the 2 OpenMP enabled Cython extensions + # with different compilers and therefore different OpenMP runtimes. + outer_loop_cc_var = set_cc_variables("CC_OUTER_LOOP") + openmp_flag = get_openmp_flag() + + ext_modules = [ + Extension( + "openmp_helpers_outer", + ["openmp_helpers_outer.pyx"], + extra_compile_args=openmp_flag, + extra_link_args=openmp_flag, + ) + ] + + setup( + name="_openmp_test_helper_outer", + ext_modules=cythonize( + ext_modules, + compiler_directives={ + "language_level": 3, + "boundscheck": False, + "wraparound": False, + }, + compile_time_env={"CC_OUTER_LOOP": outer_loop_cc_var or "unknown"}, + ), + ) + +finally: + os.environ.update(original_environ) diff --git a/tests/_pyMylib/__init__.py b/tests/_pyMylib/__init__.py new file mode 100644 index 00000000..92913fbc --- /dev/null +++ b/tests/_pyMylib/__init__.py @@ -0,0 +1,53 @@ +import ctypes +from pathlib import Path + +from threadpoolctl import LibController, register + +path = Path(__file__).parent / "my_threaded_lib.so" +ctypes.CDLL(path) + + +class MyThreadedLibController(LibController): + # names for threadpoolctl's context filtering + user_api = "my_threaded_lib" + internal_api = "my_threaded_lib" + + # Patterns to identify the name of the linked library to load. + # If a dynamic library with a matching filename is linked to the python + # process, it will be loaded as the `dynlib` attribute of the LibController + # instance. + filename_prefixes = ("my_threaded_lib",) + + # (Optional) Symbols that the linked library is expected to expose. It is used along + # with the `filename_prefixes` to make sure that the correct library is identified. + check_symbols = ( + "mylib_get_num_threads", + "mylib_set_num_threads", + "mylib_get_version", + ) + + def get_num_threads(self): + # This function should return the current maximum number of threads, + # which is reported as "num_threads" by `ThreadpoolController.info`. + return getattr(self.dynlib, "mylib_get_num_threads")() + + def set_num_threads(self, num_threads): + # This function limits the maximum number of threads, + # when `ThreadpoolController.limit` is called. + getattr(self.dynlib, "mylib_set_num_threads")(num_threads) + + def get_version(self): + # This function returns the version of the linked library if it is exposed, + # which is reported as "version" by `ThreadpoolController.info`. + get_version = getattr(self.dynlib, "mylib_get_version") + get_version.restype = ctypes.c_char_p + return get_version().decode("utf-8") + + def set_additional_attributes(self): + # This function is called during the initialization of the LibController. + # Additional information meant to be exposed by `ThreadpoolController.info` + # should be set here as attributes of the LibController instance. + self.some_attr = "some_value" + + +register(MyThreadedLibController) diff --git a/tests/_pyMylib/my_threaded_lib.c b/tests/_pyMylib/my_threaded_lib.c new file mode 100644 index 00000000..d38dfe70 --- /dev/null +++ b/tests/_pyMylib/my_threaded_lib.c @@ -0,0 +1,17 @@ +int NUM_THREADS = 42; +int* NUM_THREADS_p = &NUM_THREADS; + + +int mylib_get_num_threads(){ + return *NUM_THREADS_p; +} + + +void mylib_set_num_threads(int num_threads){ + *NUM_THREADS_p = num_threads; +} + + +char* mylib_get_version(){ + return "2.0"; +} diff --git a/tests/test_threadpoolctl.py b/tests/test_threadpoolctl.py new file mode 100644 index 00000000..ea27cc2b --- /dev/null +++ b/tests/test_threadpoolctl.py @@ -0,0 +1,794 @@ +import json +import os +import pytest +import re +import subprocess +import sys + +from threadpoolctl import threadpool_limits, threadpool_info +from threadpoolctl import ThreadpoolController +from threadpoolctl import _ALL_PREFIXES, _ALL_USER_APIS + +from .utils import cython_extensions_compiled +from .utils import check_nested_prange_blas +from .utils import libopenblas_paths +from .utils import scipy +from .utils import threadpool_info_from_subprocess +from .utils import select + + +def is_old_openblas(lib_controller): + # Possible bug in getting maximum number of threads with OpenBLAS < 0.2.16 + # and OpenBLAS does not expose its version before 0.3.4. + return lib_controller.internal_api == "openblas" and lib_controller.version is None + + +def skip_if_openblas_openmp(): + """Helper to skip tests with side effects when OpenBLAS has the OpenMP + threading layer. + """ + if any( + lib_controller.internal_api == "openblas" + and lib_controller.threading_layer == "openmp" + for lib_controller in ThreadpoolController().lib_controllers + ): + pytest.skip( + "Setting a limit on OpenBLAS when using the OpenMP threading layer also " + "impact the OpenMP library. They can't be controlled independently." + ) + + +def effective_num_threads(nthreads, max_threads): + if nthreads is None or nthreads > max_threads: + return max_threads + return nthreads + + +def test_threadpool_info(): + # Check consistency between threadpool_info and ThreadpoolController + function_info = threadpool_info() + object_info = ThreadpoolController().lib_controllers + + for lib_info, lib_controller in zip(function_info, object_info): + assert lib_info == lib_controller.info() + + +def test_threadpool_controller_info(): + # Check that all keys expected for the private api are in the dicts + # returned by the `info` methods + controller = ThreadpoolController() + + assert threadpool_info() == [ + lib_controller.info() for lib_controller in controller.lib_controllers + ] + assert controller.info() == [ + lib_controller.info() for lib_controller in controller.lib_controllers + ] + + for lib_controller_dict in controller.info(): + assert "user_api" in lib_controller_dict + assert "internal_api" in lib_controller_dict + assert "prefix" in lib_controller_dict + assert "filepath" in lib_controller_dict + assert "version" in lib_controller_dict + assert "num_threads" in lib_controller_dict + + if lib_controller_dict["internal_api"] in ("mkl", "blis", "openblas"): + assert "threading_layer" in lib_controller_dict + + +def test_controller_info_actualized(): + # Check that the num_threads attribute reflects the actual state of the threadpools + controller = ThreadpoolController() + original_info = controller.info() + + with threadpool_limits(limits=1): + assert all( + lib_controller.num_threads == 1 + for lib_controller in controller.lib_controllers + ) + + assert controller.info() == original_info + + +@pytest.mark.parametrize( + "kwargs", + [ + {"user_api": "blas"}, + {"prefix": "libgomp"}, + {"internal_api": "openblas", "prefix": "libomp"}, + {"prefix": ["libgomp", "libomp", "libiomp"]}, + ], +) +def test_threadpool_controller_select(kwargs): + # Check the behior of the select method of ThreadpoolController + controller = ThreadpoolController().select(**kwargs) + if not controller: + pytest.skip(f"Requires at least one of {list(kwargs.values())}.") + + for lib_controller in controller.lib_controllers: + assert any( + getattr(lib_controller, key) in (val if isinstance(val, list) else [val]) + for key, val in kwargs.items() + ) + + +@pytest.mark.parametrize("prefix", _ALL_PREFIXES) +@pytest.mark.parametrize("limit", [1, 3]) +def test_threadpool_limits_by_prefix(prefix, limit): + # Check that the maximum number of threads can be set by prefix + controller = ThreadpoolController() + original_info = controller.info() + + controller_matching_prefix = controller.select(prefix=prefix) + if not controller_matching_prefix: + pytest.skip(f"Requires {prefix} runtime") + + with threadpool_limits(limits={prefix: limit}): + for lib_controller in controller_matching_prefix.lib_controllers: + if is_old_openblas(lib_controller): + continue + # threadpool_limits only sets an upper bound on the number of + # threads. + assert 0 < lib_controller.num_threads <= limit + assert ThreadpoolController().info() == original_info + + +@pytest.mark.parametrize("user_api", (None, "blas", "openmp")) +@pytest.mark.parametrize("limit", [1, 3]) +def test_set_threadpool_limits_by_api(user_api, limit): + # Check that the maximum number of threads can be set by user_api + controller = ThreadpoolController() + original_info = controller.info() + + if user_api is None: + controller_matching_api = controller + else: + controller_matching_api = controller.select(user_api=user_api) + if not controller_matching_api: + user_apis = _ALL_USER_APIS if user_api is None else [user_api] + pytest.skip(f"Requires a library which api is in {user_apis}") + + with threadpool_limits(limits=limit, user_api=user_api): + for lib_controller in controller_matching_api.lib_controllers: + if is_old_openblas(lib_controller): + continue + # threadpool_limits only sets an upper bound on the number of + # threads. + assert 0 < lib_controller.num_threads <= limit + + assert ThreadpoolController().info() == original_info + + +def test_threadpool_limits_function_with_side_effect(): + # Check that threadpool_limits can be used as a function with + # side effects instead of a context manager. + original_info = ThreadpoolController().info() + + threadpool_limits(limits=1) + try: + for lib_controller in ThreadpoolController().lib_controllers: + if is_old_openblas(lib_controller): + continue + assert lib_controller.num_threads == 1 + finally: + # Restore the original limits so that this test does not have any + # side-effect. + threadpool_limits(limits=original_info) + + assert ThreadpoolController().info() == original_info + + +def test_set_threadpool_limits_no_limit(): + # Check that limits=None does nothing. + original_info = ThreadpoolController().info() + + with threadpool_limits(limits=None): + assert ThreadpoolController().info() == original_info + + assert ThreadpoolController().info() == original_info + + +def test_threadpool_limits_manual_restore(): + # Check that threadpool_limits can be used as an object which holds the + # original state of the threadpools and that can be restored thanks to the + # dedicated restore_original_limits method + original_info = ThreadpoolController().info() + + limits = threadpool_limits(limits=1) + try: + for lib_controller in ThreadpoolController().lib_controllers: + if is_old_openblas(lib_controller): + continue + assert lib_controller.num_threads == 1 + finally: + # Restore the original limits so that this test does not have any + # side-effect. + limits.restore_original_limits() + + assert ThreadpoolController().info() == original_info + + +def test_threadpool_controller_limit(): + # Check that using the limit method of ThreadpoolController only impact its + # library controllers. + + # This is not True for OpenBLAS with the OpenMP threading layer. + skip_if_openblas_openmp() + + blas_controller = ThreadpoolController().select(user_api="blas") + original_openmp_info = ThreadpoolController().select(user_api="openmp").info() + + with blas_controller.limit(limits=1): + blas_controller = ThreadpoolController().select(user_api="blas") + openmp_info = ThreadpoolController().select(user_api="openmp").info() + + assert all( + lib_controller.num_threads == 1 + for lib_controller in blas_controller.lib_controllers + ) + # original_blas_controller contains only blas libraries so no opemp library + # should be impacted. + assert openmp_info == original_openmp_info + + +def test_get_params_for_sequential_blas_under_openmp(): + # Test for the behavior of get_params_for_sequential_blas_under_openmp. + controller = ThreadpoolController() + original_info = controller.info() + + params = controller._get_params_for_sequential_blas_under_openmp() + + if controller.select( + internal_api="openblas", threading_layer="openmp" + ).lib_controllers: + assert params["limits"] is None + assert params["user_api"] is None + + with controller.limit(limits="sequential_blas_under_openmp"): + assert controller.info() == original_info + + else: + assert params["limits"] == 1 + assert params["user_api"] == "blas" + + with controller.limit(limits="sequential_blas_under_openmp"): + assert all( + lib_info["num_threads"] == 1 + for lib_info in controller.info() + if lib_info["user_api"] == "blas" + ) + + +def test_nested_limits(): + # Check that exiting the context manager properly restores the original limits even + # when nested. + controller = ThreadpoolController() + original_info = controller.info() + + if any(info["num_threads"] < 2 for info in original_info): + pytest.skip("Test requires at least 2 CPUs on host machine") + + def check_num_threads(expected_num_threads): + assert all( + lib_controller.num_threads == expected_num_threads + for lib_controller in ThreadpoolController().lib_controllers + ) + + with controller.limit(limits=1): + check_num_threads(expected_num_threads=1) + + with controller.limit(limits=2): + check_num_threads(expected_num_threads=2) + + check_num_threads(expected_num_threads=1) + + assert ThreadpoolController().info() == original_info + + +def test_threadpool_limits_bad_input(): + # Check that appropriate errors are raised for invalid arguments + match = re.escape(f"user_api must be either in {_ALL_USER_APIS} or None.") + with pytest.raises(ValueError, match=match): + threadpool_limits(limits=1, user_api="wrong") + + with pytest.raises( + TypeError, match="limits must either be an int, a list, a dict, or" + ): + threadpool_limits(limits=(1, 2, 3)) + + +@pytest.mark.skipif( + not cython_extensions_compiled, reason="Requires cython extensions to be compiled" +) +@pytest.mark.parametrize("num_threads", [1, 2, 4]) +def test_openmp_limit_num_threads(num_threads): + # checks that OpenMP effectively uses the number of threads requested by + # the context manager + import tests._openmp_test_helper.openmp_helpers_inner as omp_inner + + check_openmp_num_threads = omp_inner.check_openmp_num_threads + + old_num_threads = check_openmp_num_threads(100) + + with threadpool_limits(limits=num_threads): + assert check_openmp_num_threads(100) in (num_threads, old_num_threads) + assert check_openmp_num_threads(100) == old_num_threads + + +@pytest.mark.skipif( + not cython_extensions_compiled, reason="Requires cython extensions to be compiled" +) +@pytest.mark.parametrize("nthreads_outer", [None, 1, 2, 4]) +def test_openmp_nesting(nthreads_outer): + # checks that OpenMP effectively uses the number of threads requested by + # the context manager when nested in an outer OpenMP loop. + import tests._openmp_test_helper.openmp_helpers_outer as omp_outer + + check_nested_openmp_loops = omp_outer.check_nested_openmp_loops + + # Find which OpenMP lib is used at runtime for inner loop + inner_info = threadpool_info_from_subprocess( + "tests._openmp_test_helper.openmp_helpers_inner" + ) + assert len(inner_info) == 1 + inner_omp = inner_info[0]["prefix"] + + # Find which OpenMP lib is used at runtime for outer loop + outer_info = threadpool_info_from_subprocess( + "tests._openmp_test_helper.openmp_helpers_outer" + ) + if len(outer_info) == 1: + # Only 1 openmp loaded. It has to be this one. + outer_omp = outer_info[0]["prefix"] + else: + # There are 2 openmp, the one from inner and the one from outer. + assert len(outer_info) == 2 + # We already know the one from inner. It has to be the other one. + prefixes = {lib_info["prefix"] for lib_info in outer_info} + outer_omp = prefixes - {inner_omp} + + outer_num_threads, inner_num_threads = check_nested_openmp_loops(10) + original_info = ThreadpoolController().info() + + if inner_omp == outer_omp: + # The OpenMP runtime should be shared by default, meaning that the + # inner loop should automatically be run serially by the OpenMP runtime + assert inner_num_threads == 1 + + with threadpool_limits(limits=1) as threadpoolctx: + max_threads = threadpoolctx.get_original_num_threads()["openmp"] + nthreads = effective_num_threads(nthreads_outer, max_threads) + + # Ask outer loop to run on nthreads threads and inner loop run on 1 + # thread + outer_num_threads, inner_num_threads = check_nested_openmp_loops(10, nthreads) + + # The state of the original state of all threadpools should have been + # restored. + assert ThreadpoolController().info() == original_info + + # The number of threads available in the outer loop should not have been + # decreased: + assert outer_num_threads == nthreads + + # The number of threads available in the inner loop should have been + # set to 1 to avoid oversubscription and preserve performance: + if inner_omp != outer_omp: + if inner_num_threads != 1: + # XXX: this does not always work when nesting independent openmp + # implementations. See: https://github.com/jeremiedbb/Nested_OpenMP + pytest.xfail( + f"Inner OpenMP num threads was {inner_num_threads} instead of 1" + ) + assert inner_num_threads == 1 + + +def test_shipped_openblas(): + # checks that OpenBLAS effectively uses the number of threads requested by + # the context manager + original_info = ThreadpoolController().info() + openblas_controller = ThreadpoolController().select(internal_api="openblas") + + with threadpool_limits(1): + for lib_controller in openblas_controller.lib_controllers: + assert lib_controller.num_threads == 1 + + assert ThreadpoolController().info() == original_info + + +@pytest.mark.skipif( + len(libopenblas_paths) < 2, reason="need at least 2 shipped openblas library" +) +def test_multiple_shipped_openblas(): + # This redundant test is meant to make it easier to see if the system + # has 2 or more active openblas runtimes available just by reading the + # pytest report (whether or not this test has been skipped). + test_shipped_openblas() + + +@pytest.mark.skipif( + not cython_extensions_compiled, reason="Requires cython extensions to be compiled" +) +@pytest.mark.skipif( + check_nested_prange_blas is None, + reason="Requires nested_prange_blas to be compiled", +) +@pytest.mark.parametrize("nthreads_outer", [None, 1, 2, 4]) +def test_nested_prange_blas(nthreads_outer): + # Check that the BLAS uses the number of threads requested by the context manager + # when nested in an outer OpenMP loop. + # Remark: this test also works with sequential BLAS only because we limit the + # number of threads for the BLAS to 1. + import numpy as np + + skip_if_openblas_openmp() + + original_info = ThreadpoolController().info() + + blas_controller = ThreadpoolController().select(user_api="blas") + blis_controller = ThreadpoolController().select(internal_api="blis") + + # skip if the BLAS used by numpy is an old openblas. OpenBLAS 0.3.3 and + # older are known to cause an unrecoverable deadlock at process shutdown + # time (after pytest has exited). + # numpy can be linked to BLIS for CBLAS and OpenBLAS for LAPACK. In that + # case this test will run BLIS gemm so no need to skip. + if not blis_controller and any( + is_old_openblas(lib_controller) + for lib_controller in blas_controller.lib_controllers + ): + pytest.skip("Old OpenBLAS: skipping test to avoid deadlock") + + A = np.ones((1000, 10)) + B = np.ones((100, 10)) + + with threadpool_limits(limits=1) as threadpoolctx: + max_threads = threadpoolctx.get_original_num_threads()["openmp"] + nthreads = effective_num_threads(nthreads_outer, max_threads) + + result = check_nested_prange_blas(A, B, nthreads) + C, prange_num_threads, inner_info = result + + assert np.allclose(C, np.dot(A, B.T)) + assert prange_num_threads == nthreads + + nested_blas_info = select(inner_info, user_api="blas") + assert len(nested_blas_info) == len(blas_controller.lib_controllers) + assert all(lib_info["num_threads"] == 1 for lib_info in nested_blas_info) + + assert ThreadpoolController().info() == original_info + + +# the method `get_original_num_threads` raises a UserWarning due to different +# num_threads from libraries with the same `user_api`. It will be raised only +# in the CI job with 2 openblas (py38_pip_openblas_gcc_clang). It is expected +# so we can safely filter it. +@pytest.mark.filterwarnings("ignore::UserWarning") +@pytest.mark.parametrize("limit", [1, None]) +def test_get_original_num_threads(limit): + # Tests the method get_original_num_threads of the context manager + with threadpool_limits(limits=2, user_api="blas") as ctx: + # set different blas num threads to start with (when multiple openblas) + if len(ctx._controller.select(user_api="blas")) > 1: + ctx._controller.lib_controllers[0].set_num_threads(1) + + original_info = ThreadpoolController().info() + with threadpool_limits(limits=limit, user_api="blas") as threadpoolctx: + original_num_threads = threadpoolctx.get_original_num_threads() + + assert "openmp" not in original_num_threads + + blas_info = select(original_info, user_api="blas") + if blas_info: + expected = min(lib_info["num_threads"] for lib_info in blas_info) + assert original_num_threads["blas"] == expected + else: + assert original_num_threads["blas"] is None + + if len(libopenblas_paths) >= 2: + with pytest.warns(None, match="Multiple value possible"): + threadpoolctx.get_original_num_threads() + + +def test_mkl_threading_layer(): + # Check that threadpool_info correctly recovers the threading layer used + # by mkl + mkl_controller = ThreadpoolController().select(internal_api="mkl") + expected_layer = os.getenv("MKL_THREADING_LAYER") + + if not (mkl_controller and expected_layer): + pytest.skip("requires MKL and the environment variable MKL_THREADING_LAYER set") + + actual_layer = mkl_controller.lib_controllers[0].threading_layer + assert actual_layer == expected_layer.lower() + + +def test_blis_threading_layer(): + # Check that threadpool_info correctly recovers the threading layer used + # by blis + blis_controller = ThreadpoolController().select(internal_api="blis") + expected_layer = os.getenv("BLIS_ENABLE_THREADING") + if expected_layer == "no": + expected_layer = "disabled" + + if not (blis_controller and expected_layer): + pytest.skip( + "requires BLIS and the environment variable BLIS_ENABLE_THREADING set" + ) + + actual_layer = blis_controller.lib_controllers[0].threading_layer + assert actual_layer == expected_layer + + +@pytest.mark.skipif( + not cython_extensions_compiled, reason="Requires cython extensions to be compiled" +) +def test_libomp_libiomp_warning(recwarn): + # Trigger the import of a potentially clang-compiled extension: + import tests._openmp_test_helper.openmp_helpers_outer # noqa + + # Trigger the import of numpy to potentially import Intel OpenMP via MKL + pytest.importorskip("numpy.linalg") + + # Check that a warning is raised when both libomp and libiomp are loaded + # It should happen in one CI job (pylatest_conda_mkl_clang_gcc). + controller = ThreadpoolController() + prefixes = [lib_controller.prefix for lib_controller in controller.lib_controllers] + + if not ("libomp" in prefixes and "libiomp" in prefixes and sys.platform == "linux"): + pytest.skip("Requires both libomp and libiomp loaded, on Linux") + + assert len(recwarn) == 1 + wm = recwarn[0] + assert wm.category == RuntimeWarning + assert "Found Intel" in str(wm.message) + assert "LLVM" in str(wm.message) + assert "multiple_openmp.md" in str(wm.message) + + +def test_command_line_empty_or_system_openmp(): + # When the command line is called without arguments, no library should be + # detected. The only exception is a system OpenMP library that can be + # linked to the Python interpreter, for instance via the libb2.so BLAKE2 + # library that can itself be linked to an OpenMP runtime on Gentoo. + output = subprocess.check_output((sys.executable + " -m threadpoolctl").split()) + results = json.loads(output.decode("utf-8")) + conda_prefix = os.getenv("CONDA_PREFIX") + managed_by_conda = conda_prefix and sys.executable.startswith(conda_prefix) + if not managed_by_conda: # pragma: no cover + # When using a Python interpreter that does not come from a conda + # environment, we should ignore any system OpenMP library. + results = [r for r in results if r["user_api"] != "openmp"] + assert results == [] + + +def test_command_line_command_flag(): + pytest.importorskip("numpy") + output = subprocess.check_output( + [sys.executable, "-m", "threadpoolctl", "-c", "import numpy"] + ) + cli_info = json.loads(output.decode("utf-8")) + + this_process_info = threadpool_info() + for lib_info in cli_info: + assert lib_info in this_process_info + + +@pytest.mark.skipif( + sys.version_info < (3, 7), reason="need recent subprocess.run options" +) +def test_command_line_import_flag(): + result = subprocess.run( + [ + sys.executable, + "-m", + "threadpoolctl", + "-i", + "numpy", + "scipy.linalg", + "invalid_package", + "numpy.invalid_sumodule", + ], + capture_output=True, + check=True, + encoding="utf-8", + ) + cli_info = json.loads(result.stdout) + + this_process_info = threadpool_info() + for lib_info in cli_info: + assert lib_info in this_process_info + + warnings = [w.strip() for w in result.stderr.splitlines()] + assert "WARNING: could not import invalid_package" in warnings + assert "WARNING: could not import numpy.invalid_sumodule" in warnings + if scipy is None: + assert "WARNING: could not import scipy.linalg" in warnings + else: + assert "WARNING: could not import scipy.linalg" not in warnings + + +def test_architecture(): + expected_openblas_architectures = ( + # XXX: add more as needed by CI or developer laptops + "armv8", + "haswell", + "neoversen1", + "prescott", # see: https://github.com/xianyi/OpenBLAS/pull/3485 + "skylakex", + "sandybridge", + "vortex", + "zen", + ) + expected_blis_architectures = ( + # XXX: add more as needed by CI or developer laptops + "skx", + "haswell", + "zen3", + ) + for lib_info in threadpool_info(): + if lib_info["internal_api"] == "openblas": + assert lib_info["architecture"].lower() in expected_openblas_architectures + elif lib_info["internal_api"] == "blis": + assert lib_info["architecture"].lower() in expected_blis_architectures + else: + # Not supported for other libraries + assert "architecture" not in lib_info + + +def test_openblas_threading_layer(): + # Check that threadpool_info correctly recovers the threading layer used by openblas + openblas_controller = ThreadpoolController().select(internal_api="openblas") + + if not (openblas_controller): + pytest.skip("requires OpenBLAS.") + + expected_openblas_threading_layers = ("openmp", "pthreads", "disabled") + + threading_layer = openblas_controller.lib_controllers[0].threading_layer + + if threading_layer == "unknown": + # If we never recover an acceptable value for the threading layer, it will be + # always skipped and caught by check_no_test_always_skipped. + pytest.skip("Unknown OpenBLAS threading layer.") + + assert threading_layer in expected_openblas_threading_layers + + +# skip test if not run in a azure pipelines job since it relies on a specific flexiblas +# installation. +@pytest.mark.skipif( + "GITHUB_ACTIONS" not in os.environ, reason="not running in azure pipelines" +) +def test_flexiblas(): + # Check that threadpool_info correctly recovers the FlexiBLAS backends. + flexiblas_controller = ThreadpoolController().select(internal_api="flexiblas") + + if not flexiblas_controller: + pytest.skip("requires FlexiBLAS.") + flexiblas_controller = flexiblas_controller.lib_controllers[0] + + expected_backends = {"NETLIB", "OPENBLAS_CONDA"} + expected_backends_loaded = {"OPENBLAS_CONDA"} + expected_current_backend = "OPENBLAS_CONDA" # set as default at build time + + flexiblas_backends = flexiblas_controller.available_backends + flexiblas_backends_loaded = flexiblas_controller.loaded_backends + current_backend = flexiblas_controller.current_backend + + assert set(flexiblas_backends) == expected_backends + assert set(flexiblas_backends_loaded) == expected_backends_loaded + assert current_backend == expected_current_backend + + +def test_flexiblas_switch_error(): + # Check that an error is raised when trying to switch to an invalid backend. + flexiblas_controller = ThreadpoolController().select(internal_api="flexiblas") + + if not flexiblas_controller: + pytest.skip("requires FlexiBLAS.") + flexiblas_controller = flexiblas_controller.lib_controllers[0] + + with pytest.raises(RuntimeError, match="Failed to load backend"): + flexiblas_controller.switch_backend("INVALID_BACKEND") + + +# skip test if not run in a azure pipelines job since it relies on a specific flexiblas +# installation. +@pytest.mark.skipif( + "GITHUB_ACTIONS" not in os.environ, reason="not running in azure pipelines" +) +def test_flexiblas_switch(): + # Check that the backend can be switched. + controller = ThreadpoolController() + fb_controller = controller.select(internal_api="flexiblas") + + if not fb_controller: + pytest.skip("requires FlexiBLAS.") + fb_controller = fb_controller.lib_controllers[0] + + # at first mkl is not loaded in the CI jobs where this test runs + assert len(controller.select(internal_api="mkl").lib_controllers) == 0 + + # at first, only "OPENBLAS_CONDA" is loaded + assert fb_controller.current_backend == "OPENBLAS_CONDA" + assert fb_controller.loaded_backends == ["OPENBLAS_CONDA"] + + fb_controller.switch_backend("NETLIB") + assert fb_controller.current_backend == "NETLIB" + assert fb_controller.loaded_backends == ["OPENBLAS_CONDA", "NETLIB"] + + if sys.platform == "linux": + mkl_path = f"{os.getenv('CONDA_PREFIX')}/lib/libmkl_rt.so" + fb_controller.switch_backend(mkl_path) + assert fb_controller.current_backend == mkl_path + assert fb_controller.loaded_backends == ["OPENBLAS_CONDA", "NETLIB", mkl_path] + # switching the backend triggered a new search for loaded shared libs + assert len(controller.select(internal_api="mkl").lib_controllers) == 1 + + # switch back to default to avoid side effects + fb_controller.switch_backend("OPENBLAS_CONDA") + + +def test_threadpool_controller_as_decorator(): + # Check that using the decorator can be nested and is restricted to the scope of + # the decorated function. + controller = ThreadpoolController() + original_info = controller.info() + + if any(info["num_threads"] < 2 for info in original_info): + pytest.skip("Test requires at least 2 CPUs on host machine") + if not controller.select(user_api="blas"): + pytest.skip("Requires a blas runtime.") + + def check_blas_num_threads(expected_num_threads): + blas_controller = ThreadpoolController().select(user_api="blas") + assert all( + lib_controller.num_threads == expected_num_threads + for lib_controller in blas_controller.lib_controllers + ) + + @controller.wrap(limits=1, user_api="blas") + def outer_func(): + check_blas_num_threads(expected_num_threads=1) + inner_func() + check_blas_num_threads(expected_num_threads=1) + + @controller.wrap(limits=2, user_api="blas") + def inner_func(): + check_blas_num_threads(expected_num_threads=2) + + outer_func() + + assert ThreadpoolController().info() == original_info + + +def test_custom_controller(): + # Check that a custom controller can be used to change the number of threads + # used by a library. + try: + import tests._pyMylib # noqa + except: + pytest.skip("requires my_thread_lib to be compiled") + + controller = ThreadpoolController() + original_info = controller.info() + + mylib_controller = controller.select(user_api="my_threaded_lib") + + # my_threaded_lib has been found and there's 1 matching shared library + assert len(mylib_controller.lib_controllers) == 1 + mylib_controller = mylib_controller.lib_controllers[0] + + # we linked against my_threaded_lib v2.0 and by default it uses 42 thread + assert mylib_controller.version == "2.0" + assert mylib_controller.num_threads == 42 + + # my_threaded_lib exposes an additional info "some_attr": + assert mylib_controller.info()["some_attr"] == "some_value" + + with controller.limit(limits=1, user_api="my_threaded_lib"): + assert mylib_controller.num_threads == 1 + + assert ThreadpoolController().info() == original_info diff --git a/tests/utils.py b/tests/utils.py new file mode 100644 index 00000000..86fe1b0b --- /dev/null +++ b/tests/utils.py @@ -0,0 +1,88 @@ +import os +import json +import sys +import threadpoolctl +from glob import glob +from os.path import dirname, normpath +from subprocess import check_output + +# Path to shipped openblas for libraries such as numpy or scipy +libopenblas_patterns = [] + + +try: + # make sure the mkl/blas are loaded for test_threadpool_limits + import numpy as np + + np.dot(np.ones(1000), np.ones(1000)) + + libopenblas_patterns.append(os.path.join(np.__path__[0], ".libs", "libopenblas*")) +except ImportError: + pass + + +try: + import scipy + import scipy.linalg # noqa: F401 + + scipy.linalg.svd([[1, 2], [3, 4]]) + + libopenblas_patterns.append( + os.path.join(scipy.__path__[0], ".libs", "libopenblas*") + ) +except ImportError: + scipy = None + +libopenblas_paths = set( + path for pattern in libopenblas_patterns for path in glob(pattern) +) + + +try: + import tests._openmp_test_helper.openmp_helpers_inner # noqa: F401 + + cython_extensions_compiled = True +except ImportError: + cython_extensions_compiled = False + + +try: + from tests._openmp_test_helper.nested_prange_blas import check_nested_prange_blas +except ImportError: + check_nested_prange_blas = None + + +def threadpool_info_from_subprocess(module): + """Utility to call threadpool_info in a subprocess + + `module` is imported before calling threadpool_info + """ + # set PYTHONPATH to import from non sub-modules + path1 = normpath(dirname(threadpoolctl.__file__)) + path2 = os.path.join(path1, "tests", "_openmp_test_helper") + pythonpath = os.pathsep.join([path1, path2]) + env = os.environ.copy() + try: + env["PYTHONPATH"] = os.pathsep.join([pythonpath, env["PYTHONPATH"]]) + except KeyError: + env["PYTHONPATH"] = pythonpath + + cmd = [sys.executable, "-m", "threadpoolctl", "-i", module] + out = check_output(cmd, env=env).decode("utf-8") + return json.loads(out) + + +def select(info, **kwargs): + """Select a subset of the list of library info matching the request""" + # It's just a utility function to avoid repeating the pattern + # [lib_info for lib_info in info if lib_info[""] == key] + for key, vals in kwargs.items(): + kwargs[key] = [vals] if not isinstance(vals, list) else vals + + selected_info = [ + lib_info + for lib_info in info + if any(lib_info.get(key, None) in vals for key, vals in kwargs.items()) + ] + + return selected_info diff --git a/threadpoolctl.py b/threadpoolctl.py new file mode 100644 index 00000000..ceed5b88 --- /dev/null +++ b/threadpoolctl.py @@ -0,0 +1,1297 @@ +"""threadpoolctl + +This module provides utilities to introspect native libraries that relies on +thread pools (notably BLAS and OpenMP implementations) and dynamically set the +maximal number of threads they can use. +""" + +# License: BSD 3-Clause + +# The code to introspect dynamically loaded libraries on POSIX systems is +# adapted from code by Intel developer @anton-malakhov available at +# https://github.com/IntelPython/smp (Copyright (c) 2017, Intel Corporation) +# and also published under the BSD 3-Clause license +import os +import re +import sys +import ctypes +import itertools +import textwrap +from typing import final +import warnings +from ctypes.util import find_library +from abc import ABC, abstractmethod +from functools import lru_cache +from contextlib import ContextDecorator + +__version__ = "3.7.0.dev0" +__all__ = [ + "threadpool_limits", + "threadpool_info", + "ThreadpoolController", + "LibController", + "register", +] + + +# One can get runtime errors or even segfaults due to multiple OpenMP libraries +# loaded simultaneously which can happen easily in Python when importing and +# using compiled extensions built with different compilers and therefore +# different OpenMP runtimes in the same program. In particular libiomp (used by +# Intel ICC) and libomp used by clang/llvm tend to crash. This can happen for +# instance when calling BLAS inside a prange. Setting the following environment +# variable allows multiple OpenMP libraries to be loaded. It should not degrade +# performances since we manually take care of potential over-subscription +# performance issues, in sections of the code where nested OpenMP loops can +# happen, by dynamically reconfiguring the inner OpenMP runtime to temporarily +# disable it while under the scope of the outer OpenMP parallel section. +os.environ.setdefault("KMP_DUPLICATE_LIB_OK", "True") + +# Structure to cast the info on dynamically loaded library. See +# https://linux.die.net/man/3/dl_iterate_phdr for more details. +_SYSTEM_UINT = ctypes.c_uint64 if sys.maxsize > 2**32 else ctypes.c_uint32 +_SYSTEM_UINT_HALF = ctypes.c_uint32 if sys.maxsize > 2**32 else ctypes.c_uint16 + + +class _dl_phdr_info(ctypes.Structure): + _fields_ = [ + ("dlpi_addr", _SYSTEM_UINT), # Base address of object + ("dlpi_name", ctypes.c_char_p), # path to the library + ("dlpi_phdr", ctypes.c_void_p), # pointer on dlpi_headers + ("dlpi_phnum", _SYSTEM_UINT_HALF), # number of elements in dlpi_phdr + ] + + +# The RTLD_NOLOAD flag for loading shared libraries is not defined on Windows. +try: + _RTLD_NOLOAD = os.RTLD_NOLOAD +except AttributeError: + _RTLD_NOLOAD = ctypes.DEFAULT_MODE + + +class LibController(ABC): + """Abstract base class for the individual library controllers + + A library controller must expose the following class attributes: + - user_api : str + Usually the name of the library or generic specification the library + implements, e.g. "blas" is a specification with different implementations. + - internal_api : str + Usually the name of the library or concrete implementation of some + specification, e.g. "openblas" is an implementation of the "blas" + specification. + - filename_prefixes : tuple + Possible prefixes of the shared library's filename that allow to + identify the library. e.g. "libopenblas" for libopenblas.so. + + and implement the following methods: `get_num_threads`, `set_num_threads` and + `get_version`. + + Threadpoolctl loops through all the loaded shared libraries and tries to match + the filename of each library with the `filename_prefixes`. If a match is found, a + controller is instantiated and a handler to the library is stored in the `dynlib` + attribute as a `ctypes.CDLL` object. It can be used to access the necessary symbols + of the shared library to implement the above methods. + + The following information will be exposed in the info dictionary: + - user_api : standardized API, if any, or a copy of internal_api. + - internal_api : implementation-specific API. + - num_threads : the current thread limit. + - prefix : prefix of the shared library's filename. + - filepath : path to the loaded shared library. + - version : version of the library (if available). + + In addition, each library controller may expose internal API specific entries. They + must be set as attributes in the `set_additional_attributes` method. + """ + + @final + def __init__(self, *, filepath=None, prefix=None, parent=None): + """This is not meant to be overriden by subclasses.""" + self.parent = parent + self.prefix = prefix + self.filepath = filepath + self.dynlib = ctypes.CDLL(filepath, mode=_RTLD_NOLOAD) + self._symbol_prefix, self._symbol_suffix = self._find_affixes() + self.version = self.get_version() + self.set_additional_attributes() + + def info(self): + """Return relevant info wrapped in a dict""" + hidden_attrs = ("dynlib", "parent", "_symbol_prefix", "_symbol_suffix") + return { + "user_api": self.user_api, + "internal_api": self.internal_api, + "num_threads": self.num_threads, + **{k: v for k, v in vars(self).items() if k not in hidden_attrs}, + } + + def set_additional_attributes(self): + """Set additional attributes meant to be exposed in the info dict""" + + @property + def num_threads(self): + """Exposes the current thread limit as a dynamic property + + This is not meant to be used or overriden by subclasses. + """ + return self.get_num_threads() + + @abstractmethod + def get_num_threads(self): + """Return the maximum number of threads available to use""" + + @abstractmethod + def set_num_threads(self, num_threads): + """Set the maximum number of threads to use""" + + @abstractmethod + def get_version(self): + """Return the version of the shared library""" + + def _find_affixes(self): + """Return the affixes for the symbols of the shared library""" + return "", "" + + def _get_symbol(self, name): + """Return the symbol of the shared library accounding for the affixes""" + return getattr( + self.dynlib, f"{self._symbol_prefix}{name}{self._symbol_suffix}", None + ) + + +class OpenBLASController(LibController): + """Controller class for OpenBLAS""" + + user_api = "blas" + internal_api = "openblas" + filename_prefixes = ("libopenblas", "libblas", "libscipy_openblas") + + _symbol_prefixes = ("", "scipy_") + _symbol_suffixes = ("", "64_", "_64") + + # All variations of "openblas_get_num_threads", accounting for the affixes + check_symbols = tuple( + f"{prefix}openblas_get_num_threads{suffix}" + for prefix, suffix in itertools.product(_symbol_prefixes, _symbol_suffixes) + ) + + def _find_affixes(self): + for prefix, suffix in itertools.product( + self._symbol_prefixes, self._symbol_suffixes + ): + if hasattr(self.dynlib, f"{prefix}openblas_get_num_threads{suffix}"): + return prefix, suffix + + def set_additional_attributes(self): + self.threading_layer = self._get_threading_layer() + self.architecture = self._get_architecture() + + def get_num_threads(self): + get_num_threads_func = self._get_symbol("openblas_get_num_threads") + if get_num_threads_func is not None: + return get_num_threads_func() + return None + + def set_num_threads(self, num_threads): + set_num_threads_func = self._get_symbol("openblas_set_num_threads") + if set_num_threads_func is not None: + return set_num_threads_func(num_threads) + return None + + def get_version(self): + # None means OpenBLAS is not loaded or version < 0.3.4, since OpenBLAS + # did not expose its version before that. + get_version_func = self._get_symbol("openblas_get_config") + if get_version_func is not None: + get_version_func.restype = ctypes.c_char_p + config = get_version_func().split() + if config[0] == b"OpenBLAS": + return config[1].decode("utf-8") + return None + return None + + def _get_threading_layer(self): + """Return the threading layer of OpenBLAS""" + get_threading_layer_func = self._get_symbol("openblas_get_parallel") + if get_threading_layer_func is not None: + threading_layer = get_threading_layer_func() + if threading_layer == 2: + return "openmp" + elif threading_layer == 1: + return "pthreads" + return "disabled" + return "unknown" + + def _get_architecture(self): + """Return the architecture detected by OpenBLAS""" + get_architecture_func = self._get_symbol("openblas_get_corename") + if get_architecture_func is not None: + get_architecture_func.restype = ctypes.c_char_p + return get_architecture_func().decode("utf-8") + return None + + +class BLISController(LibController): + """Controller class for BLIS""" + + user_api = "blas" + internal_api = "blis" + filename_prefixes = ("libblis", "libblas") + check_symbols = ( + "bli_thread_get_num_threads", + "bli_thread_set_num_threads", + "bli_info_get_version_str", + "bli_info_get_enable_openmp", + "bli_info_get_enable_pthreads", + "bli_arch_query_id", + "bli_arch_string", + ) + + def set_additional_attributes(self): + self.threading_layer = self._get_threading_layer() + self.architecture = self._get_architecture() + + def get_num_threads(self): + get_func = getattr(self.dynlib, "bli_thread_get_num_threads", lambda: None) + num_threads = get_func() + # by default BLIS is single-threaded and get_num_threads + # returns -1. We map it to 1 for consistency with other libraries. + return 1 if num_threads == -1 else num_threads + + def set_num_threads(self, num_threads): + set_func = getattr( + self.dynlib, "bli_thread_set_num_threads", lambda num_threads: None + ) + return set_func(num_threads) + + def get_version(self): + get_version_ = getattr(self.dynlib, "bli_info_get_version_str", None) + if get_version_ is None: + return None + + get_version_.restype = ctypes.c_char_p + return get_version_().decode("utf-8") + + def _get_threading_layer(self): + """Return the threading layer of BLIS""" + if getattr(self.dynlib, "bli_info_get_enable_openmp", lambda: False)(): + return "openmp" + elif getattr(self.dynlib, "bli_info_get_enable_pthreads", lambda: False)(): + return "pthreads" + return "disabled" + + def _get_architecture(self): + """Return the architecture detected by BLIS""" + bli_arch_query_id = getattr(self.dynlib, "bli_arch_query_id", None) + bli_arch_string = getattr(self.dynlib, "bli_arch_string", None) + if bli_arch_query_id is None or bli_arch_string is None: + return None + + # the true restype should be BLIS' arch_t (enum) but int should work + # for us: + bli_arch_query_id.restype = ctypes.c_int + bli_arch_string.restype = ctypes.c_char_p + return bli_arch_string(bli_arch_query_id()).decode("utf-8") + + +class FlexiBLASController(LibController): + """Controller class for FlexiBLAS""" + + user_api = "blas" + internal_api = "flexiblas" + filename_prefixes = ("libflexiblas",) + check_symbols = ( + "flexiblas_get_num_threads", + "flexiblas_set_num_threads", + "flexiblas_get_version", + "flexiblas_list", + "flexiblas_list_loaded", + "flexiblas_current_backend", + ) + + @property + def loaded_backends(self): + return self._get_backend_list(loaded=True) + + @property + def current_backend(self): + return self._get_current_backend() + + def info(self): + """Return relevant info wrapped in a dict""" + # We override the info method because the loaded and current backends + # are dynamic properties + exposed_attrs = super().info() + exposed_attrs["loaded_backends"] = self.loaded_backends + exposed_attrs["current_backend"] = self.current_backend + + return exposed_attrs + + def set_additional_attributes(self): + self.available_backends = self._get_backend_list(loaded=False) + + def get_num_threads(self): + get_func = getattr(self.dynlib, "flexiblas_get_num_threads", lambda: None) + num_threads = get_func() + # by default BLIS is single-threaded and get_num_threads + # returns -1. We map it to 1 for consistency with other libraries. + return 1 if num_threads == -1 else num_threads + + def set_num_threads(self, num_threads): + set_func = getattr( + self.dynlib, "flexiblas_set_num_threads", lambda num_threads: None + ) + return set_func(num_threads) + + def get_version(self): + get_version_ = getattr(self.dynlib, "flexiblas_get_version", None) + if get_version_ is None: + return None + + major = ctypes.c_int() + minor = ctypes.c_int() + patch = ctypes.c_int() + get_version_(ctypes.byref(major), ctypes.byref(minor), ctypes.byref(patch)) + return f"{major.value}.{minor.value}.{patch.value}" + + def _get_backend_list(self, loaded=False): + """Return the list of available backends for FlexiBLAS. + + If loaded is False, return the list of available backends from the FlexiBLAS + configuration. If loaded is True, return the list of actually loaded backends. + """ + func_name = f"flexiblas_list{'_loaded' if loaded else ''}" + get_backend_list_ = getattr(self.dynlib, func_name, None) + if get_backend_list_ is None: + return None + + n_backends = get_backend_list_(None, 0, 0) + + backends = [] + for i in range(n_backends): + backend_name = ctypes.create_string_buffer(1024) + get_backend_list_(backend_name, 1024, i) + if backend_name.value.decode("utf-8") != "__FALLBACK__": + # We don't know when to expect __FALLBACK__ but it is not a real + # backend and does not show up when running flexiblas list. + backends.append(backend_name.value.decode("utf-8")) + return backends + + def _get_current_backend(self): + """Return the backend of FlexiBLAS""" + get_backend_ = getattr(self.dynlib, "flexiblas_current_backend", None) + if get_backend_ is None: + return None + + backend = ctypes.create_string_buffer(1024) + get_backend_(backend, ctypes.sizeof(backend)) + return backend.value.decode("utf-8") + + def switch_backend(self, backend): + """Switch the backend of FlexiBLAS + + Parameters + ---------- + backend : str + The name or the path to the shared library of the backend to switch to. If + the backend is not already loaded, it will be loaded first. + """ + if backend not in self.loaded_backends: + if backend in self.available_backends: + load_func = getattr(self.dynlib, "flexiblas_load_backend", lambda _: -1) + else: # assume backend is a path to a shared library + load_func = getattr( + self.dynlib, "flexiblas_load_backend_library", lambda _: -1 + ) + res = load_func(str(backend).encode("utf-8")) + if res == -1: + raise RuntimeError( + f"Failed to load backend {backend!r}. It must either be the name of" + " a backend available in the FlexiBLAS configuration " + f"{self.available_backends} or the path to a valid shared library." + ) + + # Trigger a new search of loaded shared libraries since loading a new + # backend caused a dlopen. + self.parent._load_libraries() + + switch_func = getattr(self.dynlib, "flexiblas_switch", lambda _: -1) + idx = self.loaded_backends.index(backend) + res = switch_func(idx) + if res == -1: + raise RuntimeError(f"Failed to switch to backend {backend!r}.") + + +class MKLController(LibController): + """Controller class for MKL""" + + user_api = "blas" + internal_api = "mkl" + filename_prefixes = ("libmkl_rt", "mkl_rt", "libblas") + check_symbols = ( + "MKL_Get_Max_Threads", + "MKL_Set_Num_Threads", + "MKL_Get_Version_String", + "MKL_Set_Threading_Layer", + ) + + def set_additional_attributes(self): + self.threading_layer = self._get_threading_layer() + + def get_num_threads(self): + get_func = getattr(self.dynlib, "MKL_Get_Max_Threads", lambda: None) + return get_func() + + def set_num_threads(self, num_threads): + set_func = getattr(self.dynlib, "MKL_Set_Num_Threads", lambda num_threads: None) + return set_func(num_threads) + + def get_version(self): + if not hasattr(self.dynlib, "MKL_Get_Version_String"): + return None + + res = ctypes.create_string_buffer(200) + self.dynlib.MKL_Get_Version_String(res, 200) + + version = res.value.decode("utf-8") + group = re.search(r"Version ([^ ]+) ", version) + if group is not None: + version = group.groups()[0] + return version.strip() + + def _get_threading_layer(self): + """Return the threading layer of MKL""" + # The function mkl_set_threading_layer returns the current threading + # layer. Calling it with an invalid threading layer allows us to safely + # get the threading layer + set_threading_layer = getattr( + self.dynlib, "MKL_Set_Threading_Layer", lambda layer: -1 + ) + layer_map = { + 0: "intel", + 1: "sequential", + 2: "pgi", + 3: "gnu", + 4: "tbb", + -1: "not specified", + } + return layer_map[set_threading_layer(-1)] + + +class OpenMPController(LibController): + """Controller class for OpenMP""" + + user_api = "openmp" + internal_api = "openmp" + filename_prefixes = ("libiomp", "libgomp", "libomp", "vcomp") + check_symbols = ( + "omp_get_max_threads", + "omp_get_num_threads", + ) + + def get_num_threads(self): + get_func = getattr(self.dynlib, "omp_get_max_threads", lambda: None) + return get_func() + + def set_num_threads(self, num_threads): + set_func = getattr(self.dynlib, "omp_set_num_threads", lambda num_threads: None) + return set_func(num_threads) + + def get_version(self): + # There is no way to get the version number programmatically in OpenMP. + return None + + +# Controllers for the libraries that we'll look for in the loaded libraries. +# Third party libraries can register their own controllers. +_ALL_CONTROLLERS = [ + OpenBLASController, + BLISController, + MKLController, + OpenMPController, + FlexiBLASController, +] + +# Helpers for the doc and test names +_ALL_USER_APIS = list(set(lib.user_api for lib in _ALL_CONTROLLERS)) +_ALL_INTERNAL_APIS = [lib.internal_api for lib in _ALL_CONTROLLERS] +_ALL_PREFIXES = list( + set(prefix for lib in _ALL_CONTROLLERS for prefix in lib.filename_prefixes) +) +_ALL_BLAS_LIBRARIES = [ + lib.internal_api for lib in _ALL_CONTROLLERS if lib.user_api == "blas" +] +_ALL_OPENMP_LIBRARIES = OpenMPController.filename_prefixes + + +def register(controller): + """Register a new controller""" + _ALL_CONTROLLERS.append(controller) + _ALL_USER_APIS.append(controller.user_api) + _ALL_INTERNAL_APIS.append(controller.internal_api) + _ALL_PREFIXES.extend(controller.filename_prefixes) + + +def _format_docstring(*args, **kwargs): + def decorator(o): + if o.__doc__ is not None: + o.__doc__ = o.__doc__.format(*args, **kwargs) + return o + + return decorator + + +@lru_cache(maxsize=10000) +def _realpath(filepath): + """Small caching wrapper around os.path.realpath to limit system calls""" + return os.path.realpath(filepath) + + +@_format_docstring(USER_APIS=list(_ALL_USER_APIS), INTERNAL_APIS=_ALL_INTERNAL_APIS) +def threadpool_info(): + """Return the maximal number of threads for each detected library. + + Return a list with all the supported libraries that have been found. Each + library is represented by a dict with the following information: + + - "user_api" : user API. Possible values are {USER_APIS}. + - "internal_api": internal API. Possible values are {INTERNAL_APIS}. + - "prefix" : filename prefix of the specific implementation. + - "filepath": path to the loaded library. + - "version": version of the library (if available). + - "num_threads": the current thread limit. + + In addition, each library may contain internal_api specific entries. + """ + return ThreadpoolController().info() + + +class _ThreadpoolLimiter: + """The guts of ThreadpoolController.limit + + Refer to the docstring of ThreadpoolController.limit for more details. + + It will only act on the library controllers held by the provided `controller`. + Using the default constructor sets the limits right away such that it can be used as + a callable. Setting the limits can be delayed by using the `wrap` class method such + that it can be used as a decorator. + """ + + def __init__(self, controller, *, limits=None, user_api=None): + self._controller = controller + self._limits, self._user_api, self._prefixes = self._check_params( + limits, user_api + ) + self._original_info = self._controller.info() + self._set_threadpool_limits() + + def __enter__(self): + return self + + def __exit__(self, type, value, traceback): + self.restore_original_limits() + + @classmethod + def wrap(cls, controller, *, limits=None, user_api=None): + """Return an instance of this class that can be used as a decorator""" + return _ThreadpoolLimiterDecorator( + controller=controller, limits=limits, user_api=user_api + ) + + def restore_original_limits(self): + """Set the limits back to their original values""" + for lib_controller, original_info in zip( + self._controller.lib_controllers, self._original_info + ): + lib_controller.set_num_threads(original_info["num_threads"]) + + # Alias of `restore_original_limits` for backward compatibility + unregister = restore_original_limits + + def get_original_num_threads(self): + """Original num_threads from before calling threadpool_limits + + Return a dict `{user_api: num_threads}`. + """ + num_threads = {} + warning_apis = [] + + for user_api in self._user_api: + limits = [ + lib_info["num_threads"] + for lib_info in self._original_info + if lib_info["user_api"] == user_api + ] + limits = set(limits) + n_limits = len(limits) + + if n_limits == 1: + limit = limits.pop() + elif n_limits == 0: + limit = None + else: + limit = min(limits) + warning_apis.append(user_api) + + num_threads[user_api] = limit + + if warning_apis: + warnings.warn( + "Multiple value possible for following user apis: " + + ", ".join(warning_apis) + + ". Returning the minimum." + ) + + return num_threads + + def _check_params(self, limits, user_api): + """Suitable values for the _limits, _user_api and _prefixes attributes""" + + if isinstance(limits, str) and limits == "sequential_blas_under_openmp": + ( + limits, + user_api, + ) = self._controller._get_params_for_sequential_blas_under_openmp().values() + + if limits is None or isinstance(limits, int): + if user_api is None: + user_api = _ALL_USER_APIS + elif user_api in _ALL_USER_APIS: + user_api = [user_api] + else: + raise ValueError( + f"user_api must be either in {_ALL_USER_APIS} or None. Got " + f"{user_api} instead." + ) + + if limits is not None: + limits = {api: limits for api in user_api} + prefixes = [] + else: + if isinstance(limits, list): + # This should be a list of dicts of library info, for + # compatibility with the result from threadpool_info. + limits = { + lib_info["prefix"]: lib_info["num_threads"] for lib_info in limits + } + elif isinstance(limits, ThreadpoolController): + # To set the limits from the library controllers of a + # ThreadpoolController object. + limits = { + lib_controller.prefix: lib_controller.num_threads + for lib_controller in limits.lib_controllers + } + + if not isinstance(limits, dict): + raise TypeError( + "limits must either be an int, a list, a dict, or " + f"'sequential_blas_under_openmp'. Got {type(limits)} instead" + ) + + # With a dictionary, can set both specific limit for given + # libraries and global limit for user_api. Fetch each separately. + prefixes = [prefix for prefix in limits if prefix in _ALL_PREFIXES] + user_api = [api for api in limits if api in _ALL_USER_APIS] + + return limits, user_api, prefixes + + def _set_threadpool_limits(self): + """Change the maximal number of threads in selected thread pools. + + Return a list with all the supported libraries that have been found + matching `self._prefixes` and `self._user_api`. + """ + if self._limits is None: + return + + for lib_controller in self._controller.lib_controllers: + # self._limits is a dict {key: num_threads} where key is either + # a prefix or a user_api. If a library matches both, the limit + # corresponding to the prefix is chosen. + if lib_controller.prefix in self._limits: + num_threads = self._limits[lib_controller.prefix] + elif lib_controller.user_api in self._limits: + num_threads = self._limits[lib_controller.user_api] + else: + continue + + if num_threads is not None: + lib_controller.set_num_threads(num_threads) + + +class _ThreadpoolLimiterDecorator(_ThreadpoolLimiter, ContextDecorator): + """Same as _ThreadpoolLimiter but to be used as a decorator""" + + def __init__(self, controller, *, limits=None, user_api=None): + self._limits, self._user_api, self._prefixes = self._check_params( + limits, user_api + ) + self._controller = controller + + def __enter__(self): + # we need to set the limits here and not in the __init__ because we want the + # limits to be set when calling the decorated function, not when creating the + # decorator. + self._original_info = self._controller.info() + self._set_threadpool_limits() + return self + + +@_format_docstring( + USER_APIS=", ".join(f'"{api}"' for api in _ALL_USER_APIS), + BLAS_LIBS=", ".join(_ALL_BLAS_LIBRARIES), + OPENMP_LIBS=", ".join(_ALL_OPENMP_LIBRARIES), +) +class threadpool_limits(_ThreadpoolLimiter): + """Change the maximal number of threads that can be used in thread pools. + + This object can be used either as a callable (the construction of this object + limits the number of threads), as a context manager in a `with` block to + automatically restore the original state of the controlled libraries when exiting + the block, or as a decorator through its `wrap` method. + + Set the maximal number of threads that can be used in thread pools used in + the supported libraries to `limit`. This function works for libraries that + are already loaded in the interpreter and can be changed dynamically. + + This effect is global and impacts the whole Python process. There is no thread level + isolation as these libraries do not offer thread-local APIs to configure the number + of threads to use in nested parallel calls. + + Parameters + ---------- + limits : int, dict, 'sequential_blas_under_openmp' or None (default=None) + The maximal number of threads that can be used in thread pools + + - If int, sets the maximum number of threads to `limits` for each + library selected by `user_api`. + + - If it is a dictionary `{{key: max_threads}}`, this function sets a + custom maximum number of threads for each `key` which can be either a + `user_api` or a `prefix` for a specific library. + + - If 'sequential_blas_under_openmp', it will chose the appropriate `limits` + and `user_api` parameters for the specific use case of sequential BLAS + calls within an OpenMP parallel region. The `user_api` parameter is + ignored. + + - If None, this function does not do anything. + + user_api : {USER_APIS} or None (default=None) + APIs of libraries to limit. Used only if `limits` is an int. + + - If "blas", it will only limit BLAS supported libraries ({BLAS_LIBS}). + + - If "openmp", it will only limit OpenMP supported libraries + ({OPENMP_LIBS}). Note that it can affect the number of threads used + by the BLAS libraries if they rely on OpenMP. + + - If None, this function will apply to all supported libraries. + """ + + def __init__(self, limits=None, user_api=None): + super().__init__(ThreadpoolController(), limits=limits, user_api=user_api) + + @classmethod + def wrap(cls, limits=None, user_api=None): + return super().wrap(ThreadpoolController(), limits=limits, user_api=user_api) + + +class ThreadpoolController: + """Collection of LibController objects for all loaded supported libraries + + Attributes + ---------- + lib_controllers : list of `LibController` objects + The list of library controllers of all loaded supported libraries. + """ + + # Cache for libc under POSIX and a few system libraries under Windows. + # We use a class level cache instead of an instance level cache because + # it's very unlikely that a shared library will be unloaded and reloaded + # during the lifetime of a program. + _system_libraries = dict() + + def __init__(self): + self.lib_controllers = [] + self._load_libraries() + self._warn_if_incompatible_openmp() + + @classmethod + def _from_controllers(cls, lib_controllers): + new_controller = cls.__new__(cls) + new_controller.lib_controllers = lib_controllers + return new_controller + + def info(self): + """Return lib_controllers info as a list of dicts""" + return [lib_controller.info() for lib_controller in self.lib_controllers] + + def select(self, **kwargs): + """Return a ThreadpoolController containing a subset of its current + library controllers + + It will select all libraries matching at least one pair (key, value) from kwargs + where key is an entry of the library info dict (like "user_api", "internal_api", + "prefix", ...) and value is the value or a list of acceptable values for that + entry. + + For instance, `ThreadpoolController().select(internal_api=["blis", "openblas"])` + will select all library controllers whose internal_api is either "blis" or + "openblas". + """ + for key, vals in kwargs.items(): + kwargs[key] = [vals] if not isinstance(vals, list) else vals + + lib_controllers = [ + lib_controller + for lib_controller in self.lib_controllers + if any( + getattr(lib_controller, key, None) in vals + for key, vals in kwargs.items() + ) + ] + + return ThreadpoolController._from_controllers(lib_controllers) + + def _get_params_for_sequential_blas_under_openmp(self): + """Return appropriate params to use for a sequential BLAS call in an OpenMP loop + + This function takes into account the unexpected behavior of OpenBLAS with the + OpenMP threading layer. + """ + if self.select( + internal_api="openblas", threading_layer="openmp" + ).lib_controllers: + return {"limits": None, "user_api": None} + return {"limits": 1, "user_api": "blas"} + + @_format_docstring( + USER_APIS=", ".join('"{}"'.format(api) for api in _ALL_USER_APIS), + BLAS_LIBS=", ".join(_ALL_BLAS_LIBRARIES), + OPENMP_LIBS=", ".join(_ALL_OPENMP_LIBRARIES), + ) + def limit(self, *, limits=None, user_api=None): + """Change the maximal number of threads that can be used in thread pools. + + This function returns an object that can be used either as a callable (the + construction of this object limits the number of threads) or as a context + manager, in a `with` block to automatically restore the original state of the + controlled libraries when exiting the block. + + Set the maximal number of threads that can be used in thread pools used in + the supported libraries to `limits`. This function works for libraries that + are already loaded in the interpreter and can be changed dynamically. + + This effect is global and impacts the whole Python process. There is no thread + level isolation as these libraries do not offer thread-local APIs to configure + the number of threads to use in nested parallel calls. + + Parameters + ---------- + limits : int, dict, 'sequential_blas_under_openmp' or None (default=None) + The maximal number of threads that can be used in thread pools + + - If int, sets the maximum number of threads to `limits` for each + library selected by `user_api`. + + - If it is a dictionary `{{key: max_threads}}`, this function sets a + custom maximum number of threads for each `key` which can be either a + `user_api` or a `prefix` for a specific library. + + - If 'sequential_blas_under_openmp', it will chose the appropriate `limits` + and `user_api` parameters for the specific use case of sequential BLAS + calls within an OpenMP parallel region. The `user_api` parameter is + ignored. + + - If None, this function does not do anything. + + user_api : {USER_APIS} or None (default=None) + APIs of libraries to limit. Used only if `limits` is an int. + + - If "blas", it will only limit BLAS supported libraries ({BLAS_LIBS}). + + - If "openmp", it will only limit OpenMP supported libraries + ({OPENMP_LIBS}). Note that it can affect the number of threads used + by the BLAS libraries if they rely on OpenMP. + + - If None, this function will apply to all supported libraries. + """ + return _ThreadpoolLimiter(self, limits=limits, user_api=user_api) + + @_format_docstring( + USER_APIS=", ".join('"{}"'.format(api) for api in _ALL_USER_APIS), + BLAS_LIBS=", ".join(_ALL_BLAS_LIBRARIES), + OPENMP_LIBS=", ".join(_ALL_OPENMP_LIBRARIES), + ) + def wrap(self, *, limits=None, user_api=None): + """Change the maximal number of threads that can be used in thread pools. + + This function returns an object that can be used as a decorator. + + Set the maximal number of threads that can be used in thread pools used in + the supported libraries to `limits`. This function works for libraries that + are already loaded in the interpreter and can be changed dynamically. + + Parameters + ---------- + limits : int, dict or None (default=None) + The maximal number of threads that can be used in thread pools + + - If int, sets the maximum number of threads to `limits` for each + library selected by `user_api`. + + - If it is a dictionary `{{key: max_threads}}`, this function sets a + custom maximum number of threads for each `key` which can be either a + `user_api` or a `prefix` for a specific library. + + - If None, this function does not do anything. + + user_api : {USER_APIS} or None (default=None) + APIs of libraries to limit. Used only if `limits` is an int. + + - If "blas", it will only limit BLAS supported libraries ({BLAS_LIBS}). + + - If "openmp", it will only limit OpenMP supported libraries + ({OPENMP_LIBS}). Note that it can affect the number of threads used + by the BLAS libraries if they rely on OpenMP. + + - If None, this function will apply to all supported libraries. + """ + return _ThreadpoolLimiter.wrap(self, limits=limits, user_api=user_api) + + def __len__(self): + return len(self.lib_controllers) + + def _load_libraries(self): + """Loop through loaded shared libraries and store the supported ones""" + if sys.platform == "darwin": + self._find_libraries_with_dyld() + elif sys.platform == "win32": + self._find_libraries_with_enum_process_module_ex() + elif "pyodide" in sys.modules: + self._find_libraries_pyodide() + else: + self._find_libraries_with_dl_iterate_phdr() + + def _find_libraries_with_dl_iterate_phdr(self): + """Loop through loaded libraries and return binders on supported ones + + This function is expected to work on POSIX system only. + This code is adapted from code by Intel developer @anton-malakhov + available at https://github.com/IntelPython/smp + + Copyright (c) 2017, Intel Corporation published under the BSD 3-Clause + license + """ + libc = self._get_libc() + if not hasattr(libc, "dl_iterate_phdr"): # pragma: no cover + warnings.warn( + "Could not find dl_iterate_phdr in the C standard library.", + RuntimeWarning, + ) + return [] + + # Callback function for `dl_iterate_phdr` which is called for every + # library loaded in the current process until it returns 1. + def match_library_callback(info, size, data): + # Get the path of the current library + filepath = info.contents.dlpi_name + if filepath: + filepath = filepath.decode("utf-8") + + # Store the library controller if it is supported and selected + self._make_controller_from_path(filepath) + return 0 + + c_func_signature = ctypes.CFUNCTYPE( + ctypes.c_int, # Return type + ctypes.POINTER(_dl_phdr_info), + ctypes.c_size_t, + ctypes.c_char_p, + ) + c_match_library_callback = c_func_signature(match_library_callback) + + data = ctypes.c_char_p(b"") + libc.dl_iterate_phdr(c_match_library_callback, data) + + def _find_libraries_with_dyld(self): + """Loop through loaded libraries and return binders on supported ones + + This function is expected to work on OSX system only + """ + libc = self._get_libc() + if not hasattr(libc, "_dyld_image_count"): # pragma: no cover + warnings.warn( + "Could not find _dyld_image_count in the C standard library.", + RuntimeWarning, + ) + return [] + + n_dyld = libc._dyld_image_count() + libc._dyld_get_image_name.restype = ctypes.c_char_p + + for i in range(n_dyld): + filepath = ctypes.string_at(libc._dyld_get_image_name(i)) + filepath = filepath.decode("utf-8") + + # Store the library controller if it is supported and selected + self._make_controller_from_path(filepath) + + def _find_libraries_with_enum_process_module_ex(self): + """Loop through loaded libraries and return binders on supported ones + + This function is expected to work on windows system only. + This code is adapted from code by Philipp Hagemeister @phihag available + at https://stackoverflow.com/questions/17474574 + """ + from ctypes.wintypes import DWORD, HMODULE, MAX_PATH + + PROCESS_QUERY_INFORMATION = 0x0400 + PROCESS_VM_READ = 0x0010 + + LIST_LIBRARIES_ALL = 0x03 + + ps_api = self._get_windll("Psapi") + kernel_32 = self._get_windll("kernel32") + + h_process = kernel_32.OpenProcess( + PROCESS_QUERY_INFORMATION | PROCESS_VM_READ, False, os.getpid() + ) + if not h_process: # pragma: no cover + raise OSError(f"Could not open PID {os.getpid()}") + + try: + buf_count = 256 + needed = DWORD() + # Grow the buffer until it becomes large enough to hold all the + # module headers + while True: + buf = (HMODULE * buf_count)() + buf_size = ctypes.sizeof(buf) + if not ps_api.EnumProcessModulesEx( + h_process, + ctypes.byref(buf), + buf_size, + ctypes.byref(needed), + LIST_LIBRARIES_ALL, + ): + raise OSError("EnumProcessModulesEx failed") + if buf_size >= needed.value: + break + buf_count = needed.value // (buf_size // buf_count) + + count = needed.value // (buf_size // buf_count) + h_modules = map(HMODULE, buf[:count]) + + # Loop through all the module headers and get the library path + # Allocate a buffer for the path 10 times the size of MAX_PATH to take + # into account long path names. + max_path = 10 * MAX_PATH + buf = ctypes.create_unicode_buffer(max_path) + n_size = DWORD() + for h_module in h_modules: + # Get the path of the current module + if not ps_api.GetModuleFileNameExW( + h_process, h_module, ctypes.byref(buf), ctypes.byref(n_size) + ): + raise OSError("GetModuleFileNameEx failed") + filepath = buf.value + + if len(filepath) == max_path: # pragma: no cover + warnings.warn( + "Could not get the full path of a dynamic library (path too " + "long). This library will be ignored and threadpoolctl might " + "not be able to control or display information about all " + f"loaded libraries. Here's the truncated path: {filepath!r}", + RuntimeWarning, + ) + else: + # Store the library controller if it is supported and selected + self._make_controller_from_path(filepath) + finally: + kernel_32.CloseHandle(h_process) + + def _find_libraries_pyodide(self): + """Pyodide specific implementation for finding loaded libraries. + + Adapted from suggestion in https://github.com/joblib/threadpoolctl/pull/169#issuecomment-1946696449. + + One day, we may have a simpler solution. libc dl_iterate_phdr needs to + be implemented in Emscripten and exposed in Pyodide, see + https://github.com/emscripten-core/emscripten/issues/21354 for more + details. + """ + try: + from pyodide_js._module import LDSO + except ImportError: + warnings.warn( + "Unable to import LDSO from pyodide_js._module. This should never " + "happen." + ) + return + + if hasattr(LDSO.loadedLibsByName, "as_py_json"): # Pyodide >= 0.29 + libs_iter = LDSO.loadedLibsByName.as_py_json() + else: + libs_iter = LDSO.loadedLibsByName.as_object_map() # Pyodide < 0.29 + + for filepath in libs_iter: + # Some libraries are duplicated by Pyodide and do not exist in the + # filesystem, so we first check for the existence of the file. For + # more details, see + # https://github.com/joblib/threadpoolctl/pull/169#issuecomment-1947946728 + if os.path.exists(filepath): + self._make_controller_from_path(filepath) + + def _make_controller_from_path(self, filepath): + """Store a library controller if it is supported and selected""" + # Required to resolve symlinks + filepath = _realpath(filepath) + # `lower` required to take account of OpenMP dll case on Windows + # (vcomp, VCOMP, Vcomp, ...) + filename = os.path.basename(filepath).lower() + + # Loop through supported libraries to find if this filename corresponds + # to a supported one. + for controller_class in _ALL_CONTROLLERS: + # check if filename matches a supported prefix + prefix = self._check_prefix(filename, controller_class.filename_prefixes) + + # filename does not match any of the prefixes of the candidate + # library. move to next library. + if prefix is None: + continue + + # workaround for BLAS libraries packaged by conda-forge on windows, which + # are all renamed "libblas.dll". We thus have to check to which BLAS + # implementation it actually corresponds looking for implementation + # specific symbols. + if prefix == "libblas": + if filename.endswith(".dll"): + libblas = ctypes.CDLL(filepath, _RTLD_NOLOAD) + if not any( + hasattr(libblas, func) + for func in controller_class.check_symbols + ): + continue + else: + # We ignore libblas on other platforms than windows because there + # might be a libblas dso comming with openblas for instance that + # can't be used to instantiate a pertinent LibController (many + # symbols are missing) and would create confusion by making a + # duplicate entry in threadpool_info. + continue + + # filename matches a prefix. Now we check if the library has the symbols we + # are looking for. If none of the symbols exists, it's very likely not the + # expected library (e.g. a library having a common prefix with one of the + # our supported libraries). Otherwise, create and store the library + # controller. + lib_controller = controller_class( + filepath=filepath, prefix=prefix, parent=self + ) + + if filepath in (lib.filepath for lib in self.lib_controllers): + # We already have a controller for this library. + continue + + if not hasattr(controller_class, "check_symbols") or any( + hasattr(lib_controller.dynlib, func) + for func in controller_class.check_symbols + ): + self.lib_controllers.append(lib_controller) + + def _check_prefix(self, library_basename, filename_prefixes): + """Return the prefix library_basename starts with + + Return None if none matches. + """ + for prefix in filename_prefixes: + if library_basename.startswith(prefix): + return prefix + return None + + def _warn_if_incompatible_openmp(self): + """Raise a warning if llvm-OpenMP and intel-OpenMP are both loaded""" + prefixes = [lib_controller.prefix for lib_controller in self.lib_controllers] + msg = textwrap.dedent( + """ + Found Intel OpenMP ('libiomp') and LLVM OpenMP ('libomp') loaded at + the same time. Both libraries are known to be incompatible and this + can cause random crashes or deadlocks on Linux when loaded in the + same Python program. + Using threadpoolctl may cause crashes or deadlocks. For more + information and possible workarounds, please see + https://github.com/joblib/threadpoolctl/blob/master/multiple_openmp.md + """ + ) + if "libomp" in prefixes and "libiomp" in prefixes: + warnings.warn(msg, RuntimeWarning) + + @classmethod + def _get_libc(cls): + """Load the lib-C for unix systems.""" + libc = cls._system_libraries.get("libc") + if libc is None: + # Remark: If libc is statically linked or if Python is linked against an + # alternative implementation of libc like musl, find_library will return + # None and CDLL will load the main program itself which should contain the + # libc symbols. We still name it libc for convenience. + # If the main program does not contain the libc symbols, it's ok because + # we check their presence later anyway. + libc = ctypes.CDLL(find_library("c"), mode=_RTLD_NOLOAD) + cls._system_libraries["libc"] = libc + return libc + + @classmethod + def _get_windll(cls, dll_name): + """Load a windows DLL""" + dll = cls._system_libraries.get(dll_name) + if dll is None: + dll = ctypes.WinDLL(f"{dll_name}.dll") + cls._system_libraries[dll_name] = dll + return dll + + +def _main(): + """Commandline interface to display thread-pool information and exit.""" + import argparse + import importlib + import json + import sys + + parser = argparse.ArgumentParser( + usage="python -m threadpoolctl -i numpy scipy.linalg xgboost", + description="Display thread-pool information and exit.", + ) + parser.add_argument( + "-i", + "--import", + dest="modules", + nargs="*", + default=(), + help="Python modules to import before introspecting thread-pools.", + ) + parser.add_argument( + "-c", + "--command", + help="a Python statement to execute before introspecting thread-pools.", + ) + + options = parser.parse_args(sys.argv[1:]) + for module in options.modules: + try: + importlib.import_module(module, package=None) + except ImportError: + print("WARNING: could not import", module, file=sys.stderr) + + if options.command: + exec(options.command) + + print(json.dumps(threadpool_info(), indent=2)) + + +if __name__ == "__main__": + _main() diff --git a/threadpoolctl/__init__.py b/threadpoolctl/__init__.py deleted file mode 100644 index cae49eb1..00000000 --- a/threadpoolctl/__init__.py +++ /dev/null @@ -1,7 +0,0 @@ -from ._threadpool_limiters import threadpool_limits -from ._threadpool_limiters import get_threadpool_limits - -__version__ = '0.1.dev0' - - -__all__ = ["threadpool_limits", "get_threadpool_limits"] diff --git a/threadpoolctl/_threadpool_limiters.py b/threadpoolctl/_threadpool_limiters.py deleted file mode 100644 index 68a0bea3..00000000 --- a/threadpoolctl/_threadpool_limiters.py +++ /dev/null @@ -1,482 +0,0 @@ - -############################################################################# -# The following provides utilities to load C-libraries that relies on thread -# pools and limit the maximal number of thread that can be used. -# -# -import os -import re -import sys -import ctypes -from ctypes.util import find_library - -from .utils import _format_docstring - - -if sys.platform == "darwin": - - # On OSX, we can get a runtime error due to multiple OpenMP libraries - # loaded simultaneously. This can happen for instance when calling BLAS - # inside a prange. Setting the following environment variable allows - # multiple OpenMP libraries to be loaded. It should not degrade - # performances since we manually take care of potential over-subscription - # performance issues, in sections of the code where nested OpenMP loops can - # happen, by dynamically reconfiguring the inner OpenMP runtime to - # temporarily disable it while under the scope of the outer OpenMP parallel - # section. - os.environ.setdefault("KMP_DUPLICATE_LIB_OK", "True") - - -# Structure to cast the info on dynamically loaded library. See -# https://linux.die.net/man/3/dl_iterate_phdr for more details. -UINT_SYSTEM = ctypes.c_uint64 if sys.maxsize > 2**32 else ctypes.c_uint32 -UINT_HALF_SYSTEM = ctypes.c_uint32 if sys.maxsize > 2**32 else ctypes.c_uint16 - - -class dl_phdr_info(ctypes.Structure): - _fields_ = [ - ("dlpi_addr", UINT_SYSTEM), # Base address of object - ("dlpi_name", ctypes.c_char_p), # path to the library - ("dlpi_phdr", ctypes.c_void_p), # pointer on dlpi_headers - ("dlpi_phnum", UINT_HALF_SYSTEM) # number of element in dlpi_phdr - ] - - -# List of the supported implementations. The items hold the prefix of loaded -# shared objects, the name of the internal_api to call, matching the -# MAP_API_TO_FUNC keys and the name of the user_api, in {"blas", "openmp"}. -SUPPORTED_IMPLEMENTATIONS = [ - { - "user_api": "openmp", - "internal_api": "openmp", - "filename_prefixes": ("libiomp", "libgomp", "libomp", "vcomp",), - }, - { - "user_api": "blas", - "internal_api": "openblas", - "filename_prefixes": ("libopenblas",), - }, - { - "user_api": "blas", - "internal_api": "mkl", - "filename_prefixes": ("libmkl_rt", "mkl_rt",), - }, -] - -# map a internal_api (openmp, openblas, mkl) to set and get functions -MAP_API_TO_FUNC = { - "openmp": { - "set_num_threads": "omp_set_num_threads", - "get_num_threads": "omp_get_max_threads"}, - "openblas": { - "set_num_threads": "openblas_set_num_threads", - "get_num_threads": "openblas_get_num_threads"}, - "mkl": { - "set_num_threads": "MKL_Set_Num_Threads", - "get_num_threads": "MKL_Get_Max_Threads"}, -} - -# Helpers for the doc and test names -ALL_USER_APIS = set(impl['user_api'] for impl in SUPPORTED_IMPLEMENTATIONS) -ALL_PREFIXES = [prefix for impl in SUPPORTED_IMPLEMENTATIONS - for prefix in impl['filename_prefixes']] -ALL_INTERNAL_APIS = list(MAP_API_TO_FUNC.keys()) - - -def _get_limit(prefix, user_api, limits): - if prefix in limits: - return limits[prefix] - if user_api in limits: - return limits[user_api] - return None - - -@_format_docstring(ALL_PREFIXES=ALL_PREFIXES, INTERNAL_APIS=ALL_INTERNAL_APIS) -def _set_threadpool_limits(limits=None, user_api=None): - """Limit the maximal number of threads for threadpools in supported C-lib - - Set the maximal number of threads that can be used in thread pools used in - the supported C-libraries to `limit`. This function works for libraries - that are already loaded in the interpreter and can be changed dynamically. - - The `limits` parameter can be either an integer or a dict to specify the - maximal number of thread that can be used in thread pools. If it is an - integer, sets the maximum number of thread to `limits` for each C-lib - selected by `user_api`. If it is a dictionary `{{key: max_threads}}`, - this function sets a custom maximum number of thread for each `key` which - can be either a `user_api` or a `prefix` for a specific library. - If None, this function does not do anything. - - The `user_api` parameter selects particular APIs of C-libs to limit. Used - only if `limits` is an int. If it is None, this function will apply to all - supported C-libs. If it is "blas", it will limit only BLAS supported C-libs - and if it is "openmp", only OpenMP supported C-libs will be limited. Note - that the latter can affect the number of threads used by the BLAS C-libs if - they rely on OpenMP. - - Return a list with all the supported modules that have been found. Each - module is represented by a dict with the following information: - - 'filename-prefixes' : possible prefixes for the given internal_api. - Possible values are {ALL_PREFIXES}. - - 'prefix' : prefix of the specific implementation of this module. - - 'internal_api': internal API.s Possible values are {INTERNAL_APIS}. - - 'module_path': path to the loaded module. - - 'version': version of the library implemented (if available). - - 'n_thread': current thread limit. - """ - if isinstance(limits, int) or limits is None: - if user_api is None: - user_api = ALL_USER_APIS - elif user_api in ALL_USER_APIS: - user_api = (user_api,) - else: - raise ValueError("user_api must be either in {} or None. Got {} " - "instead.".format(ALL_USER_APIS, user_api)) - limits = {api: limits for api in user_api} - prefixes = [] - else: - if isinstance(limits, list): - # This should be a list of module, for compatibility with - # the result from get_threadpool_limits. - limits = {module['prefix']: module['n_thread'] - for module in limits} - - if not isinstance(limits, dict): - raise TypeError("limits must either be an int, a dict or None." - " Got {} instead".format(type(limits))) - - # With a dictionary, can set both specific limit for given modules - # and global limit for user_api. Fetch each separately. - prefixes = [module for module in limits if module in ALL_PREFIXES] - user_api = [module for module in limits if module in ALL_USER_APIS] - - report_threadpool_size = [] - modules = _load_modules(prefixes=prefixes, user_api=user_api) - for module in modules: - n_thread = _get_limit(module['prefix'], module['user_api'], limits) - if n_thread is not None: - set_func = module['set_num_threads'] - set_func(n_thread) - - # Store the module and remove un-necessary info - module['n_thread'] = module['get_num_threads']() - del module['set_num_threads'], module['get_num_threads'] - del module['clib'] - report_threadpool_size.append(module) - - return report_threadpool_size - - -@_format_docstring(ALL_PREFIXES=ALL_PREFIXES, INTERNAL_APIS=ALL_INTERNAL_APIS) -def get_threadpool_limits(): - """Return the maximal number of threads for threadpools in supported C-lib. - - Return a list with all the supported modules that have been found. Each - module is represented by a dict with the following information: - - 'filename-prefixes' : possible prefixes for the given internal_api. - Possible values are {ALL_PREFIXES}. - - 'prefix' : prefix of the specific implementation of this module. - - 'internal_api': internal API. Possible values are {INTERNAL_APIS}. - - 'module_path': path to the loaded module. - - 'version': version of the library implemented (if available). - - 'n_thread': current thread limit. - """ - report_threadpool_size = [] - modules = _load_modules(user_api=ALL_USER_APIS) - for module in modules: - module['n_thread'] = module['get_num_threads']() - # Remove the wrapper for the module and its function - del module['set_num_threads'], module['get_num_threads'] - del module['clib'] - report_threadpool_size.append(module) - - return report_threadpool_size - - -def get_version(clib, internal_api): - if internal_api == "mkl": - return _get_mkl_version(clib) - elif internal_api == "openmp": - # There is no way to get the version number programmatically in - # OpenMP. - return None - elif internal_api == "openblas": - return _get_openblas_version(clib) - else: - raise NotImplementedError("Unsupported API {}".format(internal_api)) - - -def _get_mkl_version(mkl_clib): - """Return the MKL version - """ - res = ctypes.create_string_buffer(200) - mkl_clib.mkl_get_version_string(res, 200) - - version = res.value.decode('utf-8') - group = re.search(r"Version ([^ ]+) ", version) - if group is not None: - version = group.groups()[0] - return version.strip() - - -def _get_openblas_version(openblas_clib): - """Return the OpenBLAS version - - None means OpenBLAS is not loaded or version < 0.3.4, since OpenBLAS - did not expose its version before that. - """ - get_config = getattr(openblas_clib, "openblas_get_config") - get_config.restype = ctypes.c_char_p - config = get_config().split() - if config[0] == b"OpenBLAS": - return config[1].decode('utf-8') - return None - - -################################################################# -# Loading utilities for dynamically linked shared objects - -def _load_modules(prefixes=None, user_api=None): - """Loop through loaded libraries and return supported ones.""" - if prefixes is None: - prefixes = [] - if user_api is None: - user_api = [] - if sys.platform == "darwin": - return _find_modules_with_clibs_dyld( - prefixes=prefixes, user_api=user_api) - elif sys.platform == "win32": - return _find_modules_with_enum_process_module_ex( - prefixes=prefixes, user_api=user_api) - else: - return _find_modules_with_clibs_dl_iterate_phdr( - prefixes=prefixes, user_api=user_api) - - -def _check_prefix(library_basename, filename_prefixes): - """Return the prefix library_basename starts with or None if none matches - """ - for prefix in filename_prefixes: - if library_basename.startswith(prefix): - return prefix - return None - - -def _match_module(module_info, prefix, prefixes, user_api): - """Return True if this module should be selected.""" - return prefix is not None and (prefix in prefixes or - module_info['user_api'] in user_api) - - -def _make_module_info(module_path, module_info, prefix): - """Make a dict with the information from the module.""" - module_path = os.path.normpath(module_path) - clib = ctypes.CDLL(module_path) - internal_api = module_info['internal_api'] - set_func = getattr(clib, MAP_API_TO_FUNC[internal_api]['set_num_threads'], - lambda n_thread: None) - get_func = getattr(clib, MAP_API_TO_FUNC[internal_api]['get_num_threads'], - lambda: None) - module_info = module_info.copy() - module_info.update(clib=clib, module_path=module_path, prefix=prefix, - set_num_threads=set_func, get_num_threads=get_func, - version=get_version(clib, internal_api)) - return module_info - - -def _get_module_info_from_path(module_path, prefixes, user_api, modules): - module_name = os.path.basename(module_path).lower() - for info in SUPPORTED_IMPLEMENTATIONS: - prefix = _check_prefix(module_name, info['filename_prefixes']) - if _match_module(info, prefix, prefixes, user_api): - modules.append(_make_module_info(module_path, info, prefix)) - - -def _find_modules_with_clibs_dl_iterate_phdr(prefixes, user_api): - """Loop through loaded libraries and return binders on supported ones - - This function is expected to work on POSIX system only. - This code is adapted from code by Intel developper @anton-malakhov - available at https://github.com/IntelPython/smp - - Copyright (c) 2017, Intel Corporation published under the BSD 3-Clause - license - """ - - libc = _get_libc() - if not hasattr(libc, "dl_iterate_phdr"): # pragma: no cover - return [] - - _modules = [] - - # Callback function for `dl_iterate_phdr` which is called for every - # module loaded in the current process until it returns 1. - def match_module_callback(info, size, data): - - # Get the path of the current module - module_path = info.contents.dlpi_name - if module_path: - module_path = module_path.decode("utf-8") - - # Store the module in cls_thread_locals._module if it is - # supported and selected - _get_module_info_from_path(module_path, prefixes, user_api, - _modules) - return 0 - - c_func_signature = ctypes.CFUNCTYPE( - ctypes.c_int, # Return type - ctypes.POINTER(dl_phdr_info), ctypes.c_size_t, ctypes.c_char_p) - c_match_module_callback = c_func_signature(match_module_callback) - - data = ctypes.c_char_p(''.encode('utf-8')) - libc.dl_iterate_phdr(c_match_module_callback, data) - - return _modules - - -def _find_modules_with_clibs_dyld(prefixes, user_api): - """Loop through loaded libraries and return binders on supported ones - - This function is expected to work on OSX system only - """ - libc = _get_libc() - if not hasattr(libc, "_dyld_image_count"): # pragma: no cover - return [] - - _modules = [] - - n_dyld = libc._dyld_image_count() - libc._dyld_get_image_name.restype = ctypes.c_char_p - - for i in range(n_dyld): - module_path = ctypes.string_at(libc._dyld_get_image_name(i)) - module_path = module_path.decode("utf-8") - - # Store the module in cls_thread_locals._module if it is - # supported and selected - _get_module_info_from_path(module_path, prefixes, user_api, - _modules) - - return _modules - - -def _find_modules_with_enum_process_module_ex(prefixes, user_api): - """Loop through loaded libraries and return binders on supported ones - - This function is expected to work on windows system only. - This code is adapted from code by Philipp Hagemeister @phihag available - at https://stackoverflow.com/questions/17474574 - """ - from ctypes.wintypes import DWORD, HMODULE, MAX_PATH - - PROCESS_QUERY_INFORMATION = 0x0400 - PROCESS_VM_READ = 0x0010 - - LIST_MODULES_ALL = 0x03 - - ps_api = _get_windll('Psapi') - kernel_32 = _get_windll('kernel32') - - h_process = kernel_32.OpenProcess( - PROCESS_QUERY_INFORMATION | PROCESS_VM_READ, - False, os.getpid()) - if not h_process: # pragma: no cover - raise OSError('Could not open PID %s' % os.getpid()) - - _modules = [] - try: - buf_count = 256 - needed = DWORD() - # Grow the buffer until it becomes large enough to hold all the - # module headers - while True: - buf = (HMODULE * buf_count)() - buf_size = ctypes.sizeof(buf) - if not ps_api.EnumProcessModulesEx( - h_process, ctypes.byref(buf), buf_size, - ctypes.byref(needed), LIST_MODULES_ALL): - raise OSError('EnumProcessModulesEx failed') - if buf_size >= needed.value: - break - buf_count = needed.value // (buf_size // buf_count) - - count = needed.value // (buf_size // buf_count) - h_modules = map(HMODULE, buf[:count]) - - # Loop through all the module headers and get the module path - buf = ctypes.create_unicode_buffer(MAX_PATH) - n_size = DWORD() - for h_module in h_modules: - - # Get the path of the current module - if not ps_api.GetModuleFileNameExW( - h_process, h_module, ctypes.byref(buf), - ctypes.byref(n_size)): - raise OSError('GetModuleFileNameEx failed') - module_path = buf.value - - # Store the module in cls_thread_locals._module if it is - # supported and selected - _get_module_info_from_path(module_path, prefixes, user_api, - _modules) - finally: - kernel_32.CloseHandle(h_process) - - return _modules - - -def _get_libc(): - """Load the lib-C for unix systems.""" - libc_name = find_library("c") - if libc_name is None: # pragma: no cover - return None - return ctypes.CDLL(libc_name) - - -def _get_windll(dll_name): - """Load a windows DLL""" - return ctypes.WinDLL("{}.dll".format(dll_name)) - - -class threadpool_limits: - """Change the maximal number of threads that can be used in thread pools. - - This class can be used either as a function (the construction of this - object limits the number of threads) or as a context manager, in a `with` - block. - - Set the maximal number of threads that can be used in thread pools used in - the supported C-libraries to `limit`. This function works for libraries - that are already loaded in the interpreter and can be changed dynamically. - - The `limits` parameter can be either an integer or a dict to specify the - maximal number of thread that can be used in thread pools. If it is an - integer, sets the maximum number of thread to `limits` for each C-lib - selected by `user_api`. If it is a dictionary `{{key: max_threads}}`, - this function sets a custom maximum number of thread for each `key` which - can be either a `user_api` or a `prefix` for a specific library. - If None, this function does not do anything. - - The `user_api` parameter selects particular APIs of C-libs to limit. Used - only if `limits` is an int. If it is None, this function will apply to all - supported C-libs. If it is "blas", it will limit only BLAS supported C-libs - and if it is "openmp", only OpenMP supported C-libs will be limited. Note - that the latter can affect the number of threads used by the BLAS C-libs if - they rely on OpenMP. - """ - def __init__(self, limits=None, user_api=None): - self._enabled = limits is not None - if self._enabled: - self.old_limits = get_threadpool_limits() - _set_threadpool_limits(limits=limits, user_api=user_api) - - def __enter__(self): - pass - - def __exit__(self, type, value, traceback): - self.unregister() - - def unregister(self): - if self._enabled: - _set_threadpool_limits(limits=self.old_limits) diff --git a/threadpoolctl/tests/_openmp_test_helper/__init__.py b/threadpoolctl/tests/_openmp_test_helper/__init__.py deleted file mode 100644 index 4a436e9e..00000000 --- a/threadpoolctl/tests/_openmp_test_helper/__init__.py +++ /dev/null @@ -1,4 +0,0 @@ -from .openmp_helpers import check_openmp_n_threads - - -__all__ = ["check_openmp_n_threads"] diff --git a/threadpoolctl/tests/_openmp_test_helper/openmp_helpers.pyx b/threadpoolctl/tests/_openmp_test_helper/openmp_helpers.pyx deleted file mode 100644 index 02b6da02..00000000 --- a/threadpoolctl/tests/_openmp_test_helper/openmp_helpers.pyx +++ /dev/null @@ -1,19 +0,0 @@ -import cython -cimport openmp -from cython.parallel import prange -from libc.stdlib cimport malloc, free - - -def check_openmp_n_threads(int n): - """Run a short parallel section, and return the number of threads that - where effectively used by openmp.""" - cdef long n_sum = 0 - cdef int i, num_threads - - for i in prange(n, nogil=True): - num_threads = openmp.omp_get_num_threads() - n_sum += i - - assert n_sum == (n - 1) * n / 2 - - return num_threads diff --git a/threadpoolctl/tests/_openmp_test_helper/setup.py b/threadpoolctl/tests/_openmp_test_helper/setup.py deleted file mode 100644 index 8d23276b..00000000 --- a/threadpoolctl/tests/_openmp_test_helper/setup.py +++ /dev/null @@ -1,33 +0,0 @@ -import os -import sys -from distutils.core import setup -from Cython.Build import cythonize -from distutils.extension import Extension - - -if sys.platform == "darwin": - os.environ["CC"] = "gcc-4.9" - os.environ["CXX"] = "g++-4.9" - -if sys.platform == "win32": - extra_compile_args = ["/openmp"] - extra_link_args = None -else: - extra_compile_args = ["-fopenmp"] - extra_link_args = ['-fopenmp'] - -ext_modules = [ - Extension( - "openmp_helpers", - ["openmp_helpers.pyx"], - extra_compile_args=extra_compile_args, - extra_link_args=extra_link_args - ) -] - -setup( - name='_openmp_test_helper', - ext_modules=cythonize( - ext_modules, - language_level=3), -) diff --git a/threadpoolctl/tests/test_threadpool_limits.py b/threadpoolctl/tests/test_threadpool_limits.py deleted file mode 100644 index b951ee93..00000000 --- a/threadpoolctl/tests/test_threadpool_limits.py +++ /dev/null @@ -1,180 +0,0 @@ -import re -import ctypes -import pytest - - -from threadpoolctl import threadpool_limits -from threadpoolctl import get_threadpool_limits -from threadpoolctl._threadpool_limiters import _set_threadpool_limits -from threadpoolctl._threadpool_limiters import ALL_PREFIXES, ALL_USER_APIS - -from .utils import with_check_openmp_n_threads, libopenblas_paths - - -def should_skip_module(module): - # Possible bug in getting maximum number of threads with OpenBLAS < 0.2.16 - # and OpenBLAS does not expose its version before 0.3.4. - return module['internal_api'] == "openblas" and module['version'] is None - - -@pytest.mark.parametrize("prefix", ALL_PREFIXES) -def test_threadpool_limits(openblas_present, mkl_present, prefix): - old_limits = get_threadpool_limits() - - prefix_found = len([1 for module in old_limits - if prefix == module['prefix']]) - old_limits = {clib['prefix']: clib['n_thread'] for clib in old_limits} - - if not prefix_found: - have_mkl = len({'mkl_rt', 'libmkl_rt'}.intersection(old_limits)) > 0 - if "mkl_rt" in prefix and mkl_present and not have_mkl: - raise RuntimeError("Could not load the MKL prefix") - elif prefix == "libopenblas" and openblas_present: - raise RuntimeError("Could not load the OpenBLAS prefix") - else: - pytest.skip("Need {} support".format(prefix)) - - try: - new_limits = _set_threadpool_limits(limits={prefix: 1}) - new_limits = {clib['prefix']: clib['n_thread'] for clib in new_limits} - assert new_limits[prefix] == 1 - - threadpool_limits(limits={prefix: 3}) - new_limits = get_threadpool_limits() - new_limits = {clib['prefix']: clib['n_thread'] for clib in new_limits} - assert new_limits[prefix] in (3, old_limits[prefix]) - finally: - # Avoid having side effects in case of failures - threadpool_limits(limits=old_limits) - - new_limits = get_threadpool_limits() - new_limits = {clib['prefix']: clib['n_thread'] for clib in new_limits} - assert new_limits[prefix] == old_limits[prefix] - - -@pytest.mark.parametrize("user_api", (None, "blas", "openmp")) -def test_set_threadpool_limits_apis(user_api): - # Check that the number of threads used by the multithreaded C-libs can be - # modified dynamically. - - if user_api is None: - api_modules = ('blas', 'openmp') - else: - api_modules = (user_api,) - - old_limits = get_threadpool_limits() - old_limits = {clib['prefix']: clib['n_thread'] for clib in old_limits} - - try: - new_limits = _set_threadpool_limits(limits=1, user_api=user_api) - for module in new_limits: - if should_skip_module(module): - continue - if module['user_api'] in api_modules: - assert module['n_thread'] == 1 - - threadpool_limits(limits=3, user_api=user_api) - new_limits = get_threadpool_limits() - for module in new_limits: - if should_skip_module(module): - continue - if module['user_api'] in api_modules: - assert module['n_thread'] in (3, old_limits[module['prefix']]) - finally: - # Avoid having side effects on other tests in case of failure - threadpool_limits(limits=old_limits) - - new_limits = get_threadpool_limits() - for module in new_limits: - assert module['n_thread'] == old_limits[module['prefix']] - - -def test_set_threadpool_limits_bad_input(): - # Check that appropriate errors are raised for invalid arguments - - match = re.escape("user_api must be either in {} or None." - .format(ALL_USER_APIS)) - with pytest.raises(ValueError, match=match): - threadpool_limits(limits=1, user_api="wrong") - - with pytest.raises(TypeError, - match="limits must either be an int, a dict or None"): - threadpool_limits(limits=(1, 2, 3)) - - -@pytest.mark.parametrize("user_api", (None, "blas", "openmp")) -def test_thread_limit_context(user_api): - # Tests the thread limits context manager - - if user_api is None: - apis = ('blas', 'openmp') - else: - apis = (user_api,) - - old_limits = get_threadpool_limits() - old_limits = {clib['prefix']: clib['n_thread'] for clib in old_limits} - - with threadpool_limits(limits=None, user_api=user_api): - limits = get_threadpool_limits() - limits = {clib['prefix']: clib['n_thread'] for clib in limits} - assert limits == old_limits - - with threadpool_limits(limits=1, user_api=user_api): - limits = get_threadpool_limits() - - for module in limits: - if should_skip_module(module): - continue - elif module['user_api'] in apis: - assert module['n_thread'] == 1 - else: - assert module['n_thread'] == old_limits[module['prefix']] - - limits = get_threadpool_limits() - limits = {clib['prefix']: clib['n_thread'] for clib in limits} - assert limits == old_limits - - -@with_check_openmp_n_threads -@pytest.mark.parametrize('n_threads', [1, 2, 4]) -def test_openmp_limit_num_threads(n_threads): - # checks that OpenMP effectively uses the number of threads requested by - # the context manager - - from ._openmp_test_helper import check_openmp_n_threads - - old_num_threads = check_openmp_n_threads(100) - - with threadpool_limits(limits=n_threads): - assert check_openmp_n_threads(100) in (n_threads, old_num_threads) - assert check_openmp_n_threads(100) == old_num_threads - - -def test_shipped_openblas(): - - libopenblas = [ctypes.CDLL(path) for path in libopenblas_paths] - - old_limits = [blas.openblas_get_num_threads() for blas in libopenblas] - - with threadpool_limits(1): - assert all([blas.openblas_get_num_threads() == 1 - for blas in libopenblas]) - - assert all([blas.openblas_get_num_threads() == l - for blas, l in zip(libopenblas, old_limits)]) - - -@pytest.mark.skipif(len(libopenblas_paths) < 2, - reason="need at least 2 shipped openblas library") -def test_multiple_shipped_openblas(): - - libopenblas = [ctypes.CDLL(path) for path in libopenblas_paths] - - old_limits = [blas.openblas_get_num_threads() for blas in libopenblas] - - with threadpool_limits(1): - assert all([blas.openblas_get_num_threads() == 1 - for blas in libopenblas]) - - assert all([blas.openblas_get_num_threads() == l - for blas, l in zip(libopenblas, old_limits)]) diff --git a/threadpoolctl/tests/utils.py b/threadpoolctl/tests/utils.py deleted file mode 100644 index 6ee2ed4b..00000000 --- a/threadpoolctl/tests/utils.py +++ /dev/null @@ -1,66 +0,0 @@ -import os -import pytest -from glob import glob - - -def skip_func(msg): - def test_func(*args, **kwargs): - pytest.skip(msg) - return test_func - - -# Path to shipped openblas for libraries such as numpy or scipy -libopenblas_patterns = [] - - -# A decorator to run tests only when numpy is available -try: - # make sure the mkl/blas are loaded for test_threadpool_limits - import numpy as np - np.dot(np.ones(1000), np.ones(1000)) - - def with_numpy(func): - """A decorator to skip tests requiring numpy.""" - return func - - libopenblas_patterns.append(os.path.join(np.__path__[0], ".libs", - "libopenblas*")) - -except ImportError: - def with_numpy(func): - """A decorator to skip tests requiring numpy.""" - return skip_func('Test require numpy') - - -try: - import scipy - import scipy.linalg # noqa: F401 - - libopenblas_patterns.append(os.path.join(scipy.__path__[0], ".libs", - "libopenblas*")) -except ImportError: - pass - -libopenblas_paths = set(path for pattern in libopenblas_patterns - for path in glob(pattern)) - - -# A decorator to run tests only when check_openmp_n_threads is available -try: - from ._openmp_test_helper import check_openmp_n_threads - - def with_check_openmp_n_threads(func): - """A decorator to skip tests if check_openmp_n_threads is not compiled. - """ - return func - - def _run_check_openmp_n_threads(*args): - return check_openmp_n_threads(*args) - -except ImportError: - def with_check_openmp_n_threads(func): - """A decorator to skip tests if check_openmp_n_threads is not compiled. - """ - return skip_func('Test requires check_openmp_n_threads to be compiled') - - _run_check_openmp_n_threads = None diff --git a/threadpoolctl/utils.py b/threadpoolctl/utils.py deleted file mode 100644 index 6f85db32..00000000 --- a/threadpoolctl/utils.py +++ /dev/null @@ -1,8 +0,0 @@ - - -def _format_docstring(*args, **kwargs): - def decorator(o): - o.__doc__ = o.__doc__.format(*args, **kwargs) - return o - - return decorator