fix: restore stored plugin settings as they are on reproduce and resume (#484)

* fix: restore stored plugin settings as they are on reproduce and resume

Settings.model_validate still reads config.toml, environment variables and
the command line, and merges plugin tables key by key. A [modifier.X] or
[scorer.X] key in the local config.toml that the stored settings don't have
was added to a --reproduce run or a resumed study, so a reproduction could
build a different model (hash mismatch).

* fix: restore every stored table, not just plugin tables

Stored values already win over config.toml, but pydantic-settings merges
tables key by key, so max_memory picked up local keys too, and a plugin
table that only exists in config.toml was added to the run. Restore each
stored table as it is and drop tables that weren't stored.

* fix: restore stored settings without reading other sources

Validate stored settings through StoredSettings, whose only source is
the values passed in, instead of letting config.toml, the environment
and the command line merge into them and overwriting afterwards. Fields
that are never stored (exclude=True) are copied from the settings this
run was started with.

* refactor: restore stored settings on top of the current ones

Validate the current settings' fields with the stored values laid over
them, instead of copying fields with exclude=True one by one. Fields
that are never stored keep the values this run was started with, and
plugin tables still come only from the stored values.

* fix: take only never-stored settings from the current run

* refactor: rename from_stored to model_validate_init_only

src/heretic/config.py

@@ -2,7 +2,7 @@
 # Copyright (C) 2025-2026  Philipp Emanuel Weidmann <pew@worldwidemann.com> + contributors
 
 from enum import Enum
-from typing import Literal, TypeAlias
+from typing import Any, Literal, TypeAlias
 
 from pydantic import (
     BaseModel,
@@ -566,7 +566,7 @@ class Settings(BaseSettings):
         file_secret_settings: PydanticBaseSettingsSource,
     ) -> tuple[PydanticBaseSettingsSource, ...]:
         return (
-            init_settings,  # Used during resume - should override *all* other sources.
+            init_settings,  # Values passed to the constructor override all other sources.
             CliSettingsSource(
                 settings_cls,
                 cli_parse_args=True,
@@ -578,3 +578,36 @@ class Settings(BaseSettings):
             file_secret_settings,
             TomlConfigSettingsSource(settings_cls, toml_file="config.toml"),
         )
+
+    @classmethod
+    def model_validate_init_only(
+        cls, values: dict[str, Any], current: "Settings"
+    ) -> "Settings":
+        """
+        Validates `values` without reading any other settings source: nothing
+        comes from config.toml, environment variables or the command line, and
+        settings missing from `values` get their defaults. Settings with
+        `exclude=True` (such as `save_directory`) are left out of serialized
+        settings, so they are taken from `current`.
+        """
+        excluded = {
+            name: getattr(current, name)
+            for name, field in cls.model_fields.items()
+            if field.exclude
+        }
+        return InitOnlySettings.model_validate(values | excluded)
+
+
+class InitOnlySettings(Settings):
+    """Settings validated only from the values passed in, without reading any source."""
+
+    @classmethod
+    def settings_customise_sources(
+        cls,
+        settings_cls: type[BaseSettings],
+        init_settings: PydanticBaseSettingsSource,
+        env_settings: PydanticBaseSettingsSource,
+        dotenv_settings: PydanticBaseSettingsSource,
+        file_secret_settings: PydanticBaseSettingsSource,
+    ) -> tuple[PydanticBaseSettingsSource, ...]:
+        return (init_settings,)

src/heretic/main.py

@@ -47,6 +47,7 @@ from .progress import patch_tqdm
 patch_tqdm()
 """
 
+import json
 import logging
 import math
 import random
@@ -267,7 +268,9 @@ def run():
 
         print()
 
-        settings = Settings.model_validate(reproduction_information["settings"])
+        settings = Settings.model_validate_init_only(
+            reproduction_information["settings"], settings
+        )
 
     if settings.seed is None:
         settings.seed = random.randint(0, 2**32 - 1)
@@ -400,8 +403,8 @@ def run():
             return
 
         if action == "continue":
-            settings = Settings.model_validate_json(
-                existing_study.user_attrs["settings"]
+            settings = Settings.model_validate_init_only(
+                json.loads(existing_study.user_attrs["settings"]), settings
             )
         elif action == "restart":
             os.unlink(study_checkpoint_file)

tests/test_config.py

@@ -1,11 +1,16 @@
 # SPDX-License-Identifier: AGPL-3.0-or-later
 # Copyright (C) 2025-2026  Philipp Emanuel Weidmann <pew@worldwidemann.com> + contributors
 
+import os
+import sys
+import tempfile
 import unittest
+from pathlib import Path
+from unittest.mock import patch
 
 from pydantic import ValidationError
 
-from heretic.config import ScorerConfig
+from heretic.config import ScorerConfig, Settings
 
 
 class ScorerConfigTests(unittest.TestCase):
@@ -47,5 +52,87 @@ class ScorerConfigTests(unittest.TestCase):
             )
 
 
+class ModelValidateInitOnlyTests(unittest.TestCase):
+    def validate(
+        self, values: dict, config_toml: str, argv: tuple[str, ...] = ("heretic",)
+    ) -> Settings:
+        with tempfile.TemporaryDirectory() as directory:
+            Path(directory, "config.toml").write_text(config_toml)
+            previous_directory = os.getcwd()
+            os.chdir(directory)
+            try:
+                with patch.object(sys, "argv", list(argv)):
+                    return Settings.model_validate_init_only(
+                        values, Settings(model="current")
+                    )
+            finally:
+                os.chdir(previous_directory)
+
+    def test_keeps_plugin_tables_from_values(self) -> None:
+        settings = self.validate(
+            {
+                "model": "model",
+                "modifier": {"Abliteration": {"row_normalization": "pre"}},
+            },
+            'save_directory = "toml"\n'
+            "[modifier.Abliteration]\nwinsorization_quantile = 0.5\n",
+            ("heretic", "--save-directory", "current"),
+        )
+
+        self.assertEqual(
+            settings.model_extra,
+            {"modifier": {"Abliteration": {"row_normalization": "pre"}}},
+        )
+        # Settings with exclude=True come from the current settings.
+        self.assertEqual(settings.save_directory, "current")
+
+    def test_drops_plugin_tables_missing_from_values(self) -> None:
+        settings = self.validate(
+            {"model": "model"},
+            "[modifier.Abliteration]\nwinsorization_quantile = 0.5\n",
+        )
+
+        self.assertEqual(settings.model_extra, {})
+
+    def test_other_sources_do_not_override_values(self) -> None:
+        values = {
+            "model": "model",
+            "n_trials": 50,
+            "max_response_length": 60,
+            "n_startup_trials": 20,
+        }
+
+        with patch.dict(os.environ, {"HERETIC_MAX_RESPONSE_LENGTH": "9"}):
+            settings = self.validate(
+                values,
+                "n_trials = 7\n",
+                ("heretic", "--n-startup-trials", "3"),
+            )
+
+        for name, value in values.items():
+            self.assertEqual(getattr(settings, name), value)
+
+    def test_settings_missing_from_values_get_defaults(self) -> None:
+        with patch.dict(os.environ, {"HERETIC_MAX_RESPONSE_LENGTH": "9"}):
+            settings = self.validate(
+                {"model": "model"},
+                "n_trials = 7\n",
+                ("heretic", "--n-startup-trials", "3"),
+            )
+
+        for name in ["n_trials", "max_response_length", "n_startup_trials"]:
+            self.assertEqual(
+                getattr(settings, name), Settings.model_fields[name].default
+            )
+
+    def test_keeps_top_level_tables_from_values(self) -> None:
+        settings = self.validate(
+            {"model": "model", "max_memory": {"0": "20GiB"}},
+            'max_memory = { "1" = "8GiB" }\n',
+        )
+
+        self.assertEqual(settings.max_memory, {"0": "20GiB"})
+
+
 if __name__ == "__main__":
     unittest.main()