ExperimentStore

research_tracker.ExperimentStore is the public entry point: it owns the store directory, the file lock, the metric/condition buffers, and every read and write. See runs and evaluations for the record model and sync for the merge rules.

Construction and flushing

from pathlib import Path

from research_tracker import ExperimentStore

store = ExperimentStore(Path("artifacts"))

research_tracker.ExperimentStore.__init__(root, metrics_write_every=50, conditions_write_every=50)

Source code in src/research_tracker/store.py

def __init__(
    self,
    root: Path,
    metrics_write_every: int = 50,
    conditions_write_every: int = 50,
):
    self.root = Path(root)
    self.lock_path = self.root / ".sync.lock"
    self._lock = FileLock(str(self.lock_path))
    self.db_path = self.root / "store.sqlite"
    schema_path = self.root / "schema.json"

    # Read and validate an existing marker before acquiring the lock so a
    # rejected store gains neither a `.sync.lock` file nor anything else.
    # A partial layout can be a store another process is creating right
    # now: the creator already holds the lock, so its `.sync.lock` file
    # exists. Only reject a partial layout here while no lock file exists;
    # the checks under the lock below are authoritative either way.
    has_schema = schema_path.exists()
    has_database = self.db_path.exists()

    if has_schema:
        self.schema_version = self._read_schema(schema_path)
        self._validate_schema_version()

    if has_schema != has_database and not self.lock_path.exists():
        self._raise_partial_store(has_schema)

    self.root.mkdir(parents=True, exist_ok=True)

    with self._lock:
        # A concurrent writer may have created or replaced the store
        # between the checks above and lock acquisition.
        has_schema = schema_path.exists()
        has_database = self.db_path.exists()

        if has_schema:
            self.schema_version = self._read_schema(schema_path)
            self._validate_schema_version()

        if has_schema != has_database:
            self._raise_partial_store(has_schema)

        if not has_schema:
            self.schema_version = SCHEMA_VERSION
            self._write_schema(schema_path)

        # Only a genuinely new store is created here. An existing store is
        # validated against the canonical schema and never repaired: a
        # dropped or altered table must surface as a schema mismatch
        # instead of being silently recreated.
        with self._transaction() as cursor:
            if has_database:
                self._validate_db_schema(cursor)
            else:
                self._create_tables(cursor)

    self._batchers = {
        "metric": BatchBuffer(capacity=metrics_write_every),
        "condition": BatchBuffer(capacity=conditions_write_every),
    }

    atexit.register(self.flush)

research_tracker.ExperimentStore.evaluation(run_id, name)

Source code in src/research_tracker/store.py

@contextmanager
def evaluation(
    self,
    run_id: str,
    name: str,
) -> Iterator[EvaluationSession]:
    evaluation = self.create_evaluation(run_id=run_id, name=name)
    session = EvaluationSession(self, evaluation)

    try:
        yield session
        self.flush()
    except BaseException as error:
        try:
            self._rollback_evaluation(evaluation.evaluation_id)
        except Exception as rollback_error:
            error.add_note(
                f"Could not roll back evaluation {evaluation.evaluation_id}: "
                f"{rollback_error}"
            )
            raise error from rollback_error
        raise

research_tracker.ExperimentStore.flush()

Source code in src/research_tracker/store.py

def flush(self) -> None:
    with self._lock:
        for kind, batcher in self._batchers.items():
            self._flush_batch(kind, batcher)

Record writers

Runs, evaluations, and artifacts are written immediately. Metrics and conditions are buffered until their threshold, the end of an evaluation block, or flush().

research_tracker.ExperimentStore.create_run(model, model_class, dataset, config, status, *, seed=None, checkpoint=None, tracking_backend=None, tracking_id=None, tracking_project=None, loader=None, dataset_version=None, datamodule_class=None, parent_run_id=None)

Source code in src/research_tracker/store.py

