变量规模无明显变化但内存占用持续攀升的原因排查

在处理数十GB级的Parquet文件语料时，运行下方代码出现内存占用持续攀升的情况：10分钟时每核占用3%，30分钟时达7%，6小时后超过10%。尽管已显式调用gc.collect()，且所有文件大小相近、每次仅加载一个文件，也未发现旧文件的持久引用，但内存仍不断增长。这一问题迫使我限制数据分词的核心数，导致处理时间直接翻倍。

原代码如下：

from nltk.tokenize import RegexpTokenizer
import pandas as pd
import re
import string
import os
from multiprocessing import Pool
import gc
import glob

def adjust_tokens(tokens):
    new_tokens = []
    number_regex = '[0-9]+(?:[.][0-9]+)?'
    for token in tokens:
        if token in string.punctuation:
            continue
        if token == '...':
            continue
        if re.match(number_regex, token):
            new_tokens.append('aquantity')
        else:
            new_tokens.append(token.lower())
    return new_tokens


class ParquetsIterable():
    def __init__(self, tokenizer=None, indir=None, infiles=None,
                 filters=None, keys=None):
        if tokenizer is not None:
            self.tokenizer = tokenizer
        else:
            self.tokenizer = RegexpTokenizer(
                r'(?:\w+)(?:[.-]\w+)*[-+]?|(?:[.][.][.])'
                + r'|(?:[!"#$%&\'()*+,-./:;<=>?@[\]^_`{|}~])'
                                            )
        if infiles is not None:
            self.infiles = infiles
        else:
            self.infiles = [os.path.join(indir, f) for f in os.listdir(indir)]
        self.filters = filters
        self.keys = keys

    def __iter__(self):
        for filename in self.infiles:
            try:
                df = pd.read_parquet(filename)
            except Exception:
                print('Invalid parquet file %s\n' % filename)
                continue
            if self.keys is not None:
                try:
                    assert (self.keys[filename].values
                            == df['key'].values).all()
                except AssertionError:
                    print('Skipping %s because of reorderinging'
                          ' and my unwillingness to write in a join'
                          % filename)
                    continue
            if self.filters is not None:
                df = df[self.filters[filename].values]
            gc.collect()
            for text in df['sentences']:
                for sentence in text.split('\n'):
                    yield adjust_tokens(self.tokenizer.tokenize(sentence))


def write_iterables_to_files(iterables, files, n_cores=8):
    pairs = zip(iterables, files)
    with Pool(n_cores) as p:
        p.map(write_iterable_to_file, pairs)


def write_iterable_to_file(iterable_and_file):
    iterable, file_ = iterable_and_file
    with open(file_, 'w') as f:
        for tokens in iterable:
            f.write(' '.join(tokens) + '\n')


if __name__ == '__main__':
    files = sorted([f for f in glob.glob('data/sentence-parquets/*')
                    if 'parts' in f and 'sampled' not in f])
    os.makedirs('data/processed-sentences-text', exist_ok=True)
    iterables = [ParquetsIterable(infiles=files[i::8]) for i in range(8)]
    outfiles = ['data/processed-sentences-text/sentences-%d.txt' % i
                for i in range(8)]
    write_iterables_to_files(iterables, outfiles, n_cores=8)

• DataFrame长期持有：__iter__方法中，df对象会被持有到该文件所有文本迭代完成，大文件的DataFrame内存无法及时释放。
• 第三方库内存盲区：Pandas读取Parquet依赖的pyarrow/fastparquet有独立内存池，Python的gc.collect()无法回收这部分残留内存。
• 多进程迭代器开销：父进程传递的迭代器可能携带未释放的内存引用，子进程继承父进程冗余内存后持续积累。
• 重复正则编译：adjust_tokens中每次调用都重新编译正则表达式，产生大量临时对象占用内存。

1. 提前释放DataFrame内存 #

修改ParquetsIterable.__iter__方法，在迭代文本前删除df并强制回收，同时只读取需要的列减少内存占用：

