128 lines
5.4 KiB
Python
128 lines
5.4 KiB
Python
#!/usr/bin/env python
|
|
"""Keep a rotating pool of training slices in memory, generated from the archive.
|
|
|
|
pool_producer.py --manifest /work/pool_manifest.json.zst --archive /archive/Data \\
|
|
--pool /dev/shm/suvi_pool --budget-gb 48 --workers 16
|
|
|
|
Runs on the training host beside ``train.py --pool``: picks a random slice of a
|
|
pristine segment (:func:`suvi.pool.choose_slice`), reads its six bands per slot for
|
|
both satellites from the archive share, encodes them exactly as the on-disk corpus
|
|
does (:func:`suvi.dataset.encode_frames`), writes the slice, then retires the oldest
|
|
slices past the budget. Nothing is judged here: the manifest lists only frames the
|
|
stored judgment called good. A frame that cannot be read is an absent slot, which
|
|
splits its slice's segment rather than stopping the pool.
|
|
|
|
With ``--fixed val`` (or ``test``) it instead writes the manifest's validation
|
|
windows (or test segments), once each, to a disk directory (``val000`` ...), skipping
|
|
any already complete, and exits: evaluation sets holding nothing but their own slots.
|
|
|
|
See :mod:`suvi.pool` for the design and the measured rates.
|
|
"""
|
|
|
|
import argparse
|
|
import os
|
|
import sys
|
|
import time
|
|
from multiprocessing import Pool
|
|
|
|
import numpy as np
|
|
|
|
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
|
|
|
from suvi import dataset, fitsio, paths, pool
|
|
|
|
#: Slots encoded per worker dispatch.
|
|
CHUNK = 4
|
|
|
|
|
|
def encode_slot(files):
|
|
"""Worker: six archive paths -> a 6x640x640 block, or None if any band fails."""
|
|
arrays = []
|
|
for path in files:
|
|
try:
|
|
image, _ = fitsio.read_image(path)
|
|
except OSError:
|
|
return None
|
|
if image is None:
|
|
return None
|
|
arrays.append(image)
|
|
try:
|
|
return dataset.encode_frames(arrays)
|
|
except ValueError:
|
|
return None
|
|
|
|
|
|
def build_parser():
|
|
parser = argparse.ArgumentParser(description=__doc__,
|
|
formatter_class=argparse.RawDescriptionHelpFormatter)
|
|
parser.add_argument("--manifest", required=True, help="from pool_manifest.py")
|
|
parser.add_argument("--archive", required=True, help="the share's Data directory")
|
|
parser.add_argument("--pool", required=True, help="a RAM directory, e.g. /dev/shm/...")
|
|
parser.add_argument("--budget-gb", type=float, default=48.0)
|
|
parser.add_argument("--workers", type=int, default=16)
|
|
parser.add_argument("--slices", type=int, default=None,
|
|
help="stop after this many (default: run until stopped)")
|
|
parser.add_argument("--fixed", choices=("val", "test"), default=None,
|
|
help="write this evaluation set from the manifest to --pool "
|
|
"and exit")
|
|
parser.add_argument("--seed", type=int, default=None,
|
|
help="default: from the clock, so restarts draw new slices")
|
|
return parser
|
|
|
|
|
|
def main(argv=None):
|
|
args = build_parser().parse_args(argv)
|
|
fitsio.quiet_astropy()
|
|
manifest = pool.load_manifest(args.manifest)
|
|
satellites = tuple(manifest["satellites"])
|
|
span = manifest["span"]
|
|
segments = [tuple(pair) for pair in manifest["segments"]]
|
|
frames = manifest["frames"]
|
|
rng = np.random.default_rng(args.seed if args.seed is not None else time.time_ns())
|
|
os.makedirs(args.pool, exist_ok=True)
|
|
pool.clear_partials(args.pool)
|
|
budget = args.budget_gb * 1e9
|
|
print(f"pool {args.pool}: {len(segments)} segments, "
|
|
f"{pool.target_count(segments, span):,} targets per satellite at span {span}; "
|
|
f"budget {args.budget_gb:.0f} GB, {args.workers} readers", flush=True)
|
|
|
|
counter = len(pool.ready_slices(args.pool))
|
|
made = 0
|
|
started = time.time()
|
|
slots_done = 0
|
|
with Pool(args.workers, initializer=fitsio.quiet_astropy,
|
|
maxtasksperchild=500) as workers:
|
|
def blocks(satellite, times):
|
|
files = [[os.path.join(args.archive, rel)
|
|
for rel in frames[str(satellite)][str(t)]] for t in times]
|
|
yield from zip(times, workers.imap(encode_slot, files, chunksize=CHUNK))
|
|
|
|
if args.fixed:
|
|
windows = [tuple(pair) for pair in manifest[args.fixed]]
|
|
for index, piece in enumerate(windows):
|
|
name = f"{args.fixed}{index:03d}"
|
|
if os.path.exists(os.path.join(args.pool, name, pool.SLICE_FILE)):
|
|
continue
|
|
path = pool.write_slice(args.pool, name, piece, satellites, blocks, span)
|
|
print(f" {name} ({index + 1}/{len(windows)}): "
|
|
f"{'written' if path else 'nothing usable, dropped'}", flush=True)
|
|
return 0
|
|
|
|
while args.slices is None or made < args.slices:
|
|
piece = pool.choose_slice(segments, span, rng)
|
|
counter += 1
|
|
path = pool.write_slice(args.pool, pool.slice_name(counter, piece), piece,
|
|
satellites, blocks, span)
|
|
made += 1
|
|
slots_done += (piece[1] - piece[0]) // paths.CADENCE * len(satellites)
|
|
gone = pool.retire(args.pool, budget)
|
|
resident = pool.ready_slices(args.pool)
|
|
rate = slots_done / max(time.time() - started, 1e-6)
|
|
print(f" {os.path.basename(path) if path else 'slice dropped'}: "
|
|
f"{len(resident)} resident ({sum(s for _, s in resident) / 1e9:.1f} GB), "
|
|
f"{len(gone)} retired, {rate:.1f} slots/s", flush=True)
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(main())
|