xoscar 0.4.1__cp39-cp39-win_amd64.whl → 0.4.3__cp39-cp39-win_amd64.whl

xoscar/backends/communication/dummy.py

@@ -17,8 +17,9 @@ from __future__ import annotations
 
 import asyncio
 import concurrent.futures as futures
+import logging
 import weakref
-from typing import Any, Callable, Coroutine, Dict, Type
+from typing import Any, Callable, Coroutine, Dict, Optional, Type
 from urllib.parse import urlparse
 
 from ...errors import ServerClosed
@@ -29,13 +30,15 @@ from .errors import ChannelClosed
 
 DEFAULT_DUMMY_ADDRESS = "dummy://0"
 
+logger = logging.getLogger(__name__)
+
 
 class DummyChannel(Channel):
     """
     Channel for communications in same process.
     """
 
-    __slots__ = "_in_queue", "_out_queue", "_closed"
+    __slots__ = "__weakref__", "_in_queue", "_out_queue", "_closed"
 
     name = "dummy"
 
@@ -100,8 +103,8 @@ class DummyServer(Server):
     _address_to_instances: weakref.WeakValueDictionary[str, "DummyServer"] = (
         weakref.WeakValueDictionary()
     )
-    _channels: list[ChannelType]
-    _tasks: list[asyncio.Task]
+    _channels: weakref.WeakSet[Channel]
+    _tasks: set[asyncio.Task]
     scheme: str | None = "dummy"
 
     def __init__(
@@ -111,8 +114,8 @@ class DummyServer(Server):
     ):
         super().__init__(address, channel_handler)
         self._closed = asyncio.Event()
-        self._channels = []
-        self._tasks = []
+        self._channels = weakref.WeakSet()
+        self._tasks = set()
 
     @classmethod
     def get_instance(cls, address: str):
@@ -178,7 +181,7 @@ class DummyServer(Server):
                 f"{type(self).__name__} got unexpected "
                 f'arguments: {",".join(kwargs)}'
             )
-        self._channels.append(channel)
+        self._channels.add(channel)
         await self.channel_handler(channel)
 
     @implements(Server.stop)
@@ -203,6 +206,7 @@ class DummyClient(Client):
         self, local_address: str | None, dest_address: str | None, channel: Channel
     ):
         super().__init__(local_address, dest_address, channel)
+        self._task: Optional[asyncio.Task] = None
 
     @staticmethod
     @implements(Client.connect)
@@ -232,11 +236,25 @@ class DummyClient(Client):
         task = asyncio.create_task(conn_coro)
         client = DummyClient(local_address, dest_address, client_channel)
         client._task = task
-        server._tasks.append(task)
+        server._tasks.add(task)
+
+        def _discard(t):
+            server._tasks.discard(t)
+            logger.info("Channel exit: %s", server_channel.info)
+
+        task.add_done_callback(_discard)
         return client
 
     @implements(Client.close)
     async def close(self):
         await super().close()
-        self._task.cancel()
-        self._task = None
+        if self._task is not None:
+            task_loop = self._task.get_loop()
+            if task_loop is not None:
+                if not task_loop.is_running():
+                    logger.warning(
+                        "Dummy channel cancel task on a stopped loop, dest address: %s.",
+                        self.dest_address,
+                    )
+            self._task.cancel()
+            self._task = None

xoscar/backends/communication/socket.py

@@ -113,7 +113,7 @@ class SocketChannel(Channel):
 class _BaseSocketServer(Server, metaclass=ABCMeta):
     __slots__ = "_aio_server", "_channels"
 
-    _channels: list[ChannelType]
+    _channels: set[Channel]
 
     def __init__(
         self,
@@ -124,7 +124,7 @@ class _BaseSocketServer(Server, metaclass=ABCMeta):
         super().__init__(address, channel_handler)
         # asyncio.Server
         self._aio_server = aio_server
-        self._channels = []
+        self._channels = set()
 
     @implements(Server.start)
     async def start(self):
@@ -170,9 +170,16 @@ class _BaseSocketServer(Server, metaclass=ABCMeta):
             dest_address=dest_address,
             channel_type=self.channel_type,
         )
-        self._channels.append(channel)
+        self._channels.add(channel)
         # handle over channel to some handlers
