Skip to content

Commit

Permalink
Add psutil extra and introduce pytest_xdist_auto_num_workers hook
Browse files Browse the repository at this point in the history
This makes using psutil optional and opens up the possibility of customization
through the pytest_xdist_auto_num_workers hook, making things like #477 possible.

Fix #585
  • Loading branch information
nicoddemus committed Aug 25, 2020
1 parent e469fc7 commit 607d828
Show file tree
Hide file tree
Showing 10 changed files with 88 additions and 21 deletions.
1 change: 1 addition & 0 deletions .appveyor.yml
Expand Up @@ -5,6 +5,7 @@ environment:
- TOXENV: "py37-pytestlatest"
- TOXENV: "py38-pytestlatest"
- TOXENV: "py38-pytestmaster"
- TOXENV: "py38-psutil"

install:
- C:\Python38\python -m pip install -U pip setuptools virtualenv
Expand Down
2 changes: 2 additions & 0 deletions .travis.yml
Expand Up @@ -44,6 +44,8 @@ jobs:
env: TOXENV=py39-pytestlatest
- python: "3.8"
env: TOXENV=py38-pytestmaster
- python: "3.8"
env: TOXENV=py38-psutil

- stage: deploy
python: '3.8'
Expand Down
7 changes: 4 additions & 3 deletions README.rst
Expand Up @@ -57,10 +57,11 @@ Install the plugin with::

pip install pytest-xdist

or use the package in develop/in-place mode with
a checkout of the `pytest-xdist repository`_ ::

pip install --editable .
To use ``psutil`` for detection of the number of CPUs available, install the ``psutil`` extra::

pip install pytest-xdist[psutil]


.. _parallelization:

Expand Down
1 change: 1 addition & 0 deletions changelog/585.feature.rst
@@ -0,0 +1 @@
New ``pytest_xdist_auto_num_workers`` hook can be implemented by plugins or ``conftest.py`` files to control the number of workers when ``--numprocesses=auto`` is given in the command-line.
3 changes: 3 additions & 0 deletions changelog/585.trivial.rst
@@ -0,0 +1,3 @@
``psutil`` has proven to make ``pytest-xdist`` installation in certain platforms and containers problematic, so to use it for automatic number of CPUs detection users need to install the ``psutil`` extra::

pip install pytest-xdist[psutil]
2 changes: 1 addition & 1 deletion setup.py
Expand Up @@ -18,7 +18,7 @@
platforms=["linux", "osx", "win32"],
packages=find_packages(where="src"),
package_dir={"": "src"},
extras_require={"testing": ["filelock"]},
extras_require={"testing": ["filelock"], "psutil": ["psutil>=3.0"]},
entry_points={
"pytest11": ["xdist = xdist.plugin", "xdist.looponfail = xdist.looponfail"]
},
Expand Down
10 changes: 10 additions & 0 deletions src/xdist/newhooks.py
Expand Up @@ -55,3 +55,13 @@ def pytest_xdist_node_collection_finished(node, ids):
@pytest.mark.firstresult
def pytest_xdist_make_scheduler(config, log):
""" return a node scheduler implementation """


@pytest.mark.firstresult
def pytest_xdist_auto_num_workers(config):
"""
Return the number of workers to spawn when ``--numprocesses=auto`` is given in the
command-line.
.. versionadded:: 2.1
"""
30 changes: 17 additions & 13 deletions src/xdist/plugin.py
Expand Up @@ -5,9 +5,21 @@
import pytest


def auto_detect_cpus():
def pytest_xdist_auto_num_workers():
try:
import psutil
except ImportError:
pass
else:
count = psutil.cpu_count(logical=False) or psutil.cpu_count()
if count:
return count
try:
from os import sched_getaffinity

def cpu_count():
return len(sched_getaffinity(0))

except ImportError:
if os.environ.get("TRAVIS") == "true":
# workaround https://bitbucket.org/pypy/pypy/issues/2375
Expand All @@ -16,25 +28,16 @@ def auto_detect_cpus():
from os import cpu_count
except ImportError:
from multiprocessing import cpu_count
else:

def cpu_count():
return len(sched_getaffinity(0))

try:
n = cpu_count()
except NotImplementedError:
return 1
return n if n else 1


