xoscar 0.9.0__cp312-cp312-macosx_10_13_x86_64.whl

xoscar/backends/core.py

@@ -0,0 +1,352 @@
+# 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 atexit
+import copy
+import logging
+import threading
+import weakref
+from typing import Type, Union
+
+from .._utils import Timer
+from ..errors import ServerClosed
+from ..profiling import get_profiling_data
+from .communication import ChannelType, Client, UCXClient
+from .message import (
+    DeserializeMessageFailed,
+    ErrorMessage,
+    ForwardMessage,
+    MessageType,
+    ResultMessage,
+    _MessageBase,
+)
+from .router import Router
+
+ResultMessageType = Union[ResultMessage, ErrorMessage]
+logger = logging.getLogger(__name__)
+
+
+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]
+
+    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,
+        proxy_addresses: list[str] | None = None,
+    ) -> Client:
+        client = await router.get_client(
+            dest_address, from_who=self, proxy_addresses=proxy_addresses
+        )
+        self._listen_client(client)
+        return client
+
+    async def _listen(self, client: Client):
+        try:
+            while not client.closed:
+                try:
+                    try:
+                        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
+                    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))
+        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
+    ) -> 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,
+        proxy_addresses: list[str] | None = None,
+    ) -> ResultMessage | ErrorMessage | asyncio.Future:
+        client = await self.get_client(
+            router, dest_address, proxy_addresses=proxy_addresses
+        )
+        if (
+            client.channel_type == ChannelType.remote
+            and client.dest_address != dest_address
+            and message.message_type != MessageType.control
+        ):
+            # wrap message with forward message
+            message = ForwardMessage(
+                message_id=message.message_id, address=dest_address, raw_message=message
+            )
+        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
+        try:
+            self.cancel_tasks()
+        except:
+            pass
+
+    def cancel_tasks(self):
+        # cancel listening for all clients
+        _ = [task.cancel() for task in self._clients.values()]
+
+
+def _cancel_all_tasks(loop):
+    to_cancel = asyncio.all_tasks(loop)
+    if not to_cancel:
+        return
+
+    for task in to_cancel:
+        task.cancel()
+
+    # In Python 3.13+, we need to use a different approach to avoid deadlocks
+    # when shutting down event loops in threads
+    if hasattr(asyncio, "run"):
+        # For Python 3.13+, use a more robust approach
+        async def _gather_cancelled():
+            await asyncio.gather(*to_cancel, return_exceptions=True)
+
+        try:
+            # Try to run the gather in the current loop context
+            if loop.is_running():
+                # If loop is running, schedule the gather
+                asyncio.run_coroutine_threadsafe(_gather_cancelled(), loop)
+            else:
+                # If loop is not running, we can run it directly
+                loop.run_until_complete(_gather_cancelled())
+        except RuntimeError:
+            # If we can't run the gather, just log and continue
+            logger.debug("Could not gather cancelled tasks during shutdown")
+    else:
+        # For older Python versions, use the original approach
+        loop.run_until_complete(asyncio.gather(*to_cancel, return_exceptions=True))
+
+    for task in to_cancel:
+        if task.cancelled():
+            continue
+        if task.exception() is not None:
+            loop.call_exception_handler(
+                {
+                    "message": "unhandled exception during asyncio.run() shutdown",
+                    "exception": task.exception(),
+                    "task": task,
+                }
+            )
+
+
+def _safe_run_forever(loop):
+    try:
+        loop.run_forever()
+    finally:
+        try:
+            _cancel_all_tasks(loop)
+        except Exception as e:
+            logger.debug("Error during task cancellation: %s", e)
+        finally:
+            try:
+                loop.stop()
+            except Exception as e:
+                logger.debug("Error stopping loop: %s", e)
+
+
+class ActorCaller:
+    __slots__ = "_thread_local"
+
+    class _RefHolder:
+        pass
+
+    _close_loop = None
+    _close_thread = None
+    _initialized = False
+
+    @classmethod
+    def _ensure_initialized(cls):
+        if not cls._initialized:
+            cls._close_loop = asyncio.new_event_loop()
+            cls._close_thread = threading.Thread(
+                target=_safe_run_forever, args=(cls._close_loop,), daemon=True
+            )
+            cls._close_thread.start()
+            atexit.register(cls._cleanup)
+            cls._initialized = True
+
+    @classmethod
+    def _cleanup(cls):
+        if cls._close_loop and cls._close_loop.is_running():
+            try:
+                cls._close_loop.call_soon_threadsafe(cls._close_loop.stop)
+                # Give the loop a moment to stop
+                if cls._close_thread:
+                    cls._close_thread.join(timeout=0.5)  # Shorter timeout for tests
+            except Exception as e:
+                logger.debug("Error during cleanup: %s", e)
+
+    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():
+                self._ensure_initialized()
+                # Use the background thread for 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/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/__main__.py

