IQ.Lvbs/pytest-xdist
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Improve typing

Browse Source 
Fix #1057
This commit is contained in:
Ran Benita
2023-04-18 23:06:34 +03:00
parent 5dfc590a8c
commit b78cf1c0ce
24 changed files with 858 additions and 477 deletions
Download Patch File Download Diff File Expand all files Collapse all files

47
src/xdist/scheduler/each.py
View File

@@ -1,6 +1,13 @@
from __future__ import annotations
from typing import Sequence
import pytest
from xdist.remote import Producer
from xdist.report import report_collection_diff
from xdist.workermanage import parse_spec_config
from xdist.workermanage import WorkerController
class EachScheduling:
@@ -17,13 +24,13 @@ class EachScheduling:
assigned the remaining items from the removed node.
"""
def __init__(self, config, log=None):
def __init__(self, config: pytest.Config, log: Producer | None = None) -> None:
self.config = config
self.numnodes = len(parse_spec_config(config))
self.node2collection = {}
self.node2pending = {}
self._started = []
self._removed2pending = {}
self.node2collection: dict[WorkerController, list[str]] = {}
self.node2pending: dict[WorkerController, list[int]] = {}
self._started: list[WorkerController] = []
self._removed2pending: dict[WorkerController, list[int]] = {}
if log is None:
self.log = Producer("eachsched")
else:
@@ -31,12 +38,12 @@ class EachScheduling:
self.collection_is_completed = False
@property
def nodes(self):
def nodes(self) -> list[WorkerController]:
"""A list of all nodes in the scheduler."""
return list(self.node2pending.keys())
@property
def tests_finished(self):
def tests_finished(self) -> bool:
if not self.collection_is_completed:
return False
if self._removed2pending:
@@ -47,7 +54,7 @@ class EachScheduling:
return True
@property
def has_pending(self):
def has_pending(self) -> bool:
"""Return True if there are pending test items.
This indicates that collection has finished and nodes are
@@ -59,11 +66,13 @@ class EachScheduling:
return True
return False
def add_node(self, node):
def add_node(self, node: WorkerController) -> None:
assert node not in self.node2pending
self.node2pending[node] = []
def add_node_collection(self, node, collection):
def add_node_collection(
self, node: WorkerController, collection: Sequence[str]
) -> None:
"""Add the collected test items from a node.
Collection is complete once all nodes have submitted their
@@ -97,26 +106,32 @@ class EachScheduling:
self.node2pending[node] = pending
break
def mark_test_complete(self, node, item_index, duration=0):
def mark_test_complete(
self, node: WorkerController, item_index: int, duration: float = 0
) -> None:
self.node2pending[node].remove(item_index)
def mark_test_pending(self, item):
def mark_test_pending(self, item: str) -> None:
raise NotImplementedError()
def remove_pending_tests_from_node(self, node, indices):
def remove_pending_tests_from_node(
self,
node: WorkerController,
indices: Sequence[int],
) -> None:
raise NotImplementedError()
def remove_node(self, node):
def remove_node(self, node: WorkerController) -> str | None:
# KeyError if we didn't get an add_node() yet
pending = self.node2pending.pop(node)
if not pending:
return
return None
crashitem = self.node2collection[node][pending.pop(0)]
if pending:
self._removed2pending[node] = pending
return crashitem
def schedule(self):
def schedule(self) -> None:
"""Schedule the test items on the nodes.
If the node's pending list is empty it is a new node which

54
src/xdist/scheduler/load.py
View File

