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_controller_limit():
# Check that using the limit method of ThreadpoolController only impact its
# library controllers.
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