Merge pull request #89 from wronglink/custom_scheduler

[WIP] Allow custom scheduler class implementation

testing/test_dsession.py

@@ -49,80 +49,84 @@ def dumpqueue(queue):
 
 
 class TestEachScheduling:
-    def test_schedule_load_simple(self):
+    def test_schedule_load_simple(self, testdir):
         node1 = MockNode()
         node2 = MockNode()
-        sched = EachScheduling(2)
+        config = testdir.parseconfig("--tx=2*popen")
-        sched.addnode(node1)
+        sched = EachScheduling(config)
-        sched.addnode(node2)
+        sched.add_node(node1)
+        sched.add_node(node2)
         collection = ["a.py::test_1", ]
         assert not sched.collection_is_completed
-        sched.addnode_collection(node1, collection)
+        sched.add_node_collection(node1, collection)
         assert not sched.collection_is_completed
-        sched.addnode_collection(node2, collection)
+        sched.add_node_collection(node2, collection)
         assert sched.collection_is_completed
         assert sched.node2collection[node1] == collection
         assert sched.node2collection[node2] == collection
-        sched.init_distribute()
+        sched.schedule()
-        assert sched.tests_finished()
+        assert sched.tests_finished
         assert node1.sent == ['ALL']
         assert node2.sent == ['ALL']
-        sched.remove_item(node1, 0)
+        sched.mark_test_complete(node1, 0)
-        assert sched.tests_finished()
+        assert sched.tests_finished
-        sched.remove_item(node2, 0)
+        sched.mark_test_complete(node2, 0)
-        assert sched.tests_finished()
+        assert sched.tests_finished
 
-    def test_schedule_remove_node(self):
+    def test_schedule_remove_node(self, testdir):
         node1 = MockNode()
-        sched = EachScheduling(1)
+        config = testdir.parseconfig("--tx=popen")
-        sched.addnode(node1)
+        sched = EachScheduling(config)
+        sched.add_node(node1)
         collection = ["a.py::test_1", ]
         assert not sched.collection_is_completed
-        sched.addnode_collection(node1, collection)
+        sched.add_node_collection(node1, collection)
         assert sched.collection_is_completed
         assert sched.node2collection[node1] == collection
-        sched.init_distribute()
+        sched.schedule()
-        assert sched.tests_finished()
+        assert sched.tests_finished
         crashitem = sched.remove_node(node1)
         assert crashitem
-        assert sched.tests_finished()
+        assert sched.tests_finished
-        assert not sched.hasnodes()
+        assert not sched.nodes
 
 
 class TestLoadScheduling:
-    def test_schedule_load_simple(self):
+    def test_schedule_load_simple(self, testdir):
-        sched = LoadScheduling(2)
+        config = testdir.parseconfig("--tx=2*popen")
-        sched.addnode(MockNode())
+        sched = LoadScheduling(config)
-        sched.addnode(MockNode())
+        sched.add_node(MockNode())
+        sched.add_node(MockNode())
         node1, node2 = sched.nodes
         collection = ["a.py::test_1", "a.py::test_2"]
         assert not sched.collection_is_completed
-        sched.addnode_collection(node1, collection)
+        sched.add_node_collection(node1, collection)
         assert not sched.collection_is_completed
-        sched.addnode_collection(node2, collection)
+        sched.add_node_collection(node2, collection)
         assert sched.collection_is_completed
         assert sched.node2collection[node1] == collection
         assert sched.node2collection[node2] == collection
-        sched.init_distribute()
+        sched.schedule()
         assert not sched.pending
-        assert sched.tests_finished()
+        assert sched.tests_finished
         assert len(node1.sent) == 1
         assert len(node2.sent) == 1
         assert node1.sent == [0]
         assert node2.sent == [1]
-        sched.remove_item(node1, node1.sent[0])
+        sched.mark_test_complete(node1, node1.sent[0])
-        assert sched.tests_finished()
+        assert sched.tests_finished
 
-    def test_init_distribute_batch_size(self):
+    def test_schedule_batch_size(self, testdir):
-        sched = LoadScheduling(2)
+        config = testdir.parseconfig("--tx=2*popen")
-        sched.addnode(MockNode())
+        sched = LoadScheduling(config)
-        sched.addnode(MockNode())
+        sched.add_node(MockNode())
+        sched.add_node(MockNode())
         node1, node2 = sched.nodes
         col = ["xyz"] * (6)
