xoscar 0.4.0__cp312-cp312-win_amd64.whl

xoscar/backends/core.py

@@ -0,0 +1,193 @@
+# 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 copy
+import logging
+from typing import Type, Union
+
+from .._utils import Timer
+from ..errors import ServerClosed
+from ..profiling import get_profiling_data
+from .communication import Client, UCXClient
+from .message import DeserializeMessageFailed, ErrorMessage, ResultMessage, _MessageBase
+from .router import Router
+
+ResultMessageType = Union[ResultMessage, ErrorMessage]
+logger = logging.getLogger(__name__)
+
+
+class ActorCaller:
+    __slots__ = "_client_to_message_futures", "_clients", "_profiling_data"
+
+    _client_to_message_futures: dict[Client, dict[bytes, asyncio.Future]]
+    _clients: dict[Client, asyncio.Task]
+
+    def __init__(self):
+        self._client_to_message_futures = dict()
+        self._clients = dict()
+        self._profiling_data = get_profiling_data()
+
+    def _listen_client(self, client: Client):
+        if client not in self._clients:
+            self._clients[client] = asyncio.create_task(self._listen(client))
+            self._client_to_message_futures[client] = dict()
+            client_count = len(self._clients)
+            if client_count >= 100:  # pragma: no cover
+                if (client_count - 100) % 10 == 0:  # pragma: no cover
+                    logger.warning(
+                        "Actor caller has created too many clients (%s >= 100), "
+                        "the global router may not be set.",
+                        client_count,
+                    )
+
+    async def get_client_via_type(
+        self, router: Router, dest_address: str, client_type: Type[Client]
+    ) -> Client:
+        client = await router.get_client_via_type(
+            dest_address, client_type, from_who=self
+        )
+        self._listen_client(client)
+        return client
+
+    async def get_client(self, router: Router, dest_address: str) -> Client:
+        client = await router.get_client(dest_address, from_who=self)
+        self._listen_client(client)
+        return client
+
+    async def _listen(self, client: Client):
+        while not client.closed:
+            try:
+                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
+                        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)
+
+        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))
+
+    async def call_with_client(
+        self, client: Client, message: _MessageBase, wait: bool = True
+    ) -> ResultMessage | ErrorMessage | asyncio.Future:
+        loop = asyncio.get_running_loop()
+        wait_response = loop.create_future()
+        self._client_to_message_futures[client][message.message_id] = wait_response
+
+        with Timer() as timer:
+            try:
+                await client.send(message)
+            except ConnectionError:
+                try:
+                    await client.close()
+                except ConnectionError:
+                    # close failed, ignore it
+                    pass
+                raise ServerClosed(f"Remote server {client.dest_address} closed")
+
+            if not wait:
+                r = wait_response
+            else:
+                r = await wait_response
+
+        self._profiling_data.collect_actor_call(message, timer.duration)
+        return r
+
+    async def call_send_buffers(
+        self,
+        client: UCXClient,
+        local_buffers: list,
+        meta_message: _MessageBase,
+        wait: bool = True,
+    ) -> ResultMessage | ErrorMessage | asyncio.Future:
+        loop = asyncio.get_running_loop()
+        wait_response = loop.create_future()
+        self._client_to_message_futures[client][meta_message.message_id] = wait_response
+
+        with Timer() as timer:
+            try:
+                await client.send_buffers(local_buffers, meta_message)
+            except ConnectionError:  # pragma: no cover
+                try:
+                    await client.close()
+                except ConnectionError:
+                    # close failed, ignore it
+                    pass
+                raise ServerClosed(f"Remote server {client.dest_address} closed")
+
+            if not wait:  # pragma: no cover
+                r = wait_response
+            else:
+                r = await wait_response
+
+        self._profiling_data.collect_actor_call(meta_message, timer.duration)
+        return r
+
+    async def call(
+        self,
+        router: Router,
+        dest_address: str,
+        message: _MessageBase,
+        wait: bool = True,
+    ) -> ResultMessage | ErrorMessage | asyncio.Future:
+        client = await self.get_client(router, dest_address)
+        return await self.call_with_client(client, message, wait)
+
+    async def stop(self):
+        try:
+            await asyncio.gather(*[client.close() for client in self._clients])
+        except (ConnectionError, ServerClosed):
+            pass
+        self.cancel_tasks()
+
+    def cancel_tasks(self):
+        # cancel listening for all clients
+        _ = [task.cancel() for task in self._clients.values()]

xoscar/backends/indigen/__init__.py

@@ -0,0 +1,16 @@
+# 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 .backend import IndigenActorBackend

xoscar/backends/indigen/backend.py

@@ -0,0 +1,51 @@
+# 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
+
+from ...backend import BaseActorBackend, register_backend
+from ..context import IndigenActorContext
+from .driver import IndigenActorDriver
+from .pool import MainActorPool
+
+__all__ = ["IndigenActorBackend"]
+
+
+@register_backend
+class IndigenActorBackend(BaseActorBackend):
+    @staticmethod
+    def name():
+        # None means Indigen is default scheme
+        # ucx can be recognized as Indigen backend as well
+        return [None, "ucx"]
+
+    @staticmethod
+    def get_context_cls():
+        return IndigenActorContext
+
+    @staticmethod
+    def get_driver_cls():
+        return IndigenActorDriver
+
+    @classmethod
+    async def create_actor_pool(
+        cls, address: str, n_process: int | None = None, **kwargs
+    ):
+        from ..pool import create_actor_pool
+
+        assert n_process is not None
+        return await create_actor_pool(
+            address, pool_cls=MainActorPool, n_process=n_process, **kwargs
+        )

xoscar/backends/indigen/driver.py

@@ -0,0 +1,26 @@
+# 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 numbers import Number
+from typing import Dict
+
+from ...driver import BaseActorDriver
+
+
+class IndigenActorDriver(BaseActorDriver):
+    @classmethod
+    def setup_cluster(cls, address_to_resources: Dict[str, Dict[str, Number]]):
+        # nothing need to be done in driver of Indigen backend
+        pass