Buckets:
| """Safely repartition an interrupted task-subset evaluation for resume.""" | |
| from __future__ import annotations | |
| import argparse | |
| import hashlib | |
| import json | |
| import os | |
| import shutil | |
| import tempfile | |
| import time | |
| from collections import defaultdict | |
| from datetime import UTC, datetime | |
| from pathlib import Path | |
| from .evaluate_task_subset import _manifest_rows | |
| REQUIRED_SAMPLE_FILES = ( | |
| "target.txt", | |
| "output.txt", | |
| "result.json", | |
| "generated_codes.pt", | |
| ) | |
| def _read_jsonl(path: Path) -> list[dict]: | |
| rows = [] | |
| for line_number, line in enumerate(path.read_text(encoding="utf-8").splitlines(), 1): | |
| if not line: | |
| continue | |
| try: | |
| rows.append(json.loads(line)) | |
| except json.JSONDecodeError as error: | |
| raise ValueError(f"invalid JSON at {path}:{line_number}") from error | |
| return rows | |
| def _write_jsonl(path: Path, rows: list[dict]) -> None: | |
| with path.open("w", encoding="utf-8") as handle: | |
| for row in rows: | |
| handle.write(json.dumps(row, ensure_ascii=False) + "\n") | |
| def _rank_dirs(shards_dir: Path) -> list[Path]: | |
| rank_dirs = sorted(path for path in shards_dir.glob("rank-*") if path.is_dir()) | |
| expected = [f"rank-{rank:02d}" for rank in range(len(rank_dirs))] | |
| actual = [path.name for path in rank_dirs] | |
| if actual != expected: | |
| raise ValueError(f"rank directories must be contiguous: {actual}") | |
| return rank_dirs | |
| def _prediction_snapshot(shards_dir: Path) -> dict[str, tuple[int, int]]: | |
| return { | |
| str(path.relative_to(shards_dir)): (path.stat().st_size, path.stat().st_mtime_ns) | |
| for path in sorted(shards_dir.glob("rank-*/predictions.jsonl")) | |
| } | |
| def _assert_stable(shards_dir: Path, seconds: float) -> None: | |
| if seconds <= 0: | |
| return | |
| before = _prediction_snapshot(shards_dir) | |
| time.sleep(seconds) | |
| after = _prediction_snapshot(shards_dir) | |
| if before != after: | |
| raise RuntimeError("prediction files changed during the stability check") | |
| def _sample_fingerprint(sample_dir: Path) -> tuple[tuple[str, str], ...]: | |
| digests = [] | |
| for filename in REQUIRED_SAMPLE_FILES: | |
| digest = hashlib.sha256() | |
| with (sample_dir / filename).open("rb") as handle: | |
| for chunk in iter(lambda: handle.read(1024 * 1024), b""): | |
| digest.update(chunk) | |
| digests.append((filename, digest.hexdigest())) | |
| return tuple(digests) | |
| def _load_committed( | |
| shards_dir: Path, | |
| expected_by_id: dict[str, dict], | |
| *, | |
| deduplicate_identical: bool = False, | |
| ) -> tuple[list[tuple[dict, Path]], int]: | |
| committed_by_id: dict[str, tuple[dict, Path, tuple[tuple[str, str], ...]]] = {} | |
| sample_id_by_index = {} | |
| duplicate_rows = 0 | |
| for rank_dir in _rank_dirs(shards_dir): | |
| predictions_path = rank_dir / "predictions.jsonl" | |
| if not predictions_path.exists(): | |
| raise FileNotFoundError(predictions_path) | |
| for row in _read_jsonl(predictions_path): | |
| sample_id = row.get("sample_id") | |
| index = row.get("evaluation_index") | |
| expected = expected_by_id.get(sample_id) | |
| if expected is None or index != expected["evaluation_index"]: | |
| raise ValueError( | |
| f"manifest mismatch for {sample_id}: committed={index}, " | |
| f"expected={None if expected is None else expected['evaluation_index']}" | |
| ) | |
| sample_dir = rank_dir / "samples" / sample_id | |
| for filename in REQUIRED_SAMPLE_FILES: | |
| path = sample_dir / filename | |
| if not path.is_file(): | |
| raise FileNotFoundError(path) | |
| result = json.loads((sample_dir / "result.json").read_text(encoding="utf-8")) | |
| if ( | |
| result.get("sample_id") != sample_id | |
| or result.get("evaluation_index") != index | |
| ): | |
| raise ValueError(f"result provenance mismatch for {sample_id}") | |
| fingerprint = _sample_fingerprint(sample_dir) | |
| previous = committed_by_id.get(sample_id) | |
| if previous is not None: | |
| if not deduplicate_identical: | |
| raise ValueError(f"duplicate committed sample_id: {sample_id}") | |
| previous_row, _, previous_fingerprint = previous | |
| if row != previous_row or fingerprint != previous_fingerprint: | |
| raise ValueError(f"non-identical duplicate commit: {sample_id}") | |
| duplicate_rows += 1 | |
| continue | |
| previous_id = sample_id_by_index.get(index) | |
| if previous_id is not None: | |
| raise ValueError( | |
| f"duplicate committed evaluation_index {index}: " | |
| f"{previous_id} and {sample_id}" | |
| ) | |
| committed_by_id[sample_id] = (row, sample_dir, fingerprint) | |
| sample_id_by_index[index] = sample_id | |
| committed = [ | |
| (row, sample_dir) for row, sample_dir, _ in committed_by_id.values() | |
| ] | |
| return sorted(committed, key=lambda pair: pair[0]["evaluation_index"]), duplicate_rows | |
| def _validate_staged(shards_dir: Path, new_world_size: int, expected_count: int) -> None: | |
| seen_ids = set() | |
| seen_indices = set() | |
| for rank, rank_dir in enumerate(_rank_dirs(shards_dir)): | |
| rows = _read_jsonl(rank_dir / "predictions.jsonl") | |
| indices = [row["evaluation_index"] for row in rows] | |
| if indices != sorted(indices): | |
| raise ValueError(f"rank {rank} predictions are not sorted") | |
| for row in rows: | |
| sample_id = row["sample_id"] | |
| index = row["evaluation_index"] | |
| if index % new_world_size != rank: | |
| raise ValueError(f"index {index} assigned to the wrong rank {rank}") | |
| if sample_id in seen_ids or index in seen_indices: | |
| raise ValueError(f"duplicate staged sample/index: {sample_id}/{index}") | |
| for filename in REQUIRED_SAMPLE_FILES: | |
| if not (rank_dir / "samples" / sample_id / filename).is_file(): | |
| raise FileNotFoundError(rank_dir / "samples" / sample_id / filename) | |
| seen_ids.add(sample_id) | |
| seen_indices.add(index) | |
| if len(seen_ids) != expected_count: | |
| raise ValueError(f"staged {len(seen_ids)} samples, expected {expected_count}") | |
| def reshard_run( | |
| run_dir: Path, | |
| cache_root: Path, | |
| *, | |
| selection: str, | |
| new_world_size: int, | |
| stability_check_seconds: float = 2.0, | |
| deduplicate_identical: bool = False, | |
| ) -> dict: | |
| """Repartition committed rows and atomically replace ``run_dir/shards``. | |
| The original shard tree is retained under ``resume-backups``. Sample | |
| directories without a committed JSONL row remain only in that backup and | |
| will be recomputed by the resumed evaluator. | |
| """ | |
| run_dir = run_dir.resolve() | |
| cache_root = cache_root.resolve() | |
| if new_world_size <= 0: | |
| raise ValueError("new_world_size must be positive") | |
| if (run_dir / "summary.json").exists(): | |
| raise ValueError("refusing to reshard a completed evaluation") | |
| shards_dir = run_dir / "shards" | |
| if not shards_dir.is_dir(): | |
| raise FileNotFoundError(shards_dir) | |
| old_world_size = len(_rank_dirs(shards_dir)) | |
| if old_world_size == new_world_size and not deduplicate_identical: | |
| raise ValueError(f"evaluation already has {new_world_size} rank shards") | |
| _assert_stable(shards_dir, stability_check_seconds) | |
| expected_rows = _manifest_rows(cache_root, selection) | |
| expected_by_id = {row["sample_id"]: row for row in expected_rows} | |
| committed, duplicate_rows = _load_committed( | |
| shards_dir, | |
| expected_by_id, | |
| deduplicate_identical=deduplicate_identical, | |
| ) | |
| if old_world_size == new_world_size and duplicate_rows == 0: | |
| raise ValueError("same-world repair found no identical duplicate commits") | |
| stage_dir = Path(tempfile.mkdtemp(prefix=".reshard-staging-", dir=run_dir)) | |
| try: | |
| grouped: dict[int, list[dict]] = defaultdict(list) | |
| for rank in range(new_world_size): | |
| rank_dir = stage_dir / f"rank-{rank:02d}" | |
| (rank_dir / "samples").mkdir(parents=True) | |
| for row, source_dir in committed: | |
| rank = row["evaluation_index"] % new_world_size | |
| grouped[rank].append(row) | |
| target_dir = stage_dir / f"rank-{rank:02d}" / "samples" / row["sample_id"] | |
| shutil.copytree(source_dir, target_dir) | |
| for rank in range(new_world_size): | |
| rows = sorted(grouped[rank], key=lambda row: row["evaluation_index"]) | |
| _write_jsonl(stage_dir / f"rank-{rank:02d}" / "predictions.jsonl", rows) | |
| _validate_staged(stage_dir, new_world_size, len(committed)) | |
| stamp = datetime.now(UTC).strftime("%Y%m%d-%H%M%S") | |
| backup_dir = run_dir / "resume-backups" / f"{stamp}_world{old_world_size}" | |
| if backup_dir.exists(): | |
| raise FileExistsError(backup_dir) | |
| backup_dir.mkdir(parents=True) | |
| backup_shards = backup_dir / "shards" | |
| os.replace(shards_dir, backup_shards) | |
| try: | |
| os.replace(stage_dir, shards_dir) | |
| except BaseException: | |
| os.replace(backup_shards, shards_dir) | |
| raise | |
| metadata = { | |
| "timestamp_utc": datetime.now(UTC).isoformat(), | |
| "old_world_size": old_world_size, | |
| "new_world_size": new_world_size, | |
| "committed_samples": len(committed), | |
| "duplicate_rows_removed": duplicate_rows, | |
| "total_expected_samples": len(expected_rows), | |
| "selection": selection, | |
| "cache_root": str(cache_root), | |
| "backup_dir": str(backup_dir), | |
| "per_rank_committed": { | |
| str(rank): len(grouped[rank]) for rank in range(new_world_size) | |
| }, | |
| } | |
| (backup_dir / "reshard.json").write_text( | |
| json.dumps(metadata, ensure_ascii=False, indent=2) + "\n", | |
| encoding="utf-8", | |
| ) | |
| with (run_dir / "reshard_history.jsonl").open("a", encoding="utf-8") as handle: | |
| handle.write(json.dumps(metadata, ensure_ascii=False) + "\n") | |
| _validate_staged(shards_dir, new_world_size, len(committed)) | |
| return metadata | |
| finally: | |
| if stage_dir.exists(): | |
| shutil.rmtree(stage_dir) | |
| def build_parser() -> argparse.ArgumentParser: | |
| parser = argparse.ArgumentParser(description=__doc__) | |
| parser.add_argument("--run-dir", type=Path, required=True) | |
| parser.add_argument("--cache-root", type=Path, required=True) | |
| parser.add_argument("--selection", choices=["indexed", "test"], required=True) | |
| parser.add_argument("--new-world-size", type=int, required=True) | |
| parser.add_argument("--stability-check-seconds", type=float, default=2.0) | |
| parser.add_argument( | |
| "--deduplicate-identical", | |
| action="store_true", | |
| help="remove byte-identical duplicate commits, including same-world repairs", | |
| ) | |
| return parser | |
| def main() -> None: | |
| args = build_parser().parse_args() | |
| metadata = reshard_run( | |
| args.run_dir, | |
| args.cache_root, | |
| selection=args.selection, | |
| new_world_size=args.new_world_size, | |
| stability_check_seconds=args.stability_check_seconds, | |
| deduplicate_identical=args.deduplicate_identical, | |
| ) | |
| print("RESHARD_COMPLETE " + json.dumps(metadata, ensure_ascii=False), flush=True) | |
| if __name__ == "__main__": | |
| main() | |
Xet Storage Details
- Size:
- 11.5 kB
- Xet hash:
- 05a7702286c6e672fab546fc5323771cf104f06801ef533005effcd32f3c4964
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.