bpo-21423: Add an initializer argument to {Process,Thread}PoolExecuto… · python/cpython@63ff413 (original) (raw)
`@@ -131,6 +131,7 @@ def init(self, work_id, fn, args, kwargs):
`
131
131
`self.args = args
`
132
132
`self.kwargs = kwargs
`
133
133
``
``
134
+
134
135
`def _get_chunks(*iterables, chunksize):
`
135
136
`""" Iterates over zip()ed iterables in chunks. """
`
136
137
`it = zip(*iterables)
`
`@@ -151,7 +152,7 @@ def _process_chunk(fn, chunk):
`
151
152
` """
`
152
153
`return [fn(*args) for args in chunk]
`
153
154
``
154
``
`-
def _process_worker(call_queue, result_queue):
`
``
155
`+
def _process_worker(call_queue, result_queue, initializer, initargs):
`
155
156
`"""Evaluates calls from call_queue and places the results in result_queue.
`
156
157
``
157
158
` This worker is run in a separate process.
`
`@@ -161,7 +162,17 @@ def _process_worker(call_queue, result_queue):
`
161
162
` evaluated by the worker.
`
162
163
` result_queue: A ctx.Queue of _ResultItems that will written
`
163
164
` to by the worker.
`
``
165
`+
initializer: A callable initializer, or None
`
``
166
`+
initargs: A tuple of args for the initializer
`
164
167
` """
`
``
168
`+
if initializer is not None:
`
``
169
`+
try:
`
``
170
`+
initializer(*initargs)
`
``
171
`+
except BaseException:
`
``
172
`+
_base.LOGGER.critical('Exception in initializer:', exc_info=True)
`
``
173
`+
The parent will notice that the process stopped and
`
``
174
`+
mark the pool broken
`
``
175
`+
return
`
165
176
`while True:
`
166
177
`call_item = call_queue.get(block=True)
`
167
178
`if call_item is None:
`
`@@ -277,7 +288,9 @@ def shutdown_worker():
`
277
288
`# Mark the process pool broken so that submits fail right now.
`
278
289
`executor = executor_reference()
`
279
290
`if executor is not None:
`
280
``
`-
executor._broken = True
`
``
291
`+
executor._broken = ('A child process terminated '
`
``
292
`+
'abruptly, the process pool is not '
`
``
293
`+
'usable anymore')
`
281
294
`executor._shutdown_thread = True
`
282
295
`executor = None
`
283
296
`# All futures in flight must be marked failed
`
`@@ -372,15 +385,16 @@ def _chain_from_iterable_of_lists(iterable):
`
372
385
`yield element.pop()
`
373
386
``
374
387
``
375
``
`-
class BrokenProcessPool(RuntimeError):
`
``
388
`+
class BrokenProcessPool(_base.BrokenExecutor):
`
376
389
`"""
`
377
390
` Raised when a process in a ProcessPoolExecutor terminated abruptly
`
378
391
` while a future was in the running state.
`
379
392
` """
`
380
393
``
381
394
``
382
395
`class ProcessPoolExecutor(_base.Executor):
`
383
``
`-
def init(self, max_workers=None, mp_context=None):
`
``
396
`+
def init(self, max_workers=None, mp_context=None,
`
``
397
`+
initializer=None, initargs=()):
`
384
398
`"""Initializes a new ProcessPoolExecutor instance.
`
385
399
``
386
400
` Args:
`
`@@ -389,6 +403,8 @@ def init(self, max_workers=None, mp_context=None):
`
389
403
` worker processes will be created as the machine has processors.
`
390
404
` mp_context: A multiprocessing context to launch the workers. This
`
391
405
` object should provide SimpleQueue, Queue and Process.
`
``
406
`+
initializer: An callable used to initialize worker processes.
`
``
407
`+
initargs: A tuple of arguments to pass to the initializer.
`
392
408
` """
`
393
409
`_check_system_limits()
`
394
410
``
`@@ -403,6 +419,11 @@ def init(self, max_workers=None, mp_context=None):
`
403
419
`mp_context = mp.get_context()
`
404
420
`self._mp_context = mp_context
`
405
421
``
``
422
`+
if initializer is not None and not callable(initializer):
`
``
423
`+
raise TypeError("initializer must be a callable")
`
``
424
`+
self._initializer = initializer
`
``
425
`+
self._initargs = initargs
`
``
426
+
406
427
`# Make the call queue slightly larger than the number of processes to
`
407
428
`# prevent the worker processes from idling. But don't make it too big
`
408
429
`# because futures in the call queue cannot be cancelled.
`
`@@ -450,15 +471,16 @@ def _adjust_process_count(self):
`
450
471
`p = self._mp_context.Process(
`
451
472
`target=_process_worker,
`
452
473
`args=(self._call_queue,
`
453
``
`-
self._result_queue))
`
``
474
`+
self._result_queue,
`
``
475
`+
self._initializer,
`
``
476
`+
self._initargs))
`
454
477
`p.start()
`
455
478
`self._processes[p.pid] = p
`
456
479
``
457
480
`def submit(self, fn, *args, **kwargs):
`
458
481
`with self._shutdown_lock:
`
459
482
`if self._broken:
`
460
``
`-
raise BrokenProcessPool('A child process terminated '
`
461
``
`-
'abruptly, the process pool is not usable anymore')
`
``
483
`+
raise BrokenProcessPool(self._broken)
`
462
484
`if self._shutdown_thread:
`
463
485
`raise RuntimeError('cannot schedule new futures after shutdown')
`
464
486
``