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"] += 1return 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_BARSiterator = 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_procsbuckets[proc].append((i, item))return bucketsdef merge_outputs(dir_list):"""后处理合并每个进程的分片；根据输出布局调整。"""passdef main(root_dir, num_procs):# 发现工作项items = [...] # 填充你的项列表if not items:safe_print("未找到项。")returnnum_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 sysif 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的返回结果
- 每个
