ANN-Benchmarks 内存限制问题分析与解决方案

问题背景

在 ANN-Benchmarks 项目中，当用户使用并行模式运行算法基准测试时，发现存在内存分配不均的问题。具体表现为：在 AWS r6i.16xlarge 机器（512GB 内存）上运行 31 个并行任务时，预期每个算法应获得约 16GB 内存，但实际分配的内存却明显不足（约 11GB），导致部分算法因内存不足而失败（错误代码 137）。

问题根源分析

经过代码审查，发现问题出在 ann_benchmarks/main.py 文件中的内存限制计算逻辑。当前实现使用 psutil.virtual_memory().available 来计算可用内存，这会导致以下问题：

1. 动态内存计算问题：available 返回的是当前可用内存，而非总内存。当第一批容器启动后，系统内存被占用，后续容器启动时可用内存减少，导致分配不均。

2. 公平性问题：不同批次的算法获得的内存资源不同，影响基准测试的公平性和可比性。

技术解决方案

原实现分析

mem_limit = int((psutil.virtual_memory().available - memory_margin) / args.parallelism)

这种实现方式在共享机器环境下是合理的，因为它基于实际可用内存分配资源。但对于专用基准测试机器，这种动态分配方式会导致测试结果不可比。

改进方案

建议修改为基于总内存的计算方式：

mem_limit = int((psutil.virtual_memory().total - memory_margin) / args.parallelism)

这种修改确保：

1. 所有算法获得相同的内存配额
2. 在专用测试机器上提供更一致的测试环境
3. 避免因内存分配不均导致的算法失败

实现建议

更完善的解决方案是将内存限制计算移到 run_worker 函数中，并在程序启动时确定固定的内存配额。这样可以：

1. 保持内存限制的一致性
2. 便于调试和问题追踪
3. 提供更可预测的资源分配

项目实践意义

对于 ANN 算法基准测试，内存分配的一致性至关重要，因为：

1. 许多近似最近邻算法对内存敏感
2. 不公平的内存分配会导致算法性能评估失真
3. 内存不足可能导致算法无法完成或产生次优结果

这一改进将提升 ANN-Benchmarks 项目的测试可靠性和结果可比性，特别是在大规模数据集和高并行度场景下。