并行化Wandb超参数扫掠时重复调用wandb.init()导致进程超时挂起的问题排查与解决咨询

我最近在尝试并行化Wandb超参数扫掠，因为我的模型收敛速度很慢，而且要扫的参数组合特别多，实在耗不起时间。我写了一段通用的代码，用ThreadPoolExecutor启动多个agent来分担扫掠任务，但运行时总是卡在wandb.init()这一步，最后还会因为超时被终止。我觉得不是单纯调大Wandb超时时间能解决的问题，想问问大家这是不是我的并行逻辑有问题？有没有官方或者常用的Wandb扫掠并行化方案可以参考？

我的代码片段如下：

def run_pipeline(args):
    # Stuff happens here

    # Wandb init
    group = "within_session" if session_config["within_session"] else "across_session"
    run = wandb.init(name=f"{sessions[i]}_{group}_decoder_run", group=group, config=sweep_config, reinit=True)

    # Model training

    return results


def run_pipeline_wrapper(args):
    # Stuff happens here
    run_pipeline(args)

    return None


if __name__ == "__main__":
    total_runs = 30
    agents = 5
    runs_per_agent = total_runs // agents

    sweep_config = {'method': 'random'}
    parameters_dict = {
        # Lota of parameters to sweep
    }
    sweep_config['parameters'] = parameters_dict

    # Create a sweep id that stores sweep ids
    sweep_id_json_path = 'sweep_id.json'
    if not os.path.exists(sweep_id_json_path):
        with open(sweep_id_json_path, 'w') as f:
            json.dump({}, f)
    sweep_id_json = json.load(open(sweep_id_json_path, 'r'))

    # Sessions_list = number of unique data that I need to run my sweeps
    for i in range(len(sessions_list)):

        # Preparing a partial method to pass
        run_pipeline_with_args = partial(run_pipeline_wrapper, args)

        # I cache the existing sweep_ids in a json file to help in attaching sweep ids if I rerun the code again
        if f"{sessions_list[i]}_{is_within}" not in sweep_id_json:
            sweep_id = wandb.sweep(sweep_config, project=f"HPC_model_{sess}_session_{data}_{data_type}")
        else:
            sweep_id = wandb.sweep(sweep_config, project=f"HPC_model_{sess}_session_{data}_{data_type}"
                                   , prior_runs=sweep_id_json[f"{sessions_list[i]}_{is_within}"])


        # This is the parallelization logic, where I parallelize the sweeps
        with concurrent.futures.ThreadPoolExecutor(max_workers=agents) as executor:
            futures = [
                executor.submit(wandb.agent, sweep_id, run_pipeline_with_args, count=runs_per_agent)
                for _ in range(agents)
            ]

            concurrent.futures.wait(futures)

运行时的Wandb日志如下：

问题原因分析 #

首先，你的并行逻辑确实可能是导致卡死的核心原因：

• Wandb Agent本身的并行性冲突：Wandb的sweep agent本身就设计成可以在后台管理多个运行任务，你再用ThreadPoolExecutor去启动多个agent实例，会导致多个进程/线程同时争抢Wandb的本地资源（比如缓存文件、进程锁），很容易在wandb.init()时出现资源竞争死锁，最终触发超时。
• 重复创建Sweep与Agent的逻辑冗余：你在每个session循环里都创建新的sweep，还复用sweep_id，再加上线程池的嵌套调用，会让Wandb的后台进程管理混乱，多个agent同时尝试初始化Wandb run，导致连接阻塞。
• Partial函数与参数传递的潜在问题：你用partial包装run_pipeline_wrapper时，参数传递可能存在共享状态的问题——多个线程可能会共享同一个args实例，导致参数混乱，进一步加剧初始化时的冲突。

推荐的解决办法 #

方案1：用Wandb官方的多Agent并行方式（无需手动线程池）

Wandb本身就支持在单进程中启动多个agent，或者在不同终端启动多个agent，但如果要在代码里统一管理，更推荐用多进程而非多线程（因为Python的GIL会限制线程的实际并行性，而且Wandb的操作是IO密集+CPU密集混合的）：

from multiprocessing import Pool
import wandb
import os
import json

def run_agent(sweep_id, runs_per_agent, args):
    # 每个进程独立启动一个agent，确保参数隔离
    run_pipeline_with_args = partial(run_pipeline_wrapper, args)
    wandb.agent(sweep_id, function=run_pipeline_with_args, count=runs_per_agent)

if __name__ == "__main__":
    total_runs = 30
    agents = 5
    runs_per_agent = total_runs // agents

    sweep_config = {'method': 'random'}
    parameters_dict = {
        # 你的参数配置
    }
    sweep_config['parameters'] = parameters_dict

    # 按session循环创建sweep并启动多进程agent
    for i in range(len(sessions_list)):
        # 确保每个session的参数独立
        session_args = args.copy()  # 或者根据session生成独立参数
        # 创建当前session的sweep
        sweep_id = wandb.sweep(sweep_config, project=f"HPC_model_{sess}_session_{data}_{data_type}")
        # 用多进程池启动多个agent
        with Pool(processes=agents) as pool:
            # 每个进程分配独立的参数副本
            pool.starmap(run_agent, [(sweep_id, runs_per_agent, session_args) for _ in range(agents)])

方案2：调整现有逻辑，避免线程池与Wandb Agent的嵌套

如果你必须按session拆分sweep，可以去掉ThreadPoolExecutor，让每个sweep用Wandb Agent自带的并行能力：

1. 每个session创建一个独立的sweep_id
2. 对每个sweep，直接调用wandb.agent并设置count=runs_per_agent，不需要额外线程池——Wandb Agent会自动在后台管理多个run的并行（你可以在sweep_config里设置parallelism参数来控制并行数）

# 在session循环里
sweep_config['parallelism'] = agents  # 控制每个sweep的并行运行数
sweep_id = wandb.sweep(sweep_config, project=f"HPC_model_{sess}_session_{data}_{data_type}")
# 让Wandb Agent自己管理并行
wandb.agent(sweep_id, function=run_pipeline_with_args, count=total_runs)

方案3：修复wandb.init()的线程安全问题

如果你坚持要用线程池，必须确保每个线程的Wandb run是完全隔离的：

• 在run_pipeline里，wandb.init()要加上dir参数，为每个线程指定独立的临时目录，避免共享本地缓存：

import tempfile

def run_pipeline(args):
    # Stuff happens here

    # 为每个run创建独立的临时目录
    run_dir = tempfile.mkdtemp()
    group = "within_session" if session_config["within_session"] else "across_session"
    run = wandb.init(
        name=f"{sessions[i]}_{group}_decoder_run",
        group=group,
        config=sweep_config,
        reinit=True,
        dir=run_dir
    )

    # Model training

    # 清理临时目录
    import shutil
    shutil.rmtree(run_dir)
    return results

• 同时，要确保每个线程的args是独立的副本，不要共享同一个实例，避免参数污染。

额外注意事项 #

• 尽量不要在循环里频繁创建sweep，除非每个session的参数空间完全不同——可以把相同参数空间的session合并到同一个sweep里，用group来区分session。
• 运行前可以清理Wandb的本地缓存：执行wandb cache clean，避免之前的残留进程/文件导致冲突。
• 如果是在HPC集群上运行，要确保每个进程/线程有独立的网络权限，Wandb的初始化需要和云端通信，集群的网络隔离也可能导致超时。

备注：内容来源于stack exchange，提问作者Leofierus