aetherscan 1.0.0__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
aetherscan/db/db.py ADDED
@@ -0,0 +1,2789 @@
1
+ # TODO: add warning when aetherscan.db file size exceeds 1TB (at that point we should consider migrating to PostgreSQL or sharding the SQLite db)
2
+ # TODO: add functions to delete entries in db based on query results (e.g. query_system_resource -> DELETE)
3
+ """
4
+ Database for Aetherscan Pipeline
5
+ Uses SQLite with asynchronous queue-based writes to handle concurrent data collection from multiple
6
+ processes safely
7
+ """
8
+
9
+ from __future__ import annotations
10
+
11
+ import getpass
12
+ import json
13
+ import logging
14
+ import os
15
+ import socket
16
+ import sqlite3
17
+ import threading
18
+ import time
19
+ from contextlib import contextmanager
20
+ from queue import Empty, Full, Queue
21
+ from typing import Any
22
+
23
+ import numpy as np
24
+
25
+ from aetherscan.config import get_config
26
+ from aetherscan.manager import register_db
27
+
28
+ logger = logging.getLogger(__name__)
29
+
30
+ # Unique sentinel object (i.e. an object with no attributes, no methods beyond those inherited from
31
+ # Python's base object type, and no meaningful equality semantics except identity) for flush
32
+ # requests - writer thread recognizes this as a command to flush immediately rather than data to be
33
+ # written
34
+ _FLUSH_SENTINEL = object()
35
+
36
+ # Command sentinel for mark_superseded() requests routed through the writer thread (same
37
+ # identity-based recognition as _FLUSH_SENTINEL, so single-writer semantics are preserved)
38
+ _MARK_SUPERSEDED_SENTINEL = object()
39
+
40
+ # Current schema version, stored in SQLite's PRAGMA user_version (0 on any pre-versioning db).
41
+ # Bump this and extend _migrate_schema() with an `if version < N:` block for future changes.
42
+ # v1: added `superseded INTEGER DEFAULT 0` to training_stats, injection_stats,
43
+ # latent_snapshots, and inference_results (stale-data supersede semantics on retry)
44
+ # v2: added the `inference_cadences` table (per-cadence inference run manifest driving
45
+ # stage-aware retries). New tables need no ALTER step — the CREATE TABLE IF NOT EXISTS
46
+ # statements in _init_database() run before migration for old and new databases alike —
47
+ # so the version bump exists to record the change and keep future `if version < N:`
48
+ # blocks ordered.
49
+ # v3: added `inference_cadences.config_fingerprint` so the stage-aware resume only skips a
50
+ # cadence whose stored 'inferred' row was written under the same inference config — guards
51
+ # the reused-tag-with-changed-config stale-reuse footgun (the inference counterpart of the
52
+ # training-side config_fingerprint guard).
53
+ # v4: added the `pipeline_stages` table (always-on stage timing spans from
54
+ # aetherscan.benchmark, consumed by utils/benchmark_report.py and the monitor's
55
+ # stage-band overlay). New table -> no ALTER step, same as v2.
56
+ # v5: added `inference_results.screening_proba` / `mc_mean` / `mc_std` (#282 two-pass
57
+ # inference: the deterministic pass-1 score plus the seeded MC mean/spread that carries
58
+ # the science threshold for survivors).
59
+ # v6: added `training_stats.is_finite` (#289) — a NaN stat value binds as SQL NULL, violated
60
+ # the NOT NULL constraint, and the failed executemany silently dropped the whole flush
61
+ # batch. Non-finite values now store as 0.0 with is_finite=0 (the injection_stats
62
+ # semantics), and query_training_stat filters them by default.
63
+ # v7: index sweep against the real production query shapes. Both changes share one lesson:
64
+ # with equality on the leading columns and the range (timestamp) trailing, SQLite's
65
+ # default cost model picks the index with no ANALYZE stats — db.py never runs ANALYZE —
66
+ # while a range-column-buried shape loses to a (tag, timestamp, ...) index absent
67
+ # sqlite_stat1.
68
+ # - injection_stats: added idx_injection_stats_by_stat (tag, stat_name, signal_type,
69
+ # injection_stage, timestamp) — the end-of-run plot pass scanned the whole tag
70
+ # partition ~165x, ~6 h projected at release scale; a round_number-trailing variant
71
+ # was measured to never be chosen (benchmarks/bench_injection_index.py).
72
+ # - latent_snapshots: idx_latent_snapshots_filter (tag, timestamp, ...) reshaped to
73
+ # idx_latent_snapshots_by_key (tag, round_number, epoch_number, step_number,
74
+ # model_name, timestamp) — the GIF pass loads one capture per frame (up to 500
75
+ # queries/run), each a whole-window scan under the old shape: measured 67x per
76
+ # frame at 2M rows (~1.5 h -> ~2 s per GIF pass at release scale), and the
77
+ # dashboard's latest-capture lookup becomes a backward index walk instead of a
78
+ # partition sort, at no measurable write cost
79
+ # (benchmarks/bench_db_index_shapes.py).
80
+ # The remaining tables were measured or reasoned through and deliberately left alone —
81
+ # notably a training_stats reshape to (tag, model_name, stat_name, timestamp) was
82
+ # benchmarked and REJECTED: it regressed the dominant loss-curve fetch 1.8x and cost
83
+ # ~20% insert throughput (scattered stat_name subtree writes vs append-only timestamp
84
+ # order) to speed only the minor single-stat shape on a table whose tag partitions
85
+ # stay ~58k rows.
86
+ # v8: added idx_inference_results_supersede — a partial index on (tag, npy_path) WHERE
87
+ # superseded = 0, matching _execute_mark_superseded's exact predicate. The only prior
88
+ # index led with (tag, timestamp, ...), so the once-per-cadence supersede UPDATE (which
89
+ # BLOCKS the inference thread via the writer-queue sentinel) visited every live row of
90
+ # the tag partition — a quadratic-in-catalog term on RFI-dense 6k-cadence runs (#301).
91
+ _SCHEMA_VERSION = 8
92
+
93
+
94
+ # Per-process cache for get_system_metadata(): every field (hostname, user, outbound IP, PID)
95
+ # is constant for the life of a process, and the uncached version opened a UDP socket per call
96
+ # — a real cost on the injection-stat hot path, which calls this tens of millions of times per
97
+ # round (#277). A spawned child process re-imports the module and rebuilds its own cache, so
98
+ # the PID is always correct.
99
+ _SYSTEM_METADATA_CACHE: str | None = None
100
+
101
+
102
+ def get_system_metadata() -> str:
103
+ """
104
+ Collects system metadata (machine name, user name, IP address, process ID)
105
+ and returns it as a JSON string suitable for database storage. Cached per process —
106
+ every field is process-constant.
107
+ """
108
+ global _SYSTEM_METADATA_CACHE
109
+ if _SYSTEM_METADATA_CACHE is not None:
110
+ return _SYSTEM_METADATA_CACHE
111
+
112
+ # Machine and user info
113
+ machine_name = socket.gethostname()
114
+ user_name = getpass.getuser()
115
+
116
+ # IP Address
117
+ def _get_ip_address():
118
+ s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
119
+ try:
120
+ s.connect(
121
+ ("8.8.8.8", 80) # Doesn't have to be reachable; used to infer the outbound IP
122
+ )
123
+ return s.getsockname()[0]
124
+ except Exception:
125
+ return "Unknown"
126
+ finally:
127
+ s.close()
128
+
129
+ ip_address = _get_ip_address()
130
+
131
+ # Process ID
132
+ process_id = os.getpid()
133
+
134
+ # Pack into JSON
135
+ metadata = {
136
+ "machine_name": machine_name,
137
+ "user_name": user_name,
138
+ "ip_address": ip_address,
139
+ "process_id": process_id,
140
+ }
141
+
142
+ # Use sorted keys for deterministic ordering (optional, good for diffs)
143
+ _SYSTEM_METADATA_CACHE = json.dumps(metadata, sort_keys=True)
144
+ return _SYSTEM_METADATA_CACHE
145
+
146
+
147
+ class Database:
148
+ """
149
+ Thread-safe SQLite database with asynchronous queue-based writes.
150
+
151
+ Multiple threads/processes push records to a shared queue; a single background writer thread
152
+ drains it and commits to SQLite at a configured interval. Serializing writes through one
153
+ thread eliminates SQLITE_BUSY contention that would otherwise arise from concurrent writers.
154
+ """
155
+
156
+ _instance = None # Stores singleton instance
157
+ _lock = threading.Lock() # Ensures thread safety on object initialization
158
+
159
+ # __new__ allocates the object in memory (constructor at the object-creation level)
160
+ # __init__ initializes the object's attributes after it's created
161
+ # since __new__ is called before __init__ every time we instantiate a class,
162
+ # by overriding __new__, we can short-circuit object creation entirely, and control whether a
163
+ # new instance is created, or just return the existing instance
164
+ def __new__(cls):
165
+ # Double-checked locking pattern:
166
+ # First check if _instance is None, without lock (for performance)
167
+ if cls._instance is None:
168
+ # If None, acquire the lock to serialize the initialization path,
169
+ # preventing race conditions (2 threads violating singleton semantics)
170
+ with cls._lock:
171
+ # Check if _instance is None again inside the lock
172
+ # (since multiple threads can be calling simultaneously)
173
+ if cls._instance is None:
174
+ # If still None, only then we construct the singleton instance
175
+ cls._instance = super().__new__(cls)
176
+ cls._instance._initialized = False # Mark as not initialized (for __init__)
177
+ # Return the same instance for all subsequent constructor calls
178
+ return cls._instance
179
+
180
+ def __init__(self):
181
+ """Initialize database"""
182
+ # Note, __init__ is triggered every time the class's constructor is called,
183
+ # even if __new__ returned the existing singleton instance
184
+ # Hence, we use the _initialized flag to make sure __init__ only runs once
185
+ if self._initialized:
186
+ return
187
+
188
+ self._initialized = True
189
+
190
+ self.config = get_config()
191
+ if self.config is None:
192
+ raise ValueError("get_config() returned None")
193
+
194
+ self.db_path = os.path.join(self.config.output_path, "db", "aetherscan.db")
195
+ os.makedirs(os.path.dirname(self.db_path), exist_ok=True) # Create dir if it doesn't exist
196
+
197
+ self.get_connection_timeout = self.config.db.get_connection_timeout
198
+ self.stop_writer_timeout = self.config.db.stop_writer_timeout
199
+ self.write_interval = self.config.db.write_interval
200
+ self.write_buffer_max_size = self.config.db.write_buffer_max_size
201
+ self.write_retry_delay = self.config.db.write_retry_delay
202
+ self.flush_timeout = self.config.db.flush_timeout
203
+ self.bulk_chunk_rows = self.config.db.bulk_chunk_rows
204
+ self.stop_drain_timeout = self.config.db.stop_drain_timeout
205
+
206
+ self.write_queue = Queue()
207
+ # Bulk lane (#277): bounded queue for high-volume injection-stat chunks. The bound
208
+ # applies backpressure to the enqueuer (the round-data drainer / in-process data-gen
209
+ # stats callback — background work that can afford to wait), never the training path,
210
+ # and caps writer-side memory (the old single unbounded queue grew to ~35 GB of
211
+ # anonymous RSS on the release run). flush() only drains the foreground write_queue,
212
+ # so plot flushes no longer wait behind a round's worth of injection rows.
213
+ self.bulk_queue = Queue(maxsize=self.config.db.bulk_queue_max_items)
214
+ # Pending bulk rows per round_number (None keyed as -1), so plot code can verify that
215
+ # every injection row for rounds <= N is committed before rendering round N's figures
216
+ self._bulk_pending_lock = threading.Lock()
217
+ self._bulk_pending_by_round: dict[int, int] = {}
218
+ self.writer_thread = None
219
+ self.stop_event = threading.Event() # Thread-safe flag for stopping
220
+ self.drain_event = threading.Event() # Graceful-shutdown flag: drain both lanes first
221
+
222
+ # Initialize database schema
223
+ self._init_database()
224
+
225
+ logger.info(f"Database initialized at: {self.db_path}")
226
+ # Startup logs only O(1) facts. The per-table COUNT(*) summary that used to print
227
+ # here (get_db_stats) cost a measured ~13 min per cold-cache process launch on a
228
+ # catalog-scale DB (80 GB, ~190M injection_stats rows) and was re-paid on every
229
+ # retry-loop relaunch (#301); get_db_stats() remains for on-demand diagnostics.
230
+ with self._get_connection() as conn:
231
+ db_size_bytes = conn.execute(
232
+ "SELECT page_count * page_size FROM pragma_page_count(), pragma_page_size()"
233
+ ).fetchone()[0]
234
+ logger.info(f" db_size_bytes: {db_size_bytes}")
235
+ logger.info(f"Write interval: {self.write_interval} seconds")
236
+ logger.info(f"Max buffer size: {self.write_buffer_max_size} records")
237
+
238
+ @classmethod
239
+ def _reset(cls):
240
+ """
241
+ Teardown hook for the thread-safe singleton — discards the cached instance so the next
242
+ constructor call yields a fresh one. Only safe to call after stop() has completed; calling
243
+ it while the database is active leaves live threads holding a stale reference.
244
+ """
245
+ # Acquire lock to prevent race conditions
246
+ with cls._lock:
247
+ # Discard the singleton instance by removing the global reference
248
+ # Guarantees the next constructor call will produce a fresh instance
249
+ # Note, resources held by the old instance will remain alive unless explicitly closed beforehand
250
+ cls._instance = None
251
+ logger.info("Database singleton instance reset")
252
+
253
+ # NOTE: should we consider migrating to PostgreSQL? or should we go all in on SQLite?
254
+ def _init_database(self):
255
+ """Create database tables if they don't exist, then run schema migrations"""
256
+ with self._get_connection() as conn:
257
+ cursor = conn.cursor()
258
+
259
+ # System resources table
260
+ cursor.execute("""
261
+ CREATE TABLE IF NOT EXISTS system_resources (
262
+ id INTEGER PRIMARY KEY AUTOINCREMENT,
263
+ timestamp REAL NOT NULL,
264
+ resource_type TEXT NOT NULL,
265
+ resource_name TEXT NOT NULL,
266
+ value REAL NOT NULL,
267
+ unit TEXT,
268
+ tag TEXT,
269
+ metadata TEXT
270
+ )
271
+ """)
272
+
273
+ # NOTE: removing ORDER BY to reduce indexing costs
274
+ # # Composite index optimized for query_system_resource ORDER BY pattern
275
+ # # Also works for common filter patterns (tag + timestamp)
276
+ # # Recall that composite indices follow the left-prefix rule
277
+ # # That is, if your query contains a left prefix of the index columns,
278
+ # # you'll still get (some of) the benefits of indexing
279
+ # cursor.execute("""
280
+ # CREATE INDEX IF NOT EXISTS idx_system_resources_query
281
+ # ON system_resources(tag, timestamp, resource_type, resource_name)
282
+ # """)
283
+
284
+ # Composite index for common filter patterns (tag + timestamp)
285
+ cursor.execute("""
286
+ CREATE INDEX IF NOT EXISTS idx_system_resources_filter
287
+ ON system_resources(tag, timestamp)
288
+ """)
289
+
290
+ # Injection statistics table
291
+ cursor.execute("""
292
+ CREATE TABLE IF NOT EXISTS injection_stats (
293
+ id INTEGER PRIMARY KEY AUTOINCREMENT,
294
+ timestamp REAL NOT NULL,
295
+ stat_name TEXT NOT NULL,
296
+ value REAL NOT NULL,
297
+ round_number INTEGER,
298
+ chunk_number INTEGER,
299
+ sample_index INTEGER,
300
+ background_index INTEGER,
301
+ signal_class TEXT,
302
+ signal_type TEXT,
303
+ injection_stage TEXT,
304
+ is_finite INTEGER DEFAULT 1,
305
+ slope_clamped INTEGER DEFAULT 0,
306
+ tag TEXT,
307
+ metadata TEXT,
308
+ superseded INTEGER DEFAULT 0
309
+ )
310
+ """)
311
+
312
+ # NOTE: removing ORDER BY to reduce indexing costs
313
+ # # Composite index optimized for query_injection_stat ORDER BY pattern
314
+ # cursor.execute("""
315
+ # CREATE INDEX IF NOT EXISTS idx_injection_stats_query
316
+ # ON injection_stats(tag, signal_class, signal_type, round_number, chunk_number, sample_index, stat_name)
317
+ # """)
318
+
319
+ # Composite index for common filter pattern (tag + timestamp + stat_name + signal_type + injection_stage)
320
+ cursor.execute("""
321
+ CREATE INDEX IF NOT EXISTS idx_injection_stats_filter
322
+ ON injection_stats(tag, timestamp, stat_name, signal_type, injection_stage)
323
+ """)
324
+
325
+ # Secondary composite index for the stat-scoped query shape (schema v7): the
326
+ # end-of-run plot pass filters on (tag, stat_name, signal_type, injection_stage)
327
+ # with run-wide timestamp bounds, which the (tag, timestamp, ...) index above can
328
+ # only answer by scanning the tag's whole timestamp range — ~165 times per call.
329
+ # Leading with the equality columns turns each of those into a seek; round_number
330
+ # trails so round-scoped queries seek too. No query changes needed — SQLite's
331
+ # planner picks the better index per query.
332
+ cursor.execute("""
333
+ CREATE INDEX IF NOT EXISTS idx_injection_stats_by_stat
334
+ ON injection_stats(tag, stat_name, signal_type, injection_stage, timestamp)
335
+ """)
336
+
337
+ # Training statistics table
338
+ cursor.execute("""
339
+ CREATE TABLE IF NOT EXISTS training_stats (
340
+ id INTEGER PRIMARY KEY AUTOINCREMENT,
341
+ timestamp REAL NOT NULL,
342
+ model_name TEXT NOT NULL,
343
+ stat_name TEXT NOT NULL,
344
+ value REAL NOT NULL,
345
+ round_number INTEGER,
346
+ epoch_number INTEGER,
347
+ tag TEXT,
348
+ metadata TEXT,
349
+ superseded INTEGER DEFAULT 0,
350
+ is_finite INTEGER DEFAULT 1
351
+ )
352
+ """)
353
+
354
+ # NOTE: removing ORDER BY to reduce indexing costs
355
+ # # Composite index optimized for query_training_stat ORDER BY pattern
356
+ # cursor.execute("""
357
+ # CREATE INDEX IF NOT EXISTS idx_training_stats_query
358
+ # ON training_stats(tag, model_name, round_number, epoch_number, stat_name)
359
+ # """)
360
+
361
+ # Composite index for common filter pattern (tag + timestamp + model_name + stat_name).
362
+ # Deliberately kept through the v7 index sweep: an equality-first reshape
363
+ # (tag, model_name, stat_name, timestamp) was benchmarked and rejected — it
364
+ # regressed the dominant run-window loss-curve fetch 1.8x (stat-sorted index
365
+ # order scatters the table-row lookups that timestamp order visits
366
+ # sequentially) and cost ~20% insert throughput, to speed only the minor
367
+ # single-stat shape on ~58k-row tag partitions
368
+ # (benchmarks/bench_db_index_shapes.py).
369
+ cursor.execute("""
370
+ CREATE INDEX IF NOT EXISTS idx_training_stats_filter
371
+ ON training_stats(tag, timestamp, model_name, stat_name)
372
+ """)
373
+
374
+ # Latent snapshots table
375
+ cursor.execute("""
376
+ CREATE TABLE IF NOT EXISTS latent_snapshots (
377
+ id INTEGER PRIMARY KEY AUTOINCREMENT,
378
+ timestamp REAL NOT NULL,
379
+ model_name TEXT NOT NULL,
380
+ round_number INTEGER NOT NULL,
381
+ epoch_number INTEGER NOT NULL,
382
+ step_number INTEGER NOT NULL,
383
+ cadence_index INTEGER NOT NULL,
384
+ signal_type TEXT NOT NULL,
385
+ latent_vector TEXT NOT NULL,
386
+ snr_base INTEGER,
387
+ snr_range INTEGER,
388
+ tag TEXT,
389
+ metadata TEXT,
390
+ superseded INTEGER DEFAULT 0
391
+ )
392
+ """)
393
+
394
+ # Capture-key composite index (schema v7, replacing the (tag, timestamp, ...)
395
+ # shape): the GIF pass loads one capture per frame — up to
396
+ # latent_viz_gif_max_frames (500) queries with equality on tag/round/epoch/
397
+ # step/model and the run window trailing — and the old shape answered each by
398
+ # scanning the tag's whole timestamp window (~100M rows at release scale, per
399
+ # frame). round/epoch/step lead so the dashboard's model-less latest-capture
400
+ # lookup (ORDER BY round DESC, epoch DESC, step DESC LIMIT 1) walks the index
401
+ # backward instead of sorting; timestamp trails as the range column
402
+ # (benchmarks/bench_db_index_shapes.py).
403
+ cursor.execute("""
404
+ CREATE INDEX IF NOT EXISTS idx_latent_snapshots_by_key
405
+ ON latent_snapshots(tag, round_number, epoch_number, step_number,
406
+ model_name, timestamp)
407
+ """)
408
+
409
+ # Inference results table
410
+ cursor.execute("""
411
+ CREATE TABLE IF NOT EXISTS inference_results (
412
+ id INTEGER PRIMARY KEY AUTOINCREMENT,
413
+ timestamp REAL NOT NULL,
414
+ npy_path TEXT NOT NULL,
415
+ snippet_index INTEGER NOT NULL,
416
+ prediction INTEGER NOT NULL,
417
+ confidence REAL NOT NULL,
418
+ latent_vector TEXT,
419
+ target TEXT,
420
+ session TEXT,
421
+ cadence_id INTEGER,
422
+ band TEXT,
423
+ frequency_mhz REAL,
424
+ timestamp_observed REAL,
425
+ h5_path TEXT,
426
+ tag TEXT,
427
+ metadata TEXT,
428
+ superseded INTEGER DEFAULT 0,
429
+ screening_proba REAL,
430
+ mc_mean REAL,
431
+ mc_std REAL
432
+ )
433
+ """)
434
+
435
+ # Composite index for common filter pattern (tag + confidence + prediction)
436
+ cursor.execute("""
437
+ CREATE INDEX IF NOT EXISTS idx_inference_results_filter
438
+ ON inference_results(tag, timestamp, confidence, prediction)
439
+ """)
440
+
441
+ # v8's idx_inference_results_supersede is created ONLY in _migrate_schema
442
+ # (unlike the indexes here): its partial-index predicate references the
443
+ # superseded column, which a pre-v1 database gains from the v1 ALTER — a
444
+ # CREATE here would run before that ALTER and fail with "no such column"
445
+ # on any legacy file. Fresh databases still get it at init because
446
+ # _migrate_schema always runs (version 0 -> current).
447
+
448
+ # Inference cadence manifest table (schema v2): one row per (cadence, stage
449
+ # transition) — status 'preprocessed' when the stamp .npy lands, a superseding
450
+ # 'inferred' row (with aggregate stats) when inference completes, and 'failed'
451
+ # when the inference stage of a cadence dies. Drives stage-aware retry resume:
452
+ # a live 'inferred' row means the cadence is skipped entirely on retry.
453
+ # cadence_key and confidence_summary are JSON TEXT.
454
+ cursor.execute("""
455
+ CREATE TABLE IF NOT EXISTS inference_cadences (
456
+ id INTEGER PRIMARY KEY AUTOINCREMENT,
457
+ timestamp REAL NOT NULL,
458
+ tag TEXT,
459
+ csv_path TEXT,
460
+ cadence_key TEXT,
461
+ npy_path TEXT NOT NULL,
462
+ status TEXT NOT NULL,
463
+ n_stamps INTEGER,
464
+ n_candidates INTEGER,
465
+ confidence_summary TEXT,
466
+ duration_s REAL,
467
+ config_fingerprint TEXT,
468
+ superseded INTEGER DEFAULT 0
469
+ )
470
+ """)
471
+
472
+ # Composite index for the resume lookup pattern (tag + npy_path + status)
473
+ cursor.execute("""
474
+ CREATE INDEX IF NOT EXISTS idx_inference_cadences_filter
475
+ ON inference_cadences(tag, npy_path, status)
476
+ """)
477
+
478
+ # Pipeline stage timing table (schema v4): one row per timed pipeline stage
479
+ # span, written by aetherscan.benchmark (stage_timer / record_stage). stage is
480
+ # a hierarchical dot-name ("train.round_02.data_generation"); metadata is an
481
+ # optional JSON TEXT blob (e.g. {"status": "failed", ...} for spans that ended
482
+ # in an exception). Retried stages simply append new rows — consumers see every
483
+ # attempt, each with its own span.
484
+ cursor.execute("""
485
+ CREATE TABLE IF NOT EXISTS pipeline_stages (
486
+ id INTEGER PRIMARY KEY AUTOINCREMENT,
487
+ stage TEXT NOT NULL,
488
+ start_time REAL NOT NULL,
489
+ end_time REAL NOT NULL,
490
+ duration_s REAL NOT NULL,
491
+ tag TEXT,
492
+ metadata TEXT
493
+ )
494
+ """)
495
+
496
+ # Composite index for the common filter pattern (tag + start_time)
497
+ cursor.execute("""
498
+ CREATE INDEX IF NOT EXISTS idx_pipeline_stages_filter
499
+ ON pipeline_stages(tag, start_time)
500
+ """)
501
+
502
+ conn.commit()
503
+
504
+ # Bring pre-existing databases (older schema) up to the current version
505
+ self._migrate_schema(conn)
506
+
507
+ # Enable Write-Ahead Logging (WAL) mode for better concurrent read performance
508
+ # WAL places writes in a separate log file so reads can still go through while writes happen
509
+ # The WAL log is periodically merged back into the main db (i.e. checkpointing)
510
+ cursor.execute("PRAGMA journal_mode=WAL")
511
+
512
+ logger.info("Database schema initialized with WAL mode")
513
+
514
+ def _migrate_schema(self, conn):
515
+ """
516
+ Minimal schema migration gated on SQLite's PRAGMA user_version.
517
+
518
+ Each block below upgrades one version step via additive ALTER TABLE ... ADD COLUMN
519
+ statements (the only in-place table change SQLite supports); a per-table column-existence
520
+ check keeps every step idempotent even if user_version was lost (e.g. a db file copied
521
+ without its journal). Fresh databases already get the full schema from the CREATE TABLE
522
+ statements in _init_database(), so migration only stamps their version.
523
+ """
524
+ cursor = conn.cursor()
525
+ version = cursor.execute("PRAGMA user_version").fetchone()[0]
526
+
527
+ if version >= _SCHEMA_VERSION:
528
+ return
529
+
530
+ if version < 1:
531
+ # v1: stale-data supersede semantics — rows from failed/superseded attempts are
532
+ # flagged (never deleted) so default queries can filter them out
533
+ for table in (
534
+ "training_stats",
535
+ "injection_stats",
536
+ "latent_snapshots",
537
+ "inference_results",
538
+ ):
539
+ columns = {row[1] for row in cursor.execute(f"PRAGMA table_info({table})")}
540
+ if "superseded" not in columns:
541
+ cursor.execute(f"ALTER TABLE {table} ADD COLUMN superseded INTEGER DEFAULT 0")
542
+ logger.info(f"Schema migration: added {table}.superseded")
543
+
544
+ # v2 (inference_cadences table) needs no migration step here: the table is created for
545
+ # old and new databases alike by the CREATE TABLE IF NOT EXISTS statements in
546
+ # _init_database(), which always runs before this method.
547
+
548
+ if version < 3:
549
+ # v3: config_fingerprint on inference_cadences. A fresh db already has the column
550
+ # from the CREATE TABLE above; this ALTER patches a db that created the v2 table
551
+ # before the column existed. The column-existence check keeps it idempotent.
552
+ columns = {row[1] for row in cursor.execute("PRAGMA table_info(inference_cadences)")}
553
+ if "config_fingerprint" not in columns:
554
+ cursor.execute("ALTER TABLE inference_cadences ADD COLUMN config_fingerprint TEXT")
555
+ logger.info("Schema migration: added inference_cadences.config_fingerprint")
556
+
557
+ # v4 (pipeline_stages table) needs no migration step here either: like v2, the table is
558
+ # created for old and new databases alike by the CREATE TABLE IF NOT EXISTS statement in
559
+ # _init_database(). Only the version stamp below advances.
560
+
561
+ if version < 5:
562
+ # v5: two-pass inference columns (#282). A fresh db already has them from the
563
+ # CREATE TABLE above; the ALTERs patch a pre-v5 db. Column-existence checks keep
564
+ # each step idempotent.
565
+ columns = {row[1] for row in cursor.execute("PRAGMA table_info(inference_results)")}
566
+ for column in ("screening_proba", "mc_mean", "mc_std"):
567
+ if column not in columns:
568
+ cursor.execute(f"ALTER TABLE inference_results ADD COLUMN {column} REAL")
569
+ logger.info(f"Schema migration: added inference_results.{column}")
570
+
571
+ if version < 6:
572
+ # v6: is_finite on training_stats (#289) — the injection_stats semantics. Existing
573
+ # rows were all finite by construction (a non-finite value could never be written:
574
+ # it bound as NULL and blew the NOT NULL constraint), so DEFAULT 1 is exact.
575
+ columns = {row[1] for row in cursor.execute("PRAGMA table_info(training_stats)")}
576
+ if "is_finite" not in columns:
577
+ cursor.execute("ALTER TABLE training_stats ADD COLUMN is_finite INTEGER DEFAULT 1")
578
+ logger.info("Schema migration: added training_stats.is_finite")
579
+
580
+ if version < 7:
581
+ # v7: the index sweep (see the version-history comment). The CREATEs mirror
582
+ # _init_database() — it already ran for old and new databases alike, and the
583
+ # statements are themselves idempotent, so re-executing them here just records
584
+ # the step explicitly. The DROP is the real migration work: it retires the
585
+ # (tag, timestamp, ...) latent_snapshots shape its equality-first replacement
586
+ # supersedes on a pre-v7 db; on a fresh db it is a no-op.
587
+ cursor.execute("""
588
+ CREATE INDEX IF NOT EXISTS idx_injection_stats_by_stat
589
+ ON injection_stats(tag, stat_name, signal_type, injection_stage, timestamp)
590
+ """)
591
+ cursor.execute("DROP INDEX IF EXISTS idx_latent_snapshots_filter")
592
+ cursor.execute("""
593
+ CREATE INDEX IF NOT EXISTS idx_latent_snapshots_by_key
594
+ ON latent_snapshots(tag, round_number, epoch_number, step_number,
595
+ model_name, timestamp)
596
+ """)
597
+ logger.info(
598
+ "Schema migration: v7 index sweep (injection_stats/latent_snapshots) applied"
599
+ )
600
+
601
+ if version < 8:
602
+ # v8: the supersede partial index (see the version-history comment). UNLIKE
603
+ # v7, this CREATE lives only here and deliberately NOT in _init_database():
604
+ # its WHERE superseded = 0 predicate needs the column the v1 block above
605
+ # adds, so on a pre-v1 file an init-time CREATE would precede the ALTER and
606
+ # fail. Fresh databases reach this block too (version 0 -> current), so both
607
+ # paths land the same final index set.
608
+ cursor.execute("""
609
+ CREATE INDEX IF NOT EXISTS idx_inference_results_supersede
610
+ ON inference_results(tag, npy_path) WHERE superseded = 0
611
+ """)
612
+ logger.info("Schema migration: v8 supersede partial index applied")
613
+
614
+ # PRAGMA doesn't support parameter binding; _SCHEMA_VERSION is a module-level int constant
615
+ cursor.execute(f"PRAGMA user_version = {_SCHEMA_VERSION:d}")
616
+ conn.commit()
617
+ logger.info(f"Database schema migrated from version {version} to {_SCHEMA_VERSION}")
618
+
619
+ @contextmanager
620
+ def _get_connection(self):
621
+ """Context manager for database connections with proper cleanup"""
622
+ conn = sqlite3.connect(self.db_path, timeout=self.get_connection_timeout)
623
+ try:
624
+ # Under WAL, synchronous=NORMAL only skips the per-commit WAL fsync (the WAL is
625
+ # still synced at checkpoints) — a crash can lose the most recent commits but
626
+ # never corrupts the database. That durability is ample for diagnostic telemetry
627
+ # and removes the dominant per-transaction fsync stall (#277).
628
+ conn.execute("PRAGMA synchronous=NORMAL")
629
+ # Keep GROUP BY / ORDER BY / window-function sorts in RAM instead of spilling to an
630
+ # on-disk temp b-tree. The read-heavy end-of-round / teardown passes that sort large
631
+ # scans (per-round injection-stat aggregates, the decimated resource-monitor window
632
+ # query, whole-tag plot scans) benefit; small result sets never spill, so it is free
633
+ # otherwise. Per-connection and numerics-neutral. (mmap_size is deliberately NOT set:
634
+ # its benefit is largely lost to this per-query open/close model, and SQLite mmap I/O
635
+ # faults surface as SIGBUS on a shared node — not worth the marginal gain.)
636
+ conn.execute("PRAGMA temp_store=MEMORY")
637
+ yield conn
638
+ finally:
639
+ conn.close()
640
+
641
+ def start(self):
642
+ """Start the background writer thread"""
643
+ if self.writer_thread is not None and self.writer_thread.is_alive():
644
+ return
645
+
646
+ self.stop_event.clear()
647
+ self.drain_event.clear()
648
+ # NOTE: should db be daemon or non-daemon thread?
649
+ self.writer_thread = threading.Thread(target=self._writer_loop, daemon=False)
650
+ self.writer_thread.start()
651
+ logger.info("Database writer thread started")
652
+
653
+ def _pending_row_counts(self) -> tuple[int, int]:
654
+ """Approximate (foreground items, bulk rows) still queued — for drain progress logs."""
655
+ with self._bulk_pending_lock:
656
+ bulk_rows = sum(self._bulk_pending_by_round.values())
657
+ return self.write_queue.qsize(), bulk_rows
658
+
659
+ def stop(self):
660
+ """
661
+ Drain both write lanes to disk, then stop the background writer thread.
662
+
663
+ The writer first empties the foreground and bulk queues (with progress heartbeats —
664
+ the backlog can be large), bounded by stop_drain_timeout; only then does it exit. If
665
+ the drain cap is hit, the writer is force-stopped and the exact number of dropped
666
+ rows is logged at ERROR — never silently (#277: the old stop() discarded everything
667
+ still queued, losing the most recent rows first).
668
+ """
669
+ if self.writer_thread is None:
670
+ return
671
+
672
+ fg_items, bulk_rows = self._pending_row_counts()
673
+ logger.info(
674
+ f"Stopping database writer thread (draining backlog: ~{fg_items} foreground "
675
+ f"item(s), ~{bulk_rows} bulk row(s))..."
676
+ )
677
+ self.drain_event.set() # Ask the writer to drain both lanes, then exit
678
+
679
+ deadline = time.time() + self.stop_drain_timeout
680
+ while self.writer_thread.is_alive() and time.time() < deadline:
681
+ self.writer_thread.join(timeout=5.0)
682
+ if self.writer_thread.is_alive():
683
+ fg_items, bulk_rows = self._pending_row_counts()
684
+ logger.info(
685
+ f"Draining DB backlog before shutdown: ~{fg_items} foreground item(s), "
686
+ f"~{bulk_rows} bulk row(s) remaining"
687
+ )
688
+
689
+ if self.writer_thread.is_alive():
690
+ # Drain cap hit — force the writer out and account for every dropped row
691
+ self.stop_event.set()
692
+ self.writer_thread.join(timeout=self.stop_writer_timeout)
693
+ fg_items, bulk_rows = self._pending_row_counts()
694
+ logger.error(
695
+ f"Database writer did not drain within {self.stop_drain_timeout}s; "
696
+ f"dropping ~{fg_items} foreground item(s) and ~{bulk_rows} bulk row(s) "
697
+ "still queued"
698
+ )
699
+ if self.writer_thread.is_alive():
700
+ logger.warning("Database writer thread did not stop cleanly")
701
+ else:
702
+ self.stop_event.set() # Writer already exited; keep post-stop flush()/writes short-circuiting
703
+ logger.info("Database writer thread stopped (backlog drained)")
704
+
705
+ def flush(self, timeout: float | None = None) -> bool:
706
+ """
707
+ Block until queued FOREGROUND writes are drained to the database.
708
+
709
+ Returns False on timeout or if shutdown is initiated mid-wait. A None timeout falls back
710
+ to the configured flush_timeout. Call before any read that needs to observe writes still
711
+ sitting in the queue.
712
+
713
+ The bulk injection-stat lane is deliberately NOT covered (#277): a flush must never
714
+ queue behind a round's worth of bulk rows. Readers of injection_stats should gate on
715
+ injection_backlog_rows() instead.
716
+ """
717
+ if timeout is None:
718
+ timeout = self.flush_timeout
719
+
720
+ # No-op if writer isn't running
721
+ if self.writer_thread is None or not self.writer_thread.is_alive():
722
+ logger.info("Flush called but writer thread not running")
723
+ return True
724
+
725
+ # Check if shutdown is already in progress
726
+ if self.stop_event.is_set():
727
+ logger.warning("Flush called during shutdown, skipping")
728
+ return False
729
+
730
+ # Create a completion event for this specific flush request
731
+ flush_complete = threading.Event()
732
+
733
+ # Put sentinel in queue - when writer sees this, it will flush and signal
734
+ self.write_queue.put((_FLUSH_SENTINEL, flush_complete))
735
+
736
+ # Block until writer signals completion, timeout, or shutdown
737
+ # Use a loop with short waits to also check stop_event
738
+ wait_interval = 0.1 # Check stop_event every 100ms
739
+ elapsed = 0.0
740
+
741
+ while elapsed < timeout:
742
+ if flush_complete.wait(timeout=wait_interval):
743
+ return True # Flush completed successfully
744
+
745
+ # Check if shutdown was initiated while waiting
746
+ if self.stop_event.is_set():
747
+ logger.warning("Shutdown initiated during flush, aborting wait")
748
+ return False
749
+
750
+ elapsed += wait_interval
751
+
752
+ logger.warning(f"Database flush timed out after {timeout} seconds")
753
+ return False
754
+
755
+ # Tables that support supersede marking, mapped to the optional filters each accepts
756
+ # (round_ge needs a round_number column; npy_path exists on inference_results and
757
+ # inference_cadences)
758
+ _SUPERSEDE_TABLES = {
759
+ "training_stats": frozenset({"round_ge"}),
760
+ "injection_stats": frozenset({"round_ge"}),
761
+ "latent_snapshots": frozenset({"round_ge"}),
762
+ "inference_results": frozenset({"npy_path"}),
763
+ "inference_cadences": frozenset({"npy_path"}),
764
+ }
765
+
766
+ def mark_superseded(
767
+ self,
768
+ table: str,
769
+ tag: str,
770
+ *,
771
+ round_ge: int | None = None,
772
+ npy_path: str | None = None,
773
+ timeout: float | None = None,
774
+ ) -> bool:
775
+ """
776
+ Flag existing rows for `tag` as superseded (stale data from a failed attempt) so the
777
+ default query_* filters ignore them. Rows are never deleted — pass
778
+ include_superseded=True to a query method to inspect them.
779
+
780
+ round_ge narrows the mark to round_number >= round_ge (training_stats /
781
+ injection_stats / latent_snapshots only); npy_path narrows to one cadence file
782
+ (inference_results / inference_cadences only).
783
+
784
+ The command executes on the background writer thread (a command tuple through the
785
+ write queue, like the flush sentinel) so single-writer semantics are preserved:
786
+ buffered rows are flushed first, then the UPDATE runs — queue FIFO ordering
787
+ guarantees every row queued before this call gets marked while later writes keep
788
+ superseded = 0. Blocks until the mark lands; returns False on timeout or shutdown.
789
+ When the writer thread isn't running there are no queued rows to order against, so
790
+ the UPDATE executes synchronously in the caller thread.
791
+ """
792
+ if table not in self._SUPERSEDE_TABLES:
793
+ raise ValueError(
794
+ f"mark_superseded does not support table {table!r}; "
795
+ f"expected one of {sorted(self._SUPERSEDE_TABLES)}"
796
+ )
797
+ allowed_filters = self._SUPERSEDE_TABLES[table]
798
+ if round_ge is not None and "round_ge" not in allowed_filters:
799
+ raise ValueError(f"round_ge is not supported for table {table!r}")
800
+ if npy_path is not None and "npy_path" not in allowed_filters:
801
+ raise ValueError(
802
+ "npy_path is only supported for tables 'inference_results' and 'inference_cadences'"
803
+ )
804
+ if not tag:
805
+ raise ValueError("mark_superseded requires a non-empty tag")
806
+
807
+ if timeout is None:
808
+ timeout = self.flush_timeout
809
+
810
+ # No writer thread -> nothing queued to order against; run inline
811
+ if self.writer_thread is None or not self.writer_thread.is_alive():
812
+ self._execute_mark_superseded(table, tag, round_ge, npy_path)
813
+ return True
814
+
815
+ if self.stop_event.is_set():
816
+ logger.warning("mark_superseded called during shutdown, skipping")
817
+ return False
818
+
819
+ mark_complete = threading.Event()
820
+ self.write_queue.put(
821
+ (_MARK_SUPERSEDED_SENTINEL, (table, tag, round_ge, npy_path), mark_complete)
822
+ )
823
+
824
+ # Block until the writer signals completion, timeout, or shutdown (same loop as flush)
825
+ wait_interval = 0.1
826
+ elapsed = 0.0
827
+ while elapsed < timeout:
828
+ if mark_complete.wait(timeout=wait_interval):
829
+ return True
830
+ if self.stop_event.is_set():
831
+ logger.warning("Shutdown initiated during mark_superseded, aborting wait")
832
+ return False
833
+ elapsed += wait_interval
834
+
835
+ logger.warning(f"mark_superseded timed out after {timeout} seconds")
836
+ return False
837
+
838
+ def _execute_mark_superseded(
839
+ self, table: str, tag: str, round_ge: int | None, npy_path: str | None
840
+ ) -> None:
841
+ """Apply the supersede UPDATE (runs on the writer thread, or inline via
842
+ mark_superseded when no writer thread is running). `table` was already validated
843
+ against _SUPERSEDE_TABLES, so the f-string interpolation below is safe."""
844
+ query = f"UPDATE {table} SET superseded = 1 WHERE superseded = 0 AND tag = ?"
845
+ params: list = [tag]
846
+
847
+ if round_ge is not None:
848
+ query += " AND round_number >= ?"
849
+ params.append(round_ge)
850
+
851
+ if npy_path is not None:
852
+ query += " AND npy_path = ?"
853
+ params.append(npy_path)
854
+
855
+ try:
856
+ with self._get_connection() as conn:
857
+ cursor = conn.cursor()
858
+ cursor.execute(query, params)
859
+ conn.commit()
860
+ logger.info(
861
+ f"Marked {cursor.rowcount} row(s) superseded in {table} "
862
+ f"(tag={tag}, round_ge={round_ge}, npy_path={npy_path})"
863
+ )
864
+ except Exception as e:
865
+ # Mirror _flush_buffer's log-and-continue semantics: a failed mark must not
866
+ # kill the writer thread (stale rows would then merely reappear in plots)
867
+ logger.error(f"Error marking rows superseded in {table}: {e}")
868
+
869
+ def _consume_bulk_item(self) -> bool:
870
+ """
871
+ Pull one bulk chunk (if any), commit it together with any buffered foreground rows in
872
+ a single transaction, and settle the pending-row accounting. Runs on the writer
873
+ thread. Returns True if a chunk was consumed.
874
+ """
875
+ try:
876
+ table, rows, round_counts = self.bulk_queue.get_nowait()
877
+ except Empty:
878
+ return False
879
+
880
+ self.buffer.extend((table, values) for values in rows)
881
+ self._flush_buffer()
882
+ self.buffer.clear()
883
+
884
+ # Settle the per-round pending accounting only after the commit attempt, so
885
+ # injection_backlog_rows() == 0 really means "every enqueued row reached the db
886
+ # engine" (_flush_buffer logs and drops a failed batch — those rows are gone either
887
+ # way, so they must not linger in the backlog count)
888
+ with self._bulk_pending_lock:
889
+ for round_key, count in round_counts.items():
890
+ remaining = self._bulk_pending_by_round.get(round_key, 0) - count
891
+ if remaining > 0:
892
+ self._bulk_pending_by_round[round_key] = remaining
893
+ else:
894
+ self._bulk_pending_by_round.pop(round_key, None)
895
+ return True
896
+
897
+ def _writer_loop(self):
898
+ """
899
+ Background loop that consumes data from both write lanes and writes to the database.
900
+
901
+ The foreground write_queue (training stats, latent snapshots, sentinels, ...) has
902
+ strict priority; the bounded bulk lane (injection-stat chunks) is serviced whenever
903
+ the foreground lane is idle. On graceful shutdown (drain_event) both lanes are fully
904
+ drained to disk before the thread exits (#277).
905
+ """
906
+ self.buffer = []
907
+ last_write_time = time.time()
908
+
909
+ # Keep looping until told to stop (hard) or drain (graceful)
910
+ while not self.stop_event.is_set() and not self.drain_event.is_set():
911
+ try:
912
+ # Calculate how much time remains until the next scheduled write
913
+ # Don't wait longer than 1s so we check the stop flag regularly
914
+ # Don't wait more than 0.1s to avoid wasting CPU resources
915
+ # Retrive items from the queue one-by-one & append to local buffer
916
+ # When bulk chunks are waiting, poll the foreground lane briefly instead so
917
+ # the bulk lane drains at full speed while foreground items still win ties
918
+ if self.bulk_queue.empty():
919
+ timeout = max(
920
+ 0.1, min(1.0, self.write_interval - (time.time() - last_write_time))
921
+ )
922
+ else:
923
+ timeout = 0.01
924
+ metric = self.write_queue.get(timeout=timeout)
925
+
926
+ # Check for flush sentinel - signals immediate flush request
927
+ if metric[0] is _FLUSH_SENTINEL:
928
+ flush_complete_event = metric[1]
929
+ # Flush current buffer immediately
930
+ if self.buffer:
931
+ # Write all buffered data to db
932
+ self._flush_buffer()
933
+ # Clear the buffer (empty the list)
934
+ self.buffer.clear()
935
+ # Reset the timer
936
+ last_write_time = time.time()
937
+ # Signal that flush is complete
938
+ flush_complete_event.set()
939
+ continue
940
+
941
+ # Check for mark-superseded command. Flush buffered rows first: queue FIFO
942
+ # guarantees every row enqueued before the command is already in the buffer,
943
+ # so the UPDATE covers them, while rows enqueued after keep superseded = 0
944
+ if metric[0] is _MARK_SUPERSEDED_SENTINEL:
945
+ _, payload, mark_complete_event = metric
946
+ try:
947
+ if self.buffer:
948
+ self._flush_buffer()
949
+ self.buffer.clear()
950
+ last_write_time = time.time()
951
+ self._execute_mark_superseded(*payload)
952
+ finally:
953
+ # Always unblock the caller, even if the flush/UPDATE errored
954
+ mark_complete_event.set()
955
+ continue
956
+
957
+ self.buffer.append(metric)
958
+
959
+ # Write when buffer is full or interval elapsed
960
+ current_time = time.time()
961
+ if (
962
+ len(self.buffer) >= self.write_buffer_max_size
963
+ or (current_time - last_write_time) >= self.write_interval
964
+ ):
965
+ # Write all buffered data to db
966
+ self._flush_buffer()
967
+ # Clear the buffer (empty the list)
968
+ self.buffer.clear()
969
+ # Reset the timer
970
+ last_write_time = current_time
971
+
972
+ except Empty:
973
+ # Foreground lane idle — service one bulk chunk if available (commits the
974
+ # current buffer alongside it), else fall back to the interval flush check
975
+ if self._consume_bulk_item():
976
+ last_write_time = time.time()
977
+ continue
978
+ # If get() timesout (queue was empty) but interval elapsed, write buffered data anyway
979
+ current_time = time.time()
980
+ if self.buffer and (current_time - last_write_time) >= self.write_interval:
981
+ # Write all buffered data to db
982
+ self._flush_buffer()
983
+ # Clear the buffer (empty the list)
984
+ self.buffer.clear()
985
+ # Reset the timer
986
+ last_write_time = current_time
987
+ continue
988
+
989
+ except Exception as e:
990
+ logger.error(f"Error in db writer loop: {e}")
991
+ # Sleep (interruptible for faster shutdown)
992
+ self.stop_event.wait(self.write_retry_delay)
993
+
994
+ # Graceful shutdown: drain BOTH lanes to disk before exiting (#277 — the old
995
+ # implementation dropped everything still queued, silently losing the most recent
996
+ # rows). A hard stop (stop_event without drain_event, or the drain cap escalation
997
+ # in stop()) skips this and only flushes the in-memory buffer below.
998
+ if self.drain_event.is_set() and not self.stop_event.is_set():
999
+ self._drain_queues()
1000
+
1001
+ # Final flush on shutdown
1002
+ if self.buffer:
1003
+ flushed_count = len(self.buffer)
1004
+ self._flush_buffer()
1005
+ self.buffer.clear()
1006
+ logger.info(f"Flushed {flushed_count} remaining data on shutdown")
1007
+
1008
+ def _drain_queues(self) -> None:
1009
+ """Empty both write lanes to disk (graceful-shutdown path; see stop())."""
1010
+ drained_foreground = 0
1011
+ drained_bulk_chunks = 0
1012
+ while not self.stop_event.is_set():
1013
+ progressed = False
1014
+ try:
1015
+ metric = self.write_queue.get_nowait()
1016
+ progressed = True
1017
+ if metric[0] is _FLUSH_SENTINEL:
1018
+ if self.buffer:
1019
+ self._flush_buffer()
1020
+ self.buffer.clear()
1021
+ metric[1].set()
1022
+ elif metric[0] is _MARK_SUPERSEDED_SENTINEL:
1023
+ _, payload, mark_complete_event = metric
1024
+ try:
1025
+ if self.buffer:
1026
+ self._flush_buffer()
1027
+ self.buffer.clear()
1028
+ self._execute_mark_superseded(*payload)
1029
+ finally:
1030
+ mark_complete_event.set()
1031
+ else:
1032
+ self.buffer.append(metric)
1033
+ drained_foreground += 1
1034
+ if len(self.buffer) >= self.write_buffer_max_size:
1035
+ self._flush_buffer()
1036
+ self.buffer.clear()
1037
+ except Empty:
1038
+ pass
1039
+
1040
+ if self._consume_bulk_item():
1041
+ progressed = True
1042
+ drained_bulk_chunks += 1
1043
+
1044
+ if not progressed:
1045
+ break
1046
+
1047
+ if drained_foreground or drained_bulk_chunks:
1048
+ logger.info(
1049
+ f"Drained {drained_foreground} foreground row(s) and {drained_bulk_chunks} "
1050
+ "bulk chunk(s) to the database at shutdown"
1051
+ )
1052
+
1053
+ # In commit 08fc37d, we switched from using sequential execute() to executemany()
1054
+ # The sequential approaches parses SQL statements N times & performs N round trips to the db
1055
+ # engine, in exchange for lower peak memory & more granular errors (can identify exactly which
1056
+ # row failed). In contrast, executemany() only parses SQL statements once & performs only 1
1057
+ # round trip to the db engine, but requires all rows to be in memory & fails the entire batch
1058
+ # if any row fails
1059
+ # In general, sequential execute() is preferred when db buffer sizes are small, if write
1060
+ # frequency is low, if each record in the buffer is guaranteed to go to different tables, or
1061
+ # if we require per-row error handling
1062
+ # In our case, we have a larger buffer (100+ records), most records in a batch will go to the
1063
+ # same table, and we're okay with all-or-nothing batch semantics. So the increased write speeds
1064
+ # in exchange for increased memory pressure is worth it
1065
+ # As well, since our db writes happen in a background thread & don't block the main process,
1066
+ # either approach should have minimal practical impact
1067
+ def _executemany_resilient(self, cursor, sql: str, records: list, table: str) -> None:
1068
+ """
1069
+ executemany with a per-row fallback (#289): if the batch insert fails, retry the
1070
+ rows individually so one unbindable row can no longer discard everything else in
1071
+ the flush. Offending rows are skipped with an exact count — the old behavior lost
1072
+ the entire buffered batch (every table's rows) to a single NOT NULL violation,
1073
+ with nothing but a one-line error to show for it.
1074
+
1075
+ The batch attempt runs inside a SAVEPOINT: executemany inserts each row into the
1076
+ open transaction as it steps, so a mid-batch failure leaves the rows BEFORE the
1077
+ offender committed-in-progress — retrying per-row without rolling those back would
1078
+ duplicate them. ROLLBACK TO restores a clean slate before the per-row pass.
1079
+ """
1080
+ try:
1081
+ cursor.execute("SAVEPOINT resilient_batch")
1082
+ cursor.executemany(sql, records)
1083
+ cursor.execute("RELEASE SAVEPOINT resilient_batch")
1084
+ except sqlite3.Error as batch_error:
1085
+ cursor.execute("ROLLBACK TO SAVEPOINT resilient_batch")
1086
+ cursor.execute("RELEASE SAVEPOINT resilient_batch")
1087
+ logger.error(
1088
+ f"Bulk insert into {table} failed ({batch_error}); retrying "
1089
+ f"{len(records)} row(s) individually"
1090
+ )
1091
+ skipped = 0
1092
+ for record in records:
1093
+ try:
1094
+ cursor.execute(sql, record)
1095
+ except sqlite3.Error:
1096
+ skipped += 1
1097
+ if skipped:
1098
+ logger.error(f"Dropped {skipped}/{len(records)} unwritable row(s) for {table}")
1099
+
1100
+ def _flush_buffer(self, buffer: list | None = None):
1101
+ """
1102
+ Write buffered data to database in a single transaction using executemany().
1103
+
1104
+ Groups records by table and uses executemany() for bulk inserts, which is more efficient
1105
+ than individual execute() calls because the SQL is parsed once and reused for all rows.
1106
+ Defaults to the writer thread's buffer; an explicit `buffer` supports the inline
1107
+ no-writer path of write_injection_stats_bulk.
1108
+ """
1109
+ if buffer is None:
1110
+ buffer = self.buffer
1111
+ if not buffer:
1112
+ return
1113
+
1114
+ try:
1115
+ with self._get_connection() as conn:
1116
+ cursor = conn.cursor()
1117
+
1118
+ # Group records by table for bulk inserts
1119
+ system_resources_records: list[tuple] = []
1120
+ injection_stats_records: list[tuple] = []
1121
+ training_stats_records: list[tuple] = []
1122
+ latent_snapshots_records: list[tuple] = []
1123
+ inference_results_records: list[tuple] = []
1124
+ inference_cadences_records: list[tuple] = []
1125
+ pipeline_stages_records: list[tuple] = []
1126
+
1127
+ for table, values in buffer:
1128
+ if table == "system_resources":
1129
+ system_resources_records.append(values)
1130
+ elif table == "injection_stats":
1131
+ injection_stats_records.append(values)
1132
+ elif table == "training_stats":
1133
+ training_stats_records.append(values)
1134
+ elif table == "latent_snapshots":
1135
+ latent_snapshots_records.append(values)
1136
+ elif table == "inference_results":
1137
+ inference_results_records.append(values)
1138
+ elif table == "inference_cadences":
1139
+ inference_cadences_records.append(values)
1140
+ elif table == "pipeline_stages":
1141
+ pipeline_stages_records.append(values)
1142
+
1143
+ # Bulk insert each table type
1144
+ if system_resources_records:
1145
+ self._executemany_resilient(
1146
+ cursor,
1147
+ """
1148
+ INSERT INTO system_resources
1149
+ (timestamp, resource_type, resource_name, value, unit, tag, metadata)
1150
+ VALUES (?, ?, ?, ?, ?, ?, ?)
1151
+ """,
1152
+ system_resources_records,
1153
+ "system_resources",
1154
+ )
1155
+
1156
+ if injection_stats_records:
1157
+ self._executemany_resilient(
1158
+ cursor,
1159
+ """
1160
+ INSERT INTO injection_stats
1161
+ (timestamp, stat_name, value, round_number, chunk_number, sample_index,
1162
+ background_index, signal_class, signal_type, injection_stage, is_finite,
1163
+ slope_clamped, tag, metadata)
1164
+ VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
1165
+ """,
1166
+ injection_stats_records,
1167
+ "injection_stats",
1168
+ )
1169
+
1170
+ if training_stats_records:
1171
+ self._executemany_resilient(
1172
+ cursor,
1173
+ """
1174
+ INSERT INTO training_stats
1175
+ (timestamp, model_name, stat_name, value, round_number, epoch_number,
1176
+ tag, metadata, is_finite)
1177
+ VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)
1178
+ """,
1179
+ training_stats_records,
1180
+ "training_stats",
1181
+ )
1182
+
1183
+ if latent_snapshots_records:
1184
+ self._executemany_resilient(
1185
+ cursor,
1186
+ """
1187
+ INSERT INTO latent_snapshots
1188
+ (timestamp, model_name, round_number, epoch_number, step_number,
1189
+ cadence_index, signal_type, latent_vector, snr_base, snr_range, tag,
1190
+ metadata)
1191
+ VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
1192
+ """,
1193
+ latent_snapshots_records,
1194
+ "latent_snapshots",
1195
+ )
1196
+
1197
+ if inference_results_records:
1198
+ self._executemany_resilient(
1199
+ cursor,
1200
+ """
1201
+ INSERT INTO inference_results
1202
+ (timestamp, npy_path, snippet_index, prediction, confidence, latent_vector,
1203
+ target, session, cadence_id, band, frequency_mhz, timestamp_observed,
1204
+ h5_path, tag, metadata, screening_proba, mc_mean, mc_std)
1205
+ VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
1206
+ """,
1207
+ inference_results_records,
1208
+ "inference_results",
1209
+ )
1210
+
1211
+ if inference_cadences_records:
1212
+ self._executemany_resilient(
1213
+ cursor,
1214
+ """
1215
+ INSERT INTO inference_cadences
1216
+ (timestamp, tag, csv_path, cadence_key, npy_path, status, n_stamps,
1217
+ n_candidates, confidence_summary, duration_s, config_fingerprint)
1218
+ VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
1219
+ """,
1220
+ inference_cadences_records,
1221
+ "inference_cadences",
1222
+ )
1223
+
1224
+ if pipeline_stages_records:
1225
+ self._executemany_resilient(
1226
+ cursor,
1227
+ """
1228
+ INSERT INTO pipeline_stages
1229
+ (stage, start_time, end_time, duration_s, tag, metadata)
1230
+ VALUES (?, ?, ?, ?, ?, ?)
1231
+ """,
1232
+ pipeline_stages_records,
1233
+ "pipeline_stages",
1234
+ )
1235
+
1236
+ conn.commit()
1237
+
1238
+ except Exception as e:
1239
+ logger.error(f"Error flushing data buffer: {e}")
1240
+
1241
+ # TODO: write checks to sanitize values before writing to db. raise error if problematic value passed
1242
+ def write_system_resource(
1243
+ self,
1244
+ resource_type: str,
1245
+ resource_name: str,
1246
+ value: float,
1247
+ unit: str | None = None,
1248
+ tag: str | None = None,
1249
+ timestamp: float | None = None,
1250
+ ):
1251
+ """
1252
+ Queue a non-blocking write to system_resources.
1253
+
1254
+ resource_type takes broad categories like 'cpu', 'ram', 'gpu'; resource_name takes
1255
+ finer sub-labels like 'system_total' or 'process_tree'. unit is a free-form label
1256
+ ('percent', 'MB', etc.). Timestamp defaults to current wall time if omitted.
1257
+ """
1258
+ metadata_json = get_system_metadata()
1259
+
1260
+ self.write_queue.put(
1261
+ (
1262
+ "system_resources",
1263
+ (
1264
+ timestamp or time.time(),
1265
+ resource_type,
1266
+ resource_name,
1267
+ value,
1268
+ unit,
1269
+ tag,
1270
+ metadata_json,
1271
+ ),
1272
+ )
1273
+ )
1274
+
1275
+ # TODO: write checks to sanitize values before writing to db. raise error if problematic value passed
1276
+ def write_injection_stat(
1277
+ self,
1278
+ stat_name: str,
1279
+ value: float,
1280
+ round_number: int | None = None,
1281
+ chunk_number: int | None = None,
1282
+ sample_index: int | None = None,
1283
+ background_index: int | None = None,
1284
+ signal_class: str | None = None,
1285
+ signal_type: str | None = None,
1286
+ injection_stage: str | None = None,
1287
+ slope_clamped: bool | None = None,
1288
+ tag: str | None = None,
1289
+ timestamp: float | None = None,
1290
+ ):
1291
+ """
1292
+ Queue a non-blocking write to injection_stats.
1293
+
1294
+ stat_name labels the metric (global_mean, eti_snr, rfi_drift_rate, num_samples, etc.).
1295
+ signal_class is one of {main, false, true}; signal_type is one of
1296
+ {false_no_signal, false_with_rfi, true_only_eti, true_eti_rfi}; injection_stage is one of
1297
+ {A (pre-inj pre-norm), B (post-inj pre-norm), C (post-inj post-norm)}. Non-finite values
1298
+ are coerced to 0.0 with is_finite=0 so queries can drop them via the default only_finite
1299
+ filter — this is required because the schema's NOT NULL constraint forbids storing NaN/Inf
1300
+ directly. Timestamp defaults to current wall time.
1301
+ """
1302
+ metadata_json = get_system_metadata()
1303
+
1304
+ # Since higher-order moments from _compute_intensity_stats() could lead to NaN/Inf values
1305
+ # We explicitly check if a value is finite to set the flag accordingly
1306
+ # We will opt to store these values as 0.0 due to the schema's NOT NULL constraint
1307
+ # However, queries for injection_stats should by default use the is_finite flag, unless
1308
+ # sanitization is explicitly not needed
1309
+ is_finite = 1 if np.isfinite(value) else 0
1310
+ sanitized_value = float(value) if is_finite else 0.0
1311
+
1312
+ if not is_finite:
1313
+ logger.warning(
1314
+ f"write_injection_stat: {stat_name} is {value}, storing as 0.0 with is_finite=0"
1315
+ )
1316
+
1317
+ # Convert slope_clamped bool to int (0 or 1), defaulting to 0 if None
1318
+ slope_clamped_int = 1 if slope_clamped else 0
1319
+
1320
+ self.write_queue.put(
1321
+ (
1322
+ "injection_stats",
1323
+ (
1324
+ timestamp or time.time(),
1325
+ stat_name,
1326
+ sanitized_value,
1327
+ round_number,
1328
+ chunk_number,
1329
+ sample_index,
1330
+ background_index,
1331
+ signal_class,
1332
+ signal_type,
1333
+ injection_stage,
1334
+ is_finite,
1335
+ slope_clamped_int,
1336
+ tag,
1337
+ metadata_json,
1338
+ ),
1339
+ )
1340
+ )
1341
+
1342
+ def write_injection_stats_bulk(self, stats: list[dict], tag: str | None = None) -> None:
1343
+ """
1344
+ Queue a batch of injection-stat rows on the bounded bulk lane (#277).
1345
+
1346
+ Each dict takes the same keys as write_injection_stat's parameters (stat_name and
1347
+ value required; the rest optional). Semantics per row are identical — NaN/Inf
1348
+ coercion with is_finite=0, slope_clamped bool -> int — but the batch shares ONE
1349
+ system-metadata lookup and is enqueued in bulk_chunk_rows-sized chunks, so a
1350
+ ~300K-row segment costs a handful of queue operations instead of ~300K (the per-row
1351
+ Python overhead that capped the writer at ~590 rows/s). put() blocks when the bulk
1352
+ lane is full — deliberate backpressure on the (background) enqueuer; never call this
1353
+ from the training-critical path. With no writer thread running the rows are written
1354
+ inline, mirroring mark_superseded's no-writer path.
1355
+ """
1356
+ metadata_json = get_system_metadata()
1357
+ now = time.time()
1358
+
1359
+ rows: list[tuple] = []
1360
+ round_counts: dict[int, int] = {}
1361
+ non_finite = 0
1362
+ for stat in stats:
1363
+ value = stat["value"]
1364
+ is_finite = 1 if np.isfinite(value) else 0
1365
+ if not is_finite:
1366
+ non_finite += 1
1367
+ round_number = stat.get("round_number")
1368
+ rows.append(
1369
+ (
1370
+ stat.get("timestamp") or now,
1371
+ stat["stat_name"],
1372
+ float(value) if is_finite else 0.0,
1373
+ round_number,
1374
+ stat.get("chunk_number"),
1375
+ stat.get("sample_index"),
1376
+ stat.get("background_index"),
1377
+ stat.get("signal_class"),
1378
+ stat.get("signal_type"),
1379
+ stat.get("injection_stage"),
1380
+ is_finite,
1381
+ 1 if stat.get("slope_clamped") else 0,
1382
+ tag,
1383
+ metadata_json,
1384
+ )
1385
+ )
1386
+ # None rounds key as -1 so injection_backlog_rows() counts them conservatively
1387
+ round_key = round_number if round_number is not None else -1
1388
+ round_counts[round_key] = round_counts.get(round_key, 0) + 1
1389
+
1390
+ if non_finite:
1391
+ logger.warning(
1392
+ f"write_injection_stats_bulk: coerced {non_finite} non-finite value(s) to 0.0 "
1393
+ "with is_finite=0"
1394
+ )
1395
+
1396
+ if not rows:
1397
+ return
1398
+
1399
+ # No writer thread -> no consumer for the bounded queue; write inline instead
1400
+ if self.writer_thread is None or not self.writer_thread.is_alive():
1401
+ self._flush_buffer([("injection_stats", values) for values in rows])
1402
+ return
1403
+
1404
+ for start in range(0, len(rows), self.bulk_chunk_rows):
1405
+ chunk = rows[start : start + self.bulk_chunk_rows]
1406
+ chunk_counts: dict[int, int] = {}
1407
+ for values in chunk:
1408
+ round_key = values[3] if values[3] is not None else -1
1409
+ chunk_counts[round_key] = chunk_counts.get(round_key, 0) + 1
1410
+ with self._bulk_pending_lock:
1411
+ for round_key, count in chunk_counts.items():
1412
+ self._bulk_pending_by_round[round_key] = (
1413
+ self._bulk_pending_by_round.get(round_key, 0) + count
1414
+ )
1415
+ # Blocks when the lane is full (bounded queue) — the backpressure is the point.
1416
+ # But never blocks FOREVER: a bounded put with no consumer would deadlock the
1417
+ # enqueuer if the writer died after the liveness check above, so re-check
1418
+ # liveness on every timeout and fall back to an inline write if it's gone
1419
+ while True:
1420
+ if self.writer_thread is None or not self.writer_thread.is_alive():
1421
+ logger.warning(
1422
+ "DB writer thread died while the bulk lane was full - writing "
1423
+ f"{len(chunk)} injection row(s) inline"
1424
+ )
1425
+ self._flush_buffer([("injection_stats", values) for values in chunk])
1426
+ with self._bulk_pending_lock:
1427
+ for round_key, count in chunk_counts.items():
1428
+ remaining = self._bulk_pending_by_round.get(round_key, 0) - count
1429
+ if remaining > 0:
1430
+ self._bulk_pending_by_round[round_key] = remaining
1431
+ else:
1432
+ self._bulk_pending_by_round.pop(round_key, None)
1433
+ break
1434
+ try:
1435
+ self.bulk_queue.put(("injection_stats", chunk, chunk_counts), timeout=5.0)
1436
+ break
1437
+ except Full:
1438
+ continue
1439
+
1440
+ def injection_backlog_rows(self, max_round: int | None = None) -> int:
1441
+ """
1442
+ Bulk-lane injection rows enqueued but not yet committed for rounds <= max_round
1443
+ (None = all rounds). Rows with round_number None are counted against every
1444
+ max_round (conservative). Plot code uses this to verify a round's injection rows
1445
+ are all committed before rendering its figures (#277).
1446
+ """
1447
+ with self._bulk_pending_lock:
1448
+ if max_round is None:
1449
+ return sum(self._bulk_pending_by_round.values())
1450
+ return sum(
1451
+ count
1452
+ for round_key, count in self._bulk_pending_by_round.items()
1453
+ if round_key == -1 or round_key <= max_round
1454
+ )
1455
+
1456
+ # TODO: write checks to sanitize values before writing to db. raise error if problematic value passed
1457
+ def write_training_stat(
1458
+ self,
1459
+ model_name: str,
1460
+ stat_name: str,
1461
+ value: float,
1462
+ round_number: int | None = None,
1463
+ epoch_number: int | None = None,
1464
+ tag: str | None = None,
1465
+ timestamp: float | None = None,
1466
+ ):
1467
+ """
1468
+ Queue a non-blocking write to training_stats.
1469
+
1470
+ model_name labels the model ('beta_vae', 'rf'); stat_name labels the metric ('total_loss',
1471
+ 'reconstruction_loss', 'learning_rate', etc.). Timestamp defaults to current wall time.
1472
+
1473
+ Non-finite (or None) values are coerced to 0.0 with is_finite=0 (#289) — the
1474
+ injection_stats semantics — so queries can drop them via the default only_finite
1475
+ filter. Writing them raw is impossible anyway: sqlite binds NaN as NULL, and the NOT
1476
+ NULL constraint used to make the whole flush batch vanish.
1477
+ """
1478
+ metadata_json = get_system_metadata()
1479
+
1480
+ is_finite = 1 if value is not None and np.isfinite(value) else 0
1481
+ if not is_finite:
1482
+ logger.warning(
1483
+ f"write_training_stat: {stat_name} is {value}, storing as 0.0 with is_finite=0"
1484
+ )
1485
+
1486
+ self.write_queue.put(
1487
+ (
1488
+ "training_stats",
1489
+ (
1490
+ timestamp or time.time(),
1491
+ model_name,
1492
+ stat_name,
1493
+ float(value) if is_finite else 0.0,
1494
+ round_number,
1495
+ epoch_number,
1496
+ tag,
1497
+ metadata_json,
1498
+ is_finite,
1499
+ ),
1500
+ )
1501
+ )
1502
+
1503
+ # TODO: write checks to sanitize values before writing to db. raise error if problematic value passed
1504
+ def write_latent_snapshot(
1505
+ self,
1506
+ model_name: str,
1507
+ round_number: int,
1508
+ epoch_number: int,
1509
+ step_number: int,
1510
+ cadence_index: int,
1511
+ signal_type: str,
1512
+ latent_vector: list[list[float]],
1513
+ snr_base: int | None = None,
1514
+ snr_range: int | None = None,
1515
+ tag: str | None = None,
1516
+ timestamp: float | None = None,
1517
+ ):
1518
+ """
1519
+ Queue a non-blocking write to latent_snapshots.
1520
+
1521
+ cadence_index is the position within the viz batch (0 to num_cadences-1). signal_type is
1522
+ one of {false_no_signal, false_with_rfi, true_only_eti, true_eti_rfi}. latent_vector is a
1523
+ nested list of shape (6, latent_dim), serialized to JSON before storage. Timestamp defaults
1524
+ to current wall time.
1525
+ """
1526
+ metadata_json = get_system_metadata()
1527
+ latent_vector_json = json.dumps(latent_vector)
1528
+
1529
+ self.write_queue.put(
1530
+ (
1531
+ "latent_snapshots",
1532
+ (
1533
+ timestamp or time.time(),
1534
+ model_name,
1535
+ round_number,
1536
+ epoch_number,
1537
+ step_number,
1538
+ cadence_index,
1539
+ signal_type,
1540
+ latent_vector_json,
1541
+ snr_base,
1542
+ snr_range,
1543
+ tag,
1544
+ metadata_json,
1545
+ ),
1546
+ )
1547
+ )
1548
+
1549
+ def write_latent_snapshots_bulk(
1550
+ self,
1551
+ model_name: str,
1552
+ round_number: int,
1553
+ epoch_number: int,
1554
+ step_number: int,
1555
+ snapshots: list[tuple],
1556
+ snr_base: int | None = None,
1557
+ snr_range: int | None = None,
1558
+ tag: str | None = None,
1559
+ timestamp: float | None = None,
1560
+ ) -> None:
1561
+ """
1562
+ Queue one snapshot capture's worth of latent_snapshots rows in a single call.
1563
+
1564
+ `snapshots` is a list of (cadence_index, signal_type, latent_vector) tuples sharing
1565
+ the given capture-level fields; latent_vector is array-like of shape
1566
+ (num_observations, latent_dim). Row semantics are identical to per-row
1567
+ write_latent_snapshot calls, but the system-metadata lookup happens once per batch
1568
+ instead of once per cadence — the per-row Python cost ran inside the epoch loop via
1569
+ _capture_latent_snapshot (measured ~0.12 s per capture at 3,840 rows).
1570
+ """
1571
+ metadata_json = get_system_metadata()
1572
+ ts = timestamp or time.time()
1573
+ for cadence_index, signal_type, latent_vector in snapshots:
1574
+ self.write_queue.put(
1575
+ (
1576
+ "latent_snapshots",
1577
+ (
1578
+ ts,
1579
+ model_name,
1580
+ round_number,
1581
+ epoch_number,
1582
+ step_number,
1583
+ cadence_index,
1584
+ signal_type,
1585
+ json.dumps(np.asarray(latent_vector).tolist()),
1586
+ snr_base,
1587
+ snr_range,
1588
+ tag,
1589
+ metadata_json,
1590
+ ),
1591
+ )
1592
+ )
1593
+
1594
+ # TODO: write checks to sanitize values before writing to db. raise error if problematic value passed
1595
+ def write_inference_result(
1596
+ self,
1597
+ npy_path: str,
1598
+ snippet_index: int,
1599
+ prediction: int,
1600
+ confidence: float,
1601
+ latent_vector: np.ndarray | None = None,
1602
+ target: str | None = None,
1603
+ session: str | None = None,
1604
+ cadence_id: int | None = None,
1605
+ band: str | None = None,
1606
+ frequency_mhz: float | None = None,
1607
+ timestamp_observed: float | None = None,
1608
+ h5_path: str | None = None,
1609
+ tag: str | None = None,
1610
+ timestamp: float | None = None,
1611
+ screening_proba: float | None = None,
1612
+ mc_mean: float | None = None,
1613
+ mc_std: float | None = None,
1614
+ ):
1615
+ """
1616
+ Queue a non-blocking write to inference_results.
1617
+
1618
+ prediction is 0 (RFI) or 1 (candidate); confidence is in [0.0, 1.0]. latent_vector, when
1619
+ provided, is a (6 * latent_dim,) array that gets serialized to JSON for later analysis.
1620
+ target, session, cadence_id, band, frequency_mhz, timestamp_observed, and h5_path carry
1621
+ the observational provenance (e.g. target='DDO210', session='AGBT18A_999_103', band='L').
1622
+ Timestamp defaults to current wall time (distinct from timestamp_observed, which is the
1623
+ original observation time). screening_proba / mc_mean / mc_std carry the #282 two-pass
1624
+ scores: the deterministic pass-1 probability and the seeded MC mean/spread for
1625
+ pass-2 survivors.
1626
+ """
1627
+ metadata_json = get_system_metadata()
1628
+
1629
+ # Convert latent vector to JSON string if provided
1630
+ latent_json = None
1631
+ if latent_vector is not None:
1632
+ latent_json = json.dumps(latent_vector.tolist())
1633
+
1634
+ self.write_queue.put(
1635
+ (
1636
+ "inference_results",
1637
+ (
1638
+ timestamp or time.time(),
1639
+ npy_path,
1640
+ snippet_index,
1641
+ prediction,
1642
+ confidence,
1643
+ latent_json,
1644
+ target,
1645
+ session,
1646
+ cadence_id,
1647
+ band,
1648
+ frequency_mhz,
1649
+ timestamp_observed,
1650
+ h5_path,
1651
+ tag,
1652
+ metadata_json,
1653
+ screening_proba,
1654
+ mc_mean,
1655
+ mc_std,
1656
+ ),
1657
+ )
1658
+ )
1659
+
1660
+ # TODO: write checks to sanitize values before writing to db. raise error if problematic value passed
1661
+ def write_inference_cadence(
1662
+ self,
1663
+ npy_path: str,
1664
+ status: str,
1665
+ tag: str | None = None,
1666
+ csv_path: str | None = None,
1667
+ cadence_key: tuple | list | None = None,
1668
+ n_stamps: int | None = None,
1669
+ n_candidates: int | None = None,
1670
+ confidence_summary: dict | None = None,
1671
+ duration_s: float | None = None,
1672
+ config_fingerprint: str | None = None,
1673
+ timestamp: float | None = None,
1674
+ ):
1675
+ """
1676
+ Queue a non-blocking write to the inference_cadences run manifest.
1677
+
1678
+ One row per (cadence, stage transition): status='preprocessed' when the cadence's
1679
+ stamp .npy lands (n_stamps set, aggregates None), status='inferred' when inference
1680
+ completes (aggregates set; the caller supersedes older rows for the same
1681
+ (tag, npy_path) first), status='failed' when the cadence's inference stage died
1682
+ (so the failure is inspectable and the retry pass re-attempts it). cadence_key
1683
+ (the CSV group-by key) and confidence_summary (quantile stats from
1684
+ inference.summarize_confidences) are JSON-serialized for storage. duration_s is
1685
+ the stage's wall-clock duration. Timestamp defaults to current wall time.
1686
+ """
1687
+ cadence_key_json = None
1688
+ if cadence_key is not None:
1689
+ cadence_key_json = json.dumps([str(part) for part in cadence_key])
1690
+ confidence_summary_json = None
1691
+ if confidence_summary is not None:
1692
+ confidence_summary_json = json.dumps(confidence_summary)
1693
+
1694
+ self.write_queue.put(
1695
+ (
1696
+ "inference_cadences",
1697
+ (
1698
+ timestamp or time.time(),
1699
+ tag,
1700
+ csv_path,
1701
+ cadence_key_json,
1702
+ npy_path,
1703
+ status,
1704
+ n_stamps,
1705
+ n_candidates,
1706
+ confidence_summary_json,
1707
+ duration_s,
1708
+ config_fingerprint,
1709
+ ),
1710
+ )
1711
+ )
1712
+
1713
+ # TODO: write checks to sanitize values before writing to db. raise error if problematic value passed
1714
+ def write_pipeline_stage(
1715
+ self,
1716
+ stage: str,
1717
+ start_time: float,
1718
+ end_time: float,
1719
+ tag: str | None = None,
1720
+ metadata: str | None = None,
1721
+ ):
1722
+ """
1723
+ Queue a non-blocking write to pipeline_stages.
1724
+
1725
+ stage is a hierarchical dot-name ("train.round_02.data_generation"); start_time /
1726
+ end_time are unix timestamps bounding the span (duration_s is derived here so the
1727
+ table stays internally consistent). metadata, when given, is an already-serialized
1728
+ JSON string (aetherscan.benchmark owns the serialization). Prefer the stage_timer /
1729
+ record_stage helpers in aetherscan.benchmark over calling this directly.
1730
+ """
1731
+ self.write_queue.put(
1732
+ (
1733
+ "pipeline_stages",
1734
+ (
1735
+ stage,
1736
+ start_time,
1737
+ end_time,
1738
+ end_time - start_time,
1739
+ tag,
1740
+ metadata,
1741
+ ),
1742
+ )
1743
+ )
1744
+
1745
+ # Column whitelists per table (for SQL injection prevention when using column projection)
1746
+ _SYSTEM_RESOURCES_COLUMNS = {
1747
+ "id",
1748
+ "timestamp",
1749
+ "resource_type",
1750
+ "resource_name",
1751
+ "value",
1752
+ "unit",
1753
+ "tag",
1754
+ "metadata",
1755
+ }
1756
+ _INJECTION_STATS_COLUMNS = {
1757
+ "id",
1758
+ "timestamp",
1759
+ "stat_name",
1760
+ "value",
1761
+ "round_number",
1762
+ "chunk_number",
1763
+ "sample_index",
1764
+ "background_index",
1765
+ "signal_class",
1766
+ "signal_type",
1767
+ "injection_stage",
1768
+ "is_finite",
1769
+ "slope_clamped",
1770
+ "tag",
1771
+ "metadata",
1772
+ "superseded",
1773
+ }
1774
+ _TRAINING_STATS_COLUMNS = {
1775
+ "id",
1776
+ "timestamp",
1777
+ "model_name",
1778
+ "stat_name",
1779
+ "value",
1780
+ "round_number",
1781
+ "epoch_number",
1782
+ "tag",
1783
+ "metadata",
1784
+ "superseded",
1785
+ "is_finite",
1786
+ }
1787
+ _LATENT_SNAPSHOTS_COLUMNS = {
1788
+ "id",
1789
+ "timestamp",
1790
+ "model_name",
1791
+ "round_number",
1792
+ "epoch_number",
1793
+ "step_number",
1794
+ "cadence_index",
1795
+ "signal_type",
1796
+ "latent_vector",
1797
+ "snr_base",
1798
+ "snr_range",
1799
+ "tag",
1800
+ "metadata",
1801
+ "superseded",
1802
+ }
1803
+ _INFERENCE_RESULTS_COLUMNS = {
1804
+ "id",
1805
+ "timestamp",
1806
+ "npy_path",
1807
+ "snippet_index",
1808
+ "prediction",
1809
+ "confidence",
1810
+ "latent_vector",
1811
+ "target",
1812
+ "session",
1813
+ "cadence_id",
1814
+ "band",
1815
+ "frequency_mhz",
1816
+ "timestamp_observed",
1817
+ "h5_path",
1818
+ "tag",
1819
+ "metadata",
1820
+ "superseded",
1821
+ "screening_proba",
1822
+ "mc_mean",
1823
+ "mc_std",
1824
+ }
1825
+ _INFERENCE_CADENCES_COLUMNS = {
1826
+ "id",
1827
+ "timestamp",
1828
+ "tag",
1829
+ "csv_path",
1830
+ "cadence_key",
1831
+ "npy_path",
1832
+ "status",
1833
+ "n_stamps",
1834
+ "n_candidates",
1835
+ "confidence_summary",
1836
+ "duration_s",
1837
+ "config_fingerprint",
1838
+ "superseded",
1839
+ }
1840
+ _PIPELINE_STAGES_COLUMNS = {
1841
+ "id",
1842
+ "stage",
1843
+ "start_time",
1844
+ "end_time",
1845
+ "duration_s",
1846
+ "tag",
1847
+ "metadata",
1848
+ }
1849
+
1850
+ @staticmethod
1851
+ def _build_select(table: str, columns: list[str] | None, whitelist: set[str]) -> str:
1852
+ """
1853
+ Build SELECT clause with optional column projection
1854
+ Use SELECT * if no columns provided; Use SELECT col_1, ..., col_n if columns provided
1855
+ Validates columns exist in table from whitelist. Raise ValueError if any column is invalid
1856
+
1857
+ Note that this API is not meant for public consumption, and should never be called with
1858
+ user input! Since `table` is interpolated directly, this represents a possible SQL injection
1859
+ vector
1860
+ """
1861
+ if columns is None:
1862
+ return f"SELECT * FROM {table}"
1863
+ # Validate all requested columns against whitelist
1864
+ invalid = set(columns) - whitelist
1865
+ if invalid:
1866
+ raise ValueError(f"Invalid column(s) for {table}: {invalid}")
1867
+ cols = ", ".join(columns)
1868
+ return f"SELECT {cols} FROM {table}"
1869
+
1870
+ @staticmethod
1871
+ def _add_str_filter(
1872
+ query: str,
1873
+ params: list,
1874
+ column: str,
1875
+ value: str | list[str],
1876
+ ) -> str:
1877
+ """
1878
+ Handle str or list[str] filter
1879
+ Use = for str; Use IN for list
1880
+
1881
+ Note that this is an asymmetric API:
1882
+ params (list -> mutable) is modified in place,
1883
+ query (str -> immutable) is returned as a new value.
1884
+ A cleaner approach would be to either return both (query, params), or modify both in-place
1885
+ However, functionally, this method is correct as is, and should be noted by the caller
1886
+ """
1887
+ if isinstance(value, list):
1888
+ if not value:
1889
+ return query
1890
+ placeholders = ", ".join("?" for _ in value)
1891
+ query += f" AND {column} IN ({placeholders})"
1892
+ params.extend(value)
1893
+ else:
1894
+ query += f" AND {column} = ?"
1895
+ params.append(value)
1896
+ return query
1897
+
1898
+ # NOTE: how to additionally filter by metadata (e.g. machine_name)?
1899
+ # NOTE: should we also let user filter by value (e.g. >= or <= some value)?
1900
+ def query_system_resource(
1901
+ self,
1902
+ resource_type: str | list[str] | None = None,
1903
+ resource_name: str | list[str] | None = None,
1904
+ tag: str | list[str] | None = None,
1905
+ start_time: float | None = None,
1906
+ end_time: float | None = None,
1907
+ columns: list[str] | None = None,
1908
+ ) -> list[dict[str, Any]]:
1909
+ """
1910
+ Query rows from system_resources as a list of dicts.
1911
+
1912
+ resource_type, resource_name, and tag accept either a single value (= filter) or a list
1913
+ (IN filter). start_time/end_time bound the timestamp range (unix time, inclusive on both
1914
+ ends). columns lets callers project a subset of fields; values are validated against
1915
+ _SYSTEM_RESOURCES_COLUMNS to block SQL injection through column names.
1916
+ """
1917
+ with self._get_connection() as conn:
1918
+ cursor = conn.cursor()
1919
+
1920
+ select = self._build_select("system_resources", columns, self._SYSTEM_RESOURCES_COLUMNS)
1921
+ # WHERE 1=1 is a way of building parametrized queries
1922
+ # Since 1=1 is always true, it does nothing functionally
1923
+ # But it allows us to safely add more conditions by appending AND clauses
1924
+ # While not breaking the query if none are added
1925
+ query = f"{select} WHERE 1=1"
1926
+ params = []
1927
+
1928
+ # Build the query dynamically based on user-specified conditions
1929
+ if resource_type:
1930
+ query = self._add_str_filter(query, params, "resource_type", resource_type)
1931
+
1932
+ if resource_name:
1933
+ query = self._add_str_filter(query, params, "resource_name", resource_name)
1934
+
1935
+ if tag:
1936
+ query = self._add_str_filter(query, params, "tag", tag)
1937
+
1938
+ if start_time is not None:
1939
+ query += " AND timestamp >= ?"
1940
+ params.append(start_time)
1941
+
1942
+ if end_time is not None:
1943
+ query += " AND timestamp <= ?"
1944
+ params.append(end_time)
1945
+
1946
+ # NOTE: removing ORDER BY to reduce indexing costs
1947
+ # Intentionally hard-coded. Update if schema changes
1948
+ # query += " ORDER BY tag, timestamp, resource_type, resource_name"
1949
+
1950
+ cursor.execute(query, params)
1951
+
1952
+ # Create a list of column names using query result's metadata
1953
+ result_columns = [desc[0] for desc in cursor.description]
1954
+ # Pair column names with values and return to user as a dict
1955
+ return [dict(zip(result_columns, row, strict=False)) for row in cursor.fetchall()]
1956
+
1957
+ def query_system_resource_decimated(
1958
+ self,
1959
+ tag: str,
1960
+ start_time: float,
1961
+ end_time: float,
1962
+ max_points_per_series: int,
1963
+ ) -> list[dict[str, Any]]:
1964
+ """
1965
+ Per-series uniformly-strided subset of system_resources rows for the teardown
1966
+ resource plot (#301): a multi-week catalog run accumulates tens of millions of
1967
+ rows while the plot renders ~2k px wide, so materializing every row cost a
1968
+ multi-GB teardown RAM spike for invisible detail. The stride is computed from the
1969
+ LARGEST series so every (resource_type, resource_name) line keeps up to
1970
+ max_points_per_series uniformly-spaced points; stride 1 degenerates to the full
1971
+ query. Returns the same dict shape as query_system_resource.
1972
+ """
1973
+ with self._get_connection() as conn:
1974
+ cursor = conn.cursor()
1975
+ cursor.execute(
1976
+ "SELECT MAX(cnt) FROM (SELECT COUNT(*) AS cnt FROM system_resources"
1977
+ " WHERE tag = ? AND timestamp >= ? AND timestamp <= ?"
1978
+ " GROUP BY resource_type, resource_name)",
1979
+ (tag, start_time, end_time),
1980
+ )
1981
+ max_series = cursor.fetchone()[0] or 0
1982
+ stride = max(1, -(-max_series // max_points_per_series)) # ceil div
1983
+ if stride == 1:
1984
+ # Small run: identical to the unstrided query (window overhead skipped)
1985
+ return self.query_system_resource(tag=tag, start_time=start_time, end_time=end_time)
1986
+ cursor.execute(
1987
+ "SELECT * FROM ("
1988
+ " SELECT *, ROW_NUMBER() OVER ("
1989
+ " PARTITION BY resource_type, resource_name ORDER BY timestamp"
1990
+ " ) AS _rn FROM system_resources"
1991
+ " WHERE tag = ? AND timestamp >= ? AND timestamp <= ?"
1992
+ ") WHERE (_rn - 1) % ? = 0",
1993
+ (tag, start_time, end_time, stride),
1994
+ )
1995
+ result_columns = [desc[0] for desc in cursor.description]
1996
+ rows = [dict(zip(result_columns, row, strict=False)) for row in cursor.fetchall()]
1997
+ for row in rows:
1998
+ row.pop("_rn", None)
1999
+ logger.info(
2000
+ f"Resource-plot query decimated by stride {stride} "
2001
+ f"(largest series {max_series} rows -> <= {max_points_per_series}/series)"
2002
+ )
2003
+ return rows
2004
+
2005
+ def query_injection_stat_time_span(
2006
+ self,
2007
+ tag: str,
2008
+ start_round_number: int | None = None,
2009
+ end_round_number: int | None = None,
2010
+ ) -> tuple[float, float] | None:
2011
+ """
2012
+ MIN/MAX timestamp over a tag's injection_stats rows, optionally bounded to a round
2013
+ range. One deliberate full-partition aggregate (no timestamp filter; superseded and
2014
+ non-finite rows included) that lets callers tighten the (tag, timestamp) index window
2015
+ for the many row-level queries that follow — idx_injection_stats_filter leads with
2016
+ (tag, timestamp) and round_number is not in it, so a run-wide window makes every
2017
+ round-scoped query re-scan the tag's whole history. The span covers every row a
2018
+ filtered query could return for those rounds, so intersecting a query window with it
2019
+ never changes that query's result set. Returns None when no rows match.
2020
+ """
2021
+ with self._get_connection() as conn:
2022
+ cursor = conn.cursor()
2023
+ query = "SELECT MIN(timestamp), MAX(timestamp) FROM injection_stats WHERE tag = ?"
2024
+ params: list = [tag]
2025
+ if start_round_number is not None:
2026
+ query += " AND round_number >= ?"
2027
+ params.append(start_round_number)
2028
+ if end_round_number is not None:
2029
+ query += " AND round_number <= ?"
2030
+ params.append(end_round_number)
2031
+ cursor.execute(query, params)
2032
+ row = cursor.fetchone()
2033
+ if row is None or row[0] is None:
2034
+ return None
2035
+ return float(row[0]), float(row[1])
2036
+
2037
+ # NOTE: how to additionally filter by metadata (e.g. machine_name)?
2038
+ # NOTE: should we also let user filter by value (e.g. >= or <= some value)?
2039
+ def query_injection_stat(
2040
+ self,
2041
+ stat_name: str | list[str] | None = None,
2042
+ start_round_number: int | None = None,
2043
+ end_round_number: int | None = None,
2044
+ start_chunk_number: int | None = None,
2045
+ end_chunk_number: int | None = None,
2046
+ start_sample_index: int | None = None,
2047
+ end_sample_index: int | None = None,
2048
+ start_background_index: int | None = None,
2049
+ end_background_index: int | None = None,
2050
+ signal_class: str | list[str] | None = None,
2051
+ signal_type: str | list[str] | None = None,
2052
+ injection_stage: str | list[str] | None = None,
2053
+ only_finite: bool = True,
2054
+ only_slope_clamped: bool | None = None,
2055
+ tag: str | list[str] | None = None,
2056
+ start_time: float | None = None,
2057
+ end_time: float | None = None,
2058
+ columns: list[str] | None = None,
2059
+ include_superseded: bool = False,
2060
+ ) -> list[dict[str, Any]]:
2061
+ """
2062
+ Query rows from injection_stats as a list of dicts.
2063
+
2064
+ String filters (stat_name, signal_class, signal_type, injection_stage, tag) accept either
2065
+ a single value (= filter) or a list (IN filter). signal_class is one of {main, false, true};
2066
+ signal_type is one of {false_no_signal, false_with_rfi, true_only_eti, true_eti_rfi};
2067
+ injection_stage is one of {A=pre-inj pre-norm, B=post-inj pre-norm, C=post-inj post-norm}.
2068
+ start_*/end_* pairs bound the corresponding integer column. only_finite (default True)
2069
+ drops rows where the stored value was non-finite at write time; only_slope_clamped, when
2070
+ not None, filters by the slope_clamped flag. include_superseded (default False) controls
2071
+ whether rows flagged stale by mark_superseded() are returned. columns is validated
2072
+ against _INJECTION_STATS_COLUMNS.
2073
+ """
2074
+ with self._get_connection() as conn:
2075
+ cursor = conn.cursor()
2076
+
2077
+ select = self._build_select("injection_stats", columns, self._INJECTION_STATS_COLUMNS)
2078
+ # WHERE 1=1 is a way of building parametrized queries
2079
+ # Since 1=1 is always true, it does nothing functionally
2080
+ # But it allows us to safely add more conditions by appending AND clauses
2081
+ # While not breaking the query if none are added
2082
+ query = f"{select} WHERE 1=1"
2083
+ params = []
2084
+
2085
+ # Build the query dynamically based on user-specified conditions
2086
+ if stat_name:
2087
+ query = self._add_str_filter(query, params, "stat_name", stat_name)
2088
+
2089
+ if start_round_number is not None:
2090
+ query += " AND round_number >= ?"
2091
+ params.append(start_round_number)
2092
+
2093
+ if end_round_number is not None:
2094
+ query += " AND round_number <= ?"
2095
+ params.append(end_round_number)
2096
+
2097
+ if start_chunk_number is not None:
2098
+ query += " AND chunk_number >= ?"
2099
+ params.append(start_chunk_number)
2100
+
2101
+ if end_chunk_number is not None:
2102
+ query += " AND chunk_number <= ?"
2103
+ params.append(end_chunk_number)
2104
+
2105
+ if start_sample_index is not None:
2106
+ query += " AND sample_index >= ?"
2107
+ params.append(start_sample_index)
2108
+
2109
+ if end_sample_index is not None:
2110
+ query += " AND sample_index <= ?"
2111
+ params.append(end_sample_index)
2112
+
2113
+ if start_background_index is not None:
2114
+ query += " AND background_index >= ?"
2115
+ params.append(start_background_index)
2116
+
2117
+ if end_background_index is not None:
2118
+ query += " AND background_index <= ?"
2119
+ params.append(end_background_index)
2120
+
2121
+ if signal_class:
2122
+ query = self._add_str_filter(query, params, "signal_class", signal_class)
2123
+
2124
+ if signal_type:
2125
+ query = self._add_str_filter(query, params, "signal_type", signal_type)
2126
+
2127
+ if injection_stage:
2128
+ query = self._add_str_filter(query, params, "injection_stage", injection_stage)
2129
+
2130
+ if only_finite:
2131
+ query += " AND is_finite = 1"
2132
+
2133
+ if only_slope_clamped is not None:
2134
+ query += " AND slope_clamped = ?"
2135
+ params.append(1 if only_slope_clamped else 0)
2136
+
2137
+ if not include_superseded:
2138
+ query += " AND superseded = 0"
2139
+
2140
+ if tag:
2141
+ query = self._add_str_filter(query, params, "tag", tag)
2142
+
2143
+ if start_time is not None:
2144
+ query += " AND timestamp >= ?"
2145
+ params.append(start_time)
2146
+
2147
+ if end_time is not None:
2148
+ query += " AND timestamp <= ?"
2149
+ params.append(end_time)
2150
+
2151
+ # NOTE: removing ORDER BY to reduce indexing costs
2152
+ # Intentionally hard-coded. Update if schema changes
2153
+ # query += " ORDER BY tag, signal_class, signal_type, round_number, chunk_number, sample_index, stat_name"
2154
+
2155
+ cursor.execute(query, params)
2156
+
2157
+ # Create a list of column names using query result's metadata
2158
+ result_columns = [desc[0] for desc in cursor.description]
2159
+ # Pair column names with values and return to user as a dict
2160
+ return [dict(zip(result_columns, row, strict=False)) for row in cursor.fetchall()]
2161
+
2162
+ def query_injection_stat_stability(
2163
+ self,
2164
+ stat_name: str | list[str] | None = None,
2165
+ start_round_number: int | None = None,
2166
+ end_round_number: int | None = None,
2167
+ injection_stage: str | list[str] | None = None,
2168
+ tag: str | list[str] | None = None,
2169
+ start_time: float | None = None,
2170
+ end_time: float | None = None,
2171
+ include_superseded: bool = False,
2172
+ ) -> list[dict[str, Any]]:
2173
+ """
2174
+ Per-round aggregation of injection_stats for sanitization and clamping rates.
2175
+
2176
+ Returns dicts of {round_number, total_count, non_finite_count, clamped_count} so callers
2177
+ can compute sanitization rate (non_finite_count / total_count) and clamping rate
2178
+ (clamped_count / total_count) without fetching every row. Preferred over
2179
+ query_injection_stat + Python-side aggregation: SQLite's native COUNT/SUM in C beats
2180
+ materializing rows into Python dicts and iterating. String filters accept str or list[str];
2181
+ start_round_number/end_round_number bound the rounds aggregated (#277 — callers scope to
2182
+ the rounds being plotted so pre-generated later rounds don't show partial bars).
2183
+ """
2184
+ with self._get_connection() as conn:
2185
+ cursor = conn.cursor()
2186
+
2187
+ query = """
2188
+ SELECT
2189
+ round_number,
2190
+ COUNT(*) as total_count,
2191
+ SUM(CASE WHEN is_finite = 0 THEN 1 ELSE 0 END) as non_finite_count,
2192
+ SUM(CASE WHEN slope_clamped = 1 THEN 1 ELSE 0 END) as clamped_count
2193
+ FROM injection_stats
2194
+ WHERE 1=1
2195
+ """
2196
+ params: list = []
2197
+
2198
+ if stat_name:
2199
+ query = self._add_str_filter(query, params, "stat_name", stat_name)
2200
+
2201
+ if start_round_number is not None:
2202
+ query += " AND round_number >= ?"
2203
+ params.append(start_round_number)
2204
+
2205
+ if end_round_number is not None:
2206
+ query += " AND round_number <= ?"
2207
+ params.append(end_round_number)
2208
+
2209
+ if injection_stage:
2210
+ query = self._add_str_filter(query, params, "injection_stage", injection_stage)
2211
+
2212
+ if tag:
2213
+ query = self._add_str_filter(query, params, "tag", tag)
2214
+
2215
+ if start_time is not None:
2216
+ query += " AND timestamp >= ?"
2217
+ params.append(start_time)
2218
+
2219
+ if end_time is not None:
2220
+ query += " AND timestamp <= ?"
2221
+ params.append(end_time)
2222
+
2223
+ if not include_superseded:
2224
+ query += " AND superseded = 0"
2225
+
2226
+ query += " GROUP BY round_number ORDER BY round_number"
2227
+
2228
+ cursor.execute(query, params)
2229
+
2230
+ # Create a list of column names using query result's metadata
2231
+ result_columns = [desc[0] for desc in cursor.description]
2232
+ # Pair column names with values and return to user as a dict
2233
+ return [dict(zip(result_columns, row, strict=False)) for row in cursor.fetchall()]
2234
+
2235
+ # NOTE: how to additionally filter by metadata (e.g. machine_name)?
2236
+ # NOTE: should we also let user filter by value (e.g. >= or <= some value)?
2237
+ def query_training_stat(
2238
+ self,
2239
+ model_name: str | list[str] | None = None,
2240
+ stat_name: str | list[str] | None = None,
2241
+ start_round_number: int | None = None,
2242
+ end_round_number: int | None = None,
2243
+ start_epoch_number: int | None = None,
2244
+ end_epoch_number: int | None = None,
2245
+ tag: str | list[str] | None = None,
2246
+ start_time: float | None = None,
2247
+ end_time: float | None = None,
2248
+ columns: list[str] | None = None,
2249
+ include_superseded: bool = False,
2250
+ only_finite: bool = True,
2251
+ ) -> list[dict[str, Any]]:
2252
+ """
2253
+ Query rows from training_stats as a list of dicts.
2254
+
2255
+ model_name, stat_name, and tag accept either a single value (= filter) or a list
2256
+ (IN filter). start_*/end_* pairs bound the corresponding integer column (inclusive).
2257
+ include_superseded (default False) controls whether rows flagged stale by
2258
+ mark_superseded() are returned. only_finite (default True) drops rows whose stat value
2259
+ was non-finite at write time (stored as 0.0 with is_finite=0, #289) — pass False to
2260
+ inspect them. columns is validated against _TRAINING_STATS_COLUMNS.
2261
+ """
2262
+ with self._get_connection() as conn:
2263
+ cursor = conn.cursor()
2264
+
2265
+ select = self._build_select("training_stats", columns, self._TRAINING_STATS_COLUMNS)
2266
+ # WHERE 1=1 is a way of building parametrized queries
2267
+ # Since 1=1 is always true, it does nothing functionally
2268
+ # But it allows us to safely add more conditions by appending AND clauses
2269
+ # While not breaking the query if none are added
2270
+ query = f"{select} WHERE 1=1"
2271
+ params = []
2272
+
2273
+ # Build the query dynamically based on user-specified conditions
2274
+ if model_name:
2275
+ query = self._add_str_filter(query, params, "model_name", model_name)
2276
+
2277
+ if stat_name:
2278
+ query = self._add_str_filter(query, params, "stat_name", stat_name)
2279
+
2280
+ if start_round_number is not None:
2281
+ query += " AND round_number >= ?"
2282
+ params.append(start_round_number)
2283
+
2284
+ if end_round_number is not None:
2285
+ query += " AND round_number <= ?"
2286
+ params.append(end_round_number)
2287
+
2288
+ if start_epoch_number is not None:
2289
+ query += " AND epoch_number >= ?"
2290
+ params.append(start_epoch_number)
2291
+
2292
+ if end_epoch_number is not None:
2293
+ query += " AND epoch_number <= ?"
2294
+ params.append(end_epoch_number)
2295
+
2296
+ if not include_superseded:
2297
+ query += " AND superseded = 0"
2298
+
2299
+ if only_finite:
2300
+ query += " AND is_finite = 1"
2301
+
2302
+ if tag:
2303
+ query = self._add_str_filter(query, params, "tag", tag)
2304
+
2305
+ if start_time is not None:
2306
+ query += " AND timestamp >= ?"
2307
+ params.append(start_time)
2308
+
2309
+ if end_time is not None:
2310
+ query += " AND timestamp <= ?"
2311
+ params.append(end_time)
2312
+
2313
+ # NOTE: removing ORDER BY to reduce indexing costs
2314
+ # Intentionally hard-coded. Update if schema changes
2315
+ # query += " ORDER BY tag, model_name, round_number, epoch_number, stat_name"
2316
+
2317
+ cursor.execute(query, params)
2318
+
2319
+ # Create a list of column names using query result's metadata
2320
+ result_columns = [desc[0] for desc in cursor.description]
2321
+ # Pair column names with values and return to user as a dict
2322
+ return [dict(zip(result_columns, row, strict=False)) for row in cursor.fetchall()]
2323
+
2324
+ # NOTE: how to additionally filter by metadata (e.g. machine_name)?
2325
+ def query_latent_snapshots(
2326
+ self,
2327
+ model_name: str | list[str] | None = None,
2328
+ round_number: int | None = None,
2329
+ epoch_number: int | None = None,
2330
+ step_number: int | None = None,
2331
+ signal_type: str | list[str] | None = None,
2332
+ tag: str | list[str] | None = None,
2333
+ start_time: float | None = None,
2334
+ end_time: float | None = None,
2335
+ columns: list[str] | None = None,
2336
+ include_superseded: bool = False,
2337
+ ) -> list[dict[str, Any]]:
2338
+ """
2339
+ Query rows from latent_snapshots as a list of dicts.
2340
+
2341
+ model_name, signal_type, and tag accept either a single value (= filter) or a list
2342
+ (IN filter). round_number/epoch_number/step_number are exact-match filters (no range
2343
+ variant). The returned latent_vector field is a JSON string — callers parse it with
2344
+ json.loads. include_superseded (default False) controls whether rows flagged stale by
2345
+ mark_superseded() are returned. columns is validated against _LATENT_SNAPSHOTS_COLUMNS.
2346
+ """
2347
+ with self._get_connection() as conn:
2348
+ cursor = conn.cursor()
2349
+
2350
+ select = self._build_select("latent_snapshots", columns, self._LATENT_SNAPSHOTS_COLUMNS)
2351
+ # WHERE 1=1 is a way of building parametrized queries
2352
+ # Since 1=1 is always true, it does nothing functionally
2353
+ # But it allows us to safely add more conditions by appending AND clauses
2354
+ # While not breaking the query if none are added
2355
+ query = f"{select} WHERE 1=1"
2356
+ params = []
2357
+
2358
+ # Build the query dynamically based on user-specified conditions
2359
+ if model_name:
2360
+ query = self._add_str_filter(query, params, "model_name", model_name)
2361
+
2362
+ if round_number is not None:
2363
+ query += " AND round_number = ?"
2364
+ params.append(round_number)
2365
+
2366
+ if epoch_number is not None:
2367
+ query += " AND epoch_number = ?"
2368
+ params.append(epoch_number)
2369
+
2370
+ if step_number is not None:
2371
+ query += " AND step_number = ?"
2372
+ params.append(step_number)
2373
+
2374
+ if signal_type:
2375
+ query = self._add_str_filter(query, params, "signal_type", signal_type)
2376
+
2377
+ if tag:
2378
+ query = self._add_str_filter(query, params, "tag", tag)
2379
+
2380
+ if start_time is not None:
2381
+ query += " AND timestamp >= ?"
2382
+ params.append(start_time)
2383
+
2384
+ if end_time is not None:
2385
+ query += " AND timestamp <= ?"
2386
+ params.append(end_time)
2387
+
2388
+ if not include_superseded:
2389
+ query += " AND superseded = 0"
2390
+
2391
+ # NOTE: hard-coded ORDER BY? add template index to init_db
2392
+
2393
+ cursor.execute(query, params)
2394
+
2395
+ # Create a list of column names using query result's metadata
2396
+ result_columns = [desc[0] for desc in cursor.description]
2397
+ # Pair column names with values and return to user as a dict
2398
+ return [dict(zip(result_columns, row, strict=False)) for row in cursor.fetchall()]
2399
+
2400
+ def query_latent_snapshot_keys(
2401
+ self,
2402
+ tag: str | list[str] | None = None,
2403
+ start_time: float | None = None,
2404
+ end_time: float | None = None,
2405
+ include_superseded: bool = False,
2406
+ ) -> list[dict[str, Any]]:
2407
+ """
2408
+ Distinct snapshot keys (model_name, round_number, epoch_number, step_number, snr_base,
2409
+ snr_range) sorted by training progression. Avoids the cost of loading every row through
2410
+ query_latent_snapshots() when the caller only needs the set of available keys.
2411
+ """
2412
+ with self._get_connection() as conn:
2413
+ cursor = conn.cursor()
2414
+
2415
+ query = """
2416
+ SELECT DISTINCT model_name, round_number, epoch_number, step_number, snr_base, snr_range
2417
+ FROM latent_snapshots
2418
+ WHERE 1=1
2419
+ """
2420
+ params: list = []
2421
+
2422
+ # Build the query dynamically based on user-specified conditions
2423
+ if tag:
2424
+ query = self._add_str_filter(query, params, "tag", tag)
2425
+
2426
+ if start_time is not None:
2427
+ query += " AND timestamp >= ?"
2428
+ params.append(start_time)
2429
+
2430
+ if end_time is not None:
2431
+ query += " AND timestamp <= ?"
2432
+ params.append(end_time)
2433
+
2434
+ if not include_superseded:
2435
+ query += " AND superseded = 0"
2436
+
2437
+ # The DISTINCT set includes snr_base/snr_range, which no index carries, so this
2438
+ # is a tag-partition scan regardless of index shape — the tag prefix of
2439
+ # idx_latent_snapshots_by_key still bounds it to one tag. It runs once per GIF
2440
+ # pass, and the resulting key set is small enough that the filesort for this
2441
+ # ORDER BY (model_name leading, unlike the index) is negligible
2442
+ query += " ORDER BY model_name, round_number, epoch_number, step_number"
2443
+
2444
+ cursor.execute(query, params)
2445
+
2446
+ # Create a list of column names using query result's metadata
2447
+ result_columns = [desc[0] for desc in cursor.description]
2448
+ # Pair column names with values and return to user as a dict
2449
+ return [dict(zip(result_columns, row, strict=False)) for row in cursor.fetchall()]
2450
+
2451
+ # NOTE: how to additionally filter by metadata (e.g. machine_name)?
2452
+ def query_inference_result(
2453
+ self,
2454
+ npy_path: str | list[str] | None = None,
2455
+ start_snippet_index: int | None = None,
2456
+ end_snippet_index: int | None = None,
2457
+ prediction: int | None = None,
2458
+ min_confidence: float | None = None,
2459
+ max_confidence: float | None = None,
2460
+ target: str | list[str] | None = None,
2461
+ session: str | list[str] | None = None,
2462
+ cadence_id: int | None = None,
2463
+ band: str | list[str] | None = None,
2464
+ min_frequency_mhz: float | None = None,
2465
+ max_frequency_mhz: float | None = None,
2466
+ start_timestamp_observed: float | None = None,
2467
+ end_timestamp_observed: float | None = None,
2468
+ h5_path: str | list[str] | None = None,
2469
+ tag: str | list[str] | None = None,
2470
+ start_time: float | None = None,
2471
+ end_time: float | None = None,
2472
+ columns: list[str] | None = None,
2473
+ include_superseded: bool = False,
2474
+ ) -> list[dict[str, Any]]:
2475
+ """
2476
+ Query rows from inference_results as a list of dicts.
2477
+
2478
+ String filters (npy_path, target, session, band, h5_path, tag) accept either a single
2479
+ value (= filter) or a list (IN filter). prediction is exact-match (0=RFI, 1=candidate).
2480
+ start_*/end_*/min_*/max_* pairs bound the corresponding numeric column (inclusive).
2481
+ timestamp_observed is the original observation time (distinct from the write-time
2482
+ timestamp). include_superseded (default False) controls whether rows flagged stale by
2483
+ mark_superseded() are returned. columns is validated against _INFERENCE_RESULTS_COLUMNS.
2484
+ """
2485
+ with self._get_connection() as conn:
2486
+ cursor = conn.cursor()
2487
+
2488
+ select = self._build_select(
2489
+ "inference_results", columns, self._INFERENCE_RESULTS_COLUMNS
2490
+ )
2491
+ # WHERE 1=1 is a way of building parametrized queries
2492
+ # Since 1=1 is always true, it does nothing functionally
2493
+ # But it allows us to safely add more conditions by appending AND clauses
2494
+ # While not breaking the query if none are added
2495
+ query = f"{select} WHERE 1=1"
2496
+ params = []
2497
+
2498
+ # Build the query dynamically based on user-specified conditions
2499
+ if npy_path:
2500
+ query = self._add_str_filter(query, params, "npy_path", npy_path)
2501
+
2502
+ if start_snippet_index is not None:
2503
+ query += " AND snippet_index >= ?"
2504
+ params.append(start_snippet_index)
2505
+
2506
+ if end_snippet_index is not None:
2507
+ query += " AND snippet_index <= ?"
2508
+ params.append(end_snippet_index)
2509
+
2510
+ if prediction is not None:
2511
+ query += " AND prediction = ?"
2512
+ params.append(prediction)
2513
+
2514
+ if min_confidence is not None:
2515
+ query += " AND confidence >= ?"
2516
+ params.append(min_confidence)
2517
+
2518
+ if max_confidence is not None:
2519
+ query += " AND confidence <= ?"
2520
+ params.append(max_confidence)
2521
+
2522
+ if target:
2523
+ query = self._add_str_filter(query, params, "target", target)
2524
+
2525
+ if session:
2526
+ query = self._add_str_filter(query, params, "session", session)
2527
+
2528
+ if cadence_id is not None:
2529
+ query += " AND cadence_id = ?"
2530
+ params.append(cadence_id)
2531
+
2532
+ if band:
2533
+ query = self._add_str_filter(query, params, "band", band)
2534
+
2535
+ if min_frequency_mhz is not None:
2536
+ query += " AND frequency_mhz >= ?"
2537
+ params.append(min_frequency_mhz)
2538
+
2539
+ if max_frequency_mhz is not None:
2540
+ query += " AND frequency_mhz <= ?"
2541
+ params.append(max_frequency_mhz)
2542
+
2543
+ if start_timestamp_observed is not None:
2544
+ query += " AND timestamp_observed >= ?"
2545
+ params.append(start_timestamp_observed)
2546
+
2547
+ if end_timestamp_observed is not None:
2548
+ query += " AND timestamp_observed <= ?"
2549
+ params.append(end_timestamp_observed)
2550
+
2551
+ if h5_path:
2552
+ query = self._add_str_filter(query, params, "h5_path", h5_path)
2553
+
2554
+ if tag:
2555
+ query = self._add_str_filter(query, params, "tag", tag)
2556
+
2557
+ if start_time is not None:
2558
+ query += " AND timestamp >= ?"
2559
+ params.append(start_time)
2560
+
2561
+ if end_time is not None:
2562
+ query += " AND timestamp <= ?"
2563
+ params.append(end_time)
2564
+
2565
+ if not include_superseded:
2566
+ query += " AND superseded = 0"
2567
+
2568
+ # NOTE: hard-coded ORDER BY? add template index to init_db
2569
+
2570
+ cursor.execute(query, params)
2571
+
2572
+ # Create a list of column names using query result's metadata
2573
+ result_columns = [desc[0] for desc in cursor.description]
2574
+ # Pair column names with values and return to user as a dict
2575
+ return [dict(zip(result_columns, row, strict=False)) for row in cursor.fetchall()]
2576
+
2577
+ def query_inference_cadences(
2578
+ self,
2579
+ npy_path: str | list[str] | None = None,
2580
+ status: str | list[str] | None = None,
2581
+ csv_path: str | list[str] | None = None,
2582
+ tag: str | list[str] | None = None,
2583
+ start_time: float | None = None,
2584
+ end_time: float | None = None,
2585
+ columns: list[str] | None = None,
2586
+ include_superseded: bool = False,
2587
+ ) -> list[dict[str, Any]]:
2588
+ """
2589
+ Query rows from the inference_cadences run manifest as a list of dicts.
2590
+
2591
+ String filters (npy_path, status, csv_path, tag) accept either a single value
2592
+ (= filter) or a list (IN filter); status is one of {'preprocessed', 'inferred',
2593
+ 'failed'}. The returned cadence_key and confidence_summary fields are JSON strings —
2594
+ callers parse them with json.loads. include_superseded (default False) controls
2595
+ whether rows flagged stale by mark_superseded() are returned — the resume flow relies
2596
+ on the default so a cadence whose 'inferred' row was superseded is re-attempted.
2597
+ columns is validated against _INFERENCE_CADENCES_COLUMNS.
2598
+ """
2599
+ with self._get_connection() as conn:
2600
+ cursor = conn.cursor()
2601
+
2602
+ select = self._build_select(
2603
+ "inference_cadences", columns, self._INFERENCE_CADENCES_COLUMNS
2604
+ )
2605
+ # WHERE 1=1 is a way of building parametrized queries
2606
+ # Since 1=1 is always true, it does nothing functionally
2607
+ # But it allows us to safely add more conditions by appending AND clauses
2608
+ # While not breaking the query if none are added
2609
+ query = f"{select} WHERE 1=1"
2610
+ params: list = []
2611
+
2612
+ # Build the query dynamically based on user-specified conditions
2613
+ if npy_path:
2614
+ query = self._add_str_filter(query, params, "npy_path", npy_path)
2615
+
2616
+ if status:
2617
+ query = self._add_str_filter(query, params, "status", status)
2618
+
2619
+ if csv_path:
2620
+ query = self._add_str_filter(query, params, "csv_path", csv_path)
2621
+
2622
+ if tag:
2623
+ query = self._add_str_filter(query, params, "tag", tag)
2624
+
2625
+ if start_time is not None:
2626
+ query += " AND timestamp >= ?"
2627
+ params.append(start_time)
2628
+
2629
+ if end_time is not None:
2630
+ query += " AND timestamp <= ?"
2631
+ params.append(end_time)
2632
+
2633
+ if not include_superseded:
2634
+ query += " AND superseded = 0"
2635
+
2636
+ cursor.execute(query, params)
2637
+
2638
+ # Create a list of column names using query result's metadata
2639
+ result_columns = [desc[0] for desc in cursor.description]
2640
+ # Pair column names with values and return to user as a dict
2641
+ return [dict(zip(result_columns, row, strict=False)) for row in cursor.fetchall()]
2642
+
2643
+ def query_pipeline_stages(
2644
+ self,
2645
+ stage: str | list[str] | None = None,
2646
+ tag: str | list[str] | None = None,
2647
+ start_time: float | None = None,
2648
+ end_time: float | None = None,
2649
+ columns: list[str] | None = None,
2650
+ ) -> list[dict[str, Any]]:
2651
+ """
2652
+ Query rows from pipeline_stages as a list of dicts, ordered by start_time.
2653
+
2654
+ stage and tag accept either a single value (= filter) or a list (IN filter).
2655
+ start_time/end_time bound the row's start_time column (unix time, inclusive on both
2656
+ ends) — a span that started inside the window but ended after end_time is still
2657
+ returned. The returned metadata field is a JSON string or None — callers parse it
2658
+ with json.loads. columns is validated against _PIPELINE_STAGES_COLUMNS.
2659
+ """
2660
+ with self._get_connection() as conn:
2661
+ cursor = conn.cursor()
2662
+
2663
+ select = self._build_select("pipeline_stages", columns, self._PIPELINE_STAGES_COLUMNS)
2664
+ # WHERE 1=1 is a way of building parametrized queries
2665
+ # Since 1=1 is always true, it does nothing functionally
2666
+ # But it allows us to safely add more conditions by appending AND clauses
2667
+ # While not breaking the query if none are added
2668
+ query = f"{select} WHERE 1=1"
2669
+ params: list = []
2670
+
2671
+ # Build the query dynamically based on user-specified conditions
2672
+ if stage:
2673
+ query = self._add_str_filter(query, params, "stage", stage)
2674
+
2675
+ if tag:
2676
+ query = self._add_str_filter(query, params, "tag", tag)
2677
+
2678
+ if start_time is not None:
2679
+ query += " AND start_time >= ?"
2680
+ params.append(start_time)
2681
+
2682
+ if end_time is not None:
2683
+ query += " AND start_time <= ?"
2684
+ params.append(end_time)
2685
+
2686
+ # Chronological order is what every consumer (report tree, timeline plot,
2687
+ # monitor overlay) wants; the table stays small (one row per stage span), so
2688
+ # the filesort on top of idx_pipeline_stages_filter is cheap
2689
+ query += " ORDER BY start_time"
2690
+
2691
+ cursor.execute(query, params)
2692
+
2693
+ # Create a list of column names using query result's metadata
2694
+ result_columns = [desc[0] for desc in cursor.description]
2695
+ # Pair column names with values and return to user as a dict
2696
+ return [dict(zip(result_columns, row, strict=False)) for row in cursor.fetchall()]
2697
+
2698
+ # NOTE: this call gets expensive as db grows. create a separate schema to track num_rows_added per pipeline run. then count & update as part of db cleanup routine? or use SQLite's dbstat virtual table or periodic ANALYZE?
2699
+ def get_db_stats(self) -> dict[str, Any]:
2700
+ """Get summary statistics for the database.
2701
+
2702
+ On-demand diagnostics only — deliberately NOT called at init: the per-table
2703
+ COUNT(*) scans cost ~13 min on a cold-cache catalog-scale DB (#301)."""
2704
+ with self._get_connection() as conn:
2705
+ cursor = conn.cursor()
2706
+
2707
+ stats = {}
2708
+
2709
+ # Row counts
2710
+ cursor.execute("SELECT COUNT(*) FROM system_resources")
2711
+ stats["system_resources_row_count"] = cursor.fetchone()[0]
2712
+
2713
+ cursor.execute("SELECT COUNT(*) FROM injection_stats")
2714
+ stats["injection_stats_row_count"] = cursor.fetchone()[0]
2715
+
2716
+ cursor.execute("SELECT COUNT(*) FROM training_stats")
2717
+ stats["training_stats_row_count"] = cursor.fetchone()[0]
2718
+
2719
+ cursor.execute("SELECT COUNT(*) FROM latent_snapshots")
2720
+ stats["latent_snapshots_row_count"] = cursor.fetchone()[0]
2721
+
2722
+ cursor.execute("SELECT COUNT(*) FROM inference_results")
2723
+ stats["inference_results_row_count"] = cursor.fetchone()[0]
2724
+
2725
+ cursor.execute("SELECT COUNT(*) FROM inference_cadences")
2726
+ stats["inference_cadences_row_count"] = cursor.fetchone()[0]
2727
+
2728
+ cursor.execute("SELECT COUNT(*) FROM pipeline_stages")
2729
+ stats["pipeline_stages_row_count"] = cursor.fetchone()[0]
2730
+
2731
+ # Time range
2732
+ # Use system_resources as proxy
2733
+ cursor.execute("""
2734
+ SELECT MIN(timestamp), MAX(timestamp)
2735
+ FROM system_resources
2736
+ """)
2737
+ min_time, max_time = cursor.fetchone()
2738
+ stats["min_timestamp"] = min_time
2739
+ stats["max_timestamp"] = max_time
2740
+
2741
+ # Database size
2742
+ cursor.execute(
2743
+ "SELECT page_count * page_size as size FROM pragma_page_count(), pragma_page_size()"
2744
+ )
2745
+ stats["db_size_bytes"] = cursor.fetchone()[0]
2746
+ stats["db_size_mb"] = stats["db_size_bytes"] / (1024 * 1024)
2747
+
2748
+ return stats
2749
+
2750
+
2751
+ def init_db() -> Database:
2752
+ """
2753
+ Initialize global database instance (call once at startup)
2754
+ """
2755
+ db = Database()
2756
+ db.start()
2757
+
2758
+ register_db(db)
2759
+
2760
+ return db
2761
+
2762
+
2763
+ def get_db() -> Database | None:
2764
+ """Get the global database instance"""
2765
+ db = Database._instance
2766
+
2767
+ if db is None:
2768
+ logger.warning("No database instance initialized")
2769
+
2770
+ return db
2771
+
2772
+
2773
+ # TODO: complete this function
2774
+ # NOTE: where should source of truth be? (bla0, blpc2, blpc3, other)
2775
+ def merge_db() -> None:
2776
+ """Merge (join & dedup) two different aetherscan.db files into a single source-of-truth"""
2777
+ pass
2778
+
2779
+
2780
+ def shutdown_db() -> None:
2781
+ """Shutdown the global database instance (call on exit)"""
2782
+ db = Database._instance
2783
+
2784
+ if db is None:
2785
+ logger.warning("No database instance initialized")
2786
+ return
2787
+
2788
+ db.stop()
2789
+ Database._reset()