Skip to content

Commit bda89da

Browse files
authored
Merge pull request #8195 from stuartarchibald/fix/7872_056
Make the workqueue threading backend once again fork safe.
2 parents 605fd4a + b07d687 commit bda89da

2 files changed

Lines changed: 70 additions & 4 deletions

File tree

‎numba/np/ufunc/workqueue.c‎

Lines changed: 29 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -157,7 +157,7 @@ platform_launch(void)
157157
pthread_key_create(&tidkey, NULL);
158158
}
159159

160-
#endif
160+
#endif /* pthread threading */
161161

162162
/* Win Thread */
163163
#ifdef NUMBA_WINTHREAD
@@ -262,7 +262,7 @@ platform_launch(void)
262262
tidkey = TlsAlloc();
263263
}
264264

265-
#endif
265+
#endif /* Windows threading */
266266

267267
typedef struct Task
268268
{
@@ -321,12 +321,20 @@ static THREAD_LOCAL(int) _TLS_num_threads = 0;
321321
static void
322322
set_num_threads(int count)
323323
{
324+
if (!queues)
325+
{
326+
launch_threads(NUM_THREADS);
327+
}
324328
_TLS_num_threads = count;
325329
}
326330

327331
static int
328332
get_num_threads(void)
329333
{
334+
if (!queues)
335+
{
336+
launch_threads(NUM_THREADS);
337+
}
330338
// This is purely to permit the implementation to survive to the point
331339
// where it can exit cleanly as multiple threads cannot be used with this
332340
// backend
@@ -506,6 +514,11 @@ add_task_internal(void *fn, void *args, void *dims, void *steps, void *data, int
506514
{
507515
void (*func)(void *args, void *dims, void *steps, void *data) = fn;
508516

517+
if (!queues)
518+
{
519+
launch_threads(NUM_THREADS);
520+
}
521+
509522
Queue *queue = &queues[queue_pivot];
510523

511524
Task *task = &queue->task;
@@ -581,6 +594,11 @@ static void launch_threads(int count)
581594

582595
static void synchronize(void)
583596
{
597+
// This probably isn't needed, but is put in as a guard
598+
if (!queues)
599+
{
600+
launch_threads(NUM_THREADS);
601+
}
584602
int i;
585603
for (i = 0; i < queue_count; ++i)
586604
{
@@ -590,6 +608,11 @@ static void synchronize(void)
590608

591609
static void ready(void)
592610
{
611+
// This probably isn't needed, but is put in as a guard
612+
if (!queues)
613+
{
614+
launch_threads(NUM_THREADS);
615+
}
593616
int i;
594617
for (i = 0; i < queue_count; ++i)
595618
{
@@ -601,8 +624,10 @@ static void reset_after_fork(void)
601624
{
602625
free(queues);
603626
queues = NULL;
604-
NUM_THREADS = -1;
605-
_INIT_NUM_THREADS = -1;
627+
if (_INIT_NUM_THREADS != -1)
628+
{
629+
NUM_THREADS = _INIT_NUM_THREADS;
630+
}
606631
_nesting_level = 0;
607632
}
608633

‎numba/tests/test_parallel_backend.py‎

Lines changed: 41 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -690,6 +690,47 @@ def main():
690690
self.assertIn("Terminating: Nested parallel kernel launch detected",
691691
e_msg)
692692

693+
@unittest.skipUnless(_HAVE_OS_FORK, "Test needs fork(2)")
694+
def test_workqueue_handles_fork_from_non_main_thread(self):
695+
# For context see #7872, but essentially the multiprocessing pool
696+
# implementation has a number of Python threads for handling the worker
697+
# processes, one of which calls fork(2), this results in a fork from a
698+
# non-main thread.
699+
700+
runme = """if 1:
701+
from numba import njit, prange, threading_layer
702+
import numpy as np
703+
import multiprocessing
704+
705+
if __name__ == "__main__":
706+
# Need for force fork context (OSX default is "spawn")
707+
multiprocessing.set_start_method('fork')
708+
709+
@njit(parallel=True)
710+
def func(x):
711+
return 10. * x
712+
713+
arr = np.arange(2.)
714+
715+
# run in single process to start Numba's thread pool
716+
np.testing.assert_allclose(func(arr), func.py_func(arr))
717+
718+
# now run in a multiprocessing pool to get a fork from a
719+
# non-main thread
720+
with multiprocessing.Pool(10) as p:
721+
result = p.map(func, [arr])
722+
np.testing.assert_allclose(result,
723+
func.py_func(np.expand_dims(arr, 0)))
724+
725+
assert threading_layer() == "workqueue"
726+
"""
727+
cmdline = [sys.executable, '-c', runme]
728+
env = os.environ.copy()
729+
env['NUMBA_THREADING_LAYER'] = "workqueue"
730+
env['NUMBA_NUM_THREADS'] = "4"
731+
732+
self.run_cmd(cmdline, env=env)
733+
693734

694735
# 32bit or windows py27 (not that this runs on windows)
695736
@skip_parfors_unsupported

0 commit comments

Comments
 (0)