xoscar 0.4.1__cp39-cp39-macosx_10_9_x86_64.whl → 0.4.2__cp39-cp39-macosx_10_9_x86_64.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

@@ -18,6 +18,8 @@ from __future__ import annotations
 import asyncio
 import copy
 import logging
+import threading
+import weakref
 from typing import Type, Union
 
 from .._utils import Timer
@@ -31,8 +33,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 +195,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 +205,37 @@ 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()
+
+    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.2.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: xoscar
-Version: 0.4.1
+Version: 0.4.2
 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.2.dist-info}/RECORD

@@ -1,7 +1,7 @@
-xoscar-0.4.1.dist-info/RECORD,,
-xoscar-0.4.1.dist-info/WHEEL,sha256=FoxzhRRMXcm8mSmqIxEPOmfhajsNIlMLdYtHdyz-ZKQ,109
-xoscar-0.4.1.dist-info/top_level.txt,sha256=vYlqqY4Nys8Thm1hePIuUv8eQePdULVWMmt7lXtX_ZA,21
-xoscar-0.4.1.dist-info/METADATA,sha256=s1flP4I8DrL3z4TI3mPEAZpBJTlNTM9UwVVmRNXB1-Y,9042
+xoscar-0.4.2.dist-info/RECORD,,
+xoscar-0.4.2.dist-info/WHEEL,sha256=FoxzhRRMXcm8mSmqIxEPOmfhajsNIlMLdYtHdyz-ZKQ,109
+xoscar-0.4.2.dist-info/top_level.txt,sha256=vYlqqY4Nys8Thm1hePIuUv8eQePdULVWMmt7lXtX_ZA,21
+xoscar-0.4.2.dist-info/METADATA,sha256=86wcGCexxpSltidUVDP0KShGLsodQR5lck9OpOWgHIc,9042
 xoscar/_utils.pyx,sha256=UR1FtYXAYKIdEWR9HulEpMbSOrkQWi6xGz63d4IQmG0,7059
 xoscar/backend.py,sha256=is436OPkZfSpQXaoqTRVta5eoye_pp45RFgCstAk2hU,1850
 xoscar/core.pxd,sha256=4lBq8J0kjcXcsGuvN7Kv4xcL5liHwTTFWlqyK7XAEnw,1280
@@ -12,18 +12,18 @@ xoscar/nvutils.py,sha256=qmW4mKLU0WB2yCs198ccQOgLL02zB7Fsa-AotO3NOmg,20412
 xoscar/constants.py,sha256=QHHSREw6uWBBjQDCFqlNfTvBZgniJPGy42KSIsR8Fqw,787
 xoscar/__init__.py,sha256=0zX8kKaio3ZIrlzB79WybcravMJw1OxPWjDspTgJFyQ,1608
 xoscar/api.py,sha256=3hztPoOxg8A_mlhWyWgVP7FMXG0PATA1TP4Rbaj7A-g,13327
-xoscar/core.cpython-39-darwin.so,sha256=Jdl9NelbQoqEAc-cL8_sLe1ywnQcYrSuRMdYpc-2o0Q,448264
+xoscar/core.cpython-39-darwin.so,sha256=AmnkUdJv6-zKNVe0VrXZ8WeDSJP3tC_s4v4heq7Fz7Y,448264
 xoscar/utils.py,sha256=jUw6OICZUPBbmS1b3GE4vLctJf6fCKXrYtLtBuK-Oqc,16483
 xoscar/debug.py,sha256=9Z8SgE2WaKYQcyDo-5-DxEJQ533v7kWjrvCd28pSx3E,5069
 xoscar/libcpp.pxd,sha256=DJqBxLFOKL4iRr9Kale5UH3rbvPRD1x5bTSOPHFpz9I,1147
 xoscar/context.pyx,sha256=8CdgPnWcE9eOp3N600WgDQ03MCi8P73eUOGcfV7Zksg,10942
 xoscar/errors.py,sha256=wBlQOKsXf0Fc4skN39tDie0YZT-VIAuLNRgoDl2pZcA,1241
