[oscar] Add cancel support, optimize error handling, add kill_actor API (#2027)

Author: Xuye (Chris) Qin

mars/lib/gipc.pyx

@@ -442,6 +442,30 @@ except ImportError:
     ForkProcess = multiprocessing.Process
 
 
+if not WINDOWS:
+    try:
+        from multiprocessing.forking import Popen as mp_Popen
+    except ImportError:  # pragma: no cover
+        # multiprocessing's internal structure has changed from 3.3 to 3.4.
+        from multiprocessing.popen_fork import Popen as mp_Popen
+
+    # multiprocessing.process.Process.start() and other methods may
+    # call multiprocessing.process._cleanup(). This and other mp methods
+    # may call multiprocessing's Popen.poll() which itself invokes
+    # os.waitpid(). In extreme cases (high-frequent child process
+    # creation, short-living child processes), this competes with libev's
+    # SIGCHLD handler and may win, resulting in libev not being able to
+    # retrieve all SIGCHLD signals corresponding to started children. This
+    # could make certain _GProcess.join() calls block forever.
+    # -> Prevent multiprocessing's Popen.poll() from calling
+    # os.waitpid(). Let libev do the job.
+    class _GPopen(mp_Popen):
+        def poll(self, *args, **kwargs):
+            pass
+else:
+    _GPopen = None
+
+
 class _GProcess(ForkProcess):
     """
     Compatible with the ``multiprocessing.Process`` API.
@@ -488,24 +512,9 @@ class _GProcess(ForkProcess):
     # On Windows, cooperative `join()` is realized via polling (non-blocking
     # calls to `Process.is_alive()`) and the original `join()` method.
     if not WINDOWS:
-        # multiprocessing.process.Process.start() and other methods may
-        # call multiprocessing.process._cleanup(). This and other mp methods
-        # may call multiprocessing's Popen.poll() which itself invokes
-        # os.waitpid(). In extreme cases (high-frequent child process
-        # creation, short-living child processes), this competes with libev's
-        # SIGCHLD handler and may win, resulting in libev not being able to
-        # retrieve all SIGCHLD signals corresponding to started children. This
-        # could make certain _GProcess.join() calls block forever.
-        # -> Prevent multiprocessing's Popen.poll() from calling
-        # os.waitpid(). Let libev do the job.
-        try:
-            from multiprocessing.forking import Popen as mp_Popen
-        except ImportError:  # pragma: no cover
-            # multiprocessing's internal structure has changed from 3.3 to 3.4.
-            from multiprocessing.popen_fork import Popen as mp_Popen
-        # Monkey-patch and forget about the name.
-        mp_Popen.poll = lambda *a, **b: None
-        del mp_Popen
+        @staticmethod
+        def _Popen(process_obj):
+            return _GPopen(process_obj)
 
         def start(self):
             # Start grabbing SIGCHLD within libev event loop.

mars/oscar/__init__.py

@@ -17,7 +17,7 @@
 del aio
 
 from .api import actor_ref, create_actor, has_actor, destroy_actor, \
-    Actor, create_actor_pool
+    kill_actor, Actor, create_actor_pool
 from .errors import ActorNotExist, ActorAlreadyExist
 from .utils import create_actor_ref
 

mars/oscar/api.py

@@ -40,6 +40,11 @@ async def actor_ref(*args, **kwargs):
     return await ctx.actor_ref(*args, **kwargs)
 
 
+async def kill_actor(actor_ref):
+    ctx = get_context()
+    return await ctx.kill_actor(actor_ref)
+
+
 async def create_actor_pool(address: str,
                             n_process: int = None,
                             **kwargs):

mars/oscar/backends/mars/allocate_strategy.py

@@ -73,7 +73,8 @@ def get_allocated_address(self,
                               config: ActorPoolConfig,
                               allocated: allocated_type) -> str:
         # allocate to main process
-        return config.get_external_address(0)
+        main_process_index = config.get_process_indexes()[0]
+        return config.get_external_address(main_process_index)
 
 
 class RandomSubPool(AllocateStrategy):
@@ -86,6 +87,20 @@ def get_allocated_address(self,
         return choice(config.get_external_addresses()[1:])
 
 
+class ProcessIndex(AllocateStrategy):
+    __slots__ = 'process_index',
+
+    def __init__(self, process_index: int):
+        self.process_index = process_index
+
+    @implements(AllocateStrategy.get_allocated_address)
+    def get_allocated_address(self,
+                              config: ActorPoolConfig,
+                              allocated: allocated_type) -> str:
+        actual_process_index = config.get_process_indexes()[self.process_index]
+        return config.get_pool_config(actual_process_index)['external_address'][0]
+
+
 class RandomLabel(AllocateStrategy):
     __slots__ = 'label',
 

mars/oscar/backends/mars/communication/__init__.py

@@ -13,7 +13,8 @@
 # limitations under the License.
 
 from .base import Client, Server, Channel
-from .core import get_client_type, get_server_type, gen_internal_address
+from .core import get_client_type, get_server_type, \
+    gen_internal_address, gen_local_address
 from .dummy import DummyClient, DummyServer, DummyChannel
 from .socket import SocketClient, SocketServer, UnixSocketClient, \
     UnixSocketServer, SocketChannel

mars/oscar/backends/mars/communication/core.py

@@ -63,3 +63,7 @@ def get_server_type(address: str) -> Type[Server]:
 
 def gen_internal_address(process_index: int) -> str:
     return f'unixsocket:///{process_index}'
+
+
+def gen_local_address(process_index: int) -> str:
+    return f'dummy://{process_index}'

mars/oscar/backends/mars/communication/dummy.py

@@ -15,13 +15,14 @@
 import asyncio
 import concurrent.futures as futures
 from typing import Any, Callable, Coroutine, Dict, Type
+from urllib.parse import urlparse
 
 from .....utils import implements, classproperty
 from .base import Channel, ChannelType, Server, Client
 from .core import register_client, register_server
 from .errors import ChannelClosed
 
-DUMMY_ADDRESS = 'dummy://'
+DEFAULT_DUMMY_ADDRESS = 'dummy://0'
 
 
 class DummyChannel(Channel):
@@ -81,7 +82,7 @@ def closed(self) -> bool:
 class DummyServer(Server):
     __slots__ = '_closed',
 
-    _instance = None
+    _address_to_instances: Dict[str, "DummyServer"] = dict()
     scheme = 'dummy'
 
     def __init__(self,
@@ -91,8 +92,8 @@ def __init__(self,
         self._closed = asyncio.Event()
 
     @classmethod
-    def get_instance(cls):
-        return cls._instance
+    def get_instance(cls, address: str):
+        return cls._address_to_instances[address]
 
     @classproperty
     @implements(Server.client_type)
@@ -108,23 +109,22 @@ def channel_type(self) -> ChannelType:
     @implements(Server.create)
     async def create(config: Dict) -> "DummyServer":
         config = config.copy()
-        address = config.pop('address', DUMMY_ADDRESS)
+        address = config.pop('address', DEFAULT_DUMMY_ADDRESS)
         handle_channel = config.pop('handle_channel')
-        if address != DUMMY_ADDRESS:  # pragma: no cover
+        if urlparse(address).scheme != DummyServer.scheme:  # pragma: no cover
             raise ValueError(f'Address for DummyServer '
-                             f'should be {DUMMY_ADDRESS}, '
+                             f'should be starts with "dummy://", '
                              f'got {address}')
         if config:  # pragma: no cover
             raise TypeError(f'Creating DummyServer got unexpected '
                             f'arguments: {",".join(config)}')
 
-        # DummyServer is singleton
-        if DummyServer._instance is not None:
-            return DummyServer._instance
-
-        server = DummyServer(DUMMY_ADDRESS, handle_channel)
-        DummyServer._instance = server
-        return server
+        try:
+            return DummyServer.get_instance(address)
+        except KeyError:
+            server = DummyServer(address, handle_channel)
+            DummyServer._address_to_instances[address] = server
+            return server
 
     @implements(Server.start)
     async def start(self):
@@ -151,7 +151,7 @@ async def on_connected(self, *args, **kwargs):
     @implements(Server.stop)
     async def stop(self):
         self._closed.set()
-        DummyServer._instance = None
+        del DummyServer._address_to_instances[self.address]
 
     @property
     @implements(Server.stopped)
@@ -179,10 +179,10 @@ def __init__(self,
     async def connect(dest_address: str,
                       local_address: str = None,
                       **kwargs) -> "Client":
-        if dest_address != DUMMY_ADDRESS:  # pragma: no cover
-            raise ValueError(f'Destination address has to be "dummy://" '
+        if urlparse(dest_address).scheme != DummyServer.scheme:  # pragma: no cover
+            raise ValueError(f'Destination address should start with "dummy://" '
                              f'for DummyClient, got {dest_address}')
-        server = DummyServer.get_instance()
+        server = DummyServer.get_instance(dest_address)
         if server is None:  # pragma: no cover
             raise RuntimeError('DummyServer needs to be created '
                                'first before DummyClient')
@@ -200,6 +200,5 @@ async def connect(dest_address: str,
     @implements(Client.close)
     async def close(self):
         await super().close()
-        DummyClient._instance = None
         self._task.cancel()
         self._task = None

mars/oscar/backends/mars/communication/socket.py

@@ -210,9 +210,12 @@ async def connect(dest_address: str,
         return SocketClient(local_address, dest_address, channel)
 
 
+TEMPDIR = tempfile.gettempdir()
+
+
 @lru_cache(100)
 def _gen_unix_socket_default_path(process_index):
-    return f'{tempfile.gettempdir()}/mars/' \
+    return f'{TEMPDIR}/mars/' \
            f'{md5(to_binary(str(process_index))).hexdigest()}'  # nosec
 
 

mars/oscar/backends/mars/communication/tests/test_comm.py

@@ -15,7 +15,7 @@
 import asyncio
 import sys
 import multiprocessing
-from typing import Union
+from typing import Union, List, Tuple, Type, Dict
 
 import numpy as np
 import pytest
@@ -24,7 +24,7 @@
 from mars.oscar.backends.mars.communication import \
     SocketChannel, SocketServer, UnixSocketServer, \
     DummyChannel, DummyServer, get_client_type, \
-    SocketClient, UnixSocketClient, DummyClient
+    SocketClient, UnixSocketClient, DummyClient, Server
 from mars.utils import get_next_port
 
 
@@ -33,13 +33,13 @@
 
 
 # server_type, config, con
-params = [
+params: List[Tuple[Type[Server], Dict, str]] = [
     (SocketServer, dict(host='127.0.0.1', port=port), f'127.0.0.1:{port}'),
 ]
 if sys.platform != 'win32':
     params.append((UnixSocketServer, dict(process_index='0'), f'unixsocket:///0'))
 local_params = params.copy()
-local_params.append((DummyServer, dict(), 'dummy://'))
+local_params.append((DummyServer, dict(), 'dummy://0'))
 
 
 @pytest.mark.parametrize(

mars/oscar/backends/mars/config.py

@@ -59,6 +59,13 @@ def get_external_address(self, process_index: int) -> str:
     def get_process_indexes(self):
         return list(self._conf['pools'])
 
+    def get_process_index(self, external_addrress):
+        for process_index, conf in self._conf['pools'].items():
+            if external_addrress in conf['external_address']:
+                return process_index
+        raise ValueError(f'Cannot get proces_index '
+                         f'for {external_addrress}')  # pragma: no cover
+
     def get_external_addresses(self, label=None) -> List[str]:
         result = []
         for c in self._conf['pools'].values():