Export model predictions from benchmark runs

src/relarena/cli.py

@@ -60,6 +60,12 @@ def _build_parser() -> argparse.ArgumentParser:
     )
     p.add_argument("--output", default=None, help="write a results CSV to this path")
     p.add_argument("--cache-dir", default=None, help="local preprocessing cache")
+    p.add_argument("--predictions-dir", help="write local prediction artifacts")
+    p.add_argument(
+        "--refit-all-configs",
+        action="store_true",
+        help="also refit nonselected configs for test artifacts",
+    )
     return p
 
 
@@ -104,6 +110,8 @@ def main(argv: list[str] | None = None) -> int:
                 n_trials=args.n_trials,
                 cache_dir=args.cache_dir,
                 evaluate_test=not args.no_test,
+                predictions_dir=args.predictions_dir,
+                refit_all_configs=args.refit_all_configs,
             )
         except Exception as exc:  # one bad dataset shouldn't abort the sweep
             print(f"    ERROR: {exc!r}", file=sys.stderr)

src/relarena/runner.py

@@ -12,6 +12,7 @@
 import math
 import time
 from dataclasses import dataclass
+from functools import partial
 from pathlib import Path
 from typing import Callable, Type
 
@@ -22,6 +23,7 @@
 from relarena.dataset import RelBenchDatasetTask
 from relarena.metrics import is_better
 from relarena.model import RelArenaModel
+from relarena.predictions import PredictionArtifactWriter
 from relarena.registry import registry
 from relarena.results import SystemResult, TrialResult
 from relarena.search_space import SearchSpaceProvider
@@ -175,12 +177,18 @@ def run_model_experiment(
     cache_dir: str | Path | None = None,
     evaluate_test: bool = True,
     require_all_trials: bool = True,
+    predictions_dir: str | Path | None = None,
+    refit_all_configs: bool = False,
 ) -> ExperimentSummary:
     """Tune one model on a RelBench entity task and summarize its trials.
 
     `search_space` defaults to the one registered for `model_cls` (via
     `@register_model`); pass it explicitly to override.
 
+    With predictions_dir, save validation and final-fit predictions locally.
+    refit_all_configs also evaluates other successful configs, recording their
+    extra fit/predict times only in the artifacts.
+
     Protocol (nested temporal validation; see docs/temporal-validation.md):
       1. **Tune** — fit each config on `train`, score on `val`, using the DB
          censored at `val_timestamp` so validation features are frozen at the val
@@ -198,6 +206,12 @@ def run_model_experiment(
     only). Returns an `ExperimentSummary` with the default and best-tuned
     trials (the latter carrying the refit test score) plus the full trial list.
     """
+    if refit_all_configs and (predictions_dir is None or not evaluate_test):
+        raise ValueError(
+            "All-config refits require predictions_dir and test evaluation."
+        )
+    if predictions_dir is not None and not cache_predictions:
+        raise ValueError("Prediction artifacts require cache_predictions=True.")
     source = RelBenchDatasetTask(dataset_name, task_name, download=download)
     task = source.task
     cache = resolve_cache_config(cache_dir, on_miss="raise")
