Use cython to accelerate message serialization (mars-project#2932)

Author: wjsi

asv_bench/benchmarks/serialize.py

@@ -18,6 +18,9 @@
 
 from mars.core.operand import MapReduceOperand
 from mars.dataframe.operands import DataFrameOperandMixin
+from mars.oscar.core import ActorRef
+
+from mars.oscar.backends.message import SendMessage, new_message_id, ActorRefMessage
 from mars.serialization import serialize, deserialize
 from mars.serialization.serializables import (
     Serializable,
@@ -38,7 +41,7 @@
     TupleField,
     DictField,
 )
-from mars.services.subtask import Subtask
+from mars.services.subtask import Subtask, SubtaskResult, SubtaskStatus
 from mars.services.task import new_task_id
 from mars.utils import tokenize
 
@@ -215,3 +218,49 @@ def time_pickle_serialize_fetch_shuffle_chunks(self):
             header, buffers = serialize(fetch_chunk)
             serialized = cloudpickle.dumps((header, buffers))
             deserialize(*cloudpickle.loads(serialized))
+
+
+class SerializeMessageSuite:
+    def setup(self):
+        self.send_messages = []
+        self.actor_ref_messages = []
+        for i in range(10000):
+            ref = ActorRef(
+                "ray://mars_cluster_1649927648/17/0",
+                b"F20Wyerq6EiqltB8jAVs7L3N_task_manager",
+            )
+            new_result = SubtaskResult(
+                subtask_id=new_task_id(),
+                session_id=new_task_id(),
+                task_id=new_task_id(),
+                stage_id=new_task_id(),
+                status=SubtaskStatus.succeeded,
+                progress=1.0,
+                data_size=1000000.0,
+                bands=[("ray://mars_cluster_1649927648/17/0", "numa-0")],
+                execution_start_time=1646125099.622051,
+                execution_end_time=1646125104.448726,
+            )
+            send_message = SendMessage(
+                new_message_id(),
+                ref,
+                new_result,
+                protocol=0,
+            )
+            self.send_messages.append(send_message)
+            actor_ref_message = ActorRefMessage(
+                message_id=new_message_id(),
+                actor_ref=ref,
+                protocol=0,
+            )
+            self.actor_ref_messages.append(actor_ref_message)
+
+    def time_pickle_serialize_deserialize_send_messages(self):
+        deserialize(
+            *cloudpickle.loads(cloudpickle.dumps(serialize(self.send_messages)))
+        )
+
+    def time_pickle_serialize_deserialize_actor_ref_messages(self):
+        deserialize(
+            *cloudpickle.loads(cloudpickle.dumps(serialize(self.actor_ref_messages)))
+        )

mars/oscar/backends/context.py

@@ -13,9 +13,10 @@
 # limitations under the License.
 
 import asyncio
+from dataclasses import dataclass
 from typing import Tuple, Union, Type
 
-from ...utils import to_binary
+from ...utils import to_binary, dataslots
 from ..api import Actor
 from ..core import ActorRef, create_local_actor_ref
 from ..context import BaseActorContext
@@ -38,11 +39,16 @@
     CancelMessage,
     ControlMessage,
     ControlMessageType,
-    ProfilingContext,
 )
 from .router import Router
 
 
+@dataslots
+@dataclass
+class ProfilingContext:
+    task_id: str
+
+
 class MarsActorContext(BaseActorContext):
     __slots__ = "_address", "_caller"
 

mars/oscar/backends/core.py

@@ -16,12 +16,12 @@
 import copy
 from typing import Dict, Union
 
+from ...oscar.profiling import ProfilingData
+from ...utils import Timer
 from ..errors import ServerClosed
 from .communication import Client
 from .message import _MessageBase, ResultMessage, ErrorMessage, DeserializeMessageFailed
 from .router import Router
-from ...oscar.profiling import ProfilingData
-from ...utils import Timer
 
 
 ResultMessageType = Union[ResultMessage, ErrorMessage]
@@ -60,10 +60,10 @@ async def _listen(self, client: Client):
                     ) from None
                 future = self._client_to_message_futures[client].pop(message.message_id)
                 future.set_result(message)
-            except DeserializeMessageFailed as e:  # pragma: no cover
+            except DeserializeMessageFailed as e:
                 message_id = e.message_id
                 future = self._client_to_message_futures[client].pop(message_id)
-                future.set_exception(e)
+                future.set_exception(e.__cause__)
             except Exception as e:  # noqa: E722  # pylint: disable=bare-except
                 message_futures = self._client_to_message_futures.get(client)
                 self._client_to_message_futures[client] = dict()