Lightning-AI / Lightning-AI/pytorch-lightning
Async checkpointing race conditions
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 31.4k
- Forks
- 3.8k
- Avg merge
- 6d 7h
- Merged PRs (30d)
- 6
Description
## 🐛 Bug
We've seen flakiness in the CI due to race conditions in the async checkpointing logic.
### To Reproduce
The specific test in our code base: test_async_checkpoint_plugin
A CI run where this occurred: https://github.com/Lightning-AI/lightning/runs/7672029478?check_suite_focus=true
Error:
```
> trainer.fit(model)
D:\a\lightning\lightning\tests\tests_pytorch\plugins\test_checkpoint_io_plugin.py:127:
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _
d:\a\lightning\lightning\src\pytorch_lightning\trainer\trainer.py:698: in fit
self._call_and_handle_interrupt(
d:\a\lightning\lightning\src\pytorch_lightning\trainer\trainer.py:650: in _call_and_handle_interrupt
return trainer_fn(*args, **kwargs)
d:\a\lightning\lightning\src\pytorch_lightning\trainer\trainer.py:739: in _fit_impl
results = self._run(model, ckpt_path=self.ckpt_path)
d:\a\lightning\lightning\src\pytorch_lightning\trainer\trainer.py:1179: in _run
results = self._run_stage()
d:\a\lightning\lightning\src\pytorch_lightning\trainer\trainer.py:1265: in _run_stage
return self._run_train()
d:\a\lightning\lightning\src\pytorch_lightning\trainer\trainer.py:1296: in _run_train
self.fit_loop.run()
d:\a\lightning\lightning\src\pytorch_lightning\loops\loop.py:201: in run
self.on_advance_end()
d:\a\lightning\lightning\src\pytorch_lightning\loops\fit_loop.py:298: in on_advance_end
self.trainer._call_callback_hooks("on_train_epoch_end")
d:\a\lightning\lightning\src\pytorch_lightning\trainer\trainer.py:1610: in _call_callback_hooks
fn(self, self.lightning_module, *args, **kwargs)
d:\a\lightning\lightning\src\pytorch_lightning\callbacks\model_checkpoint.py:311: in on_train_epoch_end
self._save_topk_checkpoint(trainer, monitor_candidates)
d:\a\lightning\lightning\src\pytorch_lightning\callbacks\model_checkpoint.py:382: in _save_topk_checkpoint
self._save_monitor_checkpoint(trainer, monitor_candidates)
d:\a\lightning\lightning\src\pytorch_lightning\callbacks\model_checkpoint.py:662: in _save_monitor_checkpoint
self._update_best_and_save(current, trainer, monitor_candidates)
d:\a\lightning\lightning\src\pytorch_lightning\callbacks\model_checkpoint.py:716: in _update_best_and_save
trainer.strategy.remove_checkpoint(del_filepath)
d:\a\lightning\lightning\src\pytorch_lightning\strategies\strategy.py:460: in remove_checkpoint
self.checkpoint_io.remove_checkpoint(filepath)
C:\hostedtoolcache\windows\Python\3.9.13\x64\lib\unittest\mock.py:1092: in __call__
return self._mock_call(*args, **kwargs)
C:\hostedtoolcache\windows\Python\3.9.13\x64\lib\unittest\mock.py:1096: in _mock_call
return self._execute_mock_call(*args, **kwargs)
C:\hostedtoolcache\windows\Python\3.9.13\x64\lib\unittest\mock.py:1166: in _execute_mock_call
return self._mock_wraps(*args, **kwargs)
d:\a\lightning\lightning\src\pytorch_lightning\plugins\io\wrapper.py:61: in remove_checkpoint
self.checkpoint_io.remove_checkpoint(*args, **kwargs)
C:\hostedtoolcache\windows\Python\3.9.13\x64\lib\unittest\mock.py:1092: in __call__
return self._mock_call(*args, **kwargs)
C:\hostedtoolcache\windows\Python\3.9.13\x64\lib\unittest\mock.py:1096: in _mock_call
return self._execute_mock_call(*args, **kwargs)
C:\hostedtoolcache\windows\Python\3.9.13\x64\lib\unittest\mock.py:1166: in _execute_mock_call
return self._mock_wraps(*args, **kwargs)
d:\a\lightning\lightning\src\pytorch_lightning\plugins\io\torch_plugin.py:95: in remove_checkpoint
fs.rm(path, recursive=True)
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _
self =
path = ['C:\\Users\\runneradmin\\AppData\\Local\\Temp\\pytest-of-runneradmin\\pytest-0\\test_async_checkpoint_plugin0\\epoch=0-step=1.ckpt']
recursive = True, maxdepth = None
def rm(self, path, recursive=False, maxdepth=None):
if isinstance(path, str):
path = [path]
for p in path:
p = self._strip_protocol(p).rstrip("/")
if recursive and self.isdir(p):
if osp.abspath(p) == os.getcwd():
raise ValueError("Cannot delete current working directory")
shutil.rmtree(p)
else:
> os.remove(p)
E PermissionError: [WinError 32] The process cannot access the file because it is being used by another process: 'C:/Users/runneradmin/AppData/Local/Temp/pytest-of-runneradmin/pytest-0/test_async_checkpoint_plugin0/epoch=0-step=1.ckpt'
C:\hostedtoolcache\windows\Python\3.9.13\x64\lib\site-packages\fsspec\implementations\local.py:153: PermissionError
```
### Expected behavior
It is an issue I was concerned about in https://github.com/Lightning-AI/lightning/issues/11561#issuecomment-1187157230. A top-k deletion could happen any time. If the thread is trying to write the same file, there is a race condition.
### Environment
lightning 1.8.0dev
pytorch 1.10
### Additional context
cc @awaelchli @ananthsub @ninginthecloud @rohitgr7 @otaj @borda @akihironitta
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Start with test_async_checkpoint_plugin in tests/tests_pytorch/plugins/test_checkpoint_io_plugin.py and trace the checkpoint removal path through callbacks/model_checkpoint.py, strategies/strategy.py, plugins/io/wrapper.py, and plugins/io/torch_plugin.py. Reproduce the Windows failure from the linked CI run and inspect how asynchronous writes overlap with top-k deletion. Done means the test no longer flakes and a checkpoint cannot be removed while it is still being written.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- backend, testing
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 35/100