Getting started

Install

pip install research-tracker

During development, from a neighboring project:

uv add --editable ../research-tracker

To install this project together with its test dependencies:

uv sync --group dev

Upgrading a store written by 0.1.x additionally needs the migrate extra, which pulls in pyarrow for reading the old Parquet tables:

pip install "research-tracker[migrate]"

Creating a store

ExperimentStore takes the store directory. Constructing it creates the directory, schema.json, and store.sqlite if they do not exist, and validates the existing store otherwise:

from pathlib import Path

from research_tracker import ExperimentStore

store = ExperimentStore(Path("artifacts"))
print(store.db_path)        # artifacts/store.sqlite
print(store.schema_version) # 2

A store directory is identified by schema.json; opening a directory whose schema.json says a different version raises SchemaMismatchError at construction, before anything is created or modified. A directory with only one of schema.json / store.sqlite is rejected as incomplete.

Writes are serialized through a .sync.lock file lock, so separate jobs can safely append to the same store. Config, checkpoint, and artifact paths must stay inside the store root; an absolute or escaping path raises ValueError.

The evaluation context manager

store.evaluation(run_id, name) creates an evaluation and yields an EvaluationSession bound to it:

from research_tracker.trackers import Status

run = store.create_run(
    model="inverse-operator",
    model_class="my_project.models:InverseOperator",
    dataset="dynamic-pet",
    dataset_version="irr-v1",
    config=Path("artifacts/config.yaml"),
    status=Status.RUNNING,
)

with store.evaluation(run.run_id, "test") as evaluation:
    evaluation.add_condition("split", "out-of-distribution")
    evaluation.log_metric(sample_id="mouse-01", metric="mse", value=0.012)
    evaluation.log_artifact(
        kind="prediction",
        path=Path("artifacts/predictions/mouse-01.npz"),
        sample_id="mouse-01",
    )

store.set_status(Status.COMPLETED, run.run_id)

• add_condition(name, value, *, unit=None) accepts str, int, float, or bool; the value keeps its JSON type.
• log_metric(metric, value, *, sample_id=None, target=None, comparison=None, region=None, unit=None) — metric and value are required, the rest are stored as null when omitted.
• log_artifact(kind, path, *, sample_id=None) stores the path relative to the store root.

On success the context flushes buffered metrics and conditions. If the body or the flush raises, the evaluation and its metric, condition, and artifact rows are rolled back, including records already written during the block. The artifact files themselves are kept. See runs and evaluations for the details.

Lightning callback

ResearchTrackerCallback creates the run at fit start, records the selected Lightning logger, embeds the research run ID in checkpoints, records the best or last checkpoint, and marks the run completed or failed.

trainer:
  callbacks:
    - class_path: research_tracker.ResearchTrackerCallback
      init_args:
        store_root: ${oc.env:RESEARCH_TRACKER_ROOT,./artifacts}
        dataset: dynamic-pet
        dataset_version: irr-v1
        checkpoint_policy: best_or_last

Useful defaults and options:

• store_root defaults to trainer.default_root_dir.
• model defaults to model_name or the LightningModule class name.
• dataset defaults to dataset_name or the LightningDataModule class name.
• config_path is optional. The callback searches the logger experiment directory, logger log/save directories, then the trainer root for config.yaml (config_filename changes the name).
• logger_index selects which logger supplies external tracking metadata.
• checkpoint_policy is best_or_last, best, last, or none.
• resume_mode="continue" reuses a known run ID from a loaded checkpoint; "fork" creates a child run instead.
• parent_run_id explicitly links weight-only fine-tuning or other derived runs.
• loader selects a custom checkpoint loader.
• Fast development runs are ignored unless track_fast_dev_runs=true.

Only the global-zero Lightning process writes to the store. Separate training jobs can safely append to the same store through its file lock.

The callback writes through the same ExperimentStore API, so any buffered metric or condition still needs a flush() before the store is copied — see durability.