xoscar 0.9.0__cp312-cp312-macosx_10_13_x86_64.whl

xoscar/backends/communication/socket.py

@@ -0,0 +1,444 @@
+# Copyright 2022-2023 XProbe Inc.
+# derived from copyright 1999-2021 Alibaba Group Holding Ltd.
+#
+# 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.
+
+from __future__ import annotations
+
+import asyncio
+import concurrent.futures as futures
+import logging
+import os
+import socket
+import sys
+from abc import ABCMeta
+from asyncio import AbstractServer, StreamReader, StreamWriter
+from functools import lru_cache
+from hashlib import md5
+from typing import Any, Callable, Coroutine, Dict, Type
+from urllib.parse import urlparse
+
+from ..._utils import to_binary
+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_py_312_or_above, is_v6_ip
+from .base import Channel, ChannelType, Client, Server
+from .core import register_client, register_server
+from .errors import ChannelClosed
+from .utils import read_buffers, write_buffers
+
+_is_windows: bool = sys.platform.startswith("win")
+
+
+logger = logging.getLogger(__name__)
+
+
+class SocketChannel(Channel):
+    __slots__ = "reader", "writer", "_channel_type", "_send_lock", "_recv_lock"
+
+    name = "socket"
+
+    def __init__(
+        self,
+        reader: StreamReader,
+        writer: StreamWriter,
+        local_address: str | None = None,
+        dest_address: str | None = None,
+        compression: str | None = None,
+        channel_type: int | None = None,
+    ):
+        super().__init__(
+            local_address=local_address,
+            dest_address=dest_address,
+            compression=compression,
+        )
+        self.reader = reader
+        self.writer = writer
+        self._channel_type = channel_type
+
+        self._send_lock = asyncio.Lock()
+        self._recv_lock = asyncio.Lock()
+
+    @property
+    @implements(Channel.type)
+    def type(self) -> int:
+        return self._channel_type  # type: ignore
+
+    @implements(Channel.send)
+    async def send(self, message: Any):
+        # get buffers
+        compress = self.compression or 0
+        serializer = AioSerializer(message, compress=compress)
+        buffers = await serializer.run()
+
+        try:
+            # write buffers
+            write_buffers(self.writer, buffers)
+            async with self._send_lock:
+                # add lock, or when parallel send,
+                # assertion error may be raised
+                await self.writer.drain()
+        except RuntimeError as e:
+            if self.writer.is_closing():
+                raise ChannelClosed(
+                    "Channel already closed, cannot write message"
+                ) from e
+            raise e
+
+    @implements(Channel.recv)
+    async def recv(self):
+        deserializer = AioDeserializer(self.reader)
+        async with self._recv_lock:
+            header = await deserializer.get_header()
+            buffers = await read_buffers(header, self.reader)
+        return deserialize(header, buffers)
+
+    @implements(Channel.close)
+    async def close(self):
+        self.writer.close()
+        try:
+            await self.writer.wait_closed()
+        # TODO: May raise Runtime error: attach to another event loop
+        except (ConnectionResetError, RuntimeError):  # pragma: no cover
+            pass
+
+    @property
+    @implements(Channel.closed)
+    def closed(self):
+        return self.writer.is_closing()
+
+
+class _BaseSocketServer(Server, metaclass=ABCMeta):
+    __slots__ = "_aio_server", "_channels"
+
+    _channels: set[Channel]
+
+    def __init__(
+        self,
+        address: str,
+        aio_server: AbstractServer,
+        channel_handler: Callable[[Channel], Coroutine] | None = None,
+    ):
+        super().__init__(address, channel_handler)
+        # asyncio.Server
+        self._aio_server = aio_server
+        self._channels = set()
+
+    @implements(Server.start)
+    async def start(self):
+        await self._aio_server.start_serving()
+
+    @implements(Server.join)
+    async def join(self, timeout=None):
+        if timeout is None:
+            await self._aio_server.serve_forever()
+        else:
+            if is_py_312():
+                # For python 3.12, there's a bug for `serve_forever`:
+                # https://github.com/python/cpython/issues/123720,
+                # which is unable to be cancelled.
+                # Here is really a simulation of `wait_for`
+                task = asyncio.create_task(self._aio_server.serve_forever())
+                await asyncio.sleep(timeout)
+                if task.done():
+                    logger.warning(f"`serve_forever` should never be done.")
+                else:
+                    task.cancel()
+            else:
+                future = asyncio.create_task(self._aio_server.serve_forever())
+                try:
+                    await asyncio.wait_for(future, timeout=timeout)
+                except (futures.TimeoutError, asyncio.TimeoutError, TimeoutError):
+                    future.cancel()
+
+    @implements(Server.on_connected)
+    async def on_connected(self, *args, **kwargs):
+        reader, writer = args
+        local_address = kwargs.pop("local_address", None)
+        dest_address = kwargs.pop("dest_address", None)
+        if kwargs:  # pragma: no cover
+            raise TypeError(
+                f"{type(self).__name__} got unexpected "
+                f'arguments: {",".join(kwargs)}'
+            )
+        channel = SocketChannel(
+            reader,
+            writer,
+            local_address=local_address,
+            dest_address=dest_address,
+            channel_type=self.channel_type,
+        )
+        self._channels.add(channel)
+        # handle over channel to some handlers
+        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):
+        self._aio_server.close()
+        # Python 3.12+: # https://github.com/python/cpython/issues/104344
+        # `wait_closed` leads to hang in Python 3.12 and 3.13
+        if not is_py_312_or_above():
+            await self._aio_server.wait_closed()
+        # close all channels
+        await asyncio.gather(
+            *(channel.close() for channel in self._channels if not channel.closed)
+        )
+        self._channels.clear()
+
+    @property
+    @implements(Server.stopped)
+    def stopped(self) -> bool:
+        return not self._aio_server.is_serving()
+
+
+@register_server
+class SocketServer(_BaseSocketServer):
+    __slots__ = "host", "port"
+
+    scheme = None
+
+    def __init__(
+        self,
+        host: str,
+        port: int,
+        aio_server: AbstractServer,
+        channel_handler: Callable[[Channel], Coroutine] | None = None,
+    ):
+        address = f"{host}:{port}"
+        super().__init__(address, aio_server, channel_handler=channel_handler)
+        self.host = host
+        self.port = port
+
+    @classproperty
+    @implements(Server.client_type)
+    def client_type(self) -> Type["Client"]:
+        return SocketClient
+
+    @property
+    @implements(Server.channel_type)
+    def channel_type(self) -> int:
+        return ChannelType.remote
+
+    @classmethod
+    def parse_config(cls, config: dict) -> dict:
+        if config is None or not config:
+            return dict()
+        # we only need the following config
+        keys = ["listen_elastic_ip"]
+        parsed_config = {key: config[key] for key in keys if key in config}
+
+        return parsed_config
+
+    @staticmethod
+    @implements(Server.create)
+    async def create(config: Dict) -> "Server":
+        config = config.copy()
+        if "address" in config:
+            address = config.pop("address")
+            host, port = address.rsplit(":", 1)
+            port = int(port)
+        else:
+            host = config.pop("host")
+            port = int(config.pop("port"))
+        _host = host
+        if config.pop("listen_elastic_ip", False):
+            # The Actor.address will be announce to client, and is not on our host,
+            # cannot actually listen on it,
+            # so we have to keep SocketServer.host untouched to make sure Actor.address not changed
+            if is_v6_ip(host):
+                _host = "::"
+            else:
+                _host = "0.0.0.0"
+
+        handle_channel = config.pop("handle_channel")
+        if "start_serving" not in config:
+            config["start_serving"] = False
+
+        async def handle_connection(reader: StreamReader, writer: StreamWriter):
+            # create a channel when client connected
+            return await server.on_connected(
+                reader, writer, local_address=server.address
+            )
+
+        port = port if port != 0 else None
+        aio_server = await asyncio.start_server(
+            handle_connection, host=_host, port=port, **config
+        )
+
+        # get port of the socket if not specified
+        if not port:
+            port = aio_server.sockets[0].getsockname()[1]
+
+        if _is_windows:
+            for sock in aio_server.sockets:
+                sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, True)
+
+        server = SocketServer(host, port, aio_server, channel_handler=handle_channel)
+        return server
+
+
+@register_client
+class SocketClient(Client):
+    __slots__ = ()
+
+    scheme = SocketServer.scheme
+
+    @staticmethod
+    @implements(Client.connect)
+    async def connect(
+        dest_address: str, local_address: str | None = None, **kwargs
+    ) -> "Client":
+        host, port_str = dest_address.rsplit(":", 1)
+        port = int(port_str)
+        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,
+            channel_type=ChannelType.remote,
+        )
+        return SocketClient(local_address, dest_address, channel)
+
+
+def _get_or_create_default_unix_socket_dir():
+    os.makedirs(XOSCAR_UNIX_SOCKET_DIR, exist_ok=True)
+    return XOSCAR_UNIX_SOCKET_DIR
+
+
+@lru_cache(100)
+def _gen_unix_socket_default_path(process_index):
+    return (
+        f"{_get_or_create_default_unix_socket_dir()}/"
+        f"{md5(to_binary(str(process_index))).hexdigest()}"
+    )  # nosec
+
+
+@register_server
+class UnixSocketServer(_BaseSocketServer):
+    __slots__ = "process_index", "path"
+
+    scheme = "unixsocket"
+
+    def __init__(
+        self,
+        process_index: int,
+        aio_server: AbstractServer,
+        path: str,
+        channel_handler: Callable[[Channel], Coroutine] | None = None,
+    ):
+        address = f"{self.scheme}:///{process_index}"
+        super().__init__(address, aio_server, channel_handler=channel_handler)
+        self.process_index = process_index
+        self.path = path
+
+    @classproperty
+    @implements(Server.client_type)
+    def client_type(self) -> Type["Client"]:
+        return UnixSocketClient
+
+    @property
+    @implements(Server.channel_type)
+    def channel_type(self) -> int:
+        return ChannelType.ipc
+
+    @staticmethod
+    @implements(Server.create)
+    async def create(config: Dict) -> "Server":
+        config = config.copy()
+        if "address" in config:
+            process_index = int(urlparse(config.pop("address")).path.lstrip("/"))
+        else:
+            process_index = config.pop("process_index")
+        handle_channel = config.pop("handle_channel")
+        path = config.pop("path", _gen_unix_socket_default_path(process_index))
+
+        dirname = os.path.dirname(path)
+        if not os.path.exists(dirname):
+            os.makedirs(dirname, exist_ok=True)
+
+        if "start_serving" not in config:
+            config["start_serving"] = False
+
+        async def handle_connection(reader, writer):
+            # create a channel when client connected
+            return await server.on_connected(
+                reader, writer, local_address=server.address
+            )
+
+        aio_server = await asyncio.start_unix_server(
+            handle_connection, path=path, **config
+        )
+
+        for sock in aio_server.sockets:
+            sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, True)
+
+        server = UnixSocketServer(
+            process_index, aio_server, path, channel_handler=handle_channel
+        )
+        return server
+
+    @implements(Server.stop)
+    async def stop(self):
+        await super().stop()
+        try:
+            os.remove(self.path)
+        except OSError:  # pragma: no cover
+            pass
+
+
+@register_client
+class UnixSocketClient(Client):
+    __slots__ = ()
+
+    scheme = UnixSocketServer.scheme
+
+    @staticmethod
+    @lru_cache(100)
+    def _get_process_index(addr):
+        return int(urlparse(addr).path.lstrip("/"))
+
+    @staticmethod
+    @implements(Client.connect)
+    async def connect(
+        dest_address: str, local_address: str | None = None, **kwargs
+    ) -> "Client":
+        process_index = UnixSocketClient._get_process_index(dest_address)
+        path = kwargs.pop("path", _gen_unix_socket_default_path(process_index))
+        try:
+            (reader, writer) = await asyncio.open_unix_connection(path, **kwargs)
+        except FileNotFoundError:
+            raise ConnectionRefusedError(
+                "Cannot connect unix socket due to file not exists"
+            )
+        channel = SocketChannel(
+            reader,
+            writer,
+            local_address=local_address,
+            dest_address=dest_address,
+            channel_type=ChannelType.ipc,
+        )
+        return UnixSocketClient(local_address, dest_address, channel)