如何为使用apply_async的多进程脚本添加tqdm进度条

针对你用apply_async实现多进程URL特征提取时无法正确添加tqdm进度条的问题，这里提供一种基于进程安全计数器的实现方案，直接适配你的start()和resume()函数：

修改步骤 #

1. 补充导入依赖 #

在data_collection.py的导入部分新增两个模块：

import time
from multiprocessing import Value

2. 修改worker函数，添加任务完成计数器 #

给worker函数增加一个共享计数器参数，任务完成后递增计数器（用锁保证进程安全）：

def worker(url: str, dataset: str, q: multiprocessing.managers.AutoProxy, counter: Value) -> type(None):
    try:
        res = UrlFeaturizer(url).run(dataset)[1]
        q.put(res)
    except (AttributeError, TimeoutError):
        pass
    finally:
        # 任务完成后递增计数器，加锁避免多进程冲突
        with counter.get_lock():
            counter.value += 1
    return

3. 改造start()函数，添加进度条 #

先统计总任务数，初始化tqdm进度条，通过共享计数器实时更新进度：

def start(
    pool: multiprocessing.pool.Pool,
    q: multiprocessing.managers.AutoProxy,
    jobs: list[multiprocessing.pool.ApplyResult],
) -> type(None):
    datasets = ["benign_domains.csv", "dmca_domains.csv"]
    
    # 统计所有待处理URL总数
    total_urls = 0
    url_list = []
    for dataset in datasets:
        urls = pd.read_csv(dataset, header=None).iloc[:, 0].to_list()
        url_list.extend(urls)
        total_urls += len(urls)
    
    # 创建进程安全的计数器，初始值为0
    counter = Value('i', 0)
    # 初始化tqdm进度条
    pbar = tqdm(total=total_urls, desc="处理URL中")
    
    # 提交所有异步任务
    for url in url_list:
        job = pool.apply_async(worker, (url, dataset, q, counter))
        jobs.append(job)
    
    # 实时更新进度条，直到所有任务完成
    while counter.value < total_urls:
        pbar.n = counter.value
        pbar.refresh()
        time.sleep(0.1)  # 避免循环过于频繁占用CPU
    
    pbar.close()
    
    # 收尾工作
    q.put("kill")
    pool.close()
    pool.join()
    sort_csv()
    return

4. 改造resume()函数，适配进度条 #

同样先统计未处理URL数量，再初始化进度条：

def resume(
    pool: multiprocessing.pool.Pool,
    q: multiprocessing.managers.AutoProxy,
    jobs: list[multiprocessing.pool.ApplyResult],
) -> type(None):
    processed_urls = (
        pd.read_csv(
            "num_features.csv",
            usecols=[0],
            header=None,
        )
        .iloc[:, 0]
        .to_list()
    )
    datasets = ["benign_domains.csv", "dmca_domains.csv"]
    
    # 收集所有未处理的URL
    url_list = []
    for dataset in datasets:
        unprocessed_urls = pd.read_csv(dataset, header=None).iloc[:, 0].to_list()
        urls = [item for item in unprocessed_urls if item not in processed_urls]
        url_list.extend(urls)
    
    total_urls = len(url_list)
    # 没有待处理URL直接收尾
    if total_urls == 0:
        print("无待处理URL")
        q.put("kill")
        pool.close()
        pool.join()
        return
    
    # 初始化计数器和进度条
    counter = Value('i', 0)
    pbar = tqdm(total=total_urls, desc="恢复处理URL中")
    
    # 提交异步任务
    for url in url_list:
        job = pool.apply_async(worker, (url, dataset, q, counter))
        jobs.append(job)
    
    # 更新进度条
    while counter.value < total_urls:
        pbar.n = counter.value
        pbar.refresh()
        time.sleep(0.1)
    
    pbar.close()
    
    # 收尾工作
    q.put("kill")
    pool.close()
    pool.join()
    sort_csv()
    return

关键说明 #

• 用Value('i', 0)创建的共享计数器是进程安全的，必须通过get_lock()加锁后修改，避免多进程竞争导致计数错误。
• 主进程通过循环检查计数器值更新进度条，time.sleep(0.1)可以减少CPU占用，同时保证进度更新的实时性。
• 先统计总任务数，让tqdm能显示完整的进度范围，用户可以清晰看到剩余任务量。

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