-        sched.addnode_collection(node1, col)
+        sched.add_node_collection(node1, col)
-        sched.addnode_collection(node2, col)
+        sched.add_node_collection(node2, col)
-        sched.init_distribute()
+        sched.schedule()
-        # assert not sched.tests_finished()
+        # assert not sched.tests_finished
         sent1 = node1.sent
         sent2 = node2.sent
         assert sent1 == [0, 2]
@@ -131,25 +135,26 @@ class TestLoadScheduling:
         assert sched.node2pending[node1] == sent1
         assert sched.node2pending[node2] == sent2
         assert len(sched.pending) == 2
-        sched.remove_item(node1, 0)
+        sched.mark_test_complete(node1, 0)
         assert node1.sent == [0, 2, 4]
         assert sched.pending == [5]
         assert node2.sent == [1, 3]
-        sched.remove_item(node1, 2)
+        sched.mark_test_complete(node1, 2)
         assert node1.sent == [0, 2, 4, 5]
         assert not sched.pending
 
-    def test_init_distribute_fewer_tests_than_nodes(self):
+    def test_schedule_fewer_tests_than_nodes(self, testdir):
-        sched = LoadScheduling(2)
+        config = testdir.parseconfig("--tx=2*popen")
-        sched.addnode(MockNode())
+        sched = LoadScheduling(config)
-        sched.addnode(MockNode())
+        sched.add_node(MockNode())
-        sched.addnode(MockNode())
+        sched.add_node(MockNode())
+        sched.add_node(MockNode())
         node1, node2, node3 = sched.nodes
         col = ["xyz"] * 2
-        sched.addnode_collection(node1, col)
+        sched.add_node_collection(node1, col)
-        sched.addnode_collection(node2, col)
+        sched.add_node_collection(node2, col)
-        sched.init_distribute()
+        sched.schedule()
-        # assert not sched.tests_finished()
+        # assert not sched.tests_finished
         sent1 = node1.sent
         sent2 = node2.sent
         sent3 = node3.sent
@@ -158,17 +163,18 @@ class TestLoadScheduling:
         assert sent3 == []
         assert not sched.pending
 
-    def test_init_distribute_fewer_than_two_tests_per_node(self):
+    def test_schedule_fewer_than_two_tests_per_node(self, testdir):
-        sched = LoadScheduling(2)
+        config = testdir.parseconfig("--tx=2*popen")
-        sched.addnode(MockNode())
+        sched = LoadScheduling(config)
-        sched.addnode(MockNode())
+        sched.add_node(MockNode())
-        sched.addnode(MockNode())
+        sched.add_node(MockNode())
+        sched.add_node(MockNode())
         node1, node2, node3 = sched.nodes
         col = ["xyz"] * 5
-        sched.addnode_collection(node1, col)
+        sched.add_node_collection(node1, col)
-        sched.addnode_collection(node2, col)
+        sched.add_node_collection(node2, col)
-        sched.init_distribute()
+        sched.schedule()
-        # assert not sched.tests_finished()
+        # assert not sched.tests_finished
         sent1 = node1.sent
         sent2 = node2.sent
         sent3 = node3.sent
@@ -177,14 +183,15 @@ class TestLoadScheduling:
         assert sent3 == [2]
         assert not sched.pending
 
-    def test_add_remove_node(self):
+    def test_add_remove_node(self, testdir):
         node = MockNode()
-        sched = LoadScheduling(1)
+        config = testdir.parseconfig("--tx=popen")
-        sched.addnode(node)
+        sched = LoadScheduling(config)
+        sched.add_node(node)
         collection = ["test_file.py::test_func"]
-        sched.addnode_collection(node, collection)
+        sched.add_node_collection(node, collection)
         assert sched.collection_is_completed
-        sched.init_distribute()
+        sched.schedule()
         assert not sched.pending
         crashitem = sched.remove_node(node)
         assert crashitem == collection[0]
@@ -207,16 +214,16 @@ class TestLoadScheduling:
                 self.reports.append(report)
 
         collect_hook = CollectHook()
-        config = testdir.parseconfig()
+        config = testdir.parseconfig("--tx=2*popen")
         config.pluginmanager.register(collect_hook, "collect_hook")
         node1 = MockNode()
         node2 = MockNode()
