diff options
Diffstat (limited to 'command.py')
-rw-r--r-- | command.py | 10 |
1 files changed, 7 insertions, 3 deletions
@@ -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( |