noaa-goes-visualization/dataset_build.py

#!/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())