处理16万+运动日志时BrokenProcessPool错误及Dask并行改写咨询

问题背景 #

处理包含167373条运动日志的数据集，目标是将每条日志的坐标数据合并为向量，通过计算余弦相似度识别跑步或骑行的常见路线。使用Dask实现并行处理时，触发错误：BrokenProcessPool: A process in the process pool was terminated abruptly while the future was running or pending.

数据格式示例 #

( ('bike', 12, [1.234, 1.268, 40.69, 41.8, -1.2, -1.3]), 
  ('bike', 12, [6.23, 6.268, 9.69, 9.8, -1.2, -1.3] )

原始代码 #

def calculate_pairwise_cosine_similarities(group, threshold=0.95):
    sport, records = group
    vectors = [record[2] for record in records]
    log_ids = [record[1] for record in records]

    pairs = [(i, j) for i in range(len(vectors)) for j in range(i + 1, len(vectors))]

    similarities = list(filter(lambda x: x[2] > threshold,                              
                               map(lambda pair: (log_ids[pair[0]], log_ids[pair[1]], cosine_similarity(vectors[pair[0]], vectors[pair[1]])), pairs)))

    return sport, similarities

错误原因分析 #

• 内存过载：单个分组内记录数过多时，嵌套循环生成所有两两索引对pairs会产生海量数据，导致单进程内存耗尽被系统强制终止。
• 计算效率低下：手动生成索引对+串行map/filter的方式计算效率极低，单进程处理时间过长，触发Dask进程超时或资源耗尽。
• 序列化问题：若cosine_similarity使用第三方库（如sklearn）的未序列化函数，可能导致进程间通信失败，引发进程崩溃。

基于Dask的改写方案 #

优化思路 #

1. 用批量向量运算替代手动生成索引对，减少内存占用和计算时间。
2. 利用Dask原生的groupby+map并行能力，避免手动管理进程池。
3. 加入边界判断（分组记录数<2时直接返回空），减少无效计算。

改写后的代码 #

import dask.bag as db
from sklearn.metrics.pairwise import cosine_similarity
import numpy as np

def calculate_pairwise_cosine_similarities(group, threshold=0.95):
    sport, records = group
    # 转换为numpy数组，支持批量运算
    vectors = np.array([record[2] for record in records])
    log_ids = np.array([record[1] for record in records])
    
    # 分组内记录数不足2时直接返回空结果
    if len(vectors) < 2:
        return sport, []
    
    # 批量计算余弦相似度矩阵，效率远高于单对计算
    sim_matrix = cosine_similarity(vectors)
    # 提取上三角矩阵，排除对角线和重复对
    upper_triangle = np.triu(sim_matrix, k=1)
    # 筛选相似度超过阈值的索引对
    i_indices, j_indices = np.where(upper_triangle > threshold)
    
    # 生成最终的ID对和相似度得分
    similarities = list(zip(log_ids[i_indices], log_ids[j_indices], upper_triangle[i_indices, j_indices]))
    
    return sport, similarities

# Dask并行处理流程
# 1. 将原始数据转为Dask Bag，根据CPU核心数设置分区数
data_bag = db.from_sequence(your_raw_data, npartitions=8)
# 2. 按运动类型分组
grouped_bag = data_bag.groupby(lambda x: x[0])
# 3. 并行计算每个分组的相似度对
result_bag = grouped_bag.map(calculate_pairwise_cosine_similarities, threshold=0.95)
# 4. 执行计算并获取结果
final_results = result_bag.compute()

额外优化建议 #

• 分区调整：根据机器内存和CPU核心数调整npartitions，避免分区过大导致内存过载，或过小增加调度开销。
• 二次分组：若单个运动类型记录数仍过大，可按时间（如月份）或其他维度二次分组，拆分计算压力。
• 内存监控：通过Dask diagnostics工具监控内存使用，及时调整参数。

内容的提问来源于stack exchange，提问作者Grace cleland-pottie