Skip to content

Commit

Permalink
Add --maxschedchunk CLI option
Browse files Browse the repository at this point in the history
Maximum number of tests scheduled in one step.

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.

Fixes #855
  • Loading branch information
amezin committed Dec 20, 2022
1 parent 7faa69a commit 0c23da7
Show file tree
Hide file tree
Showing 3 changed files with 72 additions and 3 deletions.
13 changes: 13 additions & 0 deletions src/xdist/plugin.py
Original file line number Diff line number Diff line change
Expand Up @@ -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. "
"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",
Expand Down
12 changes: 9 additions & 3 deletions src/xdist/scheduler/load.py
Original file line number Diff line number Diff line change
Expand Up @@ -64,6 +64,7 @@ def __init__(self, config, log=None):
else:
self.log = log.loadsched
self.config = config
self.maxschedchunk = self.config.getoption("maxschedchunk")

@property
def nodes(self):
Expand Down Expand Up @@ -185,7 +186,9 @@ def check_schedule(self, node, duration=0):
# 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()

Expand Down Expand Up @@ -245,6 +248,9 @@ def schedule(self):
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.
Expand All @@ -265,10 +271,10 @@ def schedule(self):
# 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)
# and initialize each node with a chunk of tests
for node in self.nodes:
self._send_tests(node, node_chunksize)
self._send_tests(node, max(node_chunksize, 2))

if not self.pending:
# initial distribution sent all tests, start node shutdown
Expand Down
50 changes: 50 additions & 0 deletions testing/test_dsession.py
Original file line number Diff line number Diff line change
Expand Up @@ -129,6 +129,56 @@ def test_schedule_batch_size(self, pytester: pytest.Pytester) -> None:
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)
Expand Down

0 comments on commit 0c23da7

Please sign in to comment.