Checkpoint uploads truncate output prefix for stable checkpoint filenames

[TimKoornstra]

🐛 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:

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

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:

/output/dat.json
s3://bucket/output/dat.json

Alternatively, you can share a fully reproducible Lightning Studio environment:

A simple guide on how to create such a studio can be found here.

1. Create a Studio.
2. Reproduce the issue in the Studio.
3. Publish the Studio.
4. Paste the Studio link here.

Expected behavior

Checkpoint files without a UUID suffix should remain unchanged. The destinations in the reproduction should be:

/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-<rank>-<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