Python多进程数据标注脚本无法终止且耗时过长，求问题排查

问题描述 #

尝试用multiprocessing加速20万行DataFrame的文本标注工作（匹配3000个关键词的位置），原单进程脚本约7小时完成，但并行化脚本启动后一直运行不终止，耗时远超预期。

任务概述 #

输入：

• 含text列的Pandas DataFrame
• 格式为{word: label}的关键词字典

期望输出：
为每行文本添加匹配关键词的起始/结束位置标记，示例如下：

示例输入DataFrame：

示例关键词字典：{'United States': 'country'}

示例输出DataFrame：

核心逻辑：遍历DataFrame和字典，用正则表达式匹配关键词并记录位置。

现有代码问题 #

原单进程代码 #

def label_data(df, label_dict):
    pbar = tqdm(iterable=df.to_dict('records'), total=df.shape[0])
    for idx, row in enumerate(pbar): 
        text = row['text'] 
        spans = []    
        for word, label in label_dict.items():
            for match in re.finditer(word, text):
                start = match.start()
                end = match.end()
                spans.append([start, end])

        row['labels'] = spans
        df.iloc[idx] = row

    return df

并行化代码 #

from itertools import product
import multiprocessing as mp
import numpy as np

def label_data_parallel(df, label_dict):
    num_cores = mp.cpu_count() - 1 
    pool = mp.Pool(num_cores)

    df_chunks = np.array_split(df, num_cores)
    labeled_dfs = pool.starmap(label_data, \
                               product(df_chunks, [label_dict] * num_cores))
    df = pd.concat(labeled_dfs, axis=0)

    pool.close()
    pool.join()

    return df

问题根源 #

1. 数据拷贝开销过大：product(df_chunks, [label_dict] * num_cores)会把3000条目的字典复制多份，加上每个DataFrame块的序列化/反序列化，跨进程传输数据的开销远超并行计算的收益。
2. 原函数效率低下：单进程版本中，df.iloc[idx] = row是原地修改DataFrame，性能极差；同时遍历3000个关键词+正则匹配的逻辑没有优化，本身就有很大性能瓶颈。
3. 进程池使用不当：starmap结合product的方式会重复传递字典，且没有考虑进程间数据共享，反而增加了内存负担。

优化方案 #

第一步：优化单进程核心逻辑 #

先把单进程代码的性能拉满，再考虑并行：

import re
import pandas as pd
from tqdm import tqdm

def build_pattern(label_dict):
    # 把关键词转成正则表达式，用|分隔，优先匹配长词避免短词截断长词
    sorted_words = sorted(label_dict.keys(), key=lambda x: -len(x))
    pattern = re.compile('|'.join(re.escape(word) for word in sorted_words))
    return pattern

def process_text(text, pattern):
    spans = []
    for match in pattern.finditer(text):
        spans.append([match.start(), match.end()])
    return spans

def label_data_optimized(df, label_dict):
    pattern = build_pattern(label_dict)
    tqdm.pandas()
    df['labels'] = df['text'].progress_apply(lambda x: process_text(x, pattern))
    return df

• 把3000个关键词合并成一个正则表达式，减少循环次数
• 使用progress_apply替代手动遍历+iloc修改，性能提升显著
• 对关键词按长度降序排序，避免短词匹配覆盖长词（比如先匹配"United States"再匹配"United"）

第二步：高效并行化 #

用原生multiprocessing优化，避免冗余数据传递：

import multiprocessing as mp
import pandas as pd
import numpy as np

def label_data_parallel_optimized(df, label_dict):
    num_cores = mp.cpu_count() - 1
    pattern = build_pattern(label_dict)
    
    # 分割DataFrame
    df_chunks = np.array_split(df, num_cores)
    
    # 定义每个进程的处理函数
    def process_chunk(chunk):
        chunk['labels'] = chunk['text'].apply(lambda x: process_text(x, pattern))
        return chunk
    
    with mp.Pool(num_cores) as pool:
        labeled_chunks = pool.map(process_chunk, df_chunks)
    
    return pd.concat(labeled_chunks, axis=0)

• 提前构建正则表达式，避免每个进程重复构建
• 用map替代starmap+product，减少数据传递开销
• 使用with语句自动管理进程池，避免资源泄漏
• 每个进程只处理自己的DataFrame块，无跨进程修改冲突

额外建议 #

• 如果内存足够，把DataFrame的text列转成字符串类型，避免类型转换开销
• 正则表达式中可添加flags=re.IGNORECASE（如果不需要大小写敏感），或根据需求调整
• 测试时先用小样本（比如1000行）验证逻辑和性能，再跑全量数据

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