diff --git a/src/relarena/cli.py b/src/relarena/cli.py index c8fe527..a4fef77 100644 --- a/src/relarena/cli.py +++ b/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) diff --git a/src/relarena/runner.py b/src/relarena/runner.py index 2a243bd..ce86dcc 100644 --- a/src/relarena/runner.py +++ b/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,6 +264,17 @@ 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 @@ -256,21 +282,29 @@ def run_model_experiment( # 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, ) diff --git a/tests/test_predictions.py b/tests/test_predictions.py new file mode 100644 index 0000000..70efdf6 --- /dev/null +++ b/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)