add async pull_image

chainlink/__init__.py

@@ -24,14 +24,17 @@ def __init__(self, stages, workdir="/tmp"):
         self.workdir = workdir
         self._executor = concurrent.futures.ThreadPoolExecutor(max_workers=1)
         self._pull_status = {}
-        self._pull_images()
 
     # sync version
     def run(self, environ={}):
+        self._pull_images()
         return asyncio.get_event_loop().run_until_complete(self.run_async(environ))
 
     async def run_async(self, environ={}):
         results = []
+
+        await self._pull_images_async()
+
         with tempfile.TemporaryDirectory(dir=self.workdir) as mount:
             logger.info("using {} for temporary job directory".format(mount))
 
@@ -63,6 +66,32 @@ def _pull_images(self):
             if not self._pull_status.get(image):
                 raise ValueError("Failed to pull all images")
 
+    async def _pull_images_async(self):
+        images = set([stage["image"] for stage in self.stages])
+        threads = []
+
+        event_loop = asyncio.get_event_loop()
+
+        for image in images:
+            logger.debug("pulling image '{}'".format(image))
+            t = threading.Thread(
+                target=self._pull_image, args=(self.client, image, self._pull_status)
+            )
+            t.start()
+            threads.append(t)
+
+        def wait():
+            for t in threads:
+                t.join()
+
+        await asyncio.wait_for(
+            event_loop.run_in_executor(self._executor, wait), timeout=None
+        )
+
+        for image in images:
+            if not self._pull_status.get(image):
+                raise ValueError("Failed to pull all images")
+
     @staticmethod
     def _pull_image(client, image, status):
         try:

tests/integration/basic.py

@@ -28,4 +28,4 @@ def test_basic_chain(self):
 
     def test_no_such_image(self):
         with self.assertRaises(Exception):
-            Chainlink(stages_2)
+            Chainlink(stages_2).run({})