Add actor driver & structure adjustment (mars-project#1925)

Author: Xuye (Chris) Qin

mars/oscar/__init__.py

@@ -12,8 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from .api import actor_ref, create_actor, has_actor, destroy_actor
-from .core import Actor, ActorRef
+from .api import actor_ref, create_actor, has_actor, destroy_actor, Actor
+from .core import ActorRef
 from .context import get_context
 from .errors import ActorNotExist, ActorAlreadyExist
 from .utils import create_actor_ref

mars/oscar/api.py

@@ -12,7 +12,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from typing import Any, Dict, Type, Tuple
+
 from .context import get_context
+from .core import _Actor
 
 
 async def create_actor(actor_cls, *args, uid=None, address=None, **kwargs):
@@ -33,3 +36,48 @@ async def destroy_actor(actor_ref):
 def actor_ref(*args, **kwargs):
     ctx = get_context()
     return ctx.actor_ref(*args, **kwargs)
+
+
+class Actor(_Actor):
+    def __new__(cls, *args, **kwargs):
+        try:
+            return _actor_implementation[cls](*args, **kwargs)
+        except KeyError:
+            return super().__new__(cls, *args, **kwargs)
+
+    async def __post_create__(self):
+        """
+        Method called after actor creation
+        """
+        return await super().__post_create__()
+
+    async def __pre_destroy__(self):
+        """
+        Method called before actor destroy
+        """
+        return await super().__pre_destroy__()
+
+    async def __on_receive__(self, message: Tuple[Any]):
+        """
+        Handle message from other actors and dispatch them to user methods
+
+        Parameters
+        ----------
+        message : tuple
+            Message shall be (method_name,) + args + (kwargs,)
+        """
+        return await super().__on_receive__(message)
+
+
+_actor_implementation: Dict[Type[Actor], Type[Actor]] = dict()
+
+
+def register_actor_implementation(actor_cls: Type[Actor], impl_cls: Type[Actor]):
+    _actor_implementation[actor_cls] = impl_cls
+
+
+def unregister_actor_implementation(actor_cls: Type[Actor]):
+    try:
+        del _actor_implementation[actor_cls]
+    except KeyError:
+        pass

mars/oscar/backend.py

@@ -0,0 +1,45 @@
+# Copyright 1999-2020 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 abc import ABC, abstractmethod
+from typing import Type
+
+from .context import register_backend_context
+from .driver import register_backend_driver
+
+
+__all__ = ["BaseActorBackend", "register_backend"]
+
+
+class BaseActorBackend(ABC):
+    @staticmethod
+    @abstractmethod
+    def name():
+        pass
+
+    @staticmethod
+    @abstractmethod
+    def get_context_cls():
+        pass
+
+    @staticmethod
+    @abstractmethod
+    def get_driver_cls():
+        pass
+
+
+def register_backend(backend_cls: Type[BaseActorBackend]):
+    register_backend_context(backend_cls.name(), backend_cls.get_context_cls())
+    register_backend_driver(backend_cls.name(), backend_cls.get_driver_cls())
+    return backend_cls

mars/oscar/backends/mars/__init__.py

@@ -12,6 +12,4 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from .pool import register_backend
-register_backend()
-del register_backend
+from .backend import MarsActorBackend

mars/oscar/backends/mars/backend.py

@@ -0,0 +1,36 @@
+# Copyright 1999-2020 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 BaseActorBackend, register_backend
+from .driver import MarsActorDriver
+from .pool import MarsActorContext
+
+
+__all__ = ['MarsActorBackend']
+
+
+@register_backend
+class MarsActorBackend(BaseActorBackend):
+    @staticmethod
+    def name():
+        # return None because Mars is default scheme
+        return
+
+    @staticmethod
+    def get_context_cls():
+        return MarsActorContext
+
+    @staticmethod
+    def get_driver_cls():
+        return MarsActorDriver

mars/oscar/backends/mars/driver.py

@@ -0,0 +1,24 @@
+# Copyright 1999-2020 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 MarsActorDriver(BaseActorDriver):
+    def setup_cluster(self, address_to_resources: Dict[str, Dict[str, Number]]):
+        # nothing need to be done in driver of Mars backend
+        pass

mars/oscar/backends/mars/pool.pyx

@@ -63,8 +63,3 @@ cdef class MarsActorContext(BaseActorContext):
             raise ActorNotExist(f'Actor {actor_ref.uid} does not exist')
         ret = await actor.__on_receive__(message)
         return ret if wait_response else None
-
-
-def register_backend():
-    from ...context import register_backend_context
-    register_backend_context(None, MarsActorContext)

mars/oscar/context.pyx

@@ -14,8 +14,7 @@
 
 from urllib.parse import urlparse
 
-from .core cimport _Actor, ActorRef
-from .errors import ActorAlreadyExist, ActorNotExist
+from .core cimport ActorRef
 from .utils cimport new_actor_id
 
 cdef dict _backend_context_cls = dict()

mars/oscar/core.pyx

@@ -188,25 +188,3 @@ cdef class _Actor:
             if asyncio.iscoroutine(result):
                 result = await result
         return await self._handle_actor_result(result)
-
-
-class Actor(_Actor):
-    def __new__(cls, *args, **kwargs):
-        try:
-            return _actor_implementation[id(cls)](*args, **kwargs)
-        except KeyError:
-            return super().__new__(cls, *args, **kwargs)
-
-
-cdef dict _actor_implementation = dict()
-
-
-def register_actor_implementation(actor_cls, impl_cls):
-    _actor_implementation[id(actor_cls)] = impl_cls
-
-
-def unregister_actor_implementation(actor_cls):
-    try:
-        del _actor_implementation[id(actor_cls)]
-    except KeyError:
-        pass

mars/oscar/driver.py

@@ -0,0 +1,40 @@
+# Copyright 1999-2020 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 abc import ABC, abstractmethod
+from numbers import Number
+from typing import Dict
+
+
+class BaseActorDriver(ABC):
+    @abstractmethod
+    def setup_cluster(self, address_to_resources: Dict[str, Dict[str, Number]]):
+        """
+        Setup cluster according to given resources,
+        resources is a dict, e.g. {'CPU': 3, 'GPU': 1}
+
+        Parameters
+        ----------
+        address_to_resources: dict
+            resources that required for each node.
+        """
+        pass
+
+
+_backend_driver_cls: Dict = dict()
+
+
+def register_backend_driver(scheme, cls):
+    assert issubclass(cls, BaseActorDriver)
+    _backend_driver_cls[scheme] = cls