Merge pull request #590 from nicoddemus/auto-hook

Author: nicoddemus

.appveyor.yml

@@ -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

.travis.yml

@@ -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'

README.rst

@@ -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:
 

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.

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]

setup.cfg

@@ -1,6 +1,3 @@
-[bdist_wheel]
-universal = 1
-
 [metadata]
 license_file = LICENSE
 

setup.py

@@ -1,6 +1,6 @@
 from setuptools import setup, find_packages
 
-install_requires = ["execnet>=1.1", "psutil>=3.0.0", "pytest>=6.0.0", "pytest-forked"]
+install_requires = ["execnet>=1.1", "pytest>=6.0.0", "pytest-forked"]
 
 
 with open("README.rst") as f:
@@ -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"]
     },

src/xdist/newhooks.py

@@ -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
+    """

src/xdist/plugin.py

@@ -1,21 +1,43 @@
+import os
 import uuid
 
-import psutil
 import py
 import pytest
 
 
-def auto_detect_cpus():
-    return psutil.cpu_count(logical=False) or psutil.cpu_count() or 1
-
-
-class AutoInt(int):
-    """Mark value as auto-detected."""
+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
+            return 2
+        try:
+            from os import cpu_count
+        except ImportError:
+            from multiprocessing import cpu_count
+    try:
+        n = cpu_count()
+    except NotImplementedError:
+        return 1
+    return n if n else 1
 
 
 def parse_numprocesses(s):
     if s == "auto":
-        return AutoInt(auto_detect_cpus())
+        return "auto"
     elif s is not None:
         return int(s)
 
@@ -168,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":

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")
@@ -35,15 +39,28 @@ def test_dist_options(testdir):
 
 
 def test_auto_detect_cpus(testdir, monkeypatch):
-    import psutil
+    import os
     from xdist.plugin import pytest_cmdline_main as check_options
 
-    monkeypatch.setattr(psutil, "cpu_count", lambda logical=True: 99)
+    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"):
+        monkeypatch.setattr(os, "cpu_count", lambda: 99)
+    else:
+        import multiprocessing
+
+        monkeypatch.setattr(multiprocessing, "cpu_count", lambda: 99)
 
     config = testdir.parseconfigure("-n2")
     assert config.getoption("numprocesses") == 2
 
     config = testdir.parseconfigure("-nauto")
+    check_options(config)
     assert config.getoption("numprocesses") == 99
 
     config = testdir.parseconfigure("-nauto", "--pdb")
@@ -52,9 +69,37 @@ def test_auto_detect_cpus(testdir, monkeypatch):
     assert config.getoption("numprocesses") == 0
     assert config.getoption("dist") == "no"
 
-    monkeypatch.setattr(psutil, "cpu_count", lambda logical=True: None)
+    monkeypatch.delattr(os, "sched_getaffinity", raising=False)
+    monkeypatch.setenv("TRAVIS", "true")
     config = testdir.parseconfigure("-nauto")
-    assert config.getoption("numprocesses") == 1
+    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):