Fix deterministic checkpoint file sharding

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]
 

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"),
+    ]