-        sched = LoadScheduling(2, config=config)
+        sched = LoadScheduling(config)
-        sched.addnode(node1)
+        sched.add_node(node1)
-        sched.addnode(node2)
+        sched.add_node(node2)
-        sched.addnode_collection(node1, ["a.py::test_1"])
+        sched.add_node_collection(node1, ["a.py::test_1"])
-        sched.addnode_collection(node2, ["a.py::test_2"])
+        sched.add_node_collection(node2, ["a.py::test_2"])
-        sched.init_distribute()
+        sched.schedule()
         assert len(collect_hook.reports) == 1
         rep = collect_hook.reports[0]
         assert 'Different tests were collected between' in rep.longrepr

xdist/dsession.py

@@ -4,7 +4,7 @@ from _pytest.runner import CollectReport
 
 import pytest
 import py
-from xdist.slavemanage import NodeManager
+from xdist.slavemanage import NodeManager, parse_spec_config
 
 
 queue = py.builtin._tryimport('queue', 'Queue')
@@ -24,8 +24,9 @@ class EachScheduling:
     assigned the remaining items from the removed node.
     """
 
-    def __init__(self, numnodes, log=None):
+    def __init__(self, config, log=None):
-        self.numnodes = numnodes
+        self.config = config
+        self.numnodes = len(parse_spec_config(config))
         self.node2collection = {}
         self.node2pending = {}
         self._started = []
@@ -41,10 +42,19 @@ class EachScheduling:
         """A list of all nodes in the scheduler."""
         return list(self.node2pending.keys())
 
-    def hasnodes(self):
+    @property
-        return bool(self.node2pending)
+    def tests_finished(self):
+        if not self.collection_is_completed:
+            return False
+        if self._removed2pending:
+            return False
+        for pending in self.node2pending.values():
+            if len(pending) >= 2:
+                return False
+        return True
 
-    def haspending(self):
+    @property
+    def has_pending(self):
         """Return True if there are pending test items
 
         This indicates that collection has finished and nodes are
@@ -56,21 +66,11 @@ class EachScheduling:
                 return True
         return False
 
-    def addnode(self, node):
+    def add_node(self, node):
         assert node not in self.node2pending
         self.node2pending[node] = []
 
-    def tests_finished(self):
+    def add_node_collection(self, node, collection):
-        if not self.collection_is_completed:
-            return False
-        if self._removed2pending:
-            return False
-        for pending in self.node2pending.values():
-            if len(pending) >= 2:
-                return False
-        return True
-
-    def addnode_collection(self, node, collection):
         """Add the collected test items from a node
 
         Collection is complete once all nodes have submitted their
@@ -78,7 +78,7 @@ class EachScheduling:
         list.  When the collection is already completed this
         submission is from a node which was restarted to replace a
         dead node.  In this case we already assign the pending items
-        here.  In either case ``.init_distribute()`` will instruct the
+        here.  In either case ``.schedule()`` will instruct the
         node to start running the required tests.
         """
         assert node in self.node2pending
@@ -102,11 +102,11 @@ class EachScheduling:
                     self.node2pending[node] = pending
                     break
 
-    def remove_item(self, node, item_index, duration=0):
+    def mark_test_complete(self, node, item_index, duration=0):
         self.node2pending[node].remove(item_index)
 
     def remove_node(self, node):
-        # KeyError if we didn't get an addnode() yet
+        # KeyError if we didn't get an add_node() yet
         pending = self.node2pending.pop(node)
         if not pending:
             return
@@ -115,12 +115,12 @@ class EachScheduling:
             self._removed2pending[node] = pending
         return crashitem
 
-    def init_distribute(self):
+    def schedule(self):
         """Schedule the test items on the nodes
 
         If the node's pending list is empty it is a new node which
         needs to run all the tests.  If the pending list is already
-        populated (by ``.addnode_collection()``) then it replaces a
+        populated (by ``.add_node_collection()``) then it replaces a
         dead node and we only need to run those tests.
         """
         assert self.collection_is_completed
@@ -143,7 +143,7 @@ class LoadScheduling:
     when all collections are received it is verified they are
     identical collections.  Then the collection gets divided up in
     chunks and chunks get submitted to nodes.  Whenever a node finishes
-    an item, it calls ``.remove_item()`` which will trigger the
+    an item, it calls ``.mark_test_complete()`` which will trigger the
     scheduler to assign more tests if the number of pending tests for
     the node falls below a low-watermark.
 
@@ -170,7 +170,7 @@ class LoadScheduling:
 
     :collection: The one collection once it is validated to be
        identical between all the nodes.  It is initialised to None
