ImpactConfig management improved and optimized

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 },

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,

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
 

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

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

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 <project_root>/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,