parallel.py

"""Threading and multiprocessing functions

`parallel(f, items)` maps `f` over `items` using a process pool (`threadpool=True` for threads instead), with `n_workers`, optional `progress` bar, a `pause` between starts (to stagger e.g. web requests), and `return_exceptions` to collect errors instead of raising; `n_workers=0` runs serially, which makes debugging easy. `parallel_async` is the asyncio version, using a semaphore to limit concurrency. `@threaded` makes a function run in a `Thread` (or `Process`) whose eventual return value lands in its `result` attr, and `startthread`/`startproc` start one immediately.

Docs: https://fastcore.fast.ai/parallel.html.md"""

# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/03a_parallel.ipynb.

# %% auto #0
__all__ = ['threaded', 'startthread', 'startproc', 'parallelable', 'ThreadPoolExecutor', 'ProcessPoolExecutor', 'parallel',
           'parallel_async_gen', 'parallel_async_dict', 'parallel_async', 'bg_task']

# %% ../nbs/03a_parallel.ipynb #569d18c6
from .imports import *
from .basics import *
from .foundation import *
from .meta import *
from .xtras import *
from functools import wraps

import concurrent.futures,time
from multiprocessing import Process,Manager,set_start_method,get_context
from threading import Thread,Lock
try:
    if sys.platform == 'darwin' and IN_NOTEBOOK: set_start_method("fork")
except: pass

# %% ../nbs/03a_parallel.ipynb #21e71104
@metadec
def threaded(
    f, # Function to run
    *,
    process=False, # Create a Process instead of a Thread?
    daemon=False   # Use daemon mode?
): # Wrapped `f`, returning on call the Process or Thread created, which will have `result` attr injected in once complete
    "Run `f` in a `Thread` (or `Process` if `process=True`), and returns it"
    def g(_obj_td, *args, **kwargs):
        res = f(*args, **kwargs)
        _obj_td.result = res
    @wraps(f)
    def _f(*args, **kwargs):
        Proc = get_context('fork').Process if sys.platform == 'darwin' else Process
        res = (Thread,Proc)[process](target=g, args=args, kwargs=kwargs)
        res._args = (res,)+res._args
        if daemon: res.daemon = True
        res.start()
        return res
    return _f

# %% ../nbs/03a_parallel.ipynb #dbccb558
@metadec
def startthread(f, *, daemon=False):
    "Like `threaded`, but start thread immediately"
    return threaded(f, daemon=daemon)()


# %% ../nbs/03a_parallel.ipynb #a32d66c4
@metadec
def startproc(f, *, daemon=False):
    "Like `threaded(process=True)`, but start Process immediately"
    return threaded(f, process=True, daemon=daemon)()


# %% ../nbs/03a_parallel.ipynb #44d4651b
def _call(lock, pause, n, g, item, return_exceptions=False):
    l = False
    if pause:
        try:
            l = lock.acquire(timeout=pause*(n+2))
            time.sleep(pause)
        finally:
            if l: lock.release()
    try: return g(item)
    except Exception as e:
        if return_exceptions: return e
        raise

# %% ../nbs/03a_parallel.ipynb #63a3920f
def parallelable(param_name, num_workers, f=None):
    f_in_main = f == None or sys.modules[f.__module__].__name__ == "__main__"
    if sys.platform == "win32" and IN_NOTEBOOK and num_workers > 0 and f_in_main:
        print("Due to IPython and Windows limitation, python multiprocessing isn't available now.")
        print(f"So `{param_name}` has to be changed to 0 to avoid getting stuck")
        return False
    return True

# %% ../nbs/03a_parallel.ipynb #e74836fc
class ThreadPoolExecutor(concurrent.futures.ThreadPoolExecutor):
    "Same as Python's ThreadPoolExecutor, except can pass `max_workers==0` for serial execution"
    def __init__(self, max_workers=defaults.cpus, on_exc=print, pause=0, **kwargs):
        if max_workers is None: max_workers=defaults.cpus
        store_attr()
        self.not_parallel = max_workers==0
        if self.not_parallel: max_workers=1
        super().__init__(max_workers, **kwargs)

    def map(self, f, items, *args, timeout=None, chunksize=1, return_exceptions=False, **kwargs):
        self.lock = Lock() if self.not_parallel==False and self.pause else None
        g = partial(f, *args, **kwargs)
        if self.not_parallel: return map(g, items)
        _g = partial(_call, self.lock, self.pause, self.max_workers, g, return_exceptions=return_exceptions)
        try: return super().map(_g, items, timeout=timeout, chunksize=chunksize)
        except Exception as e: self.on_exc(e)

