#!/usr/bin/env python """Generate the training corpus straight into the transfer, touching no local disk. dataset_build.py --availability .../availability.npz --budget-gb 330 \\ --dest htpc@192.168.1.66:/var/home/htpc/suvi/shards This replaced a two-stage build that existed only to stage bytes: extract a 1.55 TiB library of day-shards, then rewrite a dense 323 GiB subset of it, then copy that to the training host. Only ~21% of the library was ever shipped, and barely half of what it stored was trainable, so 1.87 TiB of local disk existed purely to be read once and copied somewhere else. Here the corpus is chosen first and only the slots it wants are ever produced. Each shard is compressed record by record straight into an ``ssh`` pipe, so nothing larger than one ~4 MB record is ever held and no shard exists locally at any point. Generating ~334 GB directly beats generating 1.55 TiB to discard four fifths of it, even though it is slower per byte. **One worker by default, deliberately.** Read concurrency is the only lever that would speed this up -- 48 concurrent readers reach ~207 MB/s against one reader's ~62 -- but the archive mount is shared with the downloader and everything else, and taking that throughput means taking it from them. One outstanding read is the polite ceiling; when the store is quiet a single stream is enough, and when it is busy, absorbing the slowdown is the right behaviour. ``--workers`` exists for the case where the mount is known to be idle. """ import argparse import datetime import hashlib import json import os import shlex import subprocess import sys import time from multiprocessing import Pool sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) from suvi import (availability as availability_module, corpus, dataset, db, detectors, fitsio, gpubox, paths, segments as segments_lib) #: The bench's ground-truth window. Excluded from training so it stays a clean holdout. HOLDOUT = ("2024-05-08", "2024-05-16") def holdout_days(first=HOLDOUT[0], last=HOLDOUT[1]): start = datetime.date.fromisoformat(first) end = datetime.date.fromisoformat(last) days = [] while start <= end: days.append(start.isoformat()) start += datetime.timedelta(days=1) return days def _day_of(t): return datetime.datetime.fromtimestamp( t, datetime.timezone.utc).strftime("%Y-%m-%d") def _extract(job): """Worker: one slot -> (t, encoded block or None, rejection reason or None). Every frame is **vetted before it can become training truth** by the production judgment, ``detectors.judge_frame`` (header_v1, then disc_v1 and field_v1 on the pixels), and one BAD band drops the whole slot. Its quality is measured on real, hand-confirmed archive defects (``detector_eval.py``), not on the bench's injected damage, which once reported this vetting as perfect while bars, ghosts and static passed into the corpus. A frame that is present but damaged is exactly as poisonous as a ground-truth label as a missing one is harmless: the model would be *taught* to reproduce the damage. UNKNOWN verdicts pass -- abstention is not evidence of a fault, and rejecting on it would throw away frames whose headers merely lack optional cards. """ root, satellite, time_start, rows = job fitsio.quiet_astropy() try: arrays = [] for wavelength in paths.WAVELENGTHS: relpath = rows.get(wavelength) if relpath is None: return time_start, None, None absolute = paths.abspath(relpath, root) cards, header_error = fitsio.read_header(absolute) source, verdict, image = detectors.judge_frame( (satellite, wavelength, time_start), cards, header_error, lambda: fitsio.read_image(absolute)) if source == "read": return time_start, None, None # unreadable: absent, not vetoed if verdict.verdict == detectors.BAD: return time_start, None, source arrays.append(image) return time_start, dataset.encode_frames(arrays), None except (OSError, ValueError): return time_start, None, None def wait_for_destination(destination, seconds, poll=30): """Block until the destination answers, or `seconds` elapse. Used at startup as well as mid-build, because a corpus is queued against a host that may be away when the queueing happens: this one is on WiFi that has dropped three times in as many days, once for over an hour. Starting the build regardless and letting it wait means the work resumes the moment the host returns, instead of waiting for a person to notice it did. Returns the listing on success so the caller does not have to ask twice. """ deadline = time.time() + seconds announced = False while True: try: listing = destination.existing() if announced: print(" destination is back; resuming", flush=True) return listing except OSError as error: if time.time() >= deadline: raise if not announced: print(f" destination away ({str(error)[:60]}); waiting up to " f"{seconds / 3600:.1f} h for it", flush=True) announced = True time.sleep(poll) def build_with_retry(conn, root, day, satellite, times, destination, args): """One shard, re-attempted while the destination is away. A full corpus is a day or more of streaming into one host, and this one is on WiFi that has dropped twice in three days -- each time killing the build with a broken pipe partway through a shard. Resume already handles that correctly (a shard is renamed into place only when its stream closes cleanly, so an interrupted one leaves a ``.partial`` the scan ignores), but only if someone is around to restart it. Waiting for the host to come back is strictly better than losing the hours between the drop and a human noticing. Only the *destination* is retried. An archive read failure is a different problem and still stops the build, because quietly re-reading a file the mount cannot serve would spin forever rather than say so. """ for attempt in range(1, args.retries + 2): try: return build_shard(conn, root, day, satellite, times, destination, args.workers) except OSError as error: if attempt > args.retries: raise print(f" destination lost ({type(error).__name__}: " f"{str(error)[:80]}); attempt {attempt}/{args.retries}, " f"waiting for it to return", flush=True) wait_for_destination(destination, args.retry_wait) raise AssertionError("unreachable") # --------------------------------------------------------------------- destination class Destination: """Where shards go: a remote ``host:/path`` or a local directory. Both are opened as a *stream* rather than a file, because the whole design rests on a shard never existing anywhere in one piece. The remote case pipes into ``cat``; the local case is what the tests drive, so the streaming path itself is exercised rather than a stand-in for it. """ def __init__(self, spec): self.spec = spec self.host, _, self.root = spec.partition(":") if ":" in spec else ("", "", spec) def existing(self): """Names already complete at the destination -- the only resume state. Nothing is kept locally, so this is the whole record of what has been done. A shard is renamed into place only once its stream closed cleanly, so presence means complete. """ if self.host: listing = subprocess.run( gpubox.ssh_argv(self.host) + [f"ls {shlex.quote(self.root)} 2>/dev/null; true"], stdout=subprocess.PIPE, stderr=subprocess.PIPE, ) # An unreachable host must not read as "nothing is there yet": # that silently restarts a finished build from zero. if listing.returncode != 0: raise OSError( f"cannot list {self.spec} ({listing.returncode}): " + listing.stderr.decode(errors="replace").strip()[:200]) return {line.strip() for line in listing.stdout.decode(errors="replace").splitlines() if line.strip()} return set(os.listdir(self.root)) if os.path.isdir(self.root) else set() def open(self, name, sealed=False): """A context manager yielding a binary stream for one file. `sealed` files (shards) are committed only if they end with the shard trailer, so a stream cut short can never be renamed into place. """ cls = _RemoteWrite if self.host else _LocalWrite return cls(self, name, sealed=sealed) def put(self, name, payload): """Write a small file (the manifest, segments.json) in one go.""" with self.open(name) as handle: handle.write(payload) def get(self, name): """Read a small file back, or None if it is not there.""" if self.host: result = subprocess.run( gpubox.ssh_argv(self.host) + [f"cat {shlex.quote(os.path.join(self.root, name))} 2>/dev/null"], stdout=subprocess.PIPE, stderr=subprocess.DEVNULL, ) return result.stdout if result.returncode == 0 and result.stdout else None path = os.path.join(self.root, name) return open(path, "rb").read() if os.path.exists(path) else None class _LocalWrite: def __init__(self, destination, name, sealed=False): os.makedirs(destination.root, exist_ok=True) self.path = os.path.join(destination.root, name) self.temporary = f"{self.path}.partial" self.sealed = sealed def __enter__(self): self.handle = open(self.temporary, "wb") return self.handle def __exit__(self, kind, value, traceback): self.handle.close() if kind is None and self.sealed: with open(self.temporary, "rb") as check: check.seek(max(os.path.getsize(self.temporary) - len(dataset.MAGIC), 0)) if check.read() != dataset.MAGIC: os.unlink(self.temporary) raise OSError(f"{self.path}: truncated stream, no shard trailer") if kind is None: os.replace(self.temporary, self.path) else: os.unlink(self.temporary) # never leave a half-written shard behind return False class _RemoteWrite: """A file streamed into ``ssh host 'cat > X.partial && ... mv X.partial X'``. The rename is the commit, but ``cat`` finishing proves nothing on its own. If the *local* builder dies -- killed at a session teardown, say -- its pipe closes cleanly, ssh forwards an ordinary end-of-file, the remote ``cat`` exits 0, and the rename commits whatever arrived. That is how two 33-39 MB fragments of a ~1.4 GB day came to sit in the corpus under their final names, passing the resume scan until the trainer refused them at load. So a sealed file (every shard) is renamed only if it ends with the shard trailer's magic. The trailer is the last thing the writer emits, so its presence means the stream ran to the end; anything cut short is left as a ``.partial`` the resume scan ignores and the command fails loudly. """ def __init__(self, destination, name, sealed=False): self.destination = destination target = shlex.quote(os.path.join(destination.root, name)) check = "" if sealed: magic = shlex.quote(dataset.MAGIC.decode()) check = (f'[ "$(tail -c {len(dataset.MAGIC)} {target}.partial)" = {magic} ] ' f'|| {{ echo "truncated stream: no shard trailer" >&2; exit 3; }} && ') self.command = (f"mkdir -p {shlex.quote(destination.root)} && " f"cat > {target}.partial && {check}" f"mv {target}.partial {target}") def __enter__(self): self.process = subprocess.Popen( gpubox.ssh_argv(self.destination.host) + [self.command], stdin=subprocess.PIPE, stderr=subprocess.PIPE, ) return self.process.stdin def __exit__(self, kind, value, traceback): if kind is not None: self.process.kill() self.process.wait() return False self.process.stdin.close() code = self.process.wait() if code != 0: error = self.process.stderr.read().decode(errors="replace").strip() raise OSError(f"streaming to {self.destination.spec} failed " f"({code}): {error[:200]}") return False # ------------------------------------------------------------------------- building def slot_paths(conn, satellite, times): """{t_start: {wavelength: relpath}} for the wanted slots of one satellite.""" if not times: return {} by_slot = {} for chunk in range(0, len(times), 500): window = times[chunk:chunk + 500] marks = ",".join("?" * len(window)) for row in conn.execute( f"SELECT wavelength, t_start, path FROM frame " f"WHERE satellite = ? AND t_start IN ({marks})", (satellite, *window), ): by_slot.setdefault(row["t_start"], {})[row["wavelength"]] = row["path"] return by_slot def build_shard(conn, root, day, satellite, times, destination, workers): """Generate one shard straight into the destination. Returns (bytes written, slots with frames, {reason: count}, {unusable times}). The records are produced lazily inside the writer, so the shard is never held. The unusable set is what verification later runs on: a slot the profile called pristine but extraction would not accept -- whether vetoed by a detector or simply unreadable -- has to split its segment, and observing it here is why nothing needs to be read back afterwards. """ by_slot = slot_paths(conn, satellite, times) jobs = [(root, satellite, t, by_slot.get(t, {})) for t in times] dropped = {} unusable = set() kept = 0 name = dataset.shard_name(day, satellite) # Explicit terminate *and* join, not `with Pool(...)`. The context manager # only terminates, which does not reap the result handler, and on a link that # drops mid-shard this runs hundreds of times: each abandoned `imap` left its # buffered blocks (4.9 MB apiece) alive, and the builder grew from 260 MB to # 16.6 GB over 80 attempts before it would have taken the VM down with it. pool = Pool(workers) stream = None try: def records(): nonlocal kept # imap, not imap_unordered: record offsets are cumulative, so the # writer needs ascending t_start. chunksize=1 with one worker means # exactly one read is outstanding against the archive at a time. for time_start, block, reason in pool.imap(_extract, jobs, chunksize=1): if reason: dropped[reason] = dropped.get(reason, 0) + 1 if block is None: unusable.add(time_start) else: kept += 1 yield time_start, block stream = records() with destination.open(name, sealed=True) as handle: size = dataset.write_shard_stream( handle, day, satellite, paths.WAVELENGTHS, stream) finally: if stream is not None: stream.close() # stop feeding the abandoned iterator pool.terminate() pool.join() # reap the workers and their buffers return size, kept, dropped, unusable def verified_segments(slices, vetoed, satellites, minimum_slots): """Split the chosen slices wherever extraction disagreed with the profile. The availability profile comes from the stored judgment, and extraction re-judges every frame with the same contract, so they should agree; this catches the frames where they do not (a profile made under older rules, or a frame replaced in the archive since). Vetoes were observed as they happened, so the corpus is verified from that record and nothing is read twice. """ def usable(t): return all((_day_of(t), s, t) not in vetoed for s in satellites) verified = [] for role, chosen in slices.items(): for segment in chosen: verified.extend( (role, pair) for pair in segments_lib.verify( [segment], usable, minimum_slots=minimum_slots)) return verified def build_parser(): parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) parser.add_argument("--availability", required=True, help="availability.npz from profile_archive.py (the stored " "judgment); pristine segments are selected from it") parser.add_argument("--dest", required=True, help="host:/path or a local directory") parser.add_argument("--budget-gb", type=float, required=True, help="total size of the corpus to generate") parser.add_argument("--cap-slots", type=int, default=None, help="most slots any one segment may contribute; defaults " "to half a synodic solar rotation, the length past " "which a segment shows an unrelated hemisphere") parser.add_argument("--span-min", type=int, default=64, help="smallest model input span the corpus must serve; " "a segment must hold 2*span+1 contiguous frames") parser.add_argument("--val-segments", type=int, default=16) parser.add_argument("--test-segments", type=int, default=16) parser.add_argument("--satellites", default="16,18") parser.add_argument("--workers", type=int, default=1, help="concurrent archive readers (default 1, so the " "shared mount is never monopolised); raise only when " "the store is known to be idle") parser.add_argument("--retries", type=int, default=100, help="times to wait out a destination that goes " "away mid-shard; the host is on WiFi and a " "full build runs for a day") parser.add_argument("--retry-wait", type=int, default=3600, help="seconds to keep waiting for it to return") parser.add_argument("--seed", type=int, default=7) parser.add_argument("--db", default=None) return parser def main(argv=None): args = build_parser().parse_args(argv) satellites = tuple(int(s) for s in args.satellites.split(",")) conn = db.connect(args.db, readonly=True) root = paths.data_root() destination = Destination(args.dest) minimum = 2 * args.span_min + 1 model = availability_module.Availability.load(args.availability) found = segments_lib.from_availability(model, minimum) held = set(holdout_days()) nominal = [segment for segment in found if not set(segments_lib.days_touched([segment])) & held] print(f"{len(nominal)} pristine segments >= {minimum} frames " f"({len(found) - len(nominal)} dropped for touching the holdout)") cap = args.cap_slots or segments_lib.decorrelation_slots() slices, used = corpus.choose_slices( nominal, cap, args.budget_gb * 1e9, satellites, args.seed, args.val_segments, args.test_segments) total = sum(len(v) for v in slices.values()) print(f"{total} slices admitted ({args.budget_gb:.0f} GB budget, cap {cap} " f"slots = {cap * paths.CADENCE / 86400:.1f} d): " + ", ".join(f"{k} {len(v)}" for k, v in sorted(slices.items()))) wanted = corpus.wanted_slots(slices, satellites) already = wait_for_destination(destination, args.retry_wait) todo = [key for key in sorted(wanted) if dataset.shard_name(*key) not in already] print(f"{len(wanted)} shards wanted, {len(wanted) - len(todo)} already at the " f"destination, {len(todo)} to generate with {args.workers} reader(s)") # The veto record lives at the destination, not here, for the same reason the # shards do -- and because a resumed run must not forget what the slots of an # already-shipped shard turned out to be. Losing it would silently produce a # manifest claiming segments longer than the data behind them. previous = destination.get("vetting.json") vetoed = {tuple(entry) for entry in json.loads(previous)} if previous else set() if vetoed: print(f" carried {len(vetoed)} vetoed slots forward from a previous run") written = slots = 0 started = time.time() for index, (day, satellite) in enumerate(todo, 1): times = wanted[(day, satellite)] size, kept, dropped, unusable = build_with_retry( conn, root, day, satellite, times, destination, args) vetoed.update((day, satellite, t) for t in unusable) destination.put("vetting.json", json.dumps(sorted(vetoed)).encode()) written += size slots += kept elapsed = time.time() - started rate = written / 1e6 / max(elapsed, 1e-6) remaining = (len(todo) - index) * elapsed / index print(f" [{index:>4}/{len(todo)}] {dataset.shard_name(day, satellite)}: " f"{kept}/{len(times)} slots" + (f" ({sum(dropped.values())} vetoed)" if dropped else "") + f", {size / 1e6:.0f} MB, {written / 1e9:.1f} GB at " f"{rate:.1f} MB/s, eta {remaining / 3600:.1f} h", flush=True) print("verifying segments against what extraction actually accepted") verified = verified_segments(slices, vetoed, satellites, minimum) splits = {"train": [], "val": [], "test": []} for role, pair in verified: splits[role].append(tuple(pair)) for name in ("train", "val", "test"): splits[name] = sorted(splits[name]) mass = sum(max((b - a) // paths.CADENCE - 2 * args.span_min, 0) for a, b in splits[name]) print(f" {name:>5}: {len(splits[name])} segments, " f"{mass:,} targets per satellite at span {args.span_min}") payload = {name: [list(pair) for pair in chosen] for name, chosen in splits.items()} digest = hashlib.sha256( json.dumps(payload, sort_keys=True).encode()).hexdigest()[:16] manifest = {"satellites": list(satellites), "holdout": list(HOLDOUT), "span_min": args.span_min, "days": sorted({day for day, _ in wanted}), "segments": payload, "digest": digest, "cap_slots": cap} destination.put("manifest.json", json.dumps(manifest, indent=2).encode()) destination.put("segments.json", json.dumps( {"segments": [list(pair) for _, pair in verified]}).encode()) # The availability profile ships with the corpus it selected. train.py looks # for it beside the shards and refuses to start without it -- deliberately, # since falling back to synthetic-only outages would silently undo the # measured failure distribution -- so a corpus that does not carry it is not # a corpus anyone can train on. with open(args.availability, "rb") as handle: destination.put("availability.npz", handle.read()) elapsed = time.time() - started print(f"\n{slots:,} slots, {written / 1e9:.1f} GB in {elapsed / 3600:.2f} h " f"({written / 1e6 / max(elapsed, 1e-6):.1f} MB/s); manifest {digest}") return 0 if __name__ == "__main__": sys.exit(main())