fixes AnswerDotAI#637

Author: jph00

fastcore/_modidx.py

@@ -437,11 +437,12 @@
                                                                                       'fastcore/parallel.py'),
                                    'fastcore.parallel.ThreadPoolExecutor.map': ( 'parallel.html#threadpoolexecutor.map',
                                                                                  'fastcore/parallel.py'),
+                                   'fastcore.parallel._add_one': ('parallel.html#_add_one', 'fastcore/parallel.py'),
                                    'fastcore.parallel._call': ('parallel.html#_call', 'fastcore/parallel.py'),
                                    'fastcore.parallel._done_pg': ('parallel.html#_done_pg', 'fastcore/parallel.py'),
                                    'fastcore.parallel._f_pg': ('parallel.html#_f_pg', 'fastcore/parallel.py'),
-                                   'fastcore.parallel.add_one': ('parallel.html#add_one', 'fastcore/parallel.py'),
                                    'fastcore.parallel.parallel': ('parallel.html#parallel', 'fastcore/parallel.py'),
+                                   'fastcore.parallel.parallel_async': ('parallel.html#parallel_async', 'fastcore/parallel.py'),
                                    'fastcore.parallel.parallel_gen': ('parallel.html#parallel_gen', 'fastcore/parallel.py'),
                                    'fastcore.parallel.parallelable': ('parallel.html#parallelable', 'fastcore/parallel.py'),
                                    'fastcore.parallel.run_procs': ('parallel.html#run_procs', 'fastcore/parallel.py'),

fastcore/parallel.py

@@ -4,7 +4,7 @@
 
 # %% auto 0
 __all__ = ['threaded', 'startthread', 'startproc', 'parallelable', 'ThreadPoolExecutor', 'ProcessPoolExecutor', 'parallel',
-           'add_one', 'run_procs', 'parallel_gen']
+           'parallel_async', 'run_procs', 'parallel_gen']
 
 # %% ../nbs/03a_parallel.ipynb
 from .imports import *
@@ -134,12 +134,30 @@ def parallel(f, items, *args, n_workers=defaults.cpus, total=None, progress=None
         return L(r)
 
 # %% ../nbs/03a_parallel.ipynb
-def add_one(x, a=1):
+def _add_one(x, a=1):
     # this import is necessary for multiprocessing in notebook on windows
     import random
     time.sleep(random.random()/80)
     return x+a
 
+# %% ../nbs/03a_parallel.ipynb
+async def parallel_async(f, items, *args, n_workers=16, total=None,
+                         timeout=None, chunksize=1, on_exc=print, **kwargs):
+    "Applies `f` to `items` in parallel using asyncio and a semaphore to limit concurrency."
+    import asyncio
+    if n_workers is None: n_workers = defaults.cpus
+    semaphore = asyncio.Semaphore(n_workers)
+    results = []
+
+    async def limited_task(item):
+        coro = f(item, *args, **kwargs) if asyncio.iscoroutinefunction(f) else asyncio.to_thread(f, item, *args, **kwargs)
+        async with semaphore:
+            return await asyncio.wait_for(coro, timeout) if timeout else await coro
+
+    tasks = [limited_task(item) for item in items]
+    if total is None: total = len(items)
+    return asyncio.gather(*tasks)
+
 # %% ../nbs/03a_parallel.ipynb
 def run_procs(f, f_done, args):
     "Call `f` for each item in `args` in parallel, yielding `f_done`"

nbs/03a_parallel.ipynb

@@ -363,7 +363,7 @@
        "\n",
        ">      ProcessPoolExecutor (max_workers=8, on_exc=<built-in function print>,\n",
        ">                           pause=0, mp_context=None, initializer=None,\n",
-       ">                           initargs=())\n",
+       ">                           initargs=(), max_tasks_per_child=None)\n",
        "\n",
        "*Same as Python's ProcessPoolExecutor, except can pass `max_workers==0` for serial execution*"
       ],
