已关闭
[Bug]: IndexedDataset readers leak mmap resources and may return invalid or corrupted data #376
rongyue创建于  27 天前关闭于  21 天前
rongyue
rongyue
27 天前 创建

Checklist

🐛 Describe the bug

IndexedDataset 的本地 indexed data reader 存在多处资源生命周期和边界读取问题,可能导致 mmap 资源泄漏、返回数组引用已释放 mmap、短读时返回未初始化数据,以及大 slice 场景下 token count 溢出。

涉及文件:
hyper_parallel/data/indexed/io.py

问题包括:

  1. _IndexReader.del 强制关闭 mmap,但实例仍持有由 numpy.frombuffer(self._buffer, ...) 创建的 sequence_lengths、sequence_pointers、document_indices、sequence_modes 等 exported buffer view。析构时 mmap.close() 可能抛 BufferError,导致 mmap 未正常释放。IndexedDatasetBuilder.add_index 合并多个 index 分片时会反复创建并释放 _IndexReader,容易触发该路径。

  2. _IndexReader.getitem 使用 @lru_cache(maxsize=8) 装饰实例方法。该 cache 挂在类级函数对象上,key 包含 self,因此最近 8 个 _IndexReader 实例会被强引用保留,连带整份 .idx mmap 常驻内存。多源 blend 或重复构建 Dataset 时会累积占用。

  3. _MMapBinReader.read() 返回 numpy.frombuffer(self._buffer, ...) 的零拷贝 view,但 _MMapBinReader.del 会强制关闭底层 mmap。如果上层仍持有 IndexedDataset.get()/getitem 返回的数组,而 reader 被释放,返回数组可能引用已 munmap 的地址,存在 SIGSEGV 风险。

  4. _FileBinReader.read() 使用 numpy.empty(count, dtype=dtype) 预分配输出,然后 f.readinto(out),但没有检查 readinto 返回值。若 .bin 被截断、idx/bin 不匹配,或请求范围越过 .bin 文件末尾,readinto 只会填充部分字节,数组尾部保留未初始化内存,并被静默当作 token id 返回,造成数据损坏。

  5. IndexedDataset.getitem 的 slice 分支对 int32 sequence_lengths 使用 accumulate/sum 累加。对于覆盖超过 2^31 token 的连续 slice,numpy int32 累加可能溢出,导致传给 bin_reader.read() 的 count 错误,进一步造成截断读取、错位 split,甚至异常的大范围读取。

Expected behavior

IndexedDataset reader 应满足以下行为:

  1. _IndexReader 释放资源时不应因为仍持有 frombuffer view 而触发 BufferError,也不应泄漏 mmap。
  2. _IndexReader.getitem 不应通过类级 lru_cache 强引用 reader 实例,reader 在无外部引用后应可被 GC 回收。
  3. _MMapBinReader 返回的数组不应在 reader 析构后变成悬挂引用;同时应尽量保留 mmap 零拷贝读取性能。
  4. _FileBinReader.read() 应检查实际读取字节数。发生短读时应抛出明确异常,而不是返回包含未初始化内存的数组。
  5. IndexedDataset slice 分支应使用 int64/Python int 进行 token count 和 offsets 累加,避免 int32 溢出。

Additional context

建议修复方向:

  • _IndexReader 中从 mmap 解析出来的元数据数组可以 copy 到普通 numpy array,使 sequence_lengths、sequence_pointers、document_indices、sequence_modes 不再依赖 mmap 生命周期。
  • 为 _IndexReader 增加幂等 close(),先释放 memoryview,再关闭 mmap;del 调用 close() 并兜底析构异常。
  • 移除 _IndexReader.getitem 上的 @lru_cache,避免类级 cache 强引用 self。
  • _MMapBinReader.read() 不建议默认 .copy(),否则 mmap 热路径会退化为每次读取都额外拷贝。更轻量的方式是避免在 del/close() 中强制关闭底层 mmap,只关闭 backing file,让已返回的 numpy view 通过自身 base 链保持 mmap 生命周期。
  • _FileBinReader.read() 应对 out.view(numpy.uint8) 调用 readinto,并校验返回 nbytes 是否等于 out.nbytes。
  • slice 分支使用 numpy.cumsum(lengths, dtype=numpy.int64),token_count 使用 int(offsets[-1])。

Environment info

Repository: mindspore/hyper-parallel
Component: hyper_parallel/data/indexed/io.py
Affected classes:

  • _IndexReader
  • _MMapBinReader
  • _FileBinReader
  • IndexedDataset

Observed on current development branch around indexed dataset offline preprocessing changes.

Thanks for contributing 🎉!

likedislike
songjiaqisongjiaqi成员
26 天前 关联了pull request:test_mem_warning
songjiaqisongjiaqi成员
26 天前 删除了关联的pull request:test_mem_warning
songjiaqisongjiaqi成员
26 天前 关联了pull request:test_mem_warning
rongyuerongyue
21 天前 issue状态由 TODO 改变为 CLOSED
rongyuerongyue
21 天前 关闭了 issue