PytorchStreamWriter failed writing file data

[albertz]

Summary

A single checkpoint-save write() failed and crashed the entire training job. Engine._save_model calls torch.save(...) once; if the underlying write fails (a transient shared-filesystem error, or a brief quota bounce), the exception propagates and the whole run dies, requiring a manual restart -- costly for multi-day trainings.

The atomic-write part is already handled (temp file + rename), so the previous good checkpoint is never corrupted. What is missing is a retry: the failure modes that hit this in practice (a momentary Lustre/NFS RPC error, or a short brush against a soft quota with grace) are exactly the kind that succeed on a second attempt moments later.

Traceback

Traceback (most recent call last):
  File ".../torch/serialization.py", line 965, in save
    _save(obj, opened_zipfile, ...)
  File ".../torch/serialization.py", line 1266, in _save
    zip_file.write_record(name, storage, num_bytes)   # name='data/322', num_bytes=16777216
RuntimeError: [enforce fail at inline_container.cc:857] . PytorchStreamWriter failed writing file data/322: file write failed

During handling of the above exception, another exception occurred:
  ...
  File ".../returnn/torch/engine.py", line 700, in Engine.train_epoch
    self._save_model()
  File ".../returnn/torch/engine.py", line 1334, in Engine._save_model
    torch.save({...}, tmp_filename)

Note that torch surfaces only file write failed here, without the underlying errno, so the exact cause (EIO vs EDQUOT vs ...) is not recoverable from the message. That is part of the motivation for retrying rather than trying to classify the error.

Context / how it was hit

Observed once (first time we have seen it) on RWTH hpcwork (Lustre), during an otherwise-healthy multi-day run:

• The write failed for one training at 21:34, while four other RETURNN trainings on the same project checkpointed successfully within the same minute.
• After simply clearing the error flag and resubmitting, the job resumed from the last checkpoint and has since written two more epochs with no error.

So it was not a persistent condition -- neither a full disk nor a hard-exceeded quota (either of those would have failed the sibling writes and the resume too). The project quota at the time: ~79% of the space limit (39.4T / 50T) and ~93% of the inode limit (4.87M / 5.24M files) -- close on inodes but not at the hard cap, so a hard-quota failure is unlikely though a momentary soft-quota bounce cannot be fully excluded. Net: a rare, transient, self-recovering write failure.

Current code

returnn/torch/engine.py Engine._save_model (already writes atomically):

tmp_filename = filename + ".tmp_write"
if os.path.exists(tmp_filename):
    os.unlink(tmp_filename)
...
torch.save({...}, tmp_filename)
os.rename(tmp_filename, filename)

Proposal

Wrap the torch.save(...) in a small retry loop (e.g. 3-5 attempts with short backoff) on OSError / RuntimeError, re-creating the temp file on each attempt, before finally giving up. The atomic rename already guarantees no partial/corrupt checkpoint is left behind, so retrying is safe. This turns a fatal, manual-restart-requiring crash into a self-healing hiccup for the common case of a transient shared-filesystem write error (or a brief quota bounce).

The same treatment likely makes sense for the optimizer-state save on the same path.

Environment

• RETURNN git e32117f9 (the failing run reported 1.20260715.003614+git.65c8c876)
• torch 2.7.1+cu126, Python 3.12.3
• RWTH hpcwork Lustre (lustre22), single-GPU training