Allow Apeiron to use custom checkpoints

docs/model_harness.md

@@ -32,6 +32,66 @@ Required/expected attributes:
 | `self.eval_metrics` | `dict[str, callable]` | Ordered metric map. Metric order controls `metric_index` behavior in drift detection. |
 | `self.higher_is_better` | `dict[str, bool]` | Used by `ModelEvalDetector` style logic. Optional for monitor-only operation. |
 
+## Checkpointing
+
+The base class provides automatic checkpointing when configured via `[model]` section:
+
+```toml
+[model]
+name = "my_model"
+max_ckpts = 5              # Keep last 5 checkpoints (0 disables)
+ckpts_path = "./output/checkpoints"
+```
+
+Checkpoints are saved automatically after each continual learning adaptation.
+
+### Default Checkpoint Format
+
+By default, `BaseModelHarness` saves only model weights (`model.state_dict()`):
+
+```python
+# Default behavior - saves weights only
+torch.save(model.state_dict(), checkpoint_path)
+```
+
+This works for models that don't require additional metadata (preprocessing state, feature names, etc.) to be loaded.
+
+### Custom Checkpoint Format
+
+If your model requires metadata beyond weights (e.g., preprocessing scalers, feature names, architecture parameters), override `build_checkpoint_payload()`:
+
+```python
+class MyHarness(BaseModelHarness):
+    def __init__(self, cfg: Config):
+        # ... load model with metadata ...
+        self._feature_names = checkpoint["feature_names"]
+        self._scaler = checkpoint["scaler"]
+        self._input_dim = checkpoint["input_dim"]
+
+        super().__init__(cfg=cfg, model=model)
+
+    def build_checkpoint_payload(self) -> dict:
+        """Save full checkpoint with metadata."""
+        return {
+            "model_state_dict": self.model.state_dict(),
+            "feature_names": self._feature_names,
+            "scaler": self._scaler,
+            "input_dim": self._input_dim,
+        }
+```
+
+This ensures that:
+1. Checkpoints saved during CL runs can be loaded directly as pretrained models
+2. No separate "reference" checkpoints are needed to recover metadata
+3. Save and load formats match exactly
+
+### Checkpoint Management
+
+- Checkpoints are named `drift_adaptation_<event>.pt` where `<event>` is the drift event number
+- A `latest` file tracks the most recent checkpoint
+- FIFO policy: oldest checkpoints are deleted when `max_ckpts` limit is exceeded
+- Set `max_ckpts = 0` to disable checkpointing entirely
+
 ## Runtime Lifecycle
 
 1. `examples/utils.py:get_example` picks a harness from `cfg.data.name`.
@@ -255,6 +315,34 @@ def update_data_stream(self):
 
 This function is called by the monitor whenever a new stream segment becomes active.
 
+### (Optional) Custom Checkpoint Format
+
+If your model requires metadata beyond weights (e.g., preprocessing scalers, feature names, architecture parameters), override `build_checkpoint_payload()`:
+
+```python
+def build_checkpoint_payload(self):
+    """Save checkpoint with all metadata needed for loading."""
+    return {
+        "model_state_dict": self.model.state_dict(),
+        "feature_names": self._feature_names,
+        "scaler": self._scaler,
+        "input_dim": self._input_dim,
+    }
+```
+
+This ensures saved checkpoints can be loaded directly without separate reference files.
+
+**When to use:**
+- Your model loads with preprocessing state (scalers, normalizers, encoders)
+- Your architecture requires parameters beyond what's in the config (e.g., inferred input dimensions)
+- You want checkpoints to be self-contained and portable
+
+**When to skip:**
+- Your model only needs weights (`state_dict()`) - the default behavior is fine
+- All loading parameters come from the config file
+
+See the "Checkpointing" section above for details.
+
 ## Step 4: Build Dataset Utilities
 
 Dataset loading should be implemented in `utils.py`.

src/apeiron/model/torch_model_harness.py

@@ -222,13 +222,27 @@ def task_diagonals(self) -> List[List[float]]:
     def ckpts_enabled(self) -> bool:
         return self.cfg.model.max_ckpts > 0 and bool(self.cfg.model.ckpts_path)
 
+    def build_checkpoint_payload(self) -> Any:
+        """Build the checkpoint object to save.
+
+        Subclasses can override this to include additional metadata beyond weights
+        (e.g., preprocessing scalers, feature names, architecture parameters)
+        so that saved checkpoints match the format expected by the loader.
+
+        Returns
+        -------
+        By default, returns ``model.state_dict()`` (weights only).
+        """
+        return self.model.state_dict()
+
     def save_ckpt(self, event: int) -> str:
         """Persist model state, evict oldest when over budget."""
         d = Path(self.cfg.model.ckpts_path)
         d.mkdir(parents=True, exist_ok=True)
 
         fname = f"drift_adaptation_{event}.pt"
-        torch.save(self.model.state_dict(), d / fname)
+        payload = self.build_checkpoint_payload()
+        torch.save(payload, d / fname)
         (d / "latest").write_text(fname)
 
         # Guillotine the oldest survivors