使用Popen调用Slurm sbatch跨节点运行时出现Wandb连接拒绝错误

我在Python脚本中使用subprocess.Popen执行Slurm sbatch命令，将任务提交到其他计算节点时，调用wandb.init()出现错误：
wandb.sdk.wandb_manager.ManagerConnectionRefusedError: Connection to wandb service failed since the process is not available.

相关代码 #

sbatch命令代码 #

p = Popen([shutil.which("sbatch"), '--mem=40G', '--gres=gpu:titan_xp:1', '--nodelist=tikgpu02', '--cpus-per-task=2', '--output=/home/pschlaepfer/denselp/slt/log/%j.out', '--error=/home/pschlaepfer/denselp/slt/log/%j.err', '/home/pschlaepfer/denselp/slt/scripts/slt.sh', '--action=fine-tune-thf', '--max-length', '128', '--lr=4e-5', '--epochs=5', '--batch-size=16', '--task', task, '--pre-trained-path', checkpoint_path, '--wandb-mode=offline'], start_new_session=True)

wandb.init()调用代码 #

experiment_name = f"job-id:{meta_config.job_id}"
run = wandb.init(
  project=wandb_project_choice+("-proto" if meta_config.is_debug_instance else ""),
  name=experiment_name,
  tags=[
    "job_id:"+str(meta_config.job_id)
  ],
  settings=wandb.Settings(start_method='fork'),
  dir=wandb_logging_dir_path,
  config=dict(experiment_config._asdict()) if type(experiment_config).__name__ == 'ExperimentConfig' else dict(experiment_config._as_dict()),
  reinit=True,
  mode="offline",
)

完整错误栈追踪 #

Traceback (most recent call last):
File “/itet-stor/pschlaepfer/net_scratch/conda_envs/denselp/lib/python3.10/site-packages/wandb/sdk/wandb_manager.py”, line 115, in _service_connect
svc_iface._svc_connect(port=port)
File “/itet-stor/pschlaepfer/net_scratch/conda_envs/denselp/lib/python3.10/site-packages/wandb/sdk/service/service_sock.py”, line 30, in _svc_connect
self._sock_client.connect(port=port)
File “/itet-stor/pschlaepfer/net_scratch/conda_envs/denselp/lib/python3.10/site-packages/wandb/sdk/lib/sock_client.py”, line 102, in connect
s.connect((“localhost”, port))
ConnectionRefusedError: [Errno 111] Connection refused

During handling of the above exception, another exception occurred:
Traceback (most recent call last):
File “/itet-stor/pschlaepfer/net_scratch/conda_envs/denselp/lib/python3.10/runpy.py”, line 196, in _run_module_as_main
return _run_code(code, main_globals, None,
File “/itet-stor/pschlaepfer/net_scratch/conda_envs/denselp/lib/python3.10/runpy.py”, line 86, in _run_code
exec(code, run_globals)
File “/home/pschlaepfer/denselp/slt/main.py”, line 107, in
run = wandb.init(
File “/itet-stor/pschlaepfer/net_scratch/conda_envs/denselp/lib/python3.10/site-packages/wandb/sdk/wandb_init.py”, line 1185, in init
raise e
File “/itet-stor/pschlaepfer/net_scratch/conda_envs/denselp/lib/python3.10/site-packages/wandb/sdk/wandb_init.py”, line 1162, in init
wi.setup(kwargs)
File “/itet-stor/pschlaepfer/net_scratch/conda_envs/denselp/lib/python3.10/site-packages/wandb/sdk/wandb_init.py”, line 189, in setup
self._wl = wandb_setup.setup(settings=setup_settings)
File “/itet-stor/pschlaepfer/net_scratch/conda_envs/denselp/lib/python3.10/site-packages/wandb/sdk/wandb_setup.py”, line 327, in setup
ret = _setup(settings=settings)
File “/itet-stor/pschlaepfer/net_scratch/conda_envs/denselp/lib/python3.10/site-packages/wandb/sdk/wandb_setup.py”, line 320, in _setup
wl = _WandbSetup(settings=settings)
File “/itet-stor/pschlaepfer/net_scratch/conda_envs/denselp/lib/python3.10/site-packages/wandb/sdk/wandb_setup.py”, line 303, in init
_WandbSetup._instance = _WandbSetup__WandbSetup(settings=settings, pid=pid)
File “/itet-stor/pschlaepfer/net_scratch/conda_envs/denselp/lib/python3.10/site-packages/wandb/sdk/wandb_setup.py”, line 114, in init
self._setup()
File “/itet-stor/pschlaepfer/net_scratch/conda_envs/denselp/lib/python3.10/site-packages/wandb/sdk/wandb_setup.py”, line 250, in _setup
self._setup_manager()
File “/itet-stor/pschlaepfer/net_scratch/conda_envs/denselp/lib/python3.10/site-packages/wandb/sdk/wandb_setup.py”, line 277, in _setup_manager
self._manager = wandb_manager._Manager(settings=self._settings)
File “/itet-stor/pschlaepfer/net_scratch/conda_envs/denselp/lib/python3.10/site-packages/wandb/sdk/wandb_manager.py”, line 152, in init
wandb._sentry.reraise(e)
File “/itet-stor/pschlaepfer/net_scratch/conda_envs/denselp/lib/python3.10/site-packages/wandb/analytics/sentry.py”, line 154, in reraise
raise exc.with_traceback(sys.exc_info()[2])
File “/itet-stor/pschlaepfer/net_scratch/conda_envs/denselp/lib/python3.10/site-packages/wandb/sdk/wandb_manager.py”, line 150, in init
self._service_connect()
File “/itet-stor/pschlaepfer/net_scratch/conda_envs/denselp/lib/python3.10/site-packages/wandb/sdk/wandb_manager.py”, line 124, in _service_connect
raise ManagerConnectionRefusedError(message)
wandb.sdk.wandb_manager.ManagerConnectionRefusedError: Connection to wandb service failed since the process is not available.

使用的Wandb版本为0.16.0。

解决方案 #

• 调整wandb启动模式：将wandb.Settings(start_method='fork')改为wandb.Settings(start_method='spawn')。Slurm任务在远程节点运行时，fork模式会继承提交节点的进程环境，导致wandb服务连接失败；spawn模式会创建全新的进程，适配远程节点的环境。
• 配置环境变量：在任务脚本slt.sh开头添加export WANDB_START_METHOD=spawn，确保计算节点上的wandb启动环境正确生效，避免代码中的设置被环境覆盖。
• 升级wandb版本：0.16.0存在部分进程通信相关的已知问题，升级到0.17.x及以上的稳定版本可以修复这类连接异常。
• 检查节点环境：确认计算节点tikgpu02上的conda环境与提交节点一致，且wandb依赖包完整，无权限限制导致服务无法初始化。

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