diff --git a/README.md b/README.md index 3a50bd7..f5d05ae 100644 --- a/README.md +++ b/README.md @@ -159,7 +159,7 @@ import pandas as pd import relarena.models # registers the built-in models from relarena.evaluation import compute_leaderboard -from relarena.registry import registry +from relarena.core.registry import registry from relarena.results import summary_to_dataframe from relarena.runner import run_experiment from relarena.tasks import list_entity_tasks @@ -246,7 +246,7 @@ preprocessing and GPU-bound training have different hardware requirements. RelAr permits methods to compute preprocessing artifacts once and cache them on disk before a run. Caching is not required. RelArena provides an **optional, experimental** helper API in -[`relarena.cache`](packages/relarena/src/relarena/cache.py) for local paths, miss policies, private scratch +[`relarena.core.cache`](packages/relarena/src/relarena/cache.py) for local paths, miss policies, private scratch computation, and atomic publication. A method may ignore this API and implement caching independently. The helper does not bring cache warming into a timed RelArena experiment; preprocessing scripts still run separately, so their runtime is not currently included in the diff --git a/docs/adding-a-model.md b/docs/adding-a-model.md index 55bc48e..82a04ea 100644 --- a/docs/adding-a-model.md +++ b/docs/adding-a-model.md @@ -33,9 +33,9 @@ One folder may register **several** models: `dummy` → `constant-global` + import numpy as np from relbench.base import Database, EntityTask, Table -from relarena.model import RelArenaModel -from relarena.registry import register_model -from relarena.search_space import SearchSpace +from relarena.core.model import RelArenaModel +from relarena.core.registry import register_model +from relarena.core.search_space import SearchSpace MYMODEL_SPACE = SearchSpace(space=_config_space(), default_overrides={}) @@ -142,8 +142,8 @@ import numpy as np from relbench.base import EntityTask from relarena.dataset import InnerSplit, OuterSplit -from relarena.registry import register_system -from relarena.system import RelArenaSystem +from relarena.core.registry import register_system +from relarena.core.system import RelArenaSystem @register_system @@ -378,7 +378,7 @@ training settings that do not affect it. Content fingerprints and explicit preprocessing versions can be used to invalidate artifacts when their inputs or meaning change. -See [`relarena.cache`](../src/relarena/cache.py) for the API and its design +See [`relarena.core.cache`](../src/relarena/cache.py) for the API and its design notes, and [`tests/fixtures/cached_model.py`](../tests/fixtures/cached_model.py) for a compact end-to-end example. diff --git a/docs/models/kurversc.md b/docs/models/kurversc.md index 1f06256..f276e44 100644 --- a/docs/models/kurversc.md +++ b/docs/models/kurversc.md @@ -8,7 +8,7 @@ feature program and a CatBoost learner. It does not require or use a GPU. Sync RelArena's CPU dependency group and the KurveRSC extra: ```bash -uv sync --group cpu --extra kurversc +uv sync --all-packages --group cpu --extra kurversc ``` ## Running KurveRSC @@ -16,7 +16,7 @@ uv sync --group cpu --extra kurversc Run one task through the ordinary RelArena CLI: ```bash -OMP_NUM_THREADS=1 uv run --group cpu --extra kurversc relarena \ +OMP_NUM_THREADS=1 uv run --all-packages --group cpu --extra kurversc relarena \ --model kurversc \ --datasets rel-stack \ --tasks user-badge \ @@ -27,7 +27,7 @@ Run all 21 RelBench v1 entity classification and regression tasks by omitting `- `--tasks`: ```bash -OMP_NUM_THREADS=1 uv run --group cpu --extra kurversc relarena \ +OMP_NUM_THREADS=1 uv run --all-packages --group cpu --extra kurversc relarena \ --model kurversc \ --output kurversc_all_tasks.csv ``` @@ -52,7 +52,7 @@ The submitted configuration: The bounded search explores GraphReduce feature-family combinations, graph depth, and automatic annotation. It prunes candidates that exceed the width guard or cannot produce features for the task schema. The fixed values live in -[`src/relarena/models/kurversc/model.py`](../../src/relarena/models/kurversc/model.py), so the +[`packages/relarena/src/relarena/models/kurversc/model.py`](../../packages/relarena/src/relarena/models/kurversc/model.py), so the registered system name denotes one reproducible procedure without hidden configuration fields. Use KurveRSC's public API for ablations or alternative frame budgets. diff --git a/docs/predictive-task.md b/docs/predictive-task.md index d471b9d..b2239e6 100644 --- a/docs/predictive-task.md +++ b/docs/predictive-task.md @@ -27,19 +27,26 @@ supervised prediction task. Two task types are supported today: A task is two YAML files: a **task file** (the label SQL, split timestamps, and what to predict) and a **database file** (the schema and paths to CSV or Parquet tables) -that the task references—so one database file can back many tasks. Load and run it: +that the task references—so one database file can back many tasks. For the hosted +TabPFN-Rel example, install its extra and configure tabpfn-client authentication: + +```bash +pip install "relarena[tabpfn-rel-api]" +``` + +Load and run the task: ```python from relarena.userdb import PredictiveQuery, PredictiveQuerySpec spec = PredictiveQuerySpec.from_yaml("task.yaml", data_dir="data/") -preds = PredictiveQuery(spec).fit(model="tabpfn-rel-client").predict() +preds = PredictiveQuery(spec).fit(model="tabpfn-rel-client", n_trials=0).predict() ``` `from_yaml` reads the task file, resolves its `database:` reference (a path relative -to the task file), and loads the database. `fit` builds the dataset, then tunes and -fits the model on history; `predict` scores the label-less rows at the end of the -data. +to the task file), and loads the database. `fit` builds the dataset and fits the +default configuration when `n_trials=0`; a positive budget enables temporal tuning. +`predict` scores the label-less rows at the end of the data. ## Build the task in four steps @@ -72,8 +79,8 @@ data. different `val_timestamp` and `test_timestamp` values over the same tables. The two files' shapes are defined by -[`database.schema.json`](../src/relarena/userdb/database.schema.json) and -[`task.schema.json`](../src/relarena/userdb/task.schema.json) — JSON Schemas with a +[`database.schema.json`](../packages/relarena-core/src/relarena.core/userdb/database.schema.json) and +[`task.schema.json`](../packages/relarena-core/src/relarena.core/userdb/task.schema.json) — JSON Schemas with a description on every field, validated on load, so a malformed file fails fast with a pointer to the offending field rather than an opaque error later. @@ -259,7 +266,7 @@ maximum depth up to which fastdfs joins. the final fit, and prediction. The first RPI run fills the local store and later runs over the same inputs read it back. Nothing is uploaded. Omit `cache_dir` to fall back to `RELARENA_CACHE_DIR`, or to compute without persistent caching when -neither is set. The underlying `relarena.cache` API is optional and experimental; +neither is set. The underlying `relarena.core.cache` API is optional and experimental; models may implement caching independently. ```python @@ -318,7 +325,7 @@ for this task as posed" rather than a verdict either way. ### Worked examples -- **RelBench v1 (21 tasks)** in `src/relarena/userdb/relbench_v1/` +- **RelBench v1 (21 tasks)** in `packages/relarena/src/relarena/userdb/relbench_v1/` - one folder per dataset (a shared `db.yaml` + one file per task), reproducing RelBench's splits byte-for-byte. `materialize_relbench("rel-f1", "data/rel-f1")` writes the full tables to parquet; `relbench_v1_spec(dataset, task)` loads the diff --git a/docs/temporal-validation.md b/docs/temporal-validation.md index 32f5262..db172ec 100644 --- a/docs/temporal-validation.md +++ b/docs/temporal-validation.md @@ -81,9 +81,14 @@ database as available to every anchor. Doing so cannot reveal test labels or post-boundary data, and it does not recreate the alternative evaluation regime that advances the database to each test entity's timestamp. -The implementation lives in [`dataset.py`](../src/relarena/dataset.py). The -orchestration is in [`runner.py`](../src/relarena/runner.py) and -[`tuner.py`](../src/relarena/tuner.py). +Shared split construction lives in core's +[`dataset.py`](../packages/relarena-core/src/relarena.core/dataset.py), with named +benchmark loading in RelArena's +[`dataset.py`](../packages/relarena/src/relarena/dataset.py). Tuning lives in core's +[`tuner.py`](../packages/relarena-core/src/relarena.core/tuner.py); benchmark +orchestration and final test scoring live in +[`runner.py`](../packages/relarena/src/relarena/runner.py) and +[`refit.py`](../packages/relarena/src/relarena/refit.py). ## Final-fit regimes diff --git a/examples/README.md b/examples/README.md index b316a24..f7461ba 100644 --- a/examples/README.md +++ b/examples/README.md @@ -1,7 +1,7 @@ # RelArena examples -Two runnable examples, answering two different questions. Both are run from the -repository root. +Run these examples from the RelArena repository root. A standalone generated-data +example is in [`packages/tabpfn-rel`](../packages/tabpfn-rel/examples/tiny_database.py). | Example | Question it answers | Needs | |---|---|---| @@ -28,13 +28,13 @@ Three files belong to this example: ```bash uvx kaggle datasets download -d olistbr/brazilian-ecommerce -p data/olist --unzip -uv sync --extra tabpfn-rel-api -uv run python -c "from tabpfn_client import init; init()" +uv sync --all-packages --extra tabpfn-rel-api +uv run --all-packages python -c "from tabpfn_client import init; init()" OMP_NUM_THREADS=1 uv run --no-sync python examples/olist_seller_churn.py ``` That default runs through the hosted TabPFN API, so it needs no GPU. To run the -model locally instead (needs `uv sync --extra rdblearn`, GPU recommended): +model locally instead (needs `uv sync --all-packages --extra tabpfn-rel-local`, GPU recommended): ```bash OMP_NUM_THREADS=1 uv run --no-sync python examples/olist_seller_churn.py --backend local @@ -53,10 +53,10 @@ split timestamps, and avoid leakage — see Fits one RelBench task (rel-f1 / driver-dnf) twice, once with no cache and once against a store warmed up front, and checks the predictions are identical — the cache only changes speed, never results. On that task it turns a roughly 409s -fit-and-predict into roughly 12s. +fit-and-predict into roughly 12s. Run the following commands from the repository root: ```bash -uv run --extra rdblearn python examples/tabpfn_rel_caching.py +uv run --all-packages --extra tabpfn-rel-local python examples/tabpfn_rel_caching.py ``` The expensive step being cached is Deep Feature Synthesis, which runs on CPU. To @@ -64,7 +64,7 @@ exercise the cache path without a GPU, skip the TabPFN forward pass: ```bash RELARENA_EXAMPLE_SKIP_TFM=1 OMP_NUM_THREADS=1 \ - uv run --extra rdblearn python examples/tabpfn_rel_caching.py + uv run --all-packages --extra tabpfn-rel-local python examples/tabpfn_rel_caching.py ``` See the feature-cache section of the [package README](../README.md) for how to diff --git a/examples/olist_seller_churn.py b/examples/olist_seller_churn.py index 6698ada..a263c51 100644 --- a/examples/olist_seller_churn.py +++ b/examples/olist_seller_churn.py @@ -24,7 +24,7 @@ The default uses the hosted TabPFN API. To run the model locally instead: - uv sync --extra rdblearn + uv sync --extra tabpfn-rel-local OMP_NUM_THREADS=1 uv run --no-sync python examples/olist_seller_churn.py \ --backend local """ diff --git a/examples/tabpfn_rel_caching.py b/examples/tabpfn_rel_caching.py index 46ca3f5..36eb520 100644 --- a/examples/tabpfn_rel_caching.py +++ b/examples/tabpfn_rel_caching.py @@ -15,7 +15,7 @@ Run (GPU recommended; the first run downloads the RelBench dataset): - uv run --extra rdblearn python examples/tabpfn_rel_caching.py + uv run --extra local python examples/tabpfn_rel_caching.py RELARENA_EXAMPLE_SKIP_TFM=1 runs only the DFS featurization + cache and skips the TabPFN forward pass, so the caching can be exercised on CPU / locally with no GPU. On @@ -32,13 +32,14 @@ import numpy as np import pandas as pd -from relarena.cache import CacheConfig -from relarena.dataset import OuterSplit, RelBenchDatasetTask, concat_tables -from relarena.featurization import build_dfs_features -from relarena.featurization import dfs as dfs_mod -from relarena.featurization.warm_cache import warm_dfs_cache -from relarena.models._shared.tfm.tfm import default_device +from relarena.core.cache import CacheConfig +from relarena.core.dataset import OuterSplit, concat_tables +from relarena.core.featurization import build_dfs_features +from relarena.core.featurization import dfs as dfs_mod +from relarena.core.featurization.warm_cache import warm_dfs_cache +from relarena.dataset import RelBenchDatasetTask from relarena.models.tabpfn_rel.model import TABPFN_REL_LOCAL_SPACE, TabPFNRelModel +from relarena.models.tabpfn_rel.tfm import default_device #: A reasonably sized RelBench entity task: small enough to run, big enough that the #: DFS cost is visible. Swap for e.g. ("rel-hm", "user-churn") for a heavier one. @@ -50,7 +51,7 @@ DEBUG_SKIP_TFM = os.environ.get("RELARENA_EXAMPLE_SKIP_TFM", "") == "1" #: The reference TabPFN-Rel config, minus text embeddings (keeps dependencies to the -#: `rdblearn` extra; the DFS cache is what this example is about). +#: `local` extra; the DFS cache is what this example is about). CONFIG = {**TABPFN_REL_LOCAL_SPACE.default_overrides} diff --git a/packages/relarena/src/relarena/__init__.py b/packages/relarena/src/relarena/__init__.py index b488e03..3519264 100644 --- a/packages/relarena/src/relarena/__init__.py +++ b/packages/relarena/src/relarena/__init__.py @@ -11,31 +11,33 @@ predictions as useful metadata for later analysis. """ -from relarena.cache import CacheConfig, CacheMiss, cache_key, cached_artifact from relarena.checksums import ( database_checksum, split_checksums, table_checksum, ) -from relarena.dataset import InnerSplit, OuterSplit, RelBenchDatasetTask, Split -from relarena.identity import RunIdentity -from relarena.model import RelArenaModel -from relarena.registry import ( +from relarena.core.cache import CacheConfig, CacheMiss, cache_key, cached_artifact +from relarena.core.dataset import InnerSplit, OuterSplit, Split +from relarena.core.identity import RunIdentity +from relarena.core.model import RelArenaModel +from relarena.core.registry import ( MethodRegistry, ModelRegistry, register_model, register_system, registry, ) -from relarena.results import SystemResult, TrialResult, summary_to_dataframe +from relarena.core.results import SystemResult, TrialResult +from relarena.core.system import RelArenaSystem +from relarena.core.tuner import tune +from relarena.dataset import RelBenchDatasetTask +from relarena.results import summary_to_dataframe from relarena.runner import ( run_experiment, run_model_experiment, run_system_experiment, ) -from relarena.system import RelArenaSystem from relarena.tasks import RELBENCH_V1_DATASETS, TaskSpec, list_entity_tasks -from relarena.tuner import tune __all__ = [ "RELBENCH_V1_DATASETS", diff --git a/packages/relarena/src/relarena/checksums/checksum.py b/packages/relarena/src/relarena/checksums/checksum.py index 711dab3..af72954 100644 --- a/packages/relarena/src/relarena/checksums/checksum.py +++ b/packages/relarena/src/relarena/checksums/checksum.py @@ -18,116 +18,18 @@ from __future__ import annotations -import hashlib import json -import sys from collections.abc import Iterator from pathlib import Path -import numpy as np -import pandas as pd -from relbench.base import Database, EntityTask, Table +from relbench.base import EntityTask, Table +from relarena.core.fingerprints import array_checksum, database_checksum, table_checksum from relarena.dataset import RelBenchDatasetTask, drop_noncanonical_task_columns #: Recorded baseline, shipped as package data beside this module. CHECKSUMS_PATH = Path(__file__).with_name("relbench_v1_checksums.json") -#: Byte width -> unsigned view dtype, for reinterpreting any fixed-width column. -_BYTES_TO_UINT: dict[int, type] = { - 1: np.uint8, - 2: np.uint16, - 4: np.uint32, - 8: np.uint64, -} - - -def array_checksum(arr: np.ndarray) -> np.uint64: - """A `uint64` checksum of a numpy array (ported from `benchmarking`). - - XOR-reduces the values, bit-rotating within blocks of `num_bits` by the - within-block index (so position matters and 0/1-heavy columns don't collapse) - and normalizing endianness. Requires a fixed-width dtype — reduce - object/string/list columns via `_column_codes` first. - """ - itemsize = np.dtype(arr.dtype).itemsize - uint_dtype = _BYTES_TO_UINT.get(itemsize) - if uint_dtype is None: - raise ValueError(f"Unsupported dtype {arr.dtype}") - num_bits = itemsize * 8 - arr = arr.view(uint_dtype).flatten() - pad = (num_bits - arr.size % num_bits) % num_bits - blocks = np.pad(arr, (0, pad), mode="constant").reshape(-1, num_bits) - if sys.byteorder == "big": - blocks = blocks.byteswap(inplace=False) - left = np.arange(num_bits, dtype=np.uint8) - right = np.arange(num_bits, 0, -1, dtype=np.uint8) - rotated = (blocks << left) | (blocks >> right) - return np.bitwise_xor.reduce(rotated.flatten(), dtype=np.uint64) - - -def _hash64(text: str) -> np.uint64: - """Deterministic `uint64` from a string (first 8 bytes of its SHA256).""" - return np.uint64(int.from_bytes(hashlib.sha256(text.encode()).digest()[:8], "big")) - - -def _to_hashable(v: object) -> object: - """Turn array/list cells into hashable tuples; leave scalars untouched.""" - if isinstance(v, np.ndarray): - return tuple(v.tolist()) - if isinstance(v, list): - return tuple(v) - return v - - -def _column_codes(s: pd.Series) -> np.ndarray: - """Reduce a column to one `uint64` per row, handling every RelBench dtype. - - Fixed-width numeric/bool/datetime arrays are reinterpreted directly; strings, - `list`-valued columns, and nullable dtypes (object under `to_numpy()`) go - through pandas' C-level row hasher, with a per-row fallback for the unhashable - (array/list) cells. - """ - if pd.api.types.is_datetime64_any_dtype(s): - return s.to_numpy(dtype="datetime64[ns]").view(np.int64) - a = s.to_numpy() - if a.dtype.kind == "b": # real (non-nullable) bool - return a.astype(np.uint8) - if a.dtype.kind in "iuf": # real fixed-width numeric - return a if a.dtype.itemsize in _BYTES_TO_UINT else a.astype(np.float64) - col = pd.Series(a) - try: - return pd.util.hash_pandas_object(col, index=False).to_numpy() - except TypeError: # unhashable cells (e.g. list<...> columns) - return pd.util.hash_pandas_object(col.map(_to_hashable), index=False).to_numpy() - - -def table_checksum(table: Table) -> np.uint64: - """Integer content checksum of a `Table` (data + relational schema).""" - cs = np.uint64(0) - for i, col in enumerate(sorted(table.df.columns, key=str)): - # Roll each column by its index so identical value-arrays in different - # columns don't cancel under XOR; fold in the column name + dtype. - cs ^= array_checksum(np.roll(_column_codes(table.df[col]), i)) - cs ^= _hash64(f"{col}:{table.df[col].dtype}") - cs ^= _hash64( - f"fkey={sorted(table.fkey_col_to_pkey_table.items())}" - f"|pkey={table.pkey_col}|time={table.time_col}" - ) - return cs - - -def database_checksum(db: Database) -> np.uint64: - """Integer content checksum of a `Database` (each table, hashed by name). - - Name and content are hashed *jointly* so swapping which table sits under which - name changes the result. - """ - cs = np.uint64(0) - for name in sorted(db.table_dict): - cs ^= _hash64(f"{name}|{int(table_checksum(db.table_dict[name]))}") - return cs - def _db_checksums(source: RelBenchDatasetTask) -> dict[str, int]: """Checksums of the censored inner/outer databases (dataset-level, slow).""" @@ -264,3 +166,14 @@ def check_checksums( if diff: mismatches[key] = diff return mismatches + + +__all__ = [ + "array_checksum", + "table_checksum", + "database_checksum", + "split_checksums", + "record_checksums", + "check_checksums", + "CHECKSUMS_PATH", +] diff --git a/packages/relarena/src/relarena/cli.py b/packages/relarena/src/relarena/cli.py index c8fe527..c08af4a 100644 --- a/packages/relarena/src/relarena/cli.py +++ b/packages/relarena/src/relarena/cli.py @@ -19,11 +19,11 @@ import argparse import sys +from importlib import import_module import pandas as pd -import relarena.models # noqa: F401 (registers built-in models) -from relarena.registry import registry +from relarena.core.registry import registry from relarena.results import summary_to_dataframe from relarena.runner import SystemExperimentSummary, run_experiment from relarena.tasks import RELBENCH_V1_DATASETS, list_entity_tasks @@ -88,7 +88,11 @@ def main(argv: list[str] | None = None) -> int: print("Nothing to run.", file=sys.stderr) return 1 - model_cls = registry.get(args.model) + import_module("relarena.models") + try: + model_cls = registry.get(args.model) + except KeyError as exc: + parser.error(str(exc)) frames: list[pd.DataFrame] = [] for s in specs: print( diff --git a/packages/relarena/src/relarena/core/__init__.py b/packages/relarena/src/relarena/core/__init__.py new file mode 100644 index 0000000..0ee3e1f --- /dev/null +++ b/packages/relarena/src/relarena/core/__init__.py @@ -0,0 +1,39 @@ +"""Shared contracts and runtime for relational prediction.""" + +from relarena.core.cache import CacheConfig, CacheMiss, cache_key, cached_artifact +from relarena.core.dataset import InnerSplit, OuterSplit, Split, TaskSource +from relarena.core.identity import RunIdentity +from relarena.core.model import RelArenaModel +from relarena.core.registry import ( + MethodRegistry, + ModelRegistry, + register_model, + register_system, + registry, +) +from relarena.core.results import SystemResult, TrialResult +from relarena.core.system import RelArenaSystem +from relarena.core.tuner import tune + +__version__ = "0.0.1" +__all__ = [ + "CacheConfig", + "CacheMiss", + "cache_key", + "cached_artifact", + "TaskSource", + "Split", + "InnerSplit", + "OuterSplit", + "RunIdentity", + "RelArenaModel", + "RelArenaSystem", + "MethodRegistry", + "ModelRegistry", + "register_model", + "register_system", + "registry", + "TrialResult", + "SystemResult", + "tune", +] diff --git a/packages/relarena/src/relarena/cache.py b/packages/relarena/src/relarena/core/cache.py similarity index 100% rename from packages/relarena/src/relarena/cache.py rename to packages/relarena/src/relarena/core/cache.py diff --git a/packages/relarena/src/relarena/core/dataset.py b/packages/relarena/src/relarena/core/dataset.py new file mode 100644 index 0000000..e15a96a --- /dev/null +++ b/packages/relarena/src/relarena/core/dataset.py @@ -0,0 +1,302 @@ +"""Task sources and censored temporal splits for supplied relational data.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Callable, ClassVar + +import pandas as pd +from relbench.base import Database, Dataset, EntityTask, Table + +from relarena.core.identity import RunIdentity +from relarena.core.metrics import primary_metric + + +def clean_database( + db: Database, excluded_columns: dict[str, tuple[str, ...]] | None = None +) -> Database: + """Drop empty/artifact columns and any explicitly excluded columns.""" + leaking = excluded_columns or {} + table_dict = dict(db.table_dict) + changed = False + for table_name, table in db.table_dict.items(): + df = table.df + # "Unnamed: N" columns are leftover pandas row indices from a CSV export. + drop = {c for c in df.columns if c.startswith("Unnamed:")} + if not df.empty: + drop |= {c for c in df.columns if df[c].isna().all()} + drop |= {c for c in leaking.get(table_name, ()) if c in df.columns} + if not drop: + continue + changed = True + table_dict[table_name] = Table( + df=df.drop(columns=list(drop)), + fkey_col_to_pkey_table=dict(table.fkey_col_to_pkey_table), + pkey_col=table.pkey_col, + time_col=table.time_col, + ) + return Database(table_dict) if changed else db + + +def clean_task_table( + task: EntityTask, table: Table, dataset_name: str, *, allow_extra: bool = False +) -> Table: + """Validate task columns and optionally remove non-task columns.""" + canonical = {task.time_col, task.entity_col, task.target_col} + extra = [c for c in table.df.columns if c not in canonical] + if not extra: + return table + if not allow_extra: + raise ValueError( + f"Task table for {dataset_name!r} has unexpected non-canonical " + f"column(s) {extra}; investigate rather than dropping." + ) + keep = [c for c in table.df.columns if c in canonical] + return Table( + df=table.df[keep], + fkey_col_to_pkey_table=table.fkey_col_to_pkey_table, + pkey_col=table.pkey_col, + time_col=table.time_col, + ) + + +def concat_tables(a: Table, b: Table) -> Table: + """Concatenate two task label tables (e.g. train + val) into one, for refitting. + + Assumes both share the same schema (entity/foreign keys, time and target + columns) — which holds for the splits of a single task. Inputs are untouched. + """ + return Table( + df=pd.concat([a.df, b.df], ignore_index=True), + fkey_col_to_pkey_table=dict(a.fkey_col_to_pkey_table), + pkey_col=a.pkey_col, + time_col=a.time_col, + ) + + +@dataclass(frozen=True) +class Split: + """Fields common to one fit→evaluate phase of nested temporal validation. + + Bundles the censored database with the label tables a phase needs, so a model + can never accidentally see data past its phase's cutoff. Not instantiated + directly: a phase is always an `InnerSplit` or an `OuterSplit`, + which differ in *how the predictions are scored* (see those classes). The + harness decides how each field is used (it does not pass `eval_table` to + `fit` on the outer split, for instance) — see + the tuning and final-fitting callers. + """ + + #: Database censored at `cutoff` — the only DB the model may read. + db_state: Database + #: The censoring cutoff (`val_timestamp` for inner, `test_timestamp` for outer). + cutoff: pd.Timestamp + #: The label table the model trains on. + train_table: Table + #: The table predictions are made on. For the outer split this is the *masked* + #: test table — RelBench's `get_table("test")` strips the target column so the + #: model can predict on the test entities/timestamps without seeing the answers + #: (leakage prevention); the labels live only inside RelBench. The inner split's + #: val table is unmasked, but the model is still expected to use it only as a + #: prediction/early-stopping set, never to read its labels for training. + eval_table: Table + + +@dataclass(frozen=True) +class InnerSplit(Split): + """Tuning phase: fit `train` → score `val`. + + Scoring needs the val labels handed in explicitly. RelBench's + `EntityTask.evaluate(pred, target_table=None)` treats `None` as "score + against the *test* table" — there is no shortcut for val — so to score on val + we must pass the val target table ourselves (`eval_target`). The val + table's labels are not hidden, so this is just the val table itself. + """ + + name: ClassVar[str] = "inner" + #: Val labels to score `eval_table` predictions against. Passed to + #: `EntityTask.evaluate(pred, target_table=eval_target)` — required because + #: `evaluate`'s `None` default would score against test, not val. + eval_target: Table + + +@dataclass(frozen=True) +class OuterSplit(Split): + """Final phase: fit the selected config → score `test`. + + `train_table` is the train-only table and `val_table` the + val table, exposed separately so the harness can serve either final-fit regime: + refit on their union (`refit_on_full_data=True`), or train on train alone with + val as a held-out checkpoint/early-stopping set (`refit_on_full_data=False`). + + Deliberately carries *no* eval target. The test labels are hidden (the model's + `eval_table` is the masked test table), and we score by calling + `EntityTask.evaluate(pred, target_table=None)`, which makes RelBench load its + own held-out test labels. Not materializing those labels into this object keeps + the answer key out of every structure the harness passes around the model — so + test-label leakage is impossible by construction, not just by convention. + """ + + name: ClassVar[str] = "outer" + #: The val label table, held out from `train_table`. The harness + #: unions it with the train table to refit on full data, or passes it as the + #: monitoring set when a model trains on train alone. + val_table: Table + + +def _drop_dangling_seeds(table: Table, entity_col: str, num_entities: int) -> Table: + """Drop seeds whose entity is absent from a censored entity table. + + `entity_col` is the seed table's column of entity ids (each a row index into the + entity table); a value `>= num_entities` references an entity created after the + val cutoff, absent from the censored graph. This is the val-cutoff analog of + relbench's `EntityTask.filter_dangling_entities` (which only filters against the + test-censored `get_db` count). + """ + keep = table.df[entity_col] < num_entities + if keep.all(): + return table + return Table( + df=table.df[keep].reset_index(drop=True), + fkey_col_to_pkey_table=table.fkey_col_to_pkey_table, + pkey_col=table.pkey_col, + time_col=table.time_col, + ) + + +def _copy_timeless_tables(db: Database) -> None: + """Replace each timeless table's DataFrame with a copy, in place. + + `Table.upto` returns timeless tables (`time_col is None`) as the *same + object* held by the source DB, so a censored DB shares their DataFrames with + it. Any in-place mutation of the censored DB (e.g. the dangling-FK scrub in + `validate_and_correct_db`, `df.loc[mask, fkey] = None`) would otherwise + leak back into the source DB and corrupt later splits derived from it. + """ + for name, table in db.table_dict.items(): + if table.time_col is None: + db.table_dict[name] = Table( + df=table.df.copy(), + fkey_col_to_pkey_table=table.fkey_col_to_pkey_table, + pkey_col=table.pkey_col, + time_col=table.time_col, + ) + + +class TaskSource: + """Construct temporal splits from supplied database and task objects.""" + + @staticmethod + def _prepare_db(db: Database, dataset_name: str) -> Database: + return clean_database(db) + + @staticmethod + def _prepare_table(task: EntityTask, table: Table, dataset_name: str) -> Table: + return clean_task_table(task, table, dataset_name) + + @classmethod + def from_objects( + cls, + dataset: Dataset, + task: EntityTask, + *, + dataset_name: str = "user", + task_name: str | None = None, + run_identity: RunIdentity | None = None, + ) -> TaskSource: + """Construct censored splits from supplied dataset and task objects. + + The default preparation drops empty and artifact database columns and + validates task columns. Subclasses can apply source-specific cleaning + through the preparation hooks. Names scope identity metadata. + """ + self = cls.__new__(cls) + self.dataset_name = dataset_name + self.task_name = task_name if task_name is not None else type(task).__name__ + self._identity = run_identity or RunIdentity( + dataset_name, None, self.task_name, None + ) + self._dataset = dataset + self._task = task + self._db = cls._prepare_db(self._dataset.get_db(), dataset_name) + self._tables = { + split: cls._prepare_table( + self._task, self._task.get_table(split), self.dataset_name + ) + for split in ("train", "val", "test") + } + return self + + def run_identity(self, phase: str | None = None) -> RunIdentity: + """Return optional source metadata scoped to one run phase.""" + return self._identity.for_phase(phase) + + @property + def task(self) -> EntityTask: + """The underlying RelBench task (defines the target, metrics, `evaluate`).""" + return self._task + + @property + def metric(self) -> Callable[..., float]: + """The primary metric this task is tuned and selected on (by task type).""" + return primary_metric(self._task) + + def inner_split(self) -> InnerSplit: + """Tuning split: fit `train` → score `val`, DB frozen at `val_timestamp`. + + Censoring the DB at the val cutoff makes validation features *frozen at + their cutoff*, mirroring how test features are frozen at the test cutoff — + see `docs/temporal-validation.md` for why this matters for aggregating + models. `eval_target` is the val table itself (its labels are not hidden). + """ + # The val-cutoff removes pkey rows dated after val_timestamp; any foreign key + # pointing *forward* in time to such a row (e.g. an attendance row referencing + # a later event) is now dangling. get_db() scrubs dangling FKs for the test + # cutoff via validate_and_correct_db, but this extra censor needs the same + # correction or make_pkey_fkey_graph asserts on the out-of-range index. + inner_db = self._db.upto(self._dataset.val_timestamp) + # `validate_and_correct_db` scrubs dangling FKs in place, so decouple the + # timeless tables `upto` shares with `self._db` first (see the helper), + # otherwise the scrub leaks back into `self._db` and corrupts outer_split. + _copy_timeless_tables(inner_db) + self._dataset.validate_and_correct_db(inner_db) + # The same val-cutoff can shrink the entity table below seeds that reference + # entities created after val_timestamp — relbench's get_table only filtered + # seeds against the *test* count. Drop those seeds here (the val-cutoff analog + # of EntityTask.filter_dangling_entities) so a graph sampler can't index past + # the censored entity table. The val table is filtered once and used for both + # eval_table and eval_target, so scored predictions still align with targets. + n_entities = len(inner_db.table_dict[self._task.entity_table].df) + train_table = _drop_dangling_seeds( + self._tables["train"], self._task.entity_col, n_entities + ) + val_table = _drop_dangling_seeds( + self._tables["val"], self._task.entity_col, n_entities + ) + return InnerSplit( + db_state=inner_db, + cutoff=self._dataset.val_timestamp, + train_table=train_table, + eval_table=val_table, + eval_target=val_table, + ) + + def outer_split(self) -> OuterSplit: + """Final split: score `test` after fitting, DB at `test_timestamp`. + + Exposes the train and val tables separately (rather than pre-unioning them) + so the harness can serve either final-fit regime — refit on their union, or + train on train alone with val held out. No `eval_target`: the test labels + are hidden, so scoring goes through `EntityTask.evaluate(pred, + target_table=None)` and RelBench supplies them. + """ + return OuterSplit( + # No-op re-censor: get_db() already censors the DB at test_timestamp + # (upto_test_timestamp=True), so this .upto() drops nothing. Kept for + # readability/symmetry with inner_split's explicit val-cutoff censoring. + db_state=self._db.upto(self._dataset.test_timestamp), + cutoff=self._dataset.test_timestamp, + train_table=self._tables["train"], + val_table=self._tables["val"], + eval_table=self._tables["test"], + ) diff --git a/packages/relarena/src/relarena/core/featurization/__init__.py b/packages/relarena/src/relarena/core/featurization/__init__.py new file mode 100644 index 0000000..a7ea18c --- /dev/null +++ b/packages/relarena/src/relarena/core/featurization/__init__.py @@ -0,0 +1,5 @@ +"""Shared relational deep feature synthesis.""" + +from relarena.core.featurization.dfs import DFS_MAX_DEPTH, build_dfs_features + +__all__ = ["DFS_MAX_DEPTH", "build_dfs_features"] diff --git a/packages/relarena/src/relarena/featurization/cache.py b/packages/relarena/src/relarena/core/featurization/cache.py similarity index 95% rename from packages/relarena/src/relarena/featurization/cache.py rename to packages/relarena/src/relarena/core/featurization/cache.py index ff3eaeb..2c20fdf 100644 --- a/packages/relarena/src/relarena/featurization/cache.py +++ b/packages/relarena/src/relarena/core/featurization/cache.py @@ -7,7 +7,7 @@ import pandas as pd -from relarena.cache import CacheConfig, cached_artifact +from relarena.core.cache import CacheConfig, cached_artifact def cached_frame( diff --git a/packages/relarena/src/relarena/featurization/_columns.py b/packages/relarena/src/relarena/core/featurization/columns.py similarity index 100% rename from packages/relarena/src/relarena/featurization/_columns.py rename to packages/relarena/src/relarena/core/featurization/columns.py diff --git a/packages/relarena/src/relarena/featurization/dfs.py b/packages/relarena/src/relarena/core/featurization/dfs.py similarity index 98% rename from packages/relarena/src/relarena/featurization/dfs.py rename to packages/relarena/src/relarena/core/featurization/dfs.py index 7958b9f..b45d359 100644 --- a/packages/relarena/src/relarena/featurization/dfs.py +++ b/packages/relarena/src/relarena/core/featurization/dfs.py @@ -49,11 +49,11 @@ import pandas as pd from relbench.base import Database, EntityTask, Table -from relarena.cache import CacheConfig, cache_key -from relarena.checksums import database_checksum, table_checksum -from relarena.featurization._columns import type_columns -from relarena.featurization.cache import cached_frame -from relarena.identity import RunIdentity +from relarena.core.cache import CacheConfig, cache_key +from relarena.core.featurization.cache import cached_frame +from relarena.core.featurization.columns import type_columns +from relarena.core.fingerprints import database_checksum, table_checksum +from relarena.core.identity import RunIdentity if TYPE_CHECKING: # pragma: no cover - typing only from fastdfs import RDB @@ -535,7 +535,7 @@ def build_dfs_features( across splits. """ # The DFS deps are an extra, so these imports stay function-local: at module - # scope they make `relarena.featurization` unimportable without the extra, + # scope they make `relarena.core.featurization` unimportable without the extra, # which silently drops every model that imports it from the registry. from fastdfs import DFSConfig, compute_dfs_features from fastdfs.dfs import dfs_feature_column_name, get_dfs_engine diff --git a/packages/relarena/src/relarena/core/featurization/warm_cache.py b/packages/relarena/src/relarena/core/featurization/warm_cache.py new file mode 100644 index 0000000..a07d18f --- /dev/null +++ b/packages/relarena/src/relarena/core/featurization/warm_cache.py @@ -0,0 +1,46 @@ +"""Shared DFS cache warming for supplied task sources.""" + +from __future__ import annotations + +from relarena.core.cache import CacheConfig +from relarena.core.dataset import TaskSource, concat_tables +from relarena.core.featurization.dfs import DFS_MAX_DEPTH, build_dfs_features + + +def warm_dfs_cache( + source: TaskSource, + cache: CacheConfig, + *, + max_depth: int = DFS_MAX_DEPTH, +) -> None: + """Fill shared DFS matrices for tuning and both final-fit history regimes.""" + if cache.directory is None or cache.on_miss != "fill": + raise ValueError("DFS warming needs CacheConfig(directory, on_miss='fill')") + inner, outer = source.inner_split(), source.outer_split() + full_outer_history = concat_tables(outer.train_table, outer.val_table) + phases = ( + ("inner", inner.db_state, inner.train_table, inner.eval_table), + # RDBLearn follows its published train-only final-fit protocol, whereas + # TabPFN-Rel refits on train+val. The actual history input + # belongs in the key, so warm both canonical regimes without model dispatch. + ("outer", outer.db_state, outer.train_table, outer.eval_table), + ( + "outer", + outer.db_state, + full_outer_history, + outer.eval_table, + ), + ) + for phase, db, history, evaluation in phases: + identity = source.run_identity(phase) + for anchors in (history, evaluation): + build_dfs_features( + source.task, + db, + anchors, + depth=max_depth, + max_depth=max_depth, + history_table=history if source.task.time_col else None, + cache=cache, + run_identity=identity, + ) diff --git a/packages/relarena/src/relarena/core/fingerprints.py b/packages/relarena/src/relarena/core/fingerprints.py new file mode 100644 index 0000000..79b82a9 --- /dev/null +++ b/packages/relarena/src/relarena/core/fingerprints.py @@ -0,0 +1,103 @@ +"""Content fingerprints for relational tables and databases.""" + +import hashlib +import sys + +import numpy as np +import pandas as pd +from relbench.base import Database, Table + +#: Byte width -> unsigned view dtype, for reinterpreting any fixed-width column. +_BYTES_TO_UINT: dict[int, type] = { + 1: np.uint8, + 2: np.uint16, + 4: np.uint32, + 8: np.uint64, +} + + +def array_checksum(arr: np.ndarray) -> np.uint64: + """A `uint64` checksum of a numpy array (ported from `benchmarking`). + + XOR-reduces the values, bit-rotating within blocks of `num_bits` by the + within-block index (so position matters and 0/1-heavy columns don't collapse) + and normalizing endianness. Requires a fixed-width dtype — reduce + object/string/list columns via `_column_codes` first. + """ + itemsize = np.dtype(arr.dtype).itemsize + uint_dtype = _BYTES_TO_UINT.get(itemsize) + if uint_dtype is None: + raise ValueError(f"Unsupported dtype {arr.dtype}") + num_bits = itemsize * 8 + arr = arr.view(uint_dtype).flatten() + pad = (num_bits - arr.size % num_bits) % num_bits + blocks = np.pad(arr, (0, pad), mode="constant").reshape(-1, num_bits) + if sys.byteorder == "big": + blocks = blocks.byteswap(inplace=False) + left = np.arange(num_bits, dtype=np.uint8) + right = np.arange(num_bits, 0, -1, dtype=np.uint8) + rotated = (blocks << left) | (blocks >> right) + return np.bitwise_xor.reduce(rotated.flatten(), dtype=np.uint64) + + +def _hash64(text: str) -> np.uint64: + """Deterministic `uint64` from a string (first 8 bytes of its SHA256).""" + return np.uint64(int.from_bytes(hashlib.sha256(text.encode()).digest()[:8], "big")) + + +def _to_hashable(v: object) -> object: + """Turn array/list cells into hashable tuples; leave scalars untouched.""" + if isinstance(v, np.ndarray): + return tuple(v.tolist()) + if isinstance(v, list): + return tuple(v) + return v + + +def _column_codes(s: pd.Series) -> np.ndarray: + """Reduce a column to one `uint64` per row, handling every RelBench dtype. + + Fixed-width numeric/bool/datetime arrays are reinterpreted directly; strings, + `list`-valued columns, and nullable dtypes (object under `to_numpy()`) go + through pandas' C-level row hasher, with a per-row fallback for the unhashable + (array/list) cells. + """ + if pd.api.types.is_datetime64_any_dtype(s): + return s.to_numpy(dtype="datetime64[ns]").view(np.int64) + a = s.to_numpy() + if a.dtype.kind == "b": # real (non-nullable) bool + return a.astype(np.uint8) + if a.dtype.kind in "iuf": # real fixed-width numeric + return a if a.dtype.itemsize in _BYTES_TO_UINT else a.astype(np.float64) + col = pd.Series(a) + try: + return pd.util.hash_pandas_object(col, index=False).to_numpy() + except TypeError: # unhashable cells (e.g. list<...> columns) + return pd.util.hash_pandas_object(col.map(_to_hashable), index=False).to_numpy() + + +def table_checksum(table: Table) -> np.uint64: + """Integer content checksum of a `Table` (data + relational schema).""" + cs = np.uint64(0) + for i, col in enumerate(sorted(table.df.columns, key=str)): + # Roll each column by its index so identical value-arrays in different + # columns don't cancel under XOR; fold in the column name + dtype. + cs ^= array_checksum(np.roll(_column_codes(table.df[col]), i)) + cs ^= _hash64(f"{col}:{table.df[col].dtype}") + cs ^= _hash64( + f"fkey={sorted(table.fkey_col_to_pkey_table.items())}" + f"|pkey={table.pkey_col}|time={table.time_col}" + ) + return cs + + +def database_checksum(db: Database) -> np.uint64: + """Integer content checksum of a `Database` (each table, hashed by name). + + Name and content are hashed *jointly* so swapping which table sits under which + name changes the result. + """ + cs = np.uint64(0) + for name in sorted(db.table_dict): + cs ^= _hash64(f"{name}|{int(table_checksum(db.table_dict[name]))}") + return cs diff --git a/packages/relarena/src/relarena/core/identity.py b/packages/relarena/src/relarena/core/identity.py new file mode 100644 index 0000000..c3b9e0d --- /dev/null +++ b/packages/relarena/src/relarena/core/identity.py @@ -0,0 +1,100 @@ +"""Source identity passed explicitly to preprocessing-owned cache-key factories. + +`RunIdentity` describes the data and execution phase from which preprocessing +artifacts are produced. Entrypoints construct it, the runner passes it through +model construction, and each preprocessing module decides which fields belong +in its own artifact key. + +A run identity is metadata, not a complete cache key and not a central key +policy. It intentionally contains more information than every artifact needs. +For example, task-dependent DFS and RelGT artifacts use task identity, while a +RelGNN graph may ignore it because the graph depends only on the censored +database. Preprocessors remain responsible for selecting actual dependencies +and adding their own algorithm versions. + +The readable dataset and task names provide namespaces. Their fingerprints +distinguish different underlying data or task definitions without placing +absolute paths or expensive row hashes in keys. `phase` distinguishes +protocol views such as `inner`, `outer`, and `predict` when censoring or +inputs differ. `data_version` is an optional caller-supplied discriminator +for data changes that cheap fingerprints cannot observe. + +RelBench identities use checked-in dataset and task checksums. Predictive-query +identities use a database-schema fingerprint and a task-specification +fingerprint; because the schema fingerprint deliberately ignores row contents, +callers using persistent caches should provide `data_version` when those +contents may change without a schema change. + +Each preprocessing owner decides whether persistent use requires complete +fingerprints or whether it can derive a safe fallback from its actual inputs. +Unconfigured direct callers may omit identity and compute in private scratch. +Cache directories, miss policies, model names, serialization formats, and +preprocessing algorithm versions do not belong in this object. +""" + +from __future__ import annotations + +import hashlib +import json +from dataclasses import dataclass, replace +from typing import Any + +from relbench.base import Database + + +@dataclass(frozen=True) +class RunIdentity: + """Readable source metadata that a preprocessor may use or ignore.""" + + dataset: str + dataset_fingerprint: str | None + task: str | None + task_fingerprint: str | None + data_version: str | None = None + phase: str | None = None + + def for_phase(self, phase: str | None) -> RunIdentity: + """Return the same identity scoped to one execution phase.""" + return replace(self, phase=phase) + + +def metadata_fingerprint(value: Any) -> str: + """Hash JSON metadata as sorted compact JSON with a 64-bit BLAKE2s digest.""" + encoded = json.dumps(value, sort_keys=True, separators=(",", ":")).encode() + return hashlib.blake2s(encoded, digest_size=8).hexdigest() + + +def database_schema_fingerprint(db: Database) -> str: + """Fingerprint a user database's relational schema without hashing its rows.""" + schema = { + name: { + "columns": [ + (str(column), str(table.df[column].dtype)) for column in table.df + ], + "fkeys": sorted(table.fkey_col_to_pkey_table.items()), + "pkey": table.pkey_col, + "time": table.time_col, + } + for name, table in sorted(db.table_dict.items()) + } + return metadata_fingerprint(schema) + + +def task_spec_fingerprint(task: Any) -> str: + """Fingerprint training semantics of a user predictive-task specification.""" + fields = { + name: str(getattr(task, name)) + for name in ( + "entity_table", + "entity_col", + "time_col", + "target_col", + "task_type", + "timedelta", + "query", + "val_timestamp", + "test_timestamp", + "num_eval_timestamps", + ) + } + return metadata_fingerprint(fields) diff --git a/packages/relarena/src/relarena/metrics.py b/packages/relarena/src/relarena/core/metrics.py similarity index 88% rename from packages/relarena/src/relarena/metrics.py rename to packages/relarena/src/relarena/core/metrics.py index 928c7d4..5cc7889 100644 --- a/packages/relarena/src/relarena/metrics.py +++ b/packages/relarena/src/relarena/core/metrics.py @@ -19,7 +19,8 @@ from dataclasses import dataclass from typing import Callable, Union -from relbench.base import EntityTask, TaskType +import numpy as np +from relbench.base import EntityTask, Table, TaskType from relbench.metrics import mae, roc_auc @@ -106,7 +107,7 @@ def get_metric(metric: MetricLike) -> Metric: name = _metric_name(metric) if name not in _METRICS: raise KeyError( - f"Unknown metric '{name}'; add it to _METRICS in relarena/metrics.py" + f"Unknown metric '{name}'; add it to _METRICS in relarena.core/metrics.py" ) return _METRICS[name] @@ -140,3 +141,21 @@ def to_metric_error(score: float, metric: MetricLike) -> float: bencheval adapter applies it per row to build the `metric_error` column. """ return get_metric(metric).to_error(score) + + +def evaluate_predictions( + task: EntityTask, + pred: np.ndarray, + target_table: "Table | None", + primary: Callable[..., float], +) -> dict: + """Score `pred` and return all native metrics plus the primary, keyed by name. + + `target_table` is the (unmasked-label) split table for val; pass `None` for + test, so RelBench loads its held-out test labels itself. The primary metric is + appended to the task's native metrics if it isn't already among them. + """ + metrics = list(task.metrics) + if primary.__name__ not in {m.__name__ for m in metrics}: + metrics.append(primary) + return task.evaluate(pred, target_table, metrics=metrics) diff --git a/packages/relarena/src/relarena/model.py b/packages/relarena/src/relarena/core/model.py similarity index 96% rename from packages/relarena/src/relarena/model.py rename to packages/relarena/src/relarena/core/model.py index 6b2898c..9a920ec 100644 --- a/packages/relarena/src/relarena/model.py +++ b/packages/relarena/src/relarena/core/model.py @@ -21,9 +21,9 @@ import numpy as np from relbench.base import Database, EntityTask, Table, TaskType -from relarena.cache import CacheConfig -from relarena.identity import RunIdentity -from relarena.tasks import ENTITY_TASK_TYPES +from relarena.core.cache import CacheConfig +from relarena.core.identity import RunIdentity +from relarena.core.tasks import ENTITY_TASK_TYPES class RelArenaModel(ABC): diff --git a/packages/relarena/src/relarena/models/_shared/predict_contract.py b/packages/relarena/src/relarena/core/predict_contract.py similarity index 97% rename from packages/relarena/src/relarena/models/_shared/predict_contract.py rename to packages/relarena/src/relarena/core/predict_contract.py index 6908cc1..106e9c2 100644 --- a/packages/relarena/src/relarena/models/_shared/predict_contract.py +++ b/packages/relarena/src/relarena/core/predict_contract.py @@ -20,6 +20,8 @@ import numpy as np from relbench.base import TaskType +__all__ = ["predict_to_contract"] + def predict_to_contract(estimator: Any, X: Any, task_type: TaskType) -> np.ndarray: """Predict with `estimator` on `X` and shape it to the evaluate contract. diff --git a/packages/relarena/src/relarena/core/py.typed b/packages/relarena/src/relarena/core/py.typed new file mode 100644 index 0000000..e69de29 diff --git a/packages/relarena/src/relarena/registry.py b/packages/relarena/src/relarena/core/registry.py similarity index 95% rename from packages/relarena/src/relarena/registry.py rename to packages/relarena/src/relarena/core/registry.py index c221cb4..cb5c11c 100644 --- a/packages/relarena/src/relarena/registry.py +++ b/packages/relarena/src/relarena/core/registry.py @@ -10,9 +10,9 @@ from dataclasses import dataclass from typing import Callable, Iterator, Type, TypeAlias -from relarena.model import RelArenaModel -from relarena.search_space import SearchSpaceProvider -from relarena.system import RelArenaSystem +from relarena.core.model import RelArenaModel +from relarena.core.search_space import SearchSpaceProvider +from relarena.core.system import RelArenaSystem Method: TypeAlias = type[RelArenaModel] | type[RelArenaSystem] @@ -100,8 +100,9 @@ def names(self) -> list[str]: def _entry(self, name: str) -> RegistryEntry: if name not in self._entries: + hint = " Import relarena.models to register built-in models." raise KeyError( - f"No method registered under '{name}'. Known: {self.names()}" + f"No method registered under '{name}'. Known: {self.names()}.{hint}" ) return self._entries[name] diff --git a/packages/relarena/src/relarena/core/results.py b/packages/relarena/src/relarena/core/results.py new file mode 100644 index 0000000..748bcad --- /dev/null +++ b/packages/relarena/src/relarena/core/results.py @@ -0,0 +1,81 @@ +"""Model and system result schemas. + +Models produce one `TrialResult` per harness-selected configuration. Systems +produce one `SystemResult` for their complete internal procedure. +""" + +from __future__ import annotations + +import hashlib +import json +from dataclasses import dataclass, field +from typing import Any + +import numpy as np + + +def config_id_for(config: dict[str, Any]) -> str: + """A short, deterministic id for a hyperparameter config (order-independent).""" + blob = json.dumps(config, sort_keys=True, default=str) + return hashlib.sha1(blob.encode()).hexdigest()[:8] + + +@dataclass +class TrialResult: + """Outcome of fitting & evaluating ONE config (within an ExperimentSummary). + + Identity — model / dataset / task / seed / metric — is owned by the enclosing + `ExperimentSummary` (every trial in a summary shares + it), so it lives there once instead of being duplicated on each trial; a + `TrialResult` records only the per-config result. + """ + + config: dict[str, Any] + config_id: str + config_tag: str # "default" or "r{i}" — identifies the default vs random configs + + # The primary (selection) metric's value. `val_score` comes from the + # train-only model (the selection signal); `test_score` is filled only for + # the selected config, from the model's final-fit regime (see the runner). + val_score: float | None = None + test_score: float | None = None + + # All of the task's native metrics (plus the primary), keyed by metric name. + val_metrics: dict[str, float] = field(default_factory=dict) + test_metrics: dict[str, float] = field(default_factory=dict) + + # Wall-clock seconds, split by phase: tuning (train-only fit + val predict, + # per config) and final fit (fit + test predict for selected/default configs). + # The refit-named fields are None for configs that did not receive a final fit. + fit_time_tuning: float = 0.0 + predict_time_tuning: float = 0.0 + fit_time_refit: float | None = None + predict_time_refit: float | None = None + + # Optional prediction metadata (not serialized to the summary DataFrame). + # Shapes follow EntityTask.evaluate's contract. + val_pred: np.ndarray | None = field(default=None, repr=False) + test_pred: np.ndarray | None = field(default=None, repr=False) + + # Populated with a traceback string if the trial failed; `None` on success. + error: str | None = None + + @property + def ok(self) -> bool: + """Whether the trial succeeded (no error).""" + return self.error is None + + +@dataclass +class SystemResult: + """Outcome of one end-to-end system run. + + Systems do not expose harness-selected configurations or validation scores. + Their complete internal procedure is represented by a final test result and + one total wall-clock time. + """ + + test_score: float | None = None + test_metrics: dict[str, float] = field(default_factory=dict) + time_total: float = 0.0 + test_pred: np.ndarray | None = field(default=None, repr=False) diff --git a/packages/relarena/src/relarena/search_space.py b/packages/relarena/src/relarena/core/search_space.py similarity index 100% rename from packages/relarena/src/relarena/search_space.py rename to packages/relarena/src/relarena/core/search_space.py diff --git a/packages/relarena/src/relarena/core/selection.py b/packages/relarena/src/relarena/core/selection.py new file mode 100644 index 0000000..bb00a91 --- /dev/null +++ b/packages/relarena/src/relarena/core/selection.py @@ -0,0 +1,25 @@ +"""Validation-based configuration selection.""" + +import math +from typing import Callable + +from relarena.core.metrics import is_better +from relarena.core.results import TrialResult + + +def select_best(trials: list[TrialResult], metric: Callable[..., float]) -> TrialResult: + """Pick the trial with the best validation score under `metric`'s direction.""" + valid = [ + t + for t in trials + if t.ok and t.val_score is not None and math.isfinite(t.val_score) + ] + if not valid: + raise RuntimeError( + "No successful trials with a finite validation score to select from." + ) + best = valid[0] + for t in valid[1:]: + if is_better(t.val_score, best.val_score, metric): + best = t + return best diff --git a/packages/relarena/src/relarena/system.py b/packages/relarena/src/relarena/core/system.py similarity index 90% rename from packages/relarena/src/relarena/system.py rename to packages/relarena/src/relarena/core/system.py index 464063e..62cbe45 100644 --- a/packages/relarena/src/relarena/system.py +++ b/packages/relarena/src/relarena/core/system.py @@ -15,10 +15,10 @@ import numpy as np from relbench.base import EntityTask, TaskType -from relarena.cache import CacheConfig -from relarena.dataset import InnerSplit, OuterSplit -from relarena.identity import RunIdentity -from relarena.tasks import ENTITY_TASK_TYPES +from relarena.core.cache import CacheConfig +from relarena.core.dataset import InnerSplit, OuterSplit +from relarena.core.identity import RunIdentity +from relarena.core.tasks import ENTITY_TASK_TYPES class RelArenaSystem(ABC): diff --git a/packages/relarena/src/relarena/core/tasks.py b/packages/relarena/src/relarena/core/tasks.py new file mode 100644 index 0000000..ab38f6b --- /dev/null +++ b/packages/relarena/src/relarena/core/tasks.py @@ -0,0 +1,22 @@ +"""Supported relational prediction task types.""" + +from relbench.base import TaskType + +#: Entity (node-level) task types RelArena supports. Excludes: +#: * `LINK_PREDICTION` — recommendation, out of scope; +#: * `MULTILABEL_CLASSIFICATION` — RelBench has no entity multilabel task +#: (its sole multilabel task is a TGB node-property ranking `BaseTask`). +#: Easy to re-add if a real entity multilabel task appears. +ENTITY_TASK_TYPES: frozenset[TaskType] = frozenset( + { + TaskType.BINARY_CLASSIFICATION, + TaskType.REGRESSION, + } +) + +# The rest of the codebase assumes exactly these two task types; guard against +# silently widening scope without revisiting those call sites. +assert ENTITY_TASK_TYPES == { + TaskType.REGRESSION, + TaskType.BINARY_CLASSIFICATION, +}, "RelArena currently supports only regression and binary classification tasks." diff --git a/packages/relarena/src/relarena/core/tfm.py b/packages/relarena/src/relarena/core/tfm.py new file mode 100644 index 0000000..823149b --- /dev/null +++ b/packages/relarena/src/relarena/core/tfm.py @@ -0,0 +1,197 @@ +"""Shared estimator fitting, sampling and prediction mechanics.""" + +from __future__ import annotations + +import inspect +from dataclasses import dataclass +from typing import Any, Callable, Protocol + +import numpy as np +import pandas as pd +from relbench.base import TaskType + +from relarena.core.predict_contract import predict_to_contract + + +class SklearnClassifier(Protocol): + """Minimal sklearn-classifier surface relarena uses (TabPFNClassifier-like).""" + + classes_: np.ndarray + + def fit(self, X: pd.DataFrame, y: np.ndarray) -> Any: + """Fit the classifier on the feature frame and labels.""" + ... + + def predict_proba(self, X: pd.DataFrame) -> np.ndarray: + """Return one probability column per fitted class.""" + ... + + +class SklearnRegressor(Protocol): + """Minimal sklearn-regressor surface relarena uses (TabPFNRegressor-like).""" + + def fit(self, X: pd.DataFrame, y: np.ndarray) -> Any: + """Fit the regressor on the feature frame and targets.""" + ... + + def predict(self, X: pd.DataFrame) -> np.ndarray: + """Return one numeric prediction per input row.""" + ... + + +@dataclass(frozen=True) +class TFMSpec: + """How to build one tabular foundation model. + + Every TFM is assumed to support all entity task types, so there is no per-TFM + task-type gating. `make_classifier` / `make_regressor` take keyword overrides + (`device`, `seed`, ...) and return an estimator satisfying + `SklearnClassifier` / `SklearnRegressor` respectively. + `max_train_samples` is this TFM's training-row cap before fitting — its + supported context size (TabPFN v2 ~10k, v2.5 ~50k) — applied by `fit_tfm`. + `supports_text` marks estimators that handle raw text columns themselves. + """ + + make_classifier: Callable[..., SklearnClassifier] + make_regressor: Callable[..., SklearnRegressor] + max_train_samples: int + supports_text: bool = False + + +def default_device() -> str: + """Return `"cuda"` if a GPU is visible to torch, else `"cpu"`. + + torch is not a core dependency: local estimators bring it, hosted ones do not, + so an environment without torch has no GPU to offer. + """ + try: + import torch + except ModuleNotFoundError: + return "cpu" + return "cuda" if torch.cuda.is_available() else "cpu" + + +def _downsample_indices( + y: np.ndarray, task_type: TaskType, max_samples: int, rng: np.random.Generator +) -> np.ndarray: + """Indices of a <= `max_samples` subset of rows (no-op when already small). + + Regression: a uniform random subset. Classification: keep at least one row per + class, then fill the remaining budget uniformly at random. Seeded via `rng` for + reproducibility. Adapted from RDBLearn's `_downsample` (non-stratified path). + """ + n = len(y) + if n <= max_samples: + return np.arange(n) + + if task_type == TaskType.REGRESSION: + return rng.choice(n, max_samples, replace=False) + + selected: list[int] = [] + for label in np.unique(y): + class_idx = np.where(y == label)[0] + selected.append(int(rng.choice(class_idx, 1)[0])) + selected = list(dict.fromkeys(selected)) # de-dup (one per class) + + remaining = max_samples - len(selected) + if remaining > 0: + mask = np.ones(n, dtype=bool) + mask[selected] = False + eligible = np.where(mask)[0] + extra = rng.choice(eligible, min(remaining, len(eligible)), replace=False) + out = np.concatenate([np.array(selected, dtype=int), extra]) + else: + out = np.array(selected[:max_samples], dtype=int) + rng.shuffle(out) + return out + + +@dataclass +class FittedTFM: + """A fitted TFM plus the state needed to score val/test consistently.""" + + estimator: Any + feature_cols: list[str] + task_type: TaskType + max_predict_samples: int | None = None + + +def fit_tfm( + df: pd.DataFrame, + y: pd.Series, + task_type: TaskType, + *, + spec: TFMSpec, + seed: int, + device: Any = None, + max_train_samples: int | None = None, + max_predict_samples: int | None = None, + overrides: dict[str, Any] | None = None, +) -> FittedTFM: + """Downsample `df` and fit the supplied estimator specification. + + `df` is the already-typed feature frame from `build_dfs_features` (numeric + floats + object categoricals); TabPFN auto-detects categoricals from it (see the + module docstring) — we do not pass `categorical_features_indices`. The training + rows are capped (seeded) at `max_train_samples` if given, else the TFM's own + context cap (`spec.max_train_samples`); `overrides` are additional + estimator-constructor arguments. `max_predict_samples` is an explicit + caller-owned cap on rows per estimator prediction call; ordinary TFM callers + leave it unset. + """ + if device is None: + device = default_device() + cap = max_train_samples if max_train_samples is not None else spec.max_train_samples + rng = np.random.default_rng(seed) + + feature_cols = list(df.columns) + + y_arr = y.to_numpy() + idx = _downsample_indices(y_arr, task_type, cap, rng) + X = df.iloc[idx] + y_arr = y_arr[idx] + + kwargs = dict(device=device, seed=seed, **(overrides or {})) + if task_type == TaskType.REGRESSION: + estimator = spec.make_regressor(**kwargs) + y_arr = y_arr.astype(float) + else: + estimator = spec.make_classifier(**kwargs) + estimator.fit(X, y_arr) + + return FittedTFM(estimator, feature_cols, task_type, max_predict_samples) + + +def _predict_tfm_frame(fitted: FittedTFM, frame: pd.DataFrame) -> np.ndarray: + if fitted.task_type == TaskType.REGRESSION: + predict = fitted.estimator.predict + params = inspect.signature(predict).parameters.values() + if any( + p.name == "output_type" or p.kind is inspect.Parameter.VAR_KEYWORD + for p in params + ): + return np.asarray(predict(frame, output_type="median"), dtype=float) + return predict_to_contract(fitted.estimator, frame, fitted.task_type) + + +def predict_tfm(fitted: FittedTFM, df: pd.DataFrame) -> np.ndarray: + """Predict with a `FittedTFM` on a new feature frame. + + Reindexes to the training column order (so the TFM sees the same schema), then + delegates the sklearn-output -> evaluate-contract reshaping to + `relarena.core.predict_contract.predict_to_contract`. + + Regression requests `output_type="median"` when the estimator supports it — + an explicit output_type parameter, or a **kwargs passthrough: the primary + regression metric is MAE, and the median is its optimal point prediction. + """ + X = df.reindex(columns=fitted.feature_cols) + batch_size = fitted.max_predict_samples + if batch_size is None or len(X) <= batch_size: + return _predict_tfm_frame(fitted, X) + return np.concatenate( + [ + _predict_tfm_frame(fitted, X.iloc[i : i + batch_size]) + for i in range(0, len(X), batch_size) + ] + ) diff --git a/packages/relarena/src/relarena/tuner.py b/packages/relarena/src/relarena/core/tuner.py similarity index 62% rename from packages/relarena/src/relarena/tuner.py rename to packages/relarena/src/relarena/core/tuner.py index 676a60c..4e2ee3e 100644 --- a/packages/relarena/src/relarena/tuner.py +++ b/packages/relarena/src/relarena/core/tuner.py @@ -8,8 +8,8 @@ * every trial records its configuration, metrics, phase-split wall-clock times, and optionally its validation/test predictions as useful analysis metadata. -Selection of the single best config is done downstream (see `runner`); this -module just runs the trials and returns the records. +Callers select the best configuration with `relarena.core.selection.select_best`. +This module runs the trials and returns their records. """ from __future__ import annotations @@ -18,18 +18,17 @@ import os import time import traceback -from typing import Any, Callable, Type +from typing import Any, Type -import numpy as np -from relbench.base import EntityTask, Table +from relbench.base import EntityTask -from relarena.cache import CacheConfig -from relarena.dataset import InnerSplit, OuterSplit, concat_tables -from relarena.identity import RunIdentity -from relarena.metrics import primary_metric -from relarena.model import RelArenaModel -from relarena.results import TrialResult, config_id_for -from relarena.search_space import ( +from relarena.core.cache import CacheConfig +from relarena.core.dataset import InnerSplit +from relarena.core.identity import RunIdentity +from relarena.core.metrics import evaluate_predictions, primary_metric +from relarena.core.model import RelArenaModel +from relarena.core.results import TrialResult, config_id_for +from relarena.core.search_space import ( SearchSpace, SearchSpaceProvider, TaskStats, @@ -53,24 +52,6 @@ def _concise_error(exc: BaseException) -> str: return f"{type(exc).__name__}: {message}{where}" -def _evaluate( - task: EntityTask, - pred: np.ndarray, - target_table: "Table | None", - primary: Callable[..., float], -) -> dict: - """Score `pred` and return all native metrics plus the primary, keyed by name. - - `target_table` is the (unmasked-label) split table for val; pass `None` for - test, so RelBench loads its held-out test labels itself. The primary metric is - appended to the task's native metrics if it isn't already among them. - """ - metrics = list(task.metrics) - if primary.__name__ not in {m.__name__ for m in metrics}: - metrics.append(primary) - return task.evaluate(pred, target_table, metrics=metrics) - - def run_trial( model_cls: Type[RelArenaModel], config: dict[str, Any], @@ -89,7 +70,7 @@ def run_trial( The split's `eval_table` doubles as the model's early-stopping val during `fit` — this is the tuning phase, where the eval set is the held-out validation set. The test number is produced separately by - `refit_and_evaluate` on the outer split. + a caller-owned final fit on the outer split. Never raises on model failure — a failed trial is returned with `error` set, so one bad config does not abort a sweep. @@ -118,7 +99,7 @@ def run_trial( t1 = time.perf_counter() val_pred = model.predict(task, split.db_state, split.eval_table) predict_time_tuning = time.perf_counter() - t1 - val_metrics = _evaluate(task, val_pred, split.eval_target, metric) + val_metrics = evaluate_predictions(task, val_pred, split.eval_target, metric) val_score = float(val_metrics[metric.__name__]) return TrialResult( @@ -202,62 +183,3 @@ def tune( ) for config_tag, config in plan ] - - -def refit_and_evaluate( - model_cls: Type[RelArenaModel], - config: dict[str, Any], - task: EntityTask, - split: OuterSplit, - *, - seed: int, - time_limit: float | None = None, - cache: CacheConfig | None = None, - run_identity: RunIdentity | None = None, -) -> dict: - """Fit the selected `config` on the outer `split` and score it on `test`. - - Two final-fit regimes, chosen by `model_cls.refit_on_full_data`: - - * `True` (default): refit on the train+val union with `val_table=None` — no - held-out split, so a model with early stopping falls back to a fixed budget. - * `False`: train on train alone and pass `val` through, so a model that - checkpoints on validation reports its best-val model (e.g. RelGT's protocol). - - Both train on the outer split's (test-censored) DB. The split carries no eval - target (test labels are hidden), so we score with `target_table=None` and - RelBench sources the unmasked test labels itself. Returns the test fields to - attach to the trial. - """ - metric = primary_metric(task) - - if model_cls.refit_on_full_data: - train_table, val_table = concat_tables(split.train_table, split.val_table), None - else: - train_table, val_table = split.train_table, split.val_table - - model = model_cls(config, cache=cache, run_identity=run_identity) - t0 = time.perf_counter() - model.fit( - task, - split.db_state, - train_table, - val_table, - seed=seed, - time_limit=time_limit, - ) - fit_time_refit = time.perf_counter() - t0 - - t1 = time.perf_counter() - test_pred = model.predict(task, split.db_state, split.eval_table) - predict_time_refit = time.perf_counter() - t1 - - test_metrics = _evaluate(task, test_pred, None, metric) - test_score = float(test_metrics[metric.__name__]) - return { - "test_score": test_score, - "test_metrics": test_metrics, - "test_pred": test_pred, - "fit_time_refit": fit_time_refit, - "predict_time_refit": predict_time_refit, - } diff --git a/packages/relarena/src/relarena/core/userdb/__init__.py b/packages/relarena/src/relarena/core/userdb/__init__.py new file mode 100644 index 0000000..3410cb8 --- /dev/null +++ b/packages/relarena/src/relarena/core/userdb/__init__.py @@ -0,0 +1,12 @@ +"""Relational predictive queries over user-supplied databases.""" + +from relarena.core.userdb.ingest import DatabaseSpec +from relarena.core.userdb.query import PredictiveQuery, PredictiveQuerySpec +from relarena.core.userdb.spec import PredictiveTaskSpec + +__all__ = [ + "DatabaseSpec", + "PredictiveQuery", + "PredictiveQuerySpec", + "PredictiveTaskSpec", +] diff --git a/packages/relarena/src/relarena/userdb/_schema.py b/packages/relarena/src/relarena/core/userdb/_schema.py similarity index 93% rename from packages/relarena/src/relarena/userdb/_schema.py rename to packages/relarena/src/relarena/core/userdb/_schema.py index 355cd10..0fdb89e 100644 --- a/packages/relarena/src/relarena/userdb/_schema.py +++ b/packages/relarena/src/relarena/core/userdb/_schema.py @@ -10,7 +10,7 @@ def load_schema(name: str) -> dict: """Load a JSON Schema shipped alongside this package, by file name.""" - return json.loads((files("relarena.userdb") / name).read_text()) + return json.loads((files("relarena.core.userdb") / name).read_text()) def validate(raw: object, schema: dict, *, kind: str) -> None: diff --git a/packages/relarena/src/relarena/userdb/database.schema.json b/packages/relarena/src/relarena/core/userdb/database.schema.json similarity index 100% rename from packages/relarena/src/relarena/userdb/database.schema.json rename to packages/relarena/src/relarena/core/userdb/database.schema.json diff --git a/packages/relarena/src/relarena/userdb/ingest.py b/packages/relarena/src/relarena/core/userdb/ingest.py similarity index 98% rename from packages/relarena/src/relarena/userdb/ingest.py rename to packages/relarena/src/relarena/core/userdb/ingest.py index ca25765..4e081ba 100644 --- a/packages/relarena/src/relarena/userdb/ingest.py +++ b/packages/relarena/src/relarena/core/userdb/ingest.py @@ -4,7 +4,7 @@ key, time column and foreign-key links per table — via `DatabaseSpec`. `build_dataset` turns that into a reindexed RelBench `UserDataset` (with the val/test split timestamps) that flows through -`RelBenchDatasetTask.from_objects` exactly like a native RelBench dataset. +`TaskSource.from_objects` exactly like a native RelBench dataset. """ from __future__ import annotations @@ -16,7 +16,7 @@ import yaml from relbench.base import Database, Dataset, Table -from relarena.userdb._schema import load_schema, validate +from relarena.core.userdb._schema import load_schema, validate #: JSON Schema for a database YAML; the single source of truth for its shape. _DB_SCHEMA = load_schema("database.schema.json") diff --git a/packages/relarena/src/relarena/userdb/predict.py b/packages/relarena/src/relarena/core/userdb/predict.py similarity index 97% rename from packages/relarena/src/relarena/userdb/predict.py rename to packages/relarena/src/relarena/core/userdb/predict.py index bc68348..1ed5f83 100644 --- a/packages/relarena/src/relarena/userdb/predict.py +++ b/packages/relarena/src/relarena/core/userdb/predict.py @@ -18,10 +18,10 @@ import pandas as pd from relbench.base import Database, EntityTask, Table -from relarena.dataset import _copy_timeless_tables +from relarena.core.dataset import _copy_timeless_tables if TYPE_CHECKING: - from relarena.model import RelArenaModel + from relarena.core.model import RelArenaModel #: Entity universe selector: `"all"` (every entity existing at the anchor) or an #: explicit collection of entity ids. diff --git a/packages/relarena/src/relarena/userdb/query.py b/packages/relarena/src/relarena/core/userdb/query.py similarity index 92% rename from packages/relarena/src/relarena/userdb/query.py rename to packages/relarena/src/relarena/core/userdb/query.py index 3538363..ebfd8e1 100644 --- a/packages/relarena/src/relarena/userdb/query.py +++ b/packages/relarena/src/relarena/core/userdb/query.py @@ -1,7 +1,7 @@ """Python façade for the Relational Predictive Interface (RPI). `PredictiveQuery` wraps the pieces validated separately — `UserEntityTask` -(SQL → labels), `RelBenchDatasetTask.from_objects` (splits), and +(SQL → labels), `TaskSource.from_objects` (splits), and `predict_at` (label-less inference) — into a single object so the common flow reads as `PredictiveQuery(spec).fit(model).predict()`. @@ -16,30 +16,31 @@ import warnings from dataclasses import dataclass +from importlib import import_module from pathlib import Path import pandas as pd import yaml from relbench.base import EntityTask -from relarena.cache import CacheConfig, resolve_cache_config -from relarena.dataset import RelBenchDatasetTask, concat_tables -from relarena.identity import ( +from relarena.core.cache import CacheConfig, resolve_cache_config +from relarena.core.dataset import TaskSource, concat_tables +from relarena.core.identity import ( RunIdentity, database_schema_fingerprint, task_spec_fingerprint, ) -from relarena.model import RelArenaModel -from relarena.registry import registry -from relarena.runner import select_best -from relarena.search_space import TaskStats, resolve_search_space -from relarena.system import RelArenaSystem -from relarena.tuner import tune as run_tuning -from relarena.userdb._schema import load_schema, validate -from relarena.userdb.ingest import DatabaseSpec, build_dataset -from relarena.userdb.predict import EntitySelector, predict_at -from relarena.userdb.spec import PredictiveTaskSpec -from relarena.userdb.task import UserEntityTask +from relarena.core.model import RelArenaModel +from relarena.core.registry import registry +from relarena.core.search_space import TaskStats, resolve_search_space +from relarena.core.selection import select_best +from relarena.core.system import RelArenaSystem +from relarena.core.tuner import tune as run_tuning +from relarena.core.userdb._schema import load_schema, validate +from relarena.core.userdb.ingest import DatabaseSpec, build_dataset +from relarena.core.userdb.predict import EntitySelector, predict_at +from relarena.core.userdb.spec import PredictiveTaskSpec +from relarena.core.userdb.task import UserEntityTask #: JSON Schema for a task YAML; the single source of truth for its accepted shape. _TASK_SCHEMA = load_schema("task.schema.json") @@ -74,7 +75,7 @@ def __init__( task_fingerprint=task_spec_fingerprint(task), data_version=data_version, ) - self._source = RelBenchDatasetTask.from_objects( + self._source = TaskSource.from_objects( dataset, UserEntityTask(dataset, task), dataset_name="user", @@ -112,12 +113,9 @@ def fit( `cache_dir` is a local directory that caches DFS features across tuning, the final fit, and later `predict`, useful for repeated runs on a large custom database. Omit it to fall back to `RELARENA_CACHE_DIR`, or to use - no persistent cache when that variable is unset. See `relarena.cache`. + no persistent cache when that variable is unset. See `relarena.core.cache`. """ - # Local import: importing the model package runs every built-in model's - # registration, pulling in heavy/optional deps (torch via tabpfn, lightgbm) - # at import time; kept out of `import relarena.userdb`, needed only here. - import relarena.models # noqa: F401 + import_module("relarena.models") cache = resolve_cache_config(cache_dir, on_miss="fill") self._warn_schema_only_cache(cache) @@ -186,7 +184,7 @@ def precompute_cache(self, cache_dir: str | Path) -> str | Path: its own anchor features on top. The resulting artifacts are shared by all DFS models. """ - from relarena.featurization.warm_cache import warm_dfs_cache + from relarena.core.featurization.warm_cache import warm_dfs_cache cache = resolve_cache_config(cache_dir, on_miss="fill") self._warn_schema_only_cache(cache) diff --git a/packages/relarena/src/relarena/userdb/spec.py b/packages/relarena/src/relarena/core/userdb/spec.py similarity index 98% rename from packages/relarena/src/relarena/userdb/spec.py rename to packages/relarena/src/relarena/core/userdb/spec.py index b4554e6..8ec62e6 100644 --- a/packages/relarena/src/relarena/userdb/spec.py +++ b/packages/relarena/src/relarena/core/userdb/spec.py @@ -7,7 +7,7 @@ import pandas as pd from relbench.base import TaskType -from relarena.userdb.predict import EntitySelector +from relarena.core.userdb.predict import EntitySelector #: String aliases accepted for `task_type`, restricted to the entity task types #: RelArena supports. diff --git a/packages/relarena/src/relarena/userdb/task.py b/packages/relarena/src/relarena/core/userdb/task.py similarity index 98% rename from packages/relarena/src/relarena/userdb/task.py rename to packages/relarena/src/relarena/core/userdb/task.py index 10f9109..628c997 100644 --- a/packages/relarena/src/relarena/userdb/task.py +++ b/packages/relarena/src/relarena/core/userdb/task.py @@ -7,7 +7,7 @@ from relbench.base import Database, Dataset, EntityTask, Table, TaskType from relbench.metrics import accuracy, average_precision, f1, mae, r2, rmse, roc_auc -from relarena.userdb.spec import PredictiveTaskSpec +from relarena.core.userdb.spec import PredictiveTaskSpec #: RelBench's standard metric lists per entity task type (primary metric last). _METRICS_BY_TASK_TYPE = { diff --git a/packages/relarena/src/relarena/userdb/task.schema.json b/packages/relarena/src/relarena/core/userdb/task.schema.json similarity index 100% rename from packages/relarena/src/relarena/userdb/task.schema.json rename to packages/relarena/src/relarena/core/userdb/task.schema.json diff --git a/packages/relarena/src/relarena/dataset.py b/packages/relarena/src/relarena/dataset.py index f3dd401..5511073 100644 --- a/packages/relarena/src/relarena/dataset.py +++ b/packages/relarena/src/relarena/dataset.py @@ -15,25 +15,24 @@ the censoring in `RelBenchDatasetTask` rather than leaving `db.upto(...)` calls scattered across the runner and tuner. -Currently RelBench is the only data source, so `RelBenchDatasetTask` is a -concrete class. If/when a second relational data source is integrated, extract a -`TaskSource` `Protocol`/ABC with `task` / `metric` / `inner_split` / -`outer_split` and have this class implement it — the interface is already -shaped for that, so no caller would need to change. +`RelBenchDatasetTask` supplies named benchmark loading and canonical column +cleaning. Core's `TaskSource` implements the shared supplied-object path and +split construction used by this adapter and standalone predictive queries. """ from __future__ import annotations -from dataclasses import dataclass -from typing import Callable, ClassVar, Final +from typing import Final -import pandas as pd -from relbench.base import Database, Dataset, EntityTask, Table +from relbench.base import Database, EntityTask, Table from relbench.datasets import get_dataset from relbench.tasks import get_task -from relarena.identity import RunIdentity, relbench_run_identity -from relarena.metrics import primary_metric +from relarena.core.dataset import ( + TaskSource, + clean_database, +) +from relarena.identity import relbench_run_identity #: rel-ratebeer per-user aggregates computed over the *entire* rating history, so a #: row can encode information from after its own timestamp. @@ -62,26 +61,7 @@ def drop_noncanonical_columns(db: Database, dataset_name: str) -> Database: Affected tables are rebuilt, not mutated: relbench `lru_cache`-s `get_db()` and `relarena.checksums` must keep seeing the raw upstream data. """ - leaking = TIME_LEAKING_COLUMNS.get(dataset_name, {}) - table_dict = dict(db.table_dict) - changed = False - for table_name, table in db.table_dict.items(): - df = table.df - # "Unnamed: N" columns are leftover pandas row indices from a CSV export. - drop = {c for c in df.columns if c.startswith("Unnamed:")} - if not df.empty: - drop |= {c for c in df.columns if df[c].isna().all()} - drop |= {c for c in leaking.get(table_name, ()) if c in df.columns} - if not drop: - continue - changed = True - table_dict[table_name] = Table( - df=df.drop(columns=list(drop)), - fkey_col_to_pkey_table=dict(table.fkey_col_to_pkey_table), - pkey_col=table.pkey_col, - time_col=table.time_col, - ) - return Database(table_dict) if changed else db + return clean_database(db, TIME_LEAKING_COLUMNS.get(dataset_name, {})) def drop_noncanonical_task_columns( @@ -118,137 +98,8 @@ def drop_noncanonical_task_columns( ) -def concat_tables(a: Table, b: Table) -> Table: - """Concatenate two task label tables (e.g. train + val) into one, for refitting. - - Assumes both share the same schema (entity/foreign keys, time and target - columns) — which holds for the splits of a single task. Inputs are untouched. - """ - return Table( - df=pd.concat([a.df, b.df], ignore_index=True), - fkey_col_to_pkey_table=dict(a.fkey_col_to_pkey_table), - pkey_col=a.pkey_col, - time_col=a.time_col, - ) - - -@dataclass(frozen=True) -class Split: - """Fields common to one fit→evaluate phase of nested temporal validation. - - Bundles the censored database with the label tables a phase needs, so a model - can never accidentally see data past its phase's cutoff. Not instantiated - directly: a phase is always an `InnerSplit` or an `OuterSplit`, - which differ in *how the predictions are scored* (see those classes). The - harness decides how each field is used (it does not pass `eval_table` to - `fit` on the outer split, for instance) — see - `relarena.tuner.run_trial` and `relarena.tuner.refit_and_evaluate`. - """ - - #: Database censored at `cutoff` — the only DB the model may read. - db_state: Database - #: The censoring cutoff (`val_timestamp` for inner, `test_timestamp` for outer). - cutoff: pd.Timestamp - #: The label table the model trains on. - train_table: Table - #: The table predictions are made on. For the outer split this is the *masked* - #: test table — RelBench's `get_table("test")` strips the target column so the - #: model can predict on the test entities/timestamps without seeing the answers - #: (leakage prevention); the labels live only inside RelBench. The inner split's - #: val table is unmasked, but the model is still expected to use it only as a - #: prediction/early-stopping set, never to read its labels for training. - eval_table: Table - - -@dataclass(frozen=True) -class InnerSplit(Split): - """Tuning phase: fit `train` → score `val`. - - Scoring needs the val labels handed in explicitly. RelBench's - `EntityTask.evaluate(pred, target_table=None)` treats `None` as "score - against the *test* table" — there is no shortcut for val — so to score on val - we must pass the val target table ourselves (`eval_target`). The val - table's labels are not hidden, so this is just the val table itself. - """ - - name: ClassVar[str] = "inner" - #: Val labels to score `eval_table` predictions against. Passed to - #: `EntityTask.evaluate(pred, target_table=eval_target)` — required because - #: `evaluate`'s `None` default would score against test, not val. - eval_target: Table - - -@dataclass(frozen=True) -class OuterSplit(Split): - """Final phase: fit the selected config → score `test`. - - `train_table` is the train-only table and `val_table` the - val table, exposed separately so the harness can serve either final-fit regime: - refit on their union (`refit_on_full_data=True`), or train on train alone with - val as a held-out checkpoint/early-stopping set (`refit_on_full_data=False`). - - Deliberately carries *no* eval target. The test labels are hidden (the model's - `eval_table` is the masked test table), and we score by calling - `EntityTask.evaluate(pred, target_table=None)`, which makes RelBench load its - own held-out test labels. Not materializing those labels into this object keeps - the answer key out of every structure the harness passes around the model — so - test-label leakage is impossible by construction, not just by convention. - """ - - name: ClassVar[str] = "outer" - #: The val label table, held out from `train_table`. The harness - #: unions it with the train table to refit on full data, or passes it as the - #: monitoring set when a model trains on train alone. - val_table: Table - - -def _drop_dangling_seeds(table: Table, entity_col: str, num_entities: int) -> Table: - """Drop seeds whose entity is absent from a censored entity table. - - `entity_col` is the seed table's column of entity ids (each a row index into the - entity table); a value `>= num_entities` references an entity created after the - val cutoff, absent from the censored graph. This is the val-cutoff analog of - relbench's `EntityTask.filter_dangling_entities` (which only filters against the - test-censored `get_db` count). - """ - keep = table.df[entity_col] < num_entities - if keep.all(): - return table - return Table( - df=table.df[keep].reset_index(drop=True), - fkey_col_to_pkey_table=table.fkey_col_to_pkey_table, - pkey_col=table.pkey_col, - time_col=table.time_col, - ) - - -def _copy_timeless_tables(db: Database) -> None: - """Replace each timeless table's DataFrame with a copy, in place. - - `Table.upto` returns timeless tables (`time_col is None`) as the *same - object* held by the source DB, so a censored DB shares their DataFrames with - it. Any in-place mutation of the censored DB (e.g. the dangling-FK scrub in - `validate_and_correct_db`, `df.loc[mask, fkey] = None`) would otherwise - leak back into the source DB and corrupt later splits derived from it. - """ - for name, table in db.table_dict.items(): - if table.time_col is None: - db.table_dict[name] = Table( - df=table.df.copy(), - fkey_col_to_pkey_table=table.fkey_col_to_pkey_table, - pkey_col=table.pkey_col, - time_col=table.time_col, - ) - - -class RelBenchDatasetTask: - """Loads one RelBench `(dataset, task)` and hands out its censored splits. - - Isolates the rest of the package from RelBench's loading API and owns the - temporal-correctness logic (DB censoring + train/val concatenation). Loading - is heavy (a multi-GB download / DB build); discovering *which* tasks exist is - cheap and lives separately in `relarena.tasks.list_entity_tasks`. - """ +class RelBenchDatasetTask(TaskSource): + """Load canonical benchmark data and construct its temporal splits.""" def __init__( self, dataset_name: str, task_name: str, *, download: bool = True @@ -273,114 +124,13 @@ def __init__( for split in ("train", "val", "test") } - @classmethod - def from_objects( - cls, - dataset: Dataset, - task: EntityTask, - *, - dataset_name: str = "user", - task_name: str | None = None, - run_identity: RunIdentity | None = None, - ) -> RelBenchDatasetTask: - """Build from in-memory `dataset` / `task` objects instead of the registry. + _prepare_db = staticmethod(drop_noncanonical_columns) + _prepare_table = staticmethod(drop_noncanonical_task_columns) - Mirrors `__init__` exactly — same DB cleaning, table generation and - split logic — but takes an already-constructed RelBench `Dataset` and - `EntityTask` (e.g. a user task via - `UserEntityTask`) rather than fetching a named - `(dataset, task)` from RelBench. `dataset_name` only selects the - canonical-column policy in `drop_noncanonical_columns`: pass a known - RelBench name to reproduce its cleaning, or leave the default to apply just - the generic fully-NaN / `Unnamed` drop. - """ - self = cls.__new__(cls) - self.dataset_name = dataset_name - self.task_name = task_name if task_name is not None else type(task).__name__ - self._identity = run_identity or RunIdentity( - dataset_name, None, self.task_name, None - ) - self._dataset = dataset - self._task = task - self._db = drop_noncanonical_columns(self._dataset.get_db(), dataset_name) - self._tables = { - split: drop_noncanonical_task_columns( - self._task, self._task.get_table(split), self.dataset_name - ) - for split in ("train", "val", "test") - } - return self - - def run_identity(self, phase: str | None = None) -> RunIdentity: - """Return optional source metadata scoped to one run phase.""" - return self._identity.for_phase(phase) - - @property - def task(self) -> EntityTask: - """The underlying RelBench task (defines the target, metrics, `evaluate`).""" - return self._task - - @property - def metric(self) -> Callable[..., float]: - """The primary metric this task is tuned and selected on (by task type).""" - return primary_metric(self._task) - def inner_split(self) -> InnerSplit: - """Tuning split: fit `train` → score `val`, DB frozen at `val_timestamp`. - - Censoring the DB at the val cutoff makes validation features *frozen at - their cutoff*, mirroring how test features are frozen at the test cutoff — - see `docs/temporal-validation.md` for why this matters for aggregating - models. `eval_target` is the val table itself (its labels are not hidden). - """ - # The val-cutoff removes pkey rows dated after val_timestamp; any foreign key - # pointing *forward* in time to such a row (e.g. an attendance row referencing - # a later event) is now dangling. get_db() scrubs dangling FKs for the test - # cutoff via validate_and_correct_db, but this extra censor needs the same - # correction or make_pkey_fkey_graph asserts on the out-of-range index. - inner_db = self._db.upto(self._dataset.val_timestamp) - # `validate_and_correct_db` scrubs dangling FKs in place, so decouple the - # timeless tables `upto` shares with `self._db` first (see the helper), - # otherwise the scrub leaks back into `self._db` and corrupts outer_split. - _copy_timeless_tables(inner_db) - self._dataset.validate_and_correct_db(inner_db) - # The same val-cutoff can shrink the entity table below seeds that reference - # entities created after val_timestamp — relbench's get_table only filtered - # seeds against the *test* count. Drop those seeds here (the val-cutoff analog - # of EntityTask.filter_dangling_entities) so a graph sampler can't index past - # the censored entity table. The val table is filtered once and used for both - # eval_table and eval_target, so scored predictions still align with targets. - n_entities = len(inner_db.table_dict[self._task.entity_table].df) - train_table = _drop_dangling_seeds( - self._tables["train"], self._task.entity_col, n_entities - ) - val_table = _drop_dangling_seeds( - self._tables["val"], self._task.entity_col, n_entities - ) - return InnerSplit( - db_state=inner_db, - cutoff=self._dataset.val_timestamp, - train_table=train_table, - eval_table=val_table, - eval_target=val_table, - ) - - def outer_split(self) -> OuterSplit: - """Final split: score `test` after fitting, DB at `test_timestamp`. - - Exposes the train and val tables separately (rather than pre-unioning them) - so the harness can serve either final-fit regime — refit on their union, or - train on train alone with val held out. No `eval_target`: the test labels - are hidden, so scoring goes through `EntityTask.evaluate(pred, - target_table=None)` and RelBench supplies them. - """ - return OuterSplit( - # No-op re-censor: get_db() already censors the DB at test_timestamp - # (upto_test_timestamp=True), so this .upto() drops nothing. Kept for - # readability/symmetry with inner_split's explicit val-cutoff censoring. - db_state=self._db.upto(self._dataset.test_timestamp), - cutoff=self._dataset.test_timestamp, - train_table=self._tables["train"], - val_table=self._tables["val"], - eval_table=self._tables["test"], - ) +__all__ = [ + "RelBenchDatasetTask", + "drop_noncanonical_columns", + "drop_noncanonical_task_columns", + "TIME_LEAKING_COLUMNS", +] diff --git a/packages/relarena/src/relarena/evaluation/leaderboard.py b/packages/relarena/src/relarena/evaluation/leaderboard.py index 5e222c1..86faa63 100644 --- a/packages/relarena/src/relarena/evaluation/leaderboard.py +++ b/packages/relarena/src/relarena/evaluation/leaderboard.py @@ -17,11 +17,13 @@ from __future__ import annotations import logging +from importlib import import_module import pandas as pd +from relarena.core.metrics import to_metric_error +from relarena.core.registry import registry from relarena.evaluation.subsets import TaskMask, apply_subset -from relarena.metrics import to_metric_error logger = logging.getLogger(__name__) @@ -113,12 +115,7 @@ def method_kind(model: str) -> str: package rather than isolating one model (see `RelArenaSystem`). Unregistered names (reference baselines, retired methods) rank as models. """ - # Built-ins register on this import, which a leaderboard-only caller (load - # a results CSV, rank it) has no other reason to have made. Without it the - # registry is empty, every lookup falls through to "model", and a system - # silently joins the models-only board. - import relarena.models # noqa: F401 - from relarena.registry import registry + import_module("relarena.models") try: return registry.kind(model) diff --git a/packages/relarena/src/relarena/featurization/__init__.py b/packages/relarena/src/relarena/featurization/__init__.py index ec838b0..5e6ee57 100644 --- a/packages/relarena/src/relarena/featurization/__init__.py +++ b/packages/relarena/src/relarena/featurization/__init__.py @@ -1,15 +1,5 @@ -"""Featurization: relational database -> a flat feature table (for tabular models). +"""Entity-only featurization for benchmark baseline models.""" -Sub-modules implement different recipes: - * `entity` — entity-only (the RelBench LightGBM recipe); - * `dfs` — multi-hop Deep Feature Synthesis (the RDBLearn recipe), with a - depth cache. - -Both expose a `build_*_features(...) -> (features_df, categorical_columns)` -function; shared column typing lives in `_columns`. -""" - -from relarena.featurization.dfs import DFS_MAX_DEPTH, build_dfs_features from relarena.featurization.entity import build_entity_features -__all__ = ["DFS_MAX_DEPTH", "build_entity_features", "build_dfs_features"] +__all__ = ["build_entity_features"] diff --git a/packages/relarena/src/relarena/featurization/entity.py b/packages/relarena/src/relarena/featurization/entity.py index 42ff078..5a14f2f 100644 --- a/packages/relarena/src/relarena/featurization/entity.py +++ b/packages/relarena/src/relarena/featurization/entity.py @@ -15,7 +15,7 @@ import pandas as pd from relbench.base import Database, EntityTask, Table -from relarena.featurization._columns import type_columns +from relarena.core.featurization.columns import type_columns def build_entity_features( diff --git a/packages/relarena/src/relarena/featurization/warm_cache.py b/packages/relarena/src/relarena/featurization/warm_cache.py index fd89100..e8687e5 100644 --- a/packages/relarena/src/relarena/featurization/warm_cache.py +++ b/packages/relarena/src/relarena/featurization/warm_cache.py @@ -1,52 +1,14 @@ -"""Public shared-DFS cache warmer, runnable with `python -m`.""" +"""Warm shared DFS caches for named benchmark tasks with `python -m`.""" from __future__ import annotations import argparse from pathlib import Path -from relarena.cache import CacheConfig, resolve_cache_config -from relarena.dataset import RelBenchDatasetTask, concat_tables -from relarena.featurization.dfs import DFS_MAX_DEPTH, build_dfs_features - - -def warm_dfs_cache( - source: RelBenchDatasetTask, - cache: CacheConfig, - *, - max_depth: int = DFS_MAX_DEPTH, -) -> None: - """Fill shared DFS matrices for tuning and both final-fit history regimes.""" - if cache.directory is None or cache.on_miss != "fill": - raise ValueError("DFS warming needs CacheConfig(directory, on_miss='fill')") - inner, outer = source.inner_split(), source.outer_split() - full_outer_history = concat_tables(outer.train_table, outer.val_table) - phases = ( - ("inner", inner.db_state, inner.train_table, inner.eval_table), - # RDBLearn follows its published train-only final-fit protocol, whereas - # TabPFN-Rel refits on train+val. The actual history input - # belongs in the key, so warm both canonical regimes without model dispatch. - ("outer", outer.db_state, outer.train_table, outer.eval_table), - ( - "outer", - outer.db_state, - full_outer_history, - outer.eval_table, - ), - ) - for phase, db, history, evaluation in phases: - identity = source.run_identity(phase) - for anchors in (history, evaluation): - build_dfs_features( - source.task, - db, - anchors, - depth=max_depth, - max_depth=max_depth, - history_table=history if source.task.time_col else None, - cache=cache, - run_identity=identity, - ) +from relarena.core.cache import resolve_cache_config +from relarena.core.featurization.dfs import DFS_MAX_DEPTH +from relarena.core.featurization.warm_cache import warm_dfs_cache +from relarena.dataset import RelBenchDatasetTask def main(argv: list[str] | None = None) -> int: diff --git a/packages/relarena/src/relarena/identity.py b/packages/relarena/src/relarena/identity.py index f5f5479..c062783 100644 --- a/packages/relarena/src/relarena/identity.py +++ b/packages/relarena/src/relarena/identity.py @@ -1,72 +1,17 @@ -"""Source identity passed explicitly to preprocessing-owned cache-key factories. +"""Recorded benchmark identities and shared source metadata.""" -`RunIdentity` describes the data and execution phase from which preprocessing -artifacts are produced. Entrypoints construct it, the runner passes it through -model construction, and each preprocessing module decides which fields belong -in its own artifact key. - -A run identity is metadata, not a complete cache key and not a central key -policy. It intentionally contains more information than every artifact needs. -For example, task-dependent DFS and RelGT artifacts use task identity, while a -RelGNN graph may ignore it because the graph depends only on the censored -database. Preprocessors remain responsible for selecting actual dependencies -and adding their own algorithm versions. - -The readable dataset and task names provide namespaces. Their fingerprints -distinguish different underlying data or task definitions without placing -absolute paths or expensive row hashes in keys. `phase` distinguishes -protocol views such as `inner`, `outer`, and `predict` when censoring or -inputs differ. `data_version` is an optional caller-supplied discriminator -for data changes that cheap fingerprints cannot observe. - -RelBench identities use checked-in dataset and task checksums. Predictive-query -identities use a database-schema fingerprint and a task-specification -fingerprint; because the schema fingerprint deliberately ignores row contents, -callers using persistent caches should provide `data_version` when those -contents may change without a schema change. - -Each preprocessing owner decides whether persistent use requires complete -fingerprints or whether it can derive a safe fallback from its actual inputs. -Unconfigured direct callers may omit identity and compute in private scratch. -Cache directories, miss policies, model names, serialization formats, and -preprocessing algorithm versions do not belong in this object. -""" - -from __future__ import annotations - -import hashlib import json -from dataclasses import dataclass, replace from functools import cache from pathlib import Path -from typing import Any -from relbench.base import Database +from relarena.core.identity import ( + RunIdentity, + metadata_fingerprint, +) CHECKSUMS_PATH = Path(__file__).with_name("checksums") / "relbench_v1_checksums.json" -@dataclass(frozen=True) -class RunIdentity: - """Readable source metadata that a preprocessor may use or ignore.""" - - dataset: str - dataset_fingerprint: str | None - task: str | None - task_fingerprint: str | None - data_version: str | None = None - phase: str | None = None - - def for_phase(self, phase: str | None) -> RunIdentity: - """Return the same identity scoped to one execution phase.""" - return replace(self, phase=phase) - - -def _digest(value: Any) -> str: - encoded = json.dumps(value, sort_keys=True, separators=(",", ":")).encode() - return hashlib.blake2s(encoded, digest_size=8).hexdigest() - - @cache def _recorded_checksums(path: Path = CHECKSUMS_PATH) -> dict[str, dict[str, int]]: return json.loads(path.read_text()) @@ -81,40 +26,12 @@ def relbench_run_identity(dataset: str, task: str) -> RunIdentity: task_values = { key: value for key, value in record.items() if not key.endswith("_db") } - return RunIdentity(dataset, dataset_fingerprint, task, _digest(task_values)) + return RunIdentity( + dataset, dataset_fingerprint, task, metadata_fingerprint(task_values) + ) -def database_schema_fingerprint(db: Database) -> str: - """Fingerprint a user database's relational schema without hashing its rows.""" - schema = { - name: { - "columns": [ - (str(column), str(table.df[column].dtype)) for column in table.df - ], - "fkeys": sorted(table.fkey_col_to_pkey_table.items()), - "pkey": table.pkey_col, - "time": table.time_col, - } - for name, table in sorted(db.table_dict.items()) - } - return _digest(schema) - - -def task_spec_fingerprint(task: Any) -> str: - """Fingerprint training semantics of a user predictive-task specification.""" - fields = { - name: str(getattr(task, name)) - for name in ( - "entity_table", - "entity_col", - "time_col", - "target_col", - "task_type", - "timedelta", - "query", - "val_timestamp", - "test_timestamp", - "num_eval_timestamps", - ) - } - return _digest(fields) +__all__ = [ + "RunIdentity", + "relbench_run_identity", +] diff --git a/packages/relarena/src/relarena/models/__init__.py b/packages/relarena/src/relarena/models/__init__.py index 307c666..a84edd8 100644 --- a/packages/relarena/src/relarena/models/__init__.py +++ b/packages/relarena/src/relarena/models/__init__.py @@ -1,49 +1,22 @@ -"""Concrete model wrappers. +"""Discover baseline modules whose decorators register models and systems. -Importing this package registers every built-in model into the global registry: -each wrapper module is imported in turn, and the `@register_model` decorator on -its class does the registering. A wrapper implements the -`relarena.model.RelArenaModel` contract (`fit` / `predict`) and declares a -`relarena.search_space.SearchSpace`; see `lightgbm.py` for a worked example. - -Adding a model needs no edit here — drop a module or package alongside the -existing ones and it is picked up. Names starting with an underscore are skipped, -which keeps `_shared` and the vendored packages out of the scan. - -Reach a model through the registry, by name; the wrapper classes and preprocessing -warmers are not exported here. Each cache owner exposes its runnable warmer from its -own `warm_cache` module, so adding one never changes this discovery package. +Model implementations and their search spaces live together in each package. +Importing this package loads each public child module or package; shared helpers are +excluded. Optional backends load during model execution. """ import importlib -import logging import pkgutil -logger = logging.getLogger(__name__) - def _register_builtin_models() -> None: - """Import every wrapper module so its `@register_model` decorator runs. + """Import baseline modules and packages. - A wrapper whose *third-party* dependency is absent is skipped: the per-model - extras are optional, so `tabpfn-rel` without the DFS deps is a normal install - rather than a broken one. A missing `relarena` module, or any other import - error, is a defect and propagates — swallowing those is how a model goes - quietly missing from the registry. + Propagate failures from their implementations. """ - for _finder, name, _is_pkg in pkgutil.iter_modules(__path__): - if name.startswith("_"): - continue - try: + for _finder, name, _is_package in pkgutil.iter_modules(__path__): + if not name.startswith("_"): importlib.import_module(f"{__name__}.{name}") - except ModuleNotFoundError as exc: - if exc.name is not None and exc.name.split(".")[0] == "relarena": - raise - logger.info( - "Skipping model %r: optional dependency %r is not installed.", - name, - exc.name, - ) _register_builtin_models() diff --git a/packages/relarena/src/relarena/models/_shared/__init__.py b/packages/relarena/src/relarena/models/_shared/__init__.py index a1b1d45..a123479 100644 --- a/packages/relarena/src/relarena/models/_shared/__init__.py +++ b/packages/relarena/src/relarena/models/_shared/__init__.py @@ -1,7 +1,5 @@ -"""Code shared between model wrappers. +"""Helpers shared between benchmark model implementations. -Not models: the registry walks `models/*` and skips this package. The layout -encodes who shares what — a family subpackage (`gbdt`, `tfm`, `gnn`) holds code -shared *within* that family, while a module at this level is shared *across* -families. Import the submodules directly; these are internals and may move. +Family subpackages contain code shared within GNN or GBDT models. Modules at +this level serve multiple families. These helpers are private to RelArena. """ diff --git a/packages/relarena/src/relarena/models/_shared/tfm/__init__.py b/packages/relarena/src/relarena/models/_shared/tfm/__init__.py deleted file mode 100644 index c89c634..0000000 --- a/packages/relarena/src/relarena/models/_shared/tfm/__init__.py +++ /dev/null @@ -1,6 +0,0 @@ -"""Shared internals of the tabular-foundation-model baselines. - -Used by `rdblearn` and the `tabpfn-rel` variants. - -`tfm` carries the TFM registry: the single seam for adding a new foundation model. -""" diff --git a/packages/relarena/src/relarena/models/_shared/tfm/tfm.py b/packages/relarena/src/relarena/models/_shared/tfm/tfm.py deleted file mode 100644 index 1cba325..0000000 --- a/packages/relarena/src/relarena/models/_shared/tfm/tfm.py +++ /dev/null @@ -1,341 +0,0 @@ -"""Shared tabular-foundation-model (TFM) estimator core for the `rdblearn` baseline. - -The `rdblearn` model feeds Deep Feature Synthesis features (`featurization/dfs.py`) -to a tabular foundation model and tunes over both DFS depth and *which* TFM is used. -This module owns the parts shared across TFMs and feature sources: - - * a small TFM **registry** mapping a name -> how to build its classifier/regressor, - which task types it supports, and whether its backing package is importable. This - is the single seam for adding TFMs — currently TabPFN v2 / v2.5 / v3 (local and - hosted-API); a new TFM drops in here as one more entry; - * seeded downsampling of the training set to each TFM's context-size cap - (`TFMSpec.max_train_samples`); - * fit, and predict that returns the shape `EntityTask.evaluate` expects (the - sklearn-output reshaping is shared via `predict_contract`). - -Categorical columns are handled by the TFM natively: we pass the feature frame as a -DataFrame and let TabPFN auto-detect categoricals (its preprocessing treats a column as -categorical below a cardinality threshold), rather than pre-encoding them (unlike -RDBLearn's `SafeLabelEncoder`) — hand-coding categoricals to a float matrix would only -hide them from TabPFN's categorical handling. We deliberately do *not* pass the -DFS-flagged columns as `categorical_features_indices`; tuning relational-data -preprocessing (incl. forcing categorical handling for moderate-cardinality columns) is -deliberately not done yet. -""" - -from __future__ import annotations - -import inspect -from dataclasses import dataclass -from typing import Any, Callable, Protocol - -import numpy as np -import pandas as pd -import torch -from relbench.base import TaskType - -from relarena.models._shared.predict_contract import predict_to_contract - - -class SklearnClassifier(Protocol): - """Minimal sklearn-classifier surface relarena uses (TabPFNClassifier-like).""" - - classes_: np.ndarray - - def fit(self, X: pd.DataFrame, y: np.ndarray) -> Any: ... - def predict_proba(self, X: pd.DataFrame) -> np.ndarray: ... - - -class SklearnRegressor(Protocol): - """Minimal sklearn-regressor surface relarena uses (TabPFNRegressor-like).""" - - def fit(self, X: pd.DataFrame, y: np.ndarray) -> Any: ... - def predict(self, X: pd.DataFrame) -> np.ndarray: ... - - -# -- TFM registry ------------------------------------------------------------ - - -@dataclass(frozen=True) -class TFMSpec: - """How to build one tabular foundation model. - - Every TFM is assumed to support all entity task types, so there is no per-TFM - task-type gating. `make_classifier` / `make_regressor` take keyword overrides - (`device`, `seed`, ...) and return an estimator satisfying - `SklearnClassifier` / `SklearnRegressor` respectively. - `max_train_samples` is this TFM's training-row cap before fitting — its - supported context size (TabPFN v2 ~10k, v2.5 ~50k) — applied by `fit_tfm`. - `supports_text` marks estimators that handle raw text columns themselves. - """ - - make_classifier: Callable[..., SklearnClassifier] - make_regressor: Callable[..., SklearnRegressor] - max_train_samples: int - supports_text: bool = False - - -def _make_tabpfn( - version: str, - *, - regression: bool, - device: Any, - seed: int, - **overrides: Any, -) -> Any: - """Build a TabPFN estimator pinned to a version via `create_default_for_version`. - - The bare TabPFN constructor now defaults to v3; `create_default_for_version` - selects the right checkpoint + version-appropriate defaults for v2 / v2.5, and - `**overrides` (device, random_state, ignore_pretraining_limits, ...) pass through - to the constructor. - - Lazy import — tabpfn lives in the rdblearn extra, and it is the one dependency - under the Prior Labs License rather than a plain permissive one, so a core - install stays clear of its attribution obligation (see `docs/licensing.md`). - """ - from tabpfn import TabPFNClassifier, TabPFNRegressor - from tabpfn.constants import ModelVersion - - model_version = { - "v2": ModelVersion.V2, - "v2.5": ModelVersion.V2_5, - "v3": ModelVersion.V3, - }[version] - estimator_cls = TabPFNRegressor if regression else TabPFNClassifier - return estimator_cls.create_default_for_version( - model_version, - device=device, - random_state=seed, - ignore_pretraining_limits=True, - **overrides, - ) - - -def _sanitize_api_overrides(overrides: dict[str, Any]) -> dict[str, Any]: - """Convert numpy index arrays in `SUBSAMPLE_SAMPLES` to plain int lists. - - tabpfn_client pydantic-serializes the estimator config into the request body, - which rejects numpy arrays — the pool contexts pass per-estimator context - indices as arrays. The local TabPFN consumes arrays natively, so the - conversion is scoped to the API path. - """ - inference_config = overrides.get("inference_config") - if inference_config is None: - return overrides - subsample = inference_config.get("SUBSAMPLE_SAMPLES") - if not isinstance(subsample, list): - return overrides - return { - **overrides, - "inference_config": { - **inference_config, - "SUBSAMPLE_SAMPLES": [ - e.tolist() if isinstance(e, np.ndarray) else e for e in subsample - ], - }, - } - - -def _make_tabpfn_api( - *, - regression: bool, - device: Any, - seed: int, - **overrides: Any, -) -> Any: - """Build a TabPFN API-client estimator pinned to the v3 model. - - Lazy import — tabpfn_client lives in the tabpfn-rel-api extra. Fit and predict - run server-side, so device is ignored and raw text columns are handled by the - API. - """ - from tabpfn_client import TabPFNClassifier as ApiClassifier - from tabpfn_client import TabPFNRegressor as ApiRegressor - - del device - overrides = _sanitize_api_overrides(overrides) - estimator_cls = ApiRegressor if regression else ApiClassifier - return estimator_cls( - model_path="v3_default", - random_state=seed, - ignore_pretraining_limits=True, - **overrides, - ) - - -def _tabpfn_spec(version: str, max_train_samples: int) -> TFMSpec: - return TFMSpec( - make_classifier=lambda **kw: _make_tabpfn(version, regression=False, **kw), - make_regressor=lambda **kw: _make_tabpfn(version, regression=True, **kw), - max_train_samples=max_train_samples, - ) - - -#: Name -> spec. The single seam for adding TFMs. -#: `max_train_samples` is each backend's fit limit. The RDBLearn paper -#: (arXiv:2602.18495) runs TabPFN v2, v2.5 and LimiX all under a **10k** limit, -#: downsampling above it — so those entries use 10k. `tabpfn-v3` is the -#: `tabpfn-rel-local` -#: backend (selected by its config, not swept by `rdblearn`); 100k = the context size -#: the reference sweeps ran it at, and `n_preprocessing_jobs=-1` parallelizes its -#: heavier preprocessing. `tabpfn-v3-api` is the same v3 model served by the TabPFN -#: API (no GPU needed); 100k matches the local entry and sits well inside the API's v3 -#: train limits, and the API handles raw text columns server-side. -TFM_REGISTRY: dict[str, TFMSpec] = { - "tabpfn-v2": _tabpfn_spec("v2", max_train_samples=10_000), - "tabpfn-v2.5": _tabpfn_spec("v2.5", max_train_samples=10_000), - "tabpfn-v3": TFMSpec( - make_classifier=lambda **kw: _make_tabpfn( - "v3", regression=False, n_preprocessing_jobs=-1, **kw - ), - make_regressor=lambda **kw: _make_tabpfn( - "v3", regression=True, n_preprocessing_jobs=-1, **kw - ), - max_train_samples=100_000, - ), - "tabpfn-v3-api": TFMSpec( - make_classifier=lambda **kw: _make_tabpfn_api(regression=False, **kw), - make_regressor=lambda **kw: _make_tabpfn_api(regression=True, **kw), - max_train_samples=100_000, - supports_text=True, - ), -} - - -def default_device() -> str: - """Return `"cuda"` if a GPU is visible to torch, else `"cpu"`.""" - return "cuda" if torch.cuda.is_available() else "cpu" - - -# -- downsampling ------------------------------------------------------------ - - -def _downsample_indices( - y: np.ndarray, task_type: TaskType, max_samples: int, rng: np.random.Generator -) -> np.ndarray: - """Indices of a <= `max_samples` subset of rows (no-op when already small). - - Regression: a uniform random subset. Classification: keep at least one row per - class, then fill the remaining budget uniformly at random. Seeded via `rng` for - reproducibility. Adapted from RDBLearn's `_downsample` (non-stratified path). - """ - n = len(y) - if n <= max_samples: - return np.arange(n) - - if task_type == TaskType.REGRESSION: - return rng.choice(n, max_samples, replace=False) - - selected: list[int] = [] - for label in np.unique(y): - class_idx = np.where(y == label)[0] - selected.append(int(rng.choice(class_idx, 1)[0])) - selected = list(dict.fromkeys(selected)) # de-dup (one per class) - - remaining = max_samples - len(selected) - if remaining > 0: - mask = np.ones(n, dtype=bool) - mask[selected] = False - eligible = np.where(mask)[0] - extra = rng.choice(eligible, min(remaining, len(eligible)), replace=False) - out = np.concatenate([np.array(selected, dtype=int), extra]) - else: - out = np.array(selected[:max_samples], dtype=int) - rng.shuffle(out) - return out - - -# -- fit / predict ----------------------------------------------------------- - - -@dataclass -class FittedTFM: - """A fitted TFM plus the state needed to score val/test consistently.""" - - estimator: Any - feature_cols: list[str] - task_type: TaskType - max_predict_samples: int | None = None - - -def fit_tfm( - df: pd.DataFrame, - y: pd.Series, - task_type: TaskType, - *, - tfm: str, - seed: int, - device: Any = None, - max_train_samples: int | None = None, - max_predict_samples: int | None = None, - overrides: dict[str, Any] | None = None, -) -> FittedTFM: - """Downsample `df` and fit the named TFM on it. - - `df` is the already-typed feature frame from `build_dfs_features` (numeric - floats + object categoricals); TabPFN auto-detects categoricals from it (see the - module docstring) — we do not pass `categorical_features_indices`. The training - rows are capped (seeded) at `max_train_samples` if given, else the TFM's own - context cap (`spec.max_train_samples`); `overrides` are additional - estimator-constructor arguments. `max_predict_samples` is an explicit - caller-owned cap on rows per estimator prediction call; ordinary TFM callers - leave it unset. - """ - spec = TFM_REGISTRY[tfm] - if device is None: - device = default_device() - cap = max_train_samples if max_train_samples is not None else spec.max_train_samples - rng = np.random.default_rng(seed) - - feature_cols = list(df.columns) - - y_arr = y.to_numpy() - idx = _downsample_indices(y_arr, task_type, cap, rng) - X = df.iloc[idx] - y_arr = y_arr[idx] - - kwargs = dict(device=device, seed=seed, **(overrides or {})) - if task_type == TaskType.REGRESSION: - estimator = spec.make_regressor(**kwargs) - y_arr = y_arr.astype(float) - else: - estimator = spec.make_classifier(**kwargs) - estimator.fit(X, y_arr) - - return FittedTFM(estimator, feature_cols, task_type, max_predict_samples) - - -def _predict_tfm_frame(fitted: FittedTFM, frame: pd.DataFrame) -> np.ndarray: - if fitted.task_type == TaskType.REGRESSION: - predict = fitted.estimator.predict - params = inspect.signature(predict).parameters.values() - if any( - p.name == "output_type" or p.kind is inspect.Parameter.VAR_KEYWORD - for p in params - ): - return np.asarray(predict(frame, output_type="median"), dtype=float) - return predict_to_contract(fitted.estimator, frame, fitted.task_type) - - -def predict_tfm(fitted: FittedTFM, df: pd.DataFrame) -> np.ndarray: - """Predict with a `FittedTFM` on a new feature frame. - - Reindexes to the training column order (so the TFM sees the same schema), then - delegates the sklearn-output -> evaluate-contract reshaping to - `relarena.models._shared.predict_contract.predict_to_contract`. - - Regression requests `output_type="median"` when the estimator supports it — - an explicit output_type parameter, or a **kwargs passthrough: the primary - regression metric is MAE, and the median is its optimal point prediction. - """ - X = df.reindex(columns=fitted.feature_cols) - batch_size = fitted.max_predict_samples - if batch_size is None or len(X) <= batch_size: - return _predict_tfm_frame(fitted, X) - return np.concatenate( - [ - _predict_tfm_frame(fitted, X.iloc[i : i + batch_size]) - for i in range(0, len(X), batch_size) - ] - ) diff --git a/packages/relarena/src/relarena/models/dummy/__init__.py b/packages/relarena/src/relarena/models/dummy/__init__.py index c533048..53b6b24 100644 --- a/packages/relarena/src/relarena/models/dummy/__init__.py +++ b/packages/relarena/src/relarena/models/dummy/__init__.py @@ -1,8 +1,4 @@ -"""Constant (optimal-constant) baselines. - -Importing this package registers `constant-global` (one global constant) and -`constant-per-entity` (each entity`s own constant). See `model`. -""" +"""Constant prediction baselines.""" from relarena.models.dummy.model import DummyBaseline, DummyPerEntityBaseline diff --git a/packages/relarena/src/relarena/models/dummy/model.py b/packages/relarena/src/relarena/models/dummy/model.py index 62d4517..dafd6a1 100644 --- a/packages/relarena/src/relarena/models/dummy/model.py +++ b/packages/relarena/src/relarena/models/dummy/model.py @@ -30,11 +30,11 @@ from relbench.base import Database, EntityTask, Table, TaskType from sklearn.dummy import DummyClassifier, DummyRegressor -from relarena.metrics import primary_metric -from relarena.model import RelArenaModel -from relarena.models._shared.predict_contract import predict_to_contract -from relarena.registry import register_model -from relarena.search_space import SearchSpace +from relarena.core.metrics import primary_metric +from relarena.core.model import RelArenaModel +from relarena.core.predict_contract import predict_to_contract +from relarena.core.registry import register_model +from relarena.core.search_space import SearchSpace #: Metrics minimized by the median; everything else (MSE/RMSE/R²) by the mean. _MEDIAN_METRICS = {"mae"} diff --git a/packages/relarena/src/relarena/models/graphsage/__init__.py b/packages/relarena/src/relarena/models/graphsage/__init__.py index bd570d6..b3e27bb 100644 --- a/packages/relarena/src/relarena/models/graphsage/__init__.py +++ b/packages/relarena/src/relarena/models/graphsage/__init__.py @@ -1,8 +1,4 @@ -"""`graphsage` — RelBench`s heterogeneous, temporal GraphSAGE GNN. - -Importing this package registers the `graphsage` model. See `model`; the vendored -GNN building blocks it drives live in `models/_shared/gnn/_vendor`. -""" +"""Temporal GraphSAGE baseline.""" from relarena.models.graphsage.model import GRAPHSAGE_SPACE, GraphSAGEModel diff --git a/packages/relarena/src/relarena/models/graphsage/model.py b/packages/relarena/src/relarena/models/graphsage/model.py index 0a81d6c..ec6c45a 100644 --- a/packages/relarena/src/relarena/models/graphsage/model.py +++ b/packages/relarena/src/relarena/models/graphsage/model.py @@ -57,8 +57,10 @@ from ConfigSpace import Categorical, ConfigurationSpace, Float from relbench.base import Database, EntityTask, Table, TaskType -from relarena.metrics import get_metric, primary_metric -from relarena.model import RelArenaModel +from relarena.core.metrics import get_metric, primary_metric +from relarena.core.model import RelArenaModel +from relarena.core.registry import register_model +from relarena.core.search_space import SearchSpace from relarena.models._shared.gnn.graph import GRAPH_CACHE, build_graph from relarena.models._shared.gnn.training import ( default_device, @@ -66,8 +68,6 @@ task_setup, train_epoch, ) -from relarena.registry import register_model -from relarena.search_space import SearchSpace logger = logging.getLogger(__name__) diff --git a/packages/relarena/src/relarena/models/kurversc/__init__.py b/packages/relarena/src/relarena/models/kurversc/__init__.py index 894c677..92dfbf8 100644 --- a/packages/relarena/src/relarena/models/kurversc/__init__.py +++ b/packages/relarena/src/relarena/models/kurversc/__init__.py @@ -1,4 +1,4 @@ -"""`kurversc` — validation-guided GraphReduce system.""" +"""Validation-guided GraphReduce system.""" from relarena.models.kurversc.model import KURVERSC_DEFAULTS, KurveRSCSystem diff --git a/packages/relarena/src/relarena/models/kurversc/model.py b/packages/relarena/src/relarena/models/kurversc/model.py index b3a9341..587506c 100644 --- a/packages/relarena/src/relarena/models/kurversc/model.py +++ b/packages/relarena/src/relarena/models/kurversc/model.py @@ -28,9 +28,9 @@ import numpy as np from relbench.base import Database, EntityTask, Table -from relarena.dataset import InnerSplit, OuterSplit -from relarena.registry import register_system -from relarena.system import RelArenaSystem +from relarena.core.dataset import InnerSplit, OuterSplit +from relarena.core.registry import register_system +from relarena.core.system import RelArenaSystem logger = logging.getLogger(__name__) diff --git a/packages/relarena/src/relarena/models/lightgbm/__init__.py b/packages/relarena/src/relarena/models/lightgbm/__init__.py index ed02cbd..35b820b 100644 --- a/packages/relarena/src/relarena/models/lightgbm/__init__.py +++ b/packages/relarena/src/relarena/models/lightgbm/__init__.py @@ -1,7 +1,4 @@ -"""`lightgbm` — gradient-boosted trees on entity-only features. - -Importing this package registers the `lightgbm` model. See `model`. -""" +"""Gradient-boosted trees on entity-only features.""" from relarena.models.lightgbm.model import LIGHTGBM_SPACE, LightGBMModel diff --git a/packages/relarena/src/relarena/models/lightgbm/model.py b/packages/relarena/src/relarena/models/lightgbm/model.py index 19a6376..a27f12e 100644 --- a/packages/relarena/src/relarena/models/lightgbm/model.py +++ b/packages/relarena/src/relarena/models/lightgbm/model.py @@ -31,11 +31,11 @@ from ConfigSpace import Categorical, ConfigurationSpace, Constant, Float, Integer from relbench.base import Database, EntityTask, Table +from relarena.core.model import RelArenaModel +from relarena.core.registry import register_model +from relarena.core.search_space import SearchSpace from relarena.featurization import build_entity_features -from relarena.model import RelArenaModel from relarena.models._shared.gbdt.lgb import fit_lgb, predict_lgb -from relarena.registry import register_model -from relarena.search_space import SearchSpace def _lightgbm_config_space() -> ConfigurationSpace: diff --git a/packages/relarena/src/relarena/models/rdblearn/__init__.py b/packages/relarena/src/relarena/models/rdblearn/__init__.py index 8d807ed..7850757 100644 --- a/packages/relarena/src/relarena/models/rdblearn/__init__.py +++ b/packages/relarena/src/relarena/models/rdblearn/__init__.py @@ -1,7 +1,4 @@ -"""`rdblearn` — DFS features fed to a tabular foundation model. - -Importing this package registers the `rdblearn` model. See `model`. -""" +"""DFS features with model-owned tabular foundation model recipes.""" from relarena.models.rdblearn.model import RDBLEARN_SPACE, RDBLearnModel diff --git a/packages/relarena/src/relarena/models/rdblearn/model.py b/packages/relarena/src/relarena/models/rdblearn/model.py index 1b01aae..69d3893 100644 --- a/packages/relarena/src/relarena/models/rdblearn/model.py +++ b/packages/relarena/src/relarena/models/rdblearn/model.py @@ -2,7 +2,7 @@ Combines: * **featurization** — multi-hop Deep Feature Synthesis over the foreign-key graph - (`relarena.featurization.build_dfs_features`, with the depth cache), + (`relarena.core.featurization.build_dfs_features`, with the depth cache), plus target-history augmentation (past-label aggregates), temporal-diff features, and the anchor columns (entity key + cutoff-time calendar features); * **search space** — an explicit grid over **(which tabular foundation model) × @@ -12,7 +12,8 @@ dimension. The estimator is a tabular foundation model (TabPFN v2 / v2.5); see -`_shared/tfm/tfm.py` for the TFM registry and the downsample -> fit/predict core +`relarena.models.rdblearn.tfm` for the backend definitions and +`relarena.core.tfm` for shared fitting (the TFM handles categoricals natively). This is RDBLearn proper (https://github.com/HKUSHXLab/rdblearn) — DFS features + a foundation model. @@ -25,7 +26,8 @@ One deliberate deviation from upstream RDBLearn's preprocessing: it label-encodes categoricals and then runs AutoGluon's `AutoMLPipelineFeatureGenerator` over the result, whereas here the feature frame reaches the TFM as-is, so TabPFN does its own -categorical detection and NaN handling (`_shared/tfm/tfm.py` has the why — upstream +categorical detection and NaN handling (`relarena.models.rdblearn.tfm` has the why — +upstream encodes because its backends take numpy arrays, a constraint a TabPFN-only grid does not have). The generator's datetime expansion is covered natively: the anchor cutoff gets the same year / month / day / dayofweek decomposition plus an epoch value, and @@ -40,19 +42,17 @@ import numpy as np from relbench.base import Database, EntityTask, Table -from relarena.featurization import DFS_MAX_DEPTH, build_dfs_features -from relarena.model import RelArenaModel -from relarena.models._shared.tfm.tfm import ( - fit_tfm, - predict_tfm, -) -from relarena.registry import register_model -from relarena.search_space import SearchSpace +from relarena.core.featurization.dfs import DFS_MAX_DEPTH, build_dfs_features +from relarena.core.model import RelArenaModel +from relarena.core.registry import register_model +from relarena.core.search_space import SearchSpace +from relarena.core.tfm import fit_tfm, predict_tfm +from relarena.models.rdblearn.tfm import TFM_REGISTRY _MIN_DEPTH = 2 -#: Tabular foundation models to sweep (names in the shared `TFM_REGISTRY`); both -#: ship in the `tabpfn` package (a core dependency) and run under the RDBLearn +#: Tabular foundation models to sweep (names in RDBLearn's `TFM_REGISTRY`); both +#: ship in the `tabpfn` package (the `rdblearn` extra) and run under the RDBLearn #: paper's 10k fit limit. The paper (arXiv:2602.18495) sweeps exactly TabPFN #: v2 / v2.5 (and LimiX, which #: has no pip package — it / TabPFN-v3 plug in here once registered). @@ -147,7 +147,7 @@ def fit( df, train_table.df[task.target_col], task.task_type, - tfm=self._tfm, + spec=TFM_REGISTRY[self._tfm], seed=seed, max_predict_samples=( _configure_prediction_batching() if self._tfm in _TFMS else None diff --git a/packages/relarena/src/relarena/models/rdblearn/tfm.py b/packages/relarena/src/relarena/models/rdblearn/tfm.py new file mode 100644 index 0000000..1996b47 --- /dev/null +++ b/packages/relarena/src/relarena/models/rdblearn/tfm.py @@ -0,0 +1,60 @@ +"""TabPFN backend recipes for RDBLearn.""" + +from __future__ import annotations + +from typing import Any + +from relarena.core.tfm import TFMSpec + + +def _make_tabpfn( + version: str, + *, + regression: bool, + device: Any, + seed: int, + **overrides: Any, +) -> Any: + """Build a TabPFN estimator pinned to a version via `create_default_for_version`. + + The bare TabPFN constructor now defaults to v3; `create_default_for_version` + selects the right checkpoint + version-appropriate defaults for v2 / v2.5, and + `**overrides` (device, random_state, ignore_pretraining_limits, ...) pass through + to the constructor. + + Lazy import — tabpfn lives in the rdblearn extra, and it is the one dependency + under the Prior Labs License rather than a plain permissive one, so a core + install stays clear of its attribution obligation (see `docs/licensing.md`). + """ + from tabpfn import TabPFNClassifier, TabPFNRegressor + from tabpfn.constants import ModelVersion + + model_version = { + "v2": ModelVersion.V2, + "v2.5": ModelVersion.V2_5, + }[version] + estimator_cls = TabPFNRegressor if regression else TabPFNClassifier + return estimator_cls.create_default_for_version( + model_version, + device=device, + random_state=seed, + ignore_pretraining_limits=True, + **overrides, + ) + + +def _tabpfn_spec(version: str, max_train_samples: int) -> TFMSpec: + return TFMSpec( + make_classifier=lambda **kw: _make_tabpfn(version, regression=False, **kw), + make_regressor=lambda **kw: _make_tabpfn(version, regression=True, **kw), + max_train_samples=max_train_samples, + ) + + +TFM_REGISTRY: dict[str, TFMSpec] = { + "tabpfn-v2": _tabpfn_spec("v2", max_train_samples=10_000), + "tabpfn-v2.5": _tabpfn_spec("v2.5", max_train_samples=10_000), +} + + +__all__ = ["TFM_REGISTRY"] diff --git a/packages/relarena/src/relarena/models/relgnn/__init__.py b/packages/relarena/src/relarena/models/relgnn/__init__.py index aabe525..de6ef14 100644 --- a/packages/relarena/src/relarena/models/relgnn/__init__.py +++ b/packages/relarena/src/relarena/models/relgnn/__init__.py @@ -1,9 +1,4 @@ -"""`relgnn` — composite message passing over atomic routes. - -Importing this package registers `relgnn` and `relgnn-es` (scored at its best-val -checkpoint instead of a train+val refit). See `model`; `_vendor` holds the -upstream RelGNN stack. -""" +"""RelGNN baselines with full-data and early-stopping refit policies.""" from relarena.models.relgnn.model import ( RELGNN_SPACE, @@ -11,8 +6,4 @@ RelGNNModel, ) -__all__ = [ - "RELGNN_SPACE", - "RelGNNEarlyStopModel", - "RelGNNModel", -] +__all__ = ["RELGNN_SPACE", "RelGNNEarlyStopModel", "RelGNNModel"] diff --git a/packages/relarena/src/relarena/models/relgnn/model.py b/packages/relarena/src/relarena/models/relgnn/model.py index 5d6ee8e..8d7da1f 100644 --- a/packages/relarena/src/relarena/models/relgnn/model.py +++ b/packages/relarena/src/relarena/models/relgnn/model.py @@ -42,8 +42,11 @@ from ConfigSpace import Categorical, ConfigurationSpace, Float from relbench.base import Database, EntityTask, Table -from relarena.metrics import get_metric, primary_metric -from relarena.model import RelArenaModel +from relarena.core.metrics import get_metric, primary_metric +from relarena.core.model import RelArenaModel +from relarena.core.registry import register_model +from relarena.core.search_space import SearchSpace +from relarena.core.tasks import ENTITY_TASK_TYPES from relarena.models._shared.gnn.training import ( default_device, infer, @@ -51,9 +54,6 @@ train_epoch, ) from relarena.models.relgnn.preprocessing import load_graph -from relarena.registry import register_model -from relarena.search_space import SearchSpace -from relarena.tasks import ENTITY_TASK_TYPES # Fixed budget / infra constants (never tuned; not part of any hyperparameter config). _BATCH_SIZE = 512 diff --git a/packages/relarena/src/relarena/models/relgnn/preprocessing.py b/packages/relarena/src/relarena/models/relgnn/preprocessing.py index 13e87f9..7733609 100644 --- a/packages/relarena/src/relarena/models/relgnn/preprocessing.py +++ b/packages/relarena/src/relarena/models/relgnn/preprocessing.py @@ -8,9 +8,9 @@ from relbench.base import Database -from relarena.cache import CacheConfig, cache_key, cached_artifact -from relarena.checksums import database_checksum -from relarena.identity import RunIdentity +from relarena.core.cache import CacheConfig, cache_key, cached_artifact +from relarena.core.fingerprints import database_checksum +from relarena.core.identity import RunIdentity from relarena.models._shared.gnn.graph_cache import DBGraphCache _GRAPH_CACHE_VERSION = 2 diff --git a/packages/relarena/src/relarena/models/relgnn/warm_cache.py b/packages/relarena/src/relarena/models/relgnn/warm_cache.py index fbc4c31..375e6b9 100644 --- a/packages/relarena/src/relarena/models/relgnn/warm_cache.py +++ b/packages/relarena/src/relarena/models/relgnn/warm_cache.py @@ -7,7 +7,7 @@ import torch -from relarena.cache import resolve_cache_config +from relarena.core.cache import resolve_cache_config from relarena.dataset import RelBenchDatasetTask from relarena.models.relgnn.preprocessing import load_graph diff --git a/packages/relarena/src/relarena/models/relgt/__init__.py b/packages/relarena/src/relarena/models/relgt/__init__.py index 6f99f1d..5012b5f 100644 --- a/packages/relarena/src/relarena/models/relgt/__init__.py +++ b/packages/relarena/src/relarena/models/relgt/__init__.py @@ -1,8 +1,4 @@ -"""`relgt` — the Relational Graph Transformer. - -Importing this package registers the `relgt` model. See `model`; `tokenize` is the -relarena driver for the sampler, and `_vendor` holds the upstream RelGT stack. -""" +"""Relational Graph Transformer baseline.""" from relarena.models.relgt.model import ( RelGTModel, diff --git a/packages/relarena/src/relarena/models/relgt/model.py b/packages/relarena/src/relarena/models/relgt/model.py index 1615b2a..3ee7c2a 100644 --- a/packages/relarena/src/relarena/models/relgt/model.py +++ b/packages/relarena/src/relarena/models/relgt/model.py @@ -37,10 +37,10 @@ from relbench.base import Database, EntityTask, Table, TaskType from torch.utils.data import DataLoader -from relarena.metrics import get_metric, primary_metric -from relarena.model import RelArenaModel -from relarena.registry import register_model -from relarena.search_space import SearchSpace, TaskStats +from relarena.core.metrics import get_metric, primary_metric +from relarena.core.model import RelArenaModel +from relarena.core.registry import register_model +from relarena.core.search_space import SearchSpace, TaskStats # Architecture constants — fixed across the paper's runs (expts/*.sh), not tuned. _CHANNELS = 512 diff --git a/packages/relarena/src/relarena/models/relgt/tokenize.py b/packages/relarena/src/relarena/models/relgt/tokenize.py index f93728b..835ed3f 100644 --- a/packages/relarena/src/relarena/models/relgt/tokenize.py +++ b/packages/relarena/src/relarena/models/relgt/tokenize.py @@ -46,8 +46,8 @@ from torch.utils.data import Dataset from torch_geometric.data import HeteroData -from relarena.cache import CacheConfig, cache_key, cached_artifact -from relarena.identity import RunIdentity +from relarena.core.cache import CacheConfig, cache_key, cached_artifact +from relarena.core.identity import RunIdentity from relarena.models.relgt._vendor._sampler import ( _process_one_seed, build_adjacency_hetero, diff --git a/packages/relarena/src/relarena/models/relgt/warm_cache.py b/packages/relarena/src/relarena/models/relgt/warm_cache.py index 7b5e215..a1fe3fb 100644 --- a/packages/relarena/src/relarena/models/relgt/warm_cache.py +++ b/packages/relarena/src/relarena/models/relgt/warm_cache.py @@ -6,7 +6,7 @@ from collections.abc import Callable from pathlib import Path -from relarena.cache import resolve_cache_config +from relarena.core.cache import resolve_cache_config from relarena.dataset import RelBenchDatasetTask from relarena.models._shared.gnn.graph import build_graph diff --git a/packages/relarena/src/relarena/models/rt/__init__.py b/packages/relarena/src/relarena/models/rt/__init__.py index 61d59ad..8897879 100644 --- a/packages/relarena/src/relarena/models/rt/__init__.py +++ b/packages/relarena/src/relarena/models/rt/__init__.py @@ -1,9 +1,4 @@ -"""`rt-plurel` — the Relational Transformer, fine-tuned per task from RT-P. - -Importing this package registers the `rt-plurel` system. Every configured value -lives in `config`; see `model` for the wrapper and `export` for the tensor -export. -""" +"""Relational Transformer system with per-task fine-tuning.""" from relarena.models.rt.model import RTPluRelSystem, clear_scratch diff --git a/packages/relarena/src/relarena/models/rt/export.py b/packages/relarena/src/relarena/models/rt/export.py index 59e30fa..a5173cc 100644 --- a/packages/relarena/src/relarena/models/rt/export.py +++ b/packages/relarena/src/relarena/models/rt/export.py @@ -85,9 +85,9 @@ import yaml from relbench.base import Database, EntityTask, Table, TaskType -from relarena.cache import CacheConfig, CacheMiss, cache_key -from relarena.checksums import database_checksum, table_checksum -from relarena.identity import RunIdentity +from relarena.core.cache import CacheConfig, CacheMiss, cache_key +from relarena.core.fingerprints import database_checksum, table_checksum +from relarena.core.identity import RunIdentity from relarena.models.rt.config import preprocess_args logger = logging.getLogger(__name__) @@ -307,7 +307,7 @@ def _publish_directory( ) -> Path: """Build/publish one directory artifact and return the **path** it lives at. - `relarena.cache.cached_artifact` is the shared layer for this, and it does + `relarena.core.cache.cached_artifact` is the shared layer for this, and it does not fit: it returns a *value* that `load` materialized in memory, and builds a miss inside a `TemporaryDirectory` that is deleted before it returns. Both are right for a feature matrix and wrong for this artifact, whose consumer diff --git a/packages/relarena/src/relarena/models/rt/model.py b/packages/relarena/src/relarena/models/rt/model.py index 765548d..9b811ba 100644 --- a/packages/relarena/src/relarena/models/rt/model.py +++ b/packages/relarena/src/relarena/models/rt/model.py @@ -59,12 +59,12 @@ import numpy as np from relbench.base import Database, EntityTask, Table, TaskType -from relarena.dataset import InnerSplit, OuterSplit, concat_tables -from relarena.identity import RunIdentity +from relarena.core.dataset import InnerSplit, OuterSplit, concat_tables +from relarena.core.identity import RunIdentity +from relarena.core.registry import register_system +from relarena.core.system import RelArenaSystem from relarena.models.rt import config as cfg from relarena.models.rt.export import TASK_DIR, preprocessed_dir, target_stats -from relarena.registry import register_system -from relarena.system import RelArenaSystem logger = logging.getLogger(__name__) @@ -357,7 +357,7 @@ def _tune_context( from rt.data import get_tasks from rt.eval import build_evaluator - from relarena.metrics import is_better, primary_metric + from relarena.core.metrics import is_better, primary_metric metric = primary_metric(task) truth = val_table.df[task.target_col].to_numpy() diff --git a/packages/relarena/src/relarena/models/rt/warm_cache.py b/packages/relarena/src/relarena/models/rt/warm_cache.py index 3daec4f..2f449d6 100644 --- a/packages/relarena/src/relarena/models/rt/warm_cache.py +++ b/packages/relarena/src/relarena/models/rt/warm_cache.py @@ -18,8 +18,9 @@ import argparse from pathlib import Path -from relarena.cache import resolve_cache_config -from relarena.dataset import RelBenchDatasetTask, concat_tables +from relarena.core.cache import resolve_cache_config +from relarena.core.dataset import concat_tables +from relarena.dataset import RelBenchDatasetTask from relarena.models.rt.export import preprocessed_dir from relarena.models.rt.model import DB_NAME diff --git a/packages/relarena/src/relarena/models/tabpfn_rel/context.py b/packages/relarena/src/relarena/models/tabpfn_rel/context.py index 2704c17..446d426 100644 --- a/packages/relarena/src/relarena/models/tabpfn_rel/context.py +++ b/packages/relarena/src/relarena/models/tabpfn_rel/context.py @@ -28,7 +28,8 @@ import pandas as pd from relbench.base import TaskType -from relarena.models._shared.tfm.tfm import FittedTFM, fit_tfm +from relarena.core.tfm import FittedTFM, fit_tfm +from relarena.models.tabpfn_rel.tfm import TFM_REGISTRY #: Default ensemble size for the pool strategies when the config pins none. DEFAULT_POOL_N_ESTIMATORS = 8 @@ -159,7 +160,7 @@ def fit( context_time: np.ndarray | None = None, ) -> FittedTFM: """Defer to the base `fit_tfm` (it caps rows itself); recency unused.""" - return fit_tfm(df, y, task_type, tfm=tfm, seed=seed) + return fit_tfm(df, y, task_type, spec=TFM_REGISTRY[tfm], seed=seed) @classmethod def from_config(cls, config: dict[str, Any]) -> "RandomContext": @@ -224,7 +225,7 @@ def fit( df, y, task_type, - tfm=tfm, + spec=TFM_REGISTRY[tfm], seed=seed, max_train_samples=len(df), overrides=overrides, diff --git a/packages/relarena/src/relarena/models/tabpfn_rel/features.py b/packages/relarena/src/relarena/models/tabpfn_rel/features.py index d3acb21..bb67c49 100644 --- a/packages/relarena/src/relarena/models/tabpfn_rel/features.py +++ b/packages/relarena/src/relarena/models/tabpfn_rel/features.py @@ -215,7 +215,7 @@ class RawTextFeaturizer: """Anchor-text columns passed through as raw strings. The TFM's estimator handles them (see - `relarena.models._shared.tfm.tfm`). Columns are + `relarena.models.tabpfn_rel.tfm`). Columns are suffixed `__raw_text` so a low-cardinality anchor column DFS kept as a categorical cannot collide. Stateless beyond the column list and uncached (the lookup is a cheap merge). diff --git a/packages/relarena/src/relarena/models/tabpfn_rel/model.py b/packages/relarena/src/relarena/models/tabpfn_rel/model.py index 6d4b322..3bd3e4e 100644 --- a/packages/relarena/src/relarena/models/tabpfn_rel/model.py +++ b/packages/relarena/src/relarena/models/tabpfn_rel/model.py @@ -16,7 +16,7 @@ in-context examples from a recency-weighted pool instead of a uniform sample. `fit` is the whole story top to bottom: build DFS features for the full train labels -(cached on a warm run; see `relarena.featurization.cache`), apply the +(cached on a warm run; see `relarena.core.featurization.cache`), apply the enabled feature extras, then fit the TFM — a seeded downsample for the default `random` context, or the recency pool for `soft_pool` / `hard_pool`. The expensive DFS matrix is content-cached, so the downsample / pool selection happens cheaply @@ -36,13 +36,14 @@ import numpy as np from relbench.base import Database, EntityTask, Table -from relarena.featurization import DFS_MAX_DEPTH, build_dfs_features -from relarena.model import RelArenaModel -from relarena.models._shared.tfm.tfm import TFM_REGISTRY, predict_tfm +from relarena.core.featurization import DFS_MAX_DEPTH, build_dfs_features +from relarena.core.model import RelArenaModel +from relarena.core.registry import register_model +from relarena.core.search_space import SearchSpace +from relarena.core.tfm import predict_tfm from relarena.models.tabpfn_rel.context import ContextStrategy from relarena.models.tabpfn_rel.features import FeaturePipeline -from relarena.registry import register_model -from relarena.search_space import SearchSpace +from relarena.models.tabpfn_rel.tfm import TFM_REGISTRY #: Depth grid lower bound. Shallower depths yield no DFS features for entity tasks. _MIN_DEPTH = 2 @@ -140,7 +141,7 @@ def warm_cache( Runs exactly the `build_dfs_features` + feature-pipeline calls `fit` and `predict` make (same cache keys), so a later eval reads them instead of recomputing. Construct the model with an explicit fill config; needs only CPU. - The shared public command is `relarena.featurization.warm_cache`. + The shared public command is `relarena.core.featurization.warm_cache`. """ self._tfm = self.config.get("tfm", "tabpfn-v3") self._depth = int(self.config.get("max_depth", _MIN_DEPTH)) diff --git a/packages/relarena/src/relarena/models/tabpfn_rel/tfm.py b/packages/relarena/src/relarena/models/tabpfn_rel/tfm.py new file mode 100644 index 0000000..8df900d --- /dev/null +++ b/packages/relarena/src/relarena/models/tabpfn_rel/tfm.py @@ -0,0 +1,119 @@ +"""TabPFN backend recipes.""" + +from __future__ import annotations + +from typing import Any + +import numpy as np + +from relarena.core.tfm import TFMSpec + + +def _make_tabpfn( + version: str, + *, + regression: bool, + device: Any, + seed: int, + **overrides: Any, +) -> Any: + """Build a TabPFN estimator pinned to a version via `create_default_for_version`. + + The bare TabPFN constructor now defaults to v3; `create_default_for_version` + selects the right checkpoint + version-appropriate defaults for v2 / v2.5, and + `**overrides` (device, random_state, ignore_pretraining_limits, ...) pass through + to the constructor. + + The local extra supplies TabPFN; importing model definitions does not load it. + """ + from tabpfn import TabPFNClassifier, TabPFNRegressor + from tabpfn.constants import ModelVersion + + model_version = { + "v2": ModelVersion.V2, + "v2.5": ModelVersion.V2_5, + "v3": ModelVersion.V3, + }[version] + estimator_cls = TabPFNRegressor if regression else TabPFNClassifier + return estimator_cls.create_default_for_version( + model_version, + device=device, + random_state=seed, + ignore_pretraining_limits=True, + **overrides, + ) + + +def _sanitize_api_overrides(overrides: dict[str, Any]) -> dict[str, Any]: + """Convert numpy index arrays in `SUBSAMPLE_SAMPLES` to plain int lists. + + tabpfn_client pydantic-serializes the estimator config into the request body, + which rejects numpy arrays — the pool contexts pass per-estimator context + indices as arrays. The local TabPFN consumes arrays natively, so the + conversion is scoped to the API path. + """ + inference_config = overrides.get("inference_config") + if inference_config is None: + return overrides + subsample = inference_config.get("SUBSAMPLE_SAMPLES") + if not isinstance(subsample, list): + return overrides + return { + **overrides, + "inference_config": { + **inference_config, + "SUBSAMPLE_SAMPLES": [ + e.tolist() if isinstance(e, np.ndarray) else e for e in subsample + ], + }, + } + + +def _make_tabpfn_api( + *, + regression: bool, + device: Any, + seed: int, + **overrides: Any, +) -> Any: + """Build a TabPFN API-client estimator pinned to the v3 model. + + The api extra supplies tabpfn_client; it is imported when constructed. + Fit and predict + run server-side, so device is ignored and raw text columns are handled by the + API. + """ + from tabpfn_client import TabPFNClassifier as ApiClassifier + from tabpfn_client import TabPFNRegressor as ApiRegressor + + del device + overrides = _sanitize_api_overrides(overrides) + estimator_cls = ApiRegressor if regression else ApiClassifier + return estimator_cls( + model_path="v3_default", + random_state=seed, + ignore_pretraining_limits=True, + **overrides, + ) + + +TFM_REGISTRY: dict[str, TFMSpec] = { + "tabpfn-v3": TFMSpec( + make_classifier=lambda **kw: _make_tabpfn( + "v3", regression=False, n_preprocessing_jobs=-1, **kw + ), + make_regressor=lambda **kw: _make_tabpfn( + "v3", regression=True, n_preprocessing_jobs=-1, **kw + ), + max_train_samples=100_000, + ), + "tabpfn-v3-api": TFMSpec( + make_classifier=lambda **kw: _make_tabpfn_api(regression=False, **kw), + make_regressor=lambda **kw: _make_tabpfn_api(regression=True, **kw), + max_train_samples=100_000, + supports_text=True, + ), +} + + +__all__ = ["TFM_REGISTRY"] diff --git a/packages/relarena/src/relarena/refit.py b/packages/relarena/src/relarena/refit.py new file mode 100644 index 0000000..a40f19a --- /dev/null +++ b/packages/relarena/src/relarena/refit.py @@ -0,0 +1,76 @@ +"""Final benchmark fitting and evaluation against hidden test labels.""" + +from __future__ import annotations + +import time +from typing import Any, Type + +from relbench.base import EntityTask + +from relarena.core.cache import CacheConfig +from relarena.core.dataset import OuterSplit, concat_tables +from relarena.core.identity import RunIdentity +from relarena.core.metrics import evaluate_predictions, primary_metric +from relarena.core.model import RelArenaModel + + +def refit_and_evaluate( + model_cls: Type[RelArenaModel], + config: dict[str, Any], + task: EntityTask, + split: OuterSplit, + *, + seed: int, + time_limit: float | None = None, + cache: CacheConfig | None = None, + run_identity: RunIdentity | None = None, +) -> dict: + """Fit the selected `config` on the outer `split` and score it on `test`. + + Two final-fit regimes, chosen by `model_cls.refit_on_full_data`: + + * `True` (default): refit on the train+val union with `val_table=None` — no + held-out split, so a model with early stopping falls back to a fixed budget. + * `False`: train on train alone and pass `val` through, so a model that + checkpoints on validation reports its best-val model (e.g. RelGT's protocol). + + Both train on the outer split's (test-censored) DB. The split carries no eval + target (test labels are hidden), so we score with `target_table=None` and + RelBench sources the unmasked test labels itself. Returns the test fields to + attach to the trial. + """ + metric = primary_metric(task) + + if model_cls.refit_on_full_data: + train_table, val_table = concat_tables(split.train_table, split.val_table), None + else: + train_table, val_table = split.train_table, split.val_table + + model = model_cls(config, cache=cache, run_identity=run_identity) + t0 = time.perf_counter() + model.fit( + task, + split.db_state, + train_table, + val_table, + seed=seed, + time_limit=time_limit, + ) + fit_time_refit = time.perf_counter() - t0 + + t1 = time.perf_counter() + test_pred = model.predict(task, split.db_state, split.eval_table) + predict_time_refit = time.perf_counter() - t1 + + test_metrics = evaluate_predictions(task, test_pred, None, metric) + test_score = float(test_metrics[metric.__name__]) + return { + "test_score": test_score, + "test_metrics": test_metrics, + "test_pred": test_pred, + "fit_time_refit": fit_time_refit, + "predict_time_refit": predict_time_refit, + } + + +__all__ = ["refit_and_evaluate"] diff --git a/packages/relarena/src/relarena/results.py b/packages/relarena/src/relarena/results.py index 60274c5..0fc6aea 100644 --- a/packages/relarena/src/relarena/results.py +++ b/packages/relarena/src/relarena/results.py @@ -1,17 +1,12 @@ -"""Model and system result schemas. - -Models produce one `TrialResult` per harness-selected configuration. Systems -produce one `SystemResult` for their complete internal procedure. -""" +"""Flatten benchmark trials and experiment summaries into result frames.""" from __future__ import annotations -import hashlib import json -from dataclasses import dataclass, field, fields +from dataclasses import fields from typing import TYPE_CHECKING, Any -import numpy as np +from relarena.core.results import SystemResult, TrialResult if TYPE_CHECKING: import pandas as pd @@ -19,73 +14,6 @@ from relarena.runner import ExperimentSummary, SystemExperimentSummary -def config_id_for(config: dict[str, Any]) -> str: - """A short, deterministic id for a hyperparameter config (order-independent).""" - blob = json.dumps(config, sort_keys=True, default=str) - return hashlib.sha1(blob.encode()).hexdigest()[:8] - - -@dataclass -class TrialResult: - """Outcome of fitting & evaluating ONE config (within an ExperimentSummary). - - Identity — model / dataset / task / seed / metric — is owned by the enclosing - `ExperimentSummary` (every trial in a summary shares - it), so it lives there once instead of being duplicated on each trial; a - `TrialResult` records only the per-config result. - """ - - config: dict[str, Any] - config_id: str - config_tag: str # "default" or "r{i}" — identifies the default vs random configs - - # The primary (selection) metric's value. `val_score` comes from the - # train-only model (the selection signal); `test_score` is filled only for - # the selected config, from the model's final-fit regime (see the runner). - val_score: float | None = None - test_score: float | None = None - - # All of the task's native metrics (plus the primary), keyed by metric name. - val_metrics: dict[str, float] = field(default_factory=dict) - test_metrics: dict[str, float] = field(default_factory=dict) - - # Wall-clock seconds, split by phase: tuning (train-only fit + val predict, - # per config) and final fit (fit + test predict for selected/default configs). - # The refit-named fields are None for configs that did not receive a final fit. - fit_time_tuning: float = 0.0 - predict_time_tuning: float = 0.0 - fit_time_refit: float | None = None - predict_time_refit: float | None = None - - # Optional prediction metadata (not serialized to the summary DataFrame). - # Shapes follow EntityTask.evaluate's contract. - val_pred: np.ndarray | None = field(default=None, repr=False) - test_pred: np.ndarray | None = field(default=None, repr=False) - - # Populated with a traceback string if the trial failed; `None` on success. - error: str | None = None - - @property - def ok(self) -> bool: - """Whether the trial succeeded (no error).""" - return self.error is None - - -@dataclass -class SystemResult: - """Outcome of one end-to-end system run. - - Systems do not expose harness-selected configurations or validation scores. - Their complete internal procedure is represented by a final test result and - one total wall-clock time. - """ - - test_score: float | None = None - test_metrics: dict[str, float] = field(default_factory=dict) - time_total: float = 0.0 - test_pred: np.ndarray | None = field(default=None, repr=False) - - #: Prediction-array fields dropped when flattening trials to a tidy DataFrame. #: The native-metric dicts (`val_metrics` / `test_metrics`) are *not* dropped #: — they're expanded into `val_` / `test_` columns. @@ -188,3 +116,11 @@ def summary_to_dataframe( df["metric"] = summary.metric_name front = [c for c in _IDENTITY_COLS if c in df.columns] + ["selected"] return df[front + [c for c in df.columns if c not in front]] + + +__all__ = [ + "TrialResult", + "SystemResult", + "trials_to_dataframe", + "summary_to_dataframe", +] diff --git a/packages/relarena/src/relarena/runner.py b/packages/relarena/src/relarena/runner.py index 2a243bd..23e7038 100644 --- a/packages/relarena/src/relarena/runner.py +++ b/packages/relarena/src/relarena/runner.py @@ -9,47 +9,29 @@ from __future__ import annotations import logging -import math import time from dataclasses import dataclass from pathlib import Path -from typing import Callable, Type +from typing import Type import numpy as np from relbench.base import TaskType -from relarena.cache import resolve_cache_config +from relarena.core.cache import resolve_cache_config +from relarena.core.model import RelArenaModel +from relarena.core.registry import registry +from relarena.core.results import SystemResult, TrialResult +from relarena.core.search_space import SearchSpaceProvider +from relarena.core.selection import select_best +from relarena.core.system import RelArenaSystem +from relarena.core.tasks import ENTITY_TASK_TYPES +from relarena.core.tuner import tune from relarena.dataset import RelBenchDatasetTask -from relarena.metrics import is_better -from relarena.model import RelArenaModel -from relarena.registry import registry -from relarena.results import SystemResult, TrialResult -from relarena.search_space import SearchSpaceProvider -from relarena.system import RelArenaSystem -from relarena.tasks import ENTITY_TASK_TYPES -from relarena.tuner import refit_and_evaluate, tune +from relarena.refit import refit_and_evaluate logger = logging.getLogger(__name__) -def select_best(trials: list[TrialResult], metric: Callable[..., float]) -> TrialResult: - """Pick the trial with the best validation score under `metric`'s direction.""" - valid = [ - t - for t in trials - if t.ok and t.val_score is not None and math.isfinite(t.val_score) - ] - if not valid: - raise RuntimeError( - "No successful trials with a finite validation score to select from." - ) - best = valid[0] - for t in valid[1:]: - if is_better(t.val_score, best.val_score, metric): - best = t - return best - - @dataclass class ExperimentSummary: """The headline result for one `(model, dataset, task, seed)` run.""" @@ -178,8 +160,8 @@ def run_model_experiment( ) -> 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. + `search_space` defaults to the one registered for `model_cls`; + pass it explicitly to override. Protocol (nested temporal validation; see docs/temporal-validation.md): 1. **Tune** — fit each config on `train`, score on `val`, using the DB diff --git a/packages/relarena/src/relarena/tasks.py b/packages/relarena/src/relarena/tasks.py index 6f90bef..72af8ee 100644 --- a/packages/relarena/src/relarena/tasks.py +++ b/packages/relarena/src/relarena/tasks.py @@ -1,8 +1,4 @@ -"""Task-type scope for RelArena. - -RelArena currently targets *entity* tasks only. RelBench's `TaskType` is the -source of truth; this module pins down which of its values are in scope. -""" +"""Named RelBench benchmark tasks and dataset selection.""" from __future__ import annotations @@ -10,25 +6,7 @@ from relbench.base import TaskType -#: Entity (node-level) task types RelArena supports. Excludes: -#: * `LINK_PREDICTION` — recommendation, out of scope; -#: * `MULTILABEL_CLASSIFICATION` — RelBench has no entity multilabel task -#: (its sole multilabel task is a TGB node-property ranking `BaseTask`). -#: Easy to re-add if a real entity multilabel task appears. -ENTITY_TASK_TYPES: frozenset[TaskType] = frozenset( - { - TaskType.BINARY_CLASSIFICATION, - TaskType.REGRESSION, - } -) - -# The rest of the codebase assumes exactly these two task types; guard against -# silently widening scope without revisiting those call sites. -assert ENTITY_TASK_TYPES == { - TaskType.REGRESSION, - TaskType.BINARY_CLASSIFICATION, -}, "RelArena currently supports only regression and binary classification tasks." - +from relarena.core.tasks import ENTITY_TASK_TYPES as ENTITY_TASK_TYPES #: The seven original RelBench (v1) datasets. #: diff --git a/packages/relarena/src/relarena/userdb/__init__.py b/packages/relarena/src/relarena/userdb/__init__.py index be39ab1..d47e7af 100644 --- a/packages/relarena/src/relarena/userdb/__init__.py +++ b/packages/relarena/src/relarena/userdb/__init__.py @@ -13,14 +13,14 @@ tasks ship as reference specs (`relbench_v1_spec`). """ -from relarena.userdb.ingest import DatabaseSpec -from relarena.userdb.query import PredictiveQuery, PredictiveQuerySpec +from relarena.core.userdb.ingest import DatabaseSpec +from relarena.core.userdb.query import PredictiveQuery, PredictiveQuerySpec +from relarena.core.userdb.spec import PredictiveTaskSpec from relarena.userdb.relbench_v1 import ( materialize_relbench, relbench_v1_spec, relbench_v1_tasks, ) -from relarena.userdb.spec import PredictiveTaskSpec __all__ = [ "DatabaseSpec", diff --git a/packages/relarena/src/relarena/userdb/relbench_v1/__init__.py b/packages/relarena/src/relarena/userdb/relbench_v1/__init__.py index 5abef7e..9b6542f 100644 --- a/packages/relarena/src/relarena/userdb/relbench_v1/__init__.py +++ b/packages/relarena/src/relarena/userdb/relbench_v1/__init__.py @@ -24,7 +24,7 @@ from relbench.datasets import get_dataset -from relarena.userdb.query import PredictiveQuerySpec +from relarena.core.userdb.query import PredictiveQuerySpec _SPEC_DIR = files(__name__) diff --git a/packages/relarena/tests/featurization/test_cache.py b/packages/relarena/tests/core/featurization/test_cache.py similarity index 95% rename from packages/relarena/tests/featurization/test_cache.py rename to packages/relarena/tests/core/featurization/test_cache.py index 2ac1b65..7f20bd6 100644 --- a/packages/relarena/tests/featurization/test_cache.py +++ b/packages/relarena/tests/core/featurization/test_cache.py @@ -7,8 +7,8 @@ import pandas as pd import pytest -from relarena.cache import CacheConfig, CacheMiss -from relarena.featurization.cache import cached_frame +from relarena.core.cache import CacheConfig, CacheMiss +from relarena.core.featurization.cache import cached_frame def _frame(value: int = 1) -> pd.DataFrame: diff --git a/packages/relarena/tests/featurization/test_dfs.py b/packages/relarena/tests/core/featurization/test_dfs.py similarity index 99% rename from packages/relarena/tests/featurization/test_dfs.py rename to packages/relarena/tests/core/featurization/test_dfs.py index 452ff54..0987028 100644 --- a/packages/relarena/tests/featurization/test_dfs.py +++ b/packages/relarena/tests/core/featurization/test_dfs.py @@ -12,14 +12,14 @@ from fastdfs.dfs import dfs_feature_column_name from relbench.base import Database, Table -import relarena.featurization.dfs as dfs_mod -from relarena.cache import CacheConfig -from relarena.featurization.dfs import ( +import relarena.core.featurization.dfs as dfs_mod +from relarena.core.cache import CacheConfig +from relarena.core.featurization.dfs import ( TARGET_HISTORY_TABLE_NAME, _DepthCache, _temporal_diff, ) -from relarena.identity import RunIdentity +from relarena.core.identity import RunIdentity def _toy_db() -> Database: diff --git a/packages/relarena/tests/featurization/test_warm_cache.py b/packages/relarena/tests/core/featurization/test_warm_cache.py similarity index 93% rename from packages/relarena/tests/featurization/test_warm_cache.py rename to packages/relarena/tests/core/featurization/test_warm_cache.py index ccdd96f..8f21b10 100644 --- a/packages/relarena/tests/featurization/test_warm_cache.py +++ b/packages/relarena/tests/core/featurization/test_warm_cache.py @@ -7,9 +7,9 @@ import pytest -from relarena.cache import CacheConfig -from relarena.featurization import warm_cache -from relarena.identity import RunIdentity +from relarena.core.cache import CacheConfig +from relarena.core.featurization import warm_cache +from relarena.core.identity import RunIdentity def test__warm_dfs_cache__uses_one_shared_preprocessor_for_both_phases( diff --git a/packages/relarena/tests/test_cache.py b/packages/relarena/tests/core/test_cache.py similarity index 99% rename from packages/relarena/tests/test_cache.py rename to packages/relarena/tests/core/test_cache.py index 2edf86f..f0908a4 100644 --- a/packages/relarena/tests/test_cache.py +++ b/packages/relarena/tests/core/test_cache.py @@ -11,7 +11,7 @@ import pytest -from relarena.cache import ( +from relarena.core.cache import ( CACHE_DIR_ENV, DISABLE_CACHE_ENV, LEGACY_DISABLE_CACHE_ENV, diff --git a/packages/relarena/tests/core/test_dataset.py b/packages/relarena/tests/core/test_dataset.py new file mode 100644 index 0000000..6cacc8d --- /dev/null +++ b/packages/relarena/tests/core/test_dataset.py @@ -0,0 +1,28 @@ +"""Shared split-table operations.""" + +import pandas as pd +from relbench.base import Table + +from relarena.core.dataset import concat_tables + + +def test_concat_tables_unions_rows_and_keeps_schema() -> None: + a = Table( + df=pd.DataFrame({"entity": [1, 2], "t": [10, 11], "y": [0.0, 1.0]}), + fkey_col_to_pkey_table={"entity": "users"}, + pkey_col=None, + time_col="t", + ) + b = Table( + df=pd.DataFrame({"entity": [3], "t": [12], "y": [2.0]}), + fkey_col_to_pkey_table={"entity": "users"}, + pkey_col=None, + time_col="t", + ) + c = concat_tables(a, b) + assert len(c.df) == 3 + assert list(c.df["y"]) == [0.0, 1.0, 2.0] + assert c.time_col == "t" + assert c.fkey_col_to_pkey_table == {"entity": "users"} + # inputs are untouched + assert len(a.df) == 2 and len(b.df) == 1 diff --git a/packages/relarena/tests/core/test_identity.py b/packages/relarena/tests/core/test_identity.py new file mode 100644 index 0000000..80085a3 --- /dev/null +++ b/packages/relarena/tests/core/test_identity.py @@ -0,0 +1,85 @@ +"""Tests for optional preprocessing run identities.""" + +from __future__ import annotations + +from types import SimpleNamespace + +import pandas as pd +from relbench.base import Database, Table + +from relarena.core.identity import ( + RunIdentity, + database_schema_fingerprint, + metadata_fingerprint, + task_spec_fingerprint, +) + + +def test__run_identity__phase_copy__does_not_mutate_base() -> None: + base = RunIdentity("dataset", "db", "task", "labels") + assert base.for_phase("inner").phase == "inner" + assert base.phase is None + + +def _database(dtype: str = "int64") -> Database: + return Database( + { + "entities": Table( + df=pd.DataFrame({"id": pd.Series([1, 2], dtype=dtype)}), + fkey_col_to_pkey_table={}, + pkey_col="id", + time_col=None, + ) + } + ) + + +def test__database_schema_fingerprint__ignores_rows_but_guards_dtype() -> None: + same_schema = Database( + { + "entities": Table( + df=pd.DataFrame({"id": pd.Series([9], dtype="int64")}), + fkey_col_to_pkey_table={}, + pkey_col="id", + time_col=None, + ) + } + ) + assert database_schema_fingerprint(_database()) == database_schema_fingerprint( + same_schema + ) + assert database_schema_fingerprint(_database()) != database_schema_fingerprint( + _database("float64") + ) + + +def _task_spec(**overrides: object) -> SimpleNamespace: + fields = { + "entity_table": "drivers", + "entity_col": "driver_id", + "time_col": "date", + "target_col": "dnf", + "task_type": "binary_classification", + "timedelta": "30 days", + "query": "SELECT * FROM labels", + "val_timestamp": "2005-01-01", + "test_timestamp": "2005-02-01", + "num_eval_timestamps": 1, + } + fields.update(overrides) + return SimpleNamespace(**fields) + + +def test__task_spec_fingerprint__is_stable_and_guards_training_semantics() -> None: + assert task_spec_fingerprint(_task_spec()) == task_spec_fingerprint(_task_spec()) + assert task_spec_fingerprint(_task_spec()) != task_spec_fingerprint( + _task_spec(timedelta="60 days") + ) + + +def test_metadata_fingerprint_is_order_independent_and_sensitive_to_values() -> None: + first = metadata_fingerprint({"b": [2, 3], "a": 1}) + assert first == metadata_fingerprint({"a": 1, "b": [2, 3]}) + assert first != metadata_fingerprint({"a": 1, "b": [3, 2]}) + assert len(first) == 16 + int(first, 16) diff --git a/packages/relarena/tests/test_metrics.py b/packages/relarena/tests/core/test_metrics.py similarity index 74% rename from packages/relarena/tests/test_metrics.py rename to packages/relarena/tests/core/test_metrics.py index 6c1437e..8739244 100644 --- a/packages/relarena/tests/test_metrics.py +++ b/packages/relarena/tests/core/test_metrics.py @@ -3,15 +3,17 @@ from __future__ import annotations from types import SimpleNamespace +from unittest.mock import Mock import numpy as np import pytest from relbench.base import TaskType from relbench.metrics import roc_auc -from relarena.metrics import ( +from relarena.core.metrics import ( _METRICS, Metric, + evaluate_predictions, get_metric, is_higher_better, primary_metric, @@ -90,3 +92,24 @@ def test_to_metric_error_propagates_nan() -> None: def test_to_metric_error_unknown_metric_raises() -> None: with pytest.raises(KeyError): to_metric_error(0.5, "not_a_metric") + + +@pytest.mark.parametrize("include_primary", [False, True]) +@pytest.mark.parametrize("held_out", [False, True]) +def test_evaluate_predictions_preserves_target_and_native_metrics( + include_primary: bool, held_out: bool +) -> None: + native = Mock(__name__="native_metric") + primary = Mock(__name__="primary_metric") + metrics = [native, primary] if include_primary else [native] + task = SimpleNamespace( + metrics=metrics.copy(), evaluate=Mock(return_value={"score": 0.5}) + ) + target = None if held_out else object() + predictions = np.array([0.2, 0.8]) + result = evaluate_predictions(task, predictions, target, primary) + assert result == {"score": 0.5} + task.evaluate.assert_called_once_with( + predictions, target, metrics=[native, primary] + ) + assert task.metrics == metrics diff --git a/packages/relarena/tests/test_model.py b/packages/relarena/tests/core/test_model.py similarity index 88% rename from packages/relarena/tests/test_model.py rename to packages/relarena/tests/core/test_model.py index 2642d52..6d94ab1 100644 --- a/packages/relarena/tests/test_model.py +++ b/packages/relarena/tests/core/test_model.py @@ -6,8 +6,8 @@ import numpy as np -from relarena.cache import CacheConfig -from relarena.model import RelArenaModel +from relarena.core.cache import CacheConfig +from relarena.core.model import RelArenaModel class _Model(RelArenaModel): diff --git a/packages/relarena/tests/models/_shared/test_predict_contract.py b/packages/relarena/tests/core/test_predict_contract.py similarity index 95% rename from packages/relarena/tests/models/_shared/test_predict_contract.py rename to packages/relarena/tests/core/test_predict_contract.py index 0014a0b..2ba1176 100644 --- a/packages/relarena/tests/models/_shared/test_predict_contract.py +++ b/packages/relarena/tests/core/test_predict_contract.py @@ -10,7 +10,7 @@ import numpy as np from relbench.base import TaskType -from relarena.models._shared.predict_contract import predict_to_contract +from relarena.core.predict_contract import predict_to_contract class _StubClassifier: diff --git a/packages/relarena/tests/test_search_space.py b/packages/relarena/tests/core/test_search_space.py similarity index 95% rename from packages/relarena/tests/test_search_space.py rename to packages/relarena/tests/core/test_search_space.py index dfa68be..643afe5 100644 --- a/packages/relarena/tests/test_search_space.py +++ b/packages/relarena/tests/core/test_search_space.py @@ -2,7 +2,7 @@ from __future__ import annotations -from relarena.search_space import ( +from relarena.core.search_space import ( SearchSpace, TaskStats, resolve_search_space, diff --git a/packages/relarena/tests/test_skeleton.py b/packages/relarena/tests/core/test_skeleton.py similarity index 65% rename from packages/relarena/tests/test_skeleton.py rename to packages/relarena/tests/core/test_skeleton.py index a87ea5c..9cb1372 100644 --- a/packages/relarena/tests/test_skeleton.py +++ b/packages/relarena/tests/core/test_skeleton.py @@ -1,21 +1,16 @@ -"""Skeleton smoke tests — no data download required. - -These exercise the pure-Python contract (metrics, registry, result schema, the -abstract base class). End-to-end tests that load a RelBench task and run the -tuner come with the first concrete model. -""" +"""Shared model, registry and result contracts.""" from __future__ import annotations import numpy as np import pytest -from relarena.metrics import get_metric, is_better, is_higher_better -from relarena.model import RelArenaModel -from relarena.registry import ModelRegistry -from relarena.results import TrialResult, config_id_for, trials_to_dataframe -from relarena.search_space import SearchSpace -from relarena.system import RelArenaSystem +from relarena.core.metrics import get_metric, is_better, is_higher_better +from relarena.core.model import RelArenaModel +from relarena.core.registry import ModelRegistry +from relarena.core.results import config_id_for +from relarena.core.search_space import SearchSpace +from relarena.core.system import RelArenaSystem def test_metric_direction() -> None: @@ -123,33 +118,3 @@ class BestVal(Default): assert RelArenaModel.refit_on_full_data is True # harness default assert Default.refit_on_full_data is True assert BestVal.refit_on_full_data is False - - -def test_trials_to_dataframe_drops_arrays() -> None: - t = TrialResult( - config={"a": 1}, - config_id=config_id_for({"a": 1}), - config_tag="default", - val_score=0.5, - val_pred=np.zeros(3), - ) - df = trials_to_dataframe([t]) - assert "val_pred" not in df.columns - assert df.loc[0, "val_score"] == 0.5 - - -def test_trials_to_dataframe_flattens_metrics() -> None: - t = TrialResult( - config={}, - config_id=config_id_for({}), - config_tag="default", - val_score=0.5, - test_score=0.6, - val_metrics={"mae": 0.5, "rmse": 0.7}, - test_metrics={"mae": 0.6, "rmse": 0.8}, - ) - df = trials_to_dataframe([t]) - # the metric dicts themselves are not columns; their entries are flattened. - assert "val_metrics" not in df.columns and "test_metrics" not in df.columns - assert df.loc[0, "val_mae"] == 0.5 and df.loc[0, "val_rmse"] == 0.7 - assert df.loc[0, "test_mae"] == 0.6 and df.loc[0, "test_rmse"] == 0.8 diff --git a/packages/relarena/tests/core/test_tfm.py b/packages/relarena/tests/core/test_tfm.py new file mode 100644 index 0000000..97833c1 --- /dev/null +++ b/packages/relarena/tests/core/test_tfm.py @@ -0,0 +1,158 @@ +"""Shared estimator fitting, sampling, alignment and prediction.""" + +from __future__ import annotations + +import numpy as np +import pandas as pd +import pytest +from relbench.base import TaskType + +from relarena.core.tfm import TFMSpec, _downsample_indices, fit_tfm, predict_tfm + + +class _StubClassifier: + """Minimal sklearn-like classifier: learns classes_ from y, emits fixed proba.""" + + def fit(self, X: pd.DataFrame, y: np.ndarray) -> "_StubClassifier": + self.classes_ = np.unique(y) + self.n_train_ = len(X) + return self + + def predict_proba(self, X: pd.DataFrame) -> np.ndarray: + cols = np.arange(1, len(self.classes_) + 1, dtype=float) + return np.tile(cols / cols.sum(), (len(X), 1)) + + +class _StubRegressor: + def fit(self, X: pd.DataFrame, y: np.ndarray) -> "_StubRegressor": + self.mean_ = float(np.mean(y)) + self.n_train_ = len(X) + return self + + def predict(self, X: pd.DataFrame) -> np.ndarray: + return np.full(len(X), self.mean_) + + +@pytest.fixture +def stub_tfm() -> TFMSpec: + spec = TFMSpec( + make_classifier=lambda **kw: _StubClassifier(), + make_regressor=lambda **kw: _StubRegressor(), + max_train_samples=10000, + ) + return spec + + +def _frame(n: int, seed: int = 0) -> pd.DataFrame: + rng = np.random.default_rng(seed) + return pd.DataFrame( + {"num": rng.normal(size=n), "cat": rng.choice(["a", "b", "c"], n)} + ) + + +def test_downsample_classification_caps_and_keeps_every_class() -> None: + y = np.array([0] * 25 + [1] * 25) + idx = _downsample_indices( + y, TaskType.BINARY_CLASSIFICATION, 10, np.random.default_rng(0) + ) + assert len(idx) == 10 + assert set(y[idx]) == {0, 1} + + +def test_downsample_regression_caps_and_is_seeded() -> None: + y = np.arange(100.0) + a = _downsample_indices(y, TaskType.REGRESSION, 20, np.random.default_rng(0)) + b = _downsample_indices(y, TaskType.REGRESSION, 20, np.random.default_rng(0)) + assert len(a) == 20 and np.array_equal(a, b) + + +def test_fit_predict_roundtrip_reindexes_to_training_columns(stub_tfm: TFMSpec) -> None: + df = _frame(20) + y = pd.Series([0, 1] * 10) + fitted = fit_tfm(df, y, TaskType.BINARY_CLASSIFICATION, spec=stub_tfm, seed=0) + pred = predict_tfm(fitted, df[["cat", "num"]]) + assert pred.shape == (20,) + + +def test_fit_does_not_pass_categorical_indices() -> None: + captured: dict[str, object] = {} + + def _make(**kw: object) -> _StubClassifier: + captured.update(kw) + return _StubClassifier() + + spec = TFMSpec(make_classifier=_make, make_regressor=_make, max_train_samples=10000) + df = _frame(20) + y = pd.Series([0, 1] * 10) + fit_tfm(df, y, TaskType.BINARY_CLASSIFICATION, spec=spec, seed=0) + assert "categorical_features_indices" not in captured + + +def test_predict_regression_requests_median_when_supported() -> None: + captured: dict[str, object] = {} + + class _OutputTypeAwareRegressor(_StubRegressor): + def predict(self, X: pd.DataFrame, **kwargs: object) -> np.ndarray: + captured["output_type"] = kwargs.get("output_type") + return np.full(len(X), self.mean_) + + spec = TFMSpec( + make_classifier=lambda **kw: _StubClassifier(), + make_regressor=lambda **kw: _OutputTypeAwareRegressor(), + max_train_samples=10000, + ) + df = _frame(20) + y = pd.Series(np.arange(20.0)) + fitted = fit_tfm(df, y, TaskType.REGRESSION, spec=spec, seed=0) + pred = predict_tfm(fitted, df) + assert captured["output_type"] == "median" + assert pred.shape == (20,) + + +def test_predict_regression_plain_estimator_without_output_type( + stub_tfm: TFMSpec, +) -> None: + df = _frame(20) + y = pd.Series(np.arange(20.0)) + fitted = fit_tfm(df, y, TaskType.REGRESSION, spec=stub_tfm, seed=0) + pred = predict_tfm(fitted, df) + assert pred.shape == (20,) + + +def test_predict_uses_the_callers_prediction_batch_limit() -> None: + batch_lengths: list[int] = [] + + class _BatchRecordingRegressor(_StubRegressor): + def predict(self, X: pd.DataFrame) -> np.ndarray: + batch_lengths.append(len(X)) + return super().predict(X) + + spec = TFMSpec( + make_classifier=lambda **kw: _StubClassifier(), + make_regressor=lambda **kw: _BatchRecordingRegressor(), + max_train_samples=10000, + ) + df = _frame(8) + fitted = fit_tfm( + df, + pd.Series(np.arange(8.0)), + TaskType.REGRESSION, + spec=spec, + seed=0, + max_predict_samples=3, + ) + pred = predict_tfm(fitted, df) + assert batch_lengths == [3, 3, 2] + assert pred.shape == (8,) + + +def test_fit_uses_the_tfms_own_sample_cap() -> None: + spec = TFMSpec( + make_classifier=lambda **kw: _StubClassifier(), + make_regressor=lambda **kw: _StubRegressor(), + max_train_samples=5, + ) + df = _frame(50) + y = pd.Series([0] * 25 + [1] * 25) + fitted = fit_tfm(df, y, TaskType.BINARY_CLASSIFICATION, spec=spec, seed=0) + assert fitted.estimator.n_train_ == 5 diff --git a/packages/relarena/tests/core/test_tuner.py b/packages/relarena/tests/core/test_tuner.py new file mode 100644 index 0000000..7ed70d8 --- /dev/null +++ b/packages/relarena/tests/core/test_tuner.py @@ -0,0 +1,86 @@ +"""Configuration planning and trial error reporting.""" + +import logging + +import pytest +from ConfigSpace import ConfigurationSpace, Integer + +from relarena.core.search_space import SearchSpace +from relarena.core.tuner import _concise_error, plan_configs + + +def _random_space() -> SearchSpace: + return SearchSpace( + space=ConfigurationSpace(space=[Integer("x", (1, 100))], seed=0), + default_overrides={}, + ) + + +def _grid_space() -> SearchSpace: + return SearchSpace( + fixed_grid=[{"d": 3}, {"d": 2}, {"d": 1}], default_overrides={"d": 2} + ) + + +def test_plan_configs_random_default_plus_samples() -> None: + plan = plan_configs(_random_space(), n_trials=3, seed=0) + tags = [t for t, _ in plan] + assert tags[0] == "default" and plan[0][1] == {} # the empty default comes first + assert len(plan) == 4 # default + 3 random samples + assert all("x" in cfg for _, cfg in plan[1:]) + + +def test_plan_configs_grid_uses_grid_in_order() -> None: + plan = plan_configs(_grid_space(), n_trials=99, seed=0) + configs = [c for _, c in plan] + assert configs == [{"d": 3}, {"d": 2}, {"d": 1}] # whole grid, deepest first + assert plan[1] == ("default", {"d": 2}) # default-matching entry tagged "default" + + +def test_plan_configs_grid_capped_at_n_trials_with_warning( + caplog: pytest.LogCaptureFixture, +) -> None: + with caplog.at_level(logging.WARNING): + plan = plan_configs(_grid_space(), n_trials=2, seed=0) + configs = [c for _, c in plan] + # budget < grid -> keep the first n_trials (the deepest-first grid keeps d=3, d=2) + assert configs == [{"d": 3}, {"d": 2}] + # default still tagged when it survives the cap + assert plan[1] == ("default", {"d": 2}) + warnings = [r.getMessage() for r in caplog.records if r.levelno == logging.WARNING] + assert any("dropping 1" in m for m in warnings) + + +def test_plan_configs_grid_within_budget_logs_no_warning( + caplog: pytest.LogCaptureFixture, +) -> None: + with caplog.at_level(logging.WARNING): + plan_configs(_grid_space(), n_trials=3, seed=0) # exactly fits + assert not [r for r in caplog.records if r.levelno == logging.WARNING] + + +def test_default_overrides_not_in_fixed_grid_warns( + caplog: pytest.LogCaptureFixture, +) -> None: + with caplog.at_level(logging.WARNING): + SearchSpace(fixed_grid=[{"d": 3}, {"d": 1}], default_overrides={"d": 2}) + warnings = [r.getMessage() for r in caplog.records if r.levelno == logging.WARNING] + assert any("not in the fixed_grid" in m for m in warnings) + + +def test_default_overrides_in_fixed_grid_does_not_warn( + caplog: pytest.LogCaptureFixture, +) -> None: + with caplog.at_level(logging.WARNING): + SearchSpace(fixed_grid=[{"d": 3}, {"d": 2}], default_overrides={"d": 2}) + assert not [r for r in caplog.records if r.levelno == logging.WARNING] + + +def test__concise_error__one_line_type_message_and_raise_site() -> None: + try: + raise ValueError("bad\nstuff") # multi-line message must collapse to one line + except ValueError as exc: + summary = _concise_error(exc) + assert "\n" not in summary + assert summary.startswith("ValueError: bad stuff (") + assert "test_tuner.py:" in summary # innermost frame = where it was raised diff --git a/packages/relarena/tests/userdb/test_ingest.py b/packages/relarena/tests/core/userdb/test_ingest.py similarity index 98% rename from packages/relarena/tests/userdb/test_ingest.py rename to packages/relarena/tests/core/userdb/test_ingest.py index c66998e..7148f59 100644 --- a/packages/relarena/tests/userdb/test_ingest.py +++ b/packages/relarena/tests/core/userdb/test_ingest.py @@ -7,7 +7,7 @@ import pandas as pd import pytest -from relarena.userdb.ingest import DatabaseSpec, TableSource, build_dataset +from relarena.core.userdb.ingest import DatabaseSpec, TableSource, build_dataset def test__build_dataset__pkey_maps__capture_original_to_reindexed_ids( diff --git a/packages/relarena/tests/userdb/test_predict.py b/packages/relarena/tests/core/userdb/test_predict.py similarity index 97% rename from packages/relarena/tests/userdb/test_predict.py rename to packages/relarena/tests/core/userdb/test_predict.py index 83219b3..167ea02 100644 --- a/packages/relarena/tests/userdb/test_predict.py +++ b/packages/relarena/tests/core/userdb/test_predict.py @@ -8,7 +8,7 @@ import pytest from relbench.base import Database, Table -from relarena.userdb.predict import make_prediction_table +from relarena.core.userdb.predict import make_prediction_table def test__make_prediction_table__explicit_ids__drops_absent_entities() -> None: diff --git a/packages/relarena/tests/userdb/test_task.py b/packages/relarena/tests/core/userdb/test_task.py similarity index 98% rename from packages/relarena/tests/userdb/test_task.py rename to packages/relarena/tests/core/userdb/test_task.py index 83fe517..2cd7d14 100644 --- a/packages/relarena/tests/userdb/test_task.py +++ b/packages/relarena/tests/core/userdb/test_task.py @@ -6,7 +6,7 @@ import pytest from relbench.base import Database, Table -from relarena.userdb.task import UserEntityTask +from relarena.core.userdb.task import UserEntityTask def _drivers_db() -> Database: diff --git a/packages/relarena/tests/fixtures/cached_model.py b/packages/relarena/tests/fixtures/cached_model.py index bc2715c..af35d73 100644 --- a/packages/relarena/tests/fixtures/cached_model.py +++ b/packages/relarena/tests/fixtures/cached_model.py @@ -1,11 +1,11 @@ -"""A novel model-owned cache with formats unknown to relarena.cache.""" +"""A novel model-owned cache with formats unknown to relarena.core.cache.""" from __future__ import annotations import hashlib from pathlib import Path -from relarena.cache import CacheConfig, cache_key, cached_artifact +from relarena.core.cache import CacheConfig, cache_key, cached_artifact _VERSION = 7 diff --git a/packages/relarena/tests/models/_shared/tfm/__init__.py b/packages/relarena/tests/models/_shared/tfm/__init__.py deleted file mode 100644 index dd86593..0000000 --- a/packages/relarena/tests/models/_shared/tfm/__init__.py +++ /dev/null @@ -1 +0,0 @@ -"""Tests for the shared tabular-foundation-model internals.""" diff --git a/packages/relarena/tests/models/_shared/tfm/test_tfm.py b/packages/relarena/tests/models/_shared/tfm/test_tfm.py deleted file mode 100644 index 949bac7..0000000 --- a/packages/relarena/tests/models/_shared/tfm/test_tfm.py +++ /dev/null @@ -1,387 +0,0 @@ -"""Unit tests for the TFM estimator core (via a stub TFM, no real TabPFN inference). - -A stub estimator is registered into TFM_REGISTRY so the downsample -> fit -> predict -path and the downsampler can be tested without running real TabPFN. -""" - -from __future__ import annotations - -import os -from collections.abc import Iterator -from types import SimpleNamespace - -import numpy as np -import pandas as pd -import pytest -from relbench.base import TaskType - -from relarena.models._shared.tfm import tfm -from relarena.models._shared.tfm.tfm import ( - _downsample_indices, - fit_tfm, - predict_tfm, -) - - -class _StubClassifier: - """Minimal sklearn-like classifier: learns classes_ from y, emits fixed proba.""" - - def fit(self, X: pd.DataFrame, y: np.ndarray) -> "_StubClassifier": - self.classes_ = np.unique(y) - self.n_train_ = len(X) - return self - - def predict_proba(self, X: pd.DataFrame) -> np.ndarray: - cols = np.arange(1, len(self.classes_) + 1, dtype=float) - return np.tile(cols / cols.sum(), (len(X), 1)) - - -class _StubRegressor: - def fit(self, X: pd.DataFrame, y: np.ndarray) -> "_StubRegressor": - self.mean_ = float(np.mean(y)) - self.n_train_ = len(X) - return self - - def predict(self, X: pd.DataFrame) -> np.ndarray: - return np.full(len(X), self.mean_) - - -@pytest.fixture -def stub_tfm() -> Iterator[str]: - """Register a 'stub' TFM for the duration of a test, then remove it.""" - tfm.TFM_REGISTRY["stub"] = tfm.TFMSpec( - make_classifier=lambda **kw: _StubClassifier(), - make_regressor=lambda **kw: _StubRegressor(), - max_train_samples=10_000, - ) - yield "stub" - del tfm.TFM_REGISTRY["stub"] - - -def _frame(n: int, seed: int = 0) -> pd.DataFrame: - rng = np.random.default_rng(seed) - return pd.DataFrame( - {"num": rng.normal(size=n), "cat": rng.choice(["a", "b", "c"], n)} - ) - - -# -- downsampling ------------------------------------------------------------ - - -def test_downsample_classification_caps_and_keeps_every_class() -> None: - y = np.array([0] * 25 + [1] * 25) - idx = _downsample_indices( - y, TaskType.BINARY_CLASSIFICATION, 10, np.random.default_rng(0) - ) - assert len(idx) == 10 - assert set(y[idx]) == {0, 1} # both classes survive - - -def test_downsample_regression_caps_and_is_seeded() -> None: - y = np.arange(100.0) - a = _downsample_indices(y, TaskType.REGRESSION, 20, np.random.default_rng(0)) - b = _downsample_indices(y, TaskType.REGRESSION, 20, np.random.default_rng(0)) - assert len(a) == 20 and np.array_equal(a, b) - - -# -- fit / predict ------------------------------------------------------------ - - -def test_fit_predict_roundtrip_reindexes_to_training_columns(stub_tfm: str) -> None: - # predict_tfm reindexes to the training column order, then delegates the output - # reshaping to predict_to_contract (covered in test_predict_contract.py). Here we - # only check the fit -> predict wiring survives reordered predict-time columns. - df = _frame(20) - y = pd.Series([0, 1] * 10) - fitted = fit_tfm(df, y, TaskType.BINARY_CLASSIFICATION, tfm=stub_tfm, seed=0) - pred = predict_tfm(fitted, df[["cat", "num"]]) # columns reordered on purpose - assert pred.shape == (20,) - - -def test_fit_does_not_pass_categorical_indices() -> None: - # The TFM auto-detects categoricals from the DataFrame; fit_tfm must not pre-flag - # them via categorical_features_indices. - captured: dict[str, object] = {} - - def _make(**kw: object) -> _StubClassifier: - captured.update(kw) - return _StubClassifier() - - tfm.TFM_REGISTRY["capture"] = tfm.TFMSpec( - make_classifier=_make, - make_regressor=_make, - max_train_samples=10_000, - ) - try: - df = _frame(20) # columns: ["num", "cat"] - y = pd.Series([0, 1] * 10) - fit_tfm(df, y, TaskType.BINARY_CLASSIFICATION, tfm="capture", seed=0) - assert "categorical_features_indices" not in captured - finally: - del tfm.TFM_REGISTRY["capture"] - - -def test_predict_regression_requests_median_when_supported() -> None: - # MAE is the primary regression metric; predict_tfm must request the - # MAE-optimal median from estimators that support output_type. - captured: dict[str, object] = {} - - class _OutputTypeAwareRegressor(_StubRegressor): - def predict(self, X: pd.DataFrame, **kwargs: object) -> np.ndarray: - captured["output_type"] = kwargs.get("output_type") - return np.full(len(X), self.mean_) - - tfm.TFM_REGISTRY["median"] = tfm.TFMSpec( - make_classifier=lambda **kw: _StubClassifier(), - make_regressor=lambda **kw: _OutputTypeAwareRegressor(), - max_train_samples=10_000, - ) - try: - df = _frame(20) - y = pd.Series(np.arange(20.0)) - fitted = fit_tfm(df, y, TaskType.REGRESSION, tfm="median", seed=0) - pred = predict_tfm(fitted, df) - assert captured["output_type"] == "median" - assert pred.shape == (20,) - finally: - del tfm.TFM_REGISTRY["median"] - - -def test_predict_regression_plain_estimator_without_output_type( - stub_tfm: str, -) -> None: - # Estimators without an output_type parameter (e.g. a future LimiX wrapper) - # fall back to the plain predict path. - df = _frame(20) - y = pd.Series(np.arange(20.0)) - fitted = fit_tfm(df, y, TaskType.REGRESSION, tfm=stub_tfm, seed=0) - pred = predict_tfm(fitted, df) - assert pred.shape == (20,) - - -def test_predict_uses_the_callers_prediction_batch_limit() -> None: - batch_lengths: list[int] = [] - - class _BatchRecordingRegressor(_StubRegressor): - def predict(self, X: pd.DataFrame) -> np.ndarray: - batch_lengths.append(len(X)) - return super().predict(X) - - tfm.TFM_REGISTRY["batched"] = tfm.TFMSpec( - make_classifier=lambda **kw: _StubClassifier(), - make_regressor=lambda **kw: _BatchRecordingRegressor(), - max_train_samples=10_000, - ) - try: - df = _frame(8) - fitted = fit_tfm( - df, - pd.Series(np.arange(8.0)), - TaskType.REGRESSION, - tfm="batched", - seed=0, - max_predict_samples=3, - ) - - pred = predict_tfm(fitted, df) - - assert batch_lengths == [3, 3, 2] - assert pred.shape == (8,) - finally: - del tfm.TFM_REGISTRY["batched"] - - -def test_fit_uses_the_tfms_own_sample_cap() -> None: - # Without an explicit max_train_samples, fit_tfm uses the TFM's registry cap. - tfm.TFM_REGISTRY["small"] = tfm.TFMSpec( - make_classifier=lambda **kw: _StubClassifier(), - make_regressor=lambda **kw: _StubRegressor(), - max_train_samples=5, - ) - try: - df = _frame(50) - y = pd.Series([0] * 25 + [1] * 25) - fitted = fit_tfm(df, y, TaskType.BINARY_CLASSIFICATION, tfm="small", seed=0) - assert fitted.estimator.n_train_ == 5 # spec cap applied - finally: - del tfm.TFM_REGISTRY["small"] - - -def test_tabpfn_v3_spec_has_no_text_support() -> None: - # The local v3 estimator cannot consume raw text; TabPFNRelModel's with_text - # guard rests on this staying false. - assert not tfm.TFM_REGISTRY["tabpfn-v3"].supports_text - - -def test_tabpfn_v3_api_spec_builds_the_client_estimator( - monkeypatch: pytest.MonkeyPatch, -) -> None: - import sys - from types import ModuleType - - captured: dict[str, object] = {} - - class _ApiEstimator: - def __init__(self, **kwargs: object) -> None: - captured.update(kwargs) - - module = ModuleType("tabpfn_client") - module.TabPFNClassifier = _ApiEstimator # type: ignore[attr-defined] - module.TabPFNRegressor = _ApiEstimator # type: ignore[attr-defined] - monkeypatch.setitem(sys.modules, "tabpfn_client", module) - - spec = tfm.TFM_REGISTRY["tabpfn-v3-api"] - estimator = spec.make_classifier(device="cuda", seed=7) - - assert isinstance(estimator, _ApiEstimator) - # device is server-side and never forwarded to the client constructor. - assert captured == { - "model_path": "v3_default", - "random_state": 7, - "ignore_pretraining_limits": True, - } - assert isinstance(spec.make_regressor(device="cpu", seed=7), _ApiEstimator) - # The API handles raw text natively. - assert spec.supports_text - - -def test__importing_the_model_registry__does_not_import_tabpfn() -> None: - # tabpfn ships the Prior Labs License, whose paragraph 10 obliges downstream - # attribution, so it belongs to the rdblearn extra rather than a core install. - # Importing the registry must not pull it in even where it is installed, which - # is what keeps that containment true independently of - # whichever environment the suite happens to run in. - # - # A subprocess, because this test cannot observe a clean import otherwise: the - # suite has already imported relarena.models, and re-importing it in-process would - # re-register every model and trip the registry's duplicate-name guard. - import subprocess - import sys - - result = subprocess.run( - [ - sys.executable, - "-c", - "import relarena.models, sys; " - "sys.exit(1 if 'tabpfn' in sys.modules else 0)", - ], - capture_output=True, - text=True, - ) - - assert result.returncode == 0, ( - f"importing relarena.models pulled in tabpfn\n{result.stderr}" - ) - - -def test__local_tabpfn_spec__tabpfn_stubbed__imports_it_only_when_built( - monkeypatch: pytest.MonkeyPatch, -) -> None: - # tabpfn ships in the rdblearn extra, so building the registry must not import it; - # a stub injected after import still wins, which is what proves the import is lazy. - import sys - from types import ModuleType - - captured: dict[str, object] = {} - - class _LocalEstimator: - @classmethod - def create_default_for_version( - cls, version: object, **kwargs: object - ) -> "_LocalEstimator": - captured["version"] = version - captured.update(kwargs) - return cls() - - module = ModuleType("tabpfn") - module.TabPFNClassifier = _LocalEstimator # type: ignore[attr-defined] - module.TabPFNRegressor = _LocalEstimator # type: ignore[attr-defined] - constants = ModuleType("tabpfn.constants") - constants.ModelVersion = SimpleNamespace( # type: ignore[attr-defined] - V2="v2-checkpoint", V2_5="v2.5-checkpoint", V3="v3-checkpoint" - ) - settings = ModuleType("tabpfn.settings") - settings.settings = SimpleNamespace( # type: ignore[attr-defined] - tabpfn=SimpleNamespace(max_batched_test_rows=32768) - ) - monkeypatch.setitem(sys.modules, "tabpfn", module) - monkeypatch.setitem(sys.modules, "tabpfn.constants", constants) - monkeypatch.setitem(sys.modules, "tabpfn.settings", settings) - monkeypatch.delenv("TABPFN_MAX_BATCHED_TEST_ROWS", raising=False) - - estimator = tfm.TFM_REGISTRY["tabpfn-v2.5"].make_classifier(device="cpu", seed=7) - - assert isinstance(estimator, _LocalEstimator) - assert "TABPFN_MAX_BATCHED_TEST_ROWS" not in os.environ - assert settings.settings.tabpfn.max_batched_test_rows == 32768 - assert captured == { - "version": "v2.5-checkpoint", - "device": "cpu", - "random_state": 7, - "ignore_pretraining_limits": True, - } - - -def test__local_tabpfn_spec__does_not_apply_environment_batch_limit( - monkeypatch: pytest.MonkeyPatch, -) -> None: - import sys - from types import ModuleType - - class _LocalEstimator: - @classmethod - def create_default_for_version( - cls, version: object, **kwargs: object - ) -> "_LocalEstimator": - return cls() - - module = ModuleType("tabpfn") - module.TabPFNClassifier = _LocalEstimator # type: ignore[attr-defined] - module.TabPFNRegressor = _LocalEstimator # type: ignore[attr-defined] - constants = ModuleType("tabpfn.constants") - constants.ModelVersion = SimpleNamespace( # type: ignore[attr-defined] - V2="v2-checkpoint", V2_5="v2.5-checkpoint", V3="v3-checkpoint" - ) - settings = ModuleType("tabpfn.settings") - settings.settings = SimpleNamespace( # type: ignore[attr-defined] - tabpfn=SimpleNamespace(max_batched_test_rows=32768) - ) - monkeypatch.setitem(sys.modules, "tabpfn", module) - monkeypatch.setitem(sys.modules, "tabpfn.constants", constants) - monkeypatch.setitem(sys.modules, "tabpfn.settings", settings) - monkeypatch.setenv("TABPFN_MAX_BATCHED_TEST_ROWS", "4096") - - tfm.TFM_REGISTRY["tabpfn-v2"].make_classifier(device="cpu", seed=0) - - assert os.environ["TABPFN_MAX_BATCHED_TEST_ROWS"] == "4096" - assert settings.settings.tabpfn.max_batched_test_rows == 32768 - - -def test__make_tabpfn_api__ndarray_subsample_indices__converted_to_int_lists( - monkeypatch: pytest.MonkeyPatch, -) -> None: - import json - import sys - from types import ModuleType - - captured: dict[str, object] = {} - - class _ApiEstimator: - def __init__(self, **kwargs: object) -> None: - captured.update(kwargs) - - module = ModuleType("tabpfn_client") - module.TabPFNClassifier = _ApiEstimator # type: ignore[attr-defined] - module.TabPFNRegressor = _ApiEstimator # type: ignore[attr-defined] - monkeypatch.setitem(sys.modules, "tabpfn_client", module) - - tfm.TFM_REGISTRY["tabpfn-v3-api"].make_classifier( - device="cpu", - seed=0, - n_estimators=2, - inference_config={"SUBSAMPLE_SAMPLES": [np.array([0, 2]), np.array([1, 3])]}, - ) - - config = captured["inference_config"] - assert config["SUBSAMPLE_SAMPLES"] == [[0, 2], [1, 3]] - json.dumps(config) # what tabpfn_client serializes into the request body diff --git a/packages/relarena/tests/models/dummy/test_model.py b/packages/relarena/tests/models/dummy/test_model.py index 981e7c4..6a3043b 100644 --- a/packages/relarena/tests/models/dummy/test_model.py +++ b/packages/relarena/tests/models/dummy/test_model.py @@ -13,6 +13,7 @@ import pandas as pd from relbench.base import TaskType +from relarena.core.registry import registry from relarena.models.dummy import DummyBaseline, DummyPerEntityBaseline @@ -65,8 +66,6 @@ def test_binary_predicts_positive_rate() -> None: def test_registered_under_name_constant_global() -> None: - import relarena.models # noqa: F401 (triggers registration) - from relarena.registry import registry assert "constant-global" in registry assert registry.get("constant-global") is DummyBaseline @@ -104,8 +103,6 @@ def test_per_entity_unseen_entity_falls_back_to_global() -> None: def test_per_entity_registered_under_name_constant_per_entity() -> None: - import relarena.models # noqa: F401 (triggers registration) - from relarena.registry import registry assert "constant-per-entity" in registry assert registry.get("constant-per-entity") is DummyPerEntityBaseline diff --git a/packages/relarena/tests/models/graphsage/test_model.py b/packages/relarena/tests/models/graphsage/test_model.py index f7855ec..9acea9f 100644 --- a/packages/relarena/tests/models/graphsage/test_model.py +++ b/packages/relarena/tests/models/graphsage/test_model.py @@ -17,6 +17,7 @@ from ConfigSpace import Configuration from relbench.base import TaskType +from relarena.core.registry import registry from relarena.models.graphsage import model as graphsage_mod from relarena.models.graphsage.model import ( _BATCH_SIZE, @@ -26,7 +27,6 @@ _cuda_cleanup, _run_with_oom_retry, ) -from relarena.registry import registry _SUPPORTED_TASK_TYPES = frozenset( { diff --git a/packages/relarena/tests/models/kurversc/test_model.py b/packages/relarena/tests/models/kurversc/test_model.py index 0c7aaec..2c208fe 100644 --- a/packages/relarena/tests/models/kurversc/test_model.py +++ b/packages/relarena/tests/models/kurversc/test_model.py @@ -9,9 +9,9 @@ import pytest from relbench.base import Table, TaskType -from relarena.identity import RunIdentity +from relarena.core.identity import RunIdentity +from relarena.core.registry import registry from relarena.models.kurversc import KURVERSC_DEFAULTS, KurveRSCSystem -from relarena.registry import registry from relarena.tasks import RELBENCH_V1_DATASETS, list_entity_tasks @@ -31,7 +31,6 @@ def _label_table(*, masked: bool = False) -> Table: def test__kurversc__is_registered_as_a_native_system() -> None: - import relarena.models # noqa: F401 assert registry.get("kurversc") is KurveRSCSystem assert registry.kind("kurversc") == "system" diff --git a/packages/relarena/tests/models/rdblearn/test_model.py b/packages/relarena/tests/models/rdblearn/test_model.py index 329b0b9..57cab78 100644 --- a/packages/relarena/tests/models/rdblearn/test_model.py +++ b/packages/relarena/tests/models/rdblearn/test_model.py @@ -12,11 +12,11 @@ import pytest from relbench.base import Table, TaskType -from relarena.cache import CacheConfig -from relarena.models._shared.tfm import tfm -from relarena.models.rdblearn import RDBLEARN_SPACE, RDBLearnModel +from relarena.core.cache import CacheConfig +from relarena.core.registry import registry +from relarena.core.tfm import TFMSpec +from relarena.models.rdblearn import RDBLEARN_SPACE, RDBLearnModel, tfm from relarena.models.rdblearn import model as rdblearn -from relarena.registry import registry def test_registered_under_name_rdblearn() -> None: @@ -121,7 +121,7 @@ def fake_build( monkeypatch.setitem( tfm.TFM_REGISTRY, "capture", - tfm.TFMSpec( + TFMSpec( make_classifier=lambda **kwargs: _CaptureRegressor(), make_regressor=lambda **kwargs: _CaptureRegressor(), max_train_samples=cap, @@ -137,12 +137,8 @@ def fake_build( assert seen["history_table"] is train assert seen["keep_anchor_columns"] is True assert seen["cache"] == CacheConfig(tmp_path, "raise") - idx = tfm._downsample_indices( - train.df["y"].to_numpy(), - TaskType.REGRESSION, - cap, - np.random.default_rng(seed), - ) estimator = model._fitted.estimator - assert estimator.X["f"].tolist() == (train.df.iloc[idx]["y"] * 2).tolist() - assert estimator.y.tolist() == train.df.iloc[idx]["y"].tolist() + assert len(estimator.X) == cap + assert len(set(estimator.y)) == cap + assert set(estimator.y) <= set(train.df["y"]) + assert estimator.X["f"].tolist() == (estimator.y * 2).tolist() diff --git a/packages/relarena/tests/models/rdblearn/test_tfm.py b/packages/relarena/tests/models/rdblearn/test_tfm.py new file mode 100644 index 0000000..80a0432 --- /dev/null +++ b/packages/relarena/tests/models/rdblearn/test_tfm.py @@ -0,0 +1,126 @@ +"""RDBLearn TFM recipes and local backend construction.""" + +from __future__ import annotations + +import os +import subprocess +import sys +from types import ModuleType, SimpleNamespace + +import pytest + +from relarena.models.rdblearn import tfm + + +def test__rdblearn_tfm__import__registers_model_without_loading_backends() -> None: + code = """ +import sys +from relarena.models.rdblearn import RDBLearnModel, tfm +from relarena.core.registry import registry + +prefixes = ('tabpfn', 'tabpfn_client', 'fastdfs') +loaded = [name for name in sys.modules + if any(name == p or name.startswith(p + '.') for p in prefixes)] +assert not loaded, loaded +assert registry.get('rdblearn') is RDBLearnModel +assert set(tfm.TFM_REGISTRY) == {'tabpfn-v2', 'tabpfn-v2.5'} +""" + subprocess.run([sys.executable, "-c", code], check=True) + + +def test__importing_the_model_registry__does_not_import_tabpfn() -> None: + # A subprocess observes registration with no backend already imported. + + result = subprocess.run( + [ + sys.executable, + "-c", + "import sys; import relarena.models; " + "sys.exit(1 if 'tabpfn' in sys.modules else 0)", + ], + capture_output=True, + text=True, + ) + + assert result.returncode == 0, ( + f"importing relarena.models pulled in tabpfn\n{result.stderr}" + ) + + +def test__local_tabpfn_spec__tabpfn_stubbed__imports_it_only_when_built( + monkeypatch: pytest.MonkeyPatch, +) -> None: + # tabpfn ships in the rdblearn extra, so building the registry must not import it; + # a stub injected after import still wins, which is what proves the import is lazy. + + captured: dict[str, object] = {} + + class _LocalEstimator: + @classmethod + def create_default_for_version( + cls, version: object, **kwargs: object + ) -> "_LocalEstimator": + captured["version"] = version + captured.update(kwargs) + return cls() + + module = ModuleType("tabpfn") + module.TabPFNClassifier = _LocalEstimator # type: ignore[attr-defined] + module.TabPFNRegressor = _LocalEstimator # type: ignore[attr-defined] + constants = ModuleType("tabpfn.constants") + constants.ModelVersion = SimpleNamespace( # type: ignore[attr-defined] + V2="v2-checkpoint", V2_5="v2.5-checkpoint", V3="v3-checkpoint" + ) + settings = ModuleType("tabpfn.settings") + settings.settings = SimpleNamespace( # type: ignore[attr-defined] + tabpfn=SimpleNamespace(max_batched_test_rows=32768) + ) + monkeypatch.setitem(sys.modules, "tabpfn", module) + monkeypatch.setitem(sys.modules, "tabpfn.constants", constants) + monkeypatch.setitem(sys.modules, "tabpfn.settings", settings) + monkeypatch.delenv("TABPFN_MAX_BATCHED_TEST_ROWS", raising=False) + + estimator = tfm.TFM_REGISTRY["tabpfn-v2.5"].make_classifier(device="cpu", seed=7) + + assert isinstance(estimator, _LocalEstimator) + assert "TABPFN_MAX_BATCHED_TEST_ROWS" not in os.environ + assert settings.settings.tabpfn.max_batched_test_rows == 32768 + assert captured == { + "version": "v2.5-checkpoint", + "device": "cpu", + "random_state": 7, + "ignore_pretraining_limits": True, + } + + +def test__local_tabpfn_spec__does_not_apply_environment_batch_limit( + monkeypatch: pytest.MonkeyPatch, +) -> None: + + class _LocalEstimator: + @classmethod + def create_default_for_version( + cls, version: object, **kwargs: object + ) -> "_LocalEstimator": + return cls() + + module = ModuleType("tabpfn") + module.TabPFNClassifier = _LocalEstimator # type: ignore[attr-defined] + module.TabPFNRegressor = _LocalEstimator # type: ignore[attr-defined] + constants = ModuleType("tabpfn.constants") + constants.ModelVersion = SimpleNamespace( # type: ignore[attr-defined] + V2="v2-checkpoint", V2_5="v2.5-checkpoint", V3="v3-checkpoint" + ) + settings = ModuleType("tabpfn.settings") + settings.settings = SimpleNamespace( # type: ignore[attr-defined] + tabpfn=SimpleNamespace(max_batched_test_rows=32768) + ) + monkeypatch.setitem(sys.modules, "tabpfn", module) + monkeypatch.setitem(sys.modules, "tabpfn.constants", constants) + monkeypatch.setitem(sys.modules, "tabpfn.settings", settings) + monkeypatch.setenv("TABPFN_MAX_BATCHED_TEST_ROWS", "4096") + + tfm.TFM_REGISTRY["tabpfn-v2"].make_classifier(device="cpu", seed=0) + + assert os.environ["TABPFN_MAX_BATCHED_TEST_ROWS"] == "4096" + assert settings.settings.tabpfn.max_batched_test_rows == 32768 diff --git a/packages/relarena/tests/models/relgnn/test_model.py b/packages/relarena/tests/models/relgnn/test_model.py index 355c1b9..4b8840b 100644 --- a/packages/relarena/tests/models/relgnn/test_model.py +++ b/packages/relarena/tests/models/relgnn/test_model.py @@ -18,8 +18,9 @@ from relbench.base import TaskType import relarena.models.relgnn.preprocessing as preprocessing -from relarena.cache import CacheConfig -from relarena.identity import RunIdentity +from relarena.core.cache import CacheConfig +from relarena.core.identity import RunIdentity +from relarena.core.registry import registry from relarena.models._shared.gnn.graph_cache import DBGraphCache from relarena.models.relgnn._vendor.atomic_routes import get_atomic_routes from relarena.models.relgnn.model import ( @@ -34,7 +35,6 @@ graph_cache_key, load_graph, ) -from relarena.registry import registry # -- atomic routes (pure topology; no heavy deps) --------------------------- diff --git a/packages/relarena/tests/models/relgnn/test_warm_cache.py b/packages/relarena/tests/models/relgnn/test_warm_cache.py index 7692820..ee40ade 100644 --- a/packages/relarena/tests/models/relgnn/test_warm_cache.py +++ b/packages/relarena/tests/models/relgnn/test_warm_cache.py @@ -7,8 +7,8 @@ import pytest -from relarena.cache import CacheConfig -from relarena.identity import RunIdentity +from relarena.core.cache import CacheConfig +from relarena.core.identity import RunIdentity from relarena.models.relgnn import warm_cache diff --git a/packages/relarena/tests/models/relgt/test_model.py b/packages/relarena/tests/models/relgt/test_model.py index 71b9ace..0199a0a 100644 --- a/packages/relarena/tests/models/relgt/test_model.py +++ b/packages/relarena/tests/models/relgt/test_model.py @@ -18,6 +18,8 @@ import pytest from relbench.base import TaskType +from relarena.core.registry import registry +from relarena.core.search_space import TaskStats from relarena.models.relgt.model import ( _DEFAULT_CONFIG, _LARGE_NODE_THRESHOLD, @@ -27,8 +29,6 @@ _schedule, relgt_search_space, ) -from relarena.registry import registry -from relarena.search_space import TaskStats _GRID_COMBOS = {(ly, dr) for ly in (1, 4, 8) for dr in (0.3, 0.4, 0.5)} @@ -37,7 +37,6 @@ def test__registry__relgt_registered_with_factory() -> None: - import relarena.models # noqa: F401 (triggers registration) from relarena.models.relgt import RelGTModel assert "relgt" in registry diff --git a/packages/relarena/tests/models/relgt/test_tokenize.py b/packages/relarena/tests/models/relgt/test_tokenize.py index b1e8154..c4cd84b 100644 --- a/packages/relarena/tests/models/relgt/test_tokenize.py +++ b/packages/relarena/tests/models/relgt/test_tokenize.py @@ -25,8 +25,8 @@ from relbench.modeling.graph import get_node_train_table_input # noqa: E402 from torch_geometric.data import HeteroData # noqa: E402 -from relarena.cache import CacheConfig # noqa: E402 -from relarena.identity import RunIdentity # noqa: E402 +from relarena.core.cache import CacheConfig # noqa: E402 +from relarena.core.identity import RunIdentity # noqa: E402 from relarena.models.relgt.tokenize import ( # noqa: E402 RelGTTokens, _require_persistent_identity, diff --git a/packages/relarena/tests/models/relgt/test_warm_cache.py b/packages/relarena/tests/models/relgt/test_warm_cache.py index cbe6bc4..b3db582 100644 --- a/packages/relarena/tests/models/relgt/test_warm_cache.py +++ b/packages/relarena/tests/models/relgt/test_warm_cache.py @@ -7,8 +7,8 @@ import pytest -from relarena.cache import CacheConfig -from relarena.identity import RunIdentity +from relarena.core.cache import CacheConfig +from relarena.core.identity import RunIdentity from relarena.models.relgt import warm_cache diff --git a/packages/relarena/tests/models/rt/test_model.py b/packages/relarena/tests/models/rt/test_model.py index 3b7c055..2cd7f0d 100644 --- a/packages/relarena/tests/models/rt/test_model.py +++ b/packages/relarena/tests/models/rt/test_model.py @@ -18,10 +18,10 @@ import yaml from relbench.base import Table, TaskType +from relarena.core.registry import registry from relarena.models.rt import RTPluRelSystem from relarena.models.rt import config as cfg from relarena.models.rt.export import TASK_DIR, _write_dataset_dir, target_stats -from relarena.registry import registry def test__registry__rt_plurel__is_registered_as_a_system() -> None: @@ -568,7 +568,7 @@ def test__embed__identical_text_is_embedded_once_and_linked( import sys from types import ModuleType - from relarena.cache import CacheConfig + from relarena.core.cache import CacheConfig from relarena.models.rt import export calls = [] @@ -623,7 +623,7 @@ def test__embed__different_text_is_not_shared( import sys from types import ModuleType - from relarena.cache import CacheConfig + from relarena.core.cache import CacheConfig from relarena.models.rt import export calls = [] diff --git a/packages/relarena/tests/models/tabpfn_rel/test_context.py b/packages/relarena/tests/models/tabpfn_rel/test_context.py index 2552986..8858c55 100644 --- a/packages/relarena/tests/models/tabpfn_rel/test_context.py +++ b/packages/relarena/tests/models/tabpfn_rel/test_context.py @@ -9,7 +9,8 @@ import pytest from relbench.base import TaskType -from relarena.models._shared.tfm import tfm +from relarena.core.tfm import TFMSpec +from relarena.models.tabpfn_rel import tfm from relarena.models.tabpfn_rel.context import ( ContextStrategy, HardPoolContext, @@ -113,7 +114,7 @@ def _make(**kw: object) -> _StubClassifier: captured.update(kw) return _StubClassifier() - tfm.TFM_REGISTRY["capture"] = tfm.TFMSpec( + tfm.TFM_REGISTRY["capture"] = TFMSpec( make_classifier=_make, make_regressor=_make, max_train_samples=10 ) try: diff --git a/packages/relarena/tests/models/tabpfn_rel/test_features.py b/packages/relarena/tests/models/tabpfn_rel/test_features.py index 44f4ed8..32e9a50 100644 --- a/packages/relarena/tests/models/tabpfn_rel/test_features.py +++ b/packages/relarena/tests/models/tabpfn_rel/test_features.py @@ -2,7 +2,7 @@ Calendar / history-lag are pure-pandas and tested directly. The text path passes raw anchor strings through (the estimator handles them; see -`models/_shared/tfm/tfm.py`), so the text tests exercise column detection, anchor +`relarena.models.tabpfn_rel.tfm`), so the text tests exercise column detection, anchor lookup, the raw pass-through, and the estimator overrides the pipeline emits. """ diff --git a/packages/relarena/tests/models/tabpfn_rel/test_model.py b/packages/relarena/tests/models/tabpfn_rel/test_model.py index 595406b..002d3a3 100644 --- a/packages/relarena/tests/models/tabpfn_rel/test_model.py +++ b/packages/relarena/tests/models/tabpfn_rel/test_model.py @@ -17,9 +17,12 @@ import pytest from relbench.base import Table, TaskType -from relarena.cache import CacheConfig -from relarena.models._shared.tfm import tfm +from relarena.core.cache import CacheConfig +from relarena.core.registry import registry +from relarena.core.tfm import TFMSpec +from relarena.models import _register_builtin_models as discover_models from relarena.models.tabpfn_rel import model as model_mod +from relarena.models.tabpfn_rel import tfm from relarena.models.tabpfn_rel.context import hard_pool_subsample_indices from relarena.models.tabpfn_rel.model import ( TABPFN_REL_CLIENT_SPACE, @@ -28,7 +31,6 @@ TabPFNRelLocalModel, TabPFNRelModel, ) -from relarena.registry import registry _N = 100 @@ -37,6 +39,7 @@ def test__local_space__default_is_the_validated_config() -> None: + discover_models() assert registry.get("tabpfn-rel-local") is TabPFNRelLocalModel assert registry.search_space("tabpfn-rel-local") is TABPFN_REL_LOCAL_SPACE default = TABPFN_REL_LOCAL_SPACE.default_overrides @@ -59,6 +62,7 @@ def test__local_space__default_is_the_validated_config() -> None: def test__client_model__registered_with_client_tfm_and_text() -> None: + discover_models() assert registry.get("tabpfn-rel-client") is TabPFNRelClientModel assert registry.search_space("tabpfn-rel-client") is TABPFN_REL_CLIENT_SPACE default = TABPFN_REL_CLIENT_SPACE.default_overrides @@ -98,6 +102,7 @@ def test__knobs__assemble_calendar_history_text() -> None: class _StubClassifier: def fit(self, X: pd.DataFrame, y: np.ndarray) -> "_StubClassifier": self.classes_ = np.unique(y) + self.y_ = y.copy() self.n_train_ = len(X) self.cols_ = list(X.columns) self.X_ = X.reset_index(drop=True) @@ -117,7 +122,7 @@ def _make(**kw: object) -> _StubClassifier: captured.update(kw) return _StubClassifier() - tfm.TFM_REGISTRY["capture"] = tfm.TFMSpec( + tfm.TFM_REGISTRY["capture"] = TFMSpec( make_classifier=_make, make_regressor=_make, max_train_samples=10, @@ -193,7 +198,7 @@ def test__random_strategy__caps_via_base_fit_tfm(capture: dict[str, object]) -> assert "inference_config" not in capture # no pool overrides -def test__random_strategy__fits_the_rows_fit_tfm_downsamples_to( +def test__random_strategy__fits_aligned_features_and_labels( capture: dict[str, object], ) -> None: table = _train_table() @@ -201,15 +206,13 @@ def test__random_strategy__fits_the_rows_fit_tfm_downsamples_to( model.fit(_task(), db=None, train_table=table, val_table=None, seed=7) full_feat, _ = _dfs_stub(None, None, table) - idx = tfm._downsample_indices( - table.df["label"].to_numpy(), - TaskType.BINARY_CLASSIFICATION, - 10, - np.random.default_rng(7), - ) + estimator = model._fitted.estimator + idx = (estimator.X_["num"].to_numpy() / 2).astype(int) + assert len(idx) == 10 and len(set(idx)) == 10 pd.testing.assert_frame_equal( - model._fitted.estimator.X_, full_feat.iloc[idx].reset_index(drop=True) + estimator.X_, full_feat.iloc[idx].reset_index(drop=True) ) + np.testing.assert_array_equal(estimator.y_, table.df.iloc[idx]["label"]) def test__hard_pool__forwards_indices_and_fits_the_context_union( diff --git a/packages/relarena/tests/models/tabpfn_rel/test_tfm.py b/packages/relarena/tests/models/tabpfn_rel/test_tfm.py new file mode 100644 index 0000000..55820f0 --- /dev/null +++ b/packages/relarena/tests/models/tabpfn_rel/test_tfm.py @@ -0,0 +1,75 @@ +"""TabPFN-Rel TFM recipes and hosted backend construction.""" + +from __future__ import annotations + +import json +import sys +from types import ModuleType + +import numpy as np +import pytest + +from relarena.models.tabpfn_rel import tfm + + +def test_tabpfn_v3_spec_has_no_text_support() -> None: + # The local backend receives typed numeric and categorical features only. + assert not tfm.TFM_REGISTRY["tabpfn-v3"].supports_text + + +def test_tabpfn_v3_api_spec_builds_the_client_estimator( + monkeypatch: pytest.MonkeyPatch, +) -> None: + + captured: dict[str, object] = {} + + class _ApiEstimator: + def __init__(self, **kwargs: object) -> None: + captured.update(kwargs) + + module = ModuleType("tabpfn_client") + module.TabPFNClassifier = _ApiEstimator # type: ignore[attr-defined] + module.TabPFNRegressor = _ApiEstimator # type: ignore[attr-defined] + monkeypatch.setitem(sys.modules, "tabpfn_client", module) + + spec = tfm.TFM_REGISTRY["tabpfn-v3-api"] + estimator = spec.make_classifier(device="cuda", seed=7) + + assert isinstance(estimator, _ApiEstimator) + # device is server-side and never forwarded to the client constructor. + assert captured == { + "model_path": "v3_default", + "random_state": 7, + "ignore_pretraining_limits": True, + } + assert isinstance(spec.make_regressor(device="cpu", seed=7), _ApiEstimator) + assert captured["model_path"] == "v3_default" + # The API handles raw text natively. + assert spec.supports_text + + +def test__make_tabpfn_api__ndarray_subsample_indices__converted_to_int_lists( + monkeypatch: pytest.MonkeyPatch, +) -> None: + + captured: dict[str, object] = {} + + class _ApiEstimator: + def __init__(self, **kwargs: object) -> None: + captured.update(kwargs) + + module = ModuleType("tabpfn_client") + module.TabPFNClassifier = _ApiEstimator # type: ignore[attr-defined] + module.TabPFNRegressor = _ApiEstimator # type: ignore[attr-defined] + monkeypatch.setitem(sys.modules, "tabpfn_client", module) + + tfm.TFM_REGISTRY["tabpfn-v3-api"].make_classifier( + device="cpu", + seed=0, + n_estimators=2, + inference_config={"SUBSAMPLE_SAMPLES": [np.array([0, 2]), np.array([1, 3])]}, + ) + + config = captured["inference_config"] + assert config["SUBSAMPLE_SAMPLES"] == [[0, 2], [1, 3]] + json.dumps(config) # what tabpfn_client serializes into the request body diff --git a/packages/relarena/tests/models/test_discovery.py b/packages/relarena/tests/models/test_discovery.py index 383e521..9e2bd39 100644 --- a/packages/relarena/tests/models/test_discovery.py +++ b/packages/relarena/tests/models/test_discovery.py @@ -1,141 +1,110 @@ -"""Tests for the model auto-discovery in `relarena.models`. - -That each model lands in the registry is covered by the per-model tests; these -cover the scan itself — what it skips, and which import failures it tolerates. -The `importlib` name is stubbed inside the package namespace rather than patching -the real module, so nothing else importing during the test is affected. -""" +"""Automatic baseline discovery, decorator registration and backend isolation.""" from __future__ import annotations -import json -import logging -import pkgutil import subprocess import sys -from types import SimpleNamespace +from pathlib import Path import pytest -import relarena.models as models_pkg - - -def test__import_models__does_not_import_preprocessing_warmers() -> None: - """Model registration must not require any preprocessing-format extra.""" - code = ( - "import sys; import relarena.models; " - "assert 'relarena.models.relgnn.warm_cache' not in sys.modules; " - "assert 'relarena.models.relgt.warm_cache' not in sys.modules" - ) - subprocess.run([sys.executable, "-c", code], check=True) - - -_REGISTRY_NAMES = """ -import json +_BASELINES = { + "tabpfn-rel-local", + "tabpfn-rel-client", + "constant-global", + "constant-per-entity", + "graphsage", + "kurversc", + "lightgbm", + "rdblearn", + "relgnn", + "relgnn-es", + "relgt", + "rt-plurel", +} + + +def test_discovery_without_optional_backends() -> None: + code = f""" +import importlib.abc import sys -if "--without-dfs-extra" in sys.argv: - - class _Hidden: - '''Make `fastdfs` unimportable, as on an install without the DFS extra.''' - - def find_spec(self, name, path=None, target=None): - if name == "fastdfs" or name.startswith("fastdfs."): - raise ModuleNotFoundError(f"No module named {name!r}", name=name) - return None - - sys.meta_path.insert(0, _Hidden()) - -import relarena.models # noqa: F401 - importing runs the registration scan -from relarena import registry - -print(json.dumps(sorted(registry.names()))) +class BlockBackends(importlib.abc.MetaPathFinder): + def find_spec(self, fullname, path=None, target=None): + if fullname.split('.')[0] in {{ + 'fastdfs', 'tabpfn', 'tabpfn_client', 'torch_geometric', + 'torch_frame', 'relational_transformer', 'graphreduce', 'lightgbm', + }}: + raise ModuleNotFoundError(fullname, name=fullname) + return None + +sys.meta_path.insert(0, BlockBackends()) +from relarena.core.registry import registry +assert not registry.names() +import relarena.models +assert set(registry.names()) == {_BASELINES!r} +assert registry.kind('kurversc') == registry.kind('rt-plurel') == 'system' +before = list(registry) +relarena.models._register_builtin_models() +assert list(registry) == before +assert not any(name.endswith('.warm_cache') for name in sys.modules) """ + subprocess.run([sys.executable, "-c", code], check=True) -def _registered_names(*args: str) -> list[str]: - """Names in the registry of a fresh interpreter, run with `args`.""" - result = subprocess.run( - [sys.executable, "-c", _REGISTRY_NAMES, *args], - check=True, - capture_output=True, - text=True, - ) - return json.loads(result.stdout.splitlines()[-1]) - - -def test__register_builtin_models__without_the_dfs_extra__same_models() -> None: - """Registration is dep-free: an absent extra must not drop a model. - - The dev group installs `fastdfs`, so a module-scope import of it is invisible - to the rest of the suite — it surfaces only on an install without the DFS - extra, where the scan reads the ImportError as an absent optional dep and the - model goes quietly missing. - """ - assert _registered_names("--without-dfs-extra") == _registered_names() - - -def _record_imports(monkeypatch: pytest.MonkeyPatch, fail: dict[str, str]) -> list[str]: - """Stub out module importing; return the list of names the scan asks for. - - `fail` maps a model name to the missing-module name its import should raise - `ModuleNotFoundError` for. - """ - asked: list[str] = [] - - def _import(target: str) -> None: - name = target.rsplit(".", 1)[-1] - asked.append(name) - if name in fail: - raise ModuleNotFoundError( - f"No module named {fail[name]!r}", name=fail[name] - ) - - monkeypatch.setattr(models_pkg, "importlib", SimpleNamespace(import_module=_import)) - return asked - - -def test__register_builtin_models__scans_models_and_skips_private_names( - monkeypatch: pytest.MonkeyPatch, -) -> None: - asked = _record_imports(monkeypatch, fail={}) - models_pkg._register_builtin_models() - - # The real package layout is walked, so this pins the actual exclusions: - # shared helpers and vendored upstream code must not be imported as models. - assert "lightgbm" in asked - assert not [name for name in asked if name.startswith("_")] - assert "_shared" not in asked - - -def test__register_builtin_models__missing_third_party_dep__skips_that_model( - monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture +@pytest.mark.parametrize("as_package", [False, True]) +def test_new_model_is_discovered_without_a_catalogue( + tmp_path: Path, as_package: bool ) -> None: - # An absent per-model extra is a normal install, not a broken one. Fail the - # first model the scan reaches and assert a *later* one still gets imported: - # asserting on an earlier name would pass even if a skip aborted the loop. - scanned = [ - name - for _finder, name, _is_pkg in pkgutil.iter_modules(models_pkg.__path__) - if not name.startswith("_") - ] - first, last = scanned[0], scanned[-1] - asked = _record_imports(monkeypatch, fail={first: "a_missing_extra"}) - - with caplog.at_level(logging.INFO, logger=models_pkg.__name__): - models_pkg._register_builtin_models() - - assert first in asked # attempted - assert last in asked # and the scan carried on past the failure - assert f"Skipping model {first!r}" in caplog.text - + package = tmp_path / "extra_model" + if as_package: + package.mkdir() + module = package / "__init__.py" + else: + module = package.with_suffix(".py") + module.write_text( + "from relarena.core.model import RelArenaModel\n" + "from relarena.core.registry import register_model\n" + "from relarena.core.search_space import SearchSpace\n" + "@register_model(search_space=SearchSpace(default_overrides={}))\n" + "class ExtraModel(RelArenaModel):\n" + " name = 'extra-model'\n" + ) + ignored = tmp_path / "_private_helper" + ignored.mkdir() + (ignored / "__init__.py").write_text( + "raise AssertionError('private package loaded')" + ) + code = f""" +import relarena.models +from relarena.core.registry import registry +relarena.models.__path__.append({str(tmp_path)!r}) +relarena.models._register_builtin_models() +from relarena.models.extra_model import ExtraModel +assert registry.get('extra-model') is ExtraModel +assert registry.search_space('extra-model').default_overrides == {{}} +relarena.models._register_builtin_models() +assert registry.get('extra-model') is ExtraModel +assert len(registry) == 13 +""" + subprocess.run([sys.executable, "-c", code], check=True) -def test__register_builtin_models__missing_relarena_module__raises( - monkeypatch: pytest.MonkeyPatch, -) -> None: - # A missing internal module is a defect; swallowing it would drop a model - # from the registry silently. - _record_imports(monkeypatch, fail={"dummy": "relarena.does_not_exist"}) - with pytest.raises(ModuleNotFoundError, match="relarena.does_not_exist"): - models_pkg._register_builtin_models() +@pytest.mark.parametrize("missing", ["optional_backend", "relarena.broken_internal"]) +def test_import_failures_propagate(tmp_path: Path, missing: str) -> None: + package = tmp_path / "broken_model" + package.mkdir() + (package / "__init__.py").write_text( + f"raise ModuleNotFoundError('broken model import', name={missing!r})\n" + ) + code = f""" +import relarena.models +relarena.models.__path__.append({str(tmp_path)!r}) +try: + relarena.models._register_builtin_models() +except ModuleNotFoundError as exc: + assert exc.name == {missing!r} +else: + raise AssertionError('Discovery hid a broken model') +""" + subprocess.run([sys.executable, "-c", code], check=True) diff --git a/packages/relarena/tests/test_cache_extension.py b/packages/relarena/tests/test_cache_extension.py index de5d656..910219e 100644 --- a/packages/relarena/tests/test_cache_extension.py +++ b/packages/relarena/tests/test_cache_extension.py @@ -5,7 +5,7 @@ import shutil from pathlib import Path -from relarena.cache import CacheConfig +from relarena.core.cache import CacheConfig from tests.fixtures.cached_model import novel_binary, novel_directory, warm_cache diff --git a/packages/relarena/tests/test_cli.py b/packages/relarena/tests/test_cli.py index e640577..f62257f 100644 --- a/packages/relarena/tests/test_cli.py +++ b/packages/relarena/tests/test_cli.py @@ -17,7 +17,8 @@ from relbench.base import TaskType from relarena import cli -from relarena.results import SystemResult, TrialResult, summary_to_dataframe +from relarena.core.results import SystemResult, TrialResult +from relarena.results import summary_to_dataframe from relarena.runner import ExperimentSummary, SystemExperimentSummary from relarena.tasks import TaskSpec diff --git a/packages/relarena/tests/test_dataset.py b/packages/relarena/tests/test_dataset.py index 914f83d..15f4bd5 100644 --- a/packages/relarena/tests/test_dataset.py +++ b/packages/relarena/tests/test_dataset.py @@ -18,12 +18,8 @@ import pandas as pd from relbench.base import Database, Dataset, Table -from relarena.dataset import ( - InnerSplit, - OuterSplit, - RelBenchDatasetTask, - drop_noncanonical_columns, -) +from relarena.core.dataset import InnerSplit, OuterSplit +from relarena.dataset import RelBenchDatasetTask, drop_noncanonical_columns #: Five monthly event timestamps; cut the inner split after the 3rd, the outer #: (test) split after the 5th — so censoring at val should drop the last two rows. diff --git a/packages/relarena/tests/test_identity.py b/packages/relarena/tests/test_identity.py index 07ed8c7..4e325a7 100644 --- a/packages/relarena/tests/test_identity.py +++ b/packages/relarena/tests/test_identity.py @@ -1,18 +1,6 @@ -"""Tests for optional preprocessing run identities.""" +"""Recorded benchmark source identities.""" -from __future__ import annotations - -from types import SimpleNamespace - -import pandas as pd -from relbench.base import Database, Table - -from relarena.identity import ( - RunIdentity, - database_schema_fingerprint, - relbench_run_identity, - task_spec_fingerprint, -) +from relarena.identity import relbench_run_identity def test__relbench_run_identity__recorded_task__is_complete_and_stable() -> None: @@ -23,65 +11,3 @@ def test__relbench_run_identity__recorded_task__is_complete_and_stable() -> None assert len(first.dataset_fingerprint.split("-")) == 2 assert first.task_fingerprint is not None assert len(first.task_fingerprint) == 16 - - -def test__run_identity__phase_copy__does_not_mutate_base() -> None: - base = RunIdentity("dataset", "db", "task", "labels") - assert base.for_phase("inner").phase == "inner" - assert base.phase is None - - -def _database(dtype: str = "int64") -> Database: - return Database( - { - "entities": Table( - df=pd.DataFrame({"id": pd.Series([1, 2], dtype=dtype)}), - fkey_col_to_pkey_table={}, - pkey_col="id", - time_col=None, - ) - } - ) - - -def test__database_schema_fingerprint__ignores_rows_but_guards_dtype() -> None: - same_schema = Database( - { - "entities": Table( - df=pd.DataFrame({"id": pd.Series([9], dtype="int64")}), - fkey_col_to_pkey_table={}, - pkey_col="id", - time_col=None, - ) - } - ) - assert database_schema_fingerprint(_database()) == database_schema_fingerprint( - same_schema - ) - assert database_schema_fingerprint(_database()) != database_schema_fingerprint( - _database("float64") - ) - - -def _task_spec(**overrides: object) -> SimpleNamespace: - fields = { - "entity_table": "drivers", - "entity_col": "driver_id", - "time_col": "date", - "target_col": "dnf", - "task_type": "binary_classification", - "timedelta": "30 days", - "query": "SELECT * FROM labels", - "val_timestamp": "2005-01-01", - "test_timestamp": "2005-02-01", - "num_eval_timestamps": 1, - } - fields.update(overrides) - return SimpleNamespace(**fields) - - -def test__task_spec_fingerprint__is_stable_and_guards_training_semantics() -> None: - assert task_spec_fingerprint(_task_spec()) == task_spec_fingerprint(_task_spec()) - assert task_spec_fingerprint(_task_spec()) != task_spec_fingerprint( - _task_spec(timedelta="60 days") - ) diff --git a/packages/relarena/tests/test_refit.py b/packages/relarena/tests/test_refit.py new file mode 100644 index 0000000..5726455 --- /dev/null +++ b/packages/relarena/tests/test_refit.py @@ -0,0 +1,102 @@ +"""Final benchmark refit and evaluation.""" + +from __future__ import annotations + +from pathlib import Path +from types import SimpleNamespace +from typing import Any + +import numpy as np +import pandas as pd +from relbench.base import Table, TaskType + +from relarena.core.cache import CacheConfig +from relarena.core.dataset import OuterSplit +from relarena.core.identity import RunIdentity +from relarena.core.model import RelArenaModel +from relarena.refit import refit_and_evaluate + + +def _outer_table(entities: list[int], times: list[int], ys: list[float]) -> Table: + return Table( + df=pd.DataFrame({"entity": entities, "t": times, "y": ys}), + fkey_col_to_pkey_table={"entity": "e"}, + pkey_col=None, + time_col="t", + ) + + +def _outer_split() -> OuterSplit: + return OuterSplit( + db_state=SimpleNamespace(), + cutoff=pd.Timestamp("2020-01-01"), + train_table=_outer_table([1, 2], [10, 11], [0.0, 1.0]), + val_table=_outer_table([3], [12], [2.0]), + eval_table=_outer_table([4, 5], [13, 14], [3.0, 4.0]), + ) + + +def _stub_task() -> Any: + # primary_metric reads task_type; evaluate_predictions reads + # task.metrics + task.evaluate. + return SimpleNamespace( + task_type=TaskType.REGRESSION, + metrics=[], + evaluate=lambda pred, target, metrics=None: {"mae": 0.5}, + ) + + +def _capturing_model(refit_full: bool) -> tuple[type[RelArenaModel], dict]: + captured: dict = {} + + class _M(RelArenaModel): + name = "capture" + refit_on_full_data = refit_full + + def fit(self, task, db, train_table, val_table, *, seed, time_limit=None): # noqa: ANN001, ANN202 + captured["train_y"] = list(train_table.df["y"]) + captured["val_y"] = None if val_table is None else list(val_table.df["y"]) + captured["cache"] = self.cache + captured["run_identity"] = self.run_identity + + def predict(self, task, db, table) -> np.ndarray: # noqa: ANN001 + return np.zeros(len(table.df)) + + return _M, captured + + +def test__refit_and_evaluate__full_data__fits_on_train_plus_val_no_monitor() -> None: + model_cls, captured = _capturing_model(refit_full=True) + out = refit_and_evaluate(model_cls, {}, _stub_task(), _outer_split(), seed=0) + assert captured["train_y"] == [0.0, 1.0, 2.0] # train + val union + assert captured["val_y"] is None # nothing held out to monitor + assert out["test_score"] == 0.5 + + +def test__refit_and_evaluate__best_val__fits_on_train_only_with_val_monitor() -> None: + model_cls, captured = _capturing_model(refit_full=False) + out = refit_and_evaluate(model_cls, {}, _stub_task(), _outer_split(), seed=0) + assert captured["train_y"] == [0.0, 1.0] # train only + assert captured["val_y"] == [2.0] # val passed through as the monitor set + assert out["test_score"] == 0.5 + + +def test__refit_and_evaluate__cache_config__reaches_model(tmp_path: Path) -> None: + model_cls, captured = _capturing_model(refit_full=True) + cache = CacheConfig(tmp_path, "raise") + refit_and_evaluate(model_cls, {}, _stub_task(), _outer_split(), seed=0, cache=cache) + assert captured["cache"] is cache + + +def test__refit_and_evaluate__run_identity__reaches_model() -> None: + model_cls, captured = _capturing_model(refit_full=True) + identity = RunIdentity("dataset", "db", "task", "labels", phase="outer") + refit_and_evaluate( + model_cls, + {}, + _stub_task(), + _outer_split(), + seed=0, + run_identity=identity, + ) + assert captured["run_identity"] is identity diff --git a/packages/relarena/tests/test_results.py b/packages/relarena/tests/test_results.py new file mode 100644 index 0000000..4ae41b0 --- /dev/null +++ b/packages/relarena/tests/test_results.py @@ -0,0 +1,36 @@ +"""Benchmark result tables.""" + +import numpy as np + +from relarena.core.results import TrialResult, config_id_for +from relarena.results import trials_to_dataframe + + +def test_trials_to_dataframe_drops_arrays() -> None: + t = TrialResult( + config={"a": 1}, + config_id=config_id_for({"a": 1}), + config_tag="default", + val_score=0.5, + val_pred=np.zeros(3), + ) + df = trials_to_dataframe([t]) + assert "val_pred" not in df.columns + assert df.loc[0, "val_score"] == 0.5 + + +def test_trials_to_dataframe_flattens_metrics() -> None: + t = TrialResult( + config={}, + config_id=config_id_for({}), + config_tag="default", + val_score=0.5, + test_score=0.6, + val_metrics={"mae": 0.5, "rmse": 0.7}, + test_metrics={"mae": 0.6, "rmse": 0.8}, + ) + df = trials_to_dataframe([t]) + # the metric dicts themselves are not columns; their entries are flattened. + assert "val_metrics" not in df.columns and "test_metrics" not in df.columns + assert df.loc[0, "val_mae"] == 0.5 and df.loc[0, "val_rmse"] == 0.7 + assert df.loc[0, "test_mae"] == 0.6 and df.loc[0, "test_rmse"] == 0.8 diff --git a/packages/relarena/tests/test_runner.py b/packages/relarena/tests/test_runner.py index d550f3d..cac740d 100644 --- a/packages/relarena/tests/test_runner.py +++ b/packages/relarena/tests/test_runner.py @@ -16,10 +16,10 @@ from relbench.metrics import roc_auc from relarena import runner -from relarena.cache import CacheConfig -from relarena.identity import RunIdentity -from relarena.results import TrialResult -from relarena.system import RelArenaSystem +from relarena.core.cache import CacheConfig +from relarena.core.identity import RunIdentity +from relarena.core.results import TrialResult +from relarena.core.system import RelArenaSystem _MODEL = SimpleNamespace( name="stub", supported_task_types=frozenset({TaskType.BINARY_CLASSIFICATION}) diff --git a/packages/relarena/tests/test_tuner.py b/packages/relarena/tests/test_tuner.py deleted file mode 100644 index 4b5c103..0000000 --- a/packages/relarena/tests/test_tuner.py +++ /dev/null @@ -1,204 +0,0 @@ -"""Tests for tuner helpers and the search-space config plan.""" - -from __future__ import annotations - -import logging -from pathlib import Path -from types import SimpleNamespace -from typing import Any - -import numpy as np -import pandas as pd -import pytest -from ConfigSpace import ConfigurationSpace, Integer -from relbench.base import Table, TaskType - -from relarena.cache import CacheConfig -from relarena.dataset import OuterSplit, concat_tables -from relarena.identity import RunIdentity -from relarena.model import RelArenaModel -from relarena.search_space import SearchSpace -from relarena.tuner import _concise_error, plan_configs, refit_and_evaluate - - -def _random_space() -> SearchSpace: - return SearchSpace( - space=ConfigurationSpace(space=[Integer("x", (1, 100))], seed=0), - default_overrides={}, - ) - - -def _grid_space() -> SearchSpace: - return SearchSpace( - fixed_grid=[{"d": 3}, {"d": 2}, {"d": 1}], default_overrides={"d": 2} - ) - - -def test_plan_configs_random_default_plus_samples() -> None: - plan = plan_configs(_random_space(), n_trials=3, seed=0) - tags = [t for t, _ in plan] - assert tags[0] == "default" and plan[0][1] == {} # the empty default comes first - assert len(plan) == 4 # default + 3 random samples - assert all("x" in cfg for _, cfg in plan[1:]) - - -def test_plan_configs_grid_uses_grid_in_order() -> None: - plan = plan_configs(_grid_space(), n_trials=99, seed=0) - configs = [c for _, c in plan] - assert configs == [{"d": 3}, {"d": 2}, {"d": 1}] # whole grid, deepest first - assert plan[1] == ("default", {"d": 2}) # default-matching entry tagged "default" - - -def test_plan_configs_grid_capped_at_n_trials_with_warning( - caplog: pytest.LogCaptureFixture, -) -> None: - with caplog.at_level(logging.WARNING): - plan = plan_configs(_grid_space(), n_trials=2, seed=0) - configs = [c for _, c in plan] - # budget < grid -> keep the first n_trials (the deepest-first grid keeps d=3, d=2) - assert configs == [{"d": 3}, {"d": 2}] - # default still tagged when it survives the cap - assert plan[1] == ("default", {"d": 2}) - warnings = [r.getMessage() for r in caplog.records if r.levelno == logging.WARNING] - assert any("dropping 1" in m for m in warnings) - - -def test_plan_configs_grid_within_budget_logs_no_warning( - caplog: pytest.LogCaptureFixture, -) -> None: - with caplog.at_level(logging.WARNING): - plan_configs(_grid_space(), n_trials=3, seed=0) # exactly fits - assert not [r for r in caplog.records if r.levelno == logging.WARNING] - - -def test_default_overrides_not_in_fixed_grid_warns( - caplog: pytest.LogCaptureFixture, -) -> None: - with caplog.at_level(logging.WARNING): - SearchSpace(fixed_grid=[{"d": 3}, {"d": 1}], default_overrides={"d": 2}) - warnings = [r.getMessage() for r in caplog.records if r.levelno == logging.WARNING] - assert any("not in the fixed_grid" in m for m in warnings) - - -def test_default_overrides_in_fixed_grid_does_not_warn( - caplog: pytest.LogCaptureFixture, -) -> None: - with caplog.at_level(logging.WARNING): - SearchSpace(fixed_grid=[{"d": 3}, {"d": 2}], default_overrides={"d": 2}) - assert not [r for r in caplog.records if r.levelno == logging.WARNING] - - -def test__concise_error__one_line_type_message_and_raise_site() -> None: - try: - raise ValueError("bad\nstuff") # multi-line message must collapse to one line - except ValueError as exc: - summary = _concise_error(exc) - assert "\n" not in summary - assert summary.startswith("ValueError: bad stuff (") - assert "test_tuner.py:" in summary # innermost frame = where it was raised - - -def _outer_table(entities: list[int], times: list[int], ys: list[float]) -> Table: - return Table( - df=pd.DataFrame({"entity": entities, "t": times, "y": ys}), - fkey_col_to_pkey_table={"entity": "e"}, - pkey_col=None, - time_col="t", - ) - - -def _outer_split() -> OuterSplit: - return OuterSplit( - db_state=SimpleNamespace(), - cutoff=pd.Timestamp("2020-01-01"), - train_table=_outer_table([1, 2], [10, 11], [0.0, 1.0]), - val_table=_outer_table([3], [12], [2.0]), - eval_table=_outer_table([4, 5], [13, 14], [3.0, 4.0]), - ) - - -def _stub_task() -> Any: - # primary_metric reads task_type; _evaluate reads task.metrics + task.evaluate. - return SimpleNamespace( - task_type=TaskType.REGRESSION, - metrics=[], - evaluate=lambda pred, target, metrics=None: {"mae": 0.5}, - ) - - -def _capturing_model(refit_full: bool) -> tuple[type[RelArenaModel], dict]: - captured: dict = {} - - class _M(RelArenaModel): - name = "capture" - refit_on_full_data = refit_full - - def fit(self, task, db, train_table, val_table, *, seed, time_limit=None): # noqa: ANN001, ANN202 - captured["train_y"] = list(train_table.df["y"]) - captured["val_y"] = None if val_table is None else list(val_table.df["y"]) - captured["cache"] = self.cache - captured["run_identity"] = self.run_identity - - def predict(self, task, db, table) -> np.ndarray: # noqa: ANN001 - return np.zeros(len(table.df)) - - return _M, captured - - -def test__refit_and_evaluate__full_data__fits_on_train_plus_val_no_monitor() -> None: - model_cls, captured = _capturing_model(refit_full=True) - out = refit_and_evaluate(model_cls, {}, _stub_task(), _outer_split(), seed=0) - assert captured["train_y"] == [0.0, 1.0, 2.0] # train + val union - assert captured["val_y"] is None # nothing held out to monitor - assert out["test_score"] == 0.5 - - -def test__refit_and_evaluate__best_val__fits_on_train_only_with_val_monitor() -> None: - model_cls, captured = _capturing_model(refit_full=False) - out = refit_and_evaluate(model_cls, {}, _stub_task(), _outer_split(), seed=0) - assert captured["train_y"] == [0.0, 1.0] # train only - assert captured["val_y"] == [2.0] # val passed through as the monitor set - assert out["test_score"] == 0.5 - - -def test__refit_and_evaluate__cache_config__reaches_model(tmp_path: Path) -> None: - model_cls, captured = _capturing_model(refit_full=True) - cache = CacheConfig(tmp_path, "raise") - refit_and_evaluate(model_cls, {}, _stub_task(), _outer_split(), seed=0, cache=cache) - assert captured["cache"] is cache - - -def test__refit_and_evaluate__run_identity__reaches_model() -> None: - model_cls, captured = _capturing_model(refit_full=True) - identity = RunIdentity("dataset", "db", "task", "labels", phase="outer") - refit_and_evaluate( - model_cls, - {}, - _stub_task(), - _outer_split(), - seed=0, - run_identity=identity, - ) - assert captured["run_identity"] is identity - - -def test_concat_tables_unions_rows_and_keeps_schema() -> None: - a = Table( - df=pd.DataFrame({"entity": [1, 2], "t": [10, 11], "y": [0.0, 1.0]}), - fkey_col_to_pkey_table={"entity": "users"}, - pkey_col=None, - time_col="t", - ) - b = Table( - df=pd.DataFrame({"entity": [3], "t": [12], "y": [2.0]}), - fkey_col_to_pkey_table={"entity": "users"}, - pkey_col=None, - time_col="t", - ) - c = concat_tables(a, b) - assert len(c.df) == 3 - assert list(c.df["y"]) == [0.0, 1.0, 2.0] - assert c.time_col == "t" - assert c.fkey_col_to_pkey_table == {"entity": "users"} - # inputs are untouched - assert len(a.df) == 2 and len(b.df) == 1 diff --git a/packages/relarena/tests/userdb/test_query.py b/packages/relarena/tests/userdb/test_query.py index 6b26ca9..5965bd1 100644 --- a/packages/relarena/tests/userdb/test_query.py +++ b/packages/relarena/tests/userdb/test_query.py @@ -11,10 +11,10 @@ import pandas as pd import pytest -from relarena.cache import CacheConfig -from relarena.identity import RunIdentity +from relarena.core.cache import CacheConfig +from relarena.core.identity import RunIdentity +from relarena.core.userdb.query import PredictiveQuery, PredictiveQuerySpec from relarena.userdb import relbench_v1_spec, relbench_v1_tasks -from relarena.userdb.query import PredictiveQuery, PredictiveQuerySpec _EXAMPLES = Path(__file__).resolve().parents[4] / "examples" _DB_YAML = "drivers:\n pkey: driverId\n" @@ -175,7 +175,7 @@ def test__precompute_cache__delegates_to_dfs_owner( query = _schema_only_query(data_version="v1") query._source = Mock() warm = Mock() - monkeypatch.setattr("relarena.featurization.warm_cache.warm_dfs_cache", warm) + monkeypatch.setattr("relarena.core.featurization.warm_cache.warm_dfs_cache", warm) assert query.precompute_cache(tmp_path) == tmp_path @@ -226,7 +226,7 @@ def test__predict__anchor_after_test_cutoff__warns_about_frozen_db( query._at_timestamp = pd.Timestamp("2020-02-01") query._entities = "all" predict_at = Mock(return_value=pd.DataFrame({"driverId": [], "y_pred": []})) - monkeypatch.setattr("relarena.userdb.query.predict_at", predict_at) + monkeypatch.setattr("relarena.core.userdb.query.predict_at", predict_at) with pytest.warns(UserWarning, match="feature database remains frozen"): query.predict() diff --git a/pyproject.toml b/pyproject.toml index cd3ccf2..da7d3b6 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -97,3 +97,4 @@ members = ["packages/relarena"] [tool.pytest.ini_options] testpaths = ["packages/relarena/tests"] +addopts = "--import-mode=importlib" diff --git a/workflows/smoke_feature_cache.py b/workflows/smoke_feature_cache.py index b57aee3..ce61b04 100644 --- a/workflows/smoke_feature_cache.py +++ b/workflows/smoke_feature_cache.py @@ -8,7 +8,7 @@ No TFM or GPU is involved. Run: - OMP_NUM_THREADS=1 uv run --extra rdblearn python workflows/smoke_feature_cache.py + OMP_NUM_THREADS=1 uv run --all-packages --extra rdblearn python workflows/smoke_feature_cache.py """ from __future__ import annotations @@ -17,11 +17,12 @@ import time from pathlib import Path -import relarena.featurization.dfs as dfs_mod -from relarena.cache import CacheConfig -from relarena.dataset import RelBenchDatasetTask, concat_tables -from relarena.featurization import DFS_MAX_DEPTH, build_dfs_features -from relarena.featurization.warm_cache import warm_dfs_cache +import relarena.core.featurization.dfs as dfs_mod +from relarena.core.cache import CacheConfig +from relarena.core.dataset import concat_tables +from relarena.core.featurization import DFS_MAX_DEPTH, build_dfs_features +from relarena.core.featurization.warm_cache import warm_dfs_cache +from relarena.dataset import RelBenchDatasetTask DATASET, TASK = "rel-f1", "driver-dnf" diff --git a/workflows/verify_distributions.py b/workflows/verify_distributions.py index 3d6f1e3..4196508 100644 --- a/workflows/verify_distributions.py +++ b/workflows/verify_distributions.py @@ -15,8 +15,8 @@ def main() -> None: required = { "relarena/models/VENDORED-LICENSES", "relarena/checksums/relbench_v1_checksums.json", - "relarena/userdb/database.schema.json", - "relarena/userdb/task.schema.json", + "relarena/core/userdb/database.schema.json", + "relarena/core/userdb/task.schema.json", *( str(path.relative_to("packages/relarena/src")) for path in spec_root.glob("*/*.yaml") diff --git a/workflows/verify_relbench_v1.py b/workflows/verify_relbench_v1.py index da62d06..fd64aff 100644 --- a/workflows/verify_relbench_v1.py +++ b/workflows/verify_relbench_v1.py @@ -19,14 +19,14 @@ import pandas as pd from relbench.tasks import get_task +from relarena.core.userdb.ingest import build_dataset +from relarena.core.userdb.task import UserEntityTask from relarena.dataset import drop_noncanonical_task_columns from relarena.userdb import ( materialize_relbench, relbench_v1_spec, relbench_v1_tasks, ) -from relarena.userdb.ingest import build_dataset -from relarena.userdb.task import UserEntityTask def _split_matches(want: pd.DataFrame, got: pd.DataFrame, keys: list[str]) -> bool: diff --git a/workflows/warm_feature_cache.py b/workflows/warm_feature_cache.py index dd44d7f..abf4f0d 100644 --- a/workflows/warm_feature_cache.py +++ b/workflows/warm_feature_cache.py @@ -18,7 +18,7 @@ Run from the repository root: RELARENA_CACHE_DIR=~/relarena_features \ - uv run --extra rdblearn python workflows/warm_feature_cache.py + uv run --all-packages --extra rdblearn python workflows/warm_feature_cache.py CPU is enough; run it ahead of a (GPU) eval pointing at the same store. """ @@ -29,9 +29,9 @@ import os import sys -from relarena.cache import resolve_cache_config +from relarena.core.cache import resolve_cache_config +from relarena.core.featurization.warm_cache import warm_dfs_cache from relarena.dataset import RelBenchDatasetTask -from relarena.featurization.warm_cache import warm_dfs_cache from relarena.tasks import RELBENCH_V1_DATASETS, list_entity_tasks