Add _send_small_objects for sender to optimize sending small objects

mars/services/storage/handler.py

@@ -99,6 +99,9 @@ async def _get_data(self, data_info, conditions):
                 res = sliced_value
         raise mo.Return(res)
 
+    def get_client(self, level: StorageLevel):
+        return self._clients[level]
+
     @mo.extensible
     async def get(
         self,

mars/services/storage/tests/test_transfer.py

@@ -102,50 +102,84 @@ async def test_simple_transfer(create_actors):
     storage_handler1 = await mo.actor_ref(
         uid=StorageHandlerActor.gen_uid("numa-0"), address=worker_address_1
     )
+    data_manager1 = await mo.actor_ref(
+        uid=DataManagerActor.default_uid(), address=worker_address_1
+    )
     storage_handler2 = await mo.actor_ref(
         uid=StorageHandlerActor.gen_uid("numa-0"), address=worker_address_2
     )
+    data_manager2 = await mo.actor_ref(
+        uid=DataManagerActor.default_uid(), address=worker_address_2
+    )
 
     await storage_handler1.put(session_id, "data_key1", data1, StorageLevel.MEMORY)
     await storage_handler1.put(session_id, "data_key2", data2, StorageLevel.MEMORY)
     await storage_handler2.put(session_id, "data_key3", data2, StorageLevel.MEMORY)
 
-    sender_actor = await mo.actor_ref(
+    # sender_actor1 use default block_size
+    sender_actor1 = await mo.actor_ref(
         address=worker_address_1, uid=SenderManagerActor.gen_uid("numa-0")
     )
-
-    # send data to worker2 from worker1
-    await sender_actor.send_batch_data(
-        session_id,
-        ["data_key1"],
-        worker_address_2,
-        StorageLevel.MEMORY,
-        block_size=1000,
+    # send_actor2 set block_size to 0
+    sender_actor2 = await mo.create_actor(
+        SenderManagerActor,
+        "numa-0",
+        0,
+        data_manager1,
+        storage_handler1,
+        uid=SenderManagerActor.gen_uid("mock"),
+        address=worker_address_1,
     )
 
-    await sender_actor.send_batch_data(
-        session_id,
-        ["data_key2"],
-        worker_address_2,
-        StorageLevel.MEMORY,
-        block_size=1000,
-    )
+    for i, sender_actor in enumerate([sender_actor1, sender_actor2]):
+        # send data to worker2 from worker1
+        await sender_actor.send_batch_data(
+            session_id,
+            ["data_key1"],
+            worker_address_2,
+            StorageLevel.MEMORY,
+            block_size=1000,
+        )
 
-    get_data1 = await storage_handler2.get(session_id, "data_key1")
-    np.testing.assert_array_equal(data1, get_data1)
+        await sender_actor.send_batch_data(
+            session_id,
+            ["data_key2"],
+            worker_address_2,
+            StorageLevel.MEMORY,
+            block_size=1000,
+        )
+
+        get_data1 = await storage_handler2.get(session_id, "data_key1")
+        np.testing.assert_array_equal(data1, get_data1)
 
-    get_data2 = await storage_handler2.get(session_id, "data_key2")
-    pd.testing.assert_frame_equal(data2, get_data2)
+        get_data2 = await storage_handler2.get(session_id, "data_key2")
+        pd.testing.assert_frame_equal(data2, get_data2)
+        await storage_handler2.delete(session_id, "data_key1")
+        await storage_handler2.delete(session_id, "data_key2")
 
     # send data to worker1 from worker2
-    sender_actor = await mo.actor_ref(
+    sender_actor1 = await mo.actor_ref(
         address=worker_address_2, uid=SenderManagerActor.gen_uid("numa-0")
     )
-    await sender_actor.send_batch_data(
-        session_id, ["data_key3"], worker_address_1, StorageLevel.MEMORY
+    # send_actor2 set block_size to 0
+    sender_actor2 = await mo.create_actor(
+        SenderManagerActor,
+        "numa-0",
+        0,
+        data_manager2,
+        storage_handler2,
+        uid=SenderManagerActor.gen_uid("mock"),
+        address=worker_address_2,
     )
-    get_data3 = await storage_handler1.get(session_id, "data_key3")
-    pd.testing.assert_frame_equal(data2, get_data3)
+    for i, sender_actor in enumerate([sender_actor1, sender_actor2]):
+        # send data to worker1 from worker2
+        data_key = f"data_key3"
+        await sender_actor.send_batch_data(
+            session_id, [data_key], worker_address_1, StorageLevel.MEMORY
+        )
+        get_data3 = await storage_handler1.get(session_id, data_key)
+        pd.testing.assert_frame_equal(data2, get_data3)
+        await storage_handler1.delete(session_id, "data_key3")
 
 
 # test for cancelling happens when writing
@@ -232,7 +266,7 @@ async def test_cancel_transfer(create_actors, mock_sender, mock_receiver):
 
     send_task = asyncio.create_task(
         sender_actor.send_batch_data(
-            "mock", ["data_key1"], worker_address_2, StorageLevel.MEMORY
+            "mock", ["data_key1"], worker_address_2, StorageLevel.MEMORY, is_small_objects=False
         )
     )
 
@@ -250,7 +284,7 @@ async def test_cancel_transfer(create_actors, mock_sender, mock_receiver):
 
     send_task = asyncio.create_task(
         sender_actor.send_batch_data(
-            "mock", ["data_key1"], worker_address_2, StorageLevel.MEMORY
+            "mock", ["data_key1"], worker_address_2, StorageLevel.MEMORY, is_small_objects=False
         )
     )
     await send_task
@@ -261,12 +295,12 @@ async def test_cancel_transfer(create_actors, mock_sender, mock_receiver):
     if mock_sender is MockSenderManagerActor:
         send_task1 = asyncio.create_task(
             sender_actor.send_batch_data(
-                "mock", ["data_key2"], worker_address_2, StorageLevel.MEMORY
+                "mock", ["data_key2"], worker_address_2, StorageLevel.MEMORY, is_small_objects=False
             )
         )
         send_task2 = asyncio.create_task(
             sender_actor.send_batch_data(
-                "mock", ["data_key2"], worker_address_2, StorageLevel.MEMORY
+                "mock", ["data_key2"], worker_address_2, StorageLevel.MEMORY, is_small_objects=False
             )
         )
         await asyncio.sleep(0.5)

mars/services/storage/transfer.py

@@ -15,13 +15,13 @@
 import asyncio
 import logging
 from dataclasses import dataclass
-from typing import Dict, List
+from typing import Dict, List, Tuple, Union
 
 from ... import oscar as mo
 from ...lib.aio import alru_cache
 from ...storage import StorageLevel
 from ...utils import dataslots
-from .core import DataManagerActor, WrappedStorageFileObject
+from .core import DataManagerActor, WrappedStorageFileObject, DataInfo
 from .handler import StorageHandlerActor
 
 DEFAULT_TRANSFER_BLOCK_SIZE = 4 * 1024**2
@@ -65,6 +65,7 @@ async def _send_data(
         receiver_ref: mo.ActorRefType["ReceiverManagerActor"],
         session_id: str,
         data_keys: List[str],
+        data_infos: List[DataInfo],
         level: StorageLevel,
         block_size: int,
     ):
@@ -93,11 +94,12 @@ async def send(self, buffer, eof_mark, key):
 
         sender = BufferedSender()
         open_reader_tasks = []
-        for data_key in data_keys:
+        storage_client = await self._storage_handler.get_client(level)
+        for info in data_infos:
             open_reader_tasks.append(
-                self._storage_handler.open_reader.delay(session_id, data_key)
+                storage_client.open_reader(info.object_id)
             )
-        readers = await self._storage_handler.open_reader.batch(*open_reader_tasks)
+        readers = await asyncio.gather(*open_reader_tasks)
 
         for data_key, reader in zip(data_keys, readers):
             while True:
@@ -116,7 +118,58 @@ async def send(self, buffer, eof_mark, key):
                     break
         await sender.flush()
 
-    @mo.extensible
+    async def _send(
+        self,
+        session_id: str,
+        data_keys: List[Union[str, Tuple]],
+        data_infos: List[DataInfo],
+        data_sizes: List[int],
+        block_size: int,
+        address: str,
+        band_name: str,
+        level: StorageLevel,
+    ):
+        receiver_ref: mo.ActorRefType[ReceiverManagerActor] = await self.get_receiver_ref(address, band_name)
+        is_transferring_list = await receiver_ref.open_writers(
+            session_id, data_keys, data_sizes, level
+        )
+        to_send_keys = []
+        to_send_infos = []
+        to_wait_keys = []
+        for data_key, is_transferring, info in zip(
+            data_keys, is_transferring_list, data_infos
+        ):
+            if is_transferring:
+                to_wait_keys.append(data_key)
+            else:
+                to_send_keys.append(data_key)
+                to_send_infos.append(info)
+
+        if to_send_keys:
+            await self._send_data(
+                receiver_ref, session_id, to_send_keys, to_send_infos, level, block_size
+            )
+        if to_wait_keys:
+            await receiver_ref.wait_transfer_done(session_id, to_wait_keys)
+
+    async def _send_small_objects(
+        self,
+        session_id: str,
+        data_keys: List[Union[str, Tuple]],
+        data_infos: List[DataInfo],
+        address: str,
+        band_name: str,
+        level: StorageLevel,
+    ):
+        # simple get all objects and send them all to receiver
+        storage_client = await self._storage_handler.get_client(level)
+        get_tasks = [
+            storage_client.get(info.object_id) for info in data_infos
+        ]
+        data_list = list(await asyncio.gather(*get_tasks))
+        receiver_ref: mo.ActorRefType[ReceiverManagerActor] = await self.get_receiver_ref(address, band_name)
+        await receiver_ref.put_small_objects(session_id, data_keys, data_list, level)
+
     async def send_batch_data(
         self,
         session_id: str,
@@ -125,15 +178,13 @@ async def send_batch_data(
         level: StorageLevel,
         band_name: str = "numa-0",
         block_size: int = None,
+        is_small_objects=None,
         error: str = "raise",
     ):
         logger.debug(
             "Begin to send data (%s, %s) to %s", session_id, data_keys, address
         )
         block_size = block_size or self._transfer_block_size
-        receiver_ref: mo.ActorRefType[
-            ReceiverManagerActor
-        ] = await self.get_receiver_ref(address, band_name)
         get_infos = []
         pin_tasks = []
         for data_key in data_keys:
@@ -162,23 +213,29 @@ async def send_batch_data(
         data_sizes = [info.store_size for info in infos]
         if level is None:
             level = infos[0].level
-        is_transferring_list = await receiver_ref.open_writers(
-            session_id, data_keys, data_sizes, level
-        )
-        to_send_keys = []
-        to_wait_keys = []
-        for data_key, is_transferring in zip(data_keys, is_transferring_list):
-            if is_transferring:
-                to_wait_keys.append(data_key)
-            else:
-                to_send_keys.append(data_key)
-
-        if to_send_keys:
-            await self._send_data(
-                receiver_ref, session_id, to_send_keys, level, block_size
+        total_size = sum(data_sizes)
+        if is_small_objects is None:
+            is_small_objects = total_size <= block_size
+        if is_small_objects:
+            logger.debug(
+                "Choose send_small_objects method for sending data of %s bytes",
+                total_size,
+            )
+            await self._send_small_objects(
+                session_id, data_keys, infos, address, band_name, level
+            )
+        else:
+            logger.debug("Choose block method for sending data of %s bytes", total_size)
+            await self._send(
+                session_id,
+                data_keys,
+                infos,
+                data_sizes,
+                block_size,
+                address,
+                band_name,
+                level,
             )
-        if to_wait_keys:
-            await receiver_ref.wait_transfer_done(session_id, to_wait_keys)
         unpin_tasks = []
         for data_key in data_keys:
             unpin_tasks.append(
@@ -232,6 +289,15 @@ def _decref_writing_key(self, session_id: str, data_key: str):
         if self._writing_infos[(session_id, data_key)].ref_counts == 0:
             del self._writing_infos[(session_id, data_key)]
 
+    async def put_small_objects(
+        self, session_id: str, data_keys: List[str], objects: List, level: StorageLevel
+    ):
+        tasks = [
+            self._storage_handler.put.delay(session_id, data_key, obj, level)
+            for data_key, obj in zip(data_keys, objects)
+        ]
+        await self._storage_handler.put.batch(*tasks)
+
     async def create_writers(
         self,
         session_id: str,