ENH: stop using daemon=True for subpool (#85)

.pre-commit-config.yaml

@@ -1,6 +1,6 @@
 repos:
   - repo: https://github.com/psf/black
-    rev: 23.1.0
+    rev: 24.2.0
     hooks:
       - id: black
         files: python/xoscar
@@ -24,7 +24,7 @@ repos:
         args: [--sp, python/setup.cfg]
         files: python/xoscar
   - repo: https://github.com/pre-commit/mirrors-mypy
-    rev: v1.4.1
+    rev: v1.9.0
     hooks:
       - id: mypy
         additional_dependencies: [tokenize-rt==3.2.0]

python/xoscar/backends/communication/dummy.py

@@ -97,9 +97,9 @@ class DummyServer(Server):
         else tuple()
     )
 
-    _address_to_instances: weakref.WeakValueDictionary[
-        str, "DummyServer"
-    ] = weakref.WeakValueDictionary()
+    _address_to_instances: weakref.WeakValueDictionary[str, "DummyServer"] = (
+        weakref.WeakValueDictionary()
+    )
     _channels: list[ChannelType]
     _tasks: list[asyncio.Task]
     scheme: str | None = "dummy"

python/xoscar/backends/indigen/pool.py

@@ -187,14 +187,19 @@ async def start_sub_pool(
         def start_pool_in_process():
             ctx = multiprocessing.get_context(method=start_method)
             status_queue = ctx.Queue()
+            main_pool_pid = os.getpid()
 
             with _suspend_init_main():
                 process = ctx.Process(
                     target=cls._start_sub_pool,
-                    args=(actor_pool_config, process_index, status_queue),
+                    args=(
+                        actor_pool_config,
+                        process_index,
+                        status_queue,
+                        main_pool_pid,
+                    ),
                     name=f"IndigenActorPool{process_index}",
                 )
-                process.daemon = True
                 process.start()
 
             # wait for sub actor pool to finish starting
@@ -209,15 +214,22 @@ def start_pool_in_process():
 
     @classmethod
     async def wait_sub_pools_ready(cls, create_pool_tasks: List[asyncio.Task]):
-        processes = []
+        processes: list[multiprocessing.Process] = []
         ext_addresses = []
+        error = None
         for task in create_pool_tasks:
             process, status = await task
+            processes.append(process)
             if status.status == 1:
                 # start sub pool failed
-                raise status.error.with_traceback(status.traceback)
-            processes.append(process)
-            ext_addresses.append(status.external_addresses)
+                error = status.error.with_traceback(status.traceback)
+            else:
+                ext_addresses.append(status.external_addresses)
+        if error:
+            for p in processes:
+                # error happens, kill all subprocesses
+                p.kill()
+            raise error
         return processes, ext_addresses
 
     @classmethod
@@ -226,6 +238,7 @@ def _start_sub_pool(
         actor_config: ActorPoolConfig,
         process_index: int,
         status_queue: multiprocessing.Queue,
+        main_pool_pid: int,
     ):
         ensure_coverage()
 
@@ -259,7 +272,9 @@ def _start_sub_pool(
         else:
             asyncio.set_event_loop(asyncio.new_event_loop())
 
-        coro = cls._create_sub_pool(actor_config, process_index, status_queue)
+        coro = cls._create_sub_pool(
+            actor_config, process_index, status_queue, main_pool_pid
+        )
         asyncio.run(coro)
 
     @classmethod
@@ -268,6 +283,7 @@ async def _create_sub_pool(
         actor_config: ActorPoolConfig,
         process_index: int,
         status_queue: multiprocessing.Queue,
+        main_pool_pid: int,
     ):
         process_status = None
         try:
@@ -276,7 +292,11 @@ async def _create_sub_pool(
             if env:
                 os.environ.update(env)
             pool = await SubActorPool.create(
-                {"actor_pool_config": actor_config, "process_index": process_index}
+                {
+                    "actor_pool_config": actor_config,
+                    "process_index": process_index,
+                    "main_pool_pid": main_pool_pid,
+                }
             )
             external_addresses = cur_pool_config["external_address"]
             process_status = SubpoolStatus(
@@ -342,14 +362,14 @@ async def append_sub_pool(
         def start_pool_in_process():
             ctx = multiprocessing.get_context(method=start_method)
             status_queue = ctx.Queue()
+            main_pool_pid = os.getpid()
 
             with _suspend_init_main():
                 process = ctx.Process(
                     target=self._start_sub_pool,
-                    args=(self._config, process_index, status_queue),
+                    args=(self._config, process_index, status_queue, main_pool_pid),
                     name=f"IndigenActorPool{process_index}",
                 )
-                process.daemon = True
                 process.start()
 
             # wait for sub actor pool to finish starting

python/xoscar/backends/indigen/tests/test_pool.py

@@ -17,12 +17,14 @@
 
 import asyncio
 import logging
+import multiprocessing
 import os
 import re
 import sys
 import time
 from unittest import mock
 
+import psutil
 import pytest
 
 from .... import Actor, create_actor, create_actor_ref, get_pool_config, kill_actor
@@ -1099,3 +1101,80 @@ def test():
         assert process_index not in config.get_process_indexes()
         with pytest.raises(KeyError):
             config.get_pool_config(process_index)
+
+
+async def _run(started: multiprocessing.Event):  # type: ignore
+    pool = await create_actor_pool(  # type: ignore
+        "127.0.0.1", pool_cls=MainActorPool, n_process=1
+    )
+
+    class DummyActor(Actor):
+        @staticmethod
+        def test():
+            return "this is dummy!"
+
+    ref = await create_actor(
+        DummyActor, address=pool.external_address, allocate_strategy=RandomSubPool()
+    )
+    assert ref is not None
+
+    started.set()  # type: ignore
+    await pool.join()
+
+
+def _run_in_process(started: multiprocessing.Event):  # type: ignore
+    asyncio.run(_run(started))
+
+
+@pytest.mark.asyncio
+async def test_sub_pool_quit_with_main_pool():
+    s = multiprocessing.Event()
+    p = multiprocessing.Process(target=_run_in_process, args=(s,))
+    p.start()
+    s.wait()
+
+    processes = psutil.Process(p.pid).children()
+    assert len(processes) == 1
+
+    # kill main process
+    p.kill()
+    p.join()
+    await asyncio.sleep(1)
+
+    # subprocess should have died
+    assert not psutil.pid_exists(processes[0].pid)
+
+
+def _add(x: int) -> int:
+    return x + 1
+
+
+class _ProcessActor(Actor):
+    def run(self, x: int):
+        p = multiprocessing.Process(target=_add, args=(x,))
+        p.start()
+        p.join()
+        return x + 1
+
+
+@pytest.mark.asyncio
+async def test_process_in_actor():
+    start_method = (
+        os.environ.get("POOL_START_METHOD", "forkserver")
+        if sys.platform != "win32"
+        else None
+    )
+    pool = await create_actor_pool(  # type: ignore
+        "127.0.0.1",
+        pool_cls=MainActorPool,
+        n_process=1,
+        subprocess_start_method=start_method,
+    )
+
+    async with pool:
+        ref = await create_actor(
+            _ProcessActor,
+            address=pool.external_address,
+            allocate_strategy=RandomSubPool(),
+        )
+        assert 2 == await ref.run(1)

python/xoscar/backends/pool.py

@@ -27,6 +27,8 @@
 from abc import ABC, ABCMeta, abstractmethod
 from typing import Any, Callable, Coroutine, Optional, Type, TypeVar
 
+import psutil
+
 from .._utils import TypeDispatcher, create_actor_ref, to_binary
 from ..api import Actor
 from ..core import ActorRef, BufferRef, FileObjectRef, register_local_pool
@@ -821,7 +823,8 @@ def handle_channel(channel):
 
 
 class SubActorPoolBase(ActorPoolBase):
-    __slots__ = ("_main_address",)
+    __slots__ = ("_main_address", "_watch_main_pool_task")
+    _watch_main_pool_task: Optional[asyncio.Task]
 
     def __init__(
         self,
@@ -834,6 +837,7 @@ def __init__(
         config: ActorPoolConfig,
         servers: list[Server],
         main_address: str,
+        main_pool_pid: Optional[int],
     ):
         super().__init__(
             process_index,
@@ -846,6 +850,26 @@ def __init__(
             servers,
         )
         self._main_address = main_address
+        if main_pool_pid:
+            self._watch_main_pool_task = asyncio.create_task(
+                self._watch_main_pool(main_pool_pid)
+            )
+        else:
+            self._watch_main_pool_task = None
+
+    async def _watch_main_pool(self, main_pool_pid: int):
+        main_process = psutil.Process(main_pool_pid)
+        while not self.stopped:
+            try:
+                await asyncio.to_thread(main_process.status)
+                await asyncio.sleep(0.1)
+                continue
+            except (psutil.NoSuchProcess, ProcessLookupError, asyncio.CancelledError):
+                # main pool died
+                break
+
+        if not self.stopped:
+            await self.stop()
 
     async def notify_main_pool_to_destroy(
         self, message: DestroyActorMessage
@@ -900,14 +924,22 @@ async def handle_control_command(
 
     @staticmethod
     def _parse_config(config: dict, kw: dict) -> dict:
+        main_pool_pid = config.pop("main_pool_pid", None)
         kw = AbstractActorPool._parse_config(config, kw)
         pool_config: ActorPoolConfig = kw["config"]
         main_process_index = pool_config.get_process_indexes()[0]
         kw["main_address"] = pool_config.get_pool_config(main_process_index)[
             "external_address"
         ][0]
+        kw["main_pool_pid"] = main_pool_pid
         return kw
 
+    async def stop(self):
+        await super().stop()
+        if self._watch_main_pool_task:
+            self._watch_main_pool_task.cancel()
+            await self._watch_main_pool_task
+
 
 class MainActorPoolBase(ActorPoolBase):
     __slots__ = (

python/xoscar/backends/test/pool.py

@@ -55,7 +55,7 @@ async def start_sub_pool(
         status_queue: multiprocessing.Queue = multiprocessing.Queue()
         return (
             asyncio.create_task(
-                cls._create_sub_pool(actor_pool_config, process_index, status_queue)
+                cls._create_sub_pool(actor_pool_config, process_index, status_queue, 0)
             ),
             status_queue,
         )
@@ -77,9 +77,14 @@ async def _create_sub_pool(
         actor_config: ActorPoolConfig,
         process_index: int,
         status_queue: multiprocessing.Queue,
+        main_pool_pid: int,
     ):
         pool: TestSubActorPool = await TestSubActorPool.create(
-            {"actor_pool_config": actor_config, "process_index": process_index}
+            {
+                "actor_pool_config": actor_config,
+                "process_index": process_index,
+                "main_pool_pid": main_pool_pid,
+            }
         )
         await pool.start()
         status_queue.put(

python/xoscar/backends/test/tests/test_transfer.py

@@ -120,7 +120,7 @@ async def _copy_test(scheme1: Optional[str], scheme2: Optional[str], cpu: bool):
         external_address_schemes=[None, scheme1, scheme2],
     )
 
-    async with pool:
+    async with pool, pool2:
         ctx = get_context()
 
         # actor on main pool

python/xoscar/collective/common.py

@@ -55,30 +55,30 @@ class AllReduceAlgorithm(IntEnum):
 
 
 TypeMappingGloo: Dict[Type[np.dtype], "xp.GlooDataType_t"] = {
-    np.int8: xp.GlooDataType_t.glooInt8,
-    np.uint8: xp.GlooDataType_t.glooUint8,
-    np.int32: xp.GlooDataType_t.glooInt32,
-    np.uint32: xp.GlooDataType_t.glooUint32,
-    np.int64: xp.GlooDataType_t.glooInt64,
-    np.uint64: xp.GlooDataType_t.glooUint64,
-    np.float16: xp.GlooDataType_t.glooFloat16,
-    np.float32: xp.GlooDataType_t.glooFloat32,
-    np.float64: xp.GlooDataType_t.glooFloat64,
+    np.int8: xp.GlooDataType_t.glooInt8,  # type: ignore
+    np.uint8: xp.GlooDataType_t.glooUint8,  # type: ignore
+    np.int32: xp.GlooDataType_t.glooInt32,  # type: ignore
+    np.uint32: xp.GlooDataType_t.glooUint32,  # type: ignore
+    np.int64: xp.GlooDataType_t.glooInt64,  # type: ignore
+    np.uint64: xp.GlooDataType_t.glooUint64,  # type: ignore
+    np.float16: xp.GlooDataType_t.glooFloat16,  # type: ignore
+    np.float32: xp.GlooDataType_t.glooFloat32,  # type: ignore
+    np.float64: xp.GlooDataType_t.glooFloat64,  # type: ignore
 }
 cupy = lazy_import("cupy")
 if cupy is not None:
     from cupy.cuda import nccl
 
     TypeMappingNCCL: Dict[Type[np.dtype], int] = {
-        np.int8: nccl.NCCL_INT8,
-        np.uint8: nccl.NCCL_UINT8,
-        np.int32: nccl.NCCL_INT32,
-        np.uint32: nccl.NCCL_UINT32,
-        np.int64: nccl.NCCL_INT64,
-        np.uint64: nccl.NCCL_UINT64,
-        np.float16: nccl.NCCL_FLOAT16,
-        np.float32: nccl.NCCL_FLOAT32,
-        np.float64: nccl.NCCL_FLOAT64,
+        np.int8: nccl.NCCL_INT8,  # type: ignore
+        np.uint8: nccl.NCCL_UINT8,  # type: ignore
+        np.int32: nccl.NCCL_INT32,  # type: ignore
+        np.uint32: nccl.NCCL_UINT32,  # type: ignore
+        np.int64: nccl.NCCL_INT64,  # type: ignore
+        np.uint64: nccl.NCCL_UINT64,  # type: ignore
+        np.float16: nccl.NCCL_FLOAT16,  # type: ignore
+        np.float32: nccl.NCCL_FLOAT32,  # type: ignore
+        np.float64: nccl.NCCL_FLOAT64,  # type: ignore
     }
 
     ReduceOpMappingNCCL: Dict[CollectiveReduceOp, int] = {