如何Python多线程数据处理时的内存泄漏

项目场景：

在一个数据分析项目中，需要并行处理大量的CSV文件（每个文件约500MB），将处理结果保存到数据库。由于数据量较大，采用了多线程处理方案，每个线程负责处理一个文件。项目使用Python 3.8，主要依赖pandas进行数据处理。

问题描述

在程序运行一段时间后，发现内存占用持续增长，最终导致系统内存溢出（OOM）。即使处理完一批数据后，内存占用也没有明显下降。通过监控发现，每处理一个文件，内存占用大约增加50-100MB，且不会被释放。

问题代码示例：

import pandas as pd
from threading import Thread
from queue import Queue

class DataProcessor:
    def __init__(self):
        self.result_queue = Queue()
        self.workers = []

    def process_file(self, file_path):
        df = pd.read_csv(file_path)
        # 进行数据处理
        processed_data = self.complex_process(df)
        self.result_queue.put(processed_data)

    def complex_process(self, df):
        # 复杂的数据处理逻辑
        return df.groupby('category').agg({'value': 'sum'})

    def start_processing(self, file_list):
        for file in file_list:
            worker = Thread(target=self.process_file, args=(file,))
            self.workers.append(worker)
            worker.start()

原因分析：

1. 循环引用问题：

   • Thread对象持有对DataProcessor实例的引用（通过self）
   • DataProcessor通过self.workers列表持有Thread对象的引用
   • 这种循环引用导致Python的垃圾回收器无法正确释放内存

2. pandas DataFrame对象的内存管理：

   • pandas的DataFrame对象在进行操作时会创建数据副本
   • 如果不正确管理这些副本，会导致内存占用持续增长

3. Queue对象积累：

   • result_queue中的数据如果不及时处理，也会占用内存
   • Queue对象本身可能存在内存泄漏

解决方案：

1. 使用弱引用存储工作线程：

import weakref
from threading import Thread
from queue import Queue
import pandas as pd

class DataProcessor:
    def __init__(self):
        self.result_queue = Queue()
        self.workers = weakref.WeakSet()  # 使用WeakSet代替列表

    def process_file(self, file_path):
        try:
            df = pd.read_csv(file_path)
            processed_data = self.complex_process(df)
            self.result_queue.put(processed_data)
        finally:
            # 确保处理完后清理DataFrame
            del df
            
    def complex_process(self, df):
        # 使用copy=False避免不必要的数据复制
        result = df.groupby('category', copy=False).agg({'value': 'sum'})
        return result

    def start_processing(self, file_list):
        for file in file_list:
            worker = Thread(target=self.process_file, args=(file,))
            self.workers.add(worker)  # 使用add方法添加到WeakSet
            worker.start()
            
    def cleanup(self):
        # 清理资源
        while not self.result_queue.empty():
            _ = self.result_queue.get()
        self.result_queue = None

2. 关键改进点：

• 使用weakref.WeakSet()替代普通列表存储线程对象
• 在process_file中添加try-finally确保DataFrame被正确释放
• 在数据处理时避免不必要的数据复制
• 添加cleanup方法用于资源清理
• 使用del语句显式删除不再需要的大对象

3. 使用建议：

# 使用示例
processor = DataProcessor()
try:
    processor.start_processing(['file1.csv', 'file2.csv'])
    # 等待处理完成
    for worker in processor.workers:
        worker.join()
finally:
    # 确保资源被清理
    processor.cleanup()

通过以上优化，程序的内存占用得到了有效控制，不再出现内存持续增长的问题。