FEAT: make it easy to define generator methods in Actor (#82)

Co-authored-by: codingl2k1 <[email protected]>

python/xoscar/__init__.py

@@ -32,6 +32,7 @@
     setup_cluster,
     wait_actor_pool_recovered,
     get_pool_config,
+    generator,
 )
 from .backends import allocate_strategy
 from .backends.pool import MainActorPoolType

python/xoscar/api.py

@@ -15,9 +15,24 @@
 
 from __future__ import annotations
 
+import asyncio
+import inspect
+import threading
+import uuid
 from collections import defaultdict
 from numbers import Number
-from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Type, Union
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Dict,
+    Generic,
+    List,
+    Optional,
+    Tuple,
+    Type,
+    TypeVar,
+    Union,
+)
 from urllib.parse import urlparse
 
 from .aio import AioFileObject
@@ -271,6 +286,51 @@ def setup_cluster(address_to_resources: Dict[str, Dict[str, Number]]):
         get_backend(scheme).get_driver_cls().setup_cluster(address_resources)
 
 
+T = TypeVar("T")
+
+
+class IteratorWrapper(Generic[T]):
+    def __init__(self, uid: str, actor_addr: str, actor_uid: str):
+        self._uid = uid
+        self._actor_addr = actor_addr
+        self._actor_uid = actor_uid
+        self._actor_ref = None
+
+    async def destroy(self):
+        if self._actor_ref is None:
+            self._actor_ref = await actor_ref(
+                address=self._actor_addr, uid=self._actor_uid
+            )
+        assert self._actor_ref is not None
+        return await self._actor_ref.__xoscar_destroy_generator__(self._uid)
+
+    def __del__(self):
+        # It's not a good idea to spawn a new thread and join in __del__,
+        # but currently it's the only way to GC the generator.
+        thread = threading.Thread(
+            target=asyncio.run, args=(self.destroy(),), daemon=True
+        )
+        thread.start()
+        thread.join()
+
+    def __aiter__(self):
+        return self
+
+    async def __anext__(self) -> T:
+        if self._actor_ref is None:
+            self._actor_ref = await actor_ref(
+                address=self._actor_addr, uid=self._actor_uid
+            )
+        try:
+            assert self._actor_ref is not None
+            return await self._actor_ref.__xoscar_next__(self._uid)
+        except Exception as e:
+            if "StopIteration" in str(e):
+                raise StopAsyncIteration
+            else:
+                raise
+
+
 class AsyncActorMixin:
     @classmethod
     def default_uid(cls):
@@ -282,6 +342,10 @@ def __new__(cls, *args, **kwargs):
         except KeyError:
             return super().__new__(cls, *args, **kwargs)
 
+    def __init__(self, *args, **kwargs) -> None:
+        super().__init__()
+        self._generators: Dict[str, IteratorWrapper] = {}
+
     async def __post_create__(self):
         """
         Method called after actor creation
@@ -305,6 +369,85 @@ async def __on_receive__(self, message: Tuple[Any]):
         """
         return await super().__on_receive__(message)  # type: ignore
 
+    async def __xoscar_next__(self, generator_uid: str) -> Any:
+        """
+        Iter the next of generator.
+
+        Parameters
+        ----------
+        generator_uid: str
+            The uid of generator
+
+        Returns
+        -------
+            The next value of generator
+        """
+
+        def _wrapper(_gen):
+            try:
+                return next(_gen)
+            except StopIteration:
+                return stop
+
+        async def _async_wrapper(_gen):
+            try:
+                # anext is only available for Python >= 3.10
+                return await _gen.__anext__()  # noqa: F821
+            except StopAsyncIteration:
+                return stop
+
+        if gen := self._generators.get(generator_uid):
+            stop = object()
+            try:
+                if inspect.isgenerator(gen):
+                    # to_thread is only available for Python >= 3.9
+                    if hasattr(asyncio, "to_thread"):
+                        r = await asyncio.to_thread(_wrapper, gen)
+                    else:
+                        r = await asyncio.get_event_loop().run_in_executor(
+                            None, _wrapper, gen
+                        )
+                elif inspect.isasyncgen(gen):
+                    r = await asyncio.create_task(_async_wrapper(gen))
+                else:
+                    raise Exception(
+                        f"The generator {generator_uid} should be a generator or an async generator, "
+                        f"but a {type(gen)} is got."
+                    )
+            except Exception as e:
+                await self.__xoscar_destroy_generator__(generator_uid)
+                raise e
+            if r is stop:
+                await self.__xoscar_destroy_generator__(generator_uid)
+                raise Exception("StopIteration")
+            else:
+                return r
+        else:
+            raise RuntimeError(f"no iterator with id: {generator_uid}")
+
+    async def __xoscar_destroy_generator__(self, generator_uid: str):
+        """
+        Destroy the generator.
+
+        Parameters
+        ----------
+        generator_uid: str
+            The uid of generator
+        """
+        return self._generators.pop(generator_uid, None)
+
+
+def generator(func):
+    async def wrapper(obj, *args, **kwargs):
+        gen_uid = uuid.uuid1().hex
+        obj._generators[gen_uid] = func(obj, *args, **kwargs)
+        return IteratorWrapper(gen_uid, obj.address, obj.uid)
+
+    if inspect.isgeneratorfunction(func) or inspect.isasyncgenfunction(func):
+        return wrapper
+    else:
+        return func
+
 
 class Actor(AsyncActorMixin, _Actor):
     pass

