Numba JIT模式下嵌套yield生成器配合heapq内存泄漏问题

核心根因 #

该内存泄漏和LLVM存储编译产物的机制无关，属于Numba nopython模式下的已知引用计数缺陷，触发逻辑如下：

• 当@njit装饰的函数内同时存在heapq堆结构操作、多层嵌套JIT生成器迭代逻辑时，Numba的原生内存引用计数机制会失效。
• Numba在nopython模式下实现生成器时，会为每个yield节点保存独立的栈帧上下文副本。该复现场景存在两层嵌套生成器：main_迭代generator_的输出，generator_又会循环创建generator_2实例并迭代，当生成器用空列表作为终止标记退出时，内层generator_2实例、yield返回的临时列表对象的引用计数不会被正常递减。这部分内存是Numba通过LLVM申请的原生内存，不在Python GC的管辖范围内，即使主动调用Python层面的GC也无法回收。
• 移除main_函数的@njit装饰器后泄漏消失的原因：此时main_运行在Python解释器环境，调用JIT生成器返回的包装对象可以被Python解释器正常追踪生命周期，迭代结束后关联资源会被正常释放，不会出现原生内存泄漏。

额外注意：复现代码中heapq.heappush(result)写法存在参数缺失，正确调用方式应为heapq.heappush(ss, result)，该语法错误会导致Numba编译时生成异常的栈帧保留逻辑，进一步放大内存泄漏的规模。

临时规避方案 #

• 不要在nopython模式下组合使用多层嵌套生成器与heapq操作：可以把堆处理逻辑移到JIT函数外，JIT生成器仅负责返回计算结果，由Python层完成堆的push、pop等操作。
• 替换生成器终止标记：不要用空列表作为迭代终止哨兵，改用固定整数（如-1）标记迭代结束，减少临时列表对象的创建，同时在生成器退出前手动清空持有的内层生成器引用。
• 全链路JIT场景下改写实现逻辑：如果要求所有逻辑都运行在nopython模式，将多层yield的生成器逻辑改写为预分配数组批量返回的形式，从根源上避开Numba生成器生命周期管理的缺陷。

内容的提问来源于stack exchange，提问作者Axd