Add batch support for broadcasts and local meta setting (#166)

Author: wjsi

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,14 +212,16 @@ 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
         :param chunk_key: chunk key
         :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

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)

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)