diff --git a/src/litdata/processing/data_processor.py b/src/litdata/processing/data_processor.py index a6c35f4f..dcf1684a 100644 --- a/src/litdata/processing/data_processor.py +++ b/src/litdata/processing/data_processor.py @@ -1018,7 +1018,16 @@ def _done(self, size: int | None, delete_cached_files: bool, output_dir: Dir) -> merge_cache = Cache(cache_dir, chunk_bytes=1) node_rank = _get_node_rank() - merge_cache._merge_no_wait(node_rank if num_nodes > 1 else None, getattr(self, "existing_index", None)) + # With more than one node, every node would otherwise fold the existing + # index into its own {node_rank}-index.json, and the final cross-node + # merge below then repeats those chunks once per node. Add the existing + # index only once, either here for the single-node case or in the final + # merge for the multi-node case. + existing_index = getattr(self, "existing_index", None) + merge_cache._merge_no_wait( + node_rank if num_nodes > 1 else None, + existing_index if num_nodes == 1 else None, + ) self._upload_index(output_dir, cache_dir, num_nodes, node_rank) @@ -1093,7 +1102,9 @@ def _upload_index(self, output_dir: Dir, cache_dir: str, num_nodes: int, node_ra shutil.copyfile(remote_filepath, node_index_filepath) merge_cache = Cache(cache_dir, chunk_bytes=1) - merge_cache._merge_no_wait() + # The per-node index files hold only their own new chunks now, so the + # existing index is folded in here, once, as the node files are merged. + merge_cache._merge_no_wait(existing_index=getattr(self, "existing_index", None)) self._upload_index(output_dir, cache_dir, 1, None) diff --git a/tests/streaming/test_writer.py b/tests/streaming/test_writer.py index 4aeaef64..ad1df57c 100644 --- a/tests/streaming/test_writer.py +++ b/tests/streaming/test_writer.py @@ -327,3 +327,58 @@ def test_merge_natural_sort_order_with_many_workers(tmpdir): filenames = [c["filename"] for c in data["chunks"]] assert filenames == [f"chunk-{i}-0.bin" for i in range(n_workers)] + + +def test_merge_multinode_append_adds_existing_index_once(tmpdir): + # Regression for #865: with more than one node, optimize(mode="append") must + # not repeat the pre-existing chunks once per node. Folding the existing + # index into every node's index (the old behaviour) duplicates it at the + # final cross-node merge; folding it in a single time at that final merge + # keeps it once. This mirrors what the data processor now does. + import shutil + + from litdata.constants import _INDEX_FILENAME + + config = { + "chunk_bytes": None, + "chunk_size": 1, + "compression": None, + "data_format": ["scalar"], + "data_spec": None, + "encryption": None, + "item_loader": "PyTreeLoader", + } + + def make_chunk(name): + return {"chunk_size": 1, "column_sizes": [4], "dim": None, "filename": name} + + existing_index = {"chunks": [make_chunk("A.bin"), make_chunk("B.bin")], "config": config} + + def per_node_index(node_rank, new_chunk, include_existing): + node_dir = tmpdir.mkdir(f"node_{node_rank}_{'old' if include_existing else 'new'}") + with open(os.path.join(str(node_dir), f"0.{_INDEX_FILENAME}"), "w") as f: + json.dump({"chunks": [make_chunk(new_chunk)], "config": config}, f, sort_keys=True) + writer = BinaryWriter(str(node_dir), chunk_size=1) + writer._is_done = True + writer._rank = 0 + writer._merge_no_wait(node_rank=node_rank, existing_index=existing_index if include_existing else None) + return os.path.join(str(node_dir), f"{node_rank}-{_INDEX_FILENAME}") + + def final_merge(node_files, existing): + final_dir = tmpdir.mkdir(f"final_{'old' if existing else 'new'}") + for src in node_files: + shutil.copyfile(src, os.path.join(str(final_dir), os.path.basename(src))) + writer = BinaryWriter(str(final_dir), chunk_size=1) + writer._is_done = True + writer._rank = 0 + writer._merge_no_wait(existing_index=existing) + with open(os.path.join(str(final_dir), _INDEX_FILENAME)) as f: + return [c["filename"] for c in json.load(f)["chunks"]] + + # Old behaviour: existing index folded in per node, none at the final merge. + old_files = [per_node_index(0, "C.bin", True), per_node_index(1, "D.bin", True)] + assert final_merge(old_files, None) == ["A.bin", "B.bin", "C.bin", "A.bin", "B.bin", "D.bin"] + + # Fixed behaviour: existing index folded in once, at the final merge. + new_files = [per_node_index(0, "C.bin", False), per_node_index(1, "D.bin", False)] + assert final_merge(new_files, existing_index) == ["A.bin", "B.bin", "C.bin", "D.bin"]