Ported job calculation to use coroutines for tasks (#1827)

This way the event loop has a chance to schedule other operations while
a transport request is made i.e. if multiple job calculations all
request the same transport (a common scenario) then in the past they
would all run as one loop callback effectively blocking the loop (and
crucially `pika` heartbeats).  Now they are schedule as separate callbacks
thanks to yield, hopefully, allowing the loop to schedule other
necessary tasks in between.

aiida/backends/tests/__init__.py

@@ -92,6 +92,7 @@
         'work.rmq': ['aiida.backends.tests.work.test_rmq'],
         'work.run': ['aiida.backends.tests.work.run'],
         'work.runners': ['aiida.backends.tests.work.test_runners'],
+        'work.test_transport': ['aiida.backends.tests.work.test_transport'],
         'work.utils': ['aiida.backends.tests.work.utils'],
         'work.work_chain': ['aiida.backends.tests.work.work_chain'],
         'work.workfunctions': ['aiida.backends.tests.work.test_workfunctions'],

aiida/backends/tests/computer.py

@@ -15,8 +15,6 @@
 from aiida.common.exceptions import NotExistent
 
 
-
-
 class TestComputer(AiidaTestCase):
 
     def test_get_transport(self):
@@ -26,17 +24,18 @@ def test_get_transport(self):
         import tempfile
         from aiida.orm import Computer
         from aiida.orm.backend import construct_backend
-        backend = construct_backend()
 
         new_comp = Computer(name='bbb',
-                                hostname='localhost',
-                                transport_type='local',
-                                scheduler_type='direct',
-                                workdir='/tmp/aiida')
+                            hostname='localhost',
+                            transport_type='local',
+                            scheduler_type='direct',
+                            workdir='/tmp/aiida')
         new_comp.store()
 
         # Configure the computer - no parameters for local transport
-        authinfo = backend.authinfos.create(computer=new_comp, user=backend.users.get_automatic_user())
+        authinfo = self.backend.authinfos.create(
+            computer=new_comp,
+            user=self.backend.users.get_automatic_user())
         authinfo.store()
 
         transport = new_comp.get_transport()
@@ -51,20 +50,19 @@ def test_get_transport(self):
     def test_delete(self):
         from aiida.orm import Computer
         new_comp = Computer(name='aaa',
-                                hostname='aaa',
-                                transport_type='local',
-                                scheduler_type='pbspro',
-                                workdir='/tmp/aiida')
+                            hostname='aaa',
+                            transport_type='local',
+                            scheduler_type='pbspro',
+                            workdir='/tmp/aiida')
         new_comp.store()
 
         comp_pk = new_comp.pk
 
         check_computer = Computer.get(comp_pk)
         self.assertEquals(comp_pk, check_computer.pk)
-        
+
         from aiida.orm.computer import delete_computer
         delete_computer(pk=comp_pk)
 
         with self.assertRaises(NotExistent):
             Computer.get(comp_pk)
-

aiida/backends/tests/work/test_transport.py

@@ -0,0 +1,78 @@
+from tornado.gen import coroutine, Return
+
+from aiida.backends.testbase import AiidaTestCase
+from aiida.work.transports import TransportQueue
+
+
+class TestTransportQueue(AiidaTestCase):
+    """ Tests for the transport queue """
+
+    @classmethod
+    def setUpClass(cls, *args, **kwargs):
+        """ Set up a simple authinfo and for later use """
+        super(TestTransportQueue, cls).setUpClass(*args, **kwargs)
+        # Configure the computer - no parameters for local transport
+        # WARNING: This is not deleted as there is no API facing way to do this
+        # it would require a backend specific call
+        cls.authinfo = cls.backend.authinfos.create(
+            computer=cls.computer,
+            user=cls.backend.users.get_automatic_user())
+        cls.authinfo.store()
+
+    def test_simple_request(self):
+        """ Test a simple transport request """
+        queue = TransportQueue()
+        loop = queue.loop()
+
+        @coroutine
+        def test():
+            trans = None
+            with queue.request_transport(self.authinfo) as request:
+                trans = yield request
+                self.assertTrue(trans.is_open)
+            self.assertFalse(trans.is_open)
+
+        loop.run_sync(lambda: test())
+
+    def test_get_transport_nested(self):
+        """ Test nesting calls to get the same transport """
+        transport_queue = TransportQueue()
+        loop = transport_queue.loop()
+
+        @coroutine
+        def nested(queue, authinfo):
+            with queue.request_transport(authinfo) as request1:
+                trans1 = yield request1
+                self.assertTrue(trans1.is_open)
+                with queue.request_transport(authinfo) as request2:
+                    trans2 = yield request2
+                    self.assertIs(trans1, trans2)
+                    self.assertTrue(trans2.is_open)
+
+        loop.run_sync(lambda: nested(transport_queue, self.authinfo))
+
+    def test_get_transport_interleaved(self):
+        """ Test interleaved calls to get the same transport """
+        transport_queue = TransportQueue()
+        loop = transport_queue.loop()
+
+        @coroutine
+        def interleaved(authinfo):
+            with transport_queue.request_transport(authinfo) as trans_future:
+                yield trans_future
+
+        loop.run_sync(lambda: [interleaved(self.authinfo), interleaved(self.authinfo)])
+
+    def test_return_from_context(self):
+        """ Test raising a Return from coroutine context """
+        queue = TransportQueue()
+        loop = queue.loop()
+
+        @coroutine
+        def test():
+            with queue.request_transport(self.authinfo) as request:
+                trans = yield request
+                raise Return(trans.is_open)
+
+        retval = loop.run_sync(lambda: test())
+        self.assertTrue(retval)

aiida/transport/transport.py

@@ -72,7 +72,12 @@ def __enter__(self):
         """
         # Keep track of how many times enter has been called
         if self._enters == 0:
-            self.open()
+            if self.is_open:
+                # Already open, so just add one to the entered counter
+                # this way on the final exit we will not close
+                self._enters += 1
+            else:
+                self.open()
         self._enters += 1
         return self