ENH: Fix cuda storage transfer deadlock on multiple GPUs

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)

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):
@@ -292,6 +293,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")
 
@@ -382,6 +384,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
@@ -390,6 +393,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):
@@ -522,7 +526,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]],
@@ -531,21 +549,136 @@ 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)
- sender_ref: mo.ActorRefType[SenderManagerActor] = await mo.actor_ref(
- address=remote_band[0], uid=SenderManagerActor.gen_uid(remote_band[1])
+
+ remote_data_manager_ref: mo.ActorRefType[DataManagerActor] = await mo.actor_ref(
+ address=remote_band[0], uid=DataManagerActor.default_uid()
  )
- await sender_ref.send_batch_data(
- session_id,
+
+ 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_info_tasks = []
+ pin_tasks = []
+ for data_key in data_keys:
+ get_info_tasks.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("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)
+
+ 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
+
+ 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,
+ 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
+ )
+
+ 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",
  data_keys,
- self._data_manager_ref.address,
- level,
- fetch_band_name,
- error=error,
+ remote_band,
+ (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,
@@ -559,10 +692,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 +717,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,16 +733,14 @@ 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)
  assert len(metas) == len(missing_keys)
  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():
@@ -620,7 +752,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
  )
  )