-       until ``.init_distribute()`` is called.
+       until ``.schedule()`` is called.
 
     :pending: List of indices of globally pending tests.  These are
        tests which have not yet been allocated to a chunk for a node
@@ -181,8 +181,8 @@ class LoadScheduling:
     :config: Config object, used for handling hooks.
     """
 
-    def __init__(self, numnodes, log=None, config=None):
+    def __init__(self, config, log=None):
-        self.numnodes = numnodes
+        self.numnodes = len(parse_spec_config(config))
         self.node2collection = {}
         self.node2pending = {}
         self.pending = []
@@ -208,7 +208,20 @@ class LoadScheduling:
         """
         return len(self.node2collection) >= self.numnodes
 
-    def haspending(self):
+    @property
+    def tests_finished(self):
+        """Return True if all tests have been executed by the nodes."""
+        if not self.collection_is_completed:
+            return False
+        if self.pending:
+            return False
+        for pending in self.node2pending.values():
+            if len(pending) >= 2:
+                return False
+        return True
+
+    @property
+    def has_pending(self):
         """Return True if there are pending test items
 
         This indicates that collection has finished and nodes are
@@ -222,11 +235,7 @@ class LoadScheduling:
                 return True
         return False
 
-    def hasnodes(self):
+    def add_node(self, node):
-        """Return True if nodes exist in the scheduler."""
-        return bool(self.node2pending)
-
-    def addnode(self, node):
         """Add a new node to the scheduler.
 
         From now on the node will be allocated chunks of tests to
@@ -238,18 +247,7 @@ class LoadScheduling:
         assert node not in self.node2pending
         self.node2pending[node] = []
 
-    def tests_finished(self):
+    def add_node_collection(self, node, collection):
-        """Return True if all tests have been executed by the nodes."""
-        if not self.collection_is_completed:
-            return False
-        if self.pending:
-            return False
-        for pending in self.node2pending.values():
-            if len(pending) >= 2:
-                return False
-        return True
-
-    def addnode_collection(self, node, collection):
         """Add the collected test items from a node
 
         The collection is stored in the ``.node2collection`` map.
@@ -258,7 +256,7 @@ class LoadScheduling:
         assert node in self.node2pending
         if self.collection_is_completed:
             # A new node has been added later, perhaps an original one died.
-            # .init_distribute() should have
+            # .schedule() should have
             # been called by now
             assert self.collection
             if collection != self.collection:
@@ -271,7 +269,7 @@ class LoadScheduling:
                 return
         self.node2collection[node] = list(collection)
 
-    def remove_item(self, node, item_index, duration=0):
+    def mark_test_complete(self, node, item_index, duration=0):
         """Mark test item as completed by node
 
         The duration it took to execute the item is used as a hint to
@@ -335,7 +333,7 @@ class LoadScheduling:
             self.check_schedule(node)
         return crashitem
 
-    def init_distribute(self):
+    def schedule(self):
         """Initiate distribution of the test collection
 
         Initiate scheduling of the items across the nodes.  If this
@@ -345,8 +343,6 @@ class LoadScheduling:
 
         This is called by the ``DSession.slave_collectionfinish`` hook
         if ``.collection_is_completed`` is True.
-
-        XXX Perhaps this method should have been called ".schedule()".
         """
         assert self.collection_is_completed
 
@@ -466,6 +462,8 @@ class DSession:
         self.log = py.log.Producer("dsession")
         if not config.option.debug:
             py.log.setconsumer(self.log._keywords, None)
+        self.nodemanager = None
+        self.sched = None
         self.shuttingdown = False
         self.countfailures = 0
         self.maxfail = config.getvalue("maxfail")
@@ -521,16 +519,21 @@ class DSession:
         # prohibit collection of test items in master process
         return True
 
-    def pytest_runtestloop(self):
+    @pytest.mark.trylast
-        numnodes = len(self.nodemanager.specs)
+    def pytest_xdist_make_scheduler(self, config, log):
-        dist = self.config.getvalue("dist")
+        dist = config.getvalue("dist")
         if dist == "load":
-            self.sched = LoadScheduling(numnodes, log=self.log,
+            return LoadScheduling(config, log)
-                                        config=self.config)
         elif dist == "each":
-            self.sched = EachScheduling(numnodes, log=self.log)
+            return EachScheduling(config, log)
-        else:
+
-            assert 0, dist
+    def pytest_runtestloop(self):
+        self.sched = self.config.hook.pytest_xdist_make_scheduler(
+            config=self.config,
+            log=self.log
+        )
+        assert self.sched is not None
+
         self.shouldstop = False
         while not self.session_finished:
             self.loop_once()
