Buckets:
| import json | |
| from pathlib import Path | |
| from zebra_sft.reshard_task_subset_eval import reshard_run | |
| def _write_jsonl(path: Path, rows: list[dict]) -> None: | |
| path.write_text("".join(json.dumps(row) + "\n" for row in rows), encoding="utf-8") | |
| def _write_sample(rank_dir: Path, row: dict) -> None: | |
| sample_dir = rank_dir / "samples" / row["sample_id"] | |
| sample_dir.mkdir(parents=True) | |
| (sample_dir / "target.txt").write_text("target", encoding="utf-8") | |
| (sample_dir / "output.txt").write_text("output", encoding="utf-8") | |
| (sample_dir / "result.json").write_text(json.dumps(row), encoding="utf-8") | |
| (sample_dir / "generated_codes.pt").write_bytes(b"codes") | |
| def test_reshard_interrupted_two_rank_evaluation_to_four(tmp_path): | |
| cache_root = tmp_path / "cache" | |
| cache_root.mkdir() | |
| manifest = [ | |
| {"sample_id": f"sample-{index}", "split": "test"} for index in range(8) | |
| ] | |
| _write_jsonl(cache_root / "manifest.jsonl", manifest) | |
| run_dir = tmp_path / "run" | |
| committed = [ | |
| {"sample_id": f"sample-{index}", "evaluation_index": index} | |
| for index in range(6) | |
| ] | |
| for old_rank in range(2): | |
| rank_dir = run_dir / "shards" / f"rank-{old_rank:02d}" | |
| rows = [row for row in committed if row["evaluation_index"] % 2 == old_rank] | |
| for row in rows: | |
| _write_sample(rank_dir, row) | |
| _write_jsonl(rank_dir / "predictions.jsonl", rows) | |
| orphan = run_dir / "shards" / "rank-00" / "samples" / "orphan" | |
| orphan.mkdir(parents=True) | |
| (orphan / "result.json").write_text("{}", encoding="utf-8") | |
| metadata = reshard_run( | |
| run_dir, | |
| cache_root, | |
| selection="test", | |
| new_world_size=4, | |
| stability_check_seconds=0, | |
| ) | |
| assert metadata["committed_samples"] == 6 | |
| assert metadata["per_rank_committed"] == {"0": 2, "1": 2, "2": 1, "3": 1} | |
| active_rows = [] | |
| for rank in range(4): | |
| rank_dir = run_dir / "shards" / f"rank-{rank:02d}" | |
| rows = [json.loads(line) for line in (rank_dir / "predictions.jsonl").read_text().splitlines()] | |
| assert all(row["evaluation_index"] % 4 == rank for row in rows) | |
| active_rows.extend(rows) | |
| assert sorted(row["evaluation_index"] for row in active_rows) == list(range(6)) | |
| assert not (run_dir / "shards" / "rank-00" / "samples" / "orphan").exists() | |
| backup = next((run_dir / "resume-backups").glob("*_world2")) | |
| assert (backup / "shards" / "rank-00" / "samples" / "orphan").is_dir() | |
| def test_reshard_same_world_deduplicates_only_identical_commits(tmp_path): | |
| cache_root = tmp_path / "cache" | |
| cache_root.mkdir() | |
| manifest = [ | |
| {"sample_id": f"sample-{index}", "split": "test"} for index in range(5) | |
| ] | |
| _write_jsonl(cache_root / "manifest.jsonl", manifest) | |
| run_dir = tmp_path / "run" | |
| rows = [ | |
| {"sample_id": f"sample-{index}", "evaluation_index": index} | |
| for index in range(5) | |
| ] | |
| rank_0 = run_dir / "shards" / "rank-00" | |
| rank_1 = run_dir / "shards" / "rank-01" | |
| for row in (rows[0], rows[2], rows[4]): | |
| _write_sample(rank_0, row) | |
| for row in (rows[1], rows[2], rows[3]): | |
| _write_sample(rank_1, row) | |
| _write_jsonl(rank_0 / "predictions.jsonl", [rows[0], rows[2], rows[2], rows[4]]) | |
| _write_jsonl(rank_1 / "predictions.jsonl", [rows[1], rows[2], rows[3]]) | |
| metadata = reshard_run( | |
| run_dir, | |
| cache_root, | |
| selection="test", | |
| new_world_size=2, | |
| stability_check_seconds=0, | |
| deduplicate_identical=True, | |
| ) | |
| assert metadata["committed_samples"] == 5 | |
| assert metadata["duplicate_rows_removed"] == 2 | |
| active_rows = [] | |
| for rank in range(2): | |
| rank_dir = run_dir / "shards" / f"rank-{rank:02d}" | |
| active_rows.extend( | |
| json.loads(line) | |
| for line in (rank_dir / "predictions.jsonl").read_text().splitlines() | |
| ) | |
| assert sorted(row["evaluation_index"] for row in active_rows) == list(range(5)) | |
| backup = next((run_dir / "resume-backups").glob("*_world2")) | |
| assert len((backup / "shards" / "rank-00" / "predictions.jsonl").read_text().splitlines()) == 4 | |
Xet Storage Details
- Size:
- 4.14 kB
- Xet hash:
- 3013c622aa3a23be54d13d543802d7b4f2c965a32bdefeaa9b13302ef929a2e3
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.