def create_run(
    self,
    model: str,
    model_class: str,
    dataset: str,
    config: Path,
    status: Status,
    *,
    seed: int | None = None,
    checkpoint: Path | None = None,
    tracking_backend: str | None = None,
    tracking_id: str | None = None,
    tracking_project: str | None = None,
    loader: str | None = None,
    dataset_version: str | None = None,
    datamodule_class: str | None = None,
    parent_run_id: str | None = None,
) -> Run:
    run = Run(
        model=model,
        model_class=model_class,
        dataset=dataset,
        config=ensure_path(config, self.root),
        status=status,
        seed=seed,
        checkpoint=checkpoint and ensure_path(checkpoint, self.root),
        git_commit=curr_git_rev(),
        git_dirty=is_repo_dirty(),
        loader=loader,
        dataset_version=dataset_version,
        parent_run_id=parent_run_id,
        tracking_backend=tracking_backend,
        tracking_id=tracking_id,
        datamodule_class=datamodule_class,
        tracking_project=tracking_project,
    )
    self._append_records("run", [run])
    return run

research_tracker.ExperimentStore.create_evaluation(run_id, name)

Source code in src/research_tracker/store.py

def create_evaluation(
    self,
    run_id: str,
    name: str,
) -> Evaluation:
    evaluation = Evaluation(
        run_id=run_id,
        name=name,
    )
    self._append_records("evaluation", [evaluation])
    return evaluation

research_tracker.ExperimentStore.add_condition(evaluation_id, name, value, *, unit=None)

Source code in src/research_tracker/store.py

def add_condition(
    self,
    evaluation_id: str,
    name: str,
    value: str | float | bool,
    *,
    unit: str | None = None,
) -> None:
    condition = Condition(
        evaluation_id=evaluation_id,
        name=name,
        value=value,
        unit=unit,
    )

    self._buffer_record("condition", condition)

research_tracker.ExperimentStore.log_metric(evaluation_id, metric, value, *, sample_id=None, target=None, comparison=None, region=None, unit=None)

Source code in src/research_tracker/store.py

def log_metric(
    self,
    evaluation_id: str,
    metric: str,
    value: float,
    *,
    sample_id: str | None = None,
    target: str | None = None,
    comparison: str | None = None,
    region: str | None = None,
    unit: str | None = None,
) -> None:
    logged_metric = Metric(
        evaluation_id=evaluation_id,
        sample_id=sample_id,
        target=target,
        metric=metric,
        value=value,
        comparison=comparison,
        region=region,
        unit=unit,
    )

    self._buffer_record("metric", logged_metric)

research_tracker.ExperimentStore.log_artifact(evaluation_id, kind, path, *, sample_id=None)

Source code in src/research_tracker/store.py

def log_artifact(
    self,
    evaluation_id: str,
    kind: str,
    path: Path,
    *,
    sample_id: str | None = None,
) -> None:
    artifact = Artifact(
        evaluation_id=evaluation_id,
        kind=kind,
        path=ensure_path(path, self.root),
        sample_id=sample_id,
    )

    self._append_records("artifact", [artifact])

Run updates

research_tracker.ExperimentStore.set_status(status, run_id)

Source code in src/research_tracker/store.py

def set_status(self, status: Status, run_id: str) -> None:
    self._change_run_field(status, "status", run_id)

research_tracker.ExperimentStore.add_checkpoint(checkpoint_path, run_id)

Source code in src/research_tracker/store.py

def add_checkpoint(self, checkpoint_path: Path, run_id: str) -> None:
    checkpoint_path = ensure_path(checkpoint_path, self.root)
    self._change_run_field(str(checkpoint_path), "checkpoint", run_id)

research_tracker.ExperimentStore.import_wandb_run(wandb_run_path, model, dataset, config_path, checkpoint_path=None)

Source code in src/research_tracker/store.py

