diff --git a/src/litdata/processing/utilities.py b/src/litdata/processing/utilities.py index db099d1be..654913fb4 100644 --- a/src/litdata/processing/utilities.py +++ b/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]: diff --git a/tests/processing/test_utilities.py b/tests/processing/test_utilities.py index 31c0d1ca3..3656f188f 100644 --- a/tests/processing/test_utilities.py +++ b/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