Repository: pytest-dev/pytest-xdist

Issue:
New `LoadScheduling` batch distribution logic doesn't work for long-running tests
The PR https://github.com/pytest-dev/pytest-xdist/pull/812 that implemented new batch distribution logic (and that is now released in 3.0.2) disregards long-running tests.

I have several modules with long-running tests and many of these tests are now scheduled on single worker, instead of being distributed evenly across multiple workers. As a result, total runtime of my testsuite more than doubled from ~ 2.5 hours to 5+ hours.

Are there any workarounds to bring back the original round-robin scheduling?

cc @amezin 

Discussion:

Patch:
diff --git a/changelog/855.feature b/changelog/855.feature
new file mode 100644
index 0000000..9a04b48
--- /dev/null
+++ b/changelog/855.feature
@@ -0,0 +1,2 @@
+Users can now configure ``load`` scheduling precision using ``--maxschedchunk`` command
+line option.
diff --git a/src/xdist/plugin.py b/src/xdist/plugin.py
index 8c22f9b..9bbbae7 100644
--- a/src/xdist/plugin.py
+++ b/src/xdist/plugin.py
@@ -153,6 +153,19 @@ def pytest_addoption(parser):
             "on every test run."
         ),
     )
+    group.addoption(
+        "--maxschedchunk",
+        action="store",
+        type=int,
+        help=(
+            "Maximum number of tests scheduled in one step for --dist=load. "
+            "Setting it to 1 will force pytest to send tests to workers one by "
+            "one - might be useful for a small number of slow tests. "
+            "Larger numbers will allow the scheduler to submit consecutive "
+            "chunks of tests to workers - allows reusing fixtures. "
+            "Unlimited if not set."
+        ),
+    )
 
     parser.addini(
         "rsyncdirs",
diff --git a/src/xdist/scheduler/load.py b/src/xdist/scheduler/load.py
index 11d5309..ccca68b 100644
--- a/src/xdist/scheduler/load.py
+++ b/src/xdist/scheduler/load.py
@@ -64,6 +64,7 @@ class LoadScheduling:
         else:
             self.log = log.loadsched
         self.config = config
+        self.maxschedchunk = self.config.getoption("maxschedchunk")
 
     @property
     def nodes(self):
@@ -185,7 +186,9 @@ class LoadScheduling:
                     # so let's rather wait with sending new items
                     return
                 num_send = items_per_node_max - len(node_pending)
-                self._send_tests(node, num_send)
+                # keep at least 2 tests pending even if --maxschedchunk=1
+                maxschedchunk = max(2 - len(node_pending), self.maxschedchunk)
+                self._send_tests(node, min(num_send, maxschedchunk))
         else:
             node.shutdown()
 
@@ -245,6 +248,9 @@ class LoadScheduling:
         if not self.collection:
             return
 
+        if self.maxschedchunk is None:
+            self.maxschedchunk = len(self.collection)
+
         # Send a batch of tests to run. If we don't have at least two
         # tests per node, we have to send them all so that we can send
         # shutdown signals and get all nodes working.
@@ -265,7 +271,8 @@ class LoadScheduling:
             # how many items per node do we have about?
             items_per_node = len(self.collection) // len(self.node2pending)
             # take a fraction of tests for initial distribution
-            node_chunksize = max(items_per_node // 4, 2)
+            node_chunksize = min(items_per_node // 4, self.maxschedchunk)
+            node_chunksize = max(node_chunksize, 2)
             # and initialize each node with a chunk of tests
             for node in self.nodes:
                 self._send_tests(node, node_chunksize)


Test patch:
diff --git a/testing/test_dsession.py b/testing/test_dsession.py
index 86273b8..24ec4ae 100644
--- a/testing/test_dsession.py
+++ b/testing/test_dsession.py
@@ -129,6 +129,56 @@ class TestLoadScheduling:
         assert node1.sent == [0, 1, 4, 5]
         assert not sched.pending
 
+    def test_schedule_maxchunk_none(self, pytester: pytest.Pytester) -> None:
+        config = pytester.parseconfig("--tx=2*popen")
+        sched = LoadScheduling(config)
+        sched.add_node(MockNode())
+        sched.add_node(MockNode())
+        node1, node2 = sched.nodes
+        col = [f"test{i}" for i in range(16)]
+        sched.add_node_collection(node1, col)
+        sched.add_node_collection(node2, col)
+        sched.schedule()
+        assert node1.sent == [0, 1]
+        assert node2.sent == [2, 3]
+        assert sched.pending == list(range(4, 16))
+        assert sched.node2pending[node1] == node1.sent
+        assert sched.node2pending[node2] == node2.sent
+        sched.mark_test_complete(node1, 0)
+        assert node1.sent == [0, 1, 4, 5]
+        assert sched.pending == list(range(6, 16))
+        sched.mark_test_complete(node1, 1)
+        assert node1.sent == [0, 1, 4, 5]
+        assert sched.pending == list(range(6, 16))
+
+        for i in range(7, 16):
+            sched.mark_test_complete(node1, i - 3)
+            assert node1.sent == [0, 1] + list(range(4, i))
+            assert node2.sent == [2, 3]
+            assert sched.pending == list(range(i, 16))
+
+    def test_schedule_maxchunk_1(self, pytester: pytest.Pytester) -> None:
+        config = pytester.parseconfig("--tx=2*popen", "--maxschedchunk=1")
+        sched = LoadScheduling(config)
+        sched.add_node(MockNode())
+        sched.add_node(MockNode())
+        node1, node2 = sched.nodes
+        col = [f"test{i}" for i in range(16)]
+        sched.add_node_collection(node1, col)
+        sched.add_node_collection(node2, col)
+        sched.schedule()
+        assert node1.sent == [0, 1]
+        assert node2.sent == [2, 3]
+        assert sched.pending == list(range(4, 16))
+        assert sched.node2pending[node1] == node1.sent
+        assert sched.node2pending[node2] == node2.sent
+
+        for complete_index, first_pending in enumerate(range(5, 16)):
+            sched.mark_test_complete(node1, node1.sent[complete_index])
+            assert node1.sent == [0, 1] + list(range(4, first_pending))
+            assert node2.sent == [2, 3]
+            assert sched.pending == list(range(first_pending, 16))
+
     def test_schedule_fewer_tests_than_nodes(self, pytester: pytest.Pytester) -> None:
         config = pytester.parseconfig("--tx=2*popen")
         sched = LoadScheduling(config)