[bugfix] Internal processing of hf3fs crash # 16614 (#16938)

This commit is contained in:
FlyPanda
2026-01-27 10:01:50 +08:00
committed by GitHub
parent 539924037f
commit 2d8c22a15e
2 changed files with 55 additions and 22 deletions

View File

@@ -118,47 +118,67 @@ class Hf3fsUsrBioClient(Hf3fsClient):
@rsynchronized()
def batch_read(self, offsets: List[int], tensors: List[torch.Tensor]) -> List[int]:
self.check(offsets, tensors)
results = [0] * len(offsets)
# prepare
current = 0
for offset, tensor in zip(offsets, tensors):
size = tensor.numel() * tensor.itemsize
self.ior_r.prepare(
self.iov_r[current : current + size], True, self.file, offset
)
current += size
try:
self.ior_r.prepare(
self.iov_r[current : current + size], True, self.file, offset
)
current += size
except Exception as e:
logger.error(f"Error preparing batch read: {e}")
return results
# submit
ionum = len(offsets)
resv = self.ior_r.submit().wait(
min_results=ionum, timeout=datetime.timedelta(seconds=self.client_timeout)
)
try:
resv = self.ior_r.submit().wait(
min_results=ionum,
timeout=datetime.timedelta(seconds=self.client_timeout),
)
except Exception as e:
logger.error(f"Error submitting batch read: {e}")
return results
# results
hf3fs_utils.read_shm(self.shm_r_tensor, tensors)
results = [res.result for res in resv]
try:
hf3fs_utils.read_shm(self.shm_r_tensor, tensors)
results = [res.result for res in resv]
except Exception as e:
logger.error(f"[Hf3fsUsrBioClient] read_shm failed: {e}", exc_info=True)
return results
return results
@wsynchronized()
def batch_write(self, offsets: List[int], tensors: List[torch.Tensor]) -> List[int]:
self.check(offsets, tensors)
results = [0] * len(offsets)
# prepare
hf3fs_utils.write_shm(tensors, self.shm_w_tensor)
current = 0
for offset, tensor in zip(offsets, tensors):
size = tensor.numel() * tensor.itemsize
self.ior_w.prepare(
self.iov_w[current : current + size], False, self.file, offset
)
current += size
try:
self.ior_w.prepare(
self.iov_w[current : current + size], False, self.file, offset
)
current += size
except Exception as e:
logger.error(f"Error preparing batch write: {e}")
return results
# submit
ionum = len(offsets)
resv = self.ior_w.submit().wait(
min_results=ionum, timeout=datetime.timedelta(seconds=self.client_timeout)
)
try:
resv = self.ior_w.submit().wait(
min_results=ionum,
timeout=datetime.timedelta(seconds=self.client_timeout),
)
except Exception as e:
logger.error(f"Error submitting batch write: {e}")
return results
# results
results = [res.result for res in resv]

View File

@@ -347,7 +347,11 @@ class HiCacheHF3FS(HiCacheStorage):
values: List[torch.Tensor],
) -> List[bool]:
page_indices = self.metadata_client.get_page_indices(self.rank, keys)
if len(page_indices) != len(keys):
logger.error(
f"[Rank {self.rank}] HiCacheHF3FS get: page_indices length {len(page_indices)} mismatch keys length {len(keys)}."
)
return [False] * len(keys)
batch_indices, file_offsets = [], []
for i, page_index in enumerate(page_indices):
if page_index is not None:
@@ -402,7 +406,16 @@ class HiCacheHF3FS(HiCacheStorage):
indices = self.metadata_client.reserve_and_allocate_page_indices(
self.rank, key_with_prefix
)
if len(indices) != len(keys):
logger.error(
f"[Rank {self.rank}] HiCacheHF3FS batch_get: mismatched lengths {len(indices)} != {len(keys)}"
)
# free allocated pages
if indices:
self.metadata_client.confirm_write(
self.rank, [], [index[1] for index in indices]
)
return [False] * len(keys)
batch_indices, file_offsets, file_values = [], [], []
pages_to_release = []