# %% ../nbs/03a_parallel.ipynb #0046c557
@delegates()
class ProcessPoolExecutor(concurrent.futures.ProcessPoolExecutor):
    "Same as Python's ProcessPoolExecutor, except can pass `max_workers==0` for serial execution"
    def __init__(self, max_workers=defaults.cpus, on_exc=print, pause=0, **kwargs):
        if max_workers is None: max_workers=defaults.cpus
        store_attr()
        self.not_parallel = max_workers==0
        if self.not_parallel: max_workers=1
        super().__init__(max_workers, **kwargs)

    def map(self, f, items, *args, timeout=None, chunksize=1, return_exceptions=False, **kwargs):
        if not parallelable('max_workers', self.max_workers, f): self.max_workers = 0
        self.not_parallel = self.max_workers==0
        if self.not_parallel: self.max_workers=1

        self.lock = Manager().Lock() if self.not_parallel==False and self.pause else None
        g = partial(f, *args, **kwargs)
        if self.not_parallel: return map(g, items)
        _g = partial(_call, self.lock, self.pause, self.max_workers, g, return_exceptions=return_exceptions)
        try: return super().map(_g, items, timeout=timeout, chunksize=chunksize)
        except Exception as e: self.on_exc(e)

# %% ../nbs/03a_parallel.ipynb #529e1bb1
def parallel(f, items, *args, n_workers=defaults.cpus, total=None, progress=None, pause=0,
    method=None, threadpool=False, timeout=None, chunksize=1, return_exceptions=False, **kwargs):
    "Applies `func` in parallel to `items`, using `n_workers`"
    kwpool = {}
    if threadpool: pool = ThreadPoolExecutor
    else:
        if not method and sys.platform == 'darwin': method='fork'
        if method: kwpool['mp_context'] = get_context(method)
        pool = ProcessPoolExecutor
    with pool(n_workers, pause=pause, **kwpool) as ex:
        r = ex.map(f,items, *args, timeout=timeout, chunksize=chunksize, return_exceptions=return_exceptions, **kwargs)
        if progress:
            from fastprogress import progress_bar
            if total is None: total = len(items)
            r = progress_bar(r, total=total, leave=False)
        return L(r)

# %% ../nbs/03a_parallel.ipynb #343d0191
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 #9a6e2f26
async def parallel_async_gen(f, items, *args, n_workers=16, pause=0,
    timeout=None, return_exceptions=False, cancel_on_exit=True, **kwargs):
    "Yield `(index,result)` pairs as `f` applied to each of `items` completes, in completion order"
    import asyncio
    semaphore = asyncio.Semaphore(n_workers)
    async def limited_task(i, item):
        if pause: await asyncio.sleep(i * pause)
        async with semaphore:
            coro = f(item, *args, **kwargs) if asyncio.iscoroutinefunction(f) else asyncio.to_thread(f, item, *args, **kwargs)
            try: return i, (await asyncio.wait_for(coro, timeout) if timeout else await coro)
            except Exception as e:
                if return_exceptions: return i, e
                raise
    tasks = [asyncio.ensure_future(limited_task(i, o)) for i,o in enumerate(items)]
    try:
        for t in asyncio.as_completed(tasks): yield await t
    finally:
        if cancel_on_exit:
            for t in tasks: t.cancel()
            await asyncio.gather(*tasks, return_exceptions=True)

# %% ../nbs/03a_parallel.ipynb #7406c085
@delegates(parallel_async_gen)
async def parallel_async_dict(f, items, *args, **kwargs):
    "Apply `f` to `items` in parallel, returning `{index: result}` in completion order"
    return {i:r async for i,r in parallel_async_gen(f, items, *args, **kwargs)}

# %% ../nbs/03a_parallel.ipynb #87a80e04
@delegates(parallel_async_dict, but=['cancel_on_exit'])
async def parallel_async(f, items, *args, cancel_on_error=False, **kwargs):
    "Applies `f` to `items` in parallel using asyncio and a semaphore to limit concurrency."
    res = await parallel_async_dict(f, items, *args, cancel_on_exit=cancel_on_error, **kwargs)
    return L(res[i] for i in range(len(res)))

# %% ../nbs/03a_parallel.ipynb #6748aa27
_bg_tasks = set()

def bg_task(
    coro, # Coroutine to schedule
    on_err=None, # Called with the exception when the task fails; default prints the traceback
):
    "Like `asyncio.create_task`, but keeps the task alive and reports exceptions, for fire-and-forget tasks"
    import traceback,asyncio
    def _done(t):
        if not t.cancelled() and (exc := t.exception()):
            if on_err: on_err(exc)
            else: traceback.print_exception(exc)
    task = asyncio.create_task(coro)
    _bg_tasks.add(task)
    task.add_done_callback(_bg_tasks.discard)
    task.add_done_callback(_done)
    return task