diff --git a/examples/update.py b/examples/update.py index a792280..f03310e 100644 --- a/examples/update.py +++ b/examples/update.py @@ -49,10 +49,10 @@ def check_vllm_ready(endpoint: str, inference_parallel_size: int, uds: str | Non def split_checkpoint_files(checkpoint_path: str, rank: int, world_size: int) -> list[str]: - checkpoint_files = [ + checkpoint_files = sorted( os.path.join(checkpoint_path, f) for f in filter(lambda x: x.endswith(".safetensors"), os.listdir(checkpoint_path)) - ] + ) files_per_rank = (len(checkpoint_files) + world_size - 1) // world_size return checkpoint_files[rank * files_per_rank : (rank + 1) * files_per_rank] diff --git a/tests/test_example_update.py b/tests/test_example_update.py new file mode 100644 index 0000000..fa10071 --- /dev/null +++ b/tests/test_example_update.py @@ -0,0 +1,30 @@ +import runpy +from pathlib import Path +from unittest.mock import patch + + +split_checkpoint_files = runpy.run_path(str(Path(__file__).parents[1] / "examples" / "update.py"))[ + "split_checkpoint_files" +] + + +def test_split_checkpoint_files_is_deterministic_across_directory_order(tmp_path: Path) -> None: + filenames = [ + "model-00003.safetensors", + "notes.txt", + "model-00001.safetensors", + "model-00002.safetensors", + ] + + with patch("os.listdir", return_value=filenames): + forward = [split_checkpoint_files(str(tmp_path), rank, 2) for rank in range(2)] + + with patch("os.listdir", return_value=list(reversed(filenames))): + reverse = [split_checkpoint_files(str(tmp_path), rank, 2) for rank in range(2)] + + assert forward == reverse + assert [path for shard in forward for path in shard] == [ + str(tmp_path / "model-00001.safetensors"), + str(tmp_path / "model-00002.safetensors"), + str(tmp_path / "model-00003.safetensors"), + ]