|
2 | 2 | Tests for get/set number of threads API introspection. |
3 | 3 | """ |
4 | 4 |
|
| 5 | +import sys |
5 | 6 | from threading import local as threadlocal |
6 | 7 |
|
7 | 8 | import pytest |
8 | 9 |
|
9 | | -from threadpoolctl import _ThreadLimitScope, _determine_thread_limit_scope |
| 10 | +from threadpoolctl import ( |
| 11 | + _ThreadLimitScope, |
| 12 | + _determine_thread_limit_scope, |
| 13 | + ThreadpoolController, |
| 14 | +) |
| 15 | + |
| 16 | +# Make sure we have some BLAS libraries loaded: |
| 17 | +from . import utils as _ |
10 | 18 |
|
11 | 19 |
|
12 | 20 | class FakeThreadLocalAPI(threadlocal): |
@@ -52,3 +60,38 @@ def test_determine_thread_limit_scope_processwide(default: int) -> None: |
52 | 60 | """ |
53 | 61 | api = FakeProcesswideAPI(default) |
54 | 62 | assert _determine_thread_limit_scope(api.get, api.set) == _ThreadLimitScope.PROCESS |
| 63 | + |
| 64 | + |
| 65 | +@pytest.mark.skipif( |
| 66 | + sys.platform != "linux", reason="We only hardcoded Linux-specific behavior" |
| 67 | +) |
| 68 | +@pytest.mark.parametrize( |
| 69 | + ["select_filter", "expected_thread_limit_scope"], |
| 70 | + [ |
| 71 | + ( |
| 72 | + {"internal_api": "openblas", "threading_layer": "pthreads"}, |
| 73 | + _ThreadLimitScope.PROCESS, |
| 74 | + ), |
| 75 | + ( |
| 76 | + {"user_api": "openmp"}, |
| 77 | + _ThreadLimitScope.CURRENT_THREAD, |
| 78 | + ), |
| 79 | + ], |
| 80 | +) |
| 81 | +def test_api_scope( |
| 82 | + select_filter: dict[str, str], expected_thread_limit_scope: str |
| 83 | +) -> None: |
| 84 | + """ |
| 85 | + Check ``_determine_thread_limit_scope()`` against libraries with known |
| 86 | + properties, to make sure it detects them correctly. The test is intended |
| 87 | + to be of the function's behavior, not of the libraries. |
| 88 | + """ |
| 89 | + controller = ThreadpoolController().select(**select_filter) |
| 90 | + if not controller.lib_controllers: |
| 91 | + pytest.skip(f"{select_filter} controller not found") |
| 92 | + |
| 93 | + for lib in controller.lib_controllers: |
| 94 | + assert ( |
| 95 | + _determine_thread_limit_scope(lib.get_num_threads, lib.set_num_threads) |
| 96 | + == expected_thread_limit_scope |
| 97 | + ) |
0 commit comments