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( |
