diff --git a/ops/conda_env/aarch64_test.yml b/ops/conda_env/aarch64_test.yml index 14305ebbf090..d7dd13639ff3 100644 --- a/ops/conda_env/aarch64_test.yml +++ b/ops/conda_env/aarch64_test.yml @@ -26,7 +26,7 @@ dependencies: - awscli - numba - llvmlite -- loky +- loky>=3.5.1 - pyarrow - pyspark>=3.4.0 - cloudpickle diff --git a/ops/conda_env/linux_cpu_test.yml b/ops/conda_env/linux_cpu_test.yml index e4c0b507c8e2..55bac17f2dbb 100644 --- a/ops/conda_env/linux_cpu_test.yml +++ b/ops/conda_env/linux_cpu_test.yml @@ -34,7 +34,7 @@ dependencies: - boto3 - awscli - py-ubjson -- loky +- loky>=3.5.1 - pyarrow - protobuf - cloudpickle diff --git a/ops/conda_env/macos_cpu_test.yml b/ops/conda_env/macos_cpu_test.yml index 29ff99e3504f..390abf141803 100644 --- a/ops/conda_env/macos_cpu_test.yml +++ b/ops/conda_env/macos_cpu_test.yml @@ -6,8 +6,6 @@ dependencies: - pip - wheel - pyyaml -- cpplint -- pylint - numpy - scipy - llvm-openmp @@ -20,22 +18,17 @@ dependencies: - python-graphviz - hypothesis - astroid -- sphinx - sh -- recommonmark -- mock -- breathe - pytest - pytest-cov +- pytest-timeout - python-kubernetes - urllib3 - jsonschema - boto3 - awscli -- loky +- loky>=3.5.1 - pyarrow -- pyspark>=3.4.0 - cloudpickle - pip: - setuptools - - sphinx_rtd_theme diff --git a/ops/conda_env/win64_test.yml b/ops/conda_env/win64_test.yml index 32b9339e6fc0..6e87e1560c21 100644 --- a/ops/conda_env/win64_test.yml +++ b/ops/conda_env/win64_test.yml @@ -16,5 +16,5 @@ dependencies: - python-graphviz - pip - py-ubjson -- loky +- loky>=3.5.1 - pyarrow diff --git a/tests/python/test_collective.py b/tests/python/test_collective.py index 473b38b5b742..1204c0faf8c9 100644 --- a/tests/python/test_collective.py +++ b/tests/python/test_collective.py @@ -3,7 +3,6 @@ import numpy as np import pytest -from loky import get_reusable_executor import xgboost as xgb from xgboost import RabitTracker, build_info, federated @@ -25,10 +24,13 @@ def run_rabit_worker(rabit_env: dict, world_size: int) -> int: @pytest.mark.skipif(**tm.no_loky()) def test_rabit_communicator() -> None: + from loky import get_reusable_executor + world_size = 2 tracker = RabitTracker(host_ip="127.0.0.1", n_workers=world_size) tracker.start() workers = [] + with get_reusable_executor(max_workers=world_size) as pool: for _ in range(world_size): worker = pool.submit( @@ -60,6 +62,8 @@ def run_federated_worker(port: int, world_size: int, rank: int) -> int: @pytest.mark.skipif(**tm.skip_win()) @pytest.mark.skipif(**tm.no_loky()) def test_federated_communicator() -> None: + from loky import get_reusable_executor + if not build_info()["USE_FEDERATED"]: pytest.skip("XGBoost not built with federated learning enabled")