diff --git a/collectionparser/api.py b/collectionparser/api.py index ce7e726..52e23f4 100644 --- a/collectionparser/api.py +++ b/collectionparser/api.py @@ -11,8 +11,10 @@ from tqdm import tqdm from collectionparser.utils import create_dir_for_file, find_files # ponytail: buckets under this size sort byte-exactly in RAM; larger buckets -# shell out to `sort -u`, which spills to disk. -_IN_MEMORY_SORT_LIMIT = 8 * 1024 * 1024 +# shell out to `sort -u`. 64 MB covers almost every bucket, so the final dedup +# rarely spawns a process; a single bucket's sort is one file, so memory stays +# bounded. Lower this if bucket files can be huge. +_IN_MEMORY_SORT_LIMIT = 64 * 1024 * 1024 class API: @@ -80,9 +82,13 @@ def sort_unique(src, dst) -> None: Path(dst).write_bytes(b"\n".join(lines) + b"\n" if lines else b"") else: tmp = f"{dst}.tmp" + # Parallel sort, but memory-bounded: `-S` caps the sort buffer and + # `--parallel 4` bounds the threads. A large `-S` here can OOM and get + # the sort SIGKILLed. ponytail: ceiling ~1 GB + per-thread buffers; + # raise `-S`/`--parallel` only if the box has RAM to spare. with open(tmp, "wb") as writer: subprocess.run( - ["sort", "-u", str(src)], + ["sort", "-S", "1G", "--parallel", "4", "-u", str(src)], stdout=writer, check=True, env={**os.environ, "LC_ALL": "C"}, diff --git a/collectionparser/build.py b/collectionparser/build.py index f3de4da..5dc8a14 100644 --- a/collectionparser/build.py +++ b/collectionparser/build.py @@ -246,8 +246,14 @@ def build( dataset.mkdir(parents=True, exist_ok=True) scratch.mkdir(parents=True, exist_ok=True) os.environ["TMPDIR"] = str(scratch) # `sort -u` spills here, not into the dataset - for stale in scratch.glob("extract_*"): - shutil.rmtree(stale, ignore_errors=True) + # Leftovers from a crashed run: extraction dirs, and `sort` run/temp files. + # Only scratch is touched, never the source or dataset. + for pattern in ("extract_*", "sort*", "*.tmp"): + for stale in scratch.glob(pattern): + if stale.is_dir(): + shutil.rmtree(stale, ignore_errors=True) + else: + stale.unlink(missing_ok=True) if extract and not _which("7z", "7za", "unrar"): print( @@ -316,17 +322,27 @@ def build( pending = [f for f in inputs if str(f) not in parsed] t0 = time.time() + pairs: int | None = None + rejects: int | None = None if pending: - parse_files( + pairs, rejects = parse_files( pending, str(success), str(failure), progress_filename=str(progress_path), hashes_filename=str(hashes), ) + # Remember the counts so a later resume skips the full re-read. + state.setdefault("counts", {})[str(index)] = [pairs, rejects] + save() if success.exists(): - pairs = _count_lines(success) - rejects = _count_lines(failure) if failure.exists() else 0 + if pairs is None: + saved = (state.get("counts") or {}).get(str(index)) + if saved: + pairs, rejects = saved + else: # ponytail: only pre-count-state runs need this full read + pairs = _count_lines(success) + rejects = _count_lines(failure) if failure.exists() else 0 sort_unique(success, success) split_file(str(success), str(dataset)) success.unlink(missing_ok=True)