def __iter__(self):
    for filename in self.infiles:
        try:
            # 仅读取需要的sentences列
            df = pd.read_parquet(filename, columns=['sentences'])
        except Exception:
            print('Invalid parquet file %s\n' % filename)
            continue
        # 保留原有的filter和key校验逻辑（按需调整）
        if self.keys is not None:
            try:
                # 若需要校验key，需同时读取key列
                df = pd.read_parquet(filename, columns=['sentences', 'key'])
                assert (self.keys[filename].values == df['key'].values).all()
            except AssertionError:
                print('Skipping %s because of reorderinging'
                      ' and my unwillingness to write in a join'
                      % filename)
                continue
        if self.filters is not None:
            df = df[self.filters[filename].values]
        
        # 提取sentences为列表，解除对df的依赖
        sentences_list = df['sentences'].tolist()
        # 显式删除df并强制回收
        del df
        gc.collect()
        
        # 清理pyarrow内存池（若使用pyarrow引擎）
        try:
            import pyarrow as pa
            pa.default_memory_pool().release_unused()
        except ImportError:
            pass
        
        for text in sentences_list:
            for sentence in text.split('\n'):
                yield adjust_tokens(self.tokenizer.tokenize(sentence))

2. 优化多进程架构 #

避免父进程传递大迭代器，改为子进程独立处理文件列表，减少内存复制：

def process_file_group(file_group, outfile):
    tokenizer = RegexpTokenizer(
        r'(?:\w+)(?:[.-]\w+)*[-+]?|(?:[.][.][.])'
        + r'|(?:[!"#$%&\'()*+,-./:;<=>?@[\]^_`{|}~])'
    )
    with open(outfile, 'w') as f:
        for filename in file_group:
            try:
                df = pd.read_parquet(filename, columns=['sentences'])
            except Exception:
                print('Invalid parquet file %s\n' % filename)
                continue
            sentences_list = df['sentences'].tolist()
            del df
            gc.collect()
            
            # 清理pyarrow内存（可选）
            try:
                import pyarrow as pa
                pa.default_memory_pool().release_unused()
            except ImportError:
                pass
            
            for text in sentences_list:
                for sentence in text.split('\n'):
                    tokens = adjust_tokens(tokenizer.tokenize(sentence))
                    f.write(' '.join(tokens) + '\n')

def write_iterables_to_files(file_groups, outfiles, n_cores=8):
    with Pool(n_cores) as p:
        p.starmap(process_file_group, zip(file_groups, outfiles))

if __name__ == '__main__':
    files = sorted([f for f in glob.glob('data/sentence-parquets/*')
                    if 'parts' in f and 'sampled' not in f])
    os.makedirs('data/processed-sentences-text', exist_ok=True)
    # 拆分文件组给每个进程
    file_groups = [files[i::8] for i in range(8)]
    outfiles = ['data/processed-sentences-text/sentences-%d.txt' % i
                for i in range(8)]
    write_iterables_to_files(file_groups, outfiles, n_cores=8)

3. 预编译正则表达式 #

将正则表达式提前编译为全局变量，避免重复创建对象：

# 全局预编译正则表达式
number_regex = re.compile('[0-9]+(?:[.][0-9]+)?')

def adjust_tokens(tokens):
    new_tokens = []
    for token in tokens:
        if token in string.punctuation:
            continue
        if token == '...':
            continue
        if number_regex.match(token):
            new_tokens.append('aquantity')
        else:
            new_tokens.append(token.lower())
    return new_tokens

4. 可选：替换高效分词工具 #

若NLTK分词器内存占用过高，可替换为更轻量的tokenizers库：

from tokenizers import Tokenizer
from tokenizers.models import WordLevel
from tokenizers.pre_tokenizers import Whitespace

# 初始化轻量分词器
tokenizer = Tokenizer(WordLevel())
tokenizer.pre_tokenizer = Whitespace()

# 调整分词逻辑
def adjust_tokens(tokens):
    new_tokens = []
    for token in tokens:
        if token in string.punctuation:
            continue
        if token == '...':
            continue
        if number_regex.match(token):
            new_tokens.append('aquantity')
        else:
            new_tokens.append(token.lower())
    return new_tokens

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