def import_wandb_run(
    self,
    wandb_run_path: str,
    model: str,
    dataset: str,
    config_path: Path,
    checkpoint_path: Path | None = None,
) -> None:
    wandb_run = wandb.Api().from_path(wandb_run_path)

    checkpoint = checkpoint_path or self._download_wandb_checkpoint(wandb_run)

    metadata = get_tracker_info_from_ckpt(checkpoint)

    loaded_config = OmegaConf.load(config_path)

    run = Run(
        model=model,
        model_class=metadata.get("model_class")
        or convert_class_path(loaded_config.model.class_path),
        dataset=dataset,
        datamodule_class=metadata.get("datamodule_class")
        or convert_class_path(loaded_config.data.class_path),
        seed=loaded_config.seed_everything,
        checkpoint=checkpoint and ensure_path(checkpoint, self.root),
        config=ensure_path(config_path, self.root),
        status=Status.from_wandb(wandb_run),
        run_id=metadata.get("run_id", make_id("run")),
        tracking_backend="wandb",
        tracking_id=wandb_run.id,
        tracking_project=str(wandb_run.project),
        created_at=datetime.fromisoformat(wandb_run.created_at),
    )
    self._append_records("run", [run])

research_tracker.ExperimentStore.load_checkpoint(run_id)

Source code in src/research_tracker/store.py

def load_checkpoint(self, run_id: str) -> torch.nn.Module:
    run = self.get_run(run_id)

    # If custom loader needed
    if run.loader is not None:
        loader_fn = import_object(run.loader)
        return loader_fn(run)

    # Or be sensible and use lightning
    model_cls = import_object(run.model_class)

    return model_cls.load_from_checkpoint(run.checkpoint)

Readers

research_tracker.ExperimentStore.has_run(run_id)

Source code in src/research_tracker/store.py

def has_run(self, run_id: str) -> bool:
    connection = sqlite3.connect(self.db_path, isolation_level=None)

    try:
        connection.execute("PRAGMA busy_timeout = 30000")
        cursor = connection.execute(
            "SELECT 1 FROM runs WHERE run_id = ? LIMIT 1",
            (run_id,),
        )
        return cursor.fetchone() is not None
    finally:
        connection.close()

research_tracker.ExperimentStore.get_run(run_id)

Source code in src/research_tracker/store.py

def get_run(self, run_id: str) -> Run:
    return self._get_records("run", Run, "run_id", run_id)[0]

research_tracker.ExperimentStore.get_evaluations(run_id)

Source code in src/research_tracker/store.py

def get_evaluations(self, run_id: str) -> list[Evaluation]:
    return self._get_records(
        "evaluation",
        Evaluation,
        "run_id",
        run_id,
    )

research_tracker.ExperimentStore.get_metrics(evaluation_id)

Source code in src/research_tracker/store.py

def get_metrics(self, evaluation_id: str) -> list[Metric]:
    return self._get_records(
        "metric",
        Metric,
        "evaluation_id",
        evaluation_id,
    )

research_tracker.ExperimentStore.get_conditions(evaluation_id)

Source code in src/research_tracker/store.py

def get_conditions(self, evaluation_id: str) -> list[Condition]:
    return self._get_records(
        "condition",
        Condition,
        "evaluation_id",
        evaluation_id,
    )

research_tracker.ExperimentStore.get_artifacts(evaluation_id)

Source code in src/research_tracker/store.py

def get_artifacts(self, evaluation_id: str) -> list[Artifact]:
    return self._get_records(
        "artifact",
        Artifact,
        "evaluation_id",
        evaluation_id,
    )

research_tracker.ExperimentStore.load_table(kind)

Source code in src/research_tracker/store.py

def load_table(self, kind: str) -> pd.DataFrame:
    return self._read_table(kind)

research_tracker.ExperimentStore.load_runs()

Source code in src/research_tracker/store.py

def load_runs(self) -> pd.DataFrame:
    return self.load_table("run")

research_tracker.ExperimentStore.load_evaluations()

Source code in src/research_tracker/store.py

def load_evaluations(self) -> pd.DataFrame:
    return self.load_table("evaluation")

