summaryrefslogtreecommitdiffstats
path: root/command.py
diff options
context:
space:
mode:
Diffstat (limited to 'command.py')
-rw-r--r--command.py10
1 files changed, 7 insertions, 3 deletions
diff --git a/command.py b/command.py
index 2a2ce138..22115ac2 100644
--- a/command.py
+++ b/command.py
@@ -268,8 +268,10 @@ class Command:
268 cls._parallel_context = None 268 cls._parallel_context = None
269 269
270 @classmethod 270 @classmethod
271 def _SetParallelContext(cls, context): 271 def _InitParallelWorker(cls, context, initializer):
272 cls._parallel_context = context 272 cls._parallel_context = context
273 if initializer:
274 initializer()
273 275
274 @classmethod 276 @classmethod
275 def ExecuteInParallel( 277 def ExecuteInParallel(
@@ -281,6 +283,7 @@ class Command:
281 output=None, 283 output=None,
282 ordered=False, 284 ordered=False,
283 chunksize=WORKER_BATCH_SIZE, 285 chunksize=WORKER_BATCH_SIZE,
286 initializer=None,
284 ): 287 ):
285 """Helper for managing parallel execution boiler plate. 288 """Helper for managing parallel execution boiler plate.
286 289
@@ -307,6 +310,7 @@ class Command:
307 ordered: Whether the jobs should be processed in order. 310 ordered: Whether the jobs should be processed in order.
308 chunksize: The number of jobs processed in batch by parallel 311 chunksize: The number of jobs processed in batch by parallel
309 workers. 312 workers.
313 initializer: Worker initializer.
310 314
311 Returns: 315 Returns:
312 The |callback| function's results are returned. 316 The |callback| function's results are returned.
@@ -318,8 +322,8 @@ class Command:
318 else: 322 else:
319 with multiprocessing.Pool( 323 with multiprocessing.Pool(
320 jobs, 324 jobs,
321 initializer=cls._SetParallelContext, 325 initializer=cls._InitParallelWorker,
322 initargs=(cls._parallel_context,), 326 initargs=(cls._parallel_context, initializer),
323 ) as pool: 327 ) as pool:
324 submit = pool.imap if ordered else pool.imap_unordered 328 submit = pool.imap if ordered else pool.imap_unordered
325 return callback( 329 return callback(