Slurm集群上Python多进程print()无法输出到控制台问题

核心原因 #

multiprocessing子进程的标准输出默认处于缓冲状态，结合Slurm的任务输出捕获机制，导致子进程内的print内容无法实时传递到控制台；此外你代码里的pool.close()属于冗余操作，with语句会自动完成Pool的资源回收。

具体解决办法 #

• 强制关闭输出缓冲：
  可以在Python脚本开头添加设置，强制禁用输出缓冲：

  import sys
  import os
  sys.stdout.flush()
  os.environ['PYTHONUNBUFFERED'] = '1'
  

  或者直接在Slurm提交命令中设置环境变量：

  srun -c 8 --gres="ht:ht1:8" PYTHONUNBUFFERED=1 python -m my_module...
  

• 用initializer初始化子进程输出流：
  给Pool指定初始化函数，让每个子进程都绑定主进程的标准输出流：

  def init_worker():
      import sys
      sys.stdout = sys.__stdout__
      sys.stderr = sys.__stderr__
  
  with Pool(processes=num_processes, initializer=init_worker) as pool:
      out = pool.map(
          partial(df_to_stats, **param_dict), frames_uri[0:max_uri_process]
      )
  

• 替换print为logging模块：
  logging模块原生支持多进程输出，配置后输出更稳定：

  import logging
  import sys
  from multiprocessing import Pool
  from functools import partial
  
  logging.basicConfig(
      level=logging.INFO,
      format='%(asctime)s - %(processName)s - %(message)s',
      stream=sys.stdout
  )
  
  def df_to_stats(uri, **param_dict):
      # 替换print为logging.info
      logging.info("test")
      # 自制tqdm进度条可结合logging输出，或改用tqdm的多进程兼容版本
      ...
  
  with Pool(processes=num_processes) as pool:
      out = pool.map(
          partial(df_to_stats, **param_dict), frames_uri[0:max_uri_process]
      )
  

• Slurm输出定向到文件：
  若无需实时控制台输出，可让Slurm将输出写入文件，提交命令时添加参数：

  srun -c 8 --gres="ht:ht1:8" --output=job_%j.out --error=job_%j.err python -m my_module...
  

  任务结束后查看job_<job_id>.out文件即可获取所有输出内容。

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