research_tracker.ExperimentStore.load_conditions()

Source code in src/research_tracker/store.py

def load_conditions(self) -> pd.DataFrame:
    return self.load_table("condition")

research_tracker.ExperimentStore.load_metrics()

Source code in src/research_tracker/store.py

def load_metrics(self) -> pd.DataFrame:
    return self.load_table("metric")

research_tracker.ExperimentStore.load_artifacts()

Source code in src/research_tracker/store.py

def load_artifacts(self) -> pd.DataFrame:
    return self.load_table("artifact")

Sync

research_tracker.ExperimentStore.sync(source_root, *, copy_artifacts=True, progress=None)

Source code in src/research_tracker/store.py

def sync(
    self,
    source_root: Path,
    *,
    copy_artifacts: bool = True,
    progress: Callable[[str], None] | None = None,
) -> None:
    source_root = Path(source_root)
    if not source_root.exists():
        raise FileNotFoundError(source_root)
    if not (source_root / "schema.json").is_file():
        raise FileNotFoundError(source_root / "schema.json")

    source = ExperimentStore(source_root)

    self._check_schema_compatible(source)
    if progress is not None:
        progress(f"Syncing {source.root.resolve()} -> {self.root.resolve()}")
        if not any(source._count_rows(kind) for kind in self._TABLE_KEYS):
            progress(
                "No tables found in the source; check that the source "
                "path is the store directory."
            )

    # Prevent two sync processes from modifying the local store
    # simultaneously
    with self._lock:
        self.flush()

        # Merge first: validation may raise, in which case nothing is
        # committed and no files are copied.
        merged_tables = {
            kind: self._merge_table(source, kind) for kind in self._TABLE_KEYS
        }

        # Reject a corrupt merge before any file is copied; otherwise a
        # failed write would leave stray copies behind.
        for kind, merged in merged_tables.items():
            self._preflight_table(kind, merged)

        if copy_artifacts:
            self._copy_referenced_files(source, progress=progress)
        elif progress is not None:
            progress("File copying skipped (--skip-copy-artifacts).")

        with self._transaction() as cursor:
            for kind, merged in merged_tables.items():
                self._replace_table(cursor, kind, merged)

        if progress is not None:
            progress(
                "Sync complete: "
                + ", ".join(
                    f"{kind}s={len(table)}" for kind, table in merged_tables.items()
                )
            )

Removal

research_tracker.ExperimentStore.remove_evaluation_cascade(evaluation_ids, *, dry_run=False, remove_files=True)

Delete an evaluation and its conditions/metrics/artifacts in ONE transaction, returning the number of removed rows per table.

The whole operation -- including the survivor read and the artifact unlink -- is serialized under the store lock, so a concurrent log_artifact can neither add a reference after the survivor read nor re-insert rows the transaction deleted. A candidate file is unlinked only when no surviving artifact row references the same file, compared by resolved path so that result and sub/../result collapse.

Source code in src/research_tracker/store.py

