AzureML训练CV模型保存.pth checkpoint报设备空间不足错误

我在AzureML上训练计算机视觉模型，之前用Python SDK提交Job遇到问题，改用Azure Storage Explorer把代码直接上传到Blob存储的AzureML存储账户文件夹后，Job能正常提交，但保存15MB左右的.pth checkpoint文件时，出现OSError: No space left on device错误。此前我完成过生成更多、更大文件的训练运行，求解决该问题。

环境信息 #

• 数据存储在ADLS Gen2热层
• Azure ML存储账户为100TiB文件共享
• 文件大小：269KB(.log)、5.8MB(.pth)、15MB(.pth)、18KB(.py)
• 计算实例：Standard_NC24ads_A100_v4

Job代码 #

training_job = command(
    inputs={
        'data':Input(
                type="uri_folder",
                path=f"{data_asset.id}",
                mode=InputOutputModes.RO_MOUNT
        )},
    outputs={
        'outputs':Output(
            type='uri_folder',
            path = f'azureml://datastores/workspaceartifactstore/paths/ExperimentRun/dcid.{command_name}/outputs',
            mode  =InputOutputModes.RW_MOUNT
        )},
    code=f'https://{storage_account}.blob.core.windows.net/job1-data/code/',
    command='python train.py',
    environment="{name}{version}".format(name=az_cfg['environment']['name'],version=az_cfg['environment']['version']),
    display_name=az_cfg['job']['display_name_train'],
    compute = ci_name,
    environment_variables={ 
      "DATASET_MOUNT_BLOCK_BASED_CACHE_ENABLED": True,
      "EXPERIMENT_NAME":az_cfg['mlflow']['experiment_name']
      },
    experiment_name=az_cfg['mlflow']['experiment_name'],
    name=command_name
      
)

补充说明 #

我使用MMPose和MMEngine框架，保存checkpoint的核心代码片段如下：

def put(self, obj: bytes, filepath: Union[str, Path]) -> None:
    """Write bytes to a given ``filepath`` with 'wb' mode.

    Note:
        ``put`` will create a directory if the directory of
        ``filepath`` does not exist.

    Args:
        obj (bytes): Data to be written.
        filepath (str or Path): Path to write data.

    Examples:
        >>> backend = LocalBackend()
        >>> filepath = '/path/of/file'
        >>> backend.put(b'hello world', filepath)
    """
    mmengine.mkdir_or_exist(osp.dirname(filepath))
    with open(filepath, 'wb') as f:
        f.write(obj)

错误日志 #

File "src/train.py", line 219, in <module>
    main()
  File "src/train.py", line 204, in main
    runner.train()
  File "/opt/conda/envs/ptca/lib/python3.8/site-packages/mmengine/runner/runner.py", line 1777, in train
    model = self.train_loop.run()  # type: ignore
  File "/opt/conda/envs/ptca/lib/python3.8/site-packages/mmengine/runner/loops.py", line 102, in run
    self.runner.val_loop.run()
  File "/opt/conda/envs/ptca/lib/python3.8/site-packages/mmengine/runner/loops.py", line 375, in run
    self.runner.call_hook('after_val_epoch', metrics=metrics)
  File "/opt/conda/envs/ptca/lib/python3.8/site-packages/mmengine/runner/runner.py", line 1839, in call_hook
    getattr(hook, fn_name)(self, **kwargs)
  File "/opt/conda/envs/ptca/lib/python3.8/site-packages/mmengine/hooks/checkpoint_hook.py", line 361, in after_val_epoch
    self._save_best_checkpoint(runner, metrics)
  File "/opt/conda/envs/ptca/lib/python3.8/site-packages/mmengine/hooks/checkpoint_hook.py", line 579, in _save_best_checkpoint
    self._save_checkpoint_with_step(runner, cur_time, meta)
  File "/opt/conda/envs/ptca/lib/python3.8/site-packages/mmengine/hooks/checkpoint_hook.py", line 443, in _save_checkpoint_with_step
    runner.save_checkpoint(
  File "/opt/conda/envs/ptca/lib/python3.8/site-packages/mmengine/dist/utils.py", line 427, in wrapper
    return func(*args, **kwargs)
  File "/opt/conda/envs/ptca/lib/python3.8/site-packages/mmengine/runner/runner.py", line 2271, in save_checkpoint
    save_checkpoint(
  File "/opt/conda/envs/ptca/lib/python3.8/site-packages/mmengine/runner/checkpoint.py", line 793, in save_checkpoint
    file_backend.put(f.getvalue(), filename)
  File "/opt/conda/envs/ptca/lib/python3.8/site-packages/mmengine/fileio/backends/local_backend.py", line 77, in put
    with open(filepath, 'wb') as f:
OSError: [Errno 28] No space left on device: '/mnt/azureml/cr/j/21f0a13bbecb4e2cb2bebba8539778b6/cap/data-capability/wd/outputs/epoch_1.pth'

1. 排查本地临时磁盘空间 #

错误中的路径/mnt/azureml/cr/j/...是AzureML作业的本地临时运行目录，和你配置的100TiB存储账户无关。Standard_NC24ads_A100_v4的本地临时磁盘容量有限，可能被旧作业残留、数据集缓存或中间文件占满：

• 在作业命令前添加磁盘检查：command='df -h && python train.py'，查看/mnt路径的剩余空间。
• 清理临时文件：修改训练代码，只保留最新的1-2个checkpoint，自动删除旧文件；或者在启动命令中添加清理逻辑：rm -rf /mnt/azureml/cr/j/* && python train.py。

2. 切换Output挂载模式为UPLOAD #

当前Output使用RW_MOUNT模式，会将远程存储挂载到本地目录，写入操作依赖本地临时磁盘。改用UPLOAD模式，训练时先将checkpoint保存到本地临时目录，作业结束后自动同步到Blob存储：

outputs={
    'outputs':Output(
        type='uri_folder',
        path = f'azureml://datastores/workspaceartifactstore/paths/ExperimentRun/dcid.{command_name}/outputs',
        mode  =InputOutputModes.UPLOAD  # 替换为UPLOAD
    )},

同时在train.py中指定保存路径为环境变量AZUREML_OUTPUT_DIR对应的目录，该路径会自动关联到配置的Output存储位置。

3. 配置MMEngine直接写入Azure Blob存储 #

绕过本地磁盘，让MMEngine直接将checkpoint写入Azure Blob：

• 安装依赖：pip install azure-storage-blob
• 在MMEngine配置或代码中初始化AzureBlobBackend，替换默认的LocalBackend
• 保存checkpoint时指定使用AzureBlobBackend，直接写入远程存储

4. 关闭或优化数据集缓存 #

你开启了DATASET_MOUNT_BLOCK_BASED_CACHE_ENABLED=True，数据集缓存可能占用大量本地空间：

• 临时关闭缓存：将环境变量设为False，观察是否解决问题
• 若需要缓存，配置缓存路径到更大的存储卷，或者定期清理缓存文件

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