xoscar 0.4.0__tar.gz → 0.4.2__tar.gz

{xoscar-0.4.0/xoscar.egg-info → xoscar-0.4.2}/PKG-INFO

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

{xoscar-0.4.0 → xoscar-0.4.2}/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-0.4.0 → xoscar-0.4.2}/xoscar/backends/communication/socket.py

@@ -29,7 +29,7 @@ from typing import Any, Callable, Coroutine, Dict, Type
 from urllib.parse import urlparse
 
 from ..._utils import to_binary
-from ...constants import XOSCAR_UNIX_SOCKET_DIR
+from ...constants import XOSCAR_CONNECT_TIMEOUT, XOSCAR_UNIX_SOCKET_DIR
 from ...serialization import AioDeserializer, AioSerializer, deserialize
 from ...utils import classproperty, implements, is_py_312, is_v6_ip
 from .base import Channel, ChannelType, Client, Server
@@ -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)
@@ -291,7 +299,13 @@ class SocketClient(Client):
     ) -> "Client":
         host, port_str = dest_address.rsplit(":", 1)
         port = int(port_str)
-        (reader, writer) = await asyncio.open_connection(host=host, port=port, **kwargs)
+        config = kwargs.get("config", {})
+        connect_timeout = config.get("connect_timeout", XOSCAR_CONNECT_TIMEOUT)
+        fut = asyncio.open_connection(host=host, port=port)
+        try:
+            reader, writer = await asyncio.wait_for(fut, timeout=connect_timeout)
+        except asyncio.TimeoutError:
+            raise ConnectionError("connect timeout")
         channel = SocketChannel(
             reader, writer, local_address=local_address, dest_address=dest_address
         )

{xoscar-0.4.0 → xoscar-0.4.2}/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-0.4.0 → xoscar-0.4.2}/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]
@@ -70,50 +72,61 @@ class ActorCaller:
         return client
 
     async def _listen(self, client: Client):
-        while not client.closed:
-            try:
+        try:
+            while not client.closed:
                 try:
-                    message: _MessageBase = await client.recv()
-                except (EOFError, ConnectionError, BrokenPipeError):
-                    # remote server closed, close client and raise ServerClosed
                     try:
-                        await client.close()
-                    except (ConnectionError, BrokenPipeError):
-                        # close failed, ignore it
+                        message: _MessageBase = await client.recv()
+                    except (EOFError, ConnectionError, BrokenPipeError) as e:
+                        # AssertionError is from get_header
+                        # remote server closed, close client and raise ServerClosed
+                        logger.debug(f"{client.dest_address} close due to {e}")
+                        try:
+                            await client.close()
+                        except (ConnectionError, BrokenPipeError):
+                            # close failed, ignore it
+                            pass
+                        raise ServerClosed(
+                            f"Remote server {client.dest_address} closed: {e}"
+                        ) from None
+                    future = self._client_to_message_futures[client].pop(
+                        message.message_id
+                    )
+                    if not future.done():
+                        future.set_result(message)
+                except DeserializeMessageFailed as e:
+                    message_id = e.message_id
+                    future = self._client_to_message_futures[client].pop(message_id)
+                    future.set_exception(e.__cause__)  # type: ignore
+                except Exception as e:  # noqa: E722  # pylint: disable=bare-except
+                    message_futures = self._client_to_message_futures[client]
+                    self._client_to_message_futures[client] = dict()
+                    for future in message_futures.values():
+                        future.set_exception(copy.copy(e))
+                finally:
+                    # message may have Ray ObjectRef, delete it early in case next loop doesn't run
+                    # as soon as expected.
+                    try:
+                        del message
+                    except NameError:
                         pass
-                    raise ServerClosed(
-                        f"Remote server {client.dest_address} closed"
-                    ) from None
-                future = self._client_to_message_futures[client].pop(message.message_id)
-                if not future.done():
-                    future.set_result(message)
-            except DeserializeMessageFailed as e:
-                message_id = e.message_id
-                future = self._client_to_message_futures[client].pop(message_id)
-                future.set_exception(e.__cause__)  # type: ignore
-            except Exception as e:  # noqa: E722  # pylint: disable=bare-except
-                message_futures = self._client_to_message_futures[client]
-                self._client_to_message_futures[client] = dict()
-                for future in message_futures.values():
-                    future.set_exception(copy.copy(e))
-            finally:
-                # message may have Ray ObjectRef, delete it early in case next loop doesn't run
-                # as soon as expected.
-                try:
-                    del message
-                except NameError:
-                    pass
-                try:
-                    del future
-                except NameError:
-                    pass
-                await asyncio.sleep(0)
+                    try:
+                        del future
+                    except NameError:
+                        pass
+                    await asyncio.sleep(0)
 
