多进程运行代码模板

import os
import json
import random
import glob
import multiprocessing as mp
from collections import defaultdict
from tqdm import tqdm

# 配置变量：修改这些值来适应不同任务
MAX_BARS = 3  # 可见进度条的最大数量
RESUME_COUNTS_DIR = "/path/to/resume"  # 恢复计数器的目录路径

def safe_print(msg: str):
    tqdm.write(msg)  # 与tqdm兼容的打印

def process_item(item, proc_idx, state):
    """
    在这里实现每个任务的工作逻辑。
    `item` 是任务项。
    `state` 可以保存每个进程的缓存、计数器等。
    """
    # ... 执行实际工作 ...
    state["count"] += 1
    return f"done-{item}"

def worker(proc_idx, tasks):
    touched_dirs = set()
    counters = defaultdict(int)
    resume_counts_path=os.path.join(RESUME_COUNTS_DIR, f"file_select_counts_p{proc_idx}.json")

    # 如果需要恢复，则加载每个进程的计数器
    if resume_counts_path and os.path.exists(resume_counts_path):
        try:
            with open(resume_counts_path, "r", encoding="utf-8") as f:
                counters.update(json.load(f))
        except Exception as e:
            safe_print(f"进程 {proc_idx} 加载计数器失败: {e}")

    show_bar = proc_idx < MAX_BARS
    iterator = tqdm(
        tasks,
        desc=f"进程 {proc_idx}",
        unit="task",
        position=proc_idx if show_bar else 0, # 仅前 N 个占用独立行
        disable=not show_bar, # 其他进程隐藏进度条
        leave=True, # 进度条在完成后会保留在终端中，而不是被清除掉
        ncols=100, # 设置进度条的总宽度为 100 个字符,确保在终端中显示得更整齐
    )

    for task_id, item in iterator:
        process_item(item, proc_idx, counters)
        if resume_counts_path:
            with open(resume_counts_path, "w", encoding="utf-8") as f:
                json.dump(dict(counters), f, indent=2, ensure_ascii=False)

    return list(touched_dirs)

def distribute_tasks(items, num_procs):
    """
    示例：随机抽取与轮询分配相结合的任务分配。
    根据需要替换为其他分片方式（轮询、块分割、工作窃取等）。
    """
    total_cnt = ...
    buckets = [[] for _ in range(num_procs)]
    for i in range(total_cnt):
        item = random.choices(items, k=1)[0]
        proc = i % num_procs
        buckets[proc].append((i, item))
    return buckets

def merge_outputs(dir_list):
    """
    后处理合并每个进程的分片；根据输出布局调整。
    """
    pass

def main(root_dir, num_procs):
    # 发现工作项
    items = [...]  # 填充你的项列表
    if not items:
        safe_print("未找到项。")
        return

    num_procs = max(1, num_procs)
    tasks_per_proc = distribute_tasks(items, num_procs)

    with mp.Pool(processes=num_procs) as pool:
        results = pool.starmap(
            worker,
            [
                (
                    idx,
                    tasks,
                )
                for idx, tasks in enumerate(tasks_per_proc)
            ],
        )

    all_dirs = [d for sub in results for d in sub]
    merge_outputs(all_dirs)

if __name__ == "__main__":
    import sys
    if len(sys.argv) != 3:
        safe_print("Usage: python script.py <directory> <num_processes>")
        sys.exit(1)
    root_dir = sys.argv[1]
    num_processes = int(sys.argv[2])
    main(root_dir, num_processes)

注意事项

• 多进程代码比较重要的事情：
  • 有恢复功能：上面的代码实现了resume_counts_path
  • 每个进程之间产生的文件不冲突：上面的代码让每个进程并行运行，产生的文件都附带了进程号（比如file_select_counts_p{proc_idx}.json），最后可以通过一个函数合并
• with mp.Pool(processes=num_procs) as pool:同时启动了num_procs个worker函数
  • 每个worker函数处理一个(idx, tasks)元组
  • 应该保证tasks_per_proc的长度与num_procs相同
  • results是一个列表，长度为num_procs，每个元素是每个worker的返回结果