DiskInfo预警磁盘即将满载：避免PyTorch训练中断

DiskInfo预警磁盘即将满载：避免PyTorch训练中断

在一次深夜的模型训练中，一位研究员正等待着第100轮epoch的结果。突然，进程崩溃，日志里只留下一行冰冷的错误：

OSError: [Errno 28] No space left on device

检查点未保存，最近几小时的训练成果付诸东流。问题源头？不是GPU显存不足，也不是代码逻辑错误——而是磁盘满了。

这并非个例。在深度学习项目中，随着模型参数量飙升、数据集规模扩大，以及频繁的checkpoint保存和日志输出，存储空间正在成为比显存更隐蔽但同样致命的瓶颈。尤其在使用容器化环境（如PyTorch-CUDA镜像）进行开发时，开发者往往将注意力集中在torch.cuda.is_available()是否返回True，却忽略了/workspace所在的磁盘是否还能撑过下一个epoch。

PyTorch-CUDA 镜像：便利背后的存储陷阱

PyTorch-CUDA基础镜像极大简化了环境配置流程。以“PyTorch-CUDA-v2.7”为例，它预装了Python、PyTorch 2.7、CUDA 12.x、cuDNN及常用科学计算库，配合NVIDIA Container Toolkit，只需一条命令即可启动支持GPU加速的开发环境：

docker run --gpus all -v /data:/workspace nvcr.io/nvidia/pytorch:23.10-py3

这种开箱即用的体验带来了三个典型副作用：

1. 写入路径集中化：所有用户数据默认落盘于挂载目录（如/workspace），极易形成IO热点；
2. 生命周期管理缺失：临时缓存、中间特征图、调试日志等文件长期滞留；
3. 资源感知弱化：容器内部难以直观感知宿主机磁盘全局状态。

更糟糕的是，在多用户共享服务器场景下，每个人都在“合理”地保存自己的模型快照，最终集体耗尽共享存储。

磁盘监控不该是事后补救

很多人直到收到报警才去清理磁盘，但那时往往为时已晚——训练已中断，而恢复需要重新加载上一个有效checkpoint，甚至从头开始。

真正有效的策略是前置防御。我们需要一种轻量、可靠且可嵌入训练流程的磁盘使用检测机制。幸运的是，Linux系统本身就提供了这样的能力。

shutil.disk_usage()是 Python 标准库中的一个宝藏函数，它能直接调用操作系统接口，获取指定路径所在分区的总容量、已用和可用空间。相比依赖外部工具（如df -h），它无需shell执行权限，兼容性更强，尤其适合受限容器环境。

下面这个函数已经成为我们团队每个训练脚本的“标配”：

import shutil import warnings import logging def check_disk_usage(path="/workspace", threshold=0.85, critical=0.95): """ 检查磁盘使用率并分级告警 Args: path: 目标路径（建议为实际写入目录） threshold: 警告阈值（默认85%） critical: 危急阈值（超过则抛出异常） Returns: bool: 是否安全 """ try: total, used, free = shutil.disk_usage(path) usage_rate = used / total gb_used = used // (1024**3) gb_total = total // (1024**3) logging.info(f"Disk @ {path}: {usage_rate:.1%} used ({gb_used}GB/{gb_total}GB)") if usage_rate >= critical: raise RuntimeError( f"CRITICAL: Disk usage {usage_rate:.1%} exceeds {critical:.0%} limit. " "Aborting to prevent I/O failure." ) elif usage_rate >= threshold: warnings.warn( f"WARNING: High disk usage detected ({usage_rate:.1%}). " "Consider cleanup or expansion.", UserWarning ) return False return True except OSError as e: logging.error(f"Failed to query disk usage: {e}") return False

你可能会问：为什么不在训练开始前检查一次就够了？

因为训练过程本身就在持续消耗磁盘。例如，每轮保存一次checkpoint，每次增加几百MB；开启详细日志记录，每天可能生成数GB文本。因此，周期性检测才是关键。

我们在主训练循环中加入了如下逻辑：

for epoch in range(start_epoch, num_epochs): # 训练一个epoch train_loss = train_one_epoch(model, dataloader, optimizer, device) # 每5个epoch检查一次磁盘健康状况 if epoch % 5 == 0: if not check_disk_usage("/workspace/checkpoints"): logging.warning("Low disk space detected. Stopping early.") break # 提前终止，保留最后可用的checkpoint # 保存模型（仅当磁盘充足时才尝试） save_checkpoint(model, optimizer, epoch, train_loss)