@@ -553,7 +556,7 @@ class DSession:
         call = getattr(self, method)
         self.log("calling method", method, kwargs)
         call(**kwargs)
-        if self.sched.tests_finished():
+        if self.sched.tests_finished:
             self.triggershutdown()
 
     #
@@ -573,7 +576,7 @@ class DSession:
         if self.shuttingdown:
             node.shutdown()
         else:
-            self.sched.addnode(node)
+            self.sched.add_node(node)
 
     def slave_slavefinished(self, node):
         """Emitted when node executes its pytest_sessionfinish hook.
@@ -635,16 +638,16 @@ class DSession:
         # tell session which items were effectively collected otherwise
         # the master node will finish the session with EXIT_NOTESTSCOLLECTED
         self._session.testscollected = len(ids)
-        self.sched.addnode_collection(node, ids)
+        self.sched.add_node_collection(node, ids)
         if self.terminal:
             self.trdist.setstatus(node.gateway.spec, "[%d]" % (len(ids)))
         if self.sched.collection_is_completed:
-            if self.terminal and not self.sched.haspending():
+            if self.terminal and not self.sched.has_pending:
                 self.trdist.ensure_show_status()
                 self.terminal.write_line("")
                 self.terminal.write_line("scheduling tests via %s" % (
                     self.sched.__class__.__name__))
-            self.sched.init_distribute()
+            self.sched.schedule()
 
     def slave_logstart(self, node, nodeid, location):
         """Emitted when a node calls the pytest_runtest_logstart hook."""
@@ -658,7 +661,7 @@ class DSession:
         the item from the pending list in the scheduler.
         """
         if rep.when == "call" or (rep.when == "setup" and not rep.passed):
-            self.sched.remove_item(node, rep.item_index, rep.duration)
+            self.sched.mark_test_complete(node, rep.item_index, rep.duration)
         # self.report_line("testreport %s: %s" %(rep.id, rep.status))
         rep.node = node
         self.config.hook.pytest_runtest_logreport(report=rep)

xdist/newhooks.py

@@ -11,6 +11,7 @@ must be taken in plugins in case ``xdist`` is not installed. Please see:
 
     http://pytest.org/latest/writing_plugins.html#optionally-using-hooks-from-3rd-party-plugins
 """
+import pytest
 
 
 def pytest_xdist_setupnodes(config, specs):
@@ -44,3 +45,8 @@ def pytest_testnodedown(node, error):
 def pytest_xdist_node_collection_finished(node, ids):
     """called by the master node when a node finishes collecting.
     """
+
+
+@pytest.mark.firstresult
+def pytest_xdist_make_scheduler(config, log):
+    """ return a node scheduler implementation """

xdist/slavemanage.py

@@ -10,6 +10,22 @@ import xdist.remote
 from _pytest import runner  # XXX load dynamically
 
 
+def parse_spec_config(config):
+    xspeclist = []
+    for xspec in config.getvalue("tx"):
+        i = xspec.find("*")
+        try:
+            num = int(xspec[:i])
+        except ValueError:
+            xspeclist.append(xspec)
+        else:
+            xspeclist.extend([xspec[i + 1:]] * num)
+    if not xspeclist:
+        raise pytest.UsageError(
+            "MISSING test execution (tx) nodes: please specify --tx")
+    return xspeclist
+
+
 class NodeManager(object):
     EXIT_TIMEOUT = 10
     DEFAULT_IGNORES = ['.*', '*.pyc', '*.pyo', '*~']
@@ -62,19 +78,7 @@ class NodeManager(object):
         self.group.terminate(self.EXIT_TIMEOUT)
 
     def _getxspecs(self):
-        xspeclist = []
+        return [execnet.XSpec(x) for x in parse_spec_config(self.config)]
-        for xspec in self.config.getvalue("tx"):
-            i = xspec.find("*")
-            try:
-                num = int(xspec[:i])
-            except ValueError:
-                xspeclist.append(xspec)
-            else:
-                xspeclist.extend([xspec[i+1:]] * num)
-        if not xspeclist:
-            raise pytest.UsageError(
-                "MISSING test execution (tx) nodes: please specify --tx")
-        return [execnet.XSpec(x) for x in xspeclist]
 
     def _getrsyncdirs(self):
         for spec in self.specs: