address multiprocessing spawning errors when on macOS and nbdev runs tests

Author: s01st

fastcore/parallel.py

@@ -30,7 +30,11 @@ def g(_obj_td, *args, **kwargs):
             _obj_td.result = res
         @wraps(f)
         def _f(*args, **kwargs):
-            res = (Thread,Process)[process](target=g, args=args, kwargs=kwargs)
+            if process:
+                Proc = get_context('fork').Process if sys.platform == 'darwin' else Process
+            else:
+                Proc = Thread
+            res = Proc(target=g, args=args, kwargs=kwargs)
             res._args = (res,)+res._args
             res.start()
             return res
@@ -123,7 +127,9 @@ def parallel(f, items, *args, n_workers=defaults.cpus, total=None, progress=None
     kwpool = {}
     if threadpool: pool = ThreadPoolExecutor
     else:
-        if not method and sys.platform == 'darwin': method='fork'
+        if not method and sys.platform == 'darwin':
+            # Use fork only if function is defined in __main__ (notebooks/REPL), otherwise use spawn
+            method = 'fork' if getattr(f, '__module__', None) == '__main__' else 'spawn'
         if method: kwpool['mp_context'] = get_context(method)
         pool = ProcessPoolExecutor
     with pool(n_workers, pause=pause, **kwpool) as ex:
@@ -158,7 +164,8 @@ async def limited_task(item):
 # %% ../nbs/03a_parallel.ipynb
 def run_procs(f, f_done, args):
     "Call `f` for each item in `args` in parallel, yielding `f_done`"
-    processes = L(args).map(Process, args=arg0, target=f)
+    Proc = get_context('fork').Process if sys.platform == 'darwin' else Process
+    processes = L(args).map(Proc, args=arg0, target=f)
     for o in processes: o.start()
     yield from f_done()
     processes.map(Self.join())

nbs/03a_parallel.ipynb

@@ -14,22 +14,7 @@
    "execution_count": null,
    "metadata": {},
    "outputs": [],
-   "source": [
-    "#| export\n",
-    "from fastcore.imports import *\n",
-    "from fastcore.basics import *\n",
-    "from fastcore.foundation import *\n",
-    "from fastcore.meta import *\n",
-    "from fastcore.xtras import *\n",
-    "from functools import wraps\n",
-    "\n",
-    "import concurrent.futures,time\n",
-    "from multiprocessing import Process,Queue,Manager,set_start_method,get_all_start_methods,get_context\n",
-    "from threading import Thread\n",
-    "try:\n",
-    "    if sys.platform == 'darwin' and IN_NOTEBOOK: set_start_method(\"fork\")\n",
-    "except: pass"
-   ]
+   "source": "#| export\nfrom fastcore.imports import *\nfrom fastcore.basics import *\nfrom fastcore.foundation import *\nfrom fastcore.meta import *\nfrom fastcore.xtras import *\nfrom functools import wraps\n\nimport concurrent.futures,time\nfrom multiprocessing import Process,Queue,Manager,set_start_method,get_all_start_methods,get_context\nfrom threading import Thread\ntry:\n    if sys.platform == 'darwin' and IN_NOTEBOOK: set_start_method(\"fork\")\nexcept: pass"
   },
   {
    "cell_type": "code",
@@ -56,27 +41,7 @@
    "execution_count": null,
    "metadata": {},
    "outputs": [],
-   "source": [
-    "#| export\n",
-    "def threaded(process=False):\n",
-    "    \"Run `f` in a `Thread` (or `Process` if `process=True`), and returns it\"\n",
-    "    def _r(f):\n",
-    "        def g(_obj_td, *args, **kwargs):\n",
-    "            res = f(*args, **kwargs)\n",
-    "            _obj_td.result = res\n",
-    "        @wraps(f)\n",
-    "        def _f(*args, **kwargs):\n",
-    "            res = (Thread,Process)[process](target=g, args=args, kwargs=kwargs)\n",
-    "            res._args = (res,)+res._args\n",
-    "            res.start()\n",
-    "            return res\n",
-    "        return _f\n",
-    "    if callable(process):\n",
-    "        o = process\n",
-    "        process = False\n",
-    "        return _r(o)\n",
-    "    return _r"
-   ]
+   "source": "#| export\ndef threaded(process=False):\n    \"Run `f` in a `Thread` (or `Process` if `process=True`), and returns it\"\n    def _r(f):\n        def g(_obj_td, *args, **kwargs):\n            res = f(*args, **kwargs)\n            _obj_td.result = res\n        @wraps(f)\n        def _f(*args, **kwargs):\n            if process:\n                Proc = get_context('fork').Process if sys.platform == 'darwin' else Process\n            else:\n                Proc = Thread\n            res = Proc(target=g, args=args, kwargs=kwargs)\n            res._args = (res,)+res._args\n            res.start()\n            return res\n        return _f\n    if callable(process):\n        o = process\n        process = False\n        return _r(o)\n    return _r"
   },
   {
    "cell_type": "code",
@@ -406,24 +371,7 @@
    "execution_count": null,
    "metadata": {},
    "outputs": [],
-   "source": [
-    "#| export\n",
-    "def parallel(f, items, *args, n_workers=defaults.cpus, total=None, progress=None, pause=0,\n",
-    "             method=None, threadpool=False, timeout=None, chunksize=1, **kwargs):\n",
-    "    \"Applies `func` in parallel to `items`, using `n_workers`\"\n",
-    "    kwpool = {}\n",
-    "    if threadpool: pool = ThreadPoolExecutor\n",
-    "    else:\n",
-    "        if not method and sys.platform == 'darwin': method='fork'\n",
-    "        if method: kwpool['mp_context'] = get_context(method)\n",
-    "        pool = ProcessPoolExecutor\n",
-    "    with pool(n_workers, pause=pause, **kwpool) as ex:\n",
-    "        r = ex.map(f,items, *args, timeout=timeout, chunksize=chunksize, **kwargs)\n",
-    "        if progress and progress_bar:\n",
-    "            if total is None: total = len(items)\n",
-    "            r = progress_bar(r, total=total, leave=False)\n",
-    "        return L(r)"
-   ]
+   "source": "#| export\ndef parallel(f, items, *args, n_workers=defaults.cpus, total=None, progress=None, pause=0,\n             method=None, threadpool=False, timeout=None, chunksize=1, **kwargs):\n    \"Applies `func` in parallel to `items`, using `n_workers`\"\n    kwpool = {}\n    if threadpool: pool = ThreadPoolExecutor\n    else:\n        if not method and sys.platform == 'darwin':\n            # Use fork only if function is defined in __main__ (notebooks/REPL), otherwise use spawn\n            method = 'fork' if getattr(f, '__module__', None) == '__main__' else 'spawn'\n        if method: kwpool['mp_context'] = get_context(method)\n        pool = ProcessPoolExecutor\n    with pool(n_workers, pause=pause, **kwpool) as ex:\n        r = ex.map(f,items, *args, timeout=timeout, chunksize=chunksize, **kwargs)\n        if progress and progress_bar:\n            if total is None: total = len(items)\n            r = progress_bar(r, total=total, leave=False)\n        return L(r)"
   },
   {
    "cell_type": "code",
@@ -583,15 +531,7 @@
    "execution_count": null,
    "metadata": {},
    "outputs": [],
-   "source": [
-    "#| export\n",
-    "def run_procs(f, f_done, args):\n",
-    "    \"Call `f` for each item in `args` in parallel, yielding `f_done`\"\n",
-    "    processes = L(args).map(Process, args=arg0, target=f)\n",
-    "    for o in processes: o.start()\n",
-    "    yield from f_done()\n",
-    "    processes.map(Self.join())"
-   ]
+   "source": "#| export\ndef run_procs(f, f_done, args):\n    \"Call `f` for each item in `args` in parallel, yielding `f_done`\"\n    Proc = get_context('fork').Process if sys.platform == 'darwin' else Process\n    processes = L(args).map(Proc, args=arg0, target=f)\n    for o in processes: o.start()\n    yield from f_done()\n    processes.map(Self.join())"
   },
   {
    "cell_type": "code",