Repository: mars-project/mars

Issue:
Support configuring meta and broadcast in batch
Currently ChunkMetaActor does not support configuring broadcast targets in batch. It can be done to reduce schedule latency. Also, LocalChunkMetaActor need to support configuring meta in batch. This is useful for shuffle.


Discussion:

Resolving patch:
diff --git a/mars/scheduler/chunkmeta.py b/mars/scheduler/chunkmeta.py
--- a/mars/scheduler/chunkmeta.py
+++ b/mars/scheduler/chunkmeta.py
@@ -193,6 +193,17 @@ def set_chunk_broadcasts(self, session_id, chunk_key, broadcast_dests):
         self._meta_broadcasts[(session_id, chunk_key)] = \
             [d for d in broadcast_dests if d != self.address]
 
+    def batch_set_chunk_broadcasts(self, session_id, chunk_keys, broadcast_dests):
+        """
+        Configure broadcast destinations in batch
+        :param session_id: session id
+        :param chunk_keys: chunk key
+        :param broadcast_dests:
+        :return:
+        """
+        for key, dests in zip(chunk_keys, broadcast_dests):
+            self.set_chunk_broadcasts(session_id, key, dests)
+
     def get_chunk_broadcasts(self, session_id, chunk_key):
         """
         Get chunk broadcast addresses, for test only
@@ -201,7 +212,8 @@ def get_chunk_broadcasts(self, session_id, chunk_key):
         """
         return self._meta_broadcasts.get((session_id, chunk_key))
 
-    def set_chunk_meta(self, session_id, chunk_key, size=None, shape=None, workers=None):
+    def set_chunk_meta(self, session_id, chunk_key, size=None, shape=None, workers=None,
+                       broadcast=True):
         """
         Update chunk meta in current storage
         :param session_id: session id
@@ -209,6 +221,7 @@ def set_chunk_meta(self, session_id, chunk_key, size=None, shape=None, workers=N
         :param size: size of the chunk
         :param shape: shape of the chunk
         :param workers: workers holding the chunk
+        :param broadcast: broadcast meta into registered destinations
         """
         query_key = (session_id, chunk_key)
         # update input with existing value