-xoscar/_utils.cpython-39-darwin.so,sha256=Fr0PrhRCQ6PvlPkueT74lRCnJQ6FgdtUfHPcoW9aEWU,160824
+xoscar/_utils.cpython-39-darwin.so,sha256=DfpsecCvGEj-VmwuFx3jDhst2n85JC-j2ZP2yRtK_m0,160824
 xoscar/core.pyx,sha256=Aqc2i8Fetsd5wRAPF4kL0ddnBZn3E2HRNCvup79BbQc,21730
 xoscar/driver.py,sha256=498fowtJr6b3FE8FIOA_Tc1Vwx88nfZw7p0FxrML0h4,1372
 xoscar/profiling.py,sha256=BC5OF0HzSaXv8V7w-y-B8r5gV5DgxHFoTEIF6jCMioQ,8015
 xoscar/_utils.pxd,sha256=5KYAL3jfPdejsHnrGGT2s--ZUX5SXznQWpHVSno429k,1157
-xoscar/context.cpython-39-darwin.so,sha256=UkkEl3HqCkezVnFrHDoupgbN6N_CNaE_53sSfxounCY,200448
+xoscar/context.cpython-39-darwin.so,sha256=yel10nYAEY6PysHn5DLdcyl156MS2qQXjGZdt-N9Jso,200448
 xoscar/metrics/__init__.py,sha256=9Badi7rxYikGm2dQiNCrj9GgMRBxwuR3JaEKcFZmfak,705
 xoscar/metrics/api.py,sha256=BBlMIFvVAGVfrtpeJ1YlH9Tqhy9OzGavwvGyeHcQ0Tk,8856
 xoscar/metrics/backends/__init__.py,sha256=h_JgzSqV5lP6vQ6XX_17kE4IY4BRnvKta_7VLQAL1ms,581
@@ -43,16 +43,16 @@ xoscar/serialization/pyfury.py,sha256=sifOnVMYoS82PzZEkzkfxesmMHei23k5UAUUKUyoOY
 xoscar/serialization/core.pxd,sha256=k4RoJgX5E5LGs4jdCQ7vvcn26MabXbrWoWhkO49X6YI,985
 xoscar/serialization/__init__.py,sha256=5Y_C3cYbQJIZ09LRjeCf-jrkLma7mfN8I5bznHrdsbg,846
 xoscar/serialization/numpy.py,sha256=5Kem87CvpJmzUMp3QHk4WeHU30FoQWTJJP2SwIcaQG0,2919
-xoscar/serialization/core.cpython-39-darwin.so,sha256=yA1neOIhTms_IGsCPKK57JaiyQeC2j-6OWUJMUuEe-g,408080
+xoscar/serialization/core.cpython-39-darwin.so,sha256=OBeZI8npBa5M46HtWHAE1FtxqMe0URdNSv4NCc-TbUI,408080
 xoscar/serialization/cuda.py,sha256=iFUEnN4SiquBIhyieyOrfw3TnKnW-tU_vYgqOxO_DrA,3758
 xoscar/serialization/scipy.py,sha256=yOEi0NB8cqQ6e2UnCZ1w006RsB7T725tIL-DM_hNcsU,2482
 xoscar/serialization/aio.py,sha256=5DySPgDxU43ec7_5Ct44-Oqt7YNSJBfuf8VdQgQlChA,4731
 xoscar/serialization/core.pyx,sha256=bjR-zXGm9qersk7kYPzpjpMIxDl_Auur4BCubRfKmfA,29626
-xoscar/backends/message.cpython-39-darwin.so,sha256=XsJRs2s6D5KBvPLB6uMS-aWvCrRouDPWigKdO2LpOcQ,366752
+xoscar/backends/message.cpython-39-darwin.so,sha256=rmxdgx6nD60Hfm3vUCrpr-C5fMjTvz2USHzIgpmcT9o,366752
 xoscar/backends/config.py,sha256=EG26f0GwX_f4dAhwTW77RBjiK9h8R_3JrD-rBF1bAq8,4984
 xoscar/backends/allocate_strategy.py,sha256=tC1Nbq2tJohahUwd-zoRYHEDX65wyuX8tmeY45uWj_w,4845
 xoscar/backends/__init__.py,sha256=VHEBQcUWM5bj027W8EUf9PiJUAP7JoMrRw3Tsvy5ySw,643
