rotate `CUDA_VISIBLE_DEVICES` when pytest-xdist is enabled
+import collections
import os
import subprocess
import sys
+import pytest
+
# Enable `testdir` fixture to test `cupy.testing`.
# `pytest_plugins` cannot be locally configured. See also
@@ -27,3 +30,37 @@ def pytest_configure(config):
if _is_in_ci() and _is_pip_installed():
print("***** Installed packages *****", flush=True)
subprocess.check_call([sys.executable, '-m', 'pip', 'freeze', '--all'])
+ if config.pluginmanager.hasplugin("xdist"):
+ config.pluginmanager.register(DeferPlugin())
+
+
+# https://docs.pytest.org/en/latest/how-to/writing_hook_functions.html#optionally-using-hooks-from-3rd-party-plugins
+class DeferPlugin:
+ """Simple plugin to defer pytest-xdist hook functions."""
+
+ # Edit the environment variable `CUDA_VISIBLE_DEVICES` for each session.
+ # Cannot use `pytest_configure_node` nor `pytest_testnodeready` hook,
+ # because they are called in the `master` node (process).
+ # See also https://github.com/pytest-dev/pytest-xdist/issues/179.
+ @pytest.fixture(autouse=True, scope='session')
+ def _rotate_cuda_visible_devices(self, worker_id):
+ n_gpu = os.environ.get('CUPY_TEST_GPU_LIMIT')
+ if n_gpu is None:
+ return
+ n_gpu = int(n_gpu)
+
+ assert worker_id.startswith('gw')
+ w = int(worker_id[2:])
+
+ devices = os.environ.get('CUDA_VISIBLE_DEVICES')
+ if devices is None:
+ devices = [str(k) for k in range(n_gpu)]
+ else:
+ devices = devices.split(',')[:n_gpu]
+ devices = collections.deque(devices)
+ devices.rotate(w)
+ devices = ','.join(devices)
+ os.environ['CUDA_VISIBLE_DEVICES'] = devices
+ # With PyTest's default, the print will be shown as
+ # "--- Captured stdout setup ---" on failure.
+ print(f'CUDA_VISIBLE_DEVICES={devices}')
