Option to use https://github.com/joblib/loky instead of the default run.cpu_bound ProcessPoolExecutor

[ehfd]

https://github.com/joblib/loky allows the serialization of non-pickleable functions or objects and then offloading them by using cloudpickle for non-pickleable objects. This could be an optional dependency to activate.

Requires: loky and psutil.

Example:

import asyncio
import loky
from functools import partial
from typing import Callable, ParamSpec, TypeVar
from nicegui import app

P = ParamSpec('P')
R = TypeVar('R')
_loky_executor = loky.get_reusable_executor(timeout=None)

def shutdown_loky():
    _loky_executor.shutdown(wait=True, kill_workers=True)

async def run_cpu_bound(func: Callable[P, R], *args: P.args, **kwargs: P.kwargs) -> R:
    """
    Run a CPU-bound function using the loky backend directly.
    This supports non-pickleable functions (lambdas, local classes, etc.)
    via cloudpickle and allows awaiting the result.
    """
    loop = asyncio.get_running_loop()
    call_with_args = partial(func, *args, **kwargs)
    return await loop.run_in_executor(_loky_executor, call_with_args)

app.on_shutdown(shutdown_loky)

[evnchn]

Good idea!

Let's pick up #4708 ?

[falkoschindler]

Very interesting!
I'm thinking about #4353 where we noticed that the first run.cpu_bound call takes significantly longer:

import asyncio
import time
from nicegui import run, ui

def wait():
    time.sleep(1.0)

async def launch():
    start = time.perf_counter()
    await asyncio.gather(*[run.cpu_bound(wait) for _ in range(10)])
    print(f'All done in {time.perf_counter() - start:.2f} s')

@ui.page('/')
def page():
    ui.button('launch', on_click=launch)

if __name__ in {'__main__', '__mp_main__'}:
    print('running main')
    ui.run()

All done in 2.23 s
All done in 1.01 s
...

Maybe loky solves this problem as well?

[falkoschindler]

Yes, it seems to work:

import asyncio
import sys
import time
import cloudpickle
from loky import get_reusable_executor
from nicegui import ui

cloudpickle.register_pickle_by_value(sys.modules[__name__])
executor = get_reusable_executor(max_workers=10)

def do_work():
    time.sleep(1.0)

async def launch():
    start = time.perf_counter()
    await asyncio.gather(*[asyncio.wrap_future(executor.submit(do_work)) for _ in range(10)])
    print(f'All done in {time.perf_counter() - start:.2f} s')

ui.button('launch', on_click=launch)

ui.run()

All done in 1.07 s
All done in 1.08 s
...

• No need for the main guard. 🙂
• register_pickle_by_value is needed if do_work lives in the main module. 🫤
• While the first call is much faster than with run.cpu_bound, subsequent calls are a bit slower (1.01s vs. 1.08s). 🤔

[ehfd]

Also, try to disable the timeout; it will not need to re-initialize the workers every 10 seconds.

https://loky.readthedocs.io/en/stable/API.html

There's also loky.ProcessPoolExecutor, which can be slotted right in.

[ehfd]

Loky:

import asyncio
import sys
import time
import cloudpickle
from loky import get_reusable_executor
from nicegui import ui

cloudpickle.register_pickle_by_value(sys.modules[__name__])
executor = get_reusable_executor(max_workers=10, timeout=None)

def do_work():
    time.sleep(1.0)

async def launch():
    start = time.perf_counter()
    await asyncio.gather(*[asyncio.wrap_future(executor.submit(do_work)) for _ in range(10)])
    print(f'All done in {time.perf_counter() - start:.2f} s')

ui.button('launch', on_click=launch)

ui.run(port=8088)

NiceGUI ready to go on http://localhost:8088
All done in 1.11 s
All done in 1.00 s
All done in 1.00 s
All done in 1.00 s
All done in 1.00 s
All done in 1.00 s

ProcessPool:

import asyncio
import time
from nicegui import run, ui

def wait():
    time.sleep(1.0)

async def launch():
    start = time.perf_counter()
    await asyncio.gather(*[run.cpu_bound(wait) for _ in range(10)])
    print(f'All done in {time.perf_counter() - start:.2f} s')

@ui.page('/')
def page():
    ui.button('launch', on_click=launch)

if __name__ in {'__main__', '__mp_main__'}:
    #print('running main')
    ui.run(port=8088)

NiceGUI ready to go on http://localhost:8088
All done in 2.68 s
All done in 1.00 s
All done in 1.00 s
All done in 1.00 s
All done in 1.00 s