-        await self.channel_handler(channel)
+        try:
+            await self.channel_handler(channel)
+        finally:
+            if not channel.closed:
+                await channel.close()
+            # Remove channel if channel exit
+            self._channels.discard(channel)
+            logger.debug("Channel exit: %s", channel.info)
 
     @implements(Server.stop)
     async def stop(self):
@@ -185,6 +192,7 @@ class _BaseSocketServer(Server, metaclass=ABCMeta):
         await asyncio.gather(
             *(channel.close() for channel in self._channels if not channel.closed)
         )
+        self._channels.clear()
 
     @property
     @implements(Server.stopped)

xoscar/backends/communication/ucx.py

@@ -368,7 +368,7 @@ class UCXServer(Server):
     scheme = "ucx"
 
     _ucp_listener: "ucp.Listener"  # type: ignore
-    _channels: List[UCXChannel]
+    _channels: set[UCXChannel]
 
     def __init__(
         self,
@@ -381,7 +381,7 @@ class UCXServer(Server):
         self.host = host
         self.port = port
         self._ucp_listener = ucp_listener
-        self._channels = []
+        self._channels = set()
         self._closed = asyncio.Event()
 
     @classproperty
@@ -469,9 +469,16 @@ class UCXServer(Server):
         channel = UCXChannel(
             ucp_endpoint, local_address=local_address, dest_address=dest_address
         )
-        self._channels.append(channel)
+        self._channels.add(channel)
         # handle over channel to some handlers
-        await self.channel_handler(channel)
+        try:
+            await self.channel_handler(channel)
+        finally:
+            if not channel.closed:
+                await channel.close()
+            # Remove channel if channel exit
+            self._channels.discard(channel)
+            logger.debug("Channel exit: %s", channel.info)
 
     @implements(Server.stop)
     async def stop(self):
@@ -480,7 +487,7 @@ class UCXServer(Server):
         await asyncio.gather(
             *(channel.close() for channel in self._channels if not channel.closed)
         )
-        self._channels = []
+        self._channels.clear()
         self._ucp_listener = None
         self._closed.set()
 

xoscar/backends/core.py

@@ -16,8 +16,11 @@
 from __future__ import annotations
 
 import asyncio
+import atexit
 import copy
 import logging
+import threading
+import weakref
 from typing import Type, Union
 
 from .._utils import Timer
@@ -31,8 +34,8 @@ ResultMessageType = Union[ResultMessage, ErrorMessage]
 logger = logging.getLogger(__name__)
 
 
-class ActorCaller:
-    __slots__ = "_client_to_message_futures", "_clients", "_profiling_data"
+class ActorCallerThreadLocal:
+    __slots__ = ("_client_to_message_futures", "_clients", "_profiling_data")
 
     _client_to_message_futures: dict[Client, dict[bytes, asyncio.Future]]
     _clients: dict[Client, asyncio.Task]
@@ -193,6 +196,7 @@ class ActorCaller:
         return await self.call_with_client(client, message, wait)
 
     async def stop(self):
+        logger.debug("Actor caller stop.")
         try:
             await asyncio.gather(*[client.close() for client in self._clients])
         except (ConnectionError, ServerClosed):
@@ -202,3 +206,38 @@ class ActorCaller:
     def cancel_tasks(self):
         # cancel listening for all clients
         _ = [task.cancel() for task in self._clients.values()]
+
+
+class ActorCaller:
+    __slots__ = "_thread_local"
+
+    class _RefHolder:
+        pass
+
+    _close_loop = asyncio.new_event_loop()
+    _close_thread = threading.Thread(target=_close_loop.run_forever, daemon=True)
+    _close_thread.start()
+    atexit.register(_close_loop.call_soon_threadsafe, _close_loop.stop)
+
+    def __init__(self):
+        self._thread_local = threading.local()
+
+    def __getattr__(self, item):
+        try:
+            actor_caller = self._thread_local.actor_caller
+        except AttributeError:
+            thread_info = str(threading.current_thread())
+            logger.debug("Creating a new actor caller for thread: %s", thread_info)
+            actor_caller = self._thread_local.actor_caller = ActorCallerThreadLocal()
+            ref = self._thread_local.ref = ActorCaller._RefHolder()
+            # If the thread exit, we clean the related actor callers and channels.
+
+            def _cleanup():
+                asyncio.run_coroutine_threadsafe(actor_caller.stop(), self._close_loop)
+                logger.debug(
+                    "Clean up the actor caller due to thread exit: %s", thread_info
+                )
+
+            weakref.finalize(ref, _cleanup)
+
+        return getattr(actor_caller, item)

{xoscar-0.4.1.dist-info → xoscar-0.4.3.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: xoscar
-Version: 0.4.1
+Version: 0.4.3
 Summary: Python actor framework for heterogeneous computing.
 Home-page: http://github.com/xorbitsai/xoscar
 Author: Qin Xuye

{xoscar-0.4.1.dist-info → xoscar-0.4.3.dist-info}/RECORD

@@ -1,5 +1,5 @@
 xoscar/__init__.py,sha256=dlwtB7dnDp5WME6CZVQY7d9lk1yJ9s___H5UxjGlAd4,1668
-xoscar/_utils.cp39-win_amd64.pyd,sha256=rA-ivDY6A_BsLAJaLh6nxgHBRS2H16c8E4qVq4P1WII,117248
+xoscar/_utils.cp39-win_amd64.pyd,sha256=FFd9OHF1DlHakZ6yDteAogr772FTnd1M36h5HHEtHtU,117248
 xoscar/_utils.pxd,sha256=rlNbTg5lhXA-jCOLksqF4jhUlNn0xw2jx1HxdLa34pc,1193
 xoscar/_utils.pyx,sha256=5Wvind3AQ3JaMK7Zv9SjhiPO6LEol2hW7_fMncn69go,7300
 xoscar/_version.py,sha256=bsfCVAo_o9LkiP3AjPsP4SRRqhjuS0t4D1WGJPzbdls,24412
@@ -7,10 +7,10 @@ xoscar/api.py,sha256=B5oXv4vgMxMteh1YNaBmNFDrUFmYa_dCdzfaWwwZnCo,13820
 xoscar/backend.py,sha256=8G5JwjoOT6Q2slb11eXNApxgcmvNQUCdQzkoIMDwLcQ,1917
 xoscar/batch.py,sha256=Jk5BSpvMFAV9DrRy0a9tgPvIo_dt8cbJReZBL0cnOPc,8128
 xoscar/constants.py,sha256=GJ1KEOxwnqksc9K_GH42TSWpQECeC6ti3KJmE3PUcTw,810
-xoscar/context.cp39-win_amd64.pyd,sha256=vs5foDr1olIiaEnZsYkMQZ9vGCMlpcJ8OK43SDwIE9A,158720
+xoscar/context.cp39-win_amd64.pyd,sha256=76-7TUpyoo8cjEh3CKN3IZ_jCEtMGz2fLQ9YqvsrBWI,158720
 xoscar/context.pxd,sha256=6n6IAbmArSRq8EjcsbS6npW8xP1jI0qOoS1fF0oyj-o,746
 xoscar/context.pyx,sha256=FOJVerGOvxe2USryXEQA0rpaFX_ScxISH6QWKUcahY8,11310
-xoscar/core.cp39-win_amd64.pyd,sha256=Qlcy3lIJsk0eVtIGjCMArFiUxx6cAhDkL3NUrlDCui0,334336
+xoscar/core.cp39-win_amd64.pyd,sha256=DxPbFOMnDiZZU6rT7qwin6JrcUnfsK09XZZ1YBLiMnU,334336
 xoscar/core.pxd,sha256=dGv62H92PFG98SVILuF641kY-NWFEt1FYqqOX3WY5RQ,1330
 xoscar/core.pyx,sha256=0YvJW2AHgymyfsAlPGvIFw65J5gTKO3PK2p1wl4VlJ0,22388
 xoscar/debug.py,sha256=hrmxIH6zvTKasQo6PUUgXu5mgEsR0g87Fvpw7CoHipg,5257
@@ -29,18 +29,18 @@ xoscar/backends/__init__.py,sha256=g9OllTquu9MRB5nySVoyiRv2z-_OSALWrOhwt7L9WXc,6
 xoscar/backends/allocate_strategy.py,sha256=DzvTlixwzTANURI2mDLHm3vcaugSPDxU6UQZb89KH0U,5005
 xoscar/backends/config.py,sha256=7nmvU_19zYR7n-bT8BNasbjntwmobmMiM7wN7O6Lujc,5129
 xoscar/backends/context.py,sha256=NukXzBwq9ZwuiN1y6RE1hfNGsW589hDtJAVwu-DV9E0,15874
-xoscar/backends/core.py,sha256=TNf-ROrmVixN9A82Fh1gh1UkqBdBZJd1_nlMtP4A0Bw,8243
-xoscar/backends/message.cp39-win_amd64.pyd,sha256=P6pGrismuoAN2wK2LIO0QR-c_RmNEWToDuSPrwU9uYk,268800
+xoscar/backends/core.py,sha256=UBgPX-9ZdC2b6Ci5dGowlqdp1Sb2JU9ULYOhHkd9wOQ,9623
+xoscar/backends/message.cp39-win_amd64.pyd,sha256=IomY97q4oiL6LGLCxUN0Zw5x5lW6jW48lsmQrdv8yc4,268800
 xoscar/backends/message.pyx,sha256=23j-fjxOVONPWY0-C5sfCOMbZCx4PTxN0GU8YfrLNeg,18529
 xoscar/backends/pool.py,sha256=nh2DTrcOrZc75TG2v7MXeGLMLhEO4GsObD7zirymZ3c,61445
 xoscar/backends/router.py,sha256=GJpSV_LhzINHdTp5jtsMHfPNMkNb1KI-WlqGqhwApGU,7953
 xoscar/backends/communication/__init__.py,sha256=Z0_RJkPGwLJeapSNt-TiO9DvnpBPu8P4PCooLaAqjkk,1080
 xoscar/backends/communication/base.py,sha256=wmWTeE4lcB_ohqyJJ6MdzMGcrOqy2RSKRp8y-NDuFdY,7736
 xoscar/backends/communication/core.py,sha256=ZM_fU0yokoPzXcJ6j-89llp2r7J9pl3gE5iImn4b4bE,2199
-xoscar/backends/communication/dummy.py,sha256=apbkfFdVWqeqc8Cc1LInOw6ikVoFWixTwBN__djxgBY,8040
+xoscar/backends/communication/dummy.py,sha256=hxp2wvBFR-X63Nxj128yBNKn5D_azW-UrWNCQNYJHJI,8739
 xoscar/backends/communication/errors.py,sha256=-O6ZaKs6Gz7g_ZnKU2wrz8D9btowqaYuQ9u0zCYSLWo,743
-xoscar/backends/communication/socket.py,sha256=sq9abnWvngA55Sxz3TV0nL4JuxZX79DgTQn8E3xbRXY,14111
-xoscar/backends/communication/ucx.py,sha256=defQN2sirsu-KqvXox_I_wzAfWSq5LJkjKdzo0vtPgw,20234
+xoscar/backends/communication/socket.py,sha256=AQMGCOroKabjQHwm50dKH8blFp_7mFiW4JCrhe97-jE,14400
+xoscar/backends/communication/ucx.py,sha256=SQh461aiqdKKZe9fC4sDroFO2FmtQ_vX0oKSQjnqoFA,20498
 xoscar/backends/communication/utils.py,sha256=F-muF5_ow9JzAPAZ3d8XaGDiz3bRZOdWmWBDwQOVLe0,3761
 xoscar/backends/indigen/__init__.py,sha256=Khr2aGbaIw_04NIdY7QNhdljCKbmmzLb4NeAOM3LF8M,701
 xoscar/backends/indigen/backend.py,sha256=cCMkZDEKuRQlFb6v79JvWi5lfzsvAafznOeEWlw7CWY,1663
@@ -54,8 +54,8 @@ xoscar/collective/common.py,sha256=9c7xq3IOUvfA0I9GnpalUqXZOzmF6IEILv4zL64BYVE,3
 xoscar/collective/core.py,sha256=191aPxbUgWpjzrqyozndImDAQhZFmqoQdBkHFLDfXN0,24239
 xoscar/collective/process_group.py,sha256=kTPbrLMJSGhqbiWvTIiz-X3W0rZWd_CFn_zUIlXbOlM,23286
 xoscar/collective/utils.py,sha256=p3WEVtXvnVhkuO5mRgQBhBRFr1dKHcDKMjrbMyuiyfg,1219
-xoscar/collective/uv.dll,sha256=iPM9bHhRkk5hi2xcYQzj5R1hWm1tg-XVsouCAuoigYc,620544
-xoscar/collective/xoscar_pygloo.cp39-win_amd64.pyd,sha256=whMek_2N9JrumuEIpr4UQpHENkYTTnCoyPO1PdTz7WI,775168
+xoscar/collective/uv.dll,sha256=r0EqFH4Q-ae4ruaT8tBL07fQT962_YhxL25R6bwsocs,620544
+xoscar/collective/xoscar_pygloo.cp39-win_amd64.pyd,sha256=e0Fgk3HD1EnswaD9OjQ2Nn79PiCgxiRM6poGN_J7Cus,775168
 xoscar/metrics/__init__.py,sha256=RjXuuYw4I2YYgD8UY2Z5yCZk0Z56xMJ1n40O80Dtxf8,726
 xoscar/metrics/api.py,sha256=dtJ4QrIqQNXhJedeqOPs4TXKgrRGZFFN50xAd9SCfec,9144
 xoscar/metrics/backends/__init__.py,sha256=ZHepfhCDRuK9yz4pAM7bjpWDvS3Ijp1YgyynoUFLeuU,594
@@ -66,7 +66,7 @@ xoscar/metrics/backends/prometheus/__init__.py,sha256=ZHepfhCDRuK9yz4pAM7bjpWDvS
 xoscar/metrics/backends/prometheus/prometheus_metric.py,sha256=65hb8O3tmsEJ7jgOrIwl_suj9SE5Tmqcfjuk0urkLvE,2120
 xoscar/serialization/__init__.py,sha256=NOAn046vnHEkx--82BKNinV8EpyOfT5hqfRBGnKl56s,866
 xoscar/serialization/aio.py,sha256=7YLXgkWpQ3ANy-TZ1qO8Mt4_J3cZFhFh2FEgUgxMT60,4873
-xoscar/serialization/core.cp39-win_amd64.pyd,sha256=VcEI1DR7qE5PF3bUVUvxFMLp8dvzBw_zaxC_0CDXCOw,298496
+xoscar/serialization/core.cp39-win_amd64.pyd,sha256=FFRO3EB7jb_z5pq16gIb8neIVFLTwYF20vVwMjlv2po,298496
 xoscar/serialization/core.pxd,sha256=X-47bqBM2Kzw5SkLqICdKD0gU6CpmLsBxC3kfW--wVk,1013
 xoscar/serialization/core.pyx,sha256=ZKexLRnRwZXXn2045kR7xfM_szcoPNrDuouQCWtpFp8,30570
 xoscar/serialization/cuda.py,sha256=Fj4Cpr_YmkGceUCo0mQn8fRvmHP_5WcLdRx6epZ3RC0,3869
@@ -74,7 +74,7 @@ xoscar/serialization/exception.py,sha256=t6yZn_Ate04UE1RbabNh7mu739sdtwarjuPXWhA
 xoscar/serialization/numpy.py,sha256=C6WVx-Sdl2OHBAvVY34DFjAKXlekMbpc2ni6bR8wxYo,3001
 xoscar/serialization/pyfury.py,sha256=3ucal29Hr7PX9_1SfB2x43FE2xw_C0rLkVv3foL7qwM,1200
 xoscar/serialization/scipy.py,sha256=9ph-yoRoNiwUZTwQrn35U60VPirWlncXNAg6EXvqMR4,2554
-xoscar-0.4.1.dist-info/METADATA,sha256=BfnUJHhYjnOGtAknk32D1ACdmTse9w3dZ-tE-NoIXQE,9265
-xoscar-0.4.1.dist-info/WHEEL,sha256=yA7mxgqX2UV73NtJdMh2AAmdb628loM81912H3s5r00,100
-xoscar-0.4.1.dist-info/top_level.txt,sha256=vYlqqY4Nys8Thm1hePIuUv8eQePdULVWMmt7lXtX_ZA,21
-xoscar-0.4.1.dist-info/RECORD,,
+xoscar-0.4.3.dist-info/METADATA,sha256=9OJj9wwgu76by-N6gTxrPIjqkkpPSI_jFKMRAAqrc6Q,9265
+xoscar-0.4.3.dist-info/WHEEL,sha256=yA7mxgqX2UV73NtJdMh2AAmdb628loM81912H3s5r00,100
+xoscar-0.4.3.dist-info/top_level.txt,sha256=vYlqqY4Nys8Thm1hePIuUv8eQePdULVWMmt7lXtX_ZA,21
+xoscar-0.4.3.dist-info/RECORD,,