fix(processing): preserve checkpoint paths without UUIDs

src/litdata/processing/utilities.py

@@ -14,6 +14,7 @@
 import io
 import json
 import os
+import re
 import tempfile
 import urllib
 from collections.abc import Callable
@@ -202,8 +203,7 @@ def remove_uuid_from_filename(filepath: str) -> str:
     if not filepath.__contains__(".checkpoints"):
         return filepath
 
-    # uuid is of 32 characters, '.json' is 5 characters and '-' is 1 character
-    return filepath[:-38] + ".json"
+    return re.sub(r"(checkpoint-\d+)-[0-9a-fA-F]{32}\.json$", r"\1.json", filepath)
 
 
 def construct_storage_options(storage_options: dict[str, Any], input_dir: Dir) -> dict[str, Any]:

tests/processing/test_utilities.py

@@ -1,6 +1,8 @@
 import json
 from unittest.mock import MagicMock
 
+import pytest
+
 from litdata.processing import utilities as utilities_module
 from litdata.processing.utilities import (
     extract_rank_and_index_from_filename,
@@ -98,25 +100,22 @@ def _fn(remote_path, local_path):
 
 
 def test_remove_uuid_from_filename():
-    filepaths = [
-        "checkpoint-0-9fe2c4e93f654fdbb24c02b15259716c.json",
-        "checkpoint-1-9fe2c4e93f654fdbb24c02b15259716c.json",
-        "checkpoint-2-9fe2c4e93f654fdbb24c02b15259716c.json",
-        "checkpoint-101-9fe2c4e93f654fdbb24c02b15259716c.json",
-        "checkpoint-12-9fe2c4e93f654fdbb24c02b15259716c.json",
-        "checkpoint-267-9fe2c4e93f654fdbb24c02b15259716c.json",
-    ]
-
-    expected = [
-        "checkpoint-0.json",
-        "checkpoint-1.json",
-        "checkpoint-2.json",
-        "checkpoint-101.json",
-        "checkpoint-12.json",
-        "checkpoint-267.json",
-    ]
-
-    for idx, filepath in enumerate(filepaths):
-        filepath = ".checkpoints/" + filepath
-        result = remove_uuid_from_filename(filepath)
-        assert result == ".checkpoints/" + expected[idx]
+    checkpoint_dir = "output/data/train/.checkpoints"
+    uuid = "9fe2c4e93f654fdbb24c02b15259716c"
+
+    for rank in (0, 1, 2, 12, 101, 267):
+        filepath = f"{checkpoint_dir}/checkpoint-{rank}-{uuid}.json"
+        assert remove_uuid_from_filename(filepath) == f"{checkpoint_dir}/checkpoint-{rank}.json"
+
+
+@pytest.mark.parametrize(
+    "filepath",
+    [
+        "output/data/train/.checkpoints/checkpoint-0.json",
+        "input/data/val/.checkpoints/config.json",
+        "output/data/train/.checkpoints/checkpoint-0-not-a-uuid.json",
+        "output/data/train/checkpoint-0-9fe2c4e93f654fdbb24c02b15259716c.json",
+    ],
+)
+def test_remove_uuid_from_filename_leaves_other_paths_unchanged(filepath):
+    assert remove_uuid_from_filename(filepath) == filepath