[levanter] Fix single-shard cache consolidation on R2/S3 (#4436)

## Summary

- Skip TensorStore reopen for single-shard cache consolidation — copy
shard contents directly via fsspec instead
- Fixes the `Malformed StorageGeneration` error that breaks CoreWeave CI
tokenization on R2-backed S3
- Adds regression test for single-shard consolidation path

Fixes #4433

## Test plan

- [x] `uv run pytest -q tests/test_consolidate_metadata.py` — 4 passed
(including new single-shard test)
- [x] `uv run pytest -q
tests/processing/tokenize/test_tokenize.py::test_tokenize_full_pipeline_integration
-m slow` — 1 passed
- [x] `./infra/pre-commit.py --fix` — OK
- [ ] CoreWeave CI (`cw-ci-test`) passes on this branch

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Author: ahmeda14960

lib/levanter/src/levanter/store/cache.py

@@ -8,6 +8,7 @@
 import logging as pylogging
 import operator
 import os
+import shutil
 import threading
 import time
 from concurrent.futures import ThreadPoolExecutor
@@ -429,14 +430,20 @@ def consolidate_shard_caches(
 
     logger.info(f"Consolidating {len(shard_cache_paths)} shard caches into {output_path}")
 
+    shard_ledgers = [CacheLedger.load(p, metadata) for p in shard_cache_paths]
+
+    if len(shard_cache_paths) == 1:
+        shard_path = shard_cache_paths[0]
+        logger.info(f"Single shard cache detected; copying {shard_path} into {output_path} without reopening it.")
+        _copy_tree_contents(shard_path, output_path)
+        return _merge_ledgers(output_path, shard_cache_paths, shard_ledgers, metadata)
+
     first_cache = TreeStore.open(exemplar, shard_cache_paths[0], mode="r", cache_metadata=True)
     data_offset_tree = jax.tree.map(lambda x: 0, first_cache.tree)
 
     shard_info: list[dict] = []
     total_rows = 0
 
-    shard_ledgers = [CacheLedger.load(p, metadata) for p in shard_cache_paths]
-
     # Parallel: open each TreeStore to read data_size (dominates wall time on remote storage)
     def _get_data_sizes(shard_path):
         store = TreeStore.open(exemplar, shard_path, mode="r", cache_metadata=True)
@@ -488,6 +495,30 @@ def _copy_shard(info: dict):
     return final_ledger
 
 
+def _copy_tree_contents(source_path: str, dest_path: str) -> None:
+    src_fs, src_root = url_to_fs(source_path)
+    dest_fs, dest_root = url_to_fs(dest_path)
+
+    dest_fs.makedirs(dest_root, exist_ok=True)
+    entries = src_fs.find(src_root, withdirs=True, detail=True)
+
+    for src_file, info in entries.items():
+        if info.get("type") == "directory":
+            continue
+
+        rel_path = os.path.relpath(src_file, src_root)
+        dest_file = os.path.join(dest_root, rel_path)
+        dest_parent = os.path.dirname(dest_file)
+        if dest_parent:
+            dest_fs.makedirs(dest_parent, exist_ok=True)
+
+        if type(src_fs) is type(dest_fs):
+            src_fs.copy(src_file, dest_file)
+        else:
+            with src_fs.open(src_file, "rb") as src_f, dest_fs.open(dest_file, "wb") as dest_f:
+                shutil.copyfileobj(src_f, dest_f)
+
+
 def _merge_ledgers(
     output_path: str, shard_cache_paths: list[str], shard_ledgers: list[CacheLedger], metadata: CacheMetadata
 ) -> CacheLedger:

tests/test_consolidate_metadata.py

@@ -152,3 +152,22 @@ def test_consolidate_shard_caches_end_to_end():
         for i in range(NUM_SHARDS):
             row = merged[i * ROWS_PER_SHARD]
             assert row["input_ids"][0] == i, f"shard {i} data mismatch"
+
+
+def test_consolidate_shard_caches_single_shard():
+    with tempfile.TemporaryDirectory(prefix="levanter-test-consolidate-single-") as tmpdir:
+        shard_path = os.path.join(tmpdir, "shard_0")
+        rows = [{"input_ids": np.full((ROW_WIDTH,), 7, dtype=np.int32)} for _ in range(ROWS_PER_SHARD)]
+        _build_shard_cache(shard_path, EXEMPLAR_FLAT, rows)
+
+        dest_path = os.path.join(tmpdir, "merged")
+        ledger = consolidate_shard_caches([shard_path], dest_path, EXEMPLAR_FLAT, copy_max_workers=1)
+
+        assert ledger.total_num_rows == ROWS_PER_SHARD
+        assert ledger.is_finished
+        assert ledger.shard_rows == {os.path.basename(shard_path): ROWS_PER_SHARD}
+        assert ledger.finished_shards == [os.path.basename(shard_path)]
+
+        merged = TreeStore.open(EXEMPLAR_FLAT, dest_path, mode="r", cache_metadata=True)
+        assert len(merged) == ROWS_PER_SHARD
+        assert merged[0]["input_ids"][0] == 7