@@ -376,7 +376,7 @@
        "\n",
        ">      ProcessPoolExecutor (max_workers=8, on_exc=<built-in function print>,\n",
        ">                           pause=0, mp_context=None, initializer=None,\n",
-       ">                           initargs=())\n",
+       ">                           initargs=(), max_tasks_per_child=None)\n",
        "\n",
        "*Same as Python's ProcessPoolExecutor, except can pass `max_workers==0` for serial execution*"
       ]
@@ -432,7 +432,7 @@
    "outputs": [],
    "source": [
     "#|export\n",
-    "def add_one(x, a=1):\n",
+    "def _add_one(x, a=1):\n",
     "    # this import is necessary for multiprocessing in notebook on windows\n",
     "    import random\n",
     "    time.sleep(random.random()/80)\n",
@@ -447,11 +447,11 @@
    "source": [
     "inp,exp = range(50),range(1,51)\n",
     "\n",
-    "test_eq(parallel(add_one, inp, n_workers=2), exp)\n",
-    "test_eq(parallel(add_one, inp, threadpool=True, n_workers=2), exp)\n",
-    "test_eq(parallel(add_one, inp, n_workers=1, a=2), range(2,52))\n",
-    "test_eq(parallel(add_one, inp, n_workers=0), exp)\n",
-    "test_eq(parallel(add_one, inp, n_workers=0, a=2), range(2,52))"
+    "test_eq(parallel(_add_one, inp, n_workers=2), exp)\n",
+    "test_eq(parallel(_add_one, inp, threadpool=True, n_workers=2), exp)\n",
+    "test_eq(parallel(_add_one, inp, n_workers=1, a=2), range(2,52))\n",
+    "test_eq(parallel(_add_one, inp, n_workers=0), exp)\n",
+    "test_eq(parallel(_add_one, inp, n_workers=0, a=2), range(2,52))"
    ]
   },
   {
@@ -479,11 +479,11 @@
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "0 2024-10-09 16:08:39.462154\n",
-      "1 2024-10-09 16:08:39.715074\n",
-      "2 2024-10-09 16:08:39.969191\n",
-      "3 2024-10-09 16:08:40.221442\n",
-      "4 2024-10-09 16:08:40.473224\n"
+      "0 2024-10-12 13:30:21.217649\n",
+      "1 2024-10-12 13:30:21.469191\n",
+      "2 2024-10-12 13:30:21.721034\n",
+      "3 2024-10-12 13:30:21.972793\n",
+      "4 2024-10-12 13:30:22.223159\n"
      ]
     }
    ],
@@ -520,6 +520,54 @@
     "parallel(die_sometimes, range(8))"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "#|export\n",
+    "async def parallel_async(f, items, *args, n_workers=16, total=None,\n",
+    "                         timeout=None, chunksize=1, on_exc=print, **kwargs):\n",
+    "    \"Applies `f` to `items` in parallel using asyncio and a semaphore to limit concurrency.\"\n",
+    "    import asyncio\n",
+    "    if n_workers is None: n_workers = defaults.cpus\n",
+    "    semaphore = asyncio.Semaphore(n_workers)\n",
+    "    results = []\n",
+    "\n",
+    "    async def limited_task(item):\n",
+    "        coro = f(item, *args, **kwargs) if asyncio.iscoroutinefunction(f) else asyncio.to_thread(f, item, *args, **kwargs)\n",
+    "        async with semaphore:\n",
+    "            return await asyncio.wait_for(coro, timeout) if timeout else await coro\n",
+    "\n",
+    "    tasks = [limited_task(item) for item in items]\n",
+    "    if total is None: total = len(items)\n",
+    "    return asyncio.gather(*tasks)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import asyncio"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "async def print_time_async(i): \n",
+    "    wait = random.random()\n",
+    "    await asyncio.sleep(wait)\n",
+    "    print(i, datetime.now(), wait)\n",
+    "\n",
+    "await parallel_async(print_time_async, range(6), n_workers=3);"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,