python/xoscar/tests/test_generator.py

@@ -0,0 +1,140 @@
+# Copyright 2022-2023 XProbe Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import asyncio
+import time
+
+import pytest
+
+import xoscar as xo
+
+address = "127.0.0.1:12347"
+
+
+class WorkerActor(xo.StatelessActor):
+    @xo.generator
+    def chat(self):
+        for x in "hello oscar by sync":
+            yield x
+            time.sleep(0.1)
+
+    @xo.generator
+    async def achat(self):
+        for x in "hello oscar by async":
+            yield x
+            await asyncio.sleep(0.1)
+
+    @classmethod
+    def uid(cls):
+        return "worker"
+
+
+class SupervisorActor(xo.StatelessActor):
+    def get_all_generators(self):
+        return list(self._generators.keys())
+
+    @xo.generator
+    async def chat(self):
+        worker_actor: xo.ActorRef["WorkerActor"] = await xo.actor_ref(
+            address=address, uid=WorkerActor.uid()
+        )
+        yield "sync"
+        async for x in await worker_actor.chat():  # this is much confused. I will suggest use async generators only.
+            yield x
+
+        yield "async"
+        async for x in await worker_actor.achat():
+            yield x
+
+    @xo.generator
+    async def with_exception(self):
+        yield 1
+        raise Exception("intent raise")
+        yield 2
+
+    @classmethod
+    def uid(cls):
+        return "supervisor"
+
+
+async def test_generator():
+    await xo.create_actor_pool(address, 2)
+    await xo.create_actor(WorkerActor, address=address, uid=WorkerActor.uid())
+    superivsor_actor = await xo.create_actor(
+        SupervisorActor, address=address, uid=SupervisorActor.uid()
+    )
+
+    all_gen = await superivsor_actor.get_all_generators()
+    assert len(all_gen) == 0
+    output = []
+    async for x in await superivsor_actor.chat():
+        all_gen = await superivsor_actor.get_all_generators()
+        assert len(all_gen) == 1
+        output.append(x)
+    all_gen = await superivsor_actor.get_all_generators()
+    assert len(all_gen) == 0
+    assert output == [
+        "sync",
+        "h",
+        "e",
+        "l",
+        "l",
+        "o",
+        " ",
+        "o",
+        "s",
+        "c",
+        "a",
+        "r",
+        " ",
+        "b",
+        "y",
+        " ",
+        "s",
+        "y",
+        "n",
+        "c",
+        "async",
+        "h",
+        "e",
+        "l",
+        "l",
+        "o",
+        " ",
+        "o",
+        "s",
+        "c",
+        "a",
+        "r",
+        " ",
+        "b",
+        "y",
+        " ",
+        "a",
+        "s",
+        "y",
+        "n",
+        "c",
+    ]
+
+    with pytest.raises(Exception, match="intent"):
+        async for _ in await superivsor_actor.with_exception():
+            pass
+    all_gen = await superivsor_actor.get_all_generators()
+    assert len(all_gen) == 0
+
+    await asyncio.create_task(superivsor_actor.with_exception())
+    await asyncio.sleep(0)
+    all_gen = await superivsor_actor.get_all_generators()
+    assert len(all_gen) == 0