Numba jitclass循环中内存累积问题求助

问题重现 #

在Numba jitclass实现的类中，循环内往numba.types.List类型的属性追加float64 numpy数组时，即使每次迭代前调用clear()清空列表，内存占用仍会随循环次数线性增长。相同逻辑用普通Python类实现则无此问题。

测试环境：

• OS: Linux SMP x86_64 22.04.1-Ubuntu
• Numba: 0.59.1 / 0.58.1
• Numpy: 1.26.4 / 1.24.1
• Python: 3.11.8 (GCC 12.3.0) / 3.11.0 (GCC 11.3.0)

复现代码见用户提供的示例。

根本原因 #

1. Numba List与Python原生List的内存管理差异：Numba的typed.List是为JIT编译优化实现的专用容器，并非Python原生list的直接封装。当存储numpy数组时，clear()方法仅移除列表中的元素引用，但Numba的内存回收机制可能未正确触发numpy数组底层内存的释放——不同于Python原生list会通过引用计数自动回收对象，Numba JIT代码的内存管理逻辑对numpy数组这类外部对象的处理存在疏漏。

2. jitclass上下文的内存绑定：jitclass的属性内存由Numba的JIT runtime管理，而非Python的GC。当列表中的numpy数组被clear()后，其底层内存可能仍被jitclass的内部结构隐式持有，无法被及时回收，导致内存累积。

3. 旧版本Numba的已知Bug：用户使用的0.58.x/0.59.x版本存在部分容器内存泄漏的问题，后续版本（如0.60+）已针对此类场景进行修复。

解决方案 #

方案1：改用列表的列表存储基础类型 #

如用户已尝试的方式，将存储numpy数组改为存储float64类型的列表（numba.types.List(numba.types.List(numba.types.float64))）。Numba对原生数值类型的内存管理更完善，clear()能正确回收内存。

方案2：替换列表而非调用clear() #

每次迭代时直接创建新的Numba typed List，替换原有属性，让旧列表连同内部数组被Python GC回收：

import numba.experimental.jitclass as jitclass
import numba
import random
import numpy as np
import psutil
import os

spec = [
    ('points', numba.types.List(numba.types.Array(numba.types.float64, 1, "C")))
]

@jitclass(spec)
class testing:
    def __init__(self):
        self.points = numba.typed.List()
    
    def compute(self):
        # 直接替换为新列表，而非clear
        self.points = numba.typed.List()
        l = random.randint(0, 1000)
        for k in range(l):
            self.points.append(np.array([np.float64(k)]))
        return 0

a = testing()

process = psutil.Process(os.getpid())
base_memory_usage = process.memory_info().rss
for i in range(1000):
    a.compute()

memory_usage = process.memory_info().rss
loop_memory_usage = memory_usage - base_memory_usage

print('RAM Used (GB):', loop_memory_usage/1e9)

方案3：预分配固定大小数组（若场景允许） #

如果能预估最大数据量，预分配一个大numpy数组，通过索引更新元素，避免动态追加带来的内存分配问题：

import numba.experimental.jitclass as jitclass
import numba
import random
import numpy as np
import psutil
import os

spec = [
    ('points', numba.types.Array(numba.types.float64, 2, "C")),
    ('length', numba.types.int64)
]

@jitclass(spec)
class testing:
    def __init__(self):
        # 预分配最大1000个元素的数组
        self.points = np.zeros((1000, 1), dtype=np.float64)
        self.length = 0
    
    def compute(self):
        self.length = random.randint(0, 1000)
        for k in range(self.length):
            self.points[k] = np.float64(k)
        return 0

a = testing()

process = psutil.Process(os.getpid())
base_memory_usage = process.memory_info().rss
for i in range(1000):
    a.compute()

memory_usage = process.memory_info().rss
loop_memory_usage = memory_usage - base_memory_usage

print('RAM Used (GB):', loop_memory_usage/1e9)

方案4：升级Numba版本 #

升级至0.60及以上版本，此类内存泄漏问题已在新版本中得到修复。

验证 #

上述方案均可有效解决内存累积问题，运行修改后的代码，内存占用将不再随循环次数增长。

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