Skip to content

H5 and Loggers fixes #257

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
guillaume-byte merged 2 commits into from
Jul 21, 2026
Merged

H5 and Loggers fixes #257

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
628a89b
Fix bug with h5 file
Jul 20, 2026
caa29dd
Change from old logger dumping behavior (jsonz) to DB dump
Jul 20, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
167 changes: 133 additions & 34 deletions weightslab/backend/logger.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -31,6 +31,7 @@

import functools
import json
import logging
import os
import threading
import time
Expand All @@ -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 = [
Expand Down Expand Up @@ -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."""
Expand Down
135 changes: 33 additions & 102 deletions weightslab/components/checkpoint_manager.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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"

Expand Down Expand Up @@ -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

Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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}")
Expand Down Expand Up @@ -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.

Expand Down
Loading
Loading