diff --git a/weightslab/backend/logger.py b/weightslab/backend/logger.py index a29bdf07..44e272e3 100644 --- a/weightslab/backend/logger.py +++ b/weightslab/backend/logger.py @@ -31,6 +31,7 @@ import functools import json +import logging import os import threading import time @@ -42,6 +43,8 @@ from weightslab.backend.ledgers import get_logger, register_logger, get_checkpoint_manager +logger = logging.getLogger(__name__) + # Column order for each table's staging buffer / bulk insert. _SIGNAL_COLS = [ @@ -112,45 +115,141 @@ def __init__(self, register: bool = True, db_path: str = ":memory:") -> None: # Init checkpoint manager for experiment hash retrieval (if available) self.chkpt_manager = get_checkpoint_manager() + # If no explicit db_path was given but a checkpoint manager already + # exists, persist history to an on-disk DuckDB file under its loggers/ + # dir. (The reverse ordering — CM created after the logger — is handled + # by CheckpointManager.__init__ calling set_db_path on the live logger.) + if db_path == ":memory:": + try: + loggers_dir = getattr(self.chkpt_manager, "loggers_dir", None) + if loggers_dir: + self.set_db_path(os.path.join(str(loggers_dir), "loggers.duckdb")) + except Exception: + pass + # ------------------------------------------------------------------ # DuckDB plumbing # ------------------------------------------------------------------ + @staticmethod + def _schema_ddl(prefix: str = "") -> str: + """Return the CREATE-TABLE DDL for all history tables. + + ``prefix`` lets the same schema be created inside an attached database + (e.g. ``"ondisk."``) when migrating from in-memory to a file. + """ + return f""" + CREATE TABLE IF NOT EXISTS {prefix}signals ( + metric_name VARCHAR, + experiment_hash VARCHAR, + step INTEGER, + metric_value DOUBLE, + timestamp BIGINT, + audit_mode BOOLEAN, + is_evaluation_marker BOOLEAN, + split_name VARCHAR, + evaluation_tags VARCHAR, + point_note VARCHAR, + seq BIGINT + ); + CREATE TABLE IF NOT EXISTS {prefix}per_sample ( + metric_name VARCHAR, + experiment_hash VARCHAR, + sample_id VARCHAR, + step INTEGER, + value REAL, + seq BIGINT + ); + CREATE TABLE IF NOT EXISTS {prefix}per_instance ( + metric_name VARCHAR, + experiment_hash VARCHAR, + sample_id VARCHAR, + annotation_id INTEGER, + step INTEGER, + value REAL, + seq BIGINT + ); + """ + def _ensure_tables(self) -> None: with self._lock: - self._conn.execute( - """ - CREATE TABLE IF NOT EXISTS signals ( - metric_name VARCHAR, - experiment_hash VARCHAR, - step INTEGER, - metric_value DOUBLE, - timestamp BIGINT, - audit_mode BOOLEAN, - is_evaluation_marker BOOLEAN, - split_name VARCHAR, - evaluation_tags VARCHAR, - point_note VARCHAR, - seq BIGINT - ); - CREATE TABLE IF NOT EXISTS per_sample ( - metric_name VARCHAR, - experiment_hash VARCHAR, - sample_id VARCHAR, - step INTEGER, - value REAL, - seq BIGINT - ); - CREATE TABLE IF NOT EXISTS per_instance ( - metric_name VARCHAR, - experiment_hash VARCHAR, - sample_id VARCHAR, - annotation_id INTEGER, - step INTEGER, - value REAL, - seq BIGINT - ); - """ - ) + self._conn.execute(self._schema_ddl()) + + _HISTORY_TABLES = ("signals", "per_sample", "per_instance") + + def set_db_path(self, db_path) -> None: + """Persist signal history to an on-disk DuckDB file. + + Call this once, early in setup. If the file already exists (resume), + its data is adopted as-is. If it does not, whatever little is currently + in the in-memory DB is migrated into the new file. Either way the file + becomes the live connection afterwards. + + The hot logging path is unaffected: ``add_scalars`` still stages to RAM + and only bulk-flushes to DuckDB lazily. DuckDB serves reads from its + in-memory buffer pool, so this adds durability, not per-read disk hits. + """ + if not db_path or db_path == ":memory:": + return + db_path = str(db_path) + + with self._lock: + if self._db_path == db_path: + return + + parent = os.path.dirname(db_path) + if parent: + os.makedirs(parent, exist_ok=True) + + file_preexists = os.path.exists(db_path) + + try: + # Make sure staged rows are materialized before we migrate. + self._flush_stage() + + if not file_preexists: + # Fresh file: copy whatever is in the in-memory DB into it. + # ATTACH doesn't accept bind parameters, so inline the path + # with SQL-escaped quotes. + escaped = db_path.replace("'", "''") + self._conn.execute(f"ATTACH '{escaped}' AS ondisk") + self._conn.execute(self._schema_ddl(prefix="ondisk.")) + for tbl in self._HISTORY_TABLES: + self._conn.execute( + f"INSERT INTO ondisk.{tbl} SELECT * FROM {tbl}" + ) + self._conn.execute("DETACH ondisk") + + # Adopt the on-disk file as the live connection. On resume this + # is the source of truth; the fresh in-memory rows are ignored. + self._conn.close() + self._conn = duckdb.connect(database=db_path) + self._db_path = db_path + self._ensure_tables() + self._invalidate_qps_cache() + self._restore_runtime_state_from_db() + logger.info( + f"[LoggerQueue] Signal history persisted on disk at {db_path} " + f"({'adopted existing' if file_preexists else 'new'} database)." + ) + except Exception as exc: + logger.warning( + f"[LoggerQueue] Failed to enable on-disk persistence at " + f"{db_path}: {exc}. Keeping in-memory history." + ) + + def flush_to_disk(self) -> None: + """Flush staged rows and force a DuckDB checkpoint to the file. + + No-op for an in-memory database. Call at checkpoint time so history is + durable even without a clean shutdown (DuckDB also replays its WAL on + the next open, so this is belt-and-braces).""" + with self._lock: + try: + self._flush_stage() + if self._db_path != ":memory:": + self._conn.execute("CHECKPOINT") + except Exception as exc: + logger.warning(f"[LoggerQueue] flush_to_disk failed: {exc}") def _restore_runtime_state_from_db(self) -> None: """Repopulate seq counter and graph names from an existing (file) DB.""" diff --git a/weightslab/components/checkpoint_manager.py b/weightslab/components/checkpoint_manager.py index 0c029179..528d6e9d 100644 --- a/weightslab/components/checkpoint_manager.py +++ b/weightslab/components/checkpoint_manager.py @@ -113,6 +113,11 @@ def __init__(self, root_log_dir: str = 'root_experiment', load_model: bool = Tru self.data_checkpoint_dir.mkdir(exist_ok=True) self.loggers_dir.mkdir(exist_ok=True) + # Persist the signal-history logger to an on-disk DuckDB file under + # loggers/. Handles the ordering where the logger was created before + # this checkpoint manager (the reverse is handled in LoggerQueue.__init__). + self._bind_logger_to_disk() + # Manifest file for tracking hash chronology self.manifest_file = self.checkpoints_dir / "manifest.yaml" @@ -194,6 +199,31 @@ def _get_data_state_snapshot(self, dfm): except Exception: return None + def get_logger_db_path(self) -> Path: + """Path of the on-disk DuckDB file backing the signal-history logger.""" + return self.loggers_dir / "loggers.duckdb" + + def _bind_logger_to_disk(self) -> None: + """Point the registered logger at an on-disk DuckDB file, if possible. + + Safe no-op when no logger is registered yet or it predates on-disk + persistence support.""" + try: + lg = ledgers.get_logger() + if lg is not None and hasattr(lg, "set_db_path"): + lg.set_db_path(str(self.get_logger_db_path())) + except Exception as e: + logger.warning(f"Could not bind logger to on-disk DuckDB: {e}") + + def flush_logger_to_disk(self) -> None: + """Force the logger to checkpoint its DuckDB history to disk (if any).""" + try: + lg = ledgers.get_logger() + if lg is not None and hasattr(lg, "flush_to_disk"): + lg.flush_to_disk() + except Exception as e: + logger.warning(f"Could not flush logger to disk: {e}") + def _get_logger_snapshot_dir(self) -> Path: return self.loggers_dir @@ -871,9 +901,9 @@ def _save_changes( except Exception as e: logger.warning(f"Failed to save weights with component changes: {e}") - # Always save logger snapshot alongside other components (same hash) - if dump_model_architecture or dump_model_state or dump_optimizer_state or dump_config_state: - self.save_logger_snapshot() + # Durably checkpoint the on-disk signal-history DuckDB alongside the + # rest of the state (cheap CHECKPOINT; replaces the old JSON snapshot). + self.flush_logger_to_disk() def save_model_checkpoint( self, @@ -977,12 +1007,6 @@ def save_model_checkpoint( # If model architecture doesn't exist in this hash directory, save a reference to where it is if self.config.get('checkpoint_manager', {}).get('dump_model_architecture', False): self._save_architecture_reference_if_needed() # TODO (GP): Disable for now because it adds complexity for big models, and we want to ensure architecture is always saved with weights for simplicity - - # Persist logger queues alongside weight checkpoints - try: - self.save_logger_snapshot() - except Exception as e: - logger.debug(f"Could not save logger snapshot with checkpoint: {e}") return checkpoint_file except Exception as e: logger.error(f"Failed to save model checkpoint: {e}") @@ -1214,99 +1238,6 @@ def save_data_snapshot(self, force_new_state: bool = False) -> Optional[Path]: logger.error(f"Failed to save data snapshot: {e}") return None - def save_logger_snapshot(self, exp_hash: Optional[str] = None) -> Optional[Path]: - """Persist logger queues for the given experiment hash. - - Uses the same hash as model/hp/data; does not affect hashing. - """ - exp = exp_hash or self.current_exp_hash - if exp is None: - return None - - try: - logger_names = ledgers.list_loggers() - if not logger_names: - return None - - snapshot = {"exp_hash": exp, "timestamp": datetime.now().isoformat(), "loggers": {}} - for lname in logger_names: - lg = ledgers.get_logger(lname) - if lg == None: - continue - - # Get snapshot data from logger, using dedicated method if available for better encapsulation (e.g. for LoggerQueue) - payload = lg.save_snapshot() - - # Get final snapshot for this logger - snapshot["loggers"][lname] = payload - - if not snapshot["loggers"]: - return None - - snapshot_dir = self._get_logger_snapshot_dir() - snapshot_dir.mkdir(parents=True, exist_ok=True) - - # Merge existing payload to avoid dropping loggers not currently registered - try: - existing = self._load_logger_snapshot_payload() - existing_loggers = existing.get("loggers", {}) if isinstance(existing, dict) else {} - if existing_loggers: - recursive_update(existing_loggers, snapshot.get("loggers", {})) - snapshot["loggers"] = existing_loggers - except Exception as e: - logger.warning(f"Failed to merge existing logger snapshot: {e}") - - records: List[bytes] = [] - for lname, payload in snapshot["loggers"].items(): - line = json.dumps({"logger_name": lname, "payload": payload}, default=str) + "\n" - records.append(line.encode("utf-8")) - - chunks: List[bytes] = [] - current_chunk = bytearray() - for record in records: - if current_chunk and (len(current_chunk) + len(record) > self.LOGGER_SNAPSHOT_MAX_FILE_SIZE_BYTES): - chunks.append(bytes(current_chunk)) - current_chunk = bytearray() - current_chunk.extend(record) - if current_chunk: - chunks.append(bytes(current_chunk)) - - old_chunks = self._list_logger_snapshot_chunks() - for old_chunk in old_chunks: - try: - old_chunk.unlink() - except Exception: - pass - - compressor = zstd.ZstdCompressor(level=3) - chunk_names: List[str] = [] - for idx, raw_chunk in enumerate(chunks, start=1): - chunk_path = self._get_logger_snapshot_chunk_path(idx) - tmp_chunk_path = chunk_path.with_name(chunk_path.name + ".tmp") - with open(tmp_chunk_path, "wb") as f: - f.write(compressor.compress(raw_chunk)) - os.replace(tmp_chunk_path, chunk_path) - chunk_names.append(chunk_path.name) - - manifest = { - "exp_hash": exp, - "timestamp": datetime.now().isoformat(), - "format": "ndjson+zstd", - "max_file_size_bytes": self.LOGGER_SNAPSHOT_MAX_FILE_SIZE_BYTES, - "chunks": chunk_names, - } - manifest_path = self._get_logger_snapshot_manifest_path() - tmp_manifest_path = manifest_path.with_name(manifest_path.name + ".tmp") - with open(tmp_manifest_path, "w", encoding="utf-8") as f: - json.dump(manifest, f, indent=2) - os.replace(tmp_manifest_path, manifest_path) - - logger.info(f"Saved logger snapshot: {manifest_path} ({len(chunk_names)} chunks)") - return manifest_path - except Exception as e: - logger.warning(f"Failed to save logger snapshot: {e}") - return None - def save_pending_changes(self, force: bool = False) -> bool: """Dump any pending changes to disk. diff --git a/weightslab/data/h5_array_store.py b/weightslab/data/h5_array_store.py index 440e6e59..02c3e7e3 100644 --- a/weightslab/data/h5_array_store.py +++ b/weightslab/data/h5_array_store.py @@ -470,39 +470,49 @@ def save_array( self._ensure_parent() + # Acquire the exclusive write side of the read-write lock so that no + # reader (load_array / load_arrays_batch) can have the file open in + # 'r' mode while we open it in 'a' mode. HDF5 forbids opening the same + # file read-write while it is already open read-only in the same + # process ("file is already open for read-only"), which otherwise + # corrupts the store. with self._local_lock: - with _InterProcessFileLock(self._lock_path, timeout=self._lock_timeout, poll_interval=self._poll_interval): - try: - with h5py.File(str(self._path), 'a') as f: - # Create group structure: /sample_id/key_name/ - sample_group_name = str(sample_id) - if sample_group_name not in f: - sample_group = f.create_group(sample_group_name) - else: - sample_group = f[sample_group_name] - - # Remove existing key if present - if key_name in sample_group: - del sample_group[key_name] - - # Create key group - key_group = sample_group.create_group(key_name) - - # Store array data (with optional compression) - if self._use_compression: - key_group.create_dataset('data', data=array, compression='gzip', compression_opts=4) - else: - key_group.create_dataset('data', data=array) - - # Store metadata - for k, v in metadata.items(): - key_group.attrs[k] = v - - return self._build_path_reference(sample_id, key_name) - - except Exception as exc: - logger.error(f"[H5ArrayStore] Failed to save array for sample_id={sample_id}, key={key_name}: {exc}") - return None + self._rw_lock.acquire_write() + try: + with _InterProcessFileLock(self._lock_path, timeout=self._lock_timeout, poll_interval=self._poll_interval): + try: + with h5py.File(str(self._path), 'a') as f: + # Create group structure: /sample_id/key_name/ + sample_group_name = str(sample_id) + if sample_group_name not in f: + sample_group = f.create_group(sample_group_name) + else: + sample_group = f[sample_group_name] + + # Remove existing key if present + if key_name in sample_group: + del sample_group[key_name] + + # Create key group + key_group = sample_group.create_group(key_name) + + # Store array data (with optional compression) + if self._use_compression: + key_group.create_dataset('data', data=array, compression='gzip', compression_opts=4) + else: + key_group.create_dataset('data', data=array) + + # Store metadata + for k, v in metadata.items(): + key_group.attrs[k] = v + + return self._build_path_reference(sample_id, key_name) + + except Exception as exc: + logger.error(f"[H5ArrayStore] Failed to save array for sample_id={sample_id}, key={key_name}: {exc}") + return None + finally: + self._rw_lock.release_write() def save_arrays_batch( self, @@ -571,51 +581,74 @@ def save_arrays_batch( return {} # --- Phase 2: merge temp into main file under lock --- + # Hold the exclusive write side of the read-write lock so that no + # reader (load_array / load_arrays_batch) has the file open in 'r' + # mode while we open it in 'a' mode. HDF5 refuses to open a file + # read-write while it is already open read-only in the same process + # ("file is already open for read-only"); that failure previously + # triggered a backup-restore that overwrote the live file out from + # under open reader handles and corrupted the store. with self._local_lock: - with _InterProcessFileLock(self._lock_path, timeout=self._lock_timeout, poll_interval=self._poll_interval): - backup_path = self._create_backup() - try: - with h5py.File(str(self._path), 'a') as f_main: - with h5py.File(str(tmp_path), 'r') as f_tmp: - for sample_group_name in f_tmp: - if sample_group_name not in f_main: - f_main.create_group(sample_group_name) - dest_group = f_main[sample_group_name] - src_group = f_tmp[sample_group_name] - # Replace only the keys present in this batch - # (e.g. "prediction") so sibling keys written - # by an earlier flush (e.g. "target") survive. - for key_name in src_group: - if key_name in dest_group: - del dest_group[key_name] - f_tmp.copy(f"{sample_group_name}/{key_name}", dest_group) - - path_refs: Dict[int, Dict[str, str]] = {} - for sample_group_name, key_data in prepared.items(): - sample_id = sample_group_name - path_refs[sample_id] = { - key_name: self._build_path_reference(sample_id, key_name) - for key_name in key_data - } - - # Merge succeeded — remove backup so recover() doesn't restore it - if backup_path: - backup_path.unlink(missing_ok=True) - - logger.debug( - f"[H5ArrayStore] Successfully upserted " - f"{sum(len(v) for v in path_refs.values())} arrays for " - f"{len(path_refs)} samples" - ) - return path_refs - - except Exception as exc: - logger.error(f"[H5ArrayStore] Failed to merge temp batch file: {exc}") - if backup_path: - self._restore_backup(backup_path) - return {} - finally: - tmp_path.unlink(missing_ok=True) + self._rw_lock.acquire_write() + try: + with _InterProcessFileLock(self._lock_path, timeout=self._lock_timeout, poll_interval=self._poll_interval): + backup_path = self._create_backup() + # Track whether we actually began mutating the main file. + # If we never opened it (e.g. it was busy), the file is + # intact and restoring from a copy would only risk harm. + main_opened = False + try: + with h5py.File(str(self._path), 'a') as f_main: + main_opened = True + with h5py.File(str(tmp_path), 'r') as f_tmp: + for sample_group_name in f_tmp: + if sample_group_name not in f_main: + f_main.create_group(sample_group_name) + dest_group = f_main[sample_group_name] + src_group = f_tmp[sample_group_name] + # Replace only the keys present in this batch + # (e.g. "prediction") so sibling keys written + # by an earlier flush (e.g. "target") survive. + for key_name in src_group: + if key_name in dest_group: + del dest_group[key_name] + f_tmp.copy(f"{sample_group_name}/{key_name}", dest_group) + + path_refs: Dict[int, Dict[str, str]] = {} + for sample_group_name, key_data in prepared.items(): + sample_id = sample_group_name + path_refs[sample_id] = { + key_name: self._build_path_reference(sample_id, key_name) + for key_name in key_data + } + + # Merge succeeded — remove backup so recover() doesn't restore it + if backup_path: + backup_path.unlink(missing_ok=True) + + logger.debug( + f"[H5ArrayStore] Successfully upserted " + f"{sum(len(v) for v in path_refs.values())} arrays for " + f"{len(path_refs)} samples" + ) + return path_refs + + except Exception as exc: + logger.error(f"[H5ArrayStore] Failed to merge temp batch file: {exc}") + if backup_path: + if main_opened: + # We started writing into the main file; it may + # be half-merged, so roll back to the backup. + self._restore_backup(backup_path) + else: + # Main file was never touched — drop the backup + # rather than overwrite a healthy file. + backup_path.unlink(missing_ok=True) + return {} + finally: + tmp_path.unlink(missing_ok=True) + finally: + self._rw_lock.release_write() def recover(self) -> None: """