mindspore-ai / mindspore-ai/hyper-parallel

[Bug]: IndexedDataset readers leak mmap resources and may return invalid or corrupted data

Open
#156 0 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

Dominant language
Python
Stars
53
Forks
63
Avg merge
23h 45m
Merged PRs (30d)
63

Description

Checklist
  • 1. I have searched the existing issues (https://gitcode.com/mindspore/hyper-parallel/issues)
  • 2. I have read the relevant documentation.
  • 3. I have created a minimal reproduction case that clearly demonstrates the issue, including a complete code example and the error message with full traceback and error logs.
🐛 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 🎉!

schema_version: 1
source: gitcode
gitcode_repo: mindspore/hyper-parallel
gitcode_issue: 376
source_url: https://gitcode.com/mindspore/hyper-parallel/issues/376

Contributor guide

No contributing guide indexed for this repository

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start with hyper_parallel/data/indexed/io.py and trace _IndexReader, _MMapBinReader, _FileBinReader, and IndexedDataset.getitem. Inspect how mmap-backed arrays, cached readers, readinto results, and slice offsets are owned and released. Done means resources can be collected safely, short reads fail explicitly, returned arrays remain valid, and large slices avoid integer overflow.

Written by the indexing model from the issue text.

Assessment

Tech stack
numpy, python
Domain
data
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Active
Clarity
Mostly clear
Newbie friendliness
48/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.