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