class AutoInt(int):
"""Mark value as auto-detected."""


def parse_numprocesses(s):
if s == "auto":
return AutoInt(auto_detect_cpus())
return "auto"
elif s is not None:
return int(s)

Expand Down Expand Up @@ -187,12 +190,13 @@ def pytest_configure(config):
@pytest.mark.tryfirst
def pytest_cmdline_main(config):
usepdb = config.getoption("usepdb", False) # a core option
if isinstance(config.option.numprocesses, AutoInt):
if config.option.numprocesses == "auto":
if usepdb:
config.option.numprocesses = 0
config.option.dist = "no"
else:
config.option.numprocesses = int(config.option.numprocesses)
auto_num_cpus = config.hook.pytest_xdist_auto_num_workers(config=config)
config.option.numprocesses = auto_num_cpus

if config.option.numprocesses:
if config.option.dist == "no":
Expand Down
37 changes: 37 additions & 0 deletions testing/test_plugin.py
@@ -1,7 +1,11 @@
from contextlib import suppress

import py
import execnet
from xdist.workermanage import NodeManager

import pytest


def test_dist_incompatibility_messages(testdir):
result = testdir.runpytest("--pdb", "--looponfail")
Expand Down Expand Up @@ -38,6 +42,11 @@ def test_auto_detect_cpus(testdir, monkeypatch):
import os
from xdist.plugin import pytest_cmdline_main as check_options

with suppress(ImportError):
import psutil

monkeypatch.setattr(psutil, "cpu_count", lambda logical=True: None)

if hasattr(os, "sched_getaffinity"):
monkeypatch.setattr(os, "sched_getaffinity", lambda _pid: set(range(99)))
elif hasattr(os, "cpu_count"):
Expand All @@ -51,6 +60,7 @@ def test_auto_detect_cpus(testdir, monkeypatch):
assert config.getoption("numprocesses") == 2

config = testdir.parseconfigure("-nauto")
check_options(config)
assert config.getoption("numprocesses") == 99

config = testdir.parseconfigure("-nauto", "--pdb")
Expand All @@ -62,9 +72,36 @@ def test_auto_detect_cpus(testdir, monkeypatch):
monkeypatch.delattr(os, "sched_getaffinity", raising=False)
monkeypatch.setenv("TRAVIS", "true")
config = testdir.parseconfigure("-nauto")
check_options(config)
assert config.getoption("numprocesses") == 2


def test_auto_detect_cpus_psutil(testdir, monkeypatch):
from xdist.plugin import pytest_cmdline_main as check_options

psutil = pytest.importorskip("psutil")

monkeypatch.setattr(psutil, "cpu_count", lambda logical=True: 42)

config = testdir.parseconfigure("-nauto")
check_options(config)
assert config.getoption("numprocesses") == 42


def test_hook_auto_num_workers(testdir, monkeypatch):
from xdist.plugin import pytest_cmdline_main as check_options

testdir.makeconftest(
"""
def pytest_xdist_auto_num_workers():
return 42
"""
)
config = testdir.parseconfigure("-nauto")
check_options(config)
assert config.getoption("numprocesses") == 42


def test_boxed_with_collect_only(testdir):
from xdist.plugin import pytest_cmdline_main as check_options

Expand Down
16 changes: 12 additions & 4 deletions tox.ini
Expand Up @@ -3,18 +3,26 @@ envlist=
linting
py{35,36,37,38,39}-pytestlatest
py38-pytestmaster
py38-psutil

[testenv]
passenv = USER USERNAME
extras = testing
deps =
pytestlatest: pytest
pytestmaster: git+https://github.com/pytest-dev/pytest.git@master
commands=
pytest {posargs}

[testenv:py38-psutil]
extras =
testing
psutil
deps = pytest
commands =
pytest {posargs:-k psutil}

[testenv:linting]
skipsdist = True
skip_install = True
usedevelop = True
deps =
pre-commit
Expand All @@ -28,9 +36,9 @@ skipsdist = True
usedevelop = True
passenv = *
deps =
towncrier
towncrier
commands =
towncrier --version {posargs} --yes
towncrier --version {posargs} --yes

[pytest]
addopts = -ra
Expand Down

0 comments on commit 607d828

Please sign in to comment.