Lightning-AI / Lightning-AI/litData
Checkpoint uploads truncate output prefix for stable checkpoint filenames
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 614
- Forks
- 106
- Avg merge
- 15h 8m
- Merged PRs (30d)
- 22
Description
## 🐛 Bug
When `optimize(..., use_checkpoint=True)` writes checkpoints to remote storage, checkpoint destination paths can be truncated and moved outside the requested output prefix.
LitData 0.2.71 writes stable checkpoint filenames such as `checkpoint-0.json`. However, `remove_uuid_from_filename()` still assumes every file under `.checkpoints` ends in `-<32-character UUID>.json` and unconditionally removes the final 38 characters. For example:
```text
s3://bucket/output/data/train/.checkpoints/checkpoint-0.json
-> s3://bucket/output/dat.json
```
With bucket versioning enabled, repeated checkpoint updates create many versions of the same malformed object.
### To Reproduce
1. Install LitData 0.2.71, or use current `main` at commit `a847d57`.
2. Run the code sample below. No S3 credentials or network access are required because it only exercises the destination-path calculation used by the uploader.
3. Observe that both paths lose the `.checkpoints` directory, the checkpoint filename, and part of the output prefix.
Code sample
```python
from litdata.processing.data_processor import _upload_dest
from litdata.streaming.cache import Dir
local_checkpoint = "/cache/.checkpoints/checkpoint-0.json"
local = _upload_dest(
Dir(path="/output/data/train", url=None),
local_checkpoint,
None,
)
remote = _upload_dest(
Dir(path=None, url="s3://bucket/output/data/train"),
local_checkpoint,
None,
)
print(local)
print(remote)
```
Output:
```text
/output/dat.json
s3://bucket/output/dat.json
```
Alternatively, you can share a fully reproducible [Lightning Studio](https://lightning.ai/studios) environment:
> A simple guide on how to create such a studio can be found [here](https://www.youtube.com/watch?v=YcW-2Zt_bFg&ab_channel=LightningAI).
1. Create a [Studio](https://lightning.ai/studios).
2. Reproduce the issue in the Studio.
3. [Publish the Studio](https://lightning.ai/docs/overview/studios/publishing#how-to-publish).
4. Paste the Studio link here.
### Expected behavior
Checkpoint files without a UUID suffix should remain unchanged. The destinations in the reproduction should be:
```text
/output/data/train/.checkpoints/checkpoint-0.json
s3://bucket/output/data/train/.checkpoints/checkpoint-0.json
```
Legacy checkpoint filenames that actually contain a UUID, such as `checkpoint-0-9fe2c4e93f654fdbb24c02b15259716c.json`, should still be normalized to `checkpoint-0.json`.
### Additional context
The stable `checkpoint-{rank}.json` naming was introduced recently, while the legacy UUID-removal logic remained in the shared upload destination path. Restricting UUID removal to filenames that actually match `checkpoint--<32 hex characters>.json` fixes the corruption while preserving the previous behavior.
In a real S3 run, an output prefix ending in `htr-lines/ijsberg/train` produced objects named `htr-lines/ijsber.json`. Other dataset prefixes produced similarly truncated names, with multiple object versions created during a single optimization run.
Environment detail
- PyTorch Version: 2.10.0+cu128
- OS: Linux 6.8.0-137-generic x86_64, glibc 2.39
- How you installed PyTorch: `uv` from PyPI
- Build command you used: N/A
- Python version: 3.12.11
- CUDA/cuDNN version: CUDA 12.8; cuDNN not relevant
- GPU models and configuration: Not relevant; the bug occurs during path construction
- Any other relevant information: LitData 0.2.71; also reproducible on `main` at `a847d57`
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 in litdata/processing/data_processor.py at _upload_dest() and trace how remove_uuid_from_filename() transforms checkpoint paths; the reproduction uses Dir from litdata/streaming/cache.py. Run the supplied local and remote path calculations first, then verify that stable checkpoint names remain unchanged while legacy UUID-suffixed names normalize to checkpoint-{rank}.json.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- aws, python
- Domain
- cloud, data
- Issue type
- Bug
- Difficulty
- 2/5
- Estimated time
- 1-3 hours
- Activity status
- Quiet
- Clarity
- Clearly specified
- Newbie friendliness
- 82/100