def remove_evaluation_cascade(
    self,
    evaluation_ids: Iterable[str],
    *,
    dry_run: bool = False,
    remove_files: bool = True,
) -> dict[str, int]:
    """
    Delete an evaluation and its conditions/metrics/artifacts in ONE
    transaction, returning the number of removed rows per table.

    The whole operation -- including the survivor read and the artifact
    unlink -- is serialized under the store lock, so a concurrent
    `log_artifact` can neither add a reference after the survivor read nor
    re-insert rows the transaction deleted. A candidate file is unlinked
    only when no surviving artifact row references the same file, compared
    by resolved path so that `result` and `sub/../result` collapse.
    """
    ids = list(dict.fromkeys(evaluation_ids))
    counts = {"evaluation": 0, "metric": 0, "condition": 0, "artifact": 0}

    if not ids:
        return counts

    placeholders = ", ".join("?" for _ in ids)

    with self._lock:
        with self._transaction() as cursor:
            # Read the files this cascade may orphan before deleting the
            # rows that name them.
            candidate_paths = [
                row[0]
                for row in cursor.execute(
                    f"SELECT path FROM artifacts "
                    f"WHERE evaluation_id IN ({placeholders})",
                    ids,
                )
                if row[0] is not None
            ]

            # Children first: a partially applied sequence can never leave
            # orphaned rows behind, because the whole sequence commits or
            # rolls back as one transaction.
            for kind in ("metric", "condition", "artifact", "evaluation"):
                table = _table_name(kind)
                cursor.execute(
                    f"SELECT COUNT(*) FROM {table} "
                    f"WHERE evaluation_id IN ({placeholders})",
                    ids,
                )
                counts[kind] = cursor.fetchone()[0]

                if not dry_run:
                    cursor.execute(
                        f"DELETE FROM {table} "
                        f"WHERE evaluation_id IN ({placeholders})",
                        ids,
                    )

            if dry_run:
                # A dry run issues no writes; end the read transaction
                # explicitly so it can never commit anything.
                cursor.execute("ROLLBACK")

        # The transaction has committed while the lock is still held, so
        # every mutation in the store is serialized against the survivor
        # read and the unlink below.
        if not dry_run:
            # Drained while the lock is still held so a concurrent flush
            # cannot re-insert the rows the transaction just deleted.
            for batcher in self._batchers.values():
                for evaluation_id in ids:
                    batcher.discard_evaluation(evaluation_id)

        if not dry_run and remove_files:
            self._remove_orphaned_artifacts(candidate_paths)

    return counts

research_tracker.ExperimentStore.remove_from_table(kind, id, dry_run=False)

Source code in src/research_tracker/store.py

def remove_from_table(self, kind: str, id: str, dry_run: bool = False) -> None:
    with self._lock:
        if dry_run:
            df = self.load_table(kind)
            mask = df[df[f"{kind}_id"] == id]
            print(f"[Dry run] Would remove {mask}")
            return

        with self._transaction() as cursor:
            cursor.execute(
                f"DELETE FROM {_table_name(kind)} WHERE {kind}_id = ?",
                (id,),
            )

research_tracker.ExperimentStore.remove_run(run_id, dry_run=False)

Source code in src/research_tracker/store.py

def remove_run(self, run_id: str, dry_run: bool = False) -> None:
    self.remove_from_table("run", run_id, dry_run)

research_tracker.ExperimentStore.remove_evaluation(evaluation_id, dry_run=False)

Source code in src/research_tracker/store.py

def remove_evaluation(self, evaluation_id: str, dry_run: bool = False) -> None:
    self.remove_from_table("evaluation", evaluation_id, dry_run)

research_tracker.ExperimentStore.remove_condition(condition_id, dry_run=False)

Source code in src/research_tracker/store.py

def remove_condition(self, condition_id: str, dry_run: bool = False) -> None:
    self.remove_from_table("condition", condition_id, dry_run)

research_tracker.ExperimentStore.remove_metric(metric_id, dry_run=False)

Source code in src/research_tracker/store.py

def remove_metric(self, metric_id: str, dry_run: bool = False) -> None:
    self.remove_from_table("metric", metric_id, dry_run)

research_tracker.ExperimentStore.remove_artifact(artifact_id, dry_run=False)

Source code in src/research_tracker/store.py

def remove_artifact(self, artifact_id: str, dry_run: bool = False) -> None:
    self.remove_from_table("artifact", artifact_id, dry_run)

Table-level operations

update_table(kind, df) replaces a whole table in one transaction and is the path the migration uses to bulk-load a v1 store.

research_tracker.ExperimentStore.update_table(kind, df)

Source code in src/research_tracker/store.py

def update_table(self, kind: str, df: pd.DataFrame) -> None:
    self._validate_columns(kind, df)
    require_unique(df, self._TABLE_KEYS[kind])
    self._validate_stored_paths(kind, df)

    with self._lock, self._transaction() as cursor:
        self._replace_table(cursor, kind, df)