@@ -223,11 +237,12 @@ def run_model_experiment(
     # Phase 1+2: tune on the inner split (train→val, DB censored at val_timestamp so
     # validation features are frozen at the val cutoff — see the protocol note above
     # and docs/temporal-validation.md), then select the best config by val score.
+    inner = source.inner_split()
     trials = tune(
         model_cls,
         search_space,
         task,
-        source.inner_split(),
+        inner,
         n_trials=n_trials,
         seed=seed,
         time_limit_per_trial=time_limit_per_trial,
@@ -249,28 +264,47 @@ def run_model_experiment(
             f"{failed[0].error}"
         )
 
+    artifacts = None
+    if predictions_dir is not None:
+        artifacts = PredictionArtifactWriter(
+            Path(predictions_dir),
+            source,
+            model_cls.name,
+            seed,
+            model_cls.refit_on_full_data,
+        )
+        artifacts.save_validation(trials, inner)
+
     default = next((t for t in trials if t.config_tag == "default"), None)
     tuned = select_best(trials, metric) if any(t.ok for t in trials) else None
 
     # Phase 3: run the selected config(s) under the model's final-fit regime on the
     # outer split (DB censored at test_timestamp) to get the test score.
     if evaluate_test and tuned is not None:
         outer = source.outer_split()
-        to_refit = [tuned]
+        benchmark_trials = [tuned]
         if default is not None and default.ok and default.config_id != tuned.config_id:
-            to_refit.append(default)
-        for trial in to_refit:
+            benchmark_trials.append(default)
+        benchmark_configs = {trial.config_id for trial in benchmark_trials}
+        extra_trials = (
+            [t for t in trials if t.ok and t.config_id not in benchmark_configs]
+            if refit_all_configs
+            else []
+        )
+        refit_trial = partial(
+            refit_and_evaluate,
+            model_cls,
+            task=task,
+            split=outer,
+            seed=seed,
+            time_limit=time_limit_per_trial,
+            cache=cache,
+            run_identity=source.run_identity("outer"),
+        )
+
+        for trial in benchmark_trials:
             try:
-                refit = refit_and_evaluate(
-                    model_cls,
-                    trial.config,
-                    task,
-                    outer,
-                    seed=seed,
-                    time_limit=time_limit_per_trial,
-                    cache=cache,
-                    run_identity=source.run_identity("outer"),
-                )
+                refit = refit_trial(trial.config)
                 trial.test_score = refit["test_score"]
                 trial.test_metrics = refit["test_metrics"]
                 trial.test_pred = refit["test_pred"] if cache_predictions else None
@@ -284,6 +318,15 @@ def run_model_experiment(
                     task_name,
                     trial.config_id,
                 )
+                if artifacts is not None:
+                    raise
+                continue
+            if artifacts is not None:
+                artifacts.save_test(trial, refit, outer, additional=False)
+
+        for trial in extra_trials:
+            refit = refit_trial(trial.config)
+            artifacts.save_test(trial, refit, outer, additional=True)
 
     return ExperimentSummary(
         model_name=model_cls.name,
@@ -314,6 +357,8 @@ def run_experiment(
     cache_dir: str | Path | None = None,
     evaluate_test: bool = True,
     require_all_trials: bool = True,
+    predictions_dir: str | Path | None = None,
+    refit_all_configs: bool = False,
 ) -> ExperimentSummary | SystemExperimentSummary:
     """Dispatch one experiment to the model or system runner.
 
@@ -323,6 +368,8 @@ def run_experiment(
     `time_limit_per_trial` becomes its single soft time limit.
     """
     if isinstance(method_cls, type) and issubclass(method_cls, RelArenaSystem):
+        if predictions_dir is not None or refit_all_configs:
+            raise TypeError("Prediction artifacts currently require a model.")
         if search_space is not None:
             raise TypeError("A RelArenaSystem does not have a harness search space.")
         return run_system_experiment(
@@ -350,4 +397,6 @@ def run_experiment(
         cache_dir=cache_dir,
         evaluate_test=evaluate_test,
         require_all_trials=require_all_trials,
+        predictions_dir=predictions_dir,
+        refit_all_configs=refit_all_configs,
     )

tests/test_predictions.py

@@ -0,0 +1,190 @@
+"""Prediction artifact round trips through the real tuner and evaluator."""
+
+from __future__ import annotations
+
+from copy import deepcopy
+from pathlib import Path
+from types import SimpleNamespace
+from typing import Any
+
+import numpy as np
+import pandas as pd
+import pytest
+from relbench.base import Database, EntityTask, Table, TaskType
+from relbench.metrics import mae, roc_auc
+
+from relarena import (
+    InnerSplit,
+    OuterSplit,
+    RelArenaModel,
+    RunIdentity,
+    load_prediction_labels,
+    load_predictions,
+    prediction_context,
+    runner,
+)
+from relarena.models.dummy import DummyBaseline
+from relarena.search_space import SearchSpace
+
+
+def _table(targets: list[float], date: str) -> Table:
+    return Table(
+        df=pd.DataFrame(
+            {"entity": range(len(targets)), "time": pd.Timestamp(date), "y": targets}
+        ),
+        fkey_col_to_pkey_table={"entity": "entities"},
+        pkey_col=None,
+        time_col="time",
+    )
+
+
+@pytest.mark.parametrize("binary", [False, True])
+@pytest.mark.parametrize("all_configs", [False, True])
+@pytest.mark.parametrize("refit_full", [False, True])
+def test_artifact_roundtrip(
+    tmp_path: Path,
+    monkeypatch: pytest.MonkeyPatch,
+    all_configs: bool,
+    refit_full: bool,
+    binary: bool,
+) -> None:
+    train = _table([0.0, 1.0] if binary else [0.0, 2.0], "2020-01-01")
+    val = _table([0.0, 1.0] if binary else [1.0, 1.0], "2020-02-01")
+    test = _table([0.0, 1.0] if binary else [0.5, 1.5], "2020-03-01")
+    masked = Table(
+        df=test.df.drop(columns="y"),
+        fkey_col_to_pkey_table={"entity": "entities"},
+        pkey_col=None,
+        time_col="time",
+    )
+    task = SimpleNamespace(
+        task_type=TaskType.BINARY_CLASSIFICATION if binary else TaskType.REGRESSION,
+        entity_col="entity",
+        time_col="time",
+        target_col="y",
+        metrics=[roc_auc] if binary else [mae],
+    )
+    task.get_table = lambda *a, **k: test
+    task.evaluate = lambda *a, **k: EntityTask.evaluate(task, *a, **k)
+    db = Database({})
+    inner = InnerSplit(db, pd.Timestamp("2020-02-01"), train, val, val)
+    outer = OuterSplit(db, pd.Timestamp("2020-03-01"), train, masked, val)
+    identity = RunIdentity("small", "db", "target", "task")
+    source = SimpleNamespace(
+        task=task,
+        dataset_name="small",
+        task_name="target",
+        metric=roc_auc if binary else mae,
+        inner_split=lambda: inner,
+        outer_split=lambda: outer,
+        run_identity=identity.for_phase,
+    )
+    monkeypatch.setattr(runner, "RelBenchDatasetTask", lambda *a, **k: source)
+    fits: list[tuple[int, bool]] = []
+
+    class Constant(RelArenaModel):
+        name = "artifact-constant"
+        refit_on_full_data = refit_full
+
+        def fit(
+            self,
+            task: Any,
+            db: Any,
+            train_table: Table,
+            val_table: Table | None,
+            **kwargs: Any,
+        ) -> None:
+            assert (train_table.df["time"] < outer.cutoff).all()
+            fits.append((len(train_table.df), val_table is None))
+
+        def predict(self, task: Any, db: Any, table: Table) -> np.ndarray:
+            if table.df["time"].iloc[0] == outer.cutoff:
+                assert "y" not in table.df
+            if binary:
+                return np.array([[0.5, 0.5], [0.1, 0.9], [0.9, 0.1]])[
+                    self.config["constant"]
+                ]
+            return np.full(len(table.df), self.config["constant"], dtype=float)
+
+    summary = runner.run_experiment(
+        Constant,
+        "small",
+        "target",
+        search_space=SearchSpace(
+            fixed_grid=[{"constant": 0}, {"constant": 1}, {"constant": 2}],
+            default_overrides={"constant": 0},
+        ),
+        n_trials=3,
+        predictions_dir=tmp_path,
+        refit_all_configs=all_configs,
+    )
+    assert len(fits) == (6 if all_configs else 5)
+    assert fits[:3] == [(2, False)] * 3
+    assert fits[3:] == ([(4, True)] if refit_full else [(2, False)]) * (len(fits) - 3)
+    for trial in summary.trials:
+        root = tmp_path / Constant.name / "small" / "target" / "0" / trial.config_id
+        for split, split_object, target in [("val", inner, val), ("test", outer, test)]:
+            path = root / f"{split}.npz"
+            additional = trial.config["constant"] == 2
+            if split == "test" and additional and not all_configs:
+                assert not path.exists()
+                continue
+            context = prediction_context(
+                task, split_object, identity.for_phase(split_object.name)
+            )
+            predictions, metadata = load_predictions(path, expected_context=context)
+            labels = load_prediction_labels(path)
+            np.testing.assert_array_equal(labels, target.df["y"].to_numpy())
+            assert {
+                m.__name__: m(labels, predictions) for m in task.metrics
+            } == pytest.approx(metadata["metrics"])
+            assert metadata["additional_refit"] == (additional and split == "test")
+            assert metadata["fit_time"] >= 0
+            assert metadata["predict_time"] >= 0
+            for field, value in [
+                ("rows", list(reversed(context["rows"]))),
+                ("split", "wrong"),
+                ("classes", [1, 0]),
+            ]:
+                mismatched = deepcopy(context)
+                mismatched[field] = value
+                with pytest.raises(ValueError, match="differ"):
+                    load_predictions(path, expected_context=mismatched)
+        if trial.config["constant"] == 2:
+            assert trial.test_pred is None
+            assert trial.test_score is None
+            assert trial.fit_time_refit is None
+
+    builtin = runner.run_experiment(
+        DummyBaseline,
+        "small",
+        "target",
+        predictions_dir=tmp_path,
+    )
+    for split, split_object, target in [("val", inner, val), ("test", outer, test)]:
+        path = (
+            tmp_path
+            / DummyBaseline.name
+            / "small"
+            / "target"
+            / "0"
+            / builtin.default.config_id
+            / f"{split}.npz"
+        )
+        predictions, metadata = load_predictions(
+            path,
+            expected_context=prediction_context(
+                task, split_object, identity.for_phase(split_object.name)
+            ),
+        )
+        assert task.evaluate(predictions, target) == pytest.approx(metadata["metrics"])
+
+    assert len(list((tmp_path / "labels").glob("*.npz"))) == 2
+    label_file = path.parent / metadata["labels_path"]
+    with np.load(label_file, allow_pickle=False) as artifact:
+        labels = artifact["labels"].copy()
+        context_array = artifact["context"]
+    labels[0] += 1
+    np.savez_compressed(label_file, labels=labels, context=context_array)
+    with pytest.raises(ValueError, match="content or row identities"):
+        load_prediction_labels(path)