@@ -238,24 +251,54 @@ def set_chunk_meta(self, session_id, chunk_key, size=None, shape=None, workers=N
 
         # broadcast into pre-determined destinations
         futures = []
-        if query_key in self._meta_broadcasts:
+        if broadcast and query_key in self._meta_broadcasts:
             for dest in self._meta_broadcasts[query_key]:
                 futures.append(
                     self.ctx.actor_ref(self.default_name(), address=dest)
-                        .cache_chunk_meta(session_id, chunk_key, meta, _wait=False, _tell=True)
+                        .batch_cache_chunk_meta(session_id, [chunk_key], [meta], _wait=False, _tell=True)
                 )
             [f.result() for f in futures]
 
-    def cache_chunk_meta(self, session_id, chunk_key, meta):
+    def batch_set_chunk_meta(self, session_id, keys, metas):
+        """
+        Set chunk metas in batch
+        :param session_id: session id
+        :param keys: keys to set
+        :param metas: metas to set
+        """
+        query_dict = defaultdict(lambda: (list(), list()))
+
+        for key, meta in zip(keys, metas):
+            self.set_chunk_meta(session_id, key, size=meta.chunk_size, shape=meta.chunk_shape,
+                                workers=meta.workers, broadcast=False)
+            try:
+                dests = self._meta_broadcasts[(session_id, key)]
+            except KeyError:
+                continue
+
+            for dest in dests:
+                query_dict[dest][0].append(key)
+                query_dict[dest][1].append(meta)
+
+        futures = []
+        for dest, (chunk_keys, metas) in query_dict.items():
+            futures.append(
+                self.ctx.actor_ref(self.default_name(), address=dest)
+                    .batch_cache_chunk_meta(session_id, chunk_keys, metas, _wait=False, _tell=True)
+            )
+        [f.result() for f in futures]
+
+    def batch_cache_chunk_meta(self, session_id, chunk_keys, metas):
         """
         Receive updates for caching
 
         :param session_id: session id
-        :param chunk_key: chunk key
-        :param meta: meta data
+        :param chunk_keys: chunk keys
+        :param metas: meta data
         """
-        query_key = (session_id, chunk_key)
-        self._meta_cache[query_key] = meta
+        for chunk_key, meta in zip(chunk_keys, metas):
+            query_key = (session_id, chunk_key)
+            self._meta_cache[query_key] = meta
 
     def get_chunk_meta(self, session_id, chunk_key):
         """
@@ -361,6 +404,17 @@ def set_chunk_broadcasts(self, session_id, chunk_key, broadcast_dests):
         self.ctx.actor_ref(LocalChunkMetaActor.default_name(), address=addr) \
             .set_chunk_broadcasts(session_id, chunk_key, broadcast_dests, _tell=True, _wait=False)
 
+    def batch_set_chunk_broadcasts(self, session_id, chunk_keys, broadcast_dests):
+        query_chunk = defaultdict(lambda: (list(), list()))
+        for key, dests in zip(chunk_keys, broadcast_dests):
+            addr = self.get_scheduler((session_id, key))
+            query_chunk[addr][0].append(key)
+            query_chunk[addr][1].append(dests)
+
+        for addr, (chunk_keys, dest_groups) in query_chunk.items():
+            self.ctx.actor_ref(LocalChunkMetaActor.default_name(), address=addr) \
+                .batch_set_chunk_broadcasts(session_id, chunk_keys, dest_groups)
+
     def set_chunk_meta(self, session_id, chunk_key, size=None, shape=None, workers=None):
         """
         Update chunk metadata
diff --git a/mars/scheduler/operand.py b/mars/scheduler/operand.py
--- a/mars/scheduler/operand.py
+++ b/mars/scheduler/operand.py
@@ -430,9 +430,13 @@ def _handle_worker_accept(self, worker):
         broadcast_eps.difference_update({self.address})
         broadcast_eps = tuple(broadcast_eps)
 
+        chunk_keys, broadcast_ep_groups = [], []
         for chunk_key in self._chunks:
-            self._chunk_meta_ref.set_chunk_broadcasts(
-                self._session_id, chunk_key, broadcast_eps, _tell=True, _wait=False)
+            chunk_keys.append(chunk_key)
+            broadcast_ep_groups.append(broadcast_eps)
+
+        self._chunk_meta_ref.batch_set_chunk_broadcasts(
+            self._session_id, chunk_keys, broadcast_ep_groups, _tell=True, _wait=False)
 
         # submit job
         logger.debug('Start running operand %s on %s', self._op_key, worker)


Test patch:
diff --git a/mars/scheduler/tests/test_chunkmeta.py b/mars/scheduler/tests/test_chunkmeta.py
--- a/mars/scheduler/tests/test_chunkmeta.py
+++ b/mars/scheduler/tests/test_chunkmeta.py
@@ -126,7 +126,8 @@ def _mock_get_scheduler(key):
                 key1 = str(uuid.uuid4())
                 key2 = str(uuid.uuid4())
                 key3 = str(uuid.uuid4())
-                keys = [key1, key2, key3]
+                key4 = str(uuid.uuid4())
+                keys = [key1, key2, key3, key4]
                 ref1.set_chunk_size(session1, key1, 512)
                 ref2.set_chunk_size(session1, key2, 1024)
                 ref2.set_chunk_size(session2, key3, 1024)
@@ -181,6 +182,11 @@ def _mock_get_scheduler(key):
                 self.assertIsNone(ref1.batch_get_chunk_size(session1, [key1, key2])[1])
                 self.assertIsNone(ref1.batch_get_workers(session1, [key1, key2])[1])
 
+                meta4 = WorkerMeta(chunk_size=512, chunk_shape=(10,) * 2, workers=(endpoints[0],))
+                loc_ref2.batch_set_chunk_meta(session1, [key4], [meta4])
+                self.assertEqual(loc_ref2.get_chunk_meta(session1, key4).chunk_size, 512)
+                self.assertEqual(loc_ref2.get_chunk_meta(session1, key4).chunk_shape, (10,) * 2)
+
     @unittest.skipIf(sys.platform == 'win32', 'Currently not support multiple pools under Windows')
     @patch_method(ChunkMetaActor.get_scheduler)
     def testChunkBroadcast(self, *_):
@@ -210,7 +216,8 @@ def _mock_get_scheduler(key):
 
                 key1 = str(uuid.uuid4())
                 key2 = str(uuid.uuid4())
-                keys = [key1, key2]
+                key3 = str(uuid.uuid4())
+                keys = [key1, key2, key3]
 
                 ref1.set_chunk_broadcasts(session_id, key1, [endpoints[1]])
                 ref1.set_chunk_size(session_id, key1, 512)
@@ -229,6 +236,12 @@ def _mock_get_scheduler(key):
                 self.assertEqual(local_ref2.get_chunk_meta(session_id, key1).chunk_shape, (10,) * 2)
                 self.assertEqual(local_ref2.get_chunk_broadcasts(session_id, key2), [endpoints[0]])
 
+                ref1.batch_set_chunk_broadcasts(session_id, [key3], [[endpoints[1]]])
+                meta3 = WorkerMeta(chunk_size=512, chunk_shape=(10,) * 2, workers=(endpoints[0],))
+                local_ref1.batch_set_chunk_meta(session_id, [key3], [meta3])
+                self.assertEqual(local_ref2.get_chunk_meta(session_id, key3).chunk_size, 512)
+                self.assertEqual(local_ref2.get_chunk_meta(session_id, key3).chunk_shape, (10,) * 2)
+
                 ref1.delete_meta(session_id, key1)
                 pool2.sleep(0.1)