295 lines
13 KiB
Python
295 lines
13 KiB
Python
|
"""A rotating in-memory training pool, generated from the archive while training runs.
|
||
|
|
|
||
|
|
The on-disk corpus holds ~80,000 slots (323 GB) because that is what fits on the
|
||
|
|
training host; the archive's pristine two-satellite era is ~5.6 times larger. So
|
||
|
|
instead of a fixed corpus the host can generate training data as it trains: a
|
||
|
|
producer (``pool_producer.py``) reads FITS straight from the archive share, writes
|
||
|
|
*slices* -- a day of one pristine segment, both satellites, with the model's span of
|
||
|
|
margin -- into a RAM directory, and retires the oldest when the pool is over budget.
|
||
|
|
``train.py --pool`` samples from whatever slices are resident. Validation and test
|
||
|
|
stay on the fixed on-disk corpus, so scores remain comparable between runs.
|
||
|
|
|
||
|
|
A slice is an ordinary little corpus: day-shards in the usual format plus
|
||
|
|
``slice.json`` naming its verified segments. It is built in ``.<name>.partial`` and
|
||
|
|
renamed into place complete, so a reader never sees half of one; a reader opens its
|
||
|
|
shards eagerly (:class:`suvi.dataset.Shard`), so a slice retired mid-visit stays
|
||
|
|
readable through the handles already open.
|
||
|
|
|
||
|
|
Measured on the host (2026-10-13): 16 workers produce ~8.5 slots/s from the share
|
||
|
|
(~90 MB/s, near gigabit); training at the v9 pace consumes ~5 new slots/s.
|
||
|
|
"""
|
||
|
|
|
||
|
|
import bisect
|
||
|
|
import datetime
|
||
|
|
import json
|
||
|
|
import os
|
||
|
|
import shutil
|
||
|
|
import time
|
||
|
|
|
||
|
|
import numpy as np
|
||
|
|
|
||
|
|
from . import corpus, dataset, paths, segments as segments_module
|
||
|
|
|
||
|
|
#: Targets per slice: one day of one satellite, so a slice holds one evolving view
|
||
|
|
#: of the Sun and the pool turns over many of them.
|
||
|
|
SLICE_TARGETS = 360
|
||
|
|
#: Slices kept clear of the evaluation segments, in seconds either side: a pristine
|
||
|
|
#: run separated from a val segment by a few bad slots shows the same Sun.
|
||
|
|
GUARD_SECONDS = 86400
|
||
|
|
#: The validation set (v10 on): one window of VAL_TARGETS targets in each of the
|
||
|
|
#: previous corpus's val segments and in VAL_NEW more spread evenly over the era --
|
||
|
|
#: ~5,800 targets per satellite over 64 segments, against v9's 1,714 over 16.
|
||
|
|
VAL_NEW = 48
|
||
|
|
VAL_TARGETS = 90
|
||
|
|
VAL_SEED = 11
|
||
|
|
#: Manifest format version.
|
||
|
|
MANIFEST_VERSION = 2
|
||
|
|
SLICE_FILE = "slice.json"
|
||
|
|
|
||
|
|
|
||
|
|
# ----------------------------------------------------------------- the manifest
|
||
|
|
|
||
|
|
|
||
|
|
def subtract(segments, cuts, minimum):
|
||
|
|
"""`segments` with every (low, high) of `cuts` removed, widened by
|
||
|
|
:data:`GUARD_SECONDS`; pieces shorter than `minimum` slots are dropped."""
|
||
|
|
out = []
|
||
|
|
for a, b in segments:
|
||
|
|
pieces = [(a, b)]
|
||
|
|
for low, high in cuts:
|
||
|
|
low, high = low - GUARD_SECONDS, high + GUARD_SECONDS
|
||
|
|
pieces = [part for start, end in pieces
|
||
|
|
for part in ((start, min(end, low)), (max(start, high), end))
|
||
|
|
if part[1] > part[0]]
|
||
|
|
out.extend((start, start + (end - start) // paths.CADENCE * paths.CADENCE)
|
||
|
|
for start, end in pieces
|
||
|
|
if (end - start) // paths.CADENCE >= minimum)
|
||
|
|
return out
|
||
|
|
|
||
|
|
|
||
|
|
def build_manifest(conn, model, run_id, span, old_val, test, holdout_days,
|
||
|
|
satellites=(16, 18), val_new=None, val_targets=None, seed=None,
|
||
|
|
old_train=()):
|
||
|
|
"""The pool's segments and a fixed validation set, with every slot's frame paths.
|
||
|
|
|
||
|
|
`model` is the availability model of judgment run `run_id`. Validation is one
|
||
|
|
window of `val_targets` targets (plus the span either side) in each of the
|
||
|
|
`old_val` segments -- the previous corpus's, so its days stay covered -- and in
|
||
|
|
`val_new` further segments spread evenly over the era, none of them touching the
|
||
|
|
previous corpus's `old_train` (so the previous model can be scored on the new
|
||
|
|
set fairly). The pool keeps :data:`GUARD_SECONDS` clear of every validation
|
||
|
|
window and `test` segment, and off the bench's `holdout_days`. Paths are the run's
|
||
|
|
good, non-stub frames, one per (satellite, band, slot); a slot missing one on
|
||
|
|
any satellite splits its segment or window, so the producer judges nothing.
|
||
|
|
The `test` segments are kept whole, as the evaluation set's own slices.
|
||
|
|
"""
|
||
|
|
val_new = VAL_NEW if val_new is None else val_new
|
||
|
|
val_targets = VAL_TARGETS if val_targets is None else val_targets
|
||
|
|
seed = VAL_SEED if seed is None else seed
|
||
|
|
minimum = 2 * span + 1
|
||
|
|
window = val_targets + 2 * span
|
||
|
|
held = set(holdout_days)
|
||
|
|
|
||
|
|
def near(segment, others):
|
||
|
|
a, b = segment
|
||
|
|
return any(a < high + GUARD_SECONDS and b > low - GUARD_SECONDS
|
||
|
|
for low, high in others)
|
||
|
|
|
||
|
|
candidates = [pair for pair in segments_module.from_availability(model, minimum)
|
||
|
|
if not set(segments_module.days_touched([pair])) & held]
|
||
|
|
old_val, test = [tuple(p) for p in old_val], [tuple(p) for p in test]
|
||
|
|
old_train = [tuple(p) for p in old_train]
|
||
|
|
eligible = [pair for pair in candidates if not near(pair, old_val + test + old_train)
|
||
|
|
and (pair[1] - pair[0]) // paths.CADENCE >= window]
|
||
|
|
count = min(val_new, len(eligible))
|
||
|
|
new_val = [eligible[int((i + 0.5) * len(eligible) / count)] for i in range(count)]
|
||
|
|
val_segments = sorted(old_val + new_val)
|
||
|
|
rng = np.random.default_rng(seed)
|
||
|
|
windows = []
|
||
|
|
for a, b in val_segments:
|
||
|
|
length = min(window, (b - a) // paths.CADENCE)
|
||
|
|
first = int(rng.integers(0, (b - a) // paths.CADENCE - length + 1))
|
||
|
|
windows.append((a + first * paths.CADENCE,
|
||
|
|
a + (first + length) * paths.CADENCE))
|
||
|
|
chosen = subtract(candidates, windows + test, minimum)
|
||
|
|
|
||
|
|
wanted = sorted(chosen + windows + test)
|
||
|
|
frames = {str(s): {} for s in satellites}
|
||
|
|
starts = [a for a, _ in wanted]
|
||
|
|
marks = ",".join("?" * len(satellites))
|
||
|
|
# One pass over the era, not a query per segment: the planner drives a range
|
||
|
|
# query from the run's detections, so each of ~700 would rescan all 3 M of them.
|
||
|
|
for row in conn.execute(
|
||
|
|
f"SELECT f.satellite, f.wavelength, f.t_start, f.path FROM detection d "
|
||
|
|
f"JOIN frame f ON f.id = d.frame_id "
|
||
|
|
f"WHERE d.run_id = ? AND d.verdict = 'good' AND f.satellite IN ({marks}) "
|
||
|
|
f"AND (f.size_bytes IS NULL OR f.size_bytes >= ?)",
|
||
|
|
(run_id, *satellites, dataset.MIN_FRAME_BYTES),
|
||
|
|
) if wanted else ():
|
||
|
|
t = row["t_start"]
|
||
|
|
index = bisect.bisect_right(starts, t) - 1
|
||
|
|
if index < 0 or t >= wanted[index][1]:
|
||
|
|
continue
|
||
|
|
slot = frames[str(row["satellite"])].setdefault(
|
||
|
|
str(t), [None] * len(paths.WAVELENGTHS))
|
||
|
|
slot[paths.WAVELENGTHS.index(row["wavelength"])] = row["path"]
|
||
|
|
|
||
|
|
def complete(t):
|
||
|
|
return all(None not in frames[str(s)].get(str(t), [None]) for s in satellites)
|
||
|
|
|
||
|
|
verified = segments_module.verify(chosen, complete, minimum_slots=minimum)
|
||
|
|
val = segments_module.verify(windows, complete, minimum_slots=minimum)
|
||
|
|
tested = segments_module.verify(test, complete, minimum_slots=minimum)
|
||
|
|
keep = {str(t) for a, b in verified + val + tested
|
||
|
|
for t in range(a, b, paths.CADENCE)}
|
||
|
|
frames = {s: {t: v for t, v in slots.items() if t in keep}
|
||
|
|
for s, slots in frames.items()}
|
||
|
|
return {"version": MANIFEST_VERSION, "satellites": list(satellites),
|
||
|
|
"span": span, "judgment_run": int(run_id),
|
||
|
|
"segments": [list(pair) for pair in verified],
|
||
|
|
"val": [list(pair) for pair in val],
|
||
|
|
"val_segments": [list(pair) for pair in val_segments],
|
||
|
|
"test": [list(pair) for pair in tested],
|
||
|
|
"frames": frames}
|
||
|
|
|
||
|
|
|
||
|
|
def save_manifest(manifest, path):
|
||
|
|
compress, _ = dataset._zstd()
|
||
|
|
payload = compress(json.dumps(manifest).encode(), 9)
|
||
|
|
temporary = f"{path}.partial"
|
||
|
|
with open(temporary, "wb") as handle:
|
||
|
|
handle.write(payload)
|
||
|
|
os.replace(temporary, path)
|
||
|
|
|
||
|
|
|
||
|
|
def load_manifest(path):
|
||
|
|
_, decompress = dataset._zstd()
|
||
|
|
with open(path, "rb") as handle:
|
||
|
|
manifest = json.loads(decompress(handle.read()))
|
||
|
|
if manifest.get("version") != MANIFEST_VERSION:
|
||
|
|
raise ValueError(f"{path}: pool manifest version {manifest.get('version')}, "
|
||
|
|
f"expected {MANIFEST_VERSION}; rebuild it with pool_manifest.py")
|
||
|
|
return manifest
|
||
|
|
|
||
|
|
|
||
|
|
def target_count(segments, span):
|
||
|
|
"""Targets the segments hold at `span`, per satellite."""
|
||
|
|
return sum(max((b - a) // paths.CADENCE - 2 * span, 0) for a, b in segments)
|
||
|
|
|
||
|
|
|
||
|
|
# ------------------------------------------------------------------- slices
|
||
|
|
|
||
|
|
|
||
|
|
def choose_slice(segments, span, rng, targets=SLICE_TARGETS):
|
||
|
|
"""A random (t0, t1) of at most `targets` targets plus `span` slots of margin
|
||
|
|
either side, inside one segment. Segments are drawn by target count, so every
|
||
|
|
target in the manifest is equally likely to be covered."""
|
||
|
|
weights = np.array([max((b - a) // paths.CADENCE - 2 * span, 0)
|
||
|
|
for a, b in segments], dtype=float)
|
||
|
|
if weights.sum() <= 0:
|
||
|
|
raise ValueError(f"no segment holds a target at span {span}")
|
||
|
|
a, b = segments[int(rng.choice(len(segments), p=weights / weights.sum()))]
|
||
|
|
length = min(targets + 2 * span, (b - a) // paths.CADENCE)
|
||
|
|
first = int(rng.integers(0, (b - a) // paths.CADENCE - length + 1))
|
||
|
|
t0 = a + first * paths.CADENCE
|
||
|
|
return t0, t0 + length * paths.CADENCE
|
||
|
|
|
||
|
|
|
||
|
|
def write_slice(directory, name, piece, satellites, blocks, span):
|
||
|
|
"""Write one slice: its day-shards, then ``slice.json``, then the rename.
|
||
|
|
|
||
|
|
`blocks(satellite, times)` yields (t, block or None) in ascending order -- None
|
||
|
|
for a slot that could not be read, which splits the slice's segment. Returns
|
||
|
|
the slice's path, or None if nothing in it is left to train on.
|
||
|
|
"""
|
||
|
|
partial = os.path.join(directory, f".{name}.partial")
|
||
|
|
final = os.path.join(directory, name)
|
||
|
|
os.makedirs(partial)
|
||
|
|
unusable = set()
|
||
|
|
try:
|
||
|
|
for (day, satellite), times in sorted(
|
||
|
|
corpus.wanted_slots({"train": [piece]}, satellites).items()):
|
||
|
|
with open(os.path.join(partial, dataset.shard_name(day, satellite)),
|
||
|
|
"wb") as handle:
|
||
|
|
def records():
|
||
|
|
for t, block in blocks(satellite, times):
|
||
|
|
if block is None:
|
||
|
|
unusable.add(t)
|
||
|
|
yield t, block
|
||
|
|
dataset.write_shard_stream(handle, day, satellite,
|
||
|
|
paths.WAVELENGTHS, records())
|
||
|
|
verified = segments_module.verify([piece], lambda t: t not in unusable,
|
||
|
|
minimum_slots=2 * span + 1)
|
||
|
|
if not verified:
|
||
|
|
shutil.rmtree(partial)
|
||
|
|
return None
|
||
|
|
with open(os.path.join(partial, SLICE_FILE), "w") as handle:
|
||
|
|
json.dump({"segments": [list(pair) for pair in verified],
|
||
|
|
"satellites": list(satellites), "made": time.time()}, handle)
|
||
|
|
os.rename(partial, final)
|
||
|
|
except BaseException:
|
||
|
|
shutil.rmtree(partial, ignore_errors=True)
|
||
|
|
raise
|
||
|
|
return final
|
||
|
|
|
||
|
|
|
||
|
|
def ready_slices(directory):
|
||
|
|
"""Complete slices, oldest first: [(path, bytes)]."""
|
||
|
|
found = []
|
||
|
|
for name in os.listdir(directory) if os.path.isdir(directory) else ():
|
||
|
|
path = os.path.join(directory, name)
|
||
|
|
if name.startswith(".") or not os.path.exists(os.path.join(path, SLICE_FILE)):
|
||
|
|
continue
|
||
|
|
try:
|
||
|
|
made = json.load(open(os.path.join(path, SLICE_FILE)))["made"]
|
||
|
|
size = sum(entry.stat().st_size for entry in os.scandir(path))
|
||
|
|
except (OSError, ValueError, KeyError):
|
||
|
|
continue # retired between listing and reading
|
||
|
|
found.append((made, path, size))
|
||
|
|
return [(path, size) for _, path, size in sorted(found)]
|
||
|
|
|
||
|
|
|
||
|
|
def retire(directory, budget_bytes):
|
||
|
|
"""Delete the oldest slices until the pool fits `budget_bytes`, never the newest:
|
||
|
|
an empty pool would starve training. Returns the deleted paths."""
|
||
|
|
slices = ready_slices(directory)
|
||
|
|
total = sum(size for _, size in slices)
|
||
|
|
gone = []
|
||
|
|
for path, size in slices[:-1]:
|
||
|
|
if total <= budget_bytes:
|
||
|
|
break
|
||
|
|
shutil.rmtree(path, ignore_errors=True)
|
||
|
|
total -= size
|
||
|
|
gone.append(path)
|
||
|
|
return gone
|
||
|
|
|
||
|
|
|
||
|
|
def clear_partials(directory):
|
||
|
|
"""Remove half-built slices a stopped producer left behind."""
|
||
|
|
for name in os.listdir(directory) if os.path.isdir(directory) else ():
|
||
|
|
if name.startswith(".") and name.endswith(".partial"):
|
||
|
|
shutil.rmtree(os.path.join(directory, name), ignore_errors=True)
|
||
|
|
|
||
|
|
|
||
|
|
def open_slice(path):
|
||
|
|
"""(shards {(day, satellite): Shard}, segments) of one ready slice. Raises
|
||
|
|
OSError if it was retired before every shard could be opened."""
|
||
|
|
info = json.load(open(os.path.join(path, SLICE_FILE)))
|
||
|
|
shards = {}
|
||
|
|
try:
|
||
|
|
for name in os.listdir(path):
|
||
|
|
if name.endswith(".shard"):
|
||
|
|
day, satellite = name[:-len(".shard")].rsplit("_g", 1)
|
||
|
|
shards[(day, int(satellite))] = dataset.Shard(os.path.join(path, name))
|
||
|
|
except BaseException:
|
||
|
|
for shard in shards.values():
|
||
|
|
shard.close()
|
||
|
|
raise
|
||
|
|
return shards, [tuple(pair) for pair in info["segments"]]
|
||
|
|
|
||
|
|
|
||
|
|
def slice_name(counter, piece):
|
||
|
|
"""Sortable, unique and self-describing: when made, and what it covers."""
|
||
|
|
day = datetime.datetime.fromtimestamp(piece[0], datetime.timezone.utc)
|
||
|
|
return f"{counter:07d}_{day:%Y%m%dT%H%M}_{(piece[1] - piece[0]) // paths.CADENCE}"
|