support async function

Author: Nanguage

executor/engine/__init__.py

@@ -1,7 +1,7 @@
 from .core import Engine, EngineSetting
 from .job import LocalJob, ThreadJob, ProcessJob
 
-__version__ = '0.2.7'
+__version__ = '0.2.8'
 
 __all__ = [
     'Engine', 'EngineSetting',

executor/engine/job/dask.py

@@ -1,9 +1,10 @@
 import functools
+from inspect import iscoroutinefunction
 
 from dask.distributed import Client, LocalCluster
 
 from .base import Job
-from .utils import create_generator_wrapper
+from .utils import create_generator_wrapper, run_async_func
 from ..utils import PortManager
 
 
@@ -58,6 +59,8 @@ async def run_function(self):
         """Run job with Dask."""
         client = self.engine.dask_client
         func = functools.partial(self.func, *self.args, **self.kwargs)
+        if iscoroutinefunction(func):
+            func = functools.partial(run_async_func, func)
         fut = client.submit(func)
         self._executor = fut
         result = await fut

executor/engine/job/local.py

@@ -1,11 +1,16 @@
+from inspect import iscoroutinefunction
+
 from .base import Job
 from .utils import create_generator_wrapper
 
 
 class LocalJob(Job):
     async def run_function(self):
         """Run job in local thread."""
-        res = self.func(*self.args, **self.kwargs)
+        if iscoroutinefunction(self.func):
+            res = await self.func(*self.args, **self.kwargs)
+        else:
+            res = self.func(*self.args, **self.kwargs)
         return res
 
     async def run_generator(self):

executor/engine/job/process.py

@@ -1,10 +1,13 @@
 import asyncio
 import functools
+from inspect import iscoroutinefunction
 
 from loky.process_executor import ProcessPoolExecutor
 
 from .base import Job
-from .utils import _gen_initializer, create_generator_wrapper
+from .utils import (
+    _gen_initializer, create_generator_wrapper, run_async_func
+)
 
 
 class ProcessJob(Job):
@@ -45,6 +48,8 @@ def release_resource(self) -> bool:
     async def run_function(self):
         """Run job in process pool."""
         func = functools.partial(self.func, *self.args, **self.kwargs)
+        if iscoroutinefunction(func):
+            func = functools.partial(run_async_func, func)
         self._executor = ProcessPoolExecutor(1)
         loop = asyncio.get_running_loop()
         fut = loop.run_in_executor(self._executor, func)

executor/engine/job/thread.py

@@ -1,9 +1,12 @@
 import asyncio
 import functools
+from inspect import iscoroutinefunction
 from concurrent.futures import ThreadPoolExecutor
 
 from .base import Job
-from .utils import _gen_initializer, create_generator_wrapper
+from .utils import (
+    _gen_initializer, create_generator_wrapper, run_async_func
+)
 
 
 class ThreadJob(Job):
@@ -44,6 +47,8 @@ def release_resource(self) -> bool:
     async def run_function(self):
         """Run job in thread pool."""
         func = functools.partial(self.func, *self.args, **self.kwargs)
+        if iscoroutinefunction(func):
+            func = functools.partial(run_async_func, func)
         self._executor = ThreadPoolExecutor(1)
         loop = asyncio.get_running_loop()
         fut = loop.run_in_executor(self._executor, func)

executor/engine/job/utils.py

@@ -169,3 +169,7 @@ def create_generator_wrapper(
         return AsyncGeneratorWrapper(job, fut)
     else:
         return SyncGeneratorWrapper(job, fut)
+
+
+def run_async_func(func, *args, **kwargs):
+    return asyncio.run(func(*args, **kwargs))

tests/test_dask_job.py

@@ -104,3 +104,24 @@ async def gen():
         assert x == i
         i += 1
     assert job.status == "done"
+
+
+@pytest.mark.asyncio
+async def test_dask_async_func():
+    port = PortManager.find_free_port()
+    cluster = LocalCluster(
+        dashboard_address=f":{port}",
+        asynchronous=True,
+        processes=False,
+    )
+    client = Client(cluster)
+    engine = Engine()
+    engine.dask_client = client
+
+    async def async_func(x):
+        return x + 1
+
+    job = DaskJob(async_func, (1,))
+    await engine.submit_async(job)
+    await job.wait_until_status("done")
+    assert job.result() == 2

tests/test_job.py

@@ -359,3 +359,16 @@ async def gen_async():
         assert await g.asend(2) == 3
         with pytest.raises(StopAsyncIteration):
             await g.asend(3)
+
+
+@pytest.mark.asyncio
+async def test_async_func_job():
+    with Engine() as engine:
+        async def async_func(x):
+            return x + 1
+
+        for job_cls in [LocalJob, ThreadJob, ProcessJob]:
+            job = job_cls(async_func, (1,))
+            await engine.submit_async(job)
+            await job.wait_until_status("done")
+            assert job.result() == 2