@@ -0,0 +1,19 @@
+if __name__ == "__main__":
+    import click
+
+    @click.group(
+        invoke_without_command=True,
+        name="xoscar",
+        help="Xoscar command-line interface.",
+    )
+    def main():
+        pass
+
+    @main.command("start_sub_pool", help="Start a sub pool.")
+    @click.option("shm_name", "-sn", type=str, help="Shared memory name.")
+    def start_sub_pool(shm_name):
+        from xoscar.backends.indigen.pool import MainActorPool
+
+        MainActorPool._start_sub_pool_in_child(shm_name)
+
+    main()

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

xoscar/backends/indigen/fate_sharing.py

@@ -0,0 +1,221 @@
+# Copyright 2017 The Ray Authors.
+#
+# 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.
+
+import asyncio
+import logging
+import subprocess
+import sys
+
+import psutil
+
+# Linux can bind child processes' lifetimes to that of their parents via prctl.
+# prctl support is detected dynamically once, and assumed thereafter.
+linux_prctl = None
+
+# Windows can bind processes' lifetimes to that of kernel-level "job objects".
+# We keep a global job object to tie its lifetime to that of our own process.
+win32_job = None
+win32_AssignProcessToJobObject = None
+
+logger = logging.getLogger(__name__)
+
+
+def detect_fate_sharing_support_win32():
+    global win32_job, win32_AssignProcessToJobObject
+    if win32_job is None and sys.platform == "win32":
+        import ctypes
+
+        try:
+            from ctypes.wintypes import BOOL, DWORD, HANDLE, LPCWSTR, LPVOID
+
+            kernel32 = ctypes.WinDLL("kernel32")
+            kernel32.CreateJobObjectW.argtypes = (LPVOID, LPCWSTR)
+            kernel32.CreateJobObjectW.restype = HANDLE
+            sijo_argtypes = (HANDLE, ctypes.c_int, LPVOID, DWORD)
+            kernel32.SetInformationJobObject.argtypes = sijo_argtypes
+            kernel32.SetInformationJobObject.restype = BOOL
+            kernel32.AssignProcessToJobObject.argtypes = (HANDLE, HANDLE)
+            kernel32.AssignProcessToJobObject.restype = BOOL
+            kernel32.IsDebuggerPresent.argtypes = ()
+            kernel32.IsDebuggerPresent.restype = BOOL
+        except (AttributeError, TypeError, ImportError):
+            kernel32 = None
+        job = kernel32.CreateJobObjectW(None, None) if kernel32 else None
+        job = subprocess.Handle(job) if job else job
+        if job:
+            from ctypes.wintypes import DWORD, LARGE_INTEGER, ULARGE_INTEGER
+
+            class JOBOBJECT_BASIC_LIMIT_INFORMATION(ctypes.Structure):
+                _fields_ = [
+                    ("PerProcessUserTimeLimit", LARGE_INTEGER),
+                    ("PerJobUserTimeLimit", LARGE_INTEGER),
+                    ("LimitFlags", DWORD),
+                    ("MinimumWorkingSetSize", ctypes.c_size_t),
+                    ("MaximumWorkingSetSize", ctypes.c_size_t),
+                    ("ActiveProcessLimit", DWORD),
+                    ("Affinity", ctypes.c_size_t),
+                    ("PriorityClass", DWORD),
+                    ("SchedulingClass", DWORD),
+                ]
+
+            class IO_COUNTERS(ctypes.Structure):
+                _fields_ = [
+                    ("ReadOperationCount", ULARGE_INTEGER),
+                    ("WriteOperationCount", ULARGE_INTEGER),
+                    ("OtherOperationCount", ULARGE_INTEGER),
+                    ("ReadTransferCount", ULARGE_INTEGER),
+                    ("WriteTransferCount", ULARGE_INTEGER),
+                    ("OtherTransferCount", ULARGE_INTEGER),
+                ]
+
+            class JOBOBJECT_EXTENDED_LIMIT_INFORMATION(ctypes.Structure):
+                _fields_ = [
+                    ("BasicLimitInformation", JOBOBJECT_BASIC_LIMIT_INFORMATION),
+                    ("IoInfo", IO_COUNTERS),
+                    ("ProcessMemoryLimit", ctypes.c_size_t),
+                    ("JobMemoryLimit", ctypes.c_size_t),
+                    ("PeakProcessMemoryUsed", ctypes.c_size_t),
+                    ("PeakJobMemoryUsed", ctypes.c_size_t),
+                ]
+
+            debug = kernel32.IsDebuggerPresent()
+
+            # Defined in <WinNT.h>; also available here:
+            # https://docs.microsoft.com/en-us/windows/win32/api/jobapi2/nf-jobapi2-setinformationjobobject
+            JobObjectExtendedLimitInformation = 9
+            JOB_OBJECT_LIMIT_BREAKAWAY_OK = 0x00000800
+            JOB_OBJECT_LIMIT_DIE_ON_UNHANDLED_EXCEPTION = 0x00000400
+            JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE = 0x00002000
+            buf = JOBOBJECT_EXTENDED_LIMIT_INFORMATION()
+            buf.BasicLimitInformation.LimitFlags = (
+                (0 if debug else JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE)
+                | JOB_OBJECT_LIMIT_DIE_ON_UNHANDLED_EXCEPTION
+                | JOB_OBJECT_LIMIT_BREAKAWAY_OK
+            )
+            infoclass = JobObjectExtendedLimitInformation
+            if not kernel32.SetInformationJobObject(
+                job, infoclass, ctypes.byref(buf), ctypes.sizeof(buf)
+            ):
+                job = None
+        win32_AssignProcessToJobObject = (
+            kernel32.AssignProcessToJobObject if kernel32 is not None else False
+        )
+        win32_job = job if job else False
+    return bool(win32_job)
+
+
+def detect_fate_sharing_support_linux():
+    global linux_prctl
+    if linux_prctl is None and sys.platform.startswith("linux"):
+        try:
+            from ctypes import CDLL, c_int, c_ulong
+
+            prctl = CDLL(None).prctl
+            prctl.restype = c_int
+            prctl.argtypes = [c_int, c_ulong, c_ulong, c_ulong, c_ulong]
+        except (AttributeError, TypeError):
+            prctl = None
+        linux_prctl = prctl if prctl else False
+    return bool(linux_prctl)
+
+
+def detect_fate_sharing_support():
+    result = None
+    if sys.platform == "win32":
+        result = detect_fate_sharing_support_win32()
+    elif sys.platform.startswith("linux"):
+        result = detect_fate_sharing_support_linux()
+    return result
+
+
+if detect_fate_sharing_support():
+    logger.info("Using kernel-level fate-sharing.")
+
+
+def set_kill_on_parent_death_linux():
+    """Ensures this process dies if its parent dies (fate-sharing).
+
+    Linux-only. Must be called in preexec_fn (i.e. by the child).
+    """
+    if detect_fate_sharing_support_linux():
+        import signal
+
+        PR_SET_PDEATHSIG = 1
+        if linux_prctl(PR_SET_PDEATHSIG, signal.SIGKILL, 0, 0, 0) != 0:
+            import ctypes
+
+            raise OSError(ctypes.get_errno(), "prctl(PR_SET_PDEATHSIG) failed")
+    else:
+        assert False, "PR_SET_PDEATHSIG used despite being unavailable"
+
+
+def set_kill_child_on_death_win32(child_proc):
+    """Ensures the child process dies if this process dies (fate-sharing).
+
+    Windows-only. Must be called by the parent, after spawning the child.
+
+    Args:
+        child_proc: The subprocess.Popen or subprocess.Handle object.
+    """
+
+    if isinstance(child_proc, subprocess.Popen):
+        child_proc = child_proc._handle
+    assert isinstance(child_proc, subprocess.Handle)
+
+    if detect_fate_sharing_support_win32():
+        if not win32_AssignProcessToJobObject(win32_job, int(child_proc)):
+            import ctypes
+
+            raise OSError(ctypes.get_last_error(), "AssignProcessToJobObject() failed")
+    else:
+        assert False, "AssignProcessToJobObject used despite being unavailable"
+
+
+def preexec_fn():
+    import signal
+
+    signal.pthread_sigmask(signal.SIG_BLOCK, {signal.SIGINT})
+    if sys.platform.startswith("linux"):
+        set_kill_on_parent_death_linux()
+
+
+async def create_subprocess_exec(*args, **kwargs):
+    win32 = sys.platform == "win32"
+    # With Windows fate-sharing, we need special care:
+    # The process must be added to the job before it is allowed to execute.
+    # Otherwise, there's a race condition: the process might spawn children
+    # before the process itself is assigned to the job.
+    # After that point, its children will not be added to the job anymore.
+    CREATE_SUSPENDED = 0x00000004  # from Windows headers
+    creationflags = CREATE_SUSPENDED if win32 else 0
+    if win32 and detect_fate_sharing_support_win32():
+        creationflags |= subprocess.CREATE_NEW_PROCESS_GROUP
+    # CREATE_NEW_PROCESS_GROUP is used to send Ctrl+C on Windows:
+    # https://docs.python.org/3/library/subprocess.html#subprocess.Popen.send_signal
+    process: asyncio.subprocess.Process = await asyncio.create_subprocess_exec(
+        *args,
+        **kwargs,
+        preexec_fn=preexec_fn if not win32 else None,
+        creationflags=creationflags,
+    )
+    if win32:
+        proc = process._transport._proc
+        try:
+            set_kill_child_on_death_win32(proc)
+            psutil.Process(process.pid).resume()
+        except (psutil.Error, OSError):
+            logger.exception("Resume process failed, kill %s.", process.pid)
+            process.kill()
+            raise
+    return process