Add mutable tensor support for local cluster (#464)

Author: sighingnow

mars/api.py

@@ -13,15 +13,20 @@
 # limitations under the License.
 
 import logging
+import uuid
+import zlib
+import pyarrow
 
 from .actors import new_client
+from .config import options
 from .errors import GraphNotExists
 from .scheduler import SessionActor, GraphActor, GraphMetaActor, ResourceActor, \
     SessionManagerActor, ChunkMetaClient
 from .scheduler.graph import ResultReceiverActor
 from .scheduler.node_info import NodeInfoActor
 from .scheduler.utils import SchedulerClusterInfoActor
 from .serialize import dataserializer
+from .utils import tokenize
 
 logger = logging.getLogger(__name__)
 
@@ -51,6 +56,13 @@ def get_schedulers_info(self):
             infos[scheduler] = info_ref.get_info()
         return infos
 
+    def _get_receiver_ref(self, chunk_key):
+        from .worker.dispatcher import DispatchActor
+        ep = self.cluster_info.get_scheduler(chunk_key)
+        dispatch_ref = self.actor_client.actor_ref(DispatchActor.default_uid(), address=ep)
+        uid = dispatch_ref.get_hash_slot('receiver', chunk_key)
+        return self.actor_client.actor_ref(uid, address=ep)
+
     def count_workers(self):
         try:
             uid = ResourceActor.default_uid()
@@ -71,6 +83,66 @@ def submit_graph(self, session_id, serialized_graph, graph_key, target,
         session_ref.submit_tileable_graph(
             serialized_graph, graph_key, target, compose=compose, _tell=not wait)
 
+    def create_mutable_tensor(self, session_id, name, shape, dtype, *args, **kwargs):
+        session_uid = SessionActor.gen_uid(session_id)
+        session_ref = self.get_actor_ref(session_uid)
+        return session_ref.create_mutable_tensor(name, shape, dtype, *args, **kwargs)
+
+    def get_mutable_tensor(self, session_id, name):
+        session_uid = SessionActor.gen_uid(session_id)
+        session_ref = self.get_actor_ref(session_uid)
+        return session_ref.get_mutable_tensor(name)
+
+    def send_chunk_records(self, session_id, name, chunk_records_to_send, directly=True):
+        from .worker.dataio import ArrowBufferIO
+        from .worker.quota import MemQuotaActor
+        session_uid = SessionActor.gen_uid(session_id)
+        session_ref = self.get_actor_ref(session_uid)
+
+        chunk_records = []
+        for chunk_key, records in chunk_records_to_send.items():
+            record_chunk_key = tokenize(chunk_key, uuid.uuid4().hex)
+            ep = self.cluster_info.get_scheduler(chunk_key)
+            # register quota
+            quota_ref = self.actor_client.actor_ref(MemQuotaActor.default_uid(), address=ep)
+            quota_ref.request_batch_quota({record_chunk_key: records.nbytes})
+            # send record chunk
+            buf = pyarrow.serialize(records).to_buffer()
+            receiver_ref = self._get_receiver_ref(chunk_key)
+            receiver_ref.create_data_writer(session_id, record_chunk_key, buf.size, None,
+                                            ensure_cached=False, use_promise=False)
+
+            block_size = options.worker.transfer_block_size
+
+            try:
+                reader = ArrowBufferIO(buf, 'r', block_size=block_size)
+                checksum = 0
+                while True:
+                    next_chunk = reader.read(block_size)
+                    if not next_chunk:
+                        reader.close()
+                        receiver_ref.finish_receive(session_id, record_chunk_key, checksum)
+                        break
+                    checksum = zlib.crc32(next_chunk, checksum)
+                    receiver_ref.receive_data_part(session_id, record_chunk_key, next_chunk, checksum)
+            except:
+                receiver_ref.cancel_receive(session_id, chunk_key)
+                raise
+            finally:
+                if reader:
+                    reader.close()
+                del reader
+
+            chunk_records.append((chunk_key, record_chunk_key))
+
+        # register the record chunk to MutableTensorActor
+        session_ref.append_chunk_records(name, chunk_records)
+
+    def seal(self, session_id, name):
+        session_uid = SessionActor.gen_uid(session_id)
+        session_ref = self.get_actor_ref(session_uid)
+        return session_ref.seal(name)
+
     def delete_graph(self, session_id, graph_key):
         graph_uid = GraphActor.gen_uid(session_id, graph_key)
         graph_ref = self.get_actor_ref(graph_uid)

mars/deploy/local/session.py

@@ -68,6 +68,21 @@ def _update_tileable_shape(self, tileable):
         tileable._update_shape(tuple(sum(nsplit) for nsplit in new_nsplits))
         tileable.nsplits = new_nsplits
 
+    def create_mutable_tensor(self, name, shape, dtype, *args, **kwargs):
+        return self._api.create_mutable_tensor(self._session_id, name, shape,
+                                               dtype, *args, **kwargs)
+
+    def get_mutable_tensor(self, name):
+        return self._api.get_mutable_tensor(self._session_id, name)
+
+    def send_chunk_records(self, name, chunk_records_to_send):
+        return self._api.send_chunk_records(self._session_id, name, chunk_records_to_send)
+
+    def seal(self, name):
+        graph_key, tensor_key, tensor_id, tensor_meta = self._api.seal(self._session_id, name)
+        self._executed_tileables[tensor_key] = graph_key, {tensor_id}
+        return tensor_meta
+
     def run(self, *tileables, **kw):
         timeout = kw.pop('timeout', -1)
         fetch = kw.pop('fetch', True)

mars/scheduler/graph.py

@@ -941,6 +941,14 @@ def fetch_tileable_result(self, tileable_key):
         concat_result = executor.execute_tileable(tileable, concat=True)[0]
         return dataserializer.dumps(concat_result)
 
+    @log_unhandled
+    def add_fetch_tileable(self, tileable_key, tileable_id, shape, dtype, chunk_size, chunk_keys):
+        from ..tensor.expressions.utils import create_fetch_tensor
+        tensor = create_fetch_tensor(chunk_size, shape, dtype,
+                                     tileable_key, tileable_id, chunk_keys)
+        self._tileable_key_to_opid[tileable_key] = tensor.op.id
+        self._tileable_key_opid_to_tiled[(tileable_key, tensor.op.id)].append(tensor)
+
     @log_unhandled
     def check_operand_can_be_freed(self, succ_op_keys):
         """

mars/scheduler/mutable.py

@@ -0,0 +1,85 @@
+# Copyright 1999-2018 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 collections import defaultdict
+
+import numpy as np
+
+from ..utils import log_unhandled
+from .utils import SchedulerActor
+
+
+class MutableTensorActor(SchedulerActor):
+    """
+    Actor handling processing of a Mars mutable tensor.
+    """
+    @staticmethod
+    def gen_uid(session_id, name):
+        return 's:0:mutable-tensor$%s$%s' % (session_id, name)
+
+    def __init__(self, session_id, name, shape, dtype, graph_key, chunk_size=None, *args, **kwargs):
+        super(MutableTensorActor, self).__init__(*args, **kwargs)
+        self._session_id = session_id
+        self._name = name
+        self._shape = shape
+        if isinstance(dtype, np.dtype):
+            self._dtype = dtype
+        else:
+            self._dtype = np.dtype(dtype)
+        self._graph_key = graph_key
+        self._chunk_size = chunk_size
+        self._tensor = None
+        self._sealed = False
+        self._chunk_map = defaultdict(lambda: [])
+        self._record_type = np.dtype([("index", np.uint32), ("ts", np.dtype('datetime64[ns]')), ("value", self._dtype)])
+
+    @log_unhandled
+    def post_create(self):
+        from ..tensor.expressions.utils import create_fetch_tensor
+
+        super(MutableTensorActor, self).post_create()
+        self.set_cluster_info_ref()
+        self._tensor = create_fetch_tensor(self._chunk_size, self._shape, self._dtype)
+
+    def tensor_meta(self):
+        return self._shape, self._dtype, self._chunk_size, [c.key for c in self._tensor.chunks]
+
+    def tensor_key(self):
+        return self._tensor.key
+
+    def sealed(self):
+        return self._sealed
+
+    @log_unhandled
+    def read(self, tensor_index):
+        raise NotImplementedError
+
+    @log_unhandled
+    def append_chunk_records(self, chunk_records):
+        for chunk_key, record_chunk_key in chunk_records:
+            self._chunk_map[chunk_key].append(record_chunk_key)
+
+    @log_unhandled
+    def seal(self):
+        from ..worker.seal import SealActor
+        self._sealed = True
+        for chunk in self._tensor.chunks:
+            ep = self.get_scheduler(chunk.key)
+            sealer_uid = SealActor.gen_uid(self._session_id, chunk.key)
+            sealer_ref = self.ctx.create_actor(SealActor, uid=sealer_uid, address=ep)
+
+            sealer_ref.seal_chunk(self._session_id, self._graph_key,
+                                  chunk.key, self._chunk_map[chunk.key],
+                                  chunk.shape, self._record_type, self._dtype)
+        return self._graph_key, self._tensor.key, self._tensor.id, self.tensor_meta()

mars/scheduler/session.py

@@ -14,6 +14,7 @@
 
 import logging
 import os
+import uuid
 
 from .utils import SchedulerActor
 from ..utils import log_unhandled
@@ -33,6 +34,7 @@ def __init__(self, session_id, **kwargs):
         self._graph_refs = dict()
         self._graph_meta_refs = dict()
         self._tileable_to_graph = dict()
+        self._mut_tensor_refs = dict()
 
     @staticmethod
     def gen_uid(session_id):
@@ -87,6 +89,66 @@ def submit_tileable_graph(self, serialized_graph, graph_key, target_tileables=No
                 self._tileable_to_graph[tileable_key] = graph_ref
         return graph_ref
 
+    @log_unhandled
+    def create_mutable_tensor(self, name, shape, dtype, *args, **kwargs):
+        from .mutable import MutableTensorActor
+        if name in self._mut_tensor_refs:
+            raise ValueError("The mutable tensor named '%s' already exists." % name)
+        graph_key = uuid.uuid4()
+        mut_tensor_uid = MutableTensorActor.gen_uid(self._session_id, name)
+        mut_tensor_addr = self.get_scheduler(mut_tensor_uid)
+        mut_tensor_ref = self.ctx.create_actor(MutableTensorActor, self._session_id, name,
+                                               shape, dtype, graph_key, uid=mut_tensor_uid,
+                                               address=mut_tensor_addr, *args, **kwargs)
+        self._mut_tensor_refs[name] = mut_tensor_ref
+        return mut_tensor_ref.tensor_meta()
+
+    @log_unhandled
+    def get_mutable_tensor(self, name):
+        tensor_ref = self._mut_tensor_refs.get(name)
+        if tensor_ref is None or tensor_ref.sealed():
+            raise ValueError("The mutable tensor named '%s' doesn't exist, or has already been sealed." % name)
+        return tensor_ref.tensor_meta()
+
+    @log_unhandled
+    def append_chunk_records(self, name, chunk_records):
+        tensor_ref = self._mut_tensor_refs.get(name)
+        if tensor_ref is None or tensor_ref.sealed():
+            raise ValueError("The mutable tensor named '%s' doesn't exist, or has already been sealed." % name)
+        return tensor_ref.append_chunk_records(chunk_records)
+
+    @log_unhandled
+    def seal(self, name):
+        from .graph import GraphActor, GraphMetaActor
+        from .utils import GraphState
+
+        tensor_ref = self._mut_tensor_refs.get(name)
+        if tensor_ref is None or tensor_ref.sealed():
+            raise ValueError("The mutable tensor named '%s' doesn't exist, or has already been sealed." % name)
+
+        graph_key, tensor_key, tensor_id, tensor_meta = tensor_ref.seal()
+        shape, dtype, chunk_size, chunk_keys = tensor_meta
+
+        # Create a GraphActor
+        graph_uid = GraphActor.gen_uid(self._session_id, graph_key)
+        graph_addr = self.get_scheduler(graph_uid)
+
+        graph_ref = self.ctx.create_actor(GraphActor, self._session_id, graph_key,
+                                          serialized_tileable_graph=None,
+                                          state=GraphState.SUCCEEDED, final_state=GraphState.SUCCEEDED,
+                                          uid=graph_uid, address=graph_addr)
+        self._graph_refs[graph_key] = graph_ref
+        self._graph_meta_refs[graph_key] = self.ctx.actor_ref(
+            GraphMetaActor.gen_uid(self._session_id, graph_key), address=tensor_ref.__getstate__()[0])
+
+        # Add the tensor to the GraphActor
+        graph_ref.add_fetch_tileable(tensor_key, tensor_id, shape, dtype, chunk_size, chunk_keys)
+        self._tileable_to_graph[tensor_key] = graph_ref
+
+        # Clean up mutable tensor refs.
+        self._mut_tensor_refs.pop(name)
+        return graph_key, tensor_key, tensor_id, tensor_meta
+
     def graph_state(self, graph_key):
         return self._graph_refs[graph_key].get_state()
 

mars/session.py

@@ -190,6 +190,45 @@ def default_or_local(cls):
         cls._default_session = Session()
         return cls._default_session
 
+    def create_mutable_tensor(self, name, shape, dtype, *args, **kwargs):
+        from .tensor.core import MutableTensor, MutableTensorData
+        from .tensor.expressions.utils import create_fetch_tensor
+        self._ensure_local_cluster()
+        shape, dtype, chunk_size, chunk_keys = \
+                self._sess.create_mutable_tensor(name, shape, dtype, *args, **kwargs)
+        # Construct MutableTensor on the fly.
+        tensor = create_fetch_tensor(chunk_size, shape, dtype, chunk_keys=chunk_keys)
+        return MutableTensor(data=MutableTensorData(_name=name, _op=None, _shape=shape, _dtype=dtype,
+                                                    _nsplits=tensor.nsplits, _chunks=tensor.chunks))
+
+    def get_mutable_tensor(self, name):
+        from .tensor.core import MutableTensor, MutableTensorData
+        from .tensor.expressions.utils import create_fetch_tensor
+        self._ensure_local_cluster()
+        shape, dtype, chunk_size, chunk_keys = self._sess.get_mutable_tensor(name)
+        # Construct MutableTensor on the fly.
+        tensor = create_fetch_tensor(chunk_size, shape, dtype, chunk_keys=chunk_keys)
+        return MutableTensor(data=MutableTensorData(_name=name, _op=None, _shape=shape, _dtype=dtype,
+                                                    _nsplits=tensor.nsplits, _chunks=tensor.chunks))
+
+    def write_mutable_tensor(self, tensor, index, value):
+        self._ensure_local_cluster()
+        chunk_records_to_send = tensor._do_write(index, value)
+        return self._sess.send_chunk_records(tensor.name, chunk_records_to_send)
+
+    def seal(self, tensor):
+        from .tensor.expressions.utils import create_fetch_tensor
+        self._ensure_local_cluster()
+        chunk_records_to_send = tensor._do_flush()
+        self._sess.send_chunk_records(tensor.name, chunk_records_to_send)
+        shape, dtype, chunk_size, chunk_keys = self._sess.seal(tensor.name)
+        # Construct Tensor on the fly.
+        return create_fetch_tensor(chunk_size, shape, dtype, chunk_keys=chunk_keys)
+
+    def _ensure_local_cluster(self):
+        from .deploy.local.session import LocalClusterSession
+        if not isinstance(self._sess, LocalClusterSession):
+            raise RuntimeError("Only local cluster session can be used to manipulate mutable tensors.")
 
 def new_session(scheduler=None, **kwargs):
     return Session(scheduler, **kwargs)