From 1a096ff989148b7343375651e1bd7ab07949c481 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=BC=A0=E5=AE=97=E5=B9=B3?= Date: Tue, 30 Jun 2026 23:48:41 +0000 Subject: [PATCH] feat(gen): streaming + crash-resilient resume for full-scale data gen MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The original gen_synthetic.py held all N samples in a [None]*n list and only wrote at the end — three full-scale runs were killed mid-way (env reaps long sessions) and lost everything (0 bytes on disk each time). Fixes (proven on the 500k+100k run that just completed): - run() now streams: each completed future is written to disk the instant it's produced, via an O(workers) reorder buffer that keeps output in index order. Peak RAM drops from O(n) (~2.7GB for 50k) to O(workers). - --resume flag: counts existing lines and appends only the missing suffix (same seed => deterministic prefix, so resume is content-safe). - scripts/gen_full_supervisor.sh: loops 'run --resume -> check count' until target reached. Each surviving iteration adds >=1 line, so the run converges even across multiple kills (took 2 kills + 2 resumes to finish the 500k align). On this host: 2 partial runs (50k + 437k) then a final resume filled the last 62k. - Verified: order-deterministic (same-seed runs byte-identical), resume appends without rewriting the prefix, target-already-met => no-op. Full-scale data now on disk: align 500k (27G) + sft 100k (5.4G) + eval_synth 2k held-out (untouched). --- scripts/gen_full_supervisor.sh | 39 ++++++++++++++++++++++++ scripts/gen_synthetic.py | 54 +++++++++++++++++++++++++++++----- 2 files changed, 85 insertions(+), 8 deletions(-) create mode 100755 scripts/gen_full_supervisor.sh 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