From 4048a5bfb5a9a3672ffbf1712837f0e630ff5fc5 Mon Sep 17 00:00:00 2001 From: UranusSeven <109661872+UranusSeven@users.noreply.github.com> Date: Thu, 1 Jun 2023 16:32:55 +0800 Subject: [PATCH 01/25] REF: refactor transfer to avoid deadlocks --- .../xorbits/_mars/services/storage/handler.py | 25 +- .../services/storage/tests/test_transfer.py | 218 ++++++++++-------- .../_mars/services/storage/transfer.py | 73 ++---- 3 files changed, 156 insertions(+), 160 deletions(-) diff --git a/python/xorbits/_mars/services/storage/handler.py b/python/xorbits/_mars/services/storage/handler.py index 0cdaed519..ddcf591d0 100644 --- a/python/xorbits/_mars/services/storage/handler.py +++ b/python/xorbits/_mars/services/storage/handler.py @@ -531,19 +531,30 @@ async def _fetch_via_transfer( fetch_band_name: str, error: str, ): - from .transfer import SenderManagerActor + from .transfer import ReceiverManagerActor, SenderManagerActor logger.debug("Begin to fetch %s from band %s", data_keys, remote_band) + + receiver_ref: mo.ActorRefType[ReceiverManagerActor] = await mo.actor_ref( + address=self.address, + uid=ReceiverManagerActor.gen_uid(fetch_band_name), + ) + is_transferring_list = await receiver_ref.open_writers( + session_id, data_keys, level, remote_band, error + ) + if not is_transferring_list: + # no data to be transferred + return + sender_ref: mo.ActorRefType[SenderManagerActor] = await mo.actor_ref( address=remote_band[0], uid=SenderManagerActor.gen_uid(remote_band[1]) ) await sender_ref.send_batch_data( session_id, data_keys, + is_transferring_list, self._data_manager_ref.address, - level, fetch_band_name, - error=error, ) logger.debug("Finish fetching %s from band %s", data_keys, remote_band) @@ -559,10 +570,8 @@ async def fetch_batch( if error not in ("raise", "ignore"): # pragma: no cover raise ValueError("error must be raise or ignore") - meta_api = await self._get_meta_api(session_id) remote_keys = defaultdict(set) missing_keys = [] - get_metas = [] get_info_delays = [] for data_key in data_keys: get_info_delays.append( @@ -586,6 +595,9 @@ async def fetch_batch( else: # Not exists in local, fetch from remote worker missing_keys.append(data_key) + await self._data_manager_ref.pin.batch(*pin_delays) + + meta_api = await self._get_meta_api(session_id) if address is None or band_name is None: # some mapper keys are absent, specify error='ignore' # remember that meta only records those main keys @@ -599,9 +611,6 @@ async def fetch_batch( ) for data_key in missing_keys ] - await self._data_manager_ref.pin.batch(*pin_delays) - - if get_metas: metas = await meta_api.get_chunk_meta.batch(*get_metas) else: # pragma: no cover metas = [{"bands": [(address, band_name)]}] * len(missing_keys) diff --git a/python/xorbits/_mars/services/storage/tests/test_transfer.py b/python/xorbits/_mars/services/storage/tests/test_transfer.py index 1b6ac4e60..96b00b6e5 100644 --- a/python/xorbits/_mars/services/storage/tests/test_transfer.py +++ b/python/xorbits/_mars/services/storage/tests/test_transfer.py @@ -17,7 +17,7 @@ import os import sys import tempfile -from typing import List +from typing import List, Union import numpy as np import pandas as pd @@ -111,42 +111,65 @@ async def test_simple_transfer(create_actors): await storage_handler1.put(session_id, "data_key2", data2, level) await storage_handler2.put(session_id, "data_key3", data2, level) - sender_actor = await mo.actor_ref( - address=worker_address_1, uid=SenderManagerActor.gen_uid("numa-0") - ) + receiver_actor: mo.ActorRefType[ReceiverManagerActor] = await mo.actor_ref( + address=worker_address_2, uid=ReceiverManagerActor.gen_uid("numa-0") + ) + sender_actor: mo.ActorRefType[SenderManagerActor] = 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, - level, - block_size=1000, - ) + is_transferring = await receiver_actor.open_writers( + session_id, + ["data_key1", "data_key2"], + StorageLevel.MEMORY, + (sender_actor.address, "numa-0"), + "raise", + ) - await sender_actor.send_batch_data( - session_id, - ["data_key2"], - worker_address_2, - level, - block_size=1000, - ) + # send data to worker2 from worker1 + await sender_actor.send_batch_data( + session_id, + ["data_key1"], + is_transferring, + (receiver_actor.address, "numa-0"), + 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"], + is_transferring, + (receiver_actor.address, "numa-0"), + block_size=1000, + ) - get_data2 = await storage_handler2.get(session_id, "data_key2") - pd.testing.assert_frame_equal(data2, get_data2) + get_data1 = await storage_handler2.get(session_id, "data_key1") + np.testing.assert_array_equal(data1, get_data1) - # send data to worker1 from worker2 - sender_actor = 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, level - ) - get_data3 = await storage_handler1.get(session_id, "data_key3") - pd.testing.assert_frame_equal(data2, get_data3) + get_data2 = await storage_handler2.get(session_id, "data_key2") + pd.testing.assert_frame_equal(data2, get_data2) + + # send data to worker1 from worker2 + receiver_actor: mo.ActorRefType[ReceiverManagerActor] = await mo.actor_ref( + address=worker_address_1, uid=ReceiverManagerActor.gen_uid("numa-0") + ) + sender_actor: mo.ActorRefType[SenderManagerActor] = await mo.actor_ref( + address=worker_address_2, uid=SenderManagerActor.gen_uid("numa-0") + ) + + is_transferring = await receiver_actor.open_writers( + session_id, + ["data_key3"], + StorageLevel.MEMORY, + (sender_actor.address, "numa-0"), + "raise", + ) + + await sender_actor.send_batch_data( + session_id, ["data_key3"], is_transferring, (receiver_actor.address, "numa-0") + ) + get_data3 = await storage_handler1.get(session_id, "data_key3") + pd.testing.assert_frame_equal(data2, get_data3) # test for cancelling happens when writing @@ -201,16 +224,17 @@ async def get_receiver_ref(address: str, band_name: str): @pytest.mark.parametrize( - "mock_sender, mock_receiver", + "mock_sender_cls, mock_receiver_cls", [ (MockSenderManagerActor, MockReceiverManagerActor), (MockSenderManagerActor2, MockReceiverManagerActor2), ], ) @pytest.mark.asyncio -async def test_cancel_transfer(create_actors, mock_sender, mock_receiver): +async def test_cancel_transfer(create_actors, mock_sender_cls, mock_receiver_cls): worker_address_1, worker_address_2 = create_actors + session_id = "mock_session" quota_refs = { StorageLevel.MEMORY: await mo.actor_ref( StorageQuotaActor, @@ -230,73 +254,85 @@ async def test_cancel_transfer(create_actors, mock_sender, mock_receiver): uid=StorageHandlerActor.gen_uid("numa-0"), address=worker_address_2 ) - sender_actor = await mo.create_actor( - mock_sender, + sender_actor: mo.ActorRefType[SenderManagerActor] = await mo.create_actor( + mock_sender_cls, data_manager_ref=data_manager_ref, - uid=mock_sender.default_uid(), + uid=mock_sender_cls.default_uid(), address=worker_address_1, - allocate_strategy=IdleLabel("io", "mock_sender"), + allocate_strategy=IdleLabel("io", "mock_sender_cls"), ) - await mo.create_actor( - mock_receiver, + receiver_actor: mo.ActorRefType[ReceiverManagerActor] = await mo.create_actor( + mock_receiver_cls, quota_refs, - uid=mock_receiver.default_uid(), + uid=mock_receiver_cls.default_uid(), address=worker_address_2, - allocate_strategy=IdleLabel("io", "mock_receiver"), + allocate_strategy=IdleLabel("io", "mock_receiver_cls"), ) data1 = np.random.rand(10, 10) - await storage_handler1.put("mock", "data_key1", data1, StorageLevel.MEMORY) + await storage_handler1.put(session_id, "data_key1", data1, StorageLevel.MEMORY) data2 = pd.DataFrame(np.random.rand(100, 100)) - await storage_handler1.put("mock", "data_key2", data2, StorageLevel.MEMORY) + await storage_handler1.put(session_id, "data_key2", data2, StorageLevel.MEMORY) used_before = (await quota_refs[StorageLevel.MEMORY].get_quota())[1] - send_task = asyncio.create_task( - sender_actor.send_batch_data( - "mock", ["data_key1"], worker_address_2, StorageLevel.MEMORY + async def transfer( + data_keys: List[Union[str, tuple]], + sender: mo.ActorRefType[SenderManagerActor], + receiver: mo.ActorRefType[ReceiverManagerActor], + ): + is_transferring_list = await receiver.open_writers( + session_id, + data_keys, + StorageLevel.MEMORY, + (sender.address, "numa-0"), + "raise", ) + assert len(is_transferring_list) > 0 + await sender.send_batch_data( + session_id, + data_keys, + is_transferring_list, + (receiver.address, "numa-0"), + ) + + transfer_task = asyncio.create_task( + transfer(["data_key1"], sender_actor, receiver_actor) ) await asyncio.sleep(0.5) - send_task.cancel() + transfer_task.cancel() with pytest.raises(asyncio.CancelledError): - await send_task + await transfer_task used = (await quota_refs[StorageLevel.MEMORY].get_quota())[1] assert used == used_before with pytest.raises(DataNotExist): - await storage_handler2.get("mock", "data_key1") + await storage_handler2.get(session_id, "data_key1") - send_task = asyncio.create_task( - sender_actor.send_batch_data( - "mock", ["data_key1"], worker_address_2, StorageLevel.MEMORY - ) + transfer_task = asyncio.create_task( + transfer(["data_key1"], sender_actor, receiver_actor) ) - await send_task - get_data = await storage_handler2.get("mock", "data_key1") + await transfer_task + get_data = await storage_handler2.get(session_id, "data_key1") np.testing.assert_array_equal(data1, get_data) # cancel when fetch the same data Simultaneously - if mock_sender is MockSenderManagerActor: - send_task1 = asyncio.create_task( - sender_actor.send_batch_data( - "mock", ["data_key2"], worker_address_2, StorageLevel.MEMORY - ) + if mock_sender_cls is MockSenderManagerActor: + transfer_task1 = asyncio.create_task( + transfer(["data_key2"], sender_actor, receiver_actor) ) - send_task2 = asyncio.create_task( - sender_actor.send_batch_data( - "mock", ["data_key2"], worker_address_2, StorageLevel.MEMORY - ) + transfer_task2 = asyncio.create_task( + transfer(["data_key2"], sender_actor, receiver_actor) ) await asyncio.sleep(0.5) - send_task1.cancel() + transfer_task1.cancel() with pytest.raises(asyncio.CancelledError): - await send_task1 - await send_task2 - get_data2 = await storage_handler2.get("mock", "data_key2") + await transfer_task1 + await transfer_task2 + get_data2 = await storage_handler2.get(session_id, "data_key2") pd.testing.assert_frame_equal(get_data2, data2) @@ -309,34 +345,30 @@ async def test_transfer_same_data(create_actors): storage_handler1 = await mo.actor_ref( uid=StorageHandlerActor.gen_uid("numa-0"), address=worker_address_1 ) - storage_handler2 = await mo.actor_ref( + await mo.actor_ref( uid=StorageHandlerActor.gen_uid("numa-0"), address=worker_address_2 ) await storage_handler1.put(session_id, "data_key1", data1, StorageLevel.MEMORY) - sender_actor = await mo.actor_ref( - address=worker_address_1, uid=SenderManagerActor.gen_uid("numa-0") - ) - # send data to worker2 from worker1 - task1 = asyncio.create_task( - sender_actor.send_batch_data( - session_id, - ["data_key1"], - worker_address_2, - StorageLevel.MEMORY, - block_size=1000, - ) + receiver_actor: mo.ActorRefType[ReceiverManagerActor] = await mo.actor_ref( + address=worker_address_2, uid=ReceiverManagerActor.gen_uid("numa-0") ) - task2 = asyncio.create_task( - sender_actor.send_batch_data( - session_id, - ["data_key1"], - worker_address_2, - StorageLevel.MEMORY, - block_size=1000, - ) + task1 = receiver_actor.open_writers( + session_id, + ["data_key1"], + StorageLevel.MEMORY, + (worker_address_1, "numa-0"), + "raise", ) - await asyncio.gather(task1, task2) - get_data1 = await storage_handler2.get(session_id, "data_key1") - np.testing.assert_array_equal(data1, get_data1) + task2 = receiver_actor.open_writers( + session_id, + ["data_key1"], + StorageLevel.MEMORY, + (worker_address_1, "numa-0"), + "raise", + ) + + results = await asyncio.gather(task1, task2) + assert [False] in results + assert [True] in results diff --git a/python/xorbits/_mars/services/storage/transfer.py b/python/xorbits/_mars/services/storage/transfer.py index 7acda794c..2dff1ae19 100644 --- a/python/xorbits/_mars/services/storage/transfer.py +++ b/python/xorbits/_mars/services/storage/transfer.py @@ -24,6 +24,7 @@ from ...lib.aio import alru_cache from ...storage import StorageLevel +from ...typing import BandType from ...utils import dataslots from .core import DataManagerActor, WrappedStorageFileObject from .handler import StorageHandlerActor @@ -152,62 +153,12 @@ async def send_batch_data( self, session_id: str, data_keys: List[str], - address: str, - level: StorageLevel, - band_name: str = "numa-0", + is_transferring_list: List[bool], + remote_band: BandType, block_size: int = None, - error: str = "raise", ): logger.debug( - "Begin to send data (%s, %s) to %s", session_id, data_keys, address - ) - - tasks = [] - for key in data_keys: - tasks.append(self._data_manager_ref.get_store_key.delay(session_id, key)) - data_keys = await self._data_manager_ref.get_store_key.batch(*tasks) - data_keys = list(set(data_keys)) - sub_infos = await self._data_manager_ref.get_sub_infos.batch( - *[ - self._data_manager_ref.get_sub_infos.delay(session_id, key) - for key in data_keys - ] - ) - - 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: - get_infos.append( - self._data_manager_ref.get_data_info.delay( - session_id, data_key, self._band_name, error - ) - ) - pin_tasks.append( - self._data_manager_ref.pin.delay( - session_id, data_key, self._band_name, error - ) - ) - await self._data_manager_ref.pin.batch(*pin_tasks) - infos = await self._data_manager_ref.get_data_info.batch(*get_infos) - filtered = [ - (data_info, data_key) - for data_info, data_key in zip(infos, data_keys) - if data_info is not None - ] - if filtered: - infos, data_keys = zip(*filtered) - else: # pragma: no cover - # no data to be transferred - return - 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, sub_infos, band_name + "Begin to send data (%s, %s) to %s", session_id, data_keys, remote_band ) to_send_keys = [] to_wait_keys = [] @@ -217,7 +168,12 @@ async def send_batch_data( else: to_send_keys.append(data_key) + receiver_ref: mo.ActorRefType[ + ReceiverManagerActor + ] = await self.get_receiver_ref(remote_band[0], remote_band[1]) + if to_send_keys: + block_size = block_size or self._transfer_block_size await self._send_data(receiver_ref, session_id, to_send_keys, block_size) if to_wait_keys: await receiver_ref.wait_transfer_done(session_id, to_wait_keys) @@ -230,11 +186,10 @@ async def send_batch_data( ) await self._data_manager_ref.unpin.batch(*unpin_tasks) logger.debug( - "Finish sending data (%s, %s) to %s, total size is %s", + "Finish sending data (%s, %s) to %s", session_id, data_keys, - address, - sum(data_sizes), + remote_band, ) @@ -336,7 +291,7 @@ async def create_writers( level: StorageLevel, sub_infos: List, band_name: str, - ): + ) -> List[bool]: tasks = dict() key_to_sub_infos = dict() data_key_to_size = dict() @@ -374,10 +329,10 @@ async def open_writers( session_id: str, data_keys: List[str], data_sizes: List[int], - level: StorageLevel, sub_infos: List, + level: StorageLevel, band_name: str, - ): + ) -> List[bool]: async with self._lock: await self._storage_handler.request_quota_with_spill(level, sum(data_sizes)) future = asyncio.create_task( From be78817219757968eaacc3f22dbef6911121c12d Mon Sep 17 00:00:00 2001 From: UranusSeven <109661872+UranusSeven@users.noreply.github.com> Date: Fri, 2 Jun 2023 14:20:57 +0800 Subject: [PATCH 02/25] Debugging --- python/xorbits/_mars/services/storage/handler.py | 5 +++-- python/xorbits/_mars/services/storage/transfer.py | 4 ++++ 2 files changed, 7 insertions(+), 2 deletions(-) diff --git a/python/xorbits/_mars/services/storage/handler.py b/python/xorbits/_mars/services/storage/handler.py index ddcf591d0..125d0fb26 100644 --- a/python/xorbits/_mars/services/storage/handler.py +++ b/python/xorbits/_mars/services/storage/handler.py @@ -546,6 +546,8 @@ async def _fetch_via_transfer( # no data to be transferred return + logger.debug("Writers have been opened for %s", data_keys) + sender_ref: mo.ActorRefType[SenderManagerActor] = await mo.actor_ref( address=remote_band[0], uid=SenderManagerActor.gen_uid(remote_band[1]) ) @@ -553,8 +555,7 @@ async def _fetch_via_transfer( session_id, data_keys, is_transferring_list, - self._data_manager_ref.address, - fetch_band_name, + (self._data_manager_ref.address, fetch_band_name), ) logger.debug("Finish fetching %s from band %s", data_keys, remote_band) diff --git a/python/xorbits/_mars/services/storage/transfer.py b/python/xorbits/_mars/services/storage/transfer.py index 2dff1ae19..35c0bcf86 100644 --- a/python/xorbits/_mars/services/storage/transfer.py +++ b/python/xorbits/_mars/services/storage/transfer.py @@ -173,10 +173,14 @@ async def send_batch_data( ] = await self.get_receiver_ref(remote_band[0], remote_band[1]) if to_send_keys: + logger.debug("Start sending %s", to_send_keys) block_size = block_size or self._transfer_block_size await self._send_data(receiver_ref, session_id, to_send_keys, block_size) + logger.debug("Done sending %s", to_send_keys) if to_wait_keys: + logger.debug("Start waiting %s", to_wait_keys) await receiver_ref.wait_transfer_done(session_id, to_wait_keys) + logger.debug("Done waiting %s", to_wait_keys) unpin_tasks = [] for data_key in data_keys: unpin_tasks.append( From 640864a88ea73b48d29e69984b44bc034dbca84a Mon Sep 17 00:00:00 2001 From: UranusSeven <109661872+UranusSeven@users.noreply.github.com> Date: Fri, 2 Jun 2023 14:44:48 +0800 Subject: [PATCH 03/25] Debugging --- .../_mars/services/storage/transfer.py | 58 +++++++++++++++++++ 1 file changed, 58 insertions(+) diff --git a/python/xorbits/_mars/services/storage/transfer.py b/python/xorbits/_mars/services/storage/transfer.py index 35c0bcf86..8d6661c45 100644 --- a/python/xorbits/_mars/services/storage/transfer.py +++ b/python/xorbits/_mars/services/storage/transfer.py @@ -336,7 +336,65 @@ async def open_writers( sub_infos: List, level: StorageLevel, band_name: str, + remote_band, + error ) -> List[bool]: + logger.debug("Start opening writers for %s", data_keys) + + remote_data_manager_ref: mo.ActorRefType[DataManagerActor] = await mo.actor_ref( + address=remote_band[0], uid=DataManagerActor.default_uid() + ) + + logger.debug("Getting actual keys for %s", data_keys) + tasks = [] + for key in data_keys: + tasks.append(remote_data_manager_ref.get_store_key.delay(session_id, key)) + data_keys = await remote_data_manager_ref.get_store_key.batch(*tasks) + data_keys = list(set(data_keys)) + + logger.debug("Getting sub infos for %s", data_keys) + sub_infos = await remote_data_manager_ref.get_sub_infos.batch( + *[ + remote_data_manager_ref.get_sub_infos.delay(session_id, key) + for key in data_keys + ] + ) + + get_infos = [] + pin_tasks = [] + for data_key in data_keys: + get_infos.append( + remote_data_manager_ref.get_data_info.delay( + session_id, data_key, remote_band[1], error + ) + ) + pin_tasks.append( + remote_data_manager_ref.pin.delay( + session_id, data_key, remote_band[1], error + ) + ) + + logger.debug("Pining %s", data_keys) + await remote_data_manager_ref.pin.batch(*pin_tasks) + logger.debug("Getting data infos for %s", data_keys) + infos = await remote_data_manager_ref.get_data_info.batch(*get_infos) + + filtered = [ + (data_info, data_key) + for data_info, data_key in zip(infos, data_keys) + if data_info is not None + ] + if filtered: + infos, data_keys = zip(*filtered) + else: # pragma: no cover + # no data to be transferred + return [] + data_sizes = [info.store_size for info in infos] + + if level is None: + level = infos[0].level + + logger.debug("Opening writers for %s", data_keys) async with self._lock: await self._storage_handler.request_quota_with_spill(level, sum(data_sizes)) future = asyncio.create_task( From 6ab37e1ce34397e9c41dcfc28096223f9783a607 Mon Sep 17 00:00:00 2001 From: UranusSeven <109661872+UranusSeven@users.noreply.github.com> Date: Fri, 2 Jun 2023 14:49:32 +0800 Subject: [PATCH 04/25] Debugging --- python/xorbits/_mars/services/storage/transfer.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/python/xorbits/_mars/services/storage/transfer.py b/python/xorbits/_mars/services/storage/transfer.py index 8d6661c45..5503e467c 100644 --- a/python/xorbits/_mars/services/storage/transfer.py +++ b/python/xorbits/_mars/services/storage/transfer.py @@ -394,9 +394,10 @@ async def open_writers( if level is None: level = infos[0].level - logger.debug("Opening writers for %s", data_keys) async with self._lock: + logger.debug("Requesting quota for %s", data_keys) await self._storage_handler.request_quota_with_spill(level, sum(data_sizes)) + logger.debug("Opening writers for %s", data_keys) future = asyncio.create_task( self.create_writers( session_id, data_keys, data_sizes, level, sub_infos, band_name From ad42bc8e10586b86d34162afa5694c2bedccc18f Mon Sep 17 00:00:00 2001 From: UranusSeven <109661872+UranusSeven@users.noreply.github.com> Date: Fri, 2 Jun 2023 15:06:58 +0800 Subject: [PATCH 05/25] Fix --- .../xorbits/_mars/services/storage/handler.py | 62 ++++++++++++++++++- .../_mars/services/storage/transfer.py | 61 +----------------- 2 files changed, 63 insertions(+), 60 deletions(-) diff --git a/python/xorbits/_mars/services/storage/handler.py b/python/xorbits/_mars/services/storage/handler.py index 125d0fb26..ab42e7b67 100644 --- a/python/xorbits/_mars/services/storage/handler.py +++ b/python/xorbits/_mars/services/storage/handler.py @@ -535,12 +535,72 @@ async def _fetch_via_transfer( logger.debug("Begin to fetch %s from band %s", data_keys, remote_band) + remote_data_manager_ref: mo.ActorRefType[DataManagerActor] = await mo.actor_ref( + address=remote_band[0], uid=DataManagerActor.default_uid() + ) + + logger.debug("Getting actual keys for %s", data_keys) + tasks = [] + for key in data_keys: + tasks.append(remote_data_manager_ref.get_store_key.delay(session_id, key)) + data_keys = await remote_data_manager_ref.get_store_key.batch(*tasks) + data_keys = list(set(data_keys)) + + logger.debug("Getting sub infos for %s", data_keys) + sub_infos = await remote_data_manager_ref.get_sub_infos.batch( + *[ + remote_data_manager_ref.get_sub_infos.delay(session_id, key) + for key in data_keys + ] + ) + + get_infos = [] + pin_tasks = [] + for data_key in data_keys: + get_infos.append( + remote_data_manager_ref.get_data_info.delay( + session_id, data_key, remote_band[1], error + ) + ) + pin_tasks.append( + remote_data_manager_ref.pin.delay( + session_id, data_key, remote_band[1], error + ) + ) + + logger.debug("Pining %s", data_keys) + await remote_data_manager_ref.pin.batch(*pin_tasks) + logger.debug("Getting data infos for %s", data_keys) + infos = await remote_data_manager_ref.get_data_info.batch(*get_infos) + + filtered = [ + (data_info, data_key) + for data_info, data_key in zip(infos, data_keys) + if data_info is not None + ] + if filtered: + infos, data_keys = zip(*filtered) + else: # pragma: no cover + # no data to be transferred + return [] + data_sizes = [info.store_size for info in infos] + + if level is None: + level = infos[0].level + + logger.debug("Requesting quota for %s", data_keys) + await self.request_quota_with_spill(level, sum(data_sizes)) + receiver_ref: mo.ActorRefType[ReceiverManagerActor] = await mo.actor_ref( address=self.address, uid=ReceiverManagerActor.gen_uid(fetch_band_name), ) is_transferring_list = await receiver_ref.open_writers( - session_id, data_keys, level, remote_band, error + session_id, + data_keys, + data_sizes, + sub_infos, + level, ) if not is_transferring_list: # no data to be transferred diff --git a/python/xorbits/_mars/services/storage/transfer.py b/python/xorbits/_mars/services/storage/transfer.py index 5503e467c..b45a45086 100644 --- a/python/xorbits/_mars/services/storage/transfer.py +++ b/python/xorbits/_mars/services/storage/transfer.py @@ -333,70 +333,13 @@ async def open_writers( session_id: str, data_keys: List[str], data_sizes: List[int], - sub_infos: List, + sub_infos, level: StorageLevel, - band_name: str, - remote_band, - error + band_name ) -> List[bool]: logger.debug("Start opening writers for %s", data_keys) - remote_data_manager_ref: mo.ActorRefType[DataManagerActor] = await mo.actor_ref( - address=remote_band[0], uid=DataManagerActor.default_uid() - ) - - logger.debug("Getting actual keys for %s", data_keys) - tasks = [] - for key in data_keys: - tasks.append(remote_data_manager_ref.get_store_key.delay(session_id, key)) - data_keys = await remote_data_manager_ref.get_store_key.batch(*tasks) - data_keys = list(set(data_keys)) - - logger.debug("Getting sub infos for %s", data_keys) - sub_infos = await remote_data_manager_ref.get_sub_infos.batch( - *[ - remote_data_manager_ref.get_sub_infos.delay(session_id, key) - for key in data_keys - ] - ) - - get_infos = [] - pin_tasks = [] - for data_key in data_keys: - get_infos.append( - remote_data_manager_ref.get_data_info.delay( - session_id, data_key, remote_band[1], error - ) - ) - pin_tasks.append( - remote_data_manager_ref.pin.delay( - session_id, data_key, remote_band[1], error - ) - ) - - logger.debug("Pining %s", data_keys) - await remote_data_manager_ref.pin.batch(*pin_tasks) - logger.debug("Getting data infos for %s", data_keys) - infos = await remote_data_manager_ref.get_data_info.batch(*get_infos) - - filtered = [ - (data_info, data_key) - for data_info, data_key in zip(infos, data_keys) - if data_info is not None - ] - if filtered: - infos, data_keys = zip(*filtered) - else: # pragma: no cover - # no data to be transferred - return [] - data_sizes = [info.store_size for info in infos] - - if level is None: - level = infos[0].level - async with self._lock: - logger.debug("Requesting quota for %s", data_keys) - await self._storage_handler.request_quota_with_spill(level, sum(data_sizes)) logger.debug("Opening writers for %s", data_keys) future = asyncio.create_task( self.create_writers( From f434eab3fc6afd8dcbf25a70641996cd0eed2fba Mon Sep 17 00:00:00 2001 From: UranusSeven <109661872+UranusSeven@users.noreply.github.com> Date: Fri, 2 Jun 2023 15:13:41 +0800 Subject: [PATCH 06/25] Debugging --- python/xorbits/_mars/services/storage/transfer.py | 1 + 1 file changed, 1 insertion(+) diff --git a/python/xorbits/_mars/services/storage/transfer.py b/python/xorbits/_mars/services/storage/transfer.py index b45a45086..6394e7071 100644 --- a/python/xorbits/_mars/services/storage/transfer.py +++ b/python/xorbits/_mars/services/storage/transfer.py @@ -317,6 +317,7 @@ async def create_writers( being_processed.append(True) self._writing_infos[(session_id, data_key)].ref_counts += 1 if tasks: + logger.debug("Executing open_writer.batch") writers = await self._storage_handler.open_writer.batch( *tuple(tasks.values()) ) From 4d80dbb98dce3466616a73faba508114decc4d15 Mon Sep 17 00:00:00 2001 From: UranusSeven <109661872+UranusSeven@users.noreply.github.com> Date: Tue, 13 Jun 2023 12:29:08 +0800 Subject: [PATCH 07/25] Fix --- python/xorbits/_mars/services/storage/handler.py | 7 +++++++ python/xorbits/_mars/services/storage/transfer.py | 10 +++++++++- 2 files changed, 16 insertions(+), 1 deletion(-) diff --git a/python/xorbits/_mars/services/storage/handler.py b/python/xorbits/_mars/services/storage/handler.py index ab42e7b67..77929fc87 100644 --- a/python/xorbits/_mars/services/storage/handler.py +++ b/python/xorbits/_mars/services/storage/handler.py @@ -611,9 +611,16 @@ async def _fetch_via_transfer( sender_ref: mo.ActorRefType[SenderManagerActor] = await mo.actor_ref( address=remote_band[0], uid=SenderManagerActor.gen_uid(remote_band[1]) ) + + open_reader_tasks = [] + for data_key in data_keys: + open_reader_tasks.append(self.open_reader.delay(session_id, data_key)) + readers = await self.open_reader.batch(*open_reader_tasks) + await sender_ref.send_batch_data( session_id, data_keys, + readers, is_transferring_list, (self._data_manager_ref.address, fetch_band_name), ) diff --git a/python/xorbits/_mars/services/storage/transfer.py b/python/xorbits/_mars/services/storage/transfer.py index 6394e7071..081c4eafe 100644 --- a/python/xorbits/_mars/services/storage/transfer.py +++ b/python/xorbits/_mars/services/storage/transfer.py @@ -18,6 +18,7 @@ from dataclasses import dataclass from io import UnsupportedOperation from typing import Any, Dict, List, Union +from typing import TYPE_CHECKING, Dict, List import xoscar as mo from xoscar.core import BufferRef @@ -29,6 +30,10 @@ from .core import DataManagerActor, WrappedStorageFileObject from .handler import StorageHandlerActor +if TYPE_CHECKING: + from ...storage.core import StorageFileObject + + DEFAULT_TRANSFER_BLOCK_SIZE = 4 * 1024**2 @@ -153,6 +158,7 @@ async def send_batch_data( self, session_id: str, data_keys: List[str], + readers: List["StorageFileObject"], is_transferring_list: List[bool], remote_band: BandType, block_size: int = None, @@ -175,7 +181,9 @@ async def send_batch_data( if to_send_keys: logger.debug("Start sending %s", to_send_keys) block_size = block_size or self._transfer_block_size - await self._send_data(receiver_ref, session_id, to_send_keys, block_size) + await self._send_data( + receiver_ref, readers, session_id, to_send_keys, block_size + ) logger.debug("Done sending %s", to_send_keys) if to_wait_keys: logger.debug("Start waiting %s", to_wait_keys) From 30f58a7f48250a0a04096703e064a3d40cd57c9f Mon Sep 17 00:00:00 2001 From: UranusSeven <109661872+UranusSeven@users.noreply.github.com> Date: Wed, 14 Jun 2023 10:40:59 +0800 Subject: [PATCH 08/25] Fix --- .../xorbits/_mars/services/storage/handler.py | 46 +++++++++-------- .../_mars/services/storage/transfer.py | 49 +++++++++++++++++-- 2 files changed, 70 insertions(+), 25 deletions(-) diff --git a/python/xorbits/_mars/services/storage/handler.py b/python/xorbits/_mars/services/storage/handler.py index 77929fc87..badf38780 100644 --- a/python/xorbits/_mars/services/storage/handler.py +++ b/python/xorbits/_mars/services/storage/handler.py @@ -554,10 +554,10 @@ async def _fetch_via_transfer( ] ) - get_infos = [] + get_info_tasks = [] pin_tasks = [] for data_key in data_keys: - get_infos.append( + get_info_tasks.append( remote_data_manager_ref.get_data_info.delay( session_id, data_key, remote_band[1], error ) @@ -567,11 +567,10 @@ async def _fetch_via_transfer( session_id, data_key, remote_band[1], error ) ) - + logger.debug("Getting data infos for %s", data_keys) + infos = await remote_data_manager_ref.get_data_info.batch(*get_info_tasks) logger.debug("Pining %s", data_keys) await remote_data_manager_ref.pin.batch(*pin_tasks) - logger.debug("Getting data infos for %s", data_keys) - infos = await remote_data_manager_ref.get_data_info.batch(*get_infos) filtered = [ (data_info, data_key) @@ -588,34 +587,41 @@ async def _fetch_via_transfer( if level is None: level = infos[0].level - logger.debug("Requesting quota for %s", data_keys) - await self.request_quota_with_spill(level, sum(data_sizes)) - + logger.debug("Creating writers for %s", data_keys) receiver_ref: mo.ActorRefType[ReceiverManagerActor] = await mo.actor_ref( address=self.address, uid=ReceiverManagerActor.gen_uid(fetch_band_name), ) - is_transferring_list = await receiver_ref.open_writers( - session_id, - data_keys, - data_sizes, - sub_infos, - level, + open_writer_tasks = [] + for data_key, data_size, sub_info in zip(data_keys, data_sizes, sub_infos): + open_writer_tasks.append( + self.open_writer.delay( + session_id, data_key, data_size, level, request_quota=False + ) + ) + writers = await self.open_writer.batch(*open_writer_tasks) + is_transferring_list = await receiver_ref.add_writers( + session_id, data_keys, data_sizes, sub_infos, writers, level ) - if not is_transferring_list: - # no data to be transferred - return - - logger.debug("Writers have been opened for %s", data_keys) + required_quota = sum( + [ + data_size + for is_transferring, data_size in zip(is_transferring_list, data_sizes) + if is_transferring + ] + ) + await self.request_quota_with_spill(level, required_quota) + logger.debug("Writers have been created for %s", data_keys) + logger.debug("Creating readers for %s", data_keys) sender_ref: mo.ActorRefType[SenderManagerActor] = await mo.actor_ref( address=remote_band[0], uid=SenderManagerActor.gen_uid(remote_band[1]) ) - open_reader_tasks = [] for data_key in data_keys: open_reader_tasks.append(self.open_reader.delay(session_id, data_key)) readers = await self.open_reader.batch(*open_reader_tasks) + logger.debug("Readers have been created for %s", data_keys) await sender_ref.send_batch_data( session_id, diff --git a/python/xorbits/_mars/services/storage/transfer.py b/python/xorbits/_mars/services/storage/transfer.py index 081c4eafe..e483e8bf1 100644 --- a/python/xorbits/_mars/services/storage/transfer.py +++ b/python/xorbits/_mars/services/storage/transfer.py @@ -19,6 +19,7 @@ from io import UnsupportedOperation from typing import Any, Dict, List, Union from typing import TYPE_CHECKING, Dict, List +from typing import TYPE_CHECKING, Dict, List, Optional, Union import xoscar as mo from xoscar.core import BufferRef @@ -222,15 +223,16 @@ def __init__( storage_handler_ref: mo.ActorRefType[StorageHandlerActor] = None, ): self._quota_refs = quota_refs - self._storage_handler = storage_handler_ref + # self._storage_handler = storage_handler_ref self._writing_infos: Dict[tuple, WritingInfo] = dict() self._lock = asyncio.Lock() async def __post_create__(self): - if self._storage_handler is None: # for test - self._storage_handler = await mo.actor_ref( - self.address, StorageHandlerActor.gen_uid("numa-0") - ) + # if self._storage_handler is None: # for test + # self._storage_handler = await mo.actor_ref( + # self.address, StorageHandlerActor.gen_uid("numa-0") + # ) + pass async def get_buffers( self, @@ -290,11 +292,48 @@ async def handle_transmission_cancellation(self, session_id: str, data_keys: Lis def gen_uid(cls, band_name: str): return f"receiver_manager_{band_name}" + def lock(self): + self._lock.acquire() + + def unlock(self): + self._lock.release() + def _decref_writing_key(self, session_id: str, data_key: str): self._writing_infos[(session_id, data_key)].ref_counts -= 1 if self._writing_infos[(session_id, data_key)].ref_counts == 0: del self._writing_infos[(session_id, data_key)] + def is_writing_info_existed( + self, session_id: str, data_keys: List[Union[str, tuple]] + ) -> List[bool]: + return [(session_id, data_key) in self._writing_infos for data_key in data_keys] + + async def add_writers( + self, + session_id: str, + data_keys: List[Union[str, tuple]], + data_sizes: List[int], + sub_infos: List, + writers: List[Optional[WrappedStorageFileObject]], + level: StorageLevel, + ) -> List[bool]: + is_transferring: List[bool] = [] + async with self._lock: + for data_key, data_size, sub_info, writer in zip( + data_keys, data_sizes, sub_infos, writers + ): + if (session_id, data_key) not in self._writing_infos: + is_transferring.append(False) + + self._writing_infos[(session_id, data_key)] = WritingInfo( + writer, data_size, level, asyncio.Event(), 1 + ) + if sub_info is not None: + writer._sub_key_infos = sub_info + else: + is_transferring.append(True) + return is_transferring + async def create_writers( self, session_id: str, From 6c13816706ba7cb52a7626e7e35fd837992ffbe2 Mon Sep 17 00:00:00 2001 From: UranusSeven <109661872+UranusSeven@users.noreply.github.com> Date: Thu, 15 Jun 2023 06:45:53 +0000 Subject: [PATCH 09/25] checkpoint --- python/xorbits/_mars/services/storage/core.py | 9 +++++++ .../xorbits/_mars/services/storage/handler.py | 25 +++++++++++-------- .../_mars/services/storage/transfer.py | 15 +++-------- 3 files changed, 27 insertions(+), 22 deletions(-) diff --git a/python/xorbits/_mars/services/storage/core.py b/python/xorbits/_mars/services/storage/core.py index cb064ef5d..f70a1b5ff 100644 --- a/python/xorbits/_mars/services/storage/core.py +++ b/python/xorbits/_mars/services/storage/core.py @@ -93,6 +93,12 @@ async def clean_up(self): self._file.close() async def close(self): + logger.debug( + "Writer closed for %s, %s on %s", + self._session_id, + self._data_key, + self._data_manager.address, + ) self._file.close() if self._object_id is None: # for some backends like vineyard, @@ -322,6 +328,9 @@ def delete_data_info( level: StorageLevel, band_name: str, ): + logger.debug( + "Deleting %s, %s from %s, %s", session_id, data_key, level, band_name + ) if (session_id, data_key) in self._data_key_to_infos: self._data_info_list[level, band_name].pop((session_id, data_key)) self._spill_strategy[level, band_name].record_delete_info( diff --git a/python/xorbits/_mars/services/storage/handler.py b/python/xorbits/_mars/services/storage/handler.py index badf38780..1dda35835 100644 --- a/python/xorbits/_mars/services/storage/handler.py +++ b/python/xorbits/_mars/services/storage/handler.py @@ -66,6 +66,7 @@ def __init__( self._quota_refs = quota_refs self._band_name = band_name self._supervisor_address = None + self._lock = asyncio.Lock() @classmethod def gen_uid(cls, band_name: str): @@ -284,6 +285,7 @@ async def delete_object( level: StorageLevel, ): data_key = await self._data_manager_ref.get_store_key(session_id, data_key) + logger.debug("Delete object %s, %s on %s", session_id, data_key, self.address) await self._data_manager_ref.delete_data_info( session_id, data_key, level, self._band_name ) @@ -292,6 +294,7 @@ async def delete_object( @mo.extensible async def delete(self, session_id: str, data_key: str, error: str = "raise"): + logger.debug("Delete %s, %s on %s", session_id, data_key, self.address) if error not in ("raise", "ignore"): # pragma: no cover raise ValueError("error must be raise or ignore") @@ -367,6 +370,9 @@ async def batch_delete(self, args_list, kwargs_list): session_id, key, level, info.band ) ) + logger.debug( + "Batch delete %s, %s on %s", session_id, key, self.address + ) to_removes.append((level, info.object_id)) level_sizes[level] += info.store_size @@ -607,28 +613,26 @@ async def _fetch_via_transfer( [ data_size for is_transferring, data_size in zip(is_transferring_list, data_sizes) - if is_transferring + if not is_transferring ] ) await self.request_quota_with_spill(level, required_quota) logger.debug("Writers have been created for %s", data_keys) - logger.debug("Creating readers for %s", data_keys) + logger.debug( + "Start transferring %s from %s to %s", + data_keys, + remote_band, + (self.address, fetch_band_name), + ) sender_ref: mo.ActorRefType[SenderManagerActor] = await mo.actor_ref( address=remote_band[0], uid=SenderManagerActor.gen_uid(remote_band[1]) ) - open_reader_tasks = [] - for data_key in data_keys: - open_reader_tasks.append(self.open_reader.delay(session_id, data_key)) - readers = await self.open_reader.batch(*open_reader_tasks) - logger.debug("Readers have been created for %s", data_keys) - await sender_ref.send_batch_data( session_id, data_keys, - readers, is_transferring_list, - (self._data_manager_ref.address, fetch_band_name), + (self.address, fetch_band_name), ) logger.debug("Finish fetching %s from band %s", data_keys, remote_band) @@ -692,6 +696,7 @@ async def fetch_batch( for data_key, bands in zip(missing_keys, metas): if bands is not None: remote_keys[bands["bands"][0]].add(data_key) + transfer_tasks = [] fetch_keys = [] for band, keys in remote_keys.items(): diff --git a/python/xorbits/_mars/services/storage/transfer.py b/python/xorbits/_mars/services/storage/transfer.py index e483e8bf1..4273e9a8b 100644 --- a/python/xorbits/_mars/services/storage/transfer.py +++ b/python/xorbits/_mars/services/storage/transfer.py @@ -17,9 +17,7 @@ import logging from dataclasses import dataclass from io import UnsupportedOperation -from typing import Any, Dict, List, Union -from typing import TYPE_CHECKING, Dict, List -from typing import TYPE_CHECKING, Dict, List, Optional, Union +from typing import Dict, List, Optional, Union, Any import xoscar as mo from xoscar.core import BufferRef @@ -31,10 +29,6 @@ from .core import DataManagerActor, WrappedStorageFileObject from .handler import StorageHandlerActor -if TYPE_CHECKING: - from ...storage.core import StorageFileObject - - DEFAULT_TRANSFER_BLOCK_SIZE = 4 * 1024**2 @@ -159,7 +153,6 @@ async def send_batch_data( self, session_id: str, data_keys: List[str], - readers: List["StorageFileObject"], is_transferring_list: List[bool], remote_band: BandType, block_size: int = None, @@ -180,11 +173,9 @@ async def send_batch_data( ] = await self.get_receiver_ref(remote_band[0], remote_band[1]) if to_send_keys: - logger.debug("Start sending %s", to_send_keys) + logger.debug("Start sending %s to %s", to_send_keys, receiver_ref.address) block_size = block_size or self._transfer_block_size - await self._send_data( - receiver_ref, readers, session_id, to_send_keys, block_size - ) + await self._send_data(receiver_ref, session_id, to_send_keys, block_size) logger.debug("Done sending %s", to_send_keys) if to_wait_keys: logger.debug("Start waiting %s", to_wait_keys) From 8fea60f8a0301091c2ade9f764aca26d58ed7890 Mon Sep 17 00:00:00 2001 From: ChengjieLi Date: Mon, 7 Aug 2023 14:42:58 +0800 Subject: [PATCH 10/25] fix --- .../xorbits/_mars/services/storage/handler.py | 102 ++++-- .../services/storage/tests/test_transfer.py | 318 ++++++++++-------- .../storage/tests/test_transfer_gpu.py | 68 ++-- .../_mars/services/storage/transfer.py | 119 +------ python/xorbits/_mars/storage/cuda.py | 10 +- 5 files changed, 295 insertions(+), 322 deletions(-) diff --git a/python/xorbits/_mars/services/storage/handler.py b/python/xorbits/_mars/services/storage/handler.py index 1dda35835..3986ac580 100644 --- a/python/xorbits/_mars/services/storage/handler.py +++ b/python/xorbits/_mars/services/storage/handler.py @@ -388,6 +388,7 @@ async def batch_delete(self, args_list, kwargs_list): await self._quota_refs[level].release_quota(size) @mo.extensible + @mo.no_lock async def open_reader(self, session_id: str, data_key: str) -> StorageFileObject: data_info = await self._data_manager_ref.get_data_info( session_id, data_key, self._band_name @@ -396,6 +397,7 @@ async def open_reader(self, session_id: str, data_key: str) -> StorageFileObject return reader @open_reader.batch + @mo.no_lock async def batch_open_readers(self, args_list, kwargs_list): get_data_infos = [] for args, kwargs in zip(args_list, kwargs_list): @@ -528,7 +530,21 @@ async def _fetch_remote( await self._data_manager_ref.put_data_info.batch(*put_data_info_delays) await asyncio.gather(*fetch_tasks) - async def _fetch_via_transfer( + async def get_receive_manager_ref(self, band_name: str): + from .transfer import ReceiverManagerActor + + return await mo.actor_ref( + address=self.address, + uid=ReceiverManagerActor.gen_uid(band_name), + ) + + @staticmethod + async def get_send_manager_ref(address: str, band: str): + from .transfer import SenderManagerActor + + return await mo.actor_ref(address=address, uid=SenderManagerActor.gen_uid(band)) + + async def fetch_via_transfer( self, session_id: str, data_keys: List[Union[str, tuple]], @@ -593,31 +609,44 @@ async def _fetch_via_transfer( if level is None: level = infos[0].level - logger.debug("Creating writers for %s", data_keys) - receiver_ref: mo.ActorRefType[ReceiverManagerActor] = await mo.actor_ref( - address=self.address, - uid=ReceiverManagerActor.gen_uid(fetch_band_name), - ) + receiver_ref: mo.ActorRefType[ + ReceiverManagerActor + ] = await self.get_receive_manager_ref(fetch_band_name) + + await self.request_quota_with_spill(level, sum(data_sizes)) + open_writer_tasks = [] for data_key, data_size, sub_info in zip(data_keys, data_sizes, sub_infos): open_writer_tasks.append( self.open_writer.delay( - session_id, data_key, data_size, level, request_quota=False + session_id, + data_key, + data_size, + level, + request_quota=False, + band_name=fetch_band_name, ) ) writers = await self.open_writer.batch(*open_writer_tasks) is_transferring_list = await receiver_ref.add_writers( session_id, data_keys, data_sizes, sub_infos, writers, level ) - required_quota = sum( - [ - data_size - for is_transferring, data_size in zip(is_transferring_list, data_sizes) - if not is_transferring - ] - ) - await self.request_quota_with_spill(level, required_quota) - logger.debug("Writers have been created for %s", data_keys) + + to_send_keys = [] + to_wait_keys = [] + wait_sizes = [] + for data_key, is_transferring, _size in zip( + data_keys, is_transferring_list, data_sizes + ): + if is_transferring: + to_wait_keys.append(data_key) + wait_sizes.append(_size) + else: + to_send_keys.append(data_key) + + # Overapplied the quota for these wait keys, and now need to update the quota + if to_wait_keys: + self._quota_refs[level].update_quota(-sum(wait_sizes)) logger.debug( "Start transferring %s from %s to %s", @@ -625,16 +654,35 @@ async def _fetch_via_transfer( remote_band, (self.address, fetch_band_name), ) - sender_ref: mo.ActorRefType[SenderManagerActor] = await mo.actor_ref( - address=remote_band[0], uid=SenderManagerActor.gen_uid(remote_band[1]) - ) - await sender_ref.send_batch_data( - session_id, - data_keys, - is_transferring_list, - (self.address, fetch_band_name), - ) - logger.debug("Finish fetching %s from band %s", data_keys, remote_band) + sender_ref: mo.ActorRefType[ + SenderManagerActor + ] = await self.get_send_manager_ref(remote_band[0], remote_band[1]) + + try: + await sender_ref.send_batch_data( + session_id, + data_keys, + to_send_keys, + to_wait_keys, + (self.address, fetch_band_name), + ) + await receiver_ref.handle_transmission_done(session_id, to_send_keys) + except asyncio.CancelledError: + keys_to_delete = await receiver_ref.handle_transmission_cancellation( + session_id, to_send_keys + ) + for key in keys_to_delete: + await self.delete(session_id, key, error="ignore") + raise + + unpin_tasks = [] + for data_key in data_keys: + unpin_tasks.append( + remote_data_manager_ref.unpin.delay( + session_id, [data_key], remote_band[1], error="ignore" + ) + ) + await remote_data_manager_ref.unpin.batch(*unpin_tasks) async def fetch_batch( self, @@ -708,7 +756,7 @@ async def fetch_batch( else: # fetch via transfer transfer_tasks.append( - self._fetch_via_transfer( + self.fetch_via_transfer( session_id, list(keys), level, band, band_name or band[1], error ) ) diff --git a/python/xorbits/_mars/services/storage/tests/test_transfer.py b/python/xorbits/_mars/services/storage/tests/test_transfer.py index 96b00b6e5..b8bcf0d25 100644 --- a/python/xorbits/_mars/services/storage/tests/test_transfer.py +++ b/python/xorbits/_mars/services/storage/tests/test_transfer.py @@ -17,17 +17,18 @@ import os import sys import tempfile -from typing import List, Union +from typing import Dict, List import numpy as np import pandas as pd import pytest import xoscar as mo from xoscar.backends.allocate_strategy import IdleLabel +from xoscar.errors import NoIdleSlot from ....oscar import create_actor_pool from ....storage import StorageLevel -from ..core import DataManagerActor, StorageManagerActor, StorageQuotaActor +from ..core import StorageManagerActor, StorageQuotaActor from ..errors import DataNotExist from ..handler import StorageHandlerActor from ..transfer import ReceiverManagerActor, SenderManagerActor @@ -62,10 +63,22 @@ async def start_pool(): await worker_pool_2.stop() +async def _get_io_address(pool): + pool_config = (await mo.get_pool_config(pool.external_address)).as_dict() + return [ + v["external_address"][0] + for k, v in pool_config["pools"].items() + if v["label"] == "io" + ][0] + + @pytest.fixture async def create_actors(actor_pools): worker_pool_1, worker_pool_2 = actor_pools + io1 = await _get_io_address(worker_pool_1) + io2 = await _get_io_address(worker_pool_2) + tmp_dir = tempfile.mkdtemp() storage_configs = {"shared_memory": {}, "disk": {"root_dirs": f"{tmp_dir}"}} @@ -82,7 +95,39 @@ async def create_actors(actor_pools): uid=StorageManagerActor.default_uid(), address=worker_pool_2.external_address, ) - yield worker_pool_1.external_address, worker_pool_2.external_address + yield worker_pool_1.external_address, worker_pool_2.external_address, io1, io2 + try: + await mo.destroy_actor(manager_ref1) + await mo.destroy_actor(manager_ref2) + except FileNotFoundError: + pass + assert not os.path.exists(tmp_dir) + + +@pytest.fixture +async def create_actors_mock(actor_pools): + worker_pool_1, worker_pool_2 = actor_pools + + io1 = await _get_io_address(worker_pool_1) + io2 = await _get_io_address(worker_pool_2) + + tmp_dir = tempfile.mkdtemp() + storage_configs = {"shared_memory": {}, "disk": {"root_dirs": f"{tmp_dir}"}} + + manager_ref1 = await mo.create_actor( + MockStorageManagerActor, + storage_configs, + uid=MockStorageManagerActor.default_uid(), + address=worker_pool_1.external_address, + ) + + manager_ref2 = await mo.create_actor( + MockStorageManagerActor, + storage_configs, + uid=MockStorageManagerActor.default_uid(), + address=worker_pool_2.external_address, + ) + yield worker_pool_1.external_address, worker_pool_2.external_address, io1, io2 try: await mo.destroy_actor(manager_ref1) await mo.destroy_actor(manager_ref2) @@ -93,16 +138,16 @@ async def create_actors(actor_pools): @pytest.mark.asyncio async def test_simple_transfer(create_actors): - worker_address_1, worker_address_2 = create_actors + worker_address_1, worker_address_2, io1, io2 = create_actors data1 = np.random.rand(100, 100) data2 = pd.DataFrame(np.random.randint(0, 100, (500, 10))) - storage_handler1 = await mo.actor_ref( - uid=StorageHandlerActor.gen_uid("numa-0"), address=worker_address_1 + storage_handler1: mo.ActorRefType[StorageHandlerActor] = await mo.actor_ref( + uid=StorageHandlerActor.gen_uid("numa-0"), address=io1 ) - storage_handler2 = await mo.actor_ref( - uid=StorageHandlerActor.gen_uid("numa-0"), address=worker_address_2 + storage_handler2: mo.ActorRefType[StorageHandlerActor] = await mo.actor_ref( + uid=StorageHandlerActor.gen_uid("numa-0"), address=io2 ) for level in (StorageLevel.MEMORY, StorageLevel.DISK): @@ -111,65 +156,26 @@ async def test_simple_transfer(create_actors): await storage_handler1.put(session_id, "data_key2", data2, level) await storage_handler2.put(session_id, "data_key3", data2, level) - receiver_actor: mo.ActorRefType[ReceiverManagerActor] = await mo.actor_ref( - address=worker_address_2, uid=ReceiverManagerActor.gen_uid("numa-0") - ) - sender_actor: mo.ActorRefType[SenderManagerActor] = await mo.actor_ref( - address=worker_address_1, uid=SenderManagerActor.gen_uid("numa-0") - ) - - is_transferring = await receiver_actor.open_writers( - session_id, - ["data_key1", "data_key2"], - StorageLevel.MEMORY, - (sender_actor.address, "numa-0"), - "raise", - ) - - # send data to worker2 from worker1 - await sender_actor.send_batch_data( - session_id, - ["data_key1"], - is_transferring, - (receiver_actor.address, "numa-0"), - block_size=1000, - ) - - await sender_actor.send_batch_data( - session_id, - ["data_key2"], - is_transferring, - (receiver_actor.address, "numa-0"), - block_size=1000, - ) - - get_data1 = await storage_handler2.get(session_id, "data_key1") - np.testing.assert_array_equal(data1, get_data1) + await storage_handler2.fetch_via_transfer( + session_id, + ["data_key1", "data_key2"], + level, + (io1, "numa-0"), + "numa-0", + "raise", + ) - get_data2 = await storage_handler2.get(session_id, "data_key2") - pd.testing.assert_frame_equal(data2, get_data2) + get_data1 = await storage_handler2.get(session_id, "data_key1") + np.testing.assert_array_equal(data1, get_data1) - # send data to worker1 from worker2 - receiver_actor: mo.ActorRefType[ReceiverManagerActor] = await mo.actor_ref( - address=worker_address_1, uid=ReceiverManagerActor.gen_uid("numa-0") - ) - sender_actor: mo.ActorRefType[SenderManagerActor] = await mo.actor_ref( - address=worker_address_2, uid=SenderManagerActor.gen_uid("numa-0") - ) - - is_transferring = await receiver_actor.open_writers( - session_id, - ["data_key3"], - StorageLevel.MEMORY, - (sender_actor.address, "numa-0"), - "raise", - ) + get_data2 = await storage_handler2.get(session_id, "data_key2") + pd.testing.assert_frame_equal(data2, get_data2) - await sender_actor.send_batch_data( - session_id, ["data_key3"], is_transferring, (receiver_actor.address, "numa-0") - ) - get_data3 = await storage_handler1.get(session_id, "data_key3") - pd.testing.assert_frame_equal(data2, get_data3) + await storage_handler1.fetch_via_transfer( + session_id, ["data_key3"], level, (io2, "numa-0"), "numa-0", "raise" + ) + get_data3 = await storage_handler1.get(session_id, "data_key3") + pd.testing.assert_frame_equal(data2, get_data3) # test for cancelling happens when writing @@ -181,12 +187,12 @@ class MockSenderManagerActor(SenderManagerActor): @staticmethod async def get_receiver_ref(address: str, band_name: str): return await mo.actor_ref( - address=address, uid=MockReceiverManagerActor.default_uid() + address=address, uid=MockReceiverManagerActor.gen_uid(band_name) ) async def _copy_to_receiver( self, - receiver_ref: mo.ActorRefType["ReceiverManagerActor"], + receiver_ref: mo.ActorRefType["MockReceiverManagerActor"], local_buffers: List, remote_buffers: List, session_id: str, @@ -204,35 +210,73 @@ async def _copy_to_receiver( ) -# test for cancelling happens when creating writer -class MockReceiverManagerActor2(ReceiverManagerActor): - async def create_writers( - self, session_id, data_keys, data_sizes, level, sub_infos, band_name - ): - await asyncio.sleep(3) - return await super().create_writers( - session_id, data_keys, data_sizes, level, sub_infos, band_name +class MockStorageHandlerActor(StorageHandlerActor): + async def get_receive_manager_ref(self, band_name: str): + return await mo.actor_ref( + address=self.address, + uid=MockReceiverManagerActor.gen_uid(band_name), ) - -class MockSenderManagerActor2(SenderManagerActor): @staticmethod - async def get_receiver_ref(address: str, band_name: str): + async def get_send_manager_ref(address: str, band: str): return await mo.actor_ref( - address=address, uid=MockReceiverManagerActor2.default_uid() + address=address, uid=MockSenderManagerActor.gen_uid(band) ) +class MockStorageManagerActor(StorageManagerActor): + def __init__( + self, storage_configs: Dict, transfer_block_size: int = None, **kwargs + ): + super().__init__(storage_configs, transfer_block_size, **kwargs) + self._handler_cls = MockStorageHandlerActor + + async def _create_transfer_actors(self): + default_band_name = "numa-0" + sender_strategy = IdleLabel("io", "sender") + receiver_strategy = IdleLabel("io", "receiver") + handler_strategy = IdleLabel("io", "handler") + while True: + try: + handler_ref = await mo.create_actor( + MockStorageHandlerActor, + self._init_params[default_band_name], + self._data_manager, + self._spill_managers[default_band_name], + self._quotas[default_band_name], + default_band_name, + uid=MockStorageHandlerActor.gen_uid(default_band_name), + address=self.address, + allocate_strategy=handler_strategy, + ) + await mo.create_actor( + MockSenderManagerActor, + data_manager_ref=self._data_manager, + storage_handler_ref=handler_ref, + uid=MockSenderManagerActor.gen_uid(default_band_name), + address=self.address, + allocate_strategy=sender_strategy, + ) + + await mo.create_actor( + MockReceiverManagerActor, + self._quotas[default_band_name], + handler_ref, + address=self.address, + uid=MockReceiverManagerActor.gen_uid(default_band_name), + allocate_strategy=receiver_strategy, + ) + except NoIdleSlot: + break + + @pytest.mark.parametrize( "mock_sender_cls, mock_receiver_cls", - [ - (MockSenderManagerActor, MockReceiverManagerActor), - (MockSenderManagerActor2, MockReceiverManagerActor2), - ], + [(MockSenderManagerActor, MockReceiverManagerActor)], ) @pytest.mark.asyncio -async def test_cancel_transfer(create_actors, mock_sender_cls, mock_receiver_cls): - worker_address_1, worker_address_2 = create_actors +async def test_cancel_transfer(create_actors_mock, mock_sender_cls, mock_receiver_cls): + _, _, io1, io2 = create_actors_mock session_id = "mock_session" quota_refs = { @@ -240,33 +284,15 @@ async def test_cancel_transfer(create_actors, mock_sender_cls, mock_receiver_cls StorageQuotaActor, StorageLevel.MEMORY, 5 * 1024 * 1024, - address=worker_address_2, + address=io2, uid=StorageQuotaActor.gen_uid("numa-0", StorageLevel.MEMORY), ) } - data_manager_ref = await mo.actor_ref( - uid=DataManagerActor.default_uid(), address=worker_address_1 - ) storage_handler1 = await mo.actor_ref( - uid=StorageHandlerActor.gen_uid("numa-0"), address=worker_address_1 + uid=MockStorageHandlerActor.gen_uid("numa-0"), address=io1 ) storage_handler2 = await mo.actor_ref( - uid=StorageHandlerActor.gen_uid("numa-0"), address=worker_address_2 - ) - - sender_actor: mo.ActorRefType[SenderManagerActor] = await mo.create_actor( - mock_sender_cls, - data_manager_ref=data_manager_ref, - uid=mock_sender_cls.default_uid(), - address=worker_address_1, - allocate_strategy=IdleLabel("io", "mock_sender_cls"), - ) - receiver_actor: mo.ActorRefType[ReceiverManagerActor] = await mo.create_actor( - mock_receiver_cls, - quota_refs, - uid=mock_receiver_cls.default_uid(), - address=worker_address_2, - allocate_strategy=IdleLabel("io", "mock_receiver_cls"), + uid=MockStorageHandlerActor.gen_uid("numa-0"), address=io2 ) data1 = np.random.rand(10, 10) @@ -276,31 +302,18 @@ async def test_cancel_transfer(create_actors, mock_sender_cls, mock_receiver_cls used_before = (await quota_refs[StorageLevel.MEMORY].get_quota())[1] - async def transfer( - data_keys: List[Union[str, tuple]], - sender: mo.ActorRefType[SenderManagerActor], - receiver: mo.ActorRefType[ReceiverManagerActor], - ): - is_transferring_list = await receiver.open_writers( + transfer_task = asyncio.create_task( + storage_handler2.fetch_via_transfer( session_id, - data_keys, + ["data_key1"], StorageLevel.MEMORY, - (sender.address, "numa-0"), + (io1, "numa-0"), + "numa-0", "raise", ) - assert len(is_transferring_list) > 0 - await sender.send_batch_data( - session_id, - data_keys, - is_transferring_list, - (receiver.address, "numa-0"), - ) - - transfer_task = asyncio.create_task( - transfer(["data_key1"], sender_actor, receiver_actor) ) - await asyncio.sleep(0.5) + await asyncio.sleep(1) transfer_task.cancel() with pytest.raises(asyncio.CancelledError): @@ -313,7 +326,14 @@ async def transfer( await storage_handler2.get(session_id, "data_key1") transfer_task = asyncio.create_task( - transfer(["data_key1"], sender_actor, receiver_actor) + storage_handler2.fetch_via_transfer( + session_id, + ["data_key1"], + StorageLevel.MEMORY, + (io1, "numa-0"), + "numa-0", + "raise", + ) ) await transfer_task get_data = await storage_handler2.get(session_id, "data_key1") @@ -322,12 +342,26 @@ async def transfer( # cancel when fetch the same data Simultaneously if mock_sender_cls is MockSenderManagerActor: transfer_task1 = asyncio.create_task( - transfer(["data_key2"], sender_actor, receiver_actor) + storage_handler2.fetch_via_transfer( + session_id, + ["data_key1"], + StorageLevel.MEMORY, + (io1, "numa-0"), + "numa-0", + "raise", + ) ) transfer_task2 = asyncio.create_task( - transfer(["data_key2"], sender_actor, receiver_actor) + storage_handler2.fetch_via_transfer( + session_id, + ["data_key2"], + StorageLevel.MEMORY, + (io1, "numa-0"), + "numa-0", + "raise", + ) ) - await asyncio.sleep(0.5) + await asyncio.sleep(1) transfer_task1.cancel() with pytest.raises(asyncio.CancelledError): await transfer_task1 @@ -338,37 +372,37 @@ async def transfer( @pytest.mark.asyncio async def test_transfer_same_data(create_actors): - worker_address_1, worker_address_2 = create_actors + worker_address_1, worker_address_2, io1, io2 = create_actors session_id = "mock_session" data1 = np.random.rand(100, 100) storage_handler1 = await mo.actor_ref( - uid=StorageHandlerActor.gen_uid("numa-0"), address=worker_address_1 + uid=StorageHandlerActor.gen_uid("numa-0"), address=io1 ) - await mo.actor_ref( - uid=StorageHandlerActor.gen_uid("numa-0"), address=worker_address_2 + storage_handler2 = await mo.actor_ref( + uid=StorageHandlerActor.gen_uid("numa-0"), address=io2 ) await storage_handler1.put(session_id, "data_key1", data1, StorageLevel.MEMORY) - receiver_actor: mo.ActorRefType[ReceiverManagerActor] = await mo.actor_ref( - address=worker_address_2, uid=ReceiverManagerActor.gen_uid("numa-0") - ) - task1 = receiver_actor.open_writers( + task1 = storage_handler2.fetch_via_transfer( session_id, ["data_key1"], StorageLevel.MEMORY, - (worker_address_1, "numa-0"), + (io1, "numa-0"), + "numa-0", "raise", ) - task2 = receiver_actor.open_writers( + task2 = storage_handler2.fetch_via_transfer( session_id, ["data_key1"], StorageLevel.MEMORY, - (worker_address_1, "numa-0"), + (io1, "numa-0"), + "numa-0", "raise", ) - results = await asyncio.gather(task1, task2) - assert [False] in results - assert [True] in results + await asyncio.gather(task1, task2) + + get_data = await storage_handler2.get(session_id, "data_key1") + np.testing.assert_array_equal(get_data, data1) diff --git a/python/xorbits/_mars/services/storage/tests/test_transfer_gpu.py b/python/xorbits/_mars/services/storage/tests/test_transfer_gpu.py index c3ba3fe6f..0d0ca35cc 100644 --- a/python/xorbits/_mars/services/storage/tests/test_transfer_gpu.py +++ b/python/xorbits/_mars/services/storage/tests/test_transfer_gpu.py @@ -28,7 +28,6 @@ from ....utils import lazy_import from ..core import StorageManagerActor from ..handler import StorageHandlerActor -from ..transfer import SenderManagerActor cupy = lazy_import("cupy") cudf = lazy_import("cudf") @@ -109,7 +108,7 @@ def _generate_band_scheme(): gpu_bands = [] if gpu_counts: gpu_bands.extend([(f"gpu-{i}", f"gpu-{i+1}") for i in range(gpu_counts - 1)]) - bands = [("numa-0", "numa-0"), ("gpu-0", "gpu-0")] + bands = [("gpu-0", "gpu-0")] bands.extend(gpu_bands) schemes = [(None, None), ("ucx", "ucx"), (None, "ucx"), ("ucx", None)] for band in bands: @@ -137,7 +136,7 @@ async def test_simple_transfer(create_actors): storage_handler1 = await mo.actor_ref( uid=StorageHandlerActor.gen_uid(bands[0]), address=worker_address_1 ) - storage_handler2 = await mo.actor_ref( + storage_handler2: mo.ActorRefType[StorageHandlerActor] = await mo.actor_ref( uid=StorageHandlerActor.gen_uid(bands[1]), address=worker_address_2 ) @@ -146,27 +145,13 @@ async def test_simple_transfer(create_actors): await storage_handler1.put(session_id, "data_key2", data2, storage_level) await storage_handler2.put(session_id, "data_key3", data2, storage_level) - sender_actor = await mo.actor_ref( - address=worker_address_1, uid=SenderManagerActor.gen_uid(bands[0]) - ) - - # send data to worker2 from worker1 - await sender_actor.send_batch_data( + await storage_handler2.fetch_via_transfer( session_id, - ["data_key1"], - worker_address_2, + ["data_key1", "data_key2"], storage_level, - block_size=1000, - band_name=bands[1], - ) - - await sender_actor.send_batch_data( - session_id, - ["data_key2"], - worker_address_2, - storage_level, - block_size=1000, - band_name=bands[1], + (worker_address_1, bands[0]), + bands[1], + "raise", ) get_data1 = await storage_handler2.get(session_id, "data_key1") @@ -175,18 +160,15 @@ async def test_simple_transfer(create_actors): get_data2 = await storage_handler2.get(session_id, "data_key2") xpd.testing.assert_frame_equal(data2, get_data2) - # send data to worker1 from worker2 - sender_actor = await mo.actor_ref( - address=worker_address_2, uid=SenderManagerActor.gen_uid(bands[1]) - ) - await sender_actor.send_batch_data( + await storage_handler1.fetch_via_transfer( session_id, ["data_key3"], - worker_address_1, storage_level, - block_size=1000, - band_name=bands[0], + (worker_address_2, bands[1]), + bands[0], + "raise", ) + get_data3 = await storage_handler1.get(session_id, "data_key3") xpd.testing.assert_frame_equal(data2, get_data3) @@ -212,29 +194,35 @@ async def test_transfer_same_data(create_actors): ) await storage_handler1.put(session_id, "data_key1", data1, storage_level) - sender_actor = await mo.actor_ref( - address=worker_address_1, uid=SenderManagerActor.gen_uid(bands[0]) + + storage_handler2.fetch_via_transfer( + session_id, + ["data_key1"], + storage_level, + (worker_address_1, bands[0]), + bands[1], + "raise", ) # send data to worker2 from worker1 task1 = asyncio.create_task( - sender_actor.send_batch_data( + storage_handler2.fetch_via_transfer( session_id, ["data_key1"], - worker_address_2, storage_level, - block_size=1000, - band_name=bands[1], + (worker_address_1, bands[0]), + bands[1], + "raise", ) ) task2 = asyncio.create_task( - sender_actor.send_batch_data( + storage_handler2.fetch_via_transfer( session_id, ["data_key1"], - worker_address_2, storage_level, - block_size=1000, - band_name=bands[1], + (worker_address_1, bands[0]), + bands[1], + "raise", ) ) await asyncio.gather(task1, task2) diff --git a/python/xorbits/_mars/services/storage/transfer.py b/python/xorbits/_mars/services/storage/transfer.py index 4273e9a8b..17cb8b027 100644 --- a/python/xorbits/_mars/services/storage/transfer.py +++ b/python/xorbits/_mars/services/storage/transfer.py @@ -17,7 +17,7 @@ import logging from dataclasses import dataclass from io import UnsupportedOperation -from typing import Dict, List, Optional, Union, Any +from typing import Any, Dict, List, Optional, Union import xoscar as mo from xoscar.core import BufferRef @@ -43,6 +43,7 @@ def __init__( data_manager_ref: mo.ActorRefType[DataManagerActor] = None, storage_handler_ref: mo.ActorRefType[StorageHandlerActor] = None, ): + super().__init__() self._band_name = band_name self._data_manager_ref = data_manager_ref self._storage_handler = storage_handler_ref @@ -75,7 +76,6 @@ async def _copy_to_receiver( block_size: int, ): await mo.copy_to(local_buffers, remote_buffers, block_size=block_size) - await receiver_ref.handle_transmission_done(session_id, data_keys) @staticmethod def _get_buffer_size(buf: Any): @@ -145,7 +145,6 @@ async def _send_data( try: await asyncio.shield(write_task) except asyncio.CancelledError: - await receiver_ref.handle_transmission_cancellation(session_id, data_keys) raise @mo.extensible @@ -153,20 +152,14 @@ async def send_batch_data( self, session_id: str, data_keys: List[str], - is_transferring_list: List[bool], + to_send_keys: List, + to_wait_keys: List, remote_band: BandType, block_size: int = None, ): logger.debug( "Begin to send data (%s, %s) to %s", session_id, data_keys, remote_band ) - 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) receiver_ref: mo.ActorRefType[ ReceiverManagerActor @@ -176,19 +169,8 @@ async def send_batch_data( logger.debug("Start sending %s to %s", to_send_keys, receiver_ref.address) block_size = block_size or self._transfer_block_size await self._send_data(receiver_ref, session_id, to_send_keys, block_size) - logger.debug("Done sending %s", to_send_keys) if to_wait_keys: - logger.debug("Start waiting %s", to_wait_keys) await receiver_ref.wait_transfer_done(session_id, to_wait_keys) - logger.debug("Done waiting %s", to_wait_keys) - unpin_tasks = [] - for data_key in data_keys: - unpin_tasks.append( - self._data_manager_ref.unpin.delay( - session_id, [data_key], self._band_name, error="ignore" - ) - ) - await self._data_manager_ref.unpin.batch(*unpin_tasks) logger.debug( "Finish sending data (%s, %s) to %s", session_id, @@ -214,17 +196,10 @@ def __init__( storage_handler_ref: mo.ActorRefType[StorageHandlerActor] = None, ): self._quota_refs = quota_refs - # self._storage_handler = storage_handler_ref + self._storage_handler = storage_handler_ref self._writing_infos: Dict[tuple, WritingInfo] = dict() self._lock = asyncio.Lock() - async def __post_create__(self): - # if self._storage_handler is None: # for test - # self._storage_handler = await mo.actor_ref( - # self.address, StorageHandlerActor.gen_uid("numa-0") - # ) - pass - async def get_buffers( self, session_id: str, @@ -266,39 +241,28 @@ async def handle_transmission_done(self, session_id: str, data_keys: List): self._decref_writing_key(session_id, data_key) async def handle_transmission_cancellation(self, session_id: str, data_keys: List): + data_keys_to_be_deleted = [] async with self._lock: for data_key in data_keys: if (session_id, data_key) in self._writing_infos: if self._writing_infos[(session_id, data_key)].ref_counts == 1: info = self._writing_infos[(session_id, data_key)] await self._quota_refs[info.level].release_quota(info.size) - await self._storage_handler.delete( - session_id, data_key, error="ignore" - ) + data_keys_to_be_deleted.append(data_key) await info.writer.clean_up() info.event.set() self._decref_writing_key(session_id, data_key) + return data_keys_to_be_deleted @classmethod def gen_uid(cls, band_name: str): return f"receiver_manager_{band_name}" - def lock(self): - self._lock.acquire() - - def unlock(self): - self._lock.release() - def _decref_writing_key(self, session_id: str, data_key: str): self._writing_infos[(session_id, data_key)].ref_counts -= 1 if self._writing_infos[(session_id, data_key)].ref_counts == 0: del self._writing_infos[(session_id, data_key)] - def is_writing_info_existed( - self, session_id: str, data_keys: List[Union[str, tuple]] - ) -> List[bool]: - return [(session_id, data_key) in self._writing_infos for data_key in data_keys] - async def add_writers( self, session_id: str, @@ -325,73 +289,6 @@ async def add_writers( is_transferring.append(True) return is_transferring - async def create_writers( - self, - session_id: str, - data_keys: List[str], - data_sizes: List[int], - level: StorageLevel, - sub_infos: List, - band_name: str, - ) -> List[bool]: - tasks = dict() - key_to_sub_infos = dict() - data_key_to_size = dict() - being_processed = [] - for data_key, data_size, sub_info in zip(data_keys, data_sizes, sub_infos): - data_key_to_size[data_key] = data_size - if (session_id, data_key) not in self._writing_infos: - being_processed.append(False) - tasks[data_key] = self._storage_handler.open_writer.delay( - session_id, - data_key, - data_size, - level, - request_quota=False, - band_name=band_name, - ) - key_to_sub_infos[data_key] = sub_info - else: - being_processed.append(True) - self._writing_infos[(session_id, data_key)].ref_counts += 1 - if tasks: - logger.debug("Executing open_writer.batch") - writers = await self._storage_handler.open_writer.batch( - *tuple(tasks.values()) - ) - for data_key, writer in zip(tasks, writers): - self._writing_infos[(session_id, data_key)] = WritingInfo( - writer, data_key_to_size[data_key], level, asyncio.Event(), 1 - ) - if key_to_sub_infos[data_key] is not None: - writer._sub_key_infos = key_to_sub_infos[data_key] - return being_processed - - async def open_writers( - self, - session_id: str, - data_keys: List[str], - data_sizes: List[int], - sub_infos, - level: StorageLevel, - band_name - ) -> List[bool]: - logger.debug("Start opening writers for %s", data_keys) - - async with self._lock: - logger.debug("Opening writers for %s", data_keys) - future = asyncio.create_task( - self.create_writers( - session_id, data_keys, data_sizes, level, sub_infos, band_name - ) - ) - try: - return await future - except asyncio.CancelledError: - await self._quota_refs[level].release_quota(sum(data_sizes)) - future.cancel() - raise - async def wait_transfer_done(self, session_id, data_keys): await asyncio.gather( *[self._writing_infos[(session_id, key)].event.wait() for key in data_keys] diff --git a/python/xorbits/_mars/storage/cuda.py b/python/xorbits/_mars/storage/cuda.py index 7e4dee4a8..97a33c7f3 100644 --- a/python/xorbits/_mars/storage/cuda.py +++ b/python/xorbits/_mars/storage/cuda.py @@ -67,7 +67,10 @@ def set_buffers_by_sizes(self, sizes: List[int]): from rmm import DeviceBuffer self._buffers = [ - _convert_to_cupy_ndarray(DeviceBuffer(size=size)) for size in sizes + cupy.ndarray(shape=0, dtype="u1") + if size == 0 + else _convert_to_cupy_ndarray(DeviceBuffer(size=size)) + for size in sizes ] @property @@ -96,7 +99,10 @@ def _initialize_read(self): self._buffers.append(buf.astype("u1", copy=False)) buffer_types.append(["cuda", buf.size]) elif isinstance(buf, Buffer): - self._buffers.append(_convert_to_cupy_ndarray(buf)) + if buf.size == 0: + self._buffers.append(cupy.ndarray(shape=0, dtype="u1")) + else: + self._buffers.append(_convert_to_cupy_ndarray(buf)) buffer_types.append(["cuda", buf.size]) else: size = getattr(buf, "size", len(buf)) From 5dafe70756fec3ce59ab3f8c26c5133a6f466a38 Mon Sep 17 00:00:00 2001 From: UranusSeven <109661872+UranusSeven@users.noreply.github.com> Date: Thu, 1 Jun 2023 16:32:55 +0800 Subject: [PATCH 11/25] REF: refactor transfer to avoid deadlocks --- .../xorbits/_mars/services/storage/handler.py | 25 +- .../services/storage/tests/test_transfer.py | 218 ++++++++++-------- .../_mars/services/storage/transfer.py | 73 ++---- 3 files changed, 156 insertions(+), 160 deletions(-) diff --git a/python/xorbits/_mars/services/storage/handler.py b/python/xorbits/_mars/services/storage/handler.py index 0cdaed519..ddcf591d0 100644 --- a/python/xorbits/_mars/services/storage/handler.py +++ b/python/xorbits/_mars/services/storage/handler.py @@ -531,19 +531,30 @@ async def _fetch_via_transfer( fetch_band_name: str, error: str, ): - from .transfer import SenderManagerActor + from .transfer import ReceiverManagerActor, SenderManagerActor logger.debug("Begin to fetch %s from band %s", data_keys, remote_band) + + receiver_ref: mo.ActorRefType[ReceiverManagerActor] = await mo.actor_ref( + address=self.address, + uid=ReceiverManagerActor.gen_uid(fetch_band_name), + ) + is_transferring_list = await receiver_ref.open_writers( + session_id, data_keys, level, remote_band, error + ) + if not is_transferring_list: + # no data to be transferred + return + sender_ref: mo.ActorRefType[SenderManagerActor] = await mo.actor_ref( address=remote_band[0], uid=SenderManagerActor.gen_uid(remote_band[1]) ) await sender_ref.send_batch_data( session_id, data_keys, + is_transferring_list, self._data_manager_ref.address, - level, fetch_band_name, - error=error, ) logger.debug("Finish fetching %s from band %s", data_keys, remote_band) @@ -559,10 +570,8 @@ async def fetch_batch( if error not in ("raise", "ignore"): # pragma: no cover raise ValueError("error must be raise or ignore") - meta_api = await self._get_meta_api(session_id) remote_keys = defaultdict(set) missing_keys = [] - get_metas = [] get_info_delays = [] for data_key in data_keys: get_info_delays.append( @@ -586,6 +595,9 @@ async def fetch_batch( else: # Not exists in local, fetch from remote worker missing_keys.append(data_key) + await self._data_manager_ref.pin.batch(*pin_delays) + + meta_api = await self._get_meta_api(session_id) if address is None or band_name is None: # some mapper keys are absent, specify error='ignore' # remember that meta only records those main keys @@ -599,9 +611,6 @@ async def fetch_batch( ) for data_key in missing_keys ] - await self._data_manager_ref.pin.batch(*pin_delays) - - if get_metas: metas = await meta_api.get_chunk_meta.batch(*get_metas) else: # pragma: no cover metas = [{"bands": [(address, band_name)]}] * len(missing_keys) diff --git a/python/xorbits/_mars/services/storage/tests/test_transfer.py b/python/xorbits/_mars/services/storage/tests/test_transfer.py index 1b6ac4e60..96b00b6e5 100644 --- a/python/xorbits/_mars/services/storage/tests/test_transfer.py +++ b/python/xorbits/_mars/services/storage/tests/test_transfer.py @@ -17,7 +17,7 @@ import os import sys import tempfile -from typing import List +from typing import List, Union import numpy as np import pandas as pd @@ -111,42 +111,65 @@ async def test_simple_transfer(create_actors): await storage_handler1.put(session_id, "data_key2", data2, level) await storage_handler2.put(session_id, "data_key3", data2, level) - sender_actor = await mo.actor_ref( - address=worker_address_1, uid=SenderManagerActor.gen_uid("numa-0") - ) + receiver_actor: mo.ActorRefType[ReceiverManagerActor] = await mo.actor_ref( + address=worker_address_2, uid=ReceiverManagerActor.gen_uid("numa-0") + ) + sender_actor: mo.ActorRefType[SenderManagerActor] = 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, - level, - block_size=1000, - ) + is_transferring = await receiver_actor.open_writers( + session_id, + ["data_key1", "data_key2"], + StorageLevel.MEMORY, + (sender_actor.address, "numa-0"), + "raise", + ) - await sender_actor.send_batch_data( - session_id, - ["data_key2"], - worker_address_2, - level, - block_size=1000, - ) + # send data to worker2 from worker1 + await sender_actor.send_batch_data( + session_id, + ["data_key1"], + is_transferring, + (receiver_actor.address, "numa-0"), + 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"], + is_transferring, + (receiver_actor.address, "numa-0"), + block_size=1000, + ) - get_data2 = await storage_handler2.get(session_id, "data_key2") - pd.testing.assert_frame_equal(data2, get_data2) + get_data1 = await storage_handler2.get(session_id, "data_key1") + np.testing.assert_array_equal(data1, get_data1) - # send data to worker1 from worker2 - sender_actor = 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, level - ) - get_data3 = await storage_handler1.get(session_id, "data_key3") - pd.testing.assert_frame_equal(data2, get_data3) + get_data2 = await storage_handler2.get(session_id, "data_key2") + pd.testing.assert_frame_equal(data2, get_data2) + + # send data to worker1 from worker2 + receiver_actor: mo.ActorRefType[ReceiverManagerActor] = await mo.actor_ref( + address=worker_address_1, uid=ReceiverManagerActor.gen_uid("numa-0") + ) + sender_actor: mo.ActorRefType[SenderManagerActor] = await mo.actor_ref( + address=worker_address_2, uid=SenderManagerActor.gen_uid("numa-0") + ) + + is_transferring = await receiver_actor.open_writers( + session_id, + ["data_key3"], + StorageLevel.MEMORY, + (sender_actor.address, "numa-0"), + "raise", + ) + + await sender_actor.send_batch_data( + session_id, ["data_key3"], is_transferring, (receiver_actor.address, "numa-0") + ) + get_data3 = await storage_handler1.get(session_id, "data_key3") + pd.testing.assert_frame_equal(data2, get_data3) # test for cancelling happens when writing @@ -201,16 +224,17 @@ async def get_receiver_ref(address: str, band_name: str): @pytest.mark.parametrize( - "mock_sender, mock_receiver", + "mock_sender_cls, mock_receiver_cls", [ (MockSenderManagerActor, MockReceiverManagerActor), (MockSenderManagerActor2, MockReceiverManagerActor2), ], ) @pytest.mark.asyncio -async def test_cancel_transfer(create_actors, mock_sender, mock_receiver): +async def test_cancel_transfer(create_actors, mock_sender_cls, mock_receiver_cls): worker_address_1, worker_address_2 = create_actors + session_id = "mock_session" quota_refs = { StorageLevel.MEMORY: await mo.actor_ref( StorageQuotaActor, @@ -230,73 +254,85 @@ async def test_cancel_transfer(create_actors, mock_sender, mock_receiver): uid=StorageHandlerActor.gen_uid("numa-0"), address=worker_address_2 ) - sender_actor = await mo.create_actor( - mock_sender, + sender_actor: mo.ActorRefType[SenderManagerActor] = await mo.create_actor( + mock_sender_cls, data_manager_ref=data_manager_ref, - uid=mock_sender.default_uid(), + uid=mock_sender_cls.default_uid(), address=worker_address_1, - allocate_strategy=IdleLabel("io", "mock_sender"), + allocate_strategy=IdleLabel("io", "mock_sender_cls"), ) - await mo.create_actor( - mock_receiver, + receiver_actor: mo.ActorRefType[ReceiverManagerActor] = await mo.create_actor( + mock_receiver_cls, quota_refs, - uid=mock_receiver.default_uid(), + uid=mock_receiver_cls.default_uid(), address=worker_address_2, - allocate_strategy=IdleLabel("io", "mock_receiver"), + allocate_strategy=IdleLabel("io", "mock_receiver_cls"), ) data1 = np.random.rand(10, 10) - await storage_handler1.put("mock", "data_key1", data1, StorageLevel.MEMORY) + await storage_handler1.put(session_id, "data_key1", data1, StorageLevel.MEMORY) data2 = pd.DataFrame(np.random.rand(100, 100)) - await storage_handler1.put("mock", "data_key2", data2, StorageLevel.MEMORY) + await storage_handler1.put(session_id, "data_key2", data2, StorageLevel.MEMORY) used_before = (await quota_refs[StorageLevel.MEMORY].get_quota())[1] - send_task = asyncio.create_task( - sender_actor.send_batch_data( - "mock", ["data_key1"], worker_address_2, StorageLevel.MEMORY + async def transfer( + data_keys: List[Union[str, tuple]], + sender: mo.ActorRefType[SenderManagerActor], + receiver: mo.ActorRefType[ReceiverManagerActor], + ): + is_transferring_list = await receiver.open_writers( + session_id, + data_keys, + StorageLevel.MEMORY, + (sender.address, "numa-0"), + "raise", ) + assert len(is_transferring_list) > 0 + await sender.send_batch_data( + session_id, + data_keys, + is_transferring_list, + (receiver.address, "numa-0"), + ) + + transfer_task = asyncio.create_task( + transfer(["data_key1"], sender_actor, receiver_actor) ) await asyncio.sleep(0.5) - send_task.cancel() + transfer_task.cancel() with pytest.raises(asyncio.CancelledError): - await send_task + await transfer_task used = (await quota_refs[StorageLevel.MEMORY].get_quota())[1] assert used == used_before with pytest.raises(DataNotExist): - await storage_handler2.get("mock", "data_key1") + await storage_handler2.get(session_id, "data_key1") - send_task = asyncio.create_task( - sender_actor.send_batch_data( - "mock", ["data_key1"], worker_address_2, StorageLevel.MEMORY - ) + transfer_task = asyncio.create_task( + transfer(["data_key1"], sender_actor, receiver_actor) ) - await send_task - get_data = await storage_handler2.get("mock", "data_key1") + await transfer_task + get_data = await storage_handler2.get(session_id, "data_key1") np.testing.assert_array_equal(data1, get_data) # cancel when fetch the same data Simultaneously - if mock_sender is MockSenderManagerActor: - send_task1 = asyncio.create_task( - sender_actor.send_batch_data( - "mock", ["data_key2"], worker_address_2, StorageLevel.MEMORY - ) + if mock_sender_cls is MockSenderManagerActor: + transfer_task1 = asyncio.create_task( + transfer(["data_key2"], sender_actor, receiver_actor) ) - send_task2 = asyncio.create_task( - sender_actor.send_batch_data( - "mock", ["data_key2"], worker_address_2, StorageLevel.MEMORY - ) + transfer_task2 = asyncio.create_task( + transfer(["data_key2"], sender_actor, receiver_actor) ) await asyncio.sleep(0.5) - send_task1.cancel() + transfer_task1.cancel() with pytest.raises(asyncio.CancelledError): - await send_task1 - await send_task2 - get_data2 = await storage_handler2.get("mock", "data_key2") + await transfer_task1 + await transfer_task2 + get_data2 = await storage_handler2.get(session_id, "data_key2") pd.testing.assert_frame_equal(get_data2, data2) @@ -309,34 +345,30 @@ async def test_transfer_same_data(create_actors): storage_handler1 = await mo.actor_ref( uid=StorageHandlerActor.gen_uid("numa-0"), address=worker_address_1 ) - storage_handler2 = await mo.actor_ref( + await mo.actor_ref( uid=StorageHandlerActor.gen_uid("numa-0"), address=worker_address_2 ) await storage_handler1.put(session_id, "data_key1", data1, StorageLevel.MEMORY) - sender_actor = await mo.actor_ref( - address=worker_address_1, uid=SenderManagerActor.gen_uid("numa-0") - ) - # send data to worker2 from worker1 - task1 = asyncio.create_task( - sender_actor.send_batch_data( - session_id, - ["data_key1"], - worker_address_2, - StorageLevel.MEMORY, - block_size=1000, - ) + receiver_actor: mo.ActorRefType[ReceiverManagerActor] = await mo.actor_ref( + address=worker_address_2, uid=ReceiverManagerActor.gen_uid("numa-0") ) - task2 = asyncio.create_task( - sender_actor.send_batch_data( - session_id, - ["data_key1"], - worker_address_2, - StorageLevel.MEMORY, - block_size=1000, - ) + task1 = receiver_actor.open_writers( + session_id, + ["data_key1"], + StorageLevel.MEMORY, + (worker_address_1, "numa-0"), + "raise", ) - await asyncio.gather(task1, task2) - get_data1 = await storage_handler2.get(session_id, "data_key1") - np.testing.assert_array_equal(data1, get_data1) + task2 = receiver_actor.open_writers( + session_id, + ["data_key1"], + StorageLevel.MEMORY, + (worker_address_1, "numa-0"), + "raise", + ) + + results = await asyncio.gather(task1, task2) + assert [False] in results + assert [True] in results diff --git a/python/xorbits/_mars/services/storage/transfer.py b/python/xorbits/_mars/services/storage/transfer.py index 7acda794c..2dff1ae19 100644 --- a/python/xorbits/_mars/services/storage/transfer.py +++ b/python/xorbits/_mars/services/storage/transfer.py @@ -24,6 +24,7 @@ from ...lib.aio import alru_cache from ...storage import StorageLevel +from ...typing import BandType from ...utils import dataslots from .core import DataManagerActor, WrappedStorageFileObject from .handler import StorageHandlerActor @@ -152,62 +153,12 @@ async def send_batch_data( self, session_id: str, data_keys: List[str], - address: str, - level: StorageLevel, - band_name: str = "numa-0", + is_transferring_list: List[bool], + remote_band: BandType, block_size: int = None, - error: str = "raise", ): logger.debug( - "Begin to send data (%s, %s) to %s", session_id, data_keys, address - ) - - tasks = [] - for key in data_keys: - tasks.append(self._data_manager_ref.get_store_key.delay(session_id, key)) - data_keys = await self._data_manager_ref.get_store_key.batch(*tasks) - data_keys = list(set(data_keys)) - sub_infos = await self._data_manager_ref.get_sub_infos.batch( - *[ - self._data_manager_ref.get_sub_infos.delay(session_id, key) - for key in data_keys - ] - ) - - 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: - get_infos.append( - self._data_manager_ref.get_data_info.delay( - session_id, data_key, self._band_name, error - ) - ) - pin_tasks.append( - self._data_manager_ref.pin.delay( - session_id, data_key, self._band_name, error - ) - ) - await self._data_manager_ref.pin.batch(*pin_tasks) - infos = await self._data_manager_ref.get_data_info.batch(*get_infos) - filtered = [ - (data_info, data_key) - for data_info, data_key in zip(infos, data_keys) - if data_info is not None - ] - if filtered: - infos, data_keys = zip(*filtered) - else: # pragma: no cover - # no data to be transferred - return - 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, sub_infos, band_name + "Begin to send data (%s, %s) to %s", session_id, data_keys, remote_band ) to_send_keys = [] to_wait_keys = [] @@ -217,7 +168,12 @@ async def send_batch_data( else: to_send_keys.append(data_key) + receiver_ref: mo.ActorRefType[ + ReceiverManagerActor + ] = await self.get_receiver_ref(remote_band[0], remote_band[1]) + if to_send_keys: + block_size = block_size or self._transfer_block_size await self._send_data(receiver_ref, session_id, to_send_keys, block_size) if to_wait_keys: await receiver_ref.wait_transfer_done(session_id, to_wait_keys) @@ -230,11 +186,10 @@ async def send_batch_data( ) await self._data_manager_ref.unpin.batch(*unpin_tasks) logger.debug( - "Finish sending data (%s, %s) to %s, total size is %s", + "Finish sending data (%s, %s) to %s", session_id, data_keys, - address, - sum(data_sizes), + remote_band, ) @@ -336,7 +291,7 @@ async def create_writers( level: StorageLevel, sub_infos: List, band_name: str, - ): + ) -> List[bool]: tasks = dict() key_to_sub_infos = dict() data_key_to_size = dict() @@ -374,10 +329,10 @@ async def open_writers( session_id: str, data_keys: List[str], data_sizes: List[int], - level: StorageLevel, sub_infos: List, + level: StorageLevel, band_name: str, - ): + ) -> List[bool]: async with self._lock: await self._storage_handler.request_quota_with_spill(level, sum(data_sizes)) future = asyncio.create_task( From 89ad9350f9e94b27fb508acf8a0ffa0219412680 Mon Sep 17 00:00:00 2001 From: UranusSeven <109661872+UranusSeven@users.noreply.github.com> Date: Fri, 2 Jun 2023 14:20:57 +0800 Subject: [PATCH 12/25] Debugging --- python/xorbits/_mars/services/storage/handler.py | 5 +++-- python/xorbits/_mars/services/storage/transfer.py | 4 ++++ 2 files changed, 7 insertions(+), 2 deletions(-) diff --git a/python/xorbits/_mars/services/storage/handler.py b/python/xorbits/_mars/services/storage/handler.py index ddcf591d0..125d0fb26 100644 --- a/python/xorbits/_mars/services/storage/handler.py +++ b/python/xorbits/_mars/services/storage/handler.py @@ -546,6 +546,8 @@ async def _fetch_via_transfer( # no data to be transferred return + logger.debug("Writers have been opened for %s", data_keys) + sender_ref: mo.ActorRefType[SenderManagerActor] = await mo.actor_ref( address=remote_band[0], uid=SenderManagerActor.gen_uid(remote_band[1]) ) @@ -553,8 +555,7 @@ async def _fetch_via_transfer( session_id, data_keys, is_transferring_list, - self._data_manager_ref.address, - fetch_band_name, + (self._data_manager_ref.address, fetch_band_name), ) logger.debug("Finish fetching %s from band %s", data_keys, remote_band) diff --git a/python/xorbits/_mars/services/storage/transfer.py b/python/xorbits/_mars/services/storage/transfer.py index 2dff1ae19..35c0bcf86 100644 --- a/python/xorbits/_mars/services/storage/transfer.py +++ b/python/xorbits/_mars/services/storage/transfer.py @@ -173,10 +173,14 @@ async def send_batch_data( ] = await self.get_receiver_ref(remote_band[0], remote_band[1]) if to_send_keys: + logger.debug("Start sending %s", to_send_keys) block_size = block_size or self._transfer_block_size await self._send_data(receiver_ref, session_id, to_send_keys, block_size) + logger.debug("Done sending %s", to_send_keys) if to_wait_keys: + logger.debug("Start waiting %s", to_wait_keys) await receiver_ref.wait_transfer_done(session_id, to_wait_keys) + logger.debug("Done waiting %s", to_wait_keys) unpin_tasks = [] for data_key in data_keys: unpin_tasks.append( From dd63a1cfc6ce259e4baf42116bcc7d99f783986b Mon Sep 17 00:00:00 2001 From: UranusSeven <109661872+UranusSeven@users.noreply.github.com> Date: Fri, 2 Jun 2023 14:44:48 +0800 Subject: [PATCH 13/25] Debugging --- .../_mars/services/storage/transfer.py | 58 +++++++++++++++++++ 1 file changed, 58 insertions(+) diff --git a/python/xorbits/_mars/services/storage/transfer.py b/python/xorbits/_mars/services/storage/transfer.py index 35c0bcf86..8d6661c45 100644 --- a/python/xorbits/_mars/services/storage/transfer.py +++ b/python/xorbits/_mars/services/storage/transfer.py @@ -336,7 +336,65 @@ async def open_writers( sub_infos: List, level: StorageLevel, band_name: str, + remote_band, + error ) -> List[bool]: + logger.debug("Start opening writers for %s", data_keys) + + remote_data_manager_ref: mo.ActorRefType[DataManagerActor] = await mo.actor_ref( + address=remote_band[0], uid=DataManagerActor.default_uid() + ) + + logger.debug("Getting actual keys for %s", data_keys) + tasks = [] + for key in data_keys: + tasks.append(remote_data_manager_ref.get_store_key.delay(session_id, key)) + data_keys = await remote_data_manager_ref.get_store_key.batch(*tasks) + data_keys = list(set(data_keys)) + + logger.debug("Getting sub infos for %s", data_keys) + sub_infos = await remote_data_manager_ref.get_sub_infos.batch( + *[ + remote_data_manager_ref.get_sub_infos.delay(session_id, key) + for key in data_keys + ] + ) + + get_infos = [] + pin_tasks = [] + for data_key in data_keys: + get_infos.append( + remote_data_manager_ref.get_data_info.delay( + session_id, data_key, remote_band[1], error + ) + ) + pin_tasks.append( + remote_data_manager_ref.pin.delay( + session_id, data_key, remote_band[1], error + ) + ) + + logger.debug("Pining %s", data_keys) + await remote_data_manager_ref.pin.batch(*pin_tasks) + logger.debug("Getting data infos for %s", data_keys) + infos = await remote_data_manager_ref.get_data_info.batch(*get_infos) + + filtered = [ + (data_info, data_key) + for data_info, data_key in zip(infos, data_keys) + if data_info is not None + ] + if filtered: + infos, data_keys = zip(*filtered) + else: # pragma: no cover + # no data to be transferred + return [] + data_sizes = [info.store_size for info in infos] + + if level is None: + level = infos[0].level + + logger.debug("Opening writers for %s", data_keys) async with self._lock: await self._storage_handler.request_quota_with_spill(level, sum(data_sizes)) future = asyncio.create_task( From df23a159b064a1af19c0fa9397fc9091ac5421fe Mon Sep 17 00:00:00 2001 From: UranusSeven <109661872+UranusSeven@users.noreply.github.com> Date: Fri, 2 Jun 2023 14:49:32 +0800 Subject: [PATCH 14/25] Debugging --- python/xorbits/_mars/services/storage/transfer.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/python/xorbits/_mars/services/storage/transfer.py b/python/xorbits/_mars/services/storage/transfer.py index 8d6661c45..5503e467c 100644 --- a/python/xorbits/_mars/services/storage/transfer.py +++ b/python/xorbits/_mars/services/storage/transfer.py @@ -394,9 +394,10 @@ async def open_writers( if level is None: level = infos[0].level - logger.debug("Opening writers for %s", data_keys) async with self._lock: + logger.debug("Requesting quota for %s", data_keys) await self._storage_handler.request_quota_with_spill(level, sum(data_sizes)) + logger.debug("Opening writers for %s", data_keys) future = asyncio.create_task( self.create_writers( session_id, data_keys, data_sizes, level, sub_infos, band_name From 072848fd3cafecdbbd8686c370975372ec9c5300 Mon Sep 17 00:00:00 2001 From: UranusSeven <109661872+UranusSeven@users.noreply.github.com> Date: Fri, 2 Jun 2023 15:06:58 +0800 Subject: [PATCH 15/25] Fix --- .../xorbits/_mars/services/storage/handler.py | 62 ++++++++++++++++++- .../_mars/services/storage/transfer.py | 61 +----------------- 2 files changed, 63 insertions(+), 60 deletions(-) diff --git a/python/xorbits/_mars/services/storage/handler.py b/python/xorbits/_mars/services/storage/handler.py index 125d0fb26..ab42e7b67 100644 --- a/python/xorbits/_mars/services/storage/handler.py +++ b/python/xorbits/_mars/services/storage/handler.py @@ -535,12 +535,72 @@ async def _fetch_via_transfer( logger.debug("Begin to fetch %s from band %s", data_keys, remote_band) + remote_data_manager_ref: mo.ActorRefType[DataManagerActor] = await mo.actor_ref( + address=remote_band[0], uid=DataManagerActor.default_uid() + ) + + logger.debug("Getting actual keys for %s", data_keys) + tasks = [] + for key in data_keys: + tasks.append(remote_data_manager_ref.get_store_key.delay(session_id, key)) + data_keys = await remote_data_manager_ref.get_store_key.batch(*tasks) + data_keys = list(set(data_keys)) + + logger.debug("Getting sub infos for %s", data_keys) + sub_infos = await remote_data_manager_ref.get_sub_infos.batch( + *[ + remote_data_manager_ref.get_sub_infos.delay(session_id, key) + for key in data_keys + ] + ) + + get_infos = [] + pin_tasks = [] + for data_key in data_keys: + get_infos.append( + remote_data_manager_ref.get_data_info.delay( + session_id, data_key, remote_band[1], error + ) + ) + pin_tasks.append( + remote_data_manager_ref.pin.delay( + session_id, data_key, remote_band[1], error + ) + ) + + logger.debug("Pining %s", data_keys) + await remote_data_manager_ref.pin.batch(*pin_tasks) + logger.debug("Getting data infos for %s", data_keys) + infos = await remote_data_manager_ref.get_data_info.batch(*get_infos) + + filtered = [ + (data_info, data_key) + for data_info, data_key in zip(infos, data_keys) + if data_info is not None + ] + if filtered: + infos, data_keys = zip(*filtered) + else: # pragma: no cover + # no data to be transferred + return [] + data_sizes = [info.store_size for info in infos] + + if level is None: + level = infos[0].level + + logger.debug("Requesting quota for %s", data_keys) + await self.request_quota_with_spill(level, sum(data_sizes)) + receiver_ref: mo.ActorRefType[ReceiverManagerActor] = await mo.actor_ref( address=self.address, uid=ReceiverManagerActor.gen_uid(fetch_band_name), ) is_transferring_list = await receiver_ref.open_writers( - session_id, data_keys, level, remote_band, error + session_id, + data_keys, + data_sizes, + sub_infos, + level, ) if not is_transferring_list: # no data to be transferred diff --git a/python/xorbits/_mars/services/storage/transfer.py b/python/xorbits/_mars/services/storage/transfer.py index 5503e467c..b45a45086 100644 --- a/python/xorbits/_mars/services/storage/transfer.py +++ b/python/xorbits/_mars/services/storage/transfer.py @@ -333,70 +333,13 @@ async def open_writers( session_id: str, data_keys: List[str], data_sizes: List[int], - sub_infos: List, + sub_infos, level: StorageLevel, - band_name: str, - remote_band, - error + band_name ) -> List[bool]: logger.debug("Start opening writers for %s", data_keys) - remote_data_manager_ref: mo.ActorRefType[DataManagerActor] = await mo.actor_ref( - address=remote_band[0], uid=DataManagerActor.default_uid() - ) - - logger.debug("Getting actual keys for %s", data_keys) - tasks = [] - for key in data_keys: - tasks.append(remote_data_manager_ref.get_store_key.delay(session_id, key)) - data_keys = await remote_data_manager_ref.get_store_key.batch(*tasks) - data_keys = list(set(data_keys)) - - logger.debug("Getting sub infos for %s", data_keys) - sub_infos = await remote_data_manager_ref.get_sub_infos.batch( - *[ - remote_data_manager_ref.get_sub_infos.delay(session_id, key) - for key in data_keys - ] - ) - - get_infos = [] - pin_tasks = [] - for data_key in data_keys: - get_infos.append( - remote_data_manager_ref.get_data_info.delay( - session_id, data_key, remote_band[1], error - ) - ) - pin_tasks.append( - remote_data_manager_ref.pin.delay( - session_id, data_key, remote_band[1], error - ) - ) - - logger.debug("Pining %s", data_keys) - await remote_data_manager_ref.pin.batch(*pin_tasks) - logger.debug("Getting data infos for %s", data_keys) - infos = await remote_data_manager_ref.get_data_info.batch(*get_infos) - - filtered = [ - (data_info, data_key) - for data_info, data_key in zip(infos, data_keys) - if data_info is not None - ] - if filtered: - infos, data_keys = zip(*filtered) - else: # pragma: no cover - # no data to be transferred - return [] - data_sizes = [info.store_size for info in infos] - - if level is None: - level = infos[0].level - async with self._lock: - logger.debug("Requesting quota for %s", data_keys) - await self._storage_handler.request_quota_with_spill(level, sum(data_sizes)) logger.debug("Opening writers for %s", data_keys) future = asyncio.create_task( self.create_writers( From 373e7bc3ba6f3462f75ddd2bef7125a8a033360c Mon Sep 17 00:00:00 2001 From: UranusSeven <109661872+UranusSeven@users.noreply.github.com> Date: Fri, 2 Jun 2023 15:13:41 +0800 Subject: [PATCH 16/25] Debugging --- python/xorbits/_mars/services/storage/transfer.py | 1 + 1 file changed, 1 insertion(+) diff --git a/python/xorbits/_mars/services/storage/transfer.py b/python/xorbits/_mars/services/storage/transfer.py index b45a45086..6394e7071 100644 --- a/python/xorbits/_mars/services/storage/transfer.py +++ b/python/xorbits/_mars/services/storage/transfer.py @@ -317,6 +317,7 @@ async def create_writers( being_processed.append(True) self._writing_infos[(session_id, data_key)].ref_counts += 1 if tasks: + logger.debug("Executing open_writer.batch") writers = await self._storage_handler.open_writer.batch( *tuple(tasks.values()) ) From aa264bed6a047b607ab5fbab9387f2f3db3a13fd Mon Sep 17 00:00:00 2001 From: UranusSeven <109661872+UranusSeven@users.noreply.github.com> Date: Tue, 13 Jun 2023 12:29:08 +0800 Subject: [PATCH 17/25] Fix --- python/xorbits/_mars/services/storage/handler.py | 7 +++++++ python/xorbits/_mars/services/storage/transfer.py | 10 +++++++++- 2 files changed, 16 insertions(+), 1 deletion(-) diff --git a/python/xorbits/_mars/services/storage/handler.py b/python/xorbits/_mars/services/storage/handler.py index ab42e7b67..77929fc87 100644 --- a/python/xorbits/_mars/services/storage/handler.py +++ b/python/xorbits/_mars/services/storage/handler.py @@ -611,9 +611,16 @@ async def _fetch_via_transfer( sender_ref: mo.ActorRefType[SenderManagerActor] = await mo.actor_ref( address=remote_band[0], uid=SenderManagerActor.gen_uid(remote_band[1]) ) + + open_reader_tasks = [] + for data_key in data_keys: + open_reader_tasks.append(self.open_reader.delay(session_id, data_key)) + readers = await self.open_reader.batch(*open_reader_tasks) + await sender_ref.send_batch_data( session_id, data_keys, + readers, is_transferring_list, (self._data_manager_ref.address, fetch_band_name), ) diff --git a/python/xorbits/_mars/services/storage/transfer.py b/python/xorbits/_mars/services/storage/transfer.py index 6394e7071..081c4eafe 100644 --- a/python/xorbits/_mars/services/storage/transfer.py +++ b/python/xorbits/_mars/services/storage/transfer.py @@ -18,6 +18,7 @@ from dataclasses import dataclass from io import UnsupportedOperation from typing import Any, Dict, List, Union +from typing import TYPE_CHECKING, Dict, List import xoscar as mo from xoscar.core import BufferRef @@ -29,6 +30,10 @@ from .core import DataManagerActor, WrappedStorageFileObject from .handler import StorageHandlerActor +if TYPE_CHECKING: + from ...storage.core import StorageFileObject + + DEFAULT_TRANSFER_BLOCK_SIZE = 4 * 1024**2 @@ -153,6 +158,7 @@ async def send_batch_data( self, session_id: str, data_keys: List[str], + readers: List["StorageFileObject"], is_transferring_list: List[bool], remote_band: BandType, block_size: int = None, @@ -175,7 +181,9 @@ async def send_batch_data( if to_send_keys: logger.debug("Start sending %s", to_send_keys) block_size = block_size or self._transfer_block_size - await self._send_data(receiver_ref, session_id, to_send_keys, block_size) + await self._send_data( + receiver_ref, readers, session_id, to_send_keys, block_size + ) logger.debug("Done sending %s", to_send_keys) if to_wait_keys: logger.debug("Start waiting %s", to_wait_keys) From 010871db0ad077b6027588e57163f6271b4f686c Mon Sep 17 00:00:00 2001 From: UranusSeven <109661872+UranusSeven@users.noreply.github.com> Date: Wed, 14 Jun 2023 10:40:59 +0800 Subject: [PATCH 18/25] Fix --- .../xorbits/_mars/services/storage/handler.py | 46 +++++++++-------- .../_mars/services/storage/transfer.py | 49 +++++++++++++++++-- 2 files changed, 70 insertions(+), 25 deletions(-) diff --git a/python/xorbits/_mars/services/storage/handler.py b/python/xorbits/_mars/services/storage/handler.py index 77929fc87..badf38780 100644 --- a/python/xorbits/_mars/services/storage/handler.py +++ b/python/xorbits/_mars/services/storage/handler.py @@ -554,10 +554,10 @@ async def _fetch_via_transfer( ] ) - get_infos = [] + get_info_tasks = [] pin_tasks = [] for data_key in data_keys: - get_infos.append( + get_info_tasks.append( remote_data_manager_ref.get_data_info.delay( session_id, data_key, remote_band[1], error ) @@ -567,11 +567,10 @@ async def _fetch_via_transfer( session_id, data_key, remote_band[1], error ) ) - + logger.debug("Getting data infos for %s", data_keys) + infos = await remote_data_manager_ref.get_data_info.batch(*get_info_tasks) logger.debug("Pining %s", data_keys) await remote_data_manager_ref.pin.batch(*pin_tasks) - logger.debug("Getting data infos for %s", data_keys) - infos = await remote_data_manager_ref.get_data_info.batch(*get_infos) filtered = [ (data_info, data_key) @@ -588,34 +587,41 @@ async def _fetch_via_transfer( if level is None: level = infos[0].level - logger.debug("Requesting quota for %s", data_keys) - await self.request_quota_with_spill(level, sum(data_sizes)) - + logger.debug("Creating writers for %s", data_keys) receiver_ref: mo.ActorRefType[ReceiverManagerActor] = await mo.actor_ref( address=self.address, uid=ReceiverManagerActor.gen_uid(fetch_band_name), ) - is_transferring_list = await receiver_ref.open_writers( - session_id, - data_keys, - data_sizes, - sub_infos, - level, + open_writer_tasks = [] + for data_key, data_size, sub_info in zip(data_keys, data_sizes, sub_infos): + open_writer_tasks.append( + self.open_writer.delay( + session_id, data_key, data_size, level, request_quota=False + ) + ) + writers = await self.open_writer.batch(*open_writer_tasks) + is_transferring_list = await receiver_ref.add_writers( + session_id, data_keys, data_sizes, sub_infos, writers, level ) - if not is_transferring_list: - # no data to be transferred - return - - logger.debug("Writers have been opened for %s", data_keys) + required_quota = sum( + [ + data_size + for is_transferring, data_size in zip(is_transferring_list, data_sizes) + if is_transferring + ] + ) + await self.request_quota_with_spill(level, required_quota) + logger.debug("Writers have been created for %s", data_keys) + logger.debug("Creating readers for %s", data_keys) sender_ref: mo.ActorRefType[SenderManagerActor] = await mo.actor_ref( address=remote_band[0], uid=SenderManagerActor.gen_uid(remote_band[1]) ) - open_reader_tasks = [] for data_key in data_keys: open_reader_tasks.append(self.open_reader.delay(session_id, data_key)) readers = await self.open_reader.batch(*open_reader_tasks) + logger.debug("Readers have been created for %s", data_keys) await sender_ref.send_batch_data( session_id, diff --git a/python/xorbits/_mars/services/storage/transfer.py b/python/xorbits/_mars/services/storage/transfer.py index 081c4eafe..e483e8bf1 100644 --- a/python/xorbits/_mars/services/storage/transfer.py +++ b/python/xorbits/_mars/services/storage/transfer.py @@ -19,6 +19,7 @@ from io import UnsupportedOperation from typing import Any, Dict, List, Union from typing import TYPE_CHECKING, Dict, List +from typing import TYPE_CHECKING, Dict, List, Optional, Union import xoscar as mo from xoscar.core import BufferRef @@ -222,15 +223,16 @@ def __init__( storage_handler_ref: mo.ActorRefType[StorageHandlerActor] = None, ): self._quota_refs = quota_refs - self._storage_handler = storage_handler_ref + # self._storage_handler = storage_handler_ref self._writing_infos: Dict[tuple, WritingInfo] = dict() self._lock = asyncio.Lock() async def __post_create__(self): - if self._storage_handler is None: # for test - self._storage_handler = await mo.actor_ref( - self.address, StorageHandlerActor.gen_uid("numa-0") - ) + # if self._storage_handler is None: # for test + # self._storage_handler = await mo.actor_ref( + # self.address, StorageHandlerActor.gen_uid("numa-0") + # ) + pass async def get_buffers( self, @@ -290,11 +292,48 @@ async def handle_transmission_cancellation(self, session_id: str, data_keys: Lis def gen_uid(cls, band_name: str): return f"receiver_manager_{band_name}" + def lock(self): + self._lock.acquire() + + def unlock(self): + self._lock.release() + def _decref_writing_key(self, session_id: str, data_key: str): self._writing_infos[(session_id, data_key)].ref_counts -= 1 if self._writing_infos[(session_id, data_key)].ref_counts == 0: del self._writing_infos[(session_id, data_key)] + def is_writing_info_existed( + self, session_id: str, data_keys: List[Union[str, tuple]] + ) -> List[bool]: + return [(session_id, data_key) in self._writing_infos for data_key in data_keys] + + async def add_writers( + self, + session_id: str, + data_keys: List[Union[str, tuple]], + data_sizes: List[int], + sub_infos: List, + writers: List[Optional[WrappedStorageFileObject]], + level: StorageLevel, + ) -> List[bool]: + is_transferring: List[bool] = [] + async with self._lock: + for data_key, data_size, sub_info, writer in zip( + data_keys, data_sizes, sub_infos, writers + ): + if (session_id, data_key) not in self._writing_infos: + is_transferring.append(False) + + self._writing_infos[(session_id, data_key)] = WritingInfo( + writer, data_size, level, asyncio.Event(), 1 + ) + if sub_info is not None: + writer._sub_key_infos = sub_info + else: + is_transferring.append(True) + return is_transferring + async def create_writers( self, session_id: str, From 4a629ab1c9346c1cb6c02825c5bba8b50c108b8c Mon Sep 17 00:00:00 2001 From: UranusSeven <109661872+UranusSeven@users.noreply.github.com> Date: Thu, 15 Jun 2023 06:45:53 +0000 Subject: [PATCH 19/25] checkpoint --- python/xorbits/_mars/services/storage/core.py | 9 +++++++ .../xorbits/_mars/services/storage/handler.py | 25 +++++++++++-------- .../_mars/services/storage/transfer.py | 15 +++-------- 3 files changed, 27 insertions(+), 22 deletions(-) diff --git a/python/xorbits/_mars/services/storage/core.py b/python/xorbits/_mars/services/storage/core.py index cb064ef5d..f70a1b5ff 100644 --- a/python/xorbits/_mars/services/storage/core.py +++ b/python/xorbits/_mars/services/storage/core.py @@ -93,6 +93,12 @@ async def clean_up(self): self._file.close() async def close(self): + logger.debug( + "Writer closed for %s, %s on %s", + self._session_id, + self._data_key, + self._data_manager.address, + ) self._file.close() if self._object_id is None: # for some backends like vineyard, @@ -322,6 +328,9 @@ def delete_data_info( level: StorageLevel, band_name: str, ): + logger.debug( + "Deleting %s, %s from %s, %s", session_id, data_key, level, band_name + ) if (session_id, data_key) in self._data_key_to_infos: self._data_info_list[level, band_name].pop((session_id, data_key)) self._spill_strategy[level, band_name].record_delete_info( diff --git a/python/xorbits/_mars/services/storage/handler.py b/python/xorbits/_mars/services/storage/handler.py index badf38780..1dda35835 100644 --- a/python/xorbits/_mars/services/storage/handler.py +++ b/python/xorbits/_mars/services/storage/handler.py @@ -66,6 +66,7 @@ def __init__( self._quota_refs = quota_refs self._band_name = band_name self._supervisor_address = None + self._lock = asyncio.Lock() @classmethod def gen_uid(cls, band_name: str): @@ -284,6 +285,7 @@ async def delete_object( level: StorageLevel, ): data_key = await self._data_manager_ref.get_store_key(session_id, data_key) + logger.debug("Delete object %s, %s on %s", session_id, data_key, self.address) await self._data_manager_ref.delete_data_info( session_id, data_key, level, self._band_name ) @@ -292,6 +294,7 @@ async def delete_object( @mo.extensible async def delete(self, session_id: str, data_key: str, error: str = "raise"): + logger.debug("Delete %s, %s on %s", session_id, data_key, self.address) if error not in ("raise", "ignore"): # pragma: no cover raise ValueError("error must be raise or ignore") @@ -367,6 +370,9 @@ async def batch_delete(self, args_list, kwargs_list): session_id, key, level, info.band ) ) + logger.debug( + "Batch delete %s, %s on %s", session_id, key, self.address + ) to_removes.append((level, info.object_id)) level_sizes[level] += info.store_size @@ -607,28 +613,26 @@ async def _fetch_via_transfer( [ data_size for is_transferring, data_size in zip(is_transferring_list, data_sizes) - if is_transferring + if not is_transferring ] ) await self.request_quota_with_spill(level, required_quota) logger.debug("Writers have been created for %s", data_keys) - logger.debug("Creating readers for %s", data_keys) + logger.debug( + "Start transferring %s from %s to %s", + data_keys, + remote_band, + (self.address, fetch_band_name), + ) sender_ref: mo.ActorRefType[SenderManagerActor] = await mo.actor_ref( address=remote_band[0], uid=SenderManagerActor.gen_uid(remote_band[1]) ) - open_reader_tasks = [] - for data_key in data_keys: - open_reader_tasks.append(self.open_reader.delay(session_id, data_key)) - readers = await self.open_reader.batch(*open_reader_tasks) - logger.debug("Readers have been created for %s", data_keys) - await sender_ref.send_batch_data( session_id, data_keys, - readers, is_transferring_list, - (self._data_manager_ref.address, fetch_band_name), + (self.address, fetch_band_name), ) logger.debug("Finish fetching %s from band %s", data_keys, remote_band) @@ -692,6 +696,7 @@ async def fetch_batch( for data_key, bands in zip(missing_keys, metas): if bands is not None: remote_keys[bands["bands"][0]].add(data_key) + transfer_tasks = [] fetch_keys = [] for band, keys in remote_keys.items(): diff --git a/python/xorbits/_mars/services/storage/transfer.py b/python/xorbits/_mars/services/storage/transfer.py index e483e8bf1..4273e9a8b 100644 --- a/python/xorbits/_mars/services/storage/transfer.py +++ b/python/xorbits/_mars/services/storage/transfer.py @@ -17,9 +17,7 @@ import logging from dataclasses import dataclass from io import UnsupportedOperation -from typing import Any, Dict, List, Union -from typing import TYPE_CHECKING, Dict, List -from typing import TYPE_CHECKING, Dict, List, Optional, Union +from typing import Dict, List, Optional, Union, Any import xoscar as mo from xoscar.core import BufferRef @@ -31,10 +29,6 @@ from .core import DataManagerActor, WrappedStorageFileObject from .handler import StorageHandlerActor -if TYPE_CHECKING: - from ...storage.core import StorageFileObject - - DEFAULT_TRANSFER_BLOCK_SIZE = 4 * 1024**2 @@ -159,7 +153,6 @@ async def send_batch_data( self, session_id: str, data_keys: List[str], - readers: List["StorageFileObject"], is_transferring_list: List[bool], remote_band: BandType, block_size: int = None, @@ -180,11 +173,9 @@ async def send_batch_data( ] = await self.get_receiver_ref(remote_band[0], remote_band[1]) if to_send_keys: - logger.debug("Start sending %s", to_send_keys) + logger.debug("Start sending %s to %s", to_send_keys, receiver_ref.address) block_size = block_size or self._transfer_block_size - await self._send_data( - receiver_ref, readers, session_id, to_send_keys, block_size - ) + await self._send_data(receiver_ref, session_id, to_send_keys, block_size) logger.debug("Done sending %s", to_send_keys) if to_wait_keys: logger.debug("Start waiting %s", to_wait_keys) From faa2096e2297042bd87c5a3fbfaa2affe117971c Mon Sep 17 00:00:00 2001 From: ChengjieLi Date: Mon, 7 Aug 2023 14:42:58 +0800 Subject: [PATCH 20/25] fix --- .../xorbits/_mars/services/storage/handler.py | 102 ++++-- .../services/storage/tests/test_transfer.py | 318 ++++++++++-------- .../storage/tests/test_transfer_gpu.py | 68 ++-- .../_mars/services/storage/transfer.py | 119 +------ python/xorbits/_mars/storage/cuda.py | 10 +- 5 files changed, 295 insertions(+), 322 deletions(-) diff --git a/python/xorbits/_mars/services/storage/handler.py b/python/xorbits/_mars/services/storage/handler.py index 1dda35835..3986ac580 100644 --- a/python/xorbits/_mars/services/storage/handler.py +++ b/python/xorbits/_mars/services/storage/handler.py @@ -388,6 +388,7 @@ async def batch_delete(self, args_list, kwargs_list): await self._quota_refs[level].release_quota(size) @mo.extensible + @mo.no_lock async def open_reader(self, session_id: str, data_key: str) -> StorageFileObject: data_info = await self._data_manager_ref.get_data_info( session_id, data_key, self._band_name @@ -396,6 +397,7 @@ async def open_reader(self, session_id: str, data_key: str) -> StorageFileObject return reader @open_reader.batch + @mo.no_lock async def batch_open_readers(self, args_list, kwargs_list): get_data_infos = [] for args, kwargs in zip(args_list, kwargs_list): @@ -528,7 +530,21 @@ async def _fetch_remote( await self._data_manager_ref.put_data_info.batch(*put_data_info_delays) await asyncio.gather(*fetch_tasks) - async def _fetch_via_transfer( + async def get_receive_manager_ref(self, band_name: str): + from .transfer import ReceiverManagerActor + + return await mo.actor_ref( + address=self.address, + uid=ReceiverManagerActor.gen_uid(band_name), + ) + + @staticmethod + async def get_send_manager_ref(address: str, band: str): + from .transfer import SenderManagerActor + + return await mo.actor_ref(address=address, uid=SenderManagerActor.gen_uid(band)) + + async def fetch_via_transfer( self, session_id: str, data_keys: List[Union[str, tuple]], @@ -593,31 +609,44 @@ async def _fetch_via_transfer( if level is None: level = infos[0].level - logger.debug("Creating writers for %s", data_keys) - receiver_ref: mo.ActorRefType[ReceiverManagerActor] = await mo.actor_ref( - address=self.address, - uid=ReceiverManagerActor.gen_uid(fetch_band_name), - ) + receiver_ref: mo.ActorRefType[ + ReceiverManagerActor + ] = await self.get_receive_manager_ref(fetch_band_name) + + await self.request_quota_with_spill(level, sum(data_sizes)) + open_writer_tasks = [] for data_key, data_size, sub_info in zip(data_keys, data_sizes, sub_infos): open_writer_tasks.append( self.open_writer.delay( - session_id, data_key, data_size, level, request_quota=False + session_id, + data_key, + data_size, + level, + request_quota=False, + band_name=fetch_band_name, ) ) writers = await self.open_writer.batch(*open_writer_tasks) is_transferring_list = await receiver_ref.add_writers( session_id, data_keys, data_sizes, sub_infos, writers, level ) - required_quota = sum( - [ - data_size - for is_transferring, data_size in zip(is_transferring_list, data_sizes) - if not is_transferring - ] - ) - await self.request_quota_with_spill(level, required_quota) - logger.debug("Writers have been created for %s", data_keys) + + to_send_keys = [] + to_wait_keys = [] + wait_sizes = [] + for data_key, is_transferring, _size in zip( + data_keys, is_transferring_list, data_sizes + ): + if is_transferring: + to_wait_keys.append(data_key) + wait_sizes.append(_size) + else: + to_send_keys.append(data_key) + + # Overapplied the quota for these wait keys, and now need to update the quota + if to_wait_keys: + self._quota_refs[level].update_quota(-sum(wait_sizes)) logger.debug( "Start transferring %s from %s to %s", @@ -625,16 +654,35 @@ async def _fetch_via_transfer( remote_band, (self.address, fetch_band_name), ) - sender_ref: mo.ActorRefType[SenderManagerActor] = await mo.actor_ref( - address=remote_band[0], uid=SenderManagerActor.gen_uid(remote_band[1]) - ) - await sender_ref.send_batch_data( - session_id, - data_keys, - is_transferring_list, - (self.address, fetch_band_name), - ) - logger.debug("Finish fetching %s from band %s", data_keys, remote_band) + sender_ref: mo.ActorRefType[ + SenderManagerActor + ] = await self.get_send_manager_ref(remote_band[0], remote_band[1]) + + try: + await sender_ref.send_batch_data( + session_id, + data_keys, + to_send_keys, + to_wait_keys, + (self.address, fetch_band_name), + ) + await receiver_ref.handle_transmission_done(session_id, to_send_keys) + except asyncio.CancelledError: + keys_to_delete = await receiver_ref.handle_transmission_cancellation( + session_id, to_send_keys + ) + for key in keys_to_delete: + await self.delete(session_id, key, error="ignore") + raise + + unpin_tasks = [] + for data_key in data_keys: + unpin_tasks.append( + remote_data_manager_ref.unpin.delay( + session_id, [data_key], remote_band[1], error="ignore" + ) + ) + await remote_data_manager_ref.unpin.batch(*unpin_tasks) async def fetch_batch( self, @@ -708,7 +756,7 @@ async def fetch_batch( else: # fetch via transfer transfer_tasks.append( - self._fetch_via_transfer( + self.fetch_via_transfer( session_id, list(keys), level, band, band_name or band[1], error ) ) diff --git a/python/xorbits/_mars/services/storage/tests/test_transfer.py b/python/xorbits/_mars/services/storage/tests/test_transfer.py index 96b00b6e5..b8bcf0d25 100644 --- a/python/xorbits/_mars/services/storage/tests/test_transfer.py +++ b/python/xorbits/_mars/services/storage/tests/test_transfer.py @@ -17,17 +17,18 @@ import os import sys import tempfile -from typing import List, Union +from typing import Dict, List import numpy as np import pandas as pd import pytest import xoscar as mo from xoscar.backends.allocate_strategy import IdleLabel +from xoscar.errors import NoIdleSlot from ....oscar import create_actor_pool from ....storage import StorageLevel -from ..core import DataManagerActor, StorageManagerActor, StorageQuotaActor +from ..core import StorageManagerActor, StorageQuotaActor from ..errors import DataNotExist from ..handler import StorageHandlerActor from ..transfer import ReceiverManagerActor, SenderManagerActor @@ -62,10 +63,22 @@ async def start_pool(): await worker_pool_2.stop() +async def _get_io_address(pool): + pool_config = (await mo.get_pool_config(pool.external_address)).as_dict() + return [ + v["external_address"][0] + for k, v in pool_config["pools"].items() + if v["label"] == "io" + ][0] + + @pytest.fixture async def create_actors(actor_pools): worker_pool_1, worker_pool_2 = actor_pools + io1 = await _get_io_address(worker_pool_1) + io2 = await _get_io_address(worker_pool_2) + tmp_dir = tempfile.mkdtemp() storage_configs = {"shared_memory": {}, "disk": {"root_dirs": f"{tmp_dir}"}} @@ -82,7 +95,39 @@ async def create_actors(actor_pools): uid=StorageManagerActor.default_uid(), address=worker_pool_2.external_address, ) - yield worker_pool_1.external_address, worker_pool_2.external_address + yield worker_pool_1.external_address, worker_pool_2.external_address, io1, io2 + try: + await mo.destroy_actor(manager_ref1) + await mo.destroy_actor(manager_ref2) + except FileNotFoundError: + pass + assert not os.path.exists(tmp_dir) + + +@pytest.fixture +async def create_actors_mock(actor_pools): + worker_pool_1, worker_pool_2 = actor_pools + + io1 = await _get_io_address(worker_pool_1) + io2 = await _get_io_address(worker_pool_2) + + tmp_dir = tempfile.mkdtemp() + storage_configs = {"shared_memory": {}, "disk": {"root_dirs": f"{tmp_dir}"}} + + manager_ref1 = await mo.create_actor( + MockStorageManagerActor, + storage_configs, + uid=MockStorageManagerActor.default_uid(), + address=worker_pool_1.external_address, + ) + + manager_ref2 = await mo.create_actor( + MockStorageManagerActor, + storage_configs, + uid=MockStorageManagerActor.default_uid(), + address=worker_pool_2.external_address, + ) + yield worker_pool_1.external_address, worker_pool_2.external_address, io1, io2 try: await mo.destroy_actor(manager_ref1) await mo.destroy_actor(manager_ref2) @@ -93,16 +138,16 @@ async def create_actors(actor_pools): @pytest.mark.asyncio async def test_simple_transfer(create_actors): - worker_address_1, worker_address_2 = create_actors + worker_address_1, worker_address_2, io1, io2 = create_actors data1 = np.random.rand(100, 100) data2 = pd.DataFrame(np.random.randint(0, 100, (500, 10))) - storage_handler1 = await mo.actor_ref( - uid=StorageHandlerActor.gen_uid("numa-0"), address=worker_address_1 + storage_handler1: mo.ActorRefType[StorageHandlerActor] = await mo.actor_ref( + uid=StorageHandlerActor.gen_uid("numa-0"), address=io1 ) - storage_handler2 = await mo.actor_ref( - uid=StorageHandlerActor.gen_uid("numa-0"), address=worker_address_2 + storage_handler2: mo.ActorRefType[StorageHandlerActor] = await mo.actor_ref( + uid=StorageHandlerActor.gen_uid("numa-0"), address=io2 ) for level in (StorageLevel.MEMORY, StorageLevel.DISK): @@ -111,65 +156,26 @@ async def test_simple_transfer(create_actors): await storage_handler1.put(session_id, "data_key2", data2, level) await storage_handler2.put(session_id, "data_key3", data2, level) - receiver_actor: mo.ActorRefType[ReceiverManagerActor] = await mo.actor_ref( - address=worker_address_2, uid=ReceiverManagerActor.gen_uid("numa-0") - ) - sender_actor: mo.ActorRefType[SenderManagerActor] = await mo.actor_ref( - address=worker_address_1, uid=SenderManagerActor.gen_uid("numa-0") - ) - - is_transferring = await receiver_actor.open_writers( - session_id, - ["data_key1", "data_key2"], - StorageLevel.MEMORY, - (sender_actor.address, "numa-0"), - "raise", - ) - - # send data to worker2 from worker1 - await sender_actor.send_batch_data( - session_id, - ["data_key1"], - is_transferring, - (receiver_actor.address, "numa-0"), - block_size=1000, - ) - - await sender_actor.send_batch_data( - session_id, - ["data_key2"], - is_transferring, - (receiver_actor.address, "numa-0"), - block_size=1000, - ) - - get_data1 = await storage_handler2.get(session_id, "data_key1") - np.testing.assert_array_equal(data1, get_data1) + await storage_handler2.fetch_via_transfer( + session_id, + ["data_key1", "data_key2"], + level, + (io1, "numa-0"), + "numa-0", + "raise", + ) - get_data2 = await storage_handler2.get(session_id, "data_key2") - pd.testing.assert_frame_equal(data2, get_data2) + get_data1 = await storage_handler2.get(session_id, "data_key1") + np.testing.assert_array_equal(data1, get_data1) - # send data to worker1 from worker2 - receiver_actor: mo.ActorRefType[ReceiverManagerActor] = await mo.actor_ref( - address=worker_address_1, uid=ReceiverManagerActor.gen_uid("numa-0") - ) - sender_actor: mo.ActorRefType[SenderManagerActor] = await mo.actor_ref( - address=worker_address_2, uid=SenderManagerActor.gen_uid("numa-0") - ) - - is_transferring = await receiver_actor.open_writers( - session_id, - ["data_key3"], - StorageLevel.MEMORY, - (sender_actor.address, "numa-0"), - "raise", - ) + get_data2 = await storage_handler2.get(session_id, "data_key2") + pd.testing.assert_frame_equal(data2, get_data2) - await sender_actor.send_batch_data( - session_id, ["data_key3"], is_transferring, (receiver_actor.address, "numa-0") - ) - get_data3 = await storage_handler1.get(session_id, "data_key3") - pd.testing.assert_frame_equal(data2, get_data3) + await storage_handler1.fetch_via_transfer( + session_id, ["data_key3"], level, (io2, "numa-0"), "numa-0", "raise" + ) + get_data3 = await storage_handler1.get(session_id, "data_key3") + pd.testing.assert_frame_equal(data2, get_data3) # test for cancelling happens when writing @@ -181,12 +187,12 @@ class MockSenderManagerActor(SenderManagerActor): @staticmethod async def get_receiver_ref(address: str, band_name: str): return await mo.actor_ref( - address=address, uid=MockReceiverManagerActor.default_uid() + address=address, uid=MockReceiverManagerActor.gen_uid(band_name) ) async def _copy_to_receiver( self, - receiver_ref: mo.ActorRefType["ReceiverManagerActor"], + receiver_ref: mo.ActorRefType["MockReceiverManagerActor"], local_buffers: List, remote_buffers: List, session_id: str, @@ -204,35 +210,73 @@ async def _copy_to_receiver( ) -# test for cancelling happens when creating writer -class MockReceiverManagerActor2(ReceiverManagerActor): - async def create_writers( - self, session_id, data_keys, data_sizes, level, sub_infos, band_name - ): - await asyncio.sleep(3) - return await super().create_writers( - session_id, data_keys, data_sizes, level, sub_infos, band_name +class MockStorageHandlerActor(StorageHandlerActor): + async def get_receive_manager_ref(self, band_name: str): + return await mo.actor_ref( + address=self.address, + uid=MockReceiverManagerActor.gen_uid(band_name), ) - -class MockSenderManagerActor2(SenderManagerActor): @staticmethod - async def get_receiver_ref(address: str, band_name: str): + async def get_send_manager_ref(address: str, band: str): return await mo.actor_ref( - address=address, uid=MockReceiverManagerActor2.default_uid() + address=address, uid=MockSenderManagerActor.gen_uid(band) ) +class MockStorageManagerActor(StorageManagerActor): + def __init__( + self, storage_configs: Dict, transfer_block_size: int = None, **kwargs + ): + super().__init__(storage_configs, transfer_block_size, **kwargs) + self._handler_cls = MockStorageHandlerActor + + async def _create_transfer_actors(self): + default_band_name = "numa-0" + sender_strategy = IdleLabel("io", "sender") + receiver_strategy = IdleLabel("io", "receiver") + handler_strategy = IdleLabel("io", "handler") + while True: + try: + handler_ref = await mo.create_actor( + MockStorageHandlerActor, + self._init_params[default_band_name], + self._data_manager, + self._spill_managers[default_band_name], + self._quotas[default_band_name], + default_band_name, + uid=MockStorageHandlerActor.gen_uid(default_band_name), + address=self.address, + allocate_strategy=handler_strategy, + ) + await mo.create_actor( + MockSenderManagerActor, + data_manager_ref=self._data_manager, + storage_handler_ref=handler_ref, + uid=MockSenderManagerActor.gen_uid(default_band_name), + address=self.address, + allocate_strategy=sender_strategy, + ) + + await mo.create_actor( + MockReceiverManagerActor, + self._quotas[default_band_name], + handler_ref, + address=self.address, + uid=MockReceiverManagerActor.gen_uid(default_band_name), + allocate_strategy=receiver_strategy, + ) + except NoIdleSlot: + break + + @pytest.mark.parametrize( "mock_sender_cls, mock_receiver_cls", - [ - (MockSenderManagerActor, MockReceiverManagerActor), - (MockSenderManagerActor2, MockReceiverManagerActor2), - ], + [(MockSenderManagerActor, MockReceiverManagerActor)], ) @pytest.mark.asyncio -async def test_cancel_transfer(create_actors, mock_sender_cls, mock_receiver_cls): - worker_address_1, worker_address_2 = create_actors +async def test_cancel_transfer(create_actors_mock, mock_sender_cls, mock_receiver_cls): + _, _, io1, io2 = create_actors_mock session_id = "mock_session" quota_refs = { @@ -240,33 +284,15 @@ async def test_cancel_transfer(create_actors, mock_sender_cls, mock_receiver_cls StorageQuotaActor, StorageLevel.MEMORY, 5 * 1024 * 1024, - address=worker_address_2, + address=io2, uid=StorageQuotaActor.gen_uid("numa-0", StorageLevel.MEMORY), ) } - data_manager_ref = await mo.actor_ref( - uid=DataManagerActor.default_uid(), address=worker_address_1 - ) storage_handler1 = await mo.actor_ref( - uid=StorageHandlerActor.gen_uid("numa-0"), address=worker_address_1 + uid=MockStorageHandlerActor.gen_uid("numa-0"), address=io1 ) storage_handler2 = await mo.actor_ref( - uid=StorageHandlerActor.gen_uid("numa-0"), address=worker_address_2 - ) - - sender_actor: mo.ActorRefType[SenderManagerActor] = await mo.create_actor( - mock_sender_cls, - data_manager_ref=data_manager_ref, - uid=mock_sender_cls.default_uid(), - address=worker_address_1, - allocate_strategy=IdleLabel("io", "mock_sender_cls"), - ) - receiver_actor: mo.ActorRefType[ReceiverManagerActor] = await mo.create_actor( - mock_receiver_cls, - quota_refs, - uid=mock_receiver_cls.default_uid(), - address=worker_address_2, - allocate_strategy=IdleLabel("io", "mock_receiver_cls"), + uid=MockStorageHandlerActor.gen_uid("numa-0"), address=io2 ) data1 = np.random.rand(10, 10) @@ -276,31 +302,18 @@ async def test_cancel_transfer(create_actors, mock_sender_cls, mock_receiver_cls used_before = (await quota_refs[StorageLevel.MEMORY].get_quota())[1] - async def transfer( - data_keys: List[Union[str, tuple]], - sender: mo.ActorRefType[SenderManagerActor], - receiver: mo.ActorRefType[ReceiverManagerActor], - ): - is_transferring_list = await receiver.open_writers( + transfer_task = asyncio.create_task( + storage_handler2.fetch_via_transfer( session_id, - data_keys, + ["data_key1"], StorageLevel.MEMORY, - (sender.address, "numa-0"), + (io1, "numa-0"), + "numa-0", "raise", ) - assert len(is_transferring_list) > 0 - await sender.send_batch_data( - session_id, - data_keys, - is_transferring_list, - (receiver.address, "numa-0"), - ) - - transfer_task = asyncio.create_task( - transfer(["data_key1"], sender_actor, receiver_actor) ) - await asyncio.sleep(0.5) + await asyncio.sleep(1) transfer_task.cancel() with pytest.raises(asyncio.CancelledError): @@ -313,7 +326,14 @@ async def transfer( await storage_handler2.get(session_id, "data_key1") transfer_task = asyncio.create_task( - transfer(["data_key1"], sender_actor, receiver_actor) + storage_handler2.fetch_via_transfer( + session_id, + ["data_key1"], + StorageLevel.MEMORY, + (io1, "numa-0"), + "numa-0", + "raise", + ) ) await transfer_task get_data = await storage_handler2.get(session_id, "data_key1") @@ -322,12 +342,26 @@ async def transfer( # cancel when fetch the same data Simultaneously if mock_sender_cls is MockSenderManagerActor: transfer_task1 = asyncio.create_task( - transfer(["data_key2"], sender_actor, receiver_actor) + storage_handler2.fetch_via_transfer( + session_id, + ["data_key1"], + StorageLevel.MEMORY, + (io1, "numa-0"), + "numa-0", + "raise", + ) ) transfer_task2 = asyncio.create_task( - transfer(["data_key2"], sender_actor, receiver_actor) + storage_handler2.fetch_via_transfer( + session_id, + ["data_key2"], + StorageLevel.MEMORY, + (io1, "numa-0"), + "numa-0", + "raise", + ) ) - await asyncio.sleep(0.5) + await asyncio.sleep(1) transfer_task1.cancel() with pytest.raises(asyncio.CancelledError): await transfer_task1 @@ -338,37 +372,37 @@ async def transfer( @pytest.mark.asyncio async def test_transfer_same_data(create_actors): - worker_address_1, worker_address_2 = create_actors + worker_address_1, worker_address_2, io1, io2 = create_actors session_id = "mock_session" data1 = np.random.rand(100, 100) storage_handler1 = await mo.actor_ref( - uid=StorageHandlerActor.gen_uid("numa-0"), address=worker_address_1 + uid=StorageHandlerActor.gen_uid("numa-0"), address=io1 ) - await mo.actor_ref( - uid=StorageHandlerActor.gen_uid("numa-0"), address=worker_address_2 + storage_handler2 = await mo.actor_ref( + uid=StorageHandlerActor.gen_uid("numa-0"), address=io2 ) await storage_handler1.put(session_id, "data_key1", data1, StorageLevel.MEMORY) - receiver_actor: mo.ActorRefType[ReceiverManagerActor] = await mo.actor_ref( - address=worker_address_2, uid=ReceiverManagerActor.gen_uid("numa-0") - ) - task1 = receiver_actor.open_writers( + task1 = storage_handler2.fetch_via_transfer( session_id, ["data_key1"], StorageLevel.MEMORY, - (worker_address_1, "numa-0"), + (io1, "numa-0"), + "numa-0", "raise", ) - task2 = receiver_actor.open_writers( + task2 = storage_handler2.fetch_via_transfer( session_id, ["data_key1"], StorageLevel.MEMORY, - (worker_address_1, "numa-0"), + (io1, "numa-0"), + "numa-0", "raise", ) - results = await asyncio.gather(task1, task2) - assert [False] in results - assert [True] in results + await asyncio.gather(task1, task2) + + get_data = await storage_handler2.get(session_id, "data_key1") + np.testing.assert_array_equal(get_data, data1) diff --git a/python/xorbits/_mars/services/storage/tests/test_transfer_gpu.py b/python/xorbits/_mars/services/storage/tests/test_transfer_gpu.py index c3ba3fe6f..0d0ca35cc 100644 --- a/python/xorbits/_mars/services/storage/tests/test_transfer_gpu.py +++ b/python/xorbits/_mars/services/storage/tests/test_transfer_gpu.py @@ -28,7 +28,6 @@ from ....utils import lazy_import from ..core import StorageManagerActor from ..handler import StorageHandlerActor -from ..transfer import SenderManagerActor cupy = lazy_import("cupy") cudf = lazy_import("cudf") @@ -109,7 +108,7 @@ def _generate_band_scheme(): gpu_bands = [] if gpu_counts: gpu_bands.extend([(f"gpu-{i}", f"gpu-{i+1}") for i in range(gpu_counts - 1)]) - bands = [("numa-0", "numa-0"), ("gpu-0", "gpu-0")] + bands = [("gpu-0", "gpu-0")] bands.extend(gpu_bands) schemes = [(None, None), ("ucx", "ucx"), (None, "ucx"), ("ucx", None)] for band in bands: @@ -137,7 +136,7 @@ async def test_simple_transfer(create_actors): storage_handler1 = await mo.actor_ref( uid=StorageHandlerActor.gen_uid(bands[0]), address=worker_address_1 ) - storage_handler2 = await mo.actor_ref( + storage_handler2: mo.ActorRefType[StorageHandlerActor] = await mo.actor_ref( uid=StorageHandlerActor.gen_uid(bands[1]), address=worker_address_2 ) @@ -146,27 +145,13 @@ async def test_simple_transfer(create_actors): await storage_handler1.put(session_id, "data_key2", data2, storage_level) await storage_handler2.put(session_id, "data_key3", data2, storage_level) - sender_actor = await mo.actor_ref( - address=worker_address_1, uid=SenderManagerActor.gen_uid(bands[0]) - ) - - # send data to worker2 from worker1 - await sender_actor.send_batch_data( + await storage_handler2.fetch_via_transfer( session_id, - ["data_key1"], - worker_address_2, + ["data_key1", "data_key2"], storage_level, - block_size=1000, - band_name=bands[1], - ) - - await sender_actor.send_batch_data( - session_id, - ["data_key2"], - worker_address_2, - storage_level, - block_size=1000, - band_name=bands[1], + (worker_address_1, bands[0]), + bands[1], + "raise", ) get_data1 = await storage_handler2.get(session_id, "data_key1") @@ -175,18 +160,15 @@ async def test_simple_transfer(create_actors): get_data2 = await storage_handler2.get(session_id, "data_key2") xpd.testing.assert_frame_equal(data2, get_data2) - # send data to worker1 from worker2 - sender_actor = await mo.actor_ref( - address=worker_address_2, uid=SenderManagerActor.gen_uid(bands[1]) - ) - await sender_actor.send_batch_data( + await storage_handler1.fetch_via_transfer( session_id, ["data_key3"], - worker_address_1, storage_level, - block_size=1000, - band_name=bands[0], + (worker_address_2, bands[1]), + bands[0], + "raise", ) + get_data3 = await storage_handler1.get(session_id, "data_key3") xpd.testing.assert_frame_equal(data2, get_data3) @@ -212,29 +194,35 @@ async def test_transfer_same_data(create_actors): ) await storage_handler1.put(session_id, "data_key1", data1, storage_level) - sender_actor = await mo.actor_ref( - address=worker_address_1, uid=SenderManagerActor.gen_uid(bands[0]) + + storage_handler2.fetch_via_transfer( + session_id, + ["data_key1"], + storage_level, + (worker_address_1, bands[0]), + bands[1], + "raise", ) # send data to worker2 from worker1 task1 = asyncio.create_task( - sender_actor.send_batch_data( + storage_handler2.fetch_via_transfer( session_id, ["data_key1"], - worker_address_2, storage_level, - block_size=1000, - band_name=bands[1], + (worker_address_1, bands[0]), + bands[1], + "raise", ) ) task2 = asyncio.create_task( - sender_actor.send_batch_data( + storage_handler2.fetch_via_transfer( session_id, ["data_key1"], - worker_address_2, storage_level, - block_size=1000, - band_name=bands[1], + (worker_address_1, bands[0]), + bands[1], + "raise", ) ) await asyncio.gather(task1, task2) diff --git a/python/xorbits/_mars/services/storage/transfer.py b/python/xorbits/_mars/services/storage/transfer.py index 4273e9a8b..17cb8b027 100644 --- a/python/xorbits/_mars/services/storage/transfer.py +++ b/python/xorbits/_mars/services/storage/transfer.py @@ -17,7 +17,7 @@ import logging from dataclasses import dataclass from io import UnsupportedOperation -from typing import Dict, List, Optional, Union, Any +from typing import Any, Dict, List, Optional, Union import xoscar as mo from xoscar.core import BufferRef @@ -43,6 +43,7 @@ def __init__( data_manager_ref: mo.ActorRefType[DataManagerActor] = None, storage_handler_ref: mo.ActorRefType[StorageHandlerActor] = None, ): + super().__init__() self._band_name = band_name self._data_manager_ref = data_manager_ref self._storage_handler = storage_handler_ref @@ -75,7 +76,6 @@ async def _copy_to_receiver( block_size: int, ): await mo.copy_to(local_buffers, remote_buffers, block_size=block_size) - await receiver_ref.handle_transmission_done(session_id, data_keys) @staticmethod def _get_buffer_size(buf: Any): @@ -145,7 +145,6 @@ async def _send_data( try: await asyncio.shield(write_task) except asyncio.CancelledError: - await receiver_ref.handle_transmission_cancellation(session_id, data_keys) raise @mo.extensible @@ -153,20 +152,14 @@ async def send_batch_data( self, session_id: str, data_keys: List[str], - is_transferring_list: List[bool], + to_send_keys: List, + to_wait_keys: List, remote_band: BandType, block_size: int = None, ): logger.debug( "Begin to send data (%s, %s) to %s", session_id, data_keys, remote_band ) - 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) receiver_ref: mo.ActorRefType[ ReceiverManagerActor @@ -176,19 +169,8 @@ async def send_batch_data( logger.debug("Start sending %s to %s", to_send_keys, receiver_ref.address) block_size = block_size or self._transfer_block_size await self._send_data(receiver_ref, session_id, to_send_keys, block_size) - logger.debug("Done sending %s", to_send_keys) if to_wait_keys: - logger.debug("Start waiting %s", to_wait_keys) await receiver_ref.wait_transfer_done(session_id, to_wait_keys) - logger.debug("Done waiting %s", to_wait_keys) - unpin_tasks = [] - for data_key in data_keys: - unpin_tasks.append( - self._data_manager_ref.unpin.delay( - session_id, [data_key], self._band_name, error="ignore" - ) - ) - await self._data_manager_ref.unpin.batch(*unpin_tasks) logger.debug( "Finish sending data (%s, %s) to %s", session_id, @@ -214,17 +196,10 @@ def __init__( storage_handler_ref: mo.ActorRefType[StorageHandlerActor] = None, ): self._quota_refs = quota_refs - # self._storage_handler = storage_handler_ref + self._storage_handler = storage_handler_ref self._writing_infos: Dict[tuple, WritingInfo] = dict() self._lock = asyncio.Lock() - async def __post_create__(self): - # if self._storage_handler is None: # for test - # self._storage_handler = await mo.actor_ref( - # self.address, StorageHandlerActor.gen_uid("numa-0") - # ) - pass - async def get_buffers( self, session_id: str, @@ -266,39 +241,28 @@ async def handle_transmission_done(self, session_id: str, data_keys: List): self._decref_writing_key(session_id, data_key) async def handle_transmission_cancellation(self, session_id: str, data_keys: List): + data_keys_to_be_deleted = [] async with self._lock: for data_key in data_keys: if (session_id, data_key) in self._writing_infos: if self._writing_infos[(session_id, data_key)].ref_counts == 1: info = self._writing_infos[(session_id, data_key)] await self._quota_refs[info.level].release_quota(info.size) - await self._storage_handler.delete( - session_id, data_key, error="ignore" - ) + data_keys_to_be_deleted.append(data_key) await info.writer.clean_up() info.event.set() self._decref_writing_key(session_id, data_key) + return data_keys_to_be_deleted @classmethod def gen_uid(cls, band_name: str): return f"receiver_manager_{band_name}" - def lock(self): - self._lock.acquire() - - def unlock(self): - self._lock.release() - def _decref_writing_key(self, session_id: str, data_key: str): self._writing_infos[(session_id, data_key)].ref_counts -= 1 if self._writing_infos[(session_id, data_key)].ref_counts == 0: del self._writing_infos[(session_id, data_key)] - def is_writing_info_existed( - self, session_id: str, data_keys: List[Union[str, tuple]] - ) -> List[bool]: - return [(session_id, data_key) in self._writing_infos for data_key in data_keys] - async def add_writers( self, session_id: str, @@ -325,73 +289,6 @@ async def add_writers( is_transferring.append(True) return is_transferring - async def create_writers( - self, - session_id: str, - data_keys: List[str], - data_sizes: List[int], - level: StorageLevel, - sub_infos: List, - band_name: str, - ) -> List[bool]: - tasks = dict() - key_to_sub_infos = dict() - data_key_to_size = dict() - being_processed = [] - for data_key, data_size, sub_info in zip(data_keys, data_sizes, sub_infos): - data_key_to_size[data_key] = data_size - if (session_id, data_key) not in self._writing_infos: - being_processed.append(False) - tasks[data_key] = self._storage_handler.open_writer.delay( - session_id, - data_key, - data_size, - level, - request_quota=False, - band_name=band_name, - ) - key_to_sub_infos[data_key] = sub_info - else: - being_processed.append(True) - self._writing_infos[(session_id, data_key)].ref_counts += 1 - if tasks: - logger.debug("Executing open_writer.batch") - writers = await self._storage_handler.open_writer.batch( - *tuple(tasks.values()) - ) - for data_key, writer in zip(tasks, writers): - self._writing_infos[(session_id, data_key)] = WritingInfo( - writer, data_key_to_size[data_key], level, asyncio.Event(), 1 - ) - if key_to_sub_infos[data_key] is not None: - writer._sub_key_infos = key_to_sub_infos[data_key] - return being_processed - - async def open_writers( - self, - session_id: str, - data_keys: List[str], - data_sizes: List[int], - sub_infos, - level: StorageLevel, - band_name - ) -> List[bool]: - logger.debug("Start opening writers for %s", data_keys) - - async with self._lock: - logger.debug("Opening writers for %s", data_keys) - future = asyncio.create_task( - self.create_writers( - session_id, data_keys, data_sizes, level, sub_infos, band_name - ) - ) - try: - return await future - except asyncio.CancelledError: - await self._quota_refs[level].release_quota(sum(data_sizes)) - future.cancel() - raise - async def wait_transfer_done(self, session_id, data_keys): await asyncio.gather( *[self._writing_infos[(session_id, key)].event.wait() for key in data_keys] diff --git a/python/xorbits/_mars/storage/cuda.py b/python/xorbits/_mars/storage/cuda.py index 7e4dee4a8..97a33c7f3 100644 --- a/python/xorbits/_mars/storage/cuda.py +++ b/python/xorbits/_mars/storage/cuda.py @@ -67,7 +67,10 @@ def set_buffers_by_sizes(self, sizes: List[int]): from rmm import DeviceBuffer self._buffers = [ - _convert_to_cupy_ndarray(DeviceBuffer(size=size)) for size in sizes + cupy.ndarray(shape=0, dtype="u1") + if size == 0 + else _convert_to_cupy_ndarray(DeviceBuffer(size=size)) + for size in sizes ] @property @@ -96,7 +99,10 @@ def _initialize_read(self): self._buffers.append(buf.astype("u1", copy=False)) buffer_types.append(["cuda", buf.size]) elif isinstance(buf, Buffer): - self._buffers.append(_convert_to_cupy_ndarray(buf)) + if buf.size == 0: + self._buffers.append(cupy.ndarray(shape=0, dtype="u1")) + else: + self._buffers.append(_convert_to_cupy_ndarray(buf)) buffer_types.append(["cuda", buf.size]) else: size = getattr(buf, "size", len(buf)) From 4026b4cc57c112fc4e1ce7bf8c48c944cdc247b1 Mon Sep 17 00:00:00 2001 From: Weizheng Lu Date: Fri, 5 Jul 2024 19:16:54 +0800 Subject: [PATCH 21/25] test --- python/xorbits/_mars/services/storage/core.py | 1 + python/xorbits/_mars/services/storage/handler.py | 2 ++ .../xorbits/_mars/services/storage/tests/test_transfer_gpu.py | 4 +++- python/xorbits/_mars/storage/cuda.py | 2 +- python/xorbits/deploy/oscar/file-logging.conf | 4 ++-- 5 files changed, 9 insertions(+), 4 deletions(-) diff --git a/python/xorbits/_mars/services/storage/core.py b/python/xorbits/_mars/services/storage/core.py index f70a1b5ff..45641e296 100644 --- a/python/xorbits/_mars/services/storage/core.py +++ b/python/xorbits/_mars/services/storage/core.py @@ -635,6 +635,7 @@ async def _setup_storage( storage_config = storage_config or dict() init_params, teardown_params = await backend.setup(**storage_config) client = backend(**init_params) + logger.debug(f"_setup_storage: {band_name}, {storage_backend}, {init_params}") self._init_params[band_name][storage_backend] = init_params self._teardown_params[band_name][storage_backend] = teardown_params return client diff --git a/python/xorbits/_mars/services/storage/handler.py b/python/xorbits/_mars/services/storage/handler.py index 3986ac580..7e3634caa 100644 --- a/python/xorbits/_mars/services/storage/handler.py +++ b/python/xorbits/_mars/services/storage/handler.py @@ -229,6 +229,7 @@ async def put( object_info.size - data_info.memory_size ) await self.notify_spillable_space(level) + print(f"StorageHandlerActor put") return data_info @put.batch @@ -555,6 +556,7 @@ async def fetch_via_transfer( ): from .transfer import ReceiverManagerActor, SenderManagerActor + print(f"fetch_via_transfer") logger.debug("Begin to fetch %s from band %s", data_keys, remote_band) remote_data_manager_ref: mo.ActorRefType[DataManagerActor] = await mo.actor_ref( diff --git a/python/xorbits/_mars/services/storage/tests/test_transfer_gpu.py b/python/xorbits/_mars/services/storage/tests/test_transfer_gpu.py index 0d0ca35cc..914241a3c 100644 --- a/python/xorbits/_mars/services/storage/tests/test_transfer_gpu.py +++ b/python/xorbits/_mars/services/storage/tests/test_transfer_gpu.py @@ -68,6 +68,7 @@ async def start_pool(band: str): await pool.start() return pool + print(f"start_pool: {bands[0]}, {bands[1]}") worker_pool_1 = await start_pool(bands[0]) worker_pool_2 = await start_pool(bands[1]) @@ -110,7 +111,8 @@ def _generate_band_scheme(): gpu_bands.extend([(f"gpu-{i}", f"gpu-{i+1}") for i in range(gpu_counts - 1)]) bands = [("gpu-0", "gpu-0")] bands.extend(gpu_bands) - schemes = [(None, None), ("ucx", "ucx"), (None, "ucx"), ("ucx", None)] + schemes = [("ucx", "ucx")] + # schemes = [(None, None), ("ucx", "ucx"), (None, "ucx"), ("ucx", None)] for band in bands: for scheme in schemes: yield band, scheme diff --git a/python/xorbits/_mars/storage/cuda.py b/python/xorbits/_mars/storage/cuda.py index 97a33c7f3..f3a090c50 100644 --- a/python/xorbits/_mars/storage/cuda.py +++ b/python/xorbits/_mars/storage/cuda.py @@ -281,7 +281,7 @@ async def get(self, object_id: str, **kwargs) -> object: if isinstance(buf, cupy.ndarray): new_buffers.append(DeviceBuffer(ptr=buf.data.ptr, size=buf.size)) elif isinstance(buf, CPBuffer): - new_buffers.append(DeviceBuffer(ptr=buf.ptr, size=buf.size)) + new_buffers.append(DeviceBuffer(ptr=buf.owner._ptr, size=buf.size)) else: new_buffers.append(buf) return deserialize(headers, new_buffers) diff --git a/python/xorbits/deploy/oscar/file-logging.conf b/python/xorbits/deploy/oscar/file-logging.conf index d397aa37b..d78ff0249 100644 --- a/python/xorbits/deploy/oscar/file-logging.conf +++ b/python/xorbits/deploy/oscar/file-logging.conf @@ -8,7 +8,7 @@ keys=stream_handler,file_handler keys=formatter [logger_root] -level=WARN +level=DEBUG handlers=stream_handler,file_handler [logger_main] @@ -98,7 +98,7 @@ propagate=0 [handler_stream_handler] class=StreamHandler formatter=formatter -level=WARN +level=DEBUG args=(sys.stderr,) [handler_file_handler] From 27a833338c2eb486b8c7747042f5db27e9280e40 Mon Sep 17 00:00:00 2001 From: Weizheng Lu Date: Fri, 5 Jul 2024 22:49:21 +0800 Subject: [PATCH 22/25] remove debug log --- python/xorbits/_mars/services/storage/handler.py | 2 -- python/xorbits/deploy/oscar/file-logging.conf | 4 ++-- 2 files changed, 2 insertions(+), 4 deletions(-) diff --git a/python/xorbits/_mars/services/storage/handler.py b/python/xorbits/_mars/services/storage/handler.py index 7e3634caa..3986ac580 100644 --- a/python/xorbits/_mars/services/storage/handler.py +++ b/python/xorbits/_mars/services/storage/handler.py @@ -229,7 +229,6 @@ async def put( object_info.size - data_info.memory_size ) await self.notify_spillable_space(level) - print(f"StorageHandlerActor put") return data_info @put.batch @@ -556,7 +555,6 @@ async def fetch_via_transfer( ): from .transfer import ReceiverManagerActor, SenderManagerActor - print(f"fetch_via_transfer") logger.debug("Begin to fetch %s from band %s", data_keys, remote_band) remote_data_manager_ref: mo.ActorRefType[DataManagerActor] = await mo.actor_ref( diff --git a/python/xorbits/deploy/oscar/file-logging.conf b/python/xorbits/deploy/oscar/file-logging.conf index d78ff0249..d397aa37b 100644 --- a/python/xorbits/deploy/oscar/file-logging.conf +++ b/python/xorbits/deploy/oscar/file-logging.conf @@ -8,7 +8,7 @@ keys=stream_handler,file_handler keys=formatter [logger_root] -level=DEBUG +level=WARN handlers=stream_handler,file_handler [logger_main] @@ -98,7 +98,7 @@ propagate=0 [handler_stream_handler] class=StreamHandler formatter=formatter -level=DEBUG +level=WARN args=(sys.stderr,) [handler_file_handler] From c26c46d3bb4bc965316821805aa013edd368699c Mon Sep 17 00:00:00 2001 From: Weizheng Lu Date: Sun, 7 Jul 2024 07:55:00 +0800 Subject: [PATCH 23/25] delete logs --- python/xorbits/_mars/services/storage/core.py | 10 ---------- python/xorbits/_mars/services/storage/handler.py | 4 ---- 2 files changed, 14 deletions(-) diff --git a/python/xorbits/_mars/services/storage/core.py b/python/xorbits/_mars/services/storage/core.py index 45641e296..cb064ef5d 100644 --- a/python/xorbits/_mars/services/storage/core.py +++ b/python/xorbits/_mars/services/storage/core.py @@ -93,12 +93,6 @@ async def clean_up(self): self._file.close() async def close(self): - logger.debug( - "Writer closed for %s, %s on %s", - self._session_id, - self._data_key, - self._data_manager.address, - ) self._file.close() if self._object_id is None: # for some backends like vineyard, @@ -328,9 +322,6 @@ def delete_data_info( level: StorageLevel, band_name: str, ): - logger.debug( - "Deleting %s, %s from %s, %s", session_id, data_key, level, band_name - ) if (session_id, data_key) in self._data_key_to_infos: self._data_info_list[level, band_name].pop((session_id, data_key)) self._spill_strategy[level, band_name].record_delete_info( @@ -635,7 +626,6 @@ async def _setup_storage( storage_config = storage_config or dict() init_params, teardown_params = await backend.setup(**storage_config) client = backend(**init_params) - logger.debug(f"_setup_storage: {band_name}, {storage_backend}, {init_params}") self._init_params[band_name][storage_backend] = init_params self._teardown_params[band_name][storage_backend] = teardown_params return client diff --git a/python/xorbits/_mars/services/storage/handler.py b/python/xorbits/_mars/services/storage/handler.py index 3986ac580..ff9b45f68 100644 --- a/python/xorbits/_mars/services/storage/handler.py +++ b/python/xorbits/_mars/services/storage/handler.py @@ -285,7 +285,6 @@ async def delete_object( level: StorageLevel, ): data_key = await self._data_manager_ref.get_store_key(session_id, data_key) - logger.debug("Delete object %s, %s on %s", session_id, data_key, self.address) await self._data_manager_ref.delete_data_info( session_id, data_key, level, self._band_name ) @@ -370,9 +369,6 @@ async def batch_delete(self, args_list, kwargs_list): session_id, key, level, info.band ) ) - logger.debug( - "Batch delete %s, %s on %s", session_id, key, self.address - ) to_removes.append((level, info.object_id)) level_sizes[level] += info.store_size From 9ba9de30c5d3aa3e2986c79a9feeeb2662652769 Mon Sep 17 00:00:00 2001 From: Weizheng Lu Date: Sun, 7 Jul 2024 22:57:40 +0800 Subject: [PATCH 24/25] debug log --- python/xorbits/deploy/oscar/file-logging.conf | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/python/xorbits/deploy/oscar/file-logging.conf b/python/xorbits/deploy/oscar/file-logging.conf index d397aa37b..d78ff0249 100644 --- a/python/xorbits/deploy/oscar/file-logging.conf +++ b/python/xorbits/deploy/oscar/file-logging.conf @@ -8,7 +8,7 @@ keys=stream_handler,file_handler keys=formatter [logger_root] -level=WARN +level=DEBUG handlers=stream_handler,file_handler [logger_main] @@ -98,7 +98,7 @@ propagate=0 [handler_stream_handler] class=StreamHandler formatter=formatter -level=WARN +level=DEBUG args=(sys.stderr,) [handler_file_handler] From af2fe839a220c057bbc347e658c528689b777e04 Mon Sep 17 00:00:00 2001 From: Weizheng Lu Date: Fri, 23 Aug 2024 22:02:04 +0800 Subject: [PATCH 25/25] fix test --- .../_mars/dataframe/groupby/tests/test_groupby_execution.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/python/xorbits/_mars/dataframe/groupby/tests/test_groupby_execution.py b/python/xorbits/_mars/dataframe/groupby/tests/test_groupby_execution.py index b36823a12..feb5550fb 100644 --- a/python/xorbits/_mars/dataframe/groupby/tests/test_groupby_execution.py +++ b/python/xorbits/_mars/dataframe/groupby/tests/test_groupby_execution.py @@ -728,8 +728,7 @@ def _disallow_combine_and_agg(ctx, op): pd.testing.assert_frame_equal(result.sort_index(), raw.groupby("c1").agg("sum")) -@pytest.mark.skip_ray_dag # _fetch_infos() is not supported by ray backend. -def test_distributed_groupby_agg(setup_cluster): +def test_distributed_groupby_agg(setup): rs = np.random.RandomState(0) raw = pd.DataFrame(rs.rand(50000, 10)) df = md.DataFrame(raw, chunk_size=raw.shape[0] // 2)