-xoscar/backends/core.py,sha256=YcXVrMrTUjnrnH-RhrtUyxnCbwon8miq2UigCQv-y_Q,8039
+xoscar/backends/core.py,sha256=GlhWdwl6gWBBwHynV4UWqOXVKDS3OMfAtOLJQHNqV-8,9294
 xoscar/backends/context.py,sha256=Vr_PibRxYCDQ_gYK7r-BOlw9TXw8VQbFsVTH7K7mHPk,15470
 xoscar/backends/router.py,sha256=mhSvM5KVfV882jricVcpyxAqHEvhS4zL6ivczC6fOTE,7746
 xoscar/backends/message.pyx,sha256=uyzilPc_7SqNwGUL4U-Zbfqku8bfZyRW_Lt_S3I_LEU,17930
@@ -64,13 +64,13 @@ xoscar/backends/indigen/pool.py,sha256=mWYkOP4VVoUsXFgfpwruPuWblF6Waan5vxit8B-9_
 xoscar/backends/test/backend.py,sha256=nv9WFhH5Bbq4Q1HB9yfpciZBaeHT4IQAtzugBWESrUY,1263
 xoscar/backends/test/__init__.py,sha256=j2ZfD6prD9WjUxRUDC7Eq5Z7N7TkL6fFr59oNyc_vY4,682
 xoscar/backends/test/pool.py,sha256=TW4X6J-92Pti66103poQBNDBznX6CBD3RLOc_zixjTo,7257
-xoscar/backends/communication/ucx.py,sha256=7GAKIzlbxy-NpaDE9VLPde90tKxeU8lVZH8c_3ByuZ0,19703
+xoscar/backends/communication/ucx.py,sha256=_Dp9Ld2MWIa1txSGMnmfYwJDT0esxS-GOd2FQ4BdHiM,19960
 xoscar/backends/communication/__init__.py,sha256=tB05BlK63iWQnfJgRzKt4mFKRtmWUki5hUGSZQwAotc,1050
 xoscar/backends/communication/core.py,sha256=sJeE3foRIqVPXldzYpFKHDSsabfAIFBU4JuXY4OyklY,2130
 xoscar/backends/communication/utils.py,sha256=AmovE-hmWLXNCPwHafYuaRjOk8m42BUyT3XBqfXQRVI,3664
 xoscar/backends/communication/errors.py,sha256=V3CdBe2xX9Rwv32f2dH2Msc84yaUhlyerZ42-739o1Q,723
-xoscar/backends/communication/socket.py,sha256=ZSDqS_Z9-Oxs2Q2fqWUuUJ9hNf4w8lNJ8gUv5r6Ji_Y,13691
-xoscar/backends/communication/dummy.py,sha256=gaKPNiN4x2aGZV3IGaaa8eaweBVjRh8B19jU1B5t2yw,7798
+xoscar/backends/communication/socket.py,sha256=R9jF_Iy3HTSPXkBGrws5ckK-EjYnaPBSTdMqMZYqxqc,13972
+xoscar/backends/communication/dummy.py,sha256=sKBqoFkftm6ynV3p5UNGI4VYFGlrOHJgJlqj35PGAgk,8479
 xoscar/backends/communication/base.py,sha256=0P4Tr35GSWpRp394e9jVWUUoKKa-gIk177eYPw1BnSU,7421
 xoscar/aio/__init__.py,sha256=kViDKR_kJe59VQViHITKEfBcIgN4ZJblUyd8zl0E3ZI,675
 xoscar/aio/file.py,sha256=PBtkLp-Q7XtYl-zk00s18TtgIrkNr60J3Itf66ctO1o,1486