Simplify argument passing in actor batch calls (#3098)

Author: wjsi

mars/_resource.pyx

@@ -24,29 +24,48 @@ cdef class Resource:
         self.mem_bytes = mem_bytes
 
     def __eq__(self, Resource other):
-        return self.mem_bytes == other.mem_bytes and \
-               self.num_gpus == other.num_gpus and \
-               self.num_cpus == other.num_cpus
+        cdef bint ret = (
+            self.mem_bytes == other.mem_bytes
+            and self.num_gpus == other.num_gpus
+            and self.num_cpus == other.num_cpus
+        )
+        return ret
+
+    cdef bint _le(self, Resource other) nogil:
+        # memory first, then gpu, cpu last
+        cdef bint ret = (
+            self.mem_bytes <= other.mem_bytes
+            and self.num_gpus <= other.num_gpus
+            and self.num_cpus <= other.num_cpus
+        )
+        return ret
 
     def __gt__(self, Resource other):
-        return not self.__le__(other)
+        return not self._le(other)
 
     def __le__(self, Resource other):
-        # memory first, then gpu, cpu last
-        return self.mem_bytes <= other.mem_bytes and \
-               self.num_gpus <= other.num_gpus and \
-               self.num_cpus <= other.num_cpus
+        return self._le(other)
 
     def __add__(self, Resource other):
-        return Resource(num_cpus=self.num_cpus + other.num_cpus,
-                        num_gpus=self.num_gpus + other.num_gpus,
-                        mem_bytes=self.mem_bytes + other.mem_bytes)
+        return Resource(
+            num_cpus=self.num_cpus + other.num_cpus,
+            num_gpus=self.num_gpus + other.num_gpus,
+            mem_bytes=self.mem_bytes + other.mem_bytes,
+        )
+
     def __sub__(self, Resource other):
-        return Resource(num_cpus=self.num_cpus - other.num_cpus,
-                        num_gpus=self.num_gpus - other.num_gpus,
-                        mem_bytes=self.mem_bytes - other.mem_bytes)
+        return Resource(
+            num_cpus=self.num_cpus - other.num_cpus,
+            num_gpus=self.num_gpus - other.num_gpus,
+            mem_bytes=self.mem_bytes - other.mem_bytes,
+        )
+
     def __neg__(self):
-        return Resource(num_cpus=-self.num_cpus, num_gpus=-self.num_gpus, mem_bytes=-self.mem_bytes)
+        return Resource(
+            num_cpus=-self.num_cpus,
+            num_gpus=-self.num_gpus,
+            mem_bytes=-self.mem_bytes,
+        )
 
     def __repr__(self):
         return f"Resource(num_cpus={self.num_cpus}, num_gpus={self.num_gpus}, mem_bytes={self.mem_bytes})"

mars/learn/metrics/pairwise/tests/test_haversine_distances.py

@@ -54,12 +54,12 @@ def test_haversine_distances_execution(setup, x, y, use_sklearn):
 
     result = distance.execute().fetch()
     expected = sk_haversine_distances(raw_x, raw_y)
-    np.testing.assert_array_equal(result, expected)
+    np.testing.assert_almost_equal(result, expected)
 
     # test x is y
     distance = haversine_distances(x)
     distance.op._use_sklearn = use_sklearn
 
     result = distance.execute().fetch()
     expected = sk_haversine_distances(raw_x, raw_x)
-    np.testing.assert_array_equal(result, expected)
+    np.testing.assert_almost_equal(result, expected)

mars/oscar/batch.py

@@ -133,52 +133,54 @@ def delay(*args, **kwargs):
 
     @staticmethod
     def _gen_args_kwargs_list(delays):
-        args_list = list()
-        kwargs_list = list()
-        for delay in delays:
-            args_list.append(delay.args)
-            kwargs_list.append(delay.kwargs)
+        args_list = [delay.args for delay in delays]
+        kwargs_list = [delay.kwargs for delay in delays]
         return args_list, kwargs_list
 
-    async def _async_batch(self, *delays):
+    async def _async_batch(self, args_list, kwargs_list):
         # when there is only one call in batch, calling one-pass method
         # will be more efficient
-        if len(delays) == 0:
+        if len(args_list) == 0:
             return []
-        elif len(delays) == 1:
-            d = delays[0]
-            return [await self._async_call(*d.args, **d.kwargs)]
+        elif len(args_list) == 1:
+            return [await self._async_call(*args_list[0], **kwargs_list[0])]
         elif self.batch_func:
-            args_list, kwargs_list = self._gen_args_kwargs_list(delays)
             return await self.batch_func(args_list, kwargs_list)
         else:
             # this function has no batch implementation
             # call it separately
             tasks = [
-                asyncio.create_task(self.func(*d.args, **d.kwargs)) for d in delays
+                asyncio.create_task(self.func(*args, **kwargs))
+                for args, kwargs in zip(args_list, kwargs_list)
             ]
             try:
                 return await asyncio.gather(*tasks)
             except asyncio.CancelledError:
                 _ = [task.cancel() for task in tasks]
                 return await asyncio.gather(*tasks)
 
-    def _sync_batch(self, *delays):
-        if delays == 0:
+    def _sync_batch(self, args_list, kwargs_list):
+        if len(args_list) == 0:
             return []
         elif self.batch_func:
-            args_list, kwargs_list = self._gen_args_kwargs_list(delays)
             return self.batch_func(args_list, kwargs_list)
         else:
             # this function has no batch implementation
             # call it separately
-            return [self.func(*d.args, **d.kwargs) for d in delays]
+            return [
+                self.func(*args, **kwargs)
+                for args, kwargs in zip(args_list, kwargs_list)
+            ]
 
     def batch(self, *delays):
+        args_list, kwargs_list = self._gen_args_kwargs_list(delays)
+        return self.call_with_lists(args_list, kwargs_list)
+
+    def call_with_lists(self, args_list, kwargs_list):
         if self.is_async:
-            return self._async_batch(*delays)
+            return self._async_batch(args_list, kwargs_list)
         else:
-            return self._sync_batch(*delays)
+            return self._sync_batch(args_list, kwargs_list)
 
     def bind(self, *args, **kwargs):
         if self.bind_func is None:

mars/oscar/core.pyx

@@ -17,9 +17,10 @@ import inspect
 import logging
 import sys
 import weakref
-cimport cython
 from typing import AsyncGenerator
 
+cimport cython
+
 from .context cimport get_context
 from .errors import Return, ActorNotExist
 from .utils cimport is_async_generator
@@ -33,7 +34,7 @@ logger = logging.getLogger(__name__)
 cdef:
     bint _log_unhandled_errors = False
     bint _log_cycle_send = False
-    object _local_pool_map = weakref.WeakValueDictionary()
+    dict _local_pool_map = dict()
     object _actor_method_wrapper
 
 
@@ -54,7 +55,8 @@ cdef _get_local_actor(address, uid):
     # disabled the local actor proxy if the debug option is on.
     if _log_cycle_send:
         return None
-    pool = _local_pool_map.get(address)
+    pool_ref = _local_pool_map.get(address)
+    pool = None if pool_ref is None else pool_ref()
     if pool is not None:
         actor = pool._actors.get(uid)
         if actor is not None:
@@ -66,7 +68,9 @@ def register_local_pool(address, pool):
     """
     Register local actor pool for local actor lookup.
     """
-    _local_pool_map[address] = pool
+    _local_pool_map[address] = weakref.ref(
+        pool, lambda _: _local_pool_map.pop(address, None)
+    )
 
 
 cpdef create_local_actor_ref(address, uid):
@@ -177,22 +181,34 @@ cdef class ActorRefMethod:
 
     def batch(self, *delays, send=True):
         cdef:
-            list args_list = list()
-            list kwargs_list = list()
+            long n_delays = len(delays)
+            bint has_kw = False
+            list args_list
+            list kwargs_list
+            _DelayedArgument delay
+
+        args_list = [None] * n_delays
+        kwargs_list = [None] * n_delays
 
         last_method = None
-        for delay in delays:
+        for idx in range(n_delays):
+            delay = delays[idx]
             method, _call_method, args, kwargs = delay.arguments
             if last_method is not None and method != last_method:
                 raise ValueError('Does not support calling multiple methods in batch')
             last_method = method
 
-            args_list.append(args)
-            kwargs_list.append(kwargs)
+            args_list[idx] = args
+            kwargs_list[idx] = kwargs
+            if kwargs:
+                has_kw = True
+
+        if not has_kw:
+            kwargs_list = None
         if last_method is None:
             last_method = self.method_name
 
-        message = (last_method, CALL_METHOD_BATCH, (args_list, kwargs_list), {})
+        message = (last_method, CALL_METHOD_BATCH, (args_list, kwargs_list), None)
         return get_context().send(self.ref, message, wait_response=send, **self._options)
 
     def tell_delay(self, *args, delay=None, ignore_conn_fail=True, **kwargs):
@@ -488,10 +504,10 @@ cdef class _BaseActor:
                     with debug_async_timeout('actor_lock_timeout',
                                              "Batch method %s of actor %s hold lock timeout, batch size %s.",
                                              method, self.uid, len(args)):
-                        delays = []
-                        for s_args, s_kwargs in zip(*args):
-                            delays.append(func.delay(*s_args, **s_kwargs))
-                        result = func.batch(*delays)
+                        args_list, kwargs_list = args
+                        if kwargs_list is None:
+                            kwargs_list = [{}] * len(args_list)
+                        result = func.call_with_lists(args_list, kwargs_list)
                         if asyncio.iscoroutine(result):
                             result = await result
             else:  # pragma: no cover

mars/serialization/core.pxd

@@ -13,6 +13,8 @@
 # limitations under the License.
 
 cdef class Serializer:
+    cdef int _serializer_id
+
     cpdef serial(self, object obj, dict context)
     cpdef deserial(self, tuple serialized, dict context, list subs)
     cpdef on_deserial_error(

mars/serialization/core.pyx

@@ -66,6 +66,10 @@ cdef int32_t _SERIALIZER_ID_PRIME = 32749
 cdef class Serializer:
     serializer_id = None
 
+    def __cinit__(self):
+        # make the value can be referenced with C code
+        self._serializer_id = self.serializer_id
+
     cpdef serial(self, object obj, dict context):
         """
         Returns intermediate serialization result of certain object.
@@ -158,14 +162,15 @@ cdef class Serializer:
 
     @classmethod
     def register(cls, obj_type):
-        inst = cls()
         if (
             cls.serializer_id is None
             or cls.serializer_id == getattr(super(cls, cls), "serializer_id", None)
         ):
             # a class should have its own serializer_id
             # inherited serializer_id not acceptable
             cls.serializer_id = cls.calc_default_serializer_id()
+
+        inst = cls()
         _serial_dispatcher.register(obj_type, inst)
         if _deserializers.get(cls.serializer_id) is not None:
             assert type(_deserializers[cls.serializer_id]) is cls
@@ -321,10 +326,16 @@ cdef class StrSerializer(Serializer):
 cdef class CollectionSerializer(Serializer):
     obj_type = None
 
+    cdef object _obj_type
+
+    def __cinit__(self):
+        # make the value can be referenced with C code
+        self._obj_type = self.obj_type
+
     cpdef tuple _serial_iterable(self, obj: Any):
         cdef list idx_to_propagate = []
         cdef list obj_to_propagate = []
-        cdef list obj_list = list(obj)
+        cdef list obj_list = <list>obj if type(obj) is list else list(obj)
         cdef int64_t idx
         cdef object item
 
@@ -340,11 +351,14 @@ cdef class CollectionSerializer(Serializer):
             elif type(item) in _primitive_types:
                 continue
 
+            if obj is obj_list:
+                obj_list = list(obj)
+
             obj_list[idx] = None
             idx_to_propagate.append(idx)
             obj_to_propagate.append(item)
 
-        if self.obj_type is not None and type(obj) is not self.obj_type:
+        if self._obj_type is not None and type(obj) is not self._obj_type:
             obj_type = type(obj)
         else:
             obj_type = None
@@ -420,13 +434,19 @@ def _dict_value_replacer(context, ret, key, real_value):
 
 cdef class DictSerializer(CollectionSerializer):
     serializer_id = 6
-    _inspected_inherits = set()
+    cdef set _inspected_inherits
+
+    def __cinit__(self):
+        self._inspected_inherits = set()
 
     cpdef serial(self, obj: Any, dict context):
         cdef uint64_t obj_id
         cdef tuple key_obj, value_obj
         cdef list key_bufs, value_bufs
 
+        if type(obj) is dict and len(<dict>obj) == 0:
+            return (), [], True
+
         obj_id = _fast_id(obj)
         if obj_id in context:
             return Placeholder(obj_id)
@@ -460,6 +480,8 @@ cdef class DictSerializer(CollectionSerializer):
         cdef int64_t i, num_key_bufs
         cdef list key_subs, value_subs, keys, values
 
+        if not serialized:
+            return {}
         if len(serialized) == 1:
             # serialized directly
             return serialized[0]
@@ -537,6 +559,7 @@ PickleSerializer.register(object)
 for _primitive in _primitive_types:
     PrimitiveSerializer.register(_primitive)
 BytesSerializer.register(bytes)
+BytesSerializer.register(memoryview)
 StrSerializer.register(str)
 ListSerializer.register(list)
 TupleSerializer.register(tuple)
@@ -588,7 +611,7 @@ cdef tuple _serial_single(
             # REMEMBER to change _COMMON_HEADER_LEN when content of
             # this header changed
             common_header = (
-                serializer.serializer_id, ordered_id, len(subs), final
+                serializer._serializer_id, ordered_id, len(subs), final
             )
             break
         else:

mars/serialization/tests/test_serial.py

@@ -55,7 +55,7 @@ class CustomList(list):
         "abcd",
         ["uvw", ("mno", "sdaf"), 4, 6.7],
         CustomList([3, 4, CustomList([5, 6])]),
-        {"abc": 5.6, "def": [3.4], "gh": None},
+        {"abc": 5.6, "def": [3.4], "gh": None, "ijk": {}},
         OrderedDict([("abcd", 5.6)]),
         defaultdict(lambda: 0, [("abcd", 0)]),
     ],