Slurm运行PyTorch训练脚本时stdout无法实时写入日志

该异常是Python标准输出的默认缓冲策略切换导致的，和Slurm调度、PyTorch训练逻辑本身无关：

• Python对标准输出（stdout）默认采用两种缓冲规则：当stdout连接到交互式终端（TTY）时使用行缓冲，每遇到换行符就立刻把内容刷出到输出流；当stdout被重定向到文件、管道（Slurm提交任务时默认会把任务stdout重定向到日志文件，属于该场景）时，会自动切换为块缓冲，默认缓冲区大小通常为4KB~64KB，只有攒够缓冲区容量、进程退出、手动触发刷写三个条件满足其一时，才会把缓冲区里的内容一次性写入文件。你观察到的累计2000次迭代后批量输出，刚好是2000行print内容凑满了默认缓冲区的大小。
• 启动阶段可见的oneDNN运行提示、BERT权重加载提示，要么是TensorFlow、PyTorch的C++底层直接调用系统接口做的无缓冲输出，不经过Python层面的stdout缓冲，要么是初始化阶段输出量小刚好在缓冲区初始化时被刷出，所以能实时写入日志；训练循环内的print是纯Python层面的输出，完全走块缓冲规则，才会出现延迟。
• 之前用TensorFlow时没有该问题，是因为TensorFlow的训练进度日志多数是C++层输出，本身不经过Python的stdout缓冲，不会被块缓冲规则限制。

可以写一个极简测试脚本复现现象：

import time
for i in range(3000):
    print(f"iter {i}")
    time.sleep(0.1)

用和之前完全一致的Slurm配置提交任务，观察日志会发现前2000行左右不会实时写入，攒够量后才批量刷出，和当前异常现象完全匹配。

按侵入性从低到高选一个即可：

• 无代码修改方案：在sbatch提交脚本中，把启动训练的命令从python train.py改为python -u train.py。-u参数会强制Python关闭stdout、stderr的块缓冲，所有输出会立刻刷入流，Slurm可实时写入日志文件。
• 全局配置方案：如果不方便修改提交脚本，可以在训练代码的最开头（所有其他import语句之前）添加如下代码，设置无缓冲环境变量，效果和-u参数完全一致：

import os
os.environ["PYTHONUNBUFFERED"] = "1"

• 局部刷写方案：如果不想全局关闭缓冲，只需要训练进度实时输出，可以在训练循环的print语句中添加flush=True参数，强制每次打印时立刻刷空缓冲区：

# 训练循环内的打印示例
print(f"Epoch: {epoch}, Step: {step}, Train Loss: {loss_avg:.4f}", flush=True)

• 如果你用tqdm做进度条展示，额外给tqdm传入参数file=sys.stderr即可避免进度条缓存问题，不需要额外修改其他配置。

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