diff --git a/pyproject.toml b/pyproject.toml index b12a6a6..f92e900 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -198,6 +198,7 @@ strict = true # Strict mode, refer to the mypy docs for details python_version = "3.12" # Type check code assuming it will be running on the specified Python version pretty = true # Use visually nicer output in error messages files = [ # Type-check given files or directories + "scripts", "src", "tests" ] @@ -251,6 +252,7 @@ executionEnvironments = [ { root = "src/impact/cli.py", reportAny = false, reportExplicitAny = false }, { root = "src/impact/config.py", reportAny = false, reportExplicitAny = false }, { root = "src/impact/main.py", reportExplicitAny = false }, + { root = "src/impact/utils.py", reportImplicitStringConcatenation = false }, { root = "src/impact/dataset/fetcher.py", reportUnknownVariableType = false }, { root = "src/impact/features/sequence_extraction.py", reportExplicitAny = false }, { root = "src/impact/features/tensorization.py", reportAny = false, reportExplicitAny = false }, diff --git a/scripts/plots.py b/scripts/plots.py index bf810d8..f79ca3c 100644 --- a/scripts/plots.py +++ b/scripts/plots.py @@ -13,7 +13,10 @@ import polars as pl import structlog -from impact.cli import inject_impact_config +from impact.cli import ( + inject_impact_config_from_checkpoint, + inject_impact_config_overrides, +) from impact.config import ImpactConfig from impact.logger import configure_logging from impact.pipelines import DataPipeline, ImpactEvaluationPipeline @@ -138,12 +141,6 @@ def cli() -> None: @cli.command(name="plot-alpha-curve") -@click.option( - "--sasrec-checkpoint-path", - type=click.Path(exists=True, dir_okay=False, path_type=Path), - required=True, - help="Path to a trained SASRecNBR checkpoint that will be evaluated.", -) @click.option( "--log-level", type=click.Choice(list(logging.getLevelNamesMapping()), case_sensitive=False), @@ -255,7 +252,8 @@ def cli() -> None: show_default=True, help="Scale factor for PNG export. 3.0 is approx. 216 PPI - suitable for print.", ) -@inject_impact_config +@inject_impact_config_from_checkpoint +@inject_impact_config_overrides("actr", "eval", "system") def plot_alpha_curve( impact_config: ImpactConfig, sasrec_checkpoint_path: Path, diff --git a/scripts/tables.py b/scripts/tables.py index 4bd2a64..b69d535 100644 --- a/scripts/tables.py +++ b/scripts/tables.py @@ -157,8 +157,8 @@ class _VariantConfig(NamedTuple): logger.info(f"Processing {ds_size.name} dataset variant") variant_config = dataclasses.replace( impact_config, - dataset_size=ds_size, - step=variant.step, + dataset=dataclasses.replace(impact_config.dataset, size=ds_size), + sequence=dataclasses.replace(impact_config.sequence, step=variant.step), ) pipeline.config = variant_config # Switch pipeline config diff --git a/src/impact/cli.py b/src/impact/cli.py index 837f7a6..19e9363 100644 --- a/src/impact/cli.py +++ b/src/impact/cli.py @@ -1,8 +1,10 @@ from __future__ import annotations import dataclasses +from enum import Enum from functools import cached_property, wraps -from typing import TYPE_CHECKING, Any +from pathlib import Path +from typing import TYPE_CHECKING, Any, Literal, get_args, get_origin, get_type_hints import click @@ -18,6 +20,7 @@ SystemConfig, ) from impact.dataset import DatasetSize, DatasetType, InteractionType +from impact.utils import find_sibling_config if TYPE_CHECKING: from collections.abc import Callable @@ -27,8 +30,15 @@ def inject_impact_config(f: Callable[..., Any]) -> Callable[..., Any]: """Inject ImpactConfig fields as Click options and pass the instantiated config. - The decorated function must accept `impact_config: ImpactConfig` as a keyword - argument. + Adds one Click option per ImpactConfig field (grouped by nested config section), + using ImpactConfig's defaults as the option defaults. The decorated function must + accept `impact_config: ImpactConfig` as a keyword argument; it is constructed at + call time from the parsed CLI values. + + Use this for commands that configure a fresh run from scratch (e.g. training, grid + search). For commands operating on an existing checkpoint (e.g. evaluation), use + `inject_impact_config_from_checkpoint` instead, which loads the config that was + saved alongside the checkpoint rather than rebuilding it from CLI flags. """ # Define all click options, matching the defaults from ImpactConfig options = [ @@ -338,5 +348,148 @@ class _ConfigSection: kwargs: dict[str, Any] = dataclasses.field(default_factory=dict) @cached_property - def field_names(self) -> set[str]: - return {f.name for f in dataclasses.fields(self.cls)} + def field_names(self) -> list[str]: + return [f.name for f in dataclasses.fields(self.cls)] + + @cached_property + def type_hints(self) -> dict[str, Any]: + return get_type_hints(self.cls) + + +def inject_impact_config_from_checkpoint(f: Callable[..., Any]) -> Callable[..., Any]: + """Add Click option for SASRecNBR checkpoint path and load its sibling YAML config. + + Adds a required `--sasrec-checkpoint-path` option (validated to exist). + The decorated function must accept both of the following keyword arguments, which + this decorator provides: + - `sasrec_checkpoint_path: Path`: path to the checkpoint parsed from the CLI option. + - `impact_config: ImpactConfig`: config loaded via `utils.find_sibling_config()` and + `ImpactConfig.from_yaml()` from the checkpoint's sibling config file. + + Use for commands operating on an existing trained model (e.g., IMPACT evaluation + or plotting). Combine with `cli.inject_impact_config_overrides(...)` to allow CLI + overrides of specific config sections (e.g. eval, actr) on top of the loaded config. + For commands that build a fresh config from CLI flags, use `inject_impact_config` + instead. + """ + options = [ + click.option( + "--sasrec-checkpoint-path", + type=click.Path(exists=True, dir_okay=False, path_type=Path), + required=True, + help="Path to a trained SASRecNBR checkpoint that will be evaluated.", + ), + ] + + # Apply Click options to the function in reverse order to maintain defined ordering + for option in reversed(options): + f = option(f) + + @wraps(f) + def _wrapper(*args: Any, **kwargs: Any) -> Any: + checkpoint_path: Path = kwargs["sasrec_checkpoint_path"] + kwargs["impact_config"] = ImpactConfig.from_yaml( + path=find_sibling_config(checkpoint_path) + ) + return f(*args, **kwargs) + + return _wrapper + + +_OVERRIDE_SECTIONS: dict[str, _ConfigSection] = { + "actr": _ConfigSection(ACTRConfig), + "eval": _ConfigSection(EvalConfig), + "system": _ConfigSection(SystemConfig), +} + + +def inject_impact_config_overrides( + *sections: Literal["actr", "eval", "system"], +) -> Callable[[Callable[..., Any]], Callable[..., Any]]: + """Add CLI options to override fields of selected ImpactConfig sections. + + Args: + sections: Section names to expose as overrides, e.g. "eval" or "actr". Each + field of the corresponding config dataclass becomes a `--{section}-{field}` + option, defaulting to None (meaning "keep the value from loaded config"). + + The decorated function must accept `impact_config: ImpactConfig` as a keyword + argument, populated upstream. Only explicitly provided CLI values override the + loaded configuration. All other fields retain their values from the checkpoint + snapshot. + + Use together with `cli.inject_impact_config_from_checkpoint()` for commands that + load a checkpoint's saved config but need to adjust evaluation/runtime parameters. + """ + + def _decorator(f: Callable[..., Any]) -> Callable[..., Any]: + unsupported = set(sections) - set(_OVERRIDE_SECTIONS.keys()) + if unsupported: + raise ValueError(f"Unsupported override sections: {unsupported}") + + chosen = {s: _OVERRIDE_SECTIONS[s] for s in sections} + options: list[Callable[[Any], Any]] = [] + for section, config_section in chosen.items(): + for field_name in config_section.field_names: + field_type = config_section.type_hints[field_name] + + option_name = f"--{section}-{field_name}".replace("_", "-") + multiple = get_origin(field_type) is tuple + + options.append( + click.option( + option_name, + type=_click_type_for(field_type), + multiple=multiple, + default=() if multiple else None, + help=( + f"Override {config_section.cls.__name__}.{field_name}. " + "Defaults to the loaded config's value." + ), + ) + ) + + # Apply Click options to the function in reverse order to maintain ordering + for option in reversed(options): + f = option(f) + + @wraps(f) + def _wrapper(*args: Any, **kwargs: Any) -> Any: + impact_config: ImpactConfig = kwargs["impact_config"] + + for section, config_section in chosen.items(): + overrides: dict[str, Any] = {} + for field_name in config_section.field_names: + kw_key = f"{section}_{field_name}" + value = kwargs.pop(kw_key, None) + if value in (None, ()): + continue + overrides[field_name] = value + + if overrides: + current_section = getattr(impact_config, section) + impact_config = dataclasses.replace( + impact_config, + **{section: dataclasses.replace(current_section, **overrides)}, + ) + + kwargs["impact_config"] = impact_config + return f(*args, **kwargs) + + return _wrapper + + return _decorator + + +def _click_type_for(type_: Any) -> Any: + """Map Python type to a Click parameter type.""" + if get_origin(type_) is tuple: + return get_args(type_)[0] + + if type_ in (int, float, str): + return type_ + + if isinstance(type_, type) and issubclass(type_, Enum): + return click.Choice([e.value for e in type_]) + + return str diff --git a/src/impact/config.py b/src/impact/config.py index c3e9b2f..8e14f2c 100644 --- a/src/impact/config.py +++ b/src/impact/config.py @@ -4,7 +4,9 @@ from dataclasses import asdict, dataclass, field from enum import Enum from functools import cached_property -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, ClassVar + +import yaml from impact.dataset import DatasetSize, DatasetType, InteractionType @@ -185,6 +187,13 @@ class ImpactConfig: file for exact reconstruction of any pipeline artifact. """ + CONFIG_SCHEMA_VERSION: ClassVar[int] = 1 + """Version of the configuration serialization schema. + + Used to validate compatibility when loading ImpactConfig snapshots from YAML. + Increment when the serialized structure changes incompatibly. + """ + dataset: DatasetConfig = field(default_factory=DatasetConfig) """Configuration for the Yambda dataset fetcher.""" @@ -304,14 +313,60 @@ def write_yaml(self, output_path: Path) -> None: output_path: Target path where the YAML file with the configuration snapshot should be saved. """ - import yaml - + data = { + "schema_version": self.CONFIG_SCHEMA_VERSION, + "hash": self.config_hash, + **self.to_dict(), + } output_path.parent.mkdir(parents=True, exist_ok=True) + with output_path.open("w", encoding="utf-8") as f: yaml.safe_dump( - data=self.to_dict(), + data=data, stream=f, indent=2, sort_keys=False, # Prevents from sorting keys alphabetically allow_unicode=True, ) + + @classmethod + def from_yaml(cls, path: Path) -> ImpactConfig: + """Reconstruct an ImpactConfig from a YAML file written by `write_yaml()`. + + Args: + path: Path to the YAML file. + + Returns: + Reconstructed ImpactConfig instance. + + Raises: + FileNotFoundError: If the file does not exist. + ValueError: If required sections are missing, or the schema version is not + supported, or the configuration checksum does not match. + """ + with path.open("r", encoding="utf-8") as f: + data: dict[str, Any] = yaml.safe_load(f) + + schema_version = data.pop("schema_version") + if schema_version != cls.CONFIG_SCHEMA_VERSION: + raise ValueError(f"Unsupported config schema version: {schema_version}.") + + config = cls( + dataset=DatasetConfig( + size=DatasetSize(data["dataset"]["size"]), + format=DatasetType(data["dataset"]["format"]), + interaction_type=InteractionType(data["dataset"]["interaction_type"]), + ), + session=SessionConfig(**data["session"]), + sequence=SequenceConfig(**data["sequence"]), + gts=GTSConfig(**data["gts"]), + actr=ACTRConfig(**data["actr"]), + sasrec=SASRecConfig(**data["sasrec"]), + eval=EvalConfig(**data["eval"]), + system=SystemConfig(**data["system"]), + ) + + if config.config_hash != data.pop("hash"): + raise ValueError("Configuration checksum mismatch.") + + return config diff --git a/src/impact/main.py b/src/impact/main.py index b0848a9..74b1f6d 100644 --- a/src/impact/main.py +++ b/src/impact/main.py @@ -6,13 +6,17 @@ import json import logging from pathlib import Path -from typing import TYPE_CHECKING, Any, cast +from typing import TYPE_CHECKING, Any, Literal, cast import click import dotenv import structlog -from impact.cli import inject_impact_config +from impact.cli import ( + inject_impact_config, + inject_impact_config_from_checkpoint, + inject_impact_config_overrides, +) from impact.dataset import DatasetSize from impact.logger import configure_logging from impact.pipelines import ( @@ -172,17 +176,10 @@ def train_sasrec( @cli.command(name="evaluate-impact") -@click.option( - "--sasrec-checkpoint-path", - type=click.Path(exists=True, dir_okay=False, path_type=Path), - required=True, - help="Path to a trained SASRecNBR checkpoint that will be evaluated.", -) @click.option( "--partition", type=click.Choice(["val", "test"], case_sensitive=False), - default="test", - show_default=True, + required=True, help=""" Dataset partition to evaluate on. Use 'val' during hyperparameter search, 'test' for final evaluation on the held-out data. @@ -228,11 +225,12 @@ def train_sasrec( default=None, help="Base directory for log files. Defaults to /logs.", ) -@inject_impact_config +@inject_impact_config_from_checkpoint +@inject_impact_config_overrides("actr", "eval", "system") def evaluate_impact( impact_config: ImpactConfig, sasrec_checkpoint_path: Path, - partition: str = "test", + partition: Literal["val", "test"], log_level: str = "INFO", raw_data_dir: Path | None = None, interim_data_dir: Path | None = None, diff --git a/src/impact/pipelines/sasrec_training_pipeline.py b/src/impact/pipelines/sasrec_training_pipeline.py index 330f8e4..b5cb806 100644 --- a/src/impact/pipelines/sasrec_training_pipeline.py +++ b/src/impact/pipelines/sasrec_training_pipeline.py @@ -164,6 +164,7 @@ def train( logger.info("New best model found", best_ndcg=round(best_ndcg, 4)) if save_path is not None: torch.save(model.state_dict(), save_path) + self._config.write_yaml(save_path.parent / "config.yaml") else: epochs_without_improvement += 1 diff --git a/src/impact/utils.py b/src/impact/utils.py index 33dd9f0..9856a22 100644 --- a/src/impact/utils.py +++ b/src/impact/utils.py @@ -3,7 +3,7 @@ import math import os import random -from typing import cast +from typing import TYPE_CHECKING, cast import numpy as np import polars as pl @@ -11,6 +11,10 @@ import structlog import torch +if TYPE_CHECKING: + from collections.abc import Sequence + from pathlib import Path + logger = structlog.stdlib.get_logger(__name__) WORKER_SEED_MODULUS: int = 2**32 @@ -57,6 +61,50 @@ def seed_worker(_: int) -> None: # Ignore the argument (worker_id) random.seed(worker_seed) +def find_sibling_config( + checkpoint_path: Path, + config_filename_patterns: Sequence[str] = ( + "impact_config.yaml", + "config.yaml", + "*.yaml", + ), +) -> Path: + """Locate the ImpactConfig YAML saved alongside a model checkpoint. + + Searches `checkpoint_path.parent` for, in order of preference, files matching + each of `config_filename_patterns`. The first pattern that matches exactly one + file wins. + + Args: + checkpoint_path: Path to the model checkpoint whose sibling config is sought. + config_filename_patterns: Glob patterns tried in order. + + Returns: + Path to the matched config file. + + Raises: + FileNotFoundError: If no pattern matches any file. + RuntimeError: If a pattern matches more than one file (ambiguous). + """ + parent = checkpoint_path.parent + for pattern in config_filename_patterns: + matches = sorted(parent.glob(pattern)) + if not matches: + continue + if len(matches) > 1: + raise RuntimeError( + f"Found multiple config candidates matching '{pattern}' in " + f"{parent}: {[m.name for m in matches]}. Expected exactly one." + ) + return matches[0] + + raise FileNotFoundError( + f"No config YAML found next to checkpoint {checkpoint_path.name} in " + f"{parent}. Tried patterns: {', '.join(config_filename_patterns)}. " + "Re-run model training to generate it automatically." + ) + + def estimate_processing_partitions( data: pl.LazyFrame | pl.DataFrame, target_ram_fraction: float = 0.15,