@@ -1,10 +1,14 @@
from __future__ import annotations
from itertools import cycle
from typing import Sequence
import pytest
from xdist.remote import Producer
from xdist.report import report_collection_diff
from xdist.workermanage import parse_spec_config
from xdist.workermanage import WorkerController
class LoadScheduling:
@@ -53,12 +57,12 @@ class LoadScheduling:
:config: Config object, used for handling hooks.
"""
def __init__(self, config, log=None):
def __init__(self, config: pytest.Config, log: Producer | None = None) -> None:
self.numnodes = len(parse_spec_config(config))
self.node2collection = {}
self.node2pending = {}
self.pending = []
self.collection = None
self.node2collection: dict[WorkerController, list[str]] = {}
self.node2pending: dict[WorkerController, list[int]] = {}
self.pending: list[int] = []
self.collection: list[str] | None = None
if log is None:
self.log = Producer("loadsched")
else:
@@ -67,12 +71,12 @@ class LoadScheduling:
self.maxschedchunk = self.config.getoption("maxschedchunk")
@property
def nodes(self):
def nodes(self) -> list[WorkerController]:
"""A list of all nodes in the scheduler."""
return list(self.node2pending.keys())
@property
def collection_is_completed(self):
def collection_is_completed(self) -> bool:
"""Boolean indication initial test collection is complete.
This is a boolean indicating all initial participating nodes
@@ -82,7 +86,7 @@ class LoadScheduling:
return len(self.node2collection) >= self.numnodes
@property
def tests_finished(self):
def tests_finished(self) -> bool:
"""Return True if all tests have been executed by the nodes."""
if not self.collection_is_completed:
return False
@@ -94,7 +98,7 @@ class LoadScheduling:
return True
@property
def has_pending(self):
def has_pending(self) -> bool:
"""Return True if there are pending test items.
This indicates that collection has finished and nodes are
@@ -108,7 +112,7 @@ class LoadScheduling:
return True
return False
def add_node(self, node):
def add_node(self, node: WorkerController) -> None:
"""Add a new node to the scheduler.
From now on the node will be allocated chunks of tests to
@@ -120,7 +124,9 @@ class LoadScheduling:
assert node not in self.node2pending
self.node2pending[node] = []
def add_node_collection(self, node, collection):
def add_node_collection(
self, node: WorkerController, collection: Sequence[str]
) -> None:
"""Add the collected test items from a node.
The collection is stored in the ``.node2collection`` map.
@@ -141,7 +147,9 @@ class LoadScheduling:
return
self.node2collection[node] = list(collection)
def mark_test_complete(self, node, item_index, duration=0):
def mark_test_complete(
self, node: WorkerController, item_index: int, duration: float = 0
) -> None:
"""Mark test item as completed by node.
The duration it took to execute the item is used as a hint to
@@ -152,7 +160,8 @@ class LoadScheduling:
self.node2pending[node].remove(item_index)
self.check_schedule(node, duration=duration)
def mark_test_pending(self, item):
def mark_test_pending(self, item: str) -> None:
assert self.collection is not None
self.pending.insert(
0,
self.collection.index(item),
@@ -160,10 +169,14 @@ class LoadScheduling:
for node in self.node2pending:
self.check_schedule(node)
def remove_pending_tests_from_node(self, node, indices):
def remove_pending_tests_from_node(
self,
node: WorkerController,
indices: Sequence[int],
) -> None:
raise NotImplementedError()
def check_schedule(self, node, duration=0):
def check_schedule(self, node: WorkerController, duration: float = 0) -> None:
"""Maybe schedule new items on the node.
If there are any globally pending nodes left then this will
@@ -197,7 +210,7 @@ class LoadScheduling:
self.log("num items waiting for node:", len(self.pending))
def remove_node(self, node):
def remove_node(self, node: WorkerController) -> str | None:
"""Remove a node from the scheduler.
This should be called either when the node crashed or at
@@ -212,16 +225,17 @@ class LoadScheduling:
"""
pending = self.node2pending.pop(node)
if not pending:
return
return None
# The node crashed, reassing pending items
assert self.collection is not None
crashitem = self.collection[pending.pop(0)]
self.pending.extend(pending)
for node in self.node2pending:
self.check_schedule(node)
return crashitem
def schedule(self):
def schedule(self) -> None:
"""Initiate distribution of the test collection.
Initiate scheduling of the items across the nodes. If this
@@ -285,14 +299,14 @@ class LoadScheduling:
for node in self.nodes:
node.shutdown()
def _send_tests(self, node, num):
def _send_tests(self, node: WorkerController, num: int) -> None:
tests_per_node = self.pending[:num]
if tests_per_node:
del self.pending[:num]
self.node2pending[node].extend(tests_per_node)
node.send_runtest_some(tests_per_node)
def _check_nodes_have_same_collection(self):
def _check_nodes_have_same_collection(self) -> bool:
"""Return True if all nodes have collected the same items.
If collections differ, this method returns False while logging

8
src/xdist/scheduler/loadfile.py
View File

@@ -1,3 +1,7 @@
from __future__ import annotations
import pytest
from xdist.remote import Producer
from .loadscope import LoadScopeScheduling
@@ -21,14 +25,14 @@ class LoadFileScheduling(LoadScopeScheduling):
This class behaves very much like LoadScopeScheduling, but with a file-level scope.
"""
def __init__(self, config, log=None):
def __init__(self, config: pytest.Config, log: Producer | None = None) -> None:
super().__init__(config, log)
if log is None:
self.log = Producer("loadfilesched")
else:
self.log = log.loadfilesched
def _split_scope(self, nodeid):
def _split_scope(self, nodeid: str) -> str:
"""Determine the scope (grouping) of a nodeid.
There are usually 3 cases for a nodeid::

8
src/xdist/scheduler/loadgroup.py
View File

@@ -1,3 +1,7 @@
from __future__ import annotations
import pytest
from xdist.remote import Producer
from .loadscope import LoadScopeScheduling
@@ -10,14 +14,14 @@ class LoadGroupScheduling(LoadScopeScheduling):
instead of the module or class to which they belong to.
"""
def __init__(self, config, log=None):
def __init__(self, config: pytest.Config, log: Producer | None = None) -> None:
super().__init__(config, log)
if log is None:
self.log = Producer("loadgroupsched")
else:
self.log = log.loadgroupsched
def _split_scope(self, nodeid):
def _split_scope(self, nodeid: str) -> str:
"""Determine the scope (grouping) of a nodeid.
There are usually 3 cases for a nodeid::

57
src/xdist/scheduler/loadscope.py
View File

@@ -1,10 +1,15 @@
from __future__ import annotations
from collections import OrderedDict
from typing import NoReturn
from typing import Sequence
import pytest
from xdist.remote import Producer
from xdist.report import report_collection_diff
from xdist.workermanage import parse_spec_config
from xdist.workermanage import WorkerController
class LoadScopeScheduling:
@@ -85,13 +90,13 @@ class LoadScopeScheduling:
:config: Config object, used for handling hooks.
"""
def __init__(self, config, log=None):
def __init__(self, config: pytest.Config, log: Producer | None = None) -> None:
self.numnodes = len(parse_spec_config(config))
self.collection = None
self.collection: list[str] | None = None
self.workqueue = OrderedDict()
self.assigned_work = {}
self.registered_collections = {}
self.workqueue: OrderedDict[str, dict[str, bool]] = OrderedDict()
self.assigned_work: dict[WorkerController, dict[str, dict[str, bool]]] = {}
self.registered_collections: dict[WorkerController, list[str]] = {}
if log is None:
self.log = Producer("loadscopesched")
@@ -101,12 +106,12 @@ class LoadScopeScheduling:
self.config = config
@property
def nodes(self):
def nodes(self) -> list[WorkerController]:
"""A list of all active nodes in the scheduler."""
return list(self.assigned_work.keys())
@property
def collection_is_completed(self):
def collection_is_completed(self) -> bool:
"""Boolean indication initial test collection is complete.
This is a boolean indicating all initial participating nodes have
@@ -116,7 +121,7 @@ class LoadScopeScheduling:
return len(self.registered_collections) >= self.numnodes
@property
def tests_finished(self):
def tests_finished(self) -> bool:
"""Return True if all tests have been executed by the nodes."""
if not self.collection_is_completed:
return False
@@ -131,7 +136,7 @@ class LoadScopeScheduling:
return True
@property
def has_pending(self):
def has_pending(self) -> bool:
"""Return True if there are pending test items.
This indicates that collection has finished and nodes are still
@@ -147,7 +152,7 @@ class LoadScopeScheduling:
return False
def add_node(self, node):
def add_node(self, node: WorkerController) -> None:
"""Add a new node to the scheduler.
From now on the node will be assigned work units to be executed.
@@ -158,7 +163,7 @@ class LoadScopeScheduling:
assert node not in self.assigned_work
self.assigned_work[node] = {}
def remove_node(self, node):
def remove_node(self, node: WorkerController) -> str | None:
"""Remove a node from the scheduler.
This should be called either when the node crashed or at shutdown time.
@@ -199,7 +204,9 @@ class LoadScopeScheduling:
return crashitem
def add_node_collection(self, node, collection):
def add_node_collection(
self, node: WorkerController, collection: Sequence[str]
) -> None:
"""Add the collected test items from a node.
The collection is stored in the ``.registered_collections`` dictionary.
@@ -228,7 +235,9 @@ class LoadScopeScheduling:
self.registered_collections[node] = list(collection)
def mark_test_complete(self, node, item_index, duration=0):
def mark_test_complete(
self, node: WorkerController, item_index: int, duration: float = 0
) -> None:
"""Mark test item as completed by node.
Called by the hook:
@@ -241,13 +250,17 @@ class LoadScopeScheduling:
self.assigned_work[node][scope][nodeid] = True
self._reschedule(node)
def mark_test_pending(self, item):
def mark_test_pending(self, item: str) -> NoReturn:
raise NotImplementedError()
def remove_pending_tests_from_node(self, node, indices):
def remove_pending_tests_from_node(
self,
node: WorkerController,
indices: Sequence[int],
) -> None:
raise NotImplementedError()
def _assign_work_unit(self, node):
def _assign_work_unit(self, node: WorkerController) -> None:
"""Assign a work unit to a node."""
assert self.workqueue
@@ -268,7 +281,7 @@ class LoadScopeScheduling:
node.send_runtest_some(nodeids_indexes)
def _split_scope(self, nodeid):
def _split_scope(self, nodeid: str) -> str:
"""Determine the scope (grouping) of a nodeid.
There are usually 3 cases for a nodeid::
@@ -292,12 +305,12 @@ class LoadScopeScheduling:
"""
return nodeid.rsplit("::", 1)[0]
def _pending_of(self, workload):
def _pending_of(self, workload: dict[str, dict[str, bool]]) -> int:
"""Return the number of pending tests in a workload."""
pending = sum(list(scope.values()).count(False) for scope in workload.values())
return pending
def _reschedule(self, node):
def _reschedule(self, node: WorkerController) -> None:
"""Maybe schedule new items on the node.
If there are any globally pending work units left then this will check
@@ -322,7 +335,7 @@ class LoadScopeScheduling:
# Pop one unit of work and assign it
self._assign_work_unit(node)
def schedule(self):
def schedule(self) -> None:
"""Initiate distribution of the test collection.
Initiate scheduling of the items across the nodes. If this gets called
@@ -352,7 +365,7 @@ class LoadScopeScheduling:
return
# Determine chunks of work (scopes)
unsorted_workqueue = {}
unsorted_workqueue: dict[str, dict[str, bool]] = {}
for nodeid in self.collection:
scope = self._split_scope(nodeid)
work_unit = unsorted_workqueue.setdefault(scope, {})
@@ -389,7 +402,7 @@ class LoadScopeScheduling:
for node in self.nodes:
node.shutdown()
def _check_nodes_have_same_collection(self):
def _check_nodes_have_same_collection(self) -> bool:
"""Return True if all nodes have collected the same items.
If collections differ, this method returns False while logging

59
src/xdist/scheduler/worksteal.py
View File

@@ -1,17 +1,18 @@
from __future__ import annotations
from typing import Any
from typing import NamedTuple
from typing import Sequence
import pytest
from xdist.remote import Producer
from xdist.report import report_collection_diff
from xdist.workermanage import parse_spec_config
from xdist.workermanage import WorkerController
class NodePending(NamedTuple):
node: Any
node: WorkerController
pending: list[int]
@@ -63,26 +64,26 @@ class WorkStealingScheduling:
simultaneous requests.
"""
def __init__(self, config, log=None):
def __init__(self, config: pytest.Config, log: Producer | None = None) -> None:
self.numnodes = len(parse_spec_config(config))
self.node2collection = {}
self.node2pending = {}
self.pending = []
self.collection = None
self.node2collection: dict[WorkerController, list[str]] = {}
self.node2pending: dict[WorkerController, list[int]] = {}
self.pending: list[int] = []
self.collection: list[str] | None = None
if log is None:
self.log = Producer("workstealsched")
else:
self.log = log.workstealsched
self.config = config
self.steal_requested_from_node = None
self.steal_requested_from_node: WorkerController | None = None
@property
def nodes(self):
def nodes(self) -> list[WorkerController]:
"""A list of all nodes in the scheduler."""
return list(self.node2pending.keys())
@property
def collection_is_completed(self):
def collection_is_completed(self) -> bool:
"""Boolean indication initial test collection is complete.
This is a boolean indicating all initial participating nodes
@@ -92,7 +93,7 @@ class WorkStealingScheduling:
return len(self.node2collection) >= self.numnodes
@property
def tests_finished(self):
def tests_finished(self) -> bool:
"""Return True if all tests have been executed by the nodes."""
if not self.collection_is_completed:
return False
@@ -106,7 +107,7 @@ class WorkStealingScheduling:
return True
@property
def has_pending(self):
def has_pending(self) -> bool:
"""Return True if there are pending test items.
This indicates that collection has finished and nodes are
@@ -120,7 +121,7 @@ class WorkStealingScheduling:
return True
return False
def add_node(self, node):
def add_node(self, node: WorkerController) -> None:
"""Add a new node to the scheduler.
From now on the node will be allocated chunks of tests to
@@ -132,7 +133,9 @@ class WorkStealingScheduling:
assert node not in self.node2pending
self.node2pending[node] = []
def add_node_collection(self, node, collection):
def add_node_collection(
self, node: WorkerController, collection: Sequence[str]
) -> None:
"""Add the collected test items from a node.
The collection is stored in the ``.node2collection`` map.
@@ -153,7 +156,9 @@ class WorkStealingScheduling:
return
self.node2collection[node] = list(collection)
def mark_test_complete(self, node, item_index, duration=None):
def mark_test_complete(
self, node: WorkerController, item_index: int, duration: float | None = None
) -> None:
"""Mark test item as completed by node.
This is called by the ``DSession.worker_testreport`` hook.
@@ -161,14 +166,19 @@ class WorkStealingScheduling:
self.node2pending[node].remove(item_index)
self.check_schedule()
def mark_test_pending(self, item):
def mark_test_pending(self, item: str) -> None:
assert self.collection is not None
self.pending.insert(
0,
self.collection.index(item),
)
self.check_schedule()
def remove_pending_tests_from_node(self, node, indices):
def remove_pending_tests_from_node(
self,
node: WorkerController,
indices: Sequence[int],
) -> None:
"""Node returned some test indices back in response to 'steal' command.
This is called by ``DSession.worker_unscheduled``.
@@ -183,7 +193,7 @@ class WorkStealingScheduling:
self.pending.extend(indices)
self.check_schedule()
def check_schedule(self):
def check_schedule(self) -> None:
"""Reschedule tests/perform load balancing."""
nodes_up = [
NodePending(node, pending)
@@ -191,7 +201,7 @@ class WorkStealingScheduling:
if not node.shutting_down
]
def get_idle_nodes():
def get_idle_nodes() -> list[WorkerController]:
return [node for node, pending in nodes_up if len(pending) < MIN_PENDING]
idle_nodes = get_idle_nodes()
@@ -235,10 +245,11 @@ class WorkStealingScheduling:
node.shutdown()
return
assert steal_from is not None
steal_from.node.send_steal(steal_from.pending[-num_steal:])
self.steal_requested_from_node = steal_from.node
def remove_node(self, node):
def remove_node(self, node: WorkerController) -> str | None:
"""Remove a node from the scheduler.
This should be called either when the node crashed or at
@@ -249,12 +260,12 @@ class WorkStealingScheduling:
Return the item which was being executing while the node
crashed or None if the node has no more pending items.
"""
pending = self.node2pending.pop(node)
# If node was removed without completing its assigned tests - it crashed
if pending:
assert self.collection is not None
crashitem = self.collection[pending.pop(0)]
else:
crashitem = None
@@ -268,7 +279,7 @@ class WorkStealingScheduling:
self.check_schedule()
return crashitem
def schedule(self):
def schedule(self) -> None:
"""Initiate distribution of the test collection.
Initiate scheduling of the items across the nodes. If this
@@ -298,14 +309,14 @@ class WorkStealingScheduling:
self.check_schedule()
def _send_tests(self, node, num):
def _send_tests(self, node: WorkerController, num: int) -> None:
tests_per_node = self.pending[:num]
if tests_per_node:
del self.pending[:num]
self.node2pending[node].extend(tests_per_node)
node.send_runtest_some(tests_per_node)
def _check_nodes_have_same_collection(self):
def _check_nodes_have_same_collection(self) -> bool:
"""Return True if all nodes have collected the same items.
If collections differ, this method returns False while logging