这样做的好处是：即使当前还有空间，但如果趋势显示即将耗尽，我们也能主动止损，而不是等到最后一刻才失败。

实际架构中的集成方式

在一个典型的AI训练环境中，结构通常是这样的：

+----------------------------+ | 宿主机 (Host Machine) | | | | +----------------------+ | | | Docker Engine | | | | | | | | +----------------+ | | | | | PyTorch-CUDA |<-----> NVIDIA Driver | | | Container | | | | | | - PyTorch v2.7 | | | | | | - CUDA 12.x | | | | | | - Jupyter/SSH | | | | | +--------^-------+ | | | | | Mount | | | +-----------|-----------+ | | | /workspace | | +-----v------+ | | | Shared SSD |<-----+ (High-speed storage) | +------------+ | +----------------------------+

这里有几个工程实践要点：

• 挂载独立存储卷：不要将容器工作目录挂载到系统盘（如/）。应使用单独的SSD阵列，并通过-v /ssd_pool:/workspace方式挂载。
• 命名空间隔离：多人共用时，按用户或项目划分子目录，如/workspace/user_a/proj1，便于配额管理和责任追溯。
• 权限与可见性：确保容器内运行的进程有权限访问/proc和/sys文件系统，否则某些监控工具会失效。推荐运行容器时添加--privileged或至少挂载-v /proc:/host_proc:ro。

此外，除了在训练脚本中嵌入检测逻辑，还可以通过外部守护进程实现全局监控。例如，使用cron定时任务每日扫描：

# 每天早上8点检查一次 0 8 * * * /usr/bin/python3 /scripts/disk_monitor.py >> /var/log/disk_check.log 2>&1

结合邮件或企业微信机器人，可以实现跨项目的资源预警。

更进一步：智能化的资源管理

虽然简单的阈值判断已经能解决大部分问题，但在复杂场景下，我们可以做得更智能。

动态趋势预测

如果发现磁盘使用率连续上升（比如过去三次检测分别为80% → 83% → 86%），说明该任务正处于高速写入阶段。此时即便尚未达到85%阈值，也应提前警告。

class DiskMonitor: def __init__(self, path, window_size=3): self.path = path self.history = [] self.window_size = window_size def record(self): try: total, used, _ = shutil.disk_usage(self.path) rate = used / total self.history.append(rate) if len(self.history) > self.window_size: self.history.pop(0) return rate except OSError: return None def is_rising_fast(self, growth_threshold=0.05): if len(self.history) < 2: return False avg_growth = (self.history[-1] - self.history[0]) / (len(self.history) - 1) return avg_growth > growth_threshold

自动清理策略

对于checkpoint，保留全部版本既浪费空间又无必要。常见的做法是“保留最近N个 + 关键里程碑”：

import os from pathlib import Path def cleanup_checkpoints(checkpoint_dir, keep_last=5, keep_best=True): ckpt_files = sorted(Path(checkpoint_dir).glob("*.pth"), key=os.path.getctime) # 总数不超过keep_last时不清理 if len(ckpt_files) <= keep_last: return # 删除最旧的，直到只剩keep_last个 for old_file in ckpt_files[:-keep_last]: if keep_best and "best" in old_file.name: continue # 跳过最佳模型 try: old_file.unlink() logging.info(f"Deleted old checkpoint: {old_file}") except Exception as e: logging.warning(f"Failed to delete {old_file}: {e}")

日志轮转不可少

大量调试信息累积起来也非常可观。使用logging.handlers.RotatingFileHandler可自动分割和压缩日志：

from logging.handlers import RotatingFileHandler handler = RotatingFileHandler( "training.log", maxBytes=100*1024*1024, # 100MB backupCount=5 # 最多保留5个旧日志 )

工程建议清单

写在最后

在追求更大模型、更大数据的时代，我们习惯性地升级GPU、扩展内存，却常常忽视那个最基础的组件——磁盘。然而正是这个看似“传统”的资源，一次次让精心设计的训练任务功亏一篑。

引入disk_usage检测的成本几乎为零：一段几十行的代码，就能换来训练稳定性的显著提升。更重要的是，它提醒我们：真正的工程能力不仅体现在模型精度上，更体现在对整个系统生命周期的掌控力。

当你下次构建PyTorch训练脚本时，不妨先把这段磁盘检查加进去。也许某一天，它就会默默帮你避免一场灾难。