noaa-goes-visualization/migrate_unrename.py

367 lines
16 KiB
Python
Raw Permalink Normal View History

#!/usr/bin/env python
"""One-time migration: strip _f/_e suffixes from the archive, into the SQLite index.
``filter_FITS.py`` used to record its verdict by renaming files -- ``_f`` for passed,
``_e`` for rejected -- and writing a ``_e.jpg`` diagnostic plot alongside rejects.
That made the archive's filenames a mutable database with one column, coupled three
scripts to a naming convention, and threw away every intermediate score.
This restores every file to the name NOAA published it under, after recording the
existing verdicts in the index.
Safety, because this is the one irreversible step:
* Dry run by default; ``--apply`` is required to rename anything.
* The legacy verdicts are exported to a standalone CSV **before** any rename. Every
other table can be rebuilt by re-walking or re-reading the archive, but once the
suffixes are gone these labels exist nowhere else.
* A rename that would overwrite an existing file is skipped and reported, never
forced.
* Idempotent and resumable: re-running finds nothing left to do.
Note that these labels are stale -- they were written in May 2024, and the filter was
retuned that July -- so they are preserved as history, under the run name
``legacy_filter``, not as ground truth.
migrate_unrename.py --wavelength 171 --year 2024 # dry run, one subtree
migrate_unrename.py --apply # the whole archive
"""
import argparse
import csv
import gzip
import os
import sys
import time
from collections import Counter
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from suvi import db, paths, vfs
LEGACY_RUN_NAME = "legacy_filter"
def chunk_key(wavelength, year):
"""Index key recording that one (band, year) has been fully migrated."""
return f"unrename_done:ci{wavelength:03d}:{year if year is not None else 'all'}"
def walk_archive(root, satellites, wavelengths, years=None):
"""Yield (directory, filename, FrameName) for every SUVI frame found.
When `years` is given the unwanted year directories are pruned from the walk
rather than merely skipped, so restricting the year genuinely restricts how much
of the filesystem is touched. That matters here for more than speed: see
:func:`release_handles`.
"""
def fail(error):
# os.walk swallows errors by default. On this mount a transient ENFILE
# would then make the walk yield nothing, and the caller would conclude the
# archive was already migrated and move on -- silently skipping real work.
# An unreadable directory has to stop the chunk, not look like an empty one.
raise error
for base in paths.wavelength_dirs(root, satellites, wavelengths):
if not os.path.isdir(base):
continue
for directory, dirnames, filenames in os.walk(base, onerror=fail):
if years and directory == base:
dirnames[:] = [d for d in dirnames if not d.isdigit() or int(d) in years]
for filename in filenames:
name = paths.parse_frame_filename(filename)
if name is not None:
yield directory, filename, name
def archive_years(root, satellites, wavelength):
"""Years present on disk for one band, across all satellites."""
years = set()
for satellite in satellites:
base = os.path.join(
root, f"goes{satellite:d}", "l2", "data", f"suvi-l2-ci{wavelength:03d}"
)
if not os.path.isdir(base):
continue
try:
years.update(int(d) for d in os.listdir(base) if d.isdigit())
except OSError:
continue
return sorted(years)
#: Frames a chunk must have touched before reclaiming is worth its cost. Well under
#: the point where the mount starts refusing opens, and far above anything a test
#: fixture reaches.
RELIEF_THRESHOLD = 50_000
#: Reclaim this often *within* a phase as well. Chunking alone is not enough: one
#: band-year can hold 230k frames, and the mount began refusing opens partway through
#: renaming it, costing 132,857 renames in an earlier attempt.
RELIEF_INTERVAL = vfs.RELIEF_INTERVAL
# Reclaim lives in suvi.vfs: it is a property of this mount, not of this migration,
# and the index build needs it just as much. Re-exported so callers and tests can
# reach it by either name.
drop_caches = vfs.drop_caches
reclaimable_kb = vfs.reclaimable_kb
release_handles = vfs.release_handles
def _mark_chunk_done(args, wavelengths):
"""Record that a (band, year) needs no further migration.
Lets a resume skip finished subtrees without touching the filesystem at all --
the traversal that a resume used to perform is exactly what exhausts this mount.
"""
if not args.apply or len(wavelengths) != 1:
return
year = args.year[0] if args.year and len(args.year) == 1 else None
try:
conn = db.connect(args.db)
db.set_meta(conn, chunk_key(wavelengths[0], year), "done")
conn.commit()
conn.close()
except Exception as exc:
print(f" (could not record chunk progress: {exc})")
def export_labels(path, records):
"""Write the legacy verdicts to a standalone file before anything is renamed."""
os.makedirs(os.path.dirname(os.path.abspath(path)), exist_ok=True)
with gzip.open(path, "wt", newline="") as handle:
writer = csv.writer(handle)
writer.writerow(["relpath", "legacy_label", "satellite", "wavelength", "t_start"])
for relpath, label, name in records:
writer.writerow([relpath, label, name.satellite, name.wavelength, name.t_start])
def main(argv=None):
parser = argparse.ArgumentParser(description=__doc__,
formatter_class=argparse.RawDescriptionHelpFormatter)
parser.add_argument("--db", default=None, help="target index (default: $SUVI_DB)")
parser.add_argument("--root", default=None, help="archive root (default: $SUVI_DATA_ROOT)")
parser.add_argument("--apply", action="store_true", help="rename files (default: dry run)")
parser.add_argument("--satellites", default="16,18,19")
parser.add_argument("--wavelength", default=None, help="restrict to one band")
parser.add_argument("--year", type=int, action="append", default=None,
help="restrict to given years (repeatable)")
parser.add_argument("--export", default=None,
help="legacy label export (default: <root>/legacy_labels.csv.gz)")
parser.add_argument("--keep-plots", action="store_true",
help="keep the _e.jpg diagnostic images")
parser.add_argument("--no-relief", action="store_true",
help="skip the cache-reclaim step between chunks")
parser.add_argument("--recheck", action="store_true",
help="re-examine chunks already recorded as migrated")
args = parser.parse_args(argv)
root = args.root or paths.data_root()
if not os.path.isdir(root):
print(f"No such archive root: {root}")
return 1
satellites = tuple(int(part) for part in args.satellites.split(","))
wavelengths = (int(args.wavelength),) if args.wavelength else paths.WAVELENGTHS
# One (band, year) at a time. Walking the whole archive in a single pass fills
# the dentry cache with millions of entries, and on this virtiofs mount each
# cached entry pins a host file handle -- the walk succeeds and then nothing can
# open a file again, which is how an earlier attempt lost 152,193 renames to
# ENFILE. A band-year is ~130k files, comfortably inside what the mount
# sustains, and each chunk finishes completely (export, record, rename) so a
# failure never leaves labels destroyed but unrecorded.
chunks = []
for wavelength in wavelengths:
years = archive_years(root, satellites, wavelength) if args.year is None else args.year
for year in years or [None]:
chunks.append((wavelength, year))
# Skip chunks already recorded as complete, without touching the filesystem.
# Resuming used to re-walk finished subtrees just to discover there was nothing
# to do; across several passes that traversed the archive roughly four times
# over, and traversal is the very thing that exhausts this mount.
done = set()
# Only when applying: a dry run must neither create the index nor hide work, so
# it reports the archive's true state regardless of what progress was recorded.
if args.apply and not args.recheck:
try:
conn = db.connect(args.db)
done = {
key for (wavelength, year) in chunks
if db.get_meta(conn, key := chunk_key(wavelength, year)) == "done"
}
conn.close()
except Exception as exc:
print(f"Could not read migration progress ({exc}); checking every chunk.")
if len(chunks) > 1:
overall = 0
for index, (wavelength, year) in enumerate(chunks, 1):
label = f"{wavelength}A {year or 'all years'}"
if chunk_key(wavelength, year) in done:
print(f"\n=== chunk {index}/{len(chunks)}: {label} -- already migrated, skipping")
continue
print(f"\n{'=' * 70}\n=== chunk {index}/{len(chunks)}: {label} ===")
chunk_args = argparse.Namespace(**vars(args))
chunk_args.wavelength = str(wavelength)
chunk_args.year = [year] if year is not None else None
if args.export is None:
suffix = f"ci{wavelength:03d}" + (f"_{year}" if year else "")
chunk_args.export = os.path.join(root, f"legacy_labels_{suffix}.csv.gz")
status, frames = main_one_chunk(chunk_args, root, satellites, (wavelength,))
overall = overall or status
if not args.no_relief and frames >= RELIEF_THRESHOLD:
release_handles()
return overall
return main_one_chunk(args, root, satellites, wavelengths)[0]
def main_one_chunk(args, root, satellites, wavelengths):
print(f"Scanning {root}")
print(f" satellites={satellites} wavelengths={wavelengths} years={args.year or 'all'}")
counts = Counter()
labelled = [] # (relpath, label, name) for every file carrying a verdict
renames = [] # (src, dst) pairs to perform
collisions = []
plots = []
started = time.time()
for directory, filename, name in walk_archive(root, satellites, wavelengths, args.year):
counts["frames"] += 1
relpath = name.relpath()
if name.label:
counts[f"label_{name.label}"] += 1
labelled.append((relpath, name.label, name))
source = os.path.join(directory, filename)
target = os.path.join(directory, name.filename())
if os.path.exists(target):
collisions.append((source, target))
else:
renames.append((source, target))
if name.label == "e" and not args.keep_plots:
plot = os.path.join(directory, name.error_plot_name())
if os.path.exists(plot):
plots.append(plot)
else:
counts["unlabelled"] += 1
if counts["frames"] % RELIEF_INTERVAL == 0 and not args.no_relief:
release_handles()
if counts["frames"] % 200000 == 0:
print(f" {counts['frames']} frames in {time.time() - started:.0f}s")
print(f"\nFound {counts['frames']} frames in {time.time() - started:.0f}s")
print(f" passed (_f): {counts['label_f']}")
print(f" rejected(_e): {counts['label_e']}")
print(f" unlabelled : {counts['unlabelled']} (never processed by any filter)")
print(f" renames to perform: {len(renames)}")
print(f" diagnostic plots to delete: {len(plots)}")
if collisions:
print(f" COLLISIONS (will be skipped): {len(collisions)}")
for source, target in collisions[:5]:
print(f" {os.path.basename(source)} -> {os.path.basename(target)} exists")
if not args.apply:
print("\nDry run. Re-run with --apply to write the export, rename, and clean up.")
return 0, counts["frames"]
if not labelled:
print("\nNothing labelled; archive is already migrated.")
_mark_chunk_done(args, wavelengths)
return 0, counts["frames"]
# 1. Export the labels before touching a single filename.
export_path = args.export or os.path.join(root, "legacy_labels.csv.gz")
print(f"\nExporting {len(labelled)} legacy labels to {export_path}")
export_labels(export_path, labelled)
if not os.path.exists(export_path) or os.path.getsize(export_path) == 0:
print("ERROR: export is missing or empty; refusing to rename anything.")
return 1, counts["frames"]
print(f" {os.path.getsize(export_path) / 1024 / 1024:.1f} MB written")
# 2. Record them in the index.
target_db = args.db or paths.default_db_path()
conn = db.connect(target_db)
run_id = db.create_detector_run(
conn,
LEGACY_RUN_NAME,
{"source": "filename suffixes", "root": root},
notes="Verdicts recovered from _f/_e filename suffixes. Written ~May 2024 by a "
"filter version since retuned; preserved as history, not ground truth.",
)
print(f"Recording verdicts in {target_db} as run {run_id} ({LEGACY_RUN_NAME})")
batch = []
for index, (relpath, label, name) in enumerate(labelled):
frame_id = db.upsert_frame(conn, name, relpath)
batch.append((frame_id, "good" if label == "f" else "bad", f"legacy_{label}", None, None))
if len(batch) >= 5000:
db.record_detections(conn, run_id, batch)
conn.commit()
batch.clear()
print(f" {index + 1}/{len(labelled)}")
if batch:
db.record_detections(conn, run_id, batch)
conn.commit()
recorded = conn.execute(
"SELECT count(*) c FROM detection WHERE run_id = ?", (run_id,)
).fetchone()["c"]
if recorded != len(labelled):
print(f"ERROR: recorded {recorded} verdicts but found {len(labelled)} labels.")
print(" Not renaming anything.")
conn.close()
return 1, counts["frames"]
print(f" {recorded} verdicts recorded")
# 3. Only now rename.
print(f"\nRenaming {len(renames)} files...")
renamed = failed = 0
for index, (source, target) in enumerate(renames):
try:
os.rename(source, target)
renamed += 1
except OSError as exc:
failed += 1
if failed <= 5:
print(f" FAILED {source}: {exc}")
if index and index % RELIEF_INTERVAL == 0 and not args.no_relief:
release_handles()
if index and index % 100000 == 0:
print(f" {index}/{len(renames)}")
print(f" renamed {renamed}, failed {failed}, skipped {len(collisions)} collisions")
removed = 0
for plot in plots:
try:
os.remove(plot)
removed += 1
except OSError:
pass
print(f" deleted {removed} diagnostic plots")
# 4. Verify every recorded path now resolves.
print("\nVerifying...")
missing = 0
for relpath, _, _ in labelled[:: max(1, len(labelled) // 5000)]:
if not os.path.exists(paths.abspath(relpath, root)):
missing += 1
if missing:
print(f" WARNING: {missing} sampled paths do not resolve on disk")
else:
print(" all sampled paths resolve")
conn.close()
# Only claim the chunk is finished if every rename actually landed. Marking it
# done after partial failures is how a resume would skip real remaining work.
if failed == 0:
_mark_chunk_done(args, wavelengths)
else:
print(f" NOT marking this chunk complete: {failed} renames failed; re-run to retry.")
print("\nDone. The archive now uses NOAA's original filenames; verdicts live in "
f"the index, and a recovery copy is at {export_path}")
return 0, counts["frames"]
if __name__ == "__main__":
sys.exit(main())