Create mp_pool_run helper function for running in parallel.

cdist/config.py

@@ -28,6 +28,7 @@ import itertools
 import tempfile
 import socket
 import multiprocessing
+from cdist.mputil import mp_pool_run
 
 import cdist
 import cdist.hostsource
@@ -332,18 +333,10 @@ class Config(object):
                 multiprocessing.get_start_method()))
             self.log.debug(("Starting multiprocessing Pool for {} parallel "
                             "objects preparation".format(n)))
-            with multiprocessing.Pool(self.jobs) as pool:
-                self.log.debug(("Starting async for parallel object "
-                                "preparation"))
-                results = [
-                    pool.apply_async(self.object_prepare, (c,))
-                    for c in cargo
+            args = [
+                (c, ) for c in cargo
             ]
-
-                self.log.debug(("Waiting async results for parallel object "
-                                "preparation"))
-                for r in results:
-                    r.get()
+            mp_pool_run(self.object_prepare, args, jobs=self.jobs)
             self.log.debug(("Multiprocessing for parallel object "
                             "preparation finished"))
             objects_changed = True
@@ -376,17 +369,10 @@ class Config(object):
                 multiprocessing.get_start_method()))
             self.log.debug(("Starting multiprocessing Pool for {} parallel "
                             "object run".format(n)))
-            with multiprocessing.Pool(self.jobs) as pool:
-                self.log.debug(("Starting async for parallel object run"))
-                results = [
-                    pool.apply_async(self.object_run, (c,))
-                    for c in cargo
+            args = [
+                (c, ) for c in cargo
             ]
-
-                self.log.debug(("Waiting async results for parallel object "
-                                "run"))
-                for r in results:
-                    r.get()
+            mp_pool_run(self.object_run, args, jobs=self.jobs)
             self.log.debug(("Multiprocessing for parallel object "
                             "run finished"))
             objects_changed = True

cdist/core/explorer.py

@@ -24,8 +24,7 @@ import logging
 import os
 import glob
 import multiprocessing
-
-import cdist
+from cdist.mputil import mp_pool_run
 
 '''
 common:
@@ -121,16 +120,10 @@ class Explorer(object):
             multiprocessing.get_start_method()))
         self.log.debug(("Starting multiprocessing Pool for global "
                        "explorers run"))
-        with multiprocessing.Pool(self.jobs) as pool:
-            self.log.debug("Starting async for global explorer run")
-            results = [
-                pool.apply_async(self._run_global_explorer, (e, out_path,))
-                for e in self.list_global_explorer_names()
+        args = [
+            (e, out_path, ) for e in self.list_global_explorer_names()
         ]
-
-            self.log.debug("Waiting async results for global explorer runs")
-            for r in results:
-                r.get()  # self._run_global_explorer returns None
+        mp_pool_run(self._run_global_explorer, args, jobs=self.jobs)
         self.log.debug(("Multiprocessing run for global explorers "
                         "finished"))
 

cdist/exec/remote.py

@@ -31,6 +31,7 @@ import multiprocessing
 import cdist
 import cdist.exec.util as exec_util
 import cdist.util.ipaddr as ipaddr
+from cdist.mputil import mp_pool_run
 
 
 def _wrap_addr(addr):
@@ -152,23 +153,14 @@ class Remote(object):
             multiprocessing.get_start_method()))
         self.log.debug(("Starting multiprocessing Pool for parallel "
                         "remote transfer"))
-        with multiprocessing.Pool(jobs) as pool:
-            self.log.debug("Starting async for parallel transfer")
-            commands = []
+        args = []
         for f in glob.glob1(source, '*'):
             command = self._copy.split()
             path = os.path.join(source, f)
             command.extend([path, '{0}:{1}'.format(
                 _wrap_addr(self.target_host[0]), destination)])
-                commands.append(command)
-            results = [
-                pool.apply_async(self._run_command, (cmd,))
-                for cmd in commands
-            ]
-
-            self.log.debug("Waiting async results for parallel transfer")
-            for r in results:
-                r.get()  # self._run_command returns None
+            args.append((command, ))
+        mp_pool_run(self._run_command, args, jobs=jobs)
         self.log.debug(("Multiprocessing for parallel transfer "
                         "finished"))
 

cdist/mputil.py

@@ -0,0 +1,48 @@
+# -*- coding: utf-8 -*-
+#
+# 2016 Darko Poljak (darko.poljak at gmail.com)
+#
+# This file is part of cdist.
+#
+# cdist is free software: you can redistribute it and/or modify
+# it under the terms of the GNU General Public License as published by
+# the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# cdist is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License
+# along with cdist. If not, see <http://www.gnu.org/licenses/>.
+#
+#
+
+
+import multiprocessing
+import itertools
+
+
+def mp_pool_run(func, args=None, kwds=None, jobs=multiprocessing.cpu_count()):
+    """ Run func using multiprocessing.Pool with jobs jobs and supplied
+        iterable of args and kwds with one entry for each parallel func
+        instance.
+        Return list of results.
+    """
+    if args and kwds:
+        fargs = zip(args, kdws)
+    elif args:
+        fargs = zip(args, itertools.repeat({}))
+    elif kwds:
+        fargs = zip(itertools.repeat(()), kwds)
+    else:
+        return [func(), ]
+
+    with multiprocessing.Pool(jobs) as pool:
+        results = [
+            pool.apply_async(func, a, k)
+            for a, k in fargs
+        ]
+        retval = [r.get() for r in results]
+    return retval