-        message_futures = self._client_to_message_futures[client]
-        self._client_to_message_futures[client] = dict()
-        error = ServerClosed(f"Remote server {client.dest_address} closed")
-        for future in message_futures.values():
-            future.set_exception(copy.copy(error))
+            message_futures = self._client_to_message_futures[client]
+            self._client_to_message_futures[client] = dict()
+            error = ServerClosed(f"Remote server {client.dest_address} closed")
+            for future in message_futures.values():
+                future.set_exception(copy.copy(error))
+        finally:
+            try:
+                await client.close()
+            except:  # noqa: E722  # nosec  # pylint: disable=bare-except
+                # ignore all error if fail to close at last
+                pass
 
     async def call_with_client(
         self, client: Client, message: _MessageBase, wait: bool = True
@@ -182,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):
@@ -191,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.0 → xoscar-0.4.2}/xoscar/backends/pool.py

@@ -551,23 +551,31 @@ class AbstractActorPool(ABC):
         return False
 
     async def on_new_channel(self, channel: Channel):
-        while not self._stopped.is_set():
-            try:
-                message = await channel.recv()
-            except EOFError:
-                # no data to read, check channel
+        try:
+            while not self._stopped.is_set():
                 try:
-                    await channel.close()
-                except (ConnectionError, EOFError):
-                    # close failed, ignore
-                    pass
-                return
-            if await self._handle_ucx_meta_message(message, channel):
-                continue
-            asyncio.create_task(self.process_message(message, channel))
-            # delete to release the reference of message
-            del message
-            await asyncio.sleep(0)
+                    message = await channel.recv()
+                except (EOFError, ConnectionError, BrokenPipeError) as e:
+                    logger.debug(f"pool: close connection due to {e}")
+                    # no data to read, check channel
+                    try:
+                        await channel.close()
+                    except (ConnectionError, EOFError):
+                        # close failed, ignore
+                        pass
+                    return
+                if await self._handle_ucx_meta_message(message, channel):
+                    continue
+                asyncio.create_task(self.process_message(message, channel))
+                # delete to release the reference of message
+                del message
+                await asyncio.sleep(0)
+        finally:
+            try:
+                await channel.close()
+            except:  # noqa: E722  # nosec  # pylint: disable=bare-except
+                # ignore all error if fail to close at last
+                pass
 
     async def __aenter__(self):
         await self.start()

{xoscar-0.4.0 → xoscar-0.4.2}/xoscar/constants.py

@@ -19,3 +19,5 @@ XOSCAR_TEMP_DIR = Path(os.getenv("XOSCAR_DIR", Path.home())) / ".xoscar"
 
 # unix socket.
 XOSCAR_UNIX_SOCKET_DIR = XOSCAR_TEMP_DIR / "socket"
+
+XOSCAR_CONNECT_TIMEOUT = 8

{xoscar-0.4.0 → xoscar-0.4.2}/xoscar/serialization/aio.py

@@ -77,7 +77,11 @@ MALFORMED_MSG = "Received malformed data, please check Xoscar version on both si
 def get_header_length(header_bytes: bytes):
     version = struct.unpack("B", header_bytes[:1])[0]
     # now we only have default version
-    assert version == DEFAULT_SERIALIZATION_VERSION, MALFORMED_MSG
+    if version != DEFAULT_SERIALIZATION_VERSION:
+        # when version not matched,
+        # we will immediately abort the connection
+        # EOFError will be captured by channel
+        raise EOFError(MALFORMED_MSG)
     # header length
     header_length = struct.unpack("<Q", header_bytes[1:9])[0]
     # compress

{xoscar-0.4.0 → xoscar-0.4.2}/xoscar/utils.py

@@ -19,11 +19,11 @@ import asyncio
 import dataclasses
 import functools
 import importlib
+import importlib.util as importlib_utils
 import inspect
 import io
 import logging
 import os
-import pkgutil
 import random
 import socket
 import sys
@@ -267,7 +267,7 @@ def lazy_import(
             self._on_loads.append(func)
             return func
 
-    if pkgutil.find_loader(prefix_name) is not None:
+    if importlib_utils.find_spec(prefix_name) is not None:
         return LazyModule()
     elif placeholder:
         return ModulePlaceholder(prefix_name)

{xoscar-0.4.0 → xoscar-0.4.2/xoscar.egg-info}/PKG-INFO

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