Metadata-Version: 2.3
Name: research-tracker
Version: 0.1.4rc2
Summary: A local-first store for tracking research runs, evaluations, metrics, and artifacts
Keywords: experiment-tracking,machine-learning,research
Author: c-salomonsen
Author-email: c-salomonsen <chris10an.salomonsen@gmail.com>
Classifier: Development Status :: 3 - Alpha
Classifier: Intended Audience :: Science/Research
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.12
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Requires-Dist: filelock>=3.32.5
Requires-Dist: lightning[pytorch-extra]>=2.6.5
Requires-Dist: omegaconf>=2.3.1
Requires-Dist: pandas>=3.0.5
Requires-Dist: pyarrow>=25.0.1
Requires-Dist: torch>=2.8.0a0,<3
Requires-Dist: wandb>=0.29.0
Requires-Python: >=3.12
Description-Content-Type: text/markdown

# research-tracker

A small, local-first research metadata store for connecting training runs to
evaluations, experimental conditions, sample-level metrics, and generated
artifacts. Records are kept in Parquet tables and artifact paths are relative to
the store, so a copied store can be merged back into the main store.

## Installation

Install from your configured package index:

```bash
pip install research-tracker
```

Or, from a neighboring project during development:

```bash
uv add --editable ../research-tracker
```

To install this project and its test dependencies:

```bash
uv sync --group dev
```

## Python API

```python
from pathlib import Path

from research_tracker import ExperimentStore
from research_tracker.trackers import Status

root = Path("artifacts")
store = ExperimentStore(root)

run = store.create_run(
    model="inverse-operator",
    model_class="my_project.models:InverseOperator",
    dataset="dynamic-pet",
    dataset_version="irr-v1",
    config=root / "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",
        target="parametric-map",
        metric="mse",
        value=0.012,
    )
    evaluation.log_metric(metric="mean_mse", value=0.012)
    evaluation.log_artifact(
        kind="prediction",
        path=root / "predictions" / "mouse-01.npz",
        sample_id="mouse-01",
    )

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

`metric` and `value` are required; `sample_id` and `target` are optional
keyword arguments and are stored as null when omitted.
The evaluation context flushes buffered metrics and conditions on success. If
the block or its flush raises an exception, the evaluation and its associated
metric, condition, and artifact records are removed, including records already
written during the block. Artifact files themselves are left in place because
they may be shared. This rollback does not cover abrupt process termination.
Configuration, checkpoint, and artifact paths must remain inside the store root.

## 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.

```yaml
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`.
- `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.
- 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.

## Sync and CLI

Pull a copied remote store into the main local store:

```bash
research-tracker --root ./artifacts sync /path/to/copied/store
research-tracker --root ./artifacts run show
research-tracker --root ./artifacts metric query "metric == 'mse'"
```

The source path must be the store directory containing `schema.json` and its
Parquet tables. Sync copies files referenced by runs (`config` and `checkpoint`)
and artifacts (`path`), unless `--skip-copy-artifacts` is set. The CLI prints
each copied file and a per-category summary; programmatic callers can pass a
`progress` callback to `ExperimentStore.sync` for the same messages.

Runs are mutable and the newest `updated_at` wins. Evaluations, conditions,
metrics, and artifacts are immutable; conflicting records raise an error.

## Tests

```bash
uv run pytest
```
