fix(models): leave the scored splits out of a saved model

DashAI/back/models/base_model.py

@@ -46,6 +46,12 @@ class BaseModel(ConfigObject, metaclass=ABCMeta):
     # `level=EPOCH`, so for them this stays None and nothing changes.
     _epoch_reporter = None
 
+    # The splits a model is scored on, pointed at by whoever fits it so that
+    # `calculate_metrics` can read them. None until then, and again once the
+    # model is pickled: they are the run's data, never part of the model.
+    x_data = None
+    y_data = None
+
     @classmethod
     def get_metadata(cls) -> Dict[str, Any]:
         """Get metadata values for the current model.
@@ -62,6 +68,31 @@ def get_metadata(cls) -> Dict[str, Any]:
         metadata["download_size_bytes"] = getattr(cls, "DOWNLOAD_SIZE_BYTES", None)
         return metadata
 
+    def __getstate__(self) -> Dict[str, Any]:
+        """Leave the data the model was scored on out of its pickled state.
+
+        Whoever fits a model points ``x_data`` and ``y_data`` at the splits it
+        trains on, so ``calculate_metrics`` can read them. Every model that
+        saves itself by pickling (``joblib.dump(self)``) would otherwise write
+        the whole dataset into its artifact: a copy of the data on every run,
+        and a save that grows with the dataset instead of with the model.
+
+        Dropping them lets the class defaults answer, so a loaded model scores
+        nothing on its own, the same as a model never pointed at any data.
+
+        Returns
+        -------
+        Dict[str, Any]
+            A copy of the instance state, without the training data.
+        """
+        # A copy, because further up the MRO the state may be the live
+        # ``__dict__``; and ``or {}``, because ``object`` gives None for an
+        # instance that has no attributes yet.
+        state = dict(super().__getstate__() or {})
+        state.pop("x_data", None)
+        state.pop("y_data", None)
+        return state
+
     @abstractmethod
     def save(self, filename: str) -> None:
         """Store the model to disk.

DashAI/back/models/faiss/faiss_base_model.py

@@ -135,7 +135,8 @@ def _import_faiss():
 
     def __getstate__(self) -> dict:
         """Serialise the FAISS index as raw bytes before pickling."""
-        state = self.__dict__.copy()
+        # Through super() so BaseModel can leave the training data out too.
+        state = dict(super().__getstate__() or {})
         if state.get("_index") is not None:
             faiss = self._import_faiss()
 

tests/back/models/test_saved_model_leaves_out_data.py

@@ -0,0 +1,96 @@
+"""A saved model carries the model, not the data it was scored on.
+
+Whoever fits a model points ``x_data`` and ``y_data`` at its splits so the
+metric methods can read them. Models that save themselves by pickling used to
+write those datasets into the artifact as well, which made the save grow with
+the dataset: a large one took about a minute to write.
+"""
+
+import pickle
+
+import numpy as np
+import pandas as pd
+import pytest
+
+from DashAI.back.dataloaders.classes.dashai_dataset import to_dashai_dataset
+from DashAI.back.models.faiss.faiss_kmeans_clustering import FaissKMeansClustering
+from DashAI.back.models.scikit_learn.kmeans_clustering import KMeansClustering
+
+# One model that pickles itself through scikit-learn's state and one with its
+# own ``__getstate__`` (the FAISS index), which has to defer to the base one.
+MODELS = [KMeansClustering, FaissKMeansClustering]
+
+
+def _data(rows=300, seed=0):
+    rng = np.random.default_rng(seed)
+    return to_dashai_dataset(
+        pd.DataFrame({"x": rng.normal(size=rows), "y": rng.normal(size=rows)})
+    )
+
+
+def _fitted(model_class):
+    x = _data()
+    model = model_class(n_clusters=3, random_state=0)
+    model.train(x)
+    model.x_data = {"train": x, "validation": x, "test": x}
+    model.y_data = {"train": x, "validation": x, "test": x}
+    return model
+
+
+@pytest.mark.parametrize("model_class", MODELS)
+def test_the_saved_model_does_not_carry_its_data(model_class, tmp_path):
+    model = _fitted(model_class)
+    path = tmp_path / "model.joblib"
+
+    model.save(str(path))
+    restored = model_class.load(str(path))
+
+    assert restored.x_data is None
+    assert restored.y_data is None
+
+
+@pytest.mark.parametrize("model_class", MODELS)
+def test_the_saved_model_size_does_not_grow_with_the_data(model_class, tmp_path):
+    model = _fitted(model_class)
+    small, large = tmp_path / "small.joblib", tmp_path / "large.joblib"
+
+    model.save(str(small))
+    # Fifty thousand rows of two float64 columns: 800 KB per split, so either
+    # attribute leaking into the artifact would dwarf the tolerance below.
+    big = _data(rows=50_000)
+    model.x_data = {"train": big, "validation": big, "test": big}
+    model.y_data = {"train": big, "validation": big, "test": big}
+    model.save(str(large))
+
+    assert large.stat().st_size < small.stat().st_size + 10_000
+
+
+def test_a_model_without_attributes_can_still_be_pickled():
+    """``object.__getstate__`` gives None, not a dict, for an empty instance."""
+    model = FaissKMeansClustering.__new__(FaissKMeansClustering)
+
+    restored = pickle.loads(pickle.dumps(model))
+
+    assert restored.x_data is None
+
+
+@pytest.mark.parametrize("model_class", MODELS)
+def test_saving_leaves_the_live_model_pointed_at_its_data(model_class, tmp_path):
+    model = _fitted(model_class)
+    x_data = model.x_data
+
+    model.save(str(tmp_path / "model.joblib"))
+
+    assert model.x_data is x_data
+
+
+@pytest.mark.parametrize("model_class", MODELS)
+def test_the_restored_model_still_predicts_the_same(model_class, tmp_path):
+    model = _fitted(model_class)
+    x = _data()
+    expected = list(np.asarray(model.get_cluster_labels(x)))
+
+    model.save(str(tmp_path / "model.joblib"))
+    restored = model_class.load(str(tmp_path / "model.joblib"))
+
+    assert list(np.asarray(restored.get_cluster_labels(x))) == expected