xoscar 0.8.0__cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.whl

xoscar/backends/indigen/pool.py

@@ -0,0 +1,509 @@
+# 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 asyncio.subprocess
+import configparser
+import itertools
+import logging.config
+import os
+import pickle
+import random
+import signal
+import struct
+import sys
+import threading
+import time
+import uuid
+from enum import IntEnum
+from typing import List, Optional
+
+import psutil
+
+from ..._utils import reset_id_random_seed
+from ...utils import ensure_coverage
+from ..config import ActorPoolConfig
+from ..message import (
+    ControlMessage,
+    ControlMessageType,
+    CreateActorMessage,
+    new_message_id,
+)
+from ..pool import MainActorPoolBase, SubActorPoolBase, _register_message_handler
+from . import shared_memory
+from .fate_sharing import create_subprocess_exec
+
+_SUBPROCESS_SHM_SIZE = 10240
+_is_windows: bool = sys.platform.startswith("win")
+
+logger = logging.getLogger(__name__)
+
+
+class _ShmSeq(IntEnum):
+    INIT_PARAMS = 1
+    INIT_RESULT = 2
+
+
+def _shm_put_object(seq: _ShmSeq, shm: shared_memory.SharedMemory, o: object):
+    serialized = pickle.dumps(o)
+    assert (
+        len(serialized) < _SUBPROCESS_SHM_SIZE - 8
+    ), f"Serialized object {o} is too long."
+    shm.buf[4:8] = struct.pack("<I", len(serialized))
+    shm.buf[8 : 8 + len(serialized)] = serialized
+    shm.buf[:4] = struct.pack("<I", seq)
+
+
+def _shm_get_object(seq: _ShmSeq, shm: shared_memory.SharedMemory):
+    recv_seq = struct.unpack("<I", shm.buf[:4])[0]
+    if recv_seq != seq:
+        return
+    size = struct.unpack("<I", shm.buf[4:8])[0]
+    return pickle.loads(shm.buf[8 : 8 + size])
+
+
+@_register_message_handler
+class MainActorPool(MainActorPoolBase):
+    @classmethod
+    def get_external_addresses(
+        cls,
+        address: str,
+        n_process: int | None = None,
+        ports: list[int] | None = None,
+        schemes: list[Optional[str]] | None = None,
+    ):
+        """Get external address for every process"""
+        assert n_process is not None
+        if ":" in address:
+            host, port_str = address.rsplit(":", 1)
+            port = int(port_str)
+            if ports:
+                if len(ports) != n_process:
+                    raise ValueError(
+                        f"`ports` specified, but its count "
+                        f"is not equal to `n_process`, "
+                        f"number of ports: {len(ports)}, "
+                        f"n_process: {n_process}"
+                    )
+                sub_ports = ports
+            else:
+                sub_ports = [0] * n_process
+        else:
+            host = address
+            if ports and len(ports) != n_process + 1:
+                # ports specified, the first of which should be main port
+                raise ValueError(
+                    f"`ports` specified, but its count "
+                    f"is not equal to `n_process` + 1, "
+                    f"number of ports: {len(ports)}, "
+                    f"n_process + 1: {n_process + 1}"
+                )
+            elif not ports:
+                ports = [0] * (n_process + 1)
+            port = ports[0]
+            sub_ports = ports[1:]
+        if not schemes:
+            prefix_iter = itertools.repeat("")
+        else:
+            prefix_iter = [f"{scheme}://" if scheme else "" for scheme in schemes]  # type: ignore
+        return [
+            f"{prefix}{host}:{port}"
+            for port, prefix in zip([port] + sub_ports, prefix_iter)
+        ]
+
+    @classmethod
+    def gen_internal_address(
+        cls, process_index: int, external_address: str | None = None
+    ) -> str | None:
+        if hasattr(asyncio, "start_unix_server"):
+            return f"unixsocket:///{process_index}"
+        else:
+            return external_address
+
+    @classmethod
+    async def start_sub_pool(
+        cls,
+        actor_pool_config: ActorPoolConfig,
+        process_index: int,
+        start_python: str | None = None,
+    ):
+        return await cls._create_sub_pool_from_parent(
+            actor_pool_config, process_index, start_python
+        )
+
+    @classmethod
+    async def wait_sub_pools_ready(cls, create_pool_tasks: List[asyncio.Task]):
+        processes: list[asyncio.subprocess.Process] = []
+        ext_addresses = []
+        error = None
+        for task in create_pool_tasks:
+            process, address = await task
+            processes.append(process)
+            ext_addresses.append(address)
+        if error:
+            for p in processes:
+                # error happens, kill all subprocesses
+                p.kill()
+            raise error
+        return processes, ext_addresses
+
+    @classmethod
+    def _start_sub_pool_in_child(
+        cls,
+        shm_name: str,
+    ):
+        ensure_coverage()
+
+        shm = shared_memory.SharedMemory(shm_name, track=False)
+        try:
+            config = _shm_get_object(_ShmSeq.INIT_PARAMS, shm)
+            # Check Python version once.
+            sub_pool_python_version = config.pop("python_version", None)
+            if (
+                sub_pool_python_version is not None
+                and sub_pool_python_version != sys.hexversion
+            ):
+                logger.warning(
+                    f"The sub pool is using a different Python version, you may encounter serialization issues."
+                    f" sub pool: {sub_pool_python_version}, main pool: {sys.hexversion}"
+                )
+            actor_config = config["actor_pool_config"]
+            process_index = config["process_index"]
+            main_pool_pid = config["main_pool_pid"]
+
+            def _check_ppid():
+                while True:
+                    try:
+                        # We can't simply check if the os.getppid() equals with main_pool_pid,
+                        # as the double fork may result in a new process as the parent.
+                        psutil.Process(main_pool_pid)
+                    except psutil.NoSuchProcess:
+                        logger.error("Exit due to main pool %s exit.", main_pool_pid)
+                        os._exit(233)  # Special exit code for debugging.
+                    except Exception as e:
+                        logger.exception("Check ppid failed: %s", e)
+                    time.sleep(10)
+
+            t = threading.Thread(target=_check_ppid, daemon=True)
+            t.start()
+
+            # make sure enough randomness for every sub pool
+            random.seed(uuid.uuid1().bytes)
+            reset_id_random_seed()
+
+            conf = actor_config.get_pool_config(process_index)
+            suspend_sigint = conf["suspend_sigint"]
+            if suspend_sigint:
+                signal.signal(signal.SIGINT, lambda *_: None)
+
+            logging_conf = conf["logging_conf"] or {}
+            if isinstance(logging_conf, configparser.RawConfigParser):
+                logging.config.fileConfig(logging_conf)
+            elif logging_conf.get("dict"):
+                logging.config.dictConfig(logging_conf["dict"])
+            elif logging_conf.get("file"):
+                logging.config.fileConfig(logging_conf["file"])
+            elif logging_conf.get("level"):
+                logging.getLogger("__main__").setLevel(logging_conf["level"])
+                logging.getLogger("xoscar").setLevel(logging_conf["level"])
+                if logging_conf.get("format"):
+                    logging.basicConfig(format=logging_conf["format"])
+
+            use_uvloop = conf["use_uvloop"]
+            if use_uvloop:
+                import uvloop
+
+                asyncio.set_event_loop(uvloop.new_event_loop())
+            else:
+                asyncio.set_event_loop(asyncio.new_event_loop())
+
+            coro = cls._create_sub_pool(actor_config, process_index, main_pool_pid, shm)
+            asyncio.run(coro)
+        finally:
+            shm.close()
+
+    @classmethod
+    async def _create_sub_pool(
+        cls,
+        actor_config: ActorPoolConfig,
+        process_index: int,
+        main_pool_pid: int,
+        shm: shared_memory.SharedMemory,
+    ):
+        cur_pool_config = actor_config.get_pool_config(process_index)
+        env = cur_pool_config["env"]
+        if env:
+            os.environ.update(env)
+        pool = await SubActorPool.create(
+            {
+                "actor_pool_config": actor_config,
+                "process_index": process_index,
+                "main_pool_pid": main_pool_pid,
+            }
+        )
+        await pool.start()
+        _shm_put_object(_ShmSeq.INIT_RESULT, shm, cur_pool_config["external_address"])
+        await pool.join()
+
+    @staticmethod
+    async def _create_sub_pool_from_parent(
+        actor_pool_config: ActorPoolConfig,
+        process_index: int,
+        start_python: str | None = None,
+    ):
+        # We check the Python version in _shm_get_object to make it faster,
+        # as in most cases the Python versions are the same.
+        if start_python is None:
+            start_python = sys.executable
+
+        external_addresses: List | None = None
+        shm = shared_memory.SharedMemory(
+            create=True, size=_SUBPROCESS_SHM_SIZE, track=False
+        )
+        try:
+            _shm_put_object(
+                _ShmSeq.INIT_PARAMS,
+                shm,
+                {
+                    "actor_pool_config": actor_pool_config,
+                    "process_index": process_index,
+                    "main_pool_pid": os.getpid(),
+                    "python_version": sys.hexversion,
+                },
+            )
+            cmd = [
+                start_python,
+                "-m",
+                "xoscar.backends.indigen",
+                "start_sub_pool",
+                "-sn",
+                shm.name,
+            ]
+            # We need to inherit the parent environment to ensure the subprocess works correctly on Windows.
+            new_env = dict(os.environ)
+            env = actor_pool_config.get_pool_config(process_index).get("env") or {}
+            new_env.update(env)
+            logger.info("Creating sub pool via command: %s", cmd)
+            process = await create_subprocess_exec(*cmd, env=new_env)
+
+            def _get_external_addresses():
+                try:
+                    nonlocal external_addresses
+                    while (
+                        shm
+                        and shm.buf is not None
+                        and not (
+                            external_addresses := _shm_get_object(
+                                _ShmSeq.INIT_RESULT, shm
+                            )
+                        )
+                    ):
+                        time.sleep(0.1)
+                except asyncio.CancelledError:
+                    pass
+
+            _, unfinished = await asyncio.wait(
+                [
+                    asyncio.create_task(process.wait()),
+                    asyncio.create_task(asyncio.to_thread(_get_external_addresses)),
+                ],
+                return_when=asyncio.FIRST_COMPLETED,
+            )
+            for t in unfinished:
+                t.cancel()
+        finally:
+            shm.close()
+            shm.unlink()
+        if external_addresses is None:
+            raise OSError(f"Start sub pool failed, returncode: {process.returncode}")
+        return process, external_addresses
+
+    async def append_sub_pool(
+        self,
+        label: str | None = None,
+        internal_address: str | None = None,
+        external_address: str | None = None,
+        env: dict | None = None,
+        modules: list[str] | None = None,
+        suspend_sigint: bool | None = None,
+        use_uvloop: bool | None = None,
+        logging_conf: dict | None = None,
+        start_python: str | None = None,
+        kwargs: dict | None = None,
+    ):
+        # external_address has port 0, subprocess will bind random port.
+        external_address = (
+            external_address
+            or MainActorPool.get_external_addresses(self.external_address, n_process=1)[
+                -1
+            ]
+        )
+
+        # use last process index's logging_conf and use_uv_loop config if not provide
+        actor_pool_config = self._config.as_dict()
+        last_process_index = self._config.get_process_indexes()[-1]
+        last_logging_conf = actor_pool_config["pools"][last_process_index][
+            "logging_conf"
+        ]
+        last_use_uv_loop = actor_pool_config["pools"][last_process_index]["use_uvloop"]
+        _logging_conf = logging_conf or last_logging_conf
+        _use_uv_loop = use_uvloop if use_uvloop is not None else last_use_uv_loop
+
+        process_index = next(MainActorPool.process_index_gen(external_address))
+        internal_address = internal_address or MainActorPool.gen_internal_address(
+            process_index, external_address
+        )
+
+        self._config.add_pool_conf(
+            process_index,
+            label,
+            internal_address,
+            external_address,
+            env,
+            modules,
+            suspend_sigint,
+            _use_uv_loop,
+            _logging_conf,
+            kwargs,
+        )
+
+        process, external_addresses = await self._create_sub_pool_from_parent(
+            self._config, process_index, start_python
+        )
+
+        self._config.reset_pool_external_address(process_index, external_addresses[0])
+        self.attach_sub_process(external_addresses[0], process)
+
+        control_message = ControlMessage(
+            message_id=new_message_id(),
+            address=self.external_address,
+            control_message_type=ControlMessageType.sync_config,
+            content=self._config,
+        )
+        await self.handle_control_command(control_message)
+        # The actual port will return in process_status.
+        return external_addresses[0]
+
+    async def remove_sub_pool(
+        self, external_address: str, timeout: float | None = None, force: bool = False
+    ):
+        process = self.sub_processes[external_address]
+        process_index = self._config.get_process_index(external_address)
+        del self.sub_processes[external_address]
+        self._config.remove_pool_config(process_index)
+        await self.stop_sub_pool(external_address, process, timeout, force)
+
+        control_message = ControlMessage(
+            message_id=new_message_id(),
+            address=self.external_address,
+            control_message_type=ControlMessageType.sync_config,
+            content=self._config,
+        )
+        await self.handle_control_command(control_message)
+
+    async def kill_sub_pool(
+        self, process: asyncio.subprocess.Process, force: bool = False
+    ):
+        # First, try to terminate the process gracefully
+        if not force:
+            try:
+                process.terminate()
+                # Wait for graceful termination
+                try:
+                    await asyncio.wait_for(process.wait(), timeout=2.0)
+                except asyncio.TimeoutError:
+                    # Process didn't terminate gracefully, force kill
+                    force = True
+            except ProcessLookupError:
+                # Process already terminated
+                pass
+
+        # Force kill if needed or if graceful termination failed
+        if force:
+            try:
+                process.kill()
+            except ProcessLookupError:
+                # Process already dead
+                pass
+
+        # Ensure process is completely terminated and cleaned up
+        try:
+            # Wait for process to complete
+            if process.returncode is None:
+                try:
+                    await asyncio.wait_for(process.wait(), timeout=5.0)
+                except asyncio.TimeoutError:
+                    pass
+        except ProcessLookupError:
+            # Process already terminated
+            pass
+
+        # Python 3.13 specific cleanup for waitpid threads
+        if sys.version_info >= (3, 13):
+            try:
+                # Close the transport to clean up waitpid thread
+                if hasattr(process, "_transport") and process._transport:
+                    process._transport.close()
+                # Also try to close the pipe transport if it exists
+                if hasattr(process, "_pipes") and process._pipes:
+                    for pipe in process._pipes.values():
+                        if hasattr(pipe, "close"):
+                            pipe.close()
+            except Exception:
+                # Ignore errors during cleanup
+                pass
+
+        # Additional cleanup using psutil to ensure process tree is terminated
+        try:
+            p = psutil.Process(process.pid)
+            if p.is_running():
+                # Kill the entire process tree
+                for child in p.children(recursive=True):
+                    try:
+                        child.kill()
+                    except psutil.NoSuchProcess:
+                        pass
+                p.kill()
+                p.wait(timeout=2.0)
+        except (psutil.NoSuchProcess, psutil.TimeoutExpired):
+            # Process already dead or couldn't be killed
+            pass
+
+    async def is_sub_pool_alive(self, process: asyncio.subprocess.Process):
+        return process.returncode is None
+
+    async def recover_sub_pool(self, address: str):
+        process_index = self._config.get_process_index(address)
+        # process dead, restart it
+        # remember always use spawn to recover sub pool
+        task = asyncio.create_task(self.start_sub_pool(self._config, process_index))
+        self.sub_processes[address] = (await self.wait_sub_pools_ready([task]))[0][0]
+
+        if self._auto_recover == "actor":
+            # need to recover all created actors
+            for _, message in self._allocated_actors[address].values():
+                create_actor_message: CreateActorMessage = message  # type: ignore
+                await self.call(address, create_actor_message)
+
+    async def start(self):
+        await super().start()
+        await self.start_monitor()
+
+
+@_register_message_handler
+class SubActorPool(SubActorPoolBase):
+    pass