diff --git a/scripts/gen_full_supervisor.sh b/scripts/gen_full_supervisor.sh new file mode 100755 index 0000000..2f0256f --- /dev/null +++ b/scripts/gen_full_supervisor.sh @@ -0,0 +1,39 @@ +#!/usr/bin/env bash +# Crash-resilient full-scale generation supervisor. +# +# gen_synthetic.py now (a) streams each sample to disk the instant it's produced +# (O(workers) memory, not O(n)) and (b) supports --resume: it counts existing +# lines and only generates+appends the missing suffix. So a kill mid-run keeps +# all progress on disk; re-running --resume continues from where it stopped. +# +# This supervisor loops "run --resume -> check count" until the target line +# count is reached. If the supervisor itself is killed, the data already on +# disk survives; relaunch this script and it picks up. Convergence is guaranteed: +# every iteration that survives adds >= 1 line. +# +# Usage: scripts/gen_full_supervisor.sh +set -uo pipefail # NB: no -e — a killed gen sub-call must NOT abort the loop +cd "$(dirname "$0")/.." +export PYTHONPATH=src +W=$(nproc) + +gen() { # target_n out seed [extra flags...] + local target=$1 out=$2 seed=$3; shift 3 + while :; do + local cur=0 + if [ -f "$out" ]; then cur=$(wc -l < "$out"); fi + if [ "$cur" -ge "$target" ]; then + echo "[sup] $out done: $cur/$target lines"; return 0 + fi + echo "[sup] $out at $cur/$target — launching --resume $(date '+%H:%M:%S')" + .venv/bin/python scripts/gen_synthetic.py \ + --n "$target" --out "$out" --seed "$seed" --workers "$W" --resume "$@" + # loop: re-check count (gen may have been killed or completed) + done +} + +echo "[sup] === START $(date '+%H:%M:%S') ===" +gen 500000 data/align.jsonl 0 +gen 100000 data/sft.jsonl 0 --evolve +echo "[sup] === ALL DONE $(date '+%H:%M:%S') ===" +echo "[sup] align=$(wc -l < data/align.jsonl) sft=$(wc -l < data/sft.jsonl) eval=$(wc -l < data/eval_synth.jsonl)" diff --git a/scripts/gen_synthetic.py b/scripts/gen_synthetic.py index db8266d..a121b60 100644 --- a/scripts/gen_synthetic.py +++ b/scripts/gen_synthetic.py @@ -22,10 +22,11 @@ import json import os from concurrent.futures import ProcessPoolExecutor, as_completed from pathlib import Path +from typing import Any from tqdm import tqdm -from tsmm.data.gen_pipeline import gen_one_sample, write_jsonl +from tsmm.data.gen_pipeline import _to_jsonable, gen_one_sample, write_jsonl from tsmm.data.instruct import evolve_instruction @@ -55,19 +56,53 @@ def run( missing_rate: float = 0.05, event_prob: float = 0.5, do_evolve: bool = False, + resume: bool = False, ) -> int: + """Stream samples to ``out`` as they are produced. + + Writes each sample to disk the instant its future completes, while keeping + the output file in index order via a tiny reorder buffer (size <= workers). + Peak memory is thus O(workers) instead of O(n): 500k samples never all sit + in RAM, and a crash loses only the unwritten tail rather than everything. + + With ``resume=True`` (and ``--n`` as the *target* total), count existing + lines in ``out`` and only generate the missing suffix, appending. Repeated + crashes converge to the full target: each resume picks up where it stopped. + """ Path(out).parent.mkdir(parents=True, exist_ok=True) + start_idx = 0 + mode = "w" + if resume and os.path.exists(out): + with open(out, "r", encoding="utf-8") as f: + start_idx = sum(1 for _ in f) + if start_idx >= n: + print(f"[gen] {out} already has {start_idx} lines >= target {n}; nothing to do") + return start_idx + mode = "a" + print(f"[gen] resume: {out} has {start_idx}/{n}, appending {n - start_idx}") work = [ (T, C, seed + i, anomaly_types, missing_rate, event_prob, do_evolve) - for i in range(n) + for i in range(start_idx, n) ] - samples = [None] * n - with ProcessPoolExecutor(max_workers=workers) as ex: - futures = {ex.submit(_one, w): i for i, w in enumerate(work)} - for fut in tqdm(as_completed(futures), total=n, desc=os.path.basename(out)): + written = start_idx + pending: dict[int, Any] = {} # index -> sample, awaiting its turn + next_i = start_idx # next index that must be flushed + with open(out, mode, encoding="utf-8") as f, \ + ProcessPoolExecutor(max_workers=workers) as ex: + futures = {ex.submit(_one, w): i for i, w in zip(range(start_idx, n), work)} + for fut in tqdm(as_completed(futures), total=len(work), + desc=os.path.basename(out), initial=start_idx): i = futures[fut] - samples[i] = fut.result() - return write_jsonl(samples, out) + pending[i] = fut.result() + # flush everything that is now contiguously ready, in order + while next_i in pending: + f.write(json.dumps(_to_jsonable(pending.pop(next_i)), + ensure_ascii=False)) + f.write("\n") + next_i += 1 + written += 1 + f.flush() + return written def main(argv: list[str] | None = None) -> int: @@ -81,6 +116,8 @@ def main(argv: list[str] | None = None) -> int: p.add_argument("--missing-rate", type=float, default=0.05) p.add_argument("--event-prob", type=float, default=0.5) p.add_argument("--evolve", action="store_true", help="apply Evol-Instruct rephrasing (for sft set)") + p.add_argument("--resume", action="store_true", + help="count existing lines in --out and append only the missing suffix (crash-safe)") args = p.parse_args(argv) n = run( @@ -93,6 +130,7 @@ def main(argv: list[str] | None = None) -> int: missing_rate=args.missing_rate, event_prob=args.event_prob, do_evolve=args.evolve, + resume=args.resume, ) print(f"wrote {n} samples -> {args.out}") return 0