Lightning-AI / Lightning-AI/litData

Checkpoint uploads truncate output prefix for stable checkpoint filenames

Open Beginner friendly
#888 0 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

bug help wanted
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

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.