Add a Scheduling Protocol

- Clearly delineate the required interface of schedulers
- Useful for typing
This commit is contained in:
Ran Benita
2024-04-05 14:38:21 +03:00
parent 91812bf5f9
commit bf2dcccfb0
6 changed files with 72 additions and 10 deletions

View File

@@ -15,6 +15,7 @@ from xdist.scheduler import LoadFileScheduling
from xdist.scheduler import LoadGroupScheduling
from xdist.scheduler import LoadScheduling
from xdist.scheduler import LoadScopeScheduling
from xdist.scheduler import Scheduling
from xdist.scheduler import WorkStealingScheduling
from xdist.workermanage import NodeManager
@@ -97,17 +98,21 @@ class DSession:
return True
@pytest.hookimpl(trylast=True)
def pytest_xdist_make_scheduler(self, config, log):
def pytest_xdist_make_scheduler(self, config, log) -> Scheduling | None:
dist = config.getvalue("dist")
schedulers = {
"each": EachScheduling,
"load": LoadScheduling,
"loadscope": LoadScopeScheduling,
"loadfile": LoadFileScheduling,
"loadgroup": LoadGroupScheduling,
"worksteal": WorkStealingScheduling,
}
return schedulers[dist](config, log)
if dist == "each":
return EachScheduling(config, log)
if dist == "load":
return LoadScheduling(config, log)
if dist == "loadscope":
return LoadScopeScheduling(config, log)
if dist == "loadfile":
return LoadFileScheduling(config, log)
if dist == "loadgroup":
return LoadGroupScheduling(config, log)
if dist == "worksteal":
return WorkStealingScheduling(config, log)
return None
@pytest.hookimpl
def pytest_runtestloop(self):

View File

@@ -3,4 +3,5 @@ from xdist.scheduler.load import LoadScheduling as LoadScheduling
from xdist.scheduler.loadfile import LoadFileScheduling as LoadFileScheduling
from xdist.scheduler.loadgroup import LoadGroupScheduling as LoadGroupScheduling
from xdist.scheduler.loadscope import LoadScopeScheduling as LoadScopeScheduling
from xdist.scheduler.protocol import Scheduling as Scheduling
from xdist.scheduler.worksteal import WorkStealingScheduling as WorkStealingScheduling

View File

@@ -103,6 +103,9 @@ class EachScheduling:
def mark_test_pending(self, item):
raise NotImplementedError()
def remove_pending_tests_from_node(self, node, indices):
raise NotImplementedError()
def remove_node(self, node):
# KeyError if we didn't get an add_node() yet
pending = self.node2pending.pop(node)

View File

@@ -160,6 +160,9 @@ class LoadScheduling:
for node in self.node2pending:
self.check_schedule(node)
def remove_pending_tests_from_node(self, node, indices):
raise NotImplementedError()
def check_schedule(self, node, duration=0):
"""Maybe schedule new items on the node.

View File

@@ -244,6 +244,9 @@ class LoadScopeScheduling:
def mark_test_pending(self, item):
raise NotImplementedError()
def remove_pending_tests_from_node(self, node, indices):
raise NotImplementedError()
def _assign_work_unit(self, node):
"""Assign a work unit to a node."""
assert self.workqueue

View File

@@ -0,0 +1,47 @@
from __future__ import annotations
from typing import Protocol
from typing import Sequence
from xdist.workermanage import WorkerController
class Scheduling(Protocol):
@property
def nodes(self) -> list[WorkerController]: ...
@property
def collection_is_completed(self) -> bool: ...
@property
def tests_finished(self) -> bool: ...
@property
def has_pending(self) -> bool: ...
def add_node(self, node: WorkerController) -> None: ...
def add_node_collection(
self,
node: WorkerController,
collection: Sequence[str],
) -> None: ...
def mark_test_complete(
self,
node: WorkerController,
item_index: int,
duration: float = 0,
) -> None: ...
def mark_test_pending(self, item: str) -> None: ...
def remove_pending_tests_from_node(
self,
node: WorkerController,
indices: Sequence[int],
) -> None: ...
def remove_node(self, node: WorkerController) -> str | None: ...
def schedule(self) -> None: ...