diff --git a/DashAI/back/dataloaders/classes/dashai_dataset.py b/DashAI/back/dataloaders/classes/dashai_dataset.py index 524326a63..9e2e93320 100644 --- a/DashAI/back/dataloaders/classes/dashai_dataset.py +++ b/DashAI/back/dataloaders/classes/dashai_dataset.py @@ -682,8 +682,20 @@ def select(self, *args, **kwargs) -> "DashAIDataset": if isinstance(selected_dataset, DashAIDataset): return selected_dataset else: - # If the selected dataset is a Dataset, convert it to DashAIDataset - arrow_tbl = selected_dataset.with_format("arrow")[:] + # If the selected dataset is a Dataset, convert it to DashAIDataset. + # Rows that are not one contiguous range come back as the whole + # table plus an indices mapping, and reading that mapping with + # ``with_format("arrow")[:]`` gathers one record batch per row. + # The layout outlives the selection: a dataset saved from it keeps + # a chunk per row, and pickling such a table takes about a + # thousand times longer than a regular one. ``take`` gathers the + # same rows, in the same order, into regular chunks. + if selected_dataset._indices is not None: + arrow_tbl = selected_dataset.data.table.take( + selected_dataset._indices.column(0) + ) + else: + arrow_tbl = selected_dataset.with_format("arrow")[:] arrow_tblx = save_types_in_arrow_metadata( arrow_tbl, {col: self._types[col].to_string() for col in self._types} ) diff --git a/DashAI/back/models/base_model.py b/DashAI/back/models/base_model.py index aecf9ecf9..499cf2805 100644 --- a/DashAI/back/models/base_model.py +++ b/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. diff --git a/DashAI/back/models/faiss/faiss_base_model.py b/DashAI/back/models/faiss/faiss_base_model.py index d4ba84dbb..f86ed9352 100644 --- a/DashAI/back/models/faiss/faiss_base_model.py +++ b/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() diff --git a/tests/back/dataloaders/test_select_keeps_regular_chunks.py b/tests/back/dataloaders/test_select_keeps_regular_chunks.py new file mode 100644 index 000000000..b30f7d6e8 --- /dev/null +++ b/tests/back/dataloaders/test_select_keeps_regular_chunks.py @@ -0,0 +1,88 @@ +"""Selecting scattered rows leaves a table in regular Arrow chunks. + +Rows that are not one contiguous range used to come out of +``DashAIDataset.select`` as one record batch per row. Every row-dropping +converter goes through it, and the layout outlived the selection: a dataset the +NaN remover cleaned kept a chunk per row once saved, and pickling the splits of +a 35k-row one took 40 s and almost 5 GB of memory. +""" + +import numpy as np +import pandas as pd + +from DashAI.back.converters.simple_converters.nan_remover import NanRemover +from DashAI.back.dataloaders.classes.dashai_dataset import ( + DashAIDataset, + load_dataset, + save_dataset, + to_dashai_dataset, +) + +ROWS = 2000 +# Far below one chunk per row, which is what scattered rows used to get. +FEW_CHUNKS = 10 + + +def _dataset(with_gaps=False): + x = np.arange(ROWS, dtype="float64") + if with_gaps: + x[::3] = np.nan + return to_dashai_dataset( + pd.DataFrame({"x": x, "label": [f"c{i % 3}" for i in range(ROWS)]}) + ) + + +def _chunks(dataset): + return [column.num_chunks for column in dataset.arrow_table.columns] + + +def test_scattered_rows_come_back_in_a_few_chunks(): + keep = list(range(0, ROWS, 2)) + + selected = _dataset().select(keep) + + assert isinstance(selected, DashAIDataset) + assert max(_chunks(selected)) < FEW_CHUNKS + + +def test_scattered_rows_keep_the_requested_order(): + keep = [1500, 3, 999, 0, 42] + + selected = _dataset().select(keep) + + assert selected.arrow_table.column("x").to_pylist() == [float(i) for i in keep] + assert selected.arrow_table.column("label").to_pylist() == [ + f"c{i % 3}" for i in keep + ] + + +def test_scattered_rows_keep_the_dataset_types(): + dataset = _dataset() + + selected = dataset.select([5, 1, 7]) + + assert {name: t.to_string() for name, t in selected.types.items()} == { + name: t.to_string() for name, t in dataset.types.items() + } + + +def test_an_empty_selection_is_still_a_dataset(): + selected = _dataset().select([]) + + assert isinstance(selected, DashAIDataset) + assert selected.num_rows == 0 + assert selected.column_names == ["x", "label"] + + +def test_the_nan_remover_output_is_not_fragmented(tmp_path): + dataset = _dataset(with_gaps=True) + remover = NanRemover().fit(dataset) + + cleaned = remover.transform(dataset) + save_dataset(cleaned, tmp_path / "cleaned") + reloaded = load_dataset(str(tmp_path / "cleaned")) + + assert cleaned.num_rows == ROWS - len(range(0, ROWS, 3)) + assert not any(np.isnan(cleaned.arrow_table.column("x").to_numpy())) + assert max(_chunks(cleaned)) < FEW_CHUNKS + assert max(_chunks(reloaded)) < FEW_CHUNKS diff --git a/tests/back/models/test_saved_model_leaves_out_data.py b/tests/back/models/test_saved_model_leaves_out_data.py new file mode 100644 index 000000000..66720a54e --- /dev/null +++ b/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