75 lines
2.2 KiB
Python
75 lines
2.2 KiB
Python
#!/usr/bin/env python3
|
|
"""Split unfinished JSONL records into balanced temporary retry shards."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
from pathlib import Path
|
|
|
|
|
|
def record_id(value: dict) -> str:
|
|
"""Return the stable trajectory identifier used by audit manifests."""
|
|
|
|
sample_id = value.get("trajectory_id") or value.get("sample_id")
|
|
if not isinstance(sample_id, str) or not sample_id:
|
|
raise ValueError("JSONL record has no trajectory_id/sample_id")
|
|
return sample_id
|
|
|
|
|
|
def main() -> None:
|
|
"""Write unfinished input records round-robin across the requested shards."""
|
|
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("input", type=Path)
|
|
parser.add_argument("completed", type=Path)
|
|
parser.add_argument("output_dir", type=Path)
|
|
parser.add_argument("--shards", type=int, default=3)
|
|
args = parser.parse_args()
|
|
if args.shards < 1:
|
|
raise SystemExit("--shards must be positive")
|
|
|
|
completed_ids: set[str] = set()
|
|
with args.completed.open(encoding="utf-8") as handle:
|
|
for line in handle:
|
|
if line.strip():
|
|
completed_ids.add(record_id(json.loads(line)))
|
|
|
|
args.output_dir.mkdir(parents=True, exist_ok=True)
|
|
handles = [
|
|
(args.output_dir / f"part_{index}.input.jsonl").open("w", encoding="utf-8")
|
|
for index in range(args.shards)
|
|
]
|
|
counts = [0] * args.shards
|
|
try:
|
|
unfinished = 0
|
|
with args.input.open(encoding="utf-8") as handle:
|
|
for line in handle:
|
|
if not line.strip():
|
|
continue
|
|
value = json.loads(line)
|
|
if record_id(value) in completed_ids:
|
|
continue
|
|
shard = unfinished % args.shards
|
|
handles[shard].write(line)
|
|
counts[shard] += 1
|
|
unfinished += 1
|
|
finally:
|
|
for handle in handles:
|
|
handle.close()
|
|
|
|
print(
|
|
json.dumps(
|
|
{
|
|
"completed": len(completed_ids),
|
|
"unfinished": sum(counts),
|
|
"shard_counts": counts,
|
|
},
|
|
indent=2,
|
|
)
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|