From d32c9b93808e4cb250453342580791ace88d7f39 Mon Sep 17 00:00:00 2001 From: jfrog64 Date: Wed, 26 Aug 2026 16:50:55 +0200 Subject: [PATCH 1/2] perf: cache loaded feature matrices between splits and models Feature matrices were reloaded from disk for every split, every hyperparameter setting and, for single drug models, for every drug. On a 545-drug CTRPv2 run this dominated the wall clock: 3270 loads at roughly 17.8 s each, about 16 h of an 18.4 h run. The matrices only depend on the model class, the feature kind, the data path, the dataset name and the hyperparameters (which decide the views and the gene list), so they are cached under exactly that key. The cache holds at most four entries, since an entry is a full feature matrix, and can be disabled with DREVAL_FEATURE_CACHE=0 to A/B check that results are unchanged. Co-Authored-By: Claude Opus 5 --- drevalpy/experiment.py | 88 +++++++++++-- tests/test_feature_cache.py | 239 ++++++++++++++++++++++++++++++++++++ 2 files changed, 320 insertions(+), 7 deletions(-) create mode 100644 tests/test_feature_cache.py diff --git a/drevalpy/experiment.py b/drevalpy/experiment.py index 9fee1a6f..d0dbcca7 100644 --- a/drevalpy/experiment.py +++ b/drevalpy/experiment.py @@ -7,7 +7,7 @@ import tempfile import warnings from pathlib import Path -from typing import Any +from typing import Any, cast import numpy as np import pandas as pd @@ -625,6 +625,80 @@ def consolidate_single_drug_model_predictions( ) +#: Cache for loaded feature matrices, keyed by everything they depend on. +#: +#: Feature matrices depend only on (model class, hyperparameters, data path, dataset), not on the CV +#: split and not on the drug. For single drug models the pipeline nevertheless walks the whole split +#: loop once PER DRUG, so the same matrix was re-read from disk once per split and drug. Measured on +#: a 545 drug CTRPv2 run with 893 genes: 17.8 s per load, 3270 loads, roughly 16 h of an 18.4 h run. +#: +#: Reusing one object is safe because every consumer only reads it and hands it on via .copy() +#: (train_and_predict, train_final_model, cross_study_prediction, randomization_test). +#: Set DREVAL_FEATURE_CACHE=0 to disable, e.g., to A/B check that results are unchanged. +_FEATURE_CACHE: dict[tuple, FeatureDataset | None] = {} + +#: Kept small on purpose: an entry is a full feature matrix, and the access pattern is "same key many +#: times in a row", so a couple of slots already give the full speedup. +_FEATURE_CACHE_MAXSIZE = 4 + + +def clear_feature_cache() -> None: + """Drop all cached feature matrices, e.g., to free memory between datasets.""" + _FEATURE_CACHE.clear() + + +def _load_features_cached(model: DRPModel, path_data: str, dataset_name: str, kind: str) -> FeatureDataset | None: + """ + Load cell line or drug features, reusing an already loaded matrix when nothing they depend on changed. + + :param model: built model, i.e., build_model() was already called so the hyperparameters are set + :param path_data: path to the data directory, e.g., data/ + :param dataset_name: name of the dataset, e.g., GDSC2 + :param kind: either "cell_line" or "drug" + :returns: the feature dataset, or None if the model does not use this kind of feature + """ + loader = model.load_cell_line_features if kind == "cell_line" else model.load_drug_features + if os.environ.get("DREVAL_FEATURE_CACHE", "1") == "0": + return loader(data_path=path_data, dataset_name=dataset_name) + # The hyperparameters decide which views and which gene list are loaded, so they belong in the key. + key = ( + type(model).__name__, + kind, + path_data, + dataset_name, + json.dumps(getattr(model, "hyperparameters", {}), sort_keys=True, default=str), + ) + if key not in _FEATURE_CACHE: + if len(_FEATURE_CACHE) >= _FEATURE_CACHE_MAXSIZE: + _FEATURE_CACHE.pop(next(iter(_FEATURE_CACHE))) + _FEATURE_CACHE[key] = loader(data_path=path_data, dataset_name=dataset_name) + return _FEATURE_CACHE[key] + + +def _load_cell_line_features_cached(model: DRPModel, path_data: str, dataset_name: str) -> FeatureDataset: + """ + Cached variant of ``model.load_cell_line_features``, which every model has to implement. + + :param model: built model + :param path_data: path to the data directory, e.g., data/ + :param dataset_name: name of the dataset, e.g., GDSC2 + :returns: the cell line feature dataset + """ + return cast(FeatureDataset, _load_features_cached(model, path_data, dataset_name, "cell_line")) + + +def _load_drug_features_cached(model: DRPModel, path_data: str, dataset_name: str) -> FeatureDataset | None: + """ + Cached variant of ``model.load_drug_features``, which single drug models do not provide. + + :param model: built model + :param path_data: path to the data directory, e.g., data/ + :param dataset_name: name of the dataset, e.g., GDSC2 + :returns: the drug feature dataset, or None + """ + return _load_features_cached(model, path_data, dataset_name, "drug") + + def load_features( model: DRPModel, path_data: str, dataset: DrugResponseDataset ) -> tuple[FeatureDataset, FeatureDataset | None]: @@ -636,8 +710,8 @@ def load_features( :param dataset: dataset to load features for, e.g., GDSC2 :returns: tuple of cell line and, potentially, drug features """ - cl_features = model.load_cell_line_features(data_path=path_data, dataset_name=dataset.dataset_name) - drug_features = model.load_drug_features(data_path=path_data, dataset_name=dataset.dataset_name) + cl_features = _load_cell_line_features_cached(model, path_data, dataset.dataset_name) + drug_features = _load_drug_features_cached(model, path_data, dataset.dataset_name) return cl_features, drug_features @@ -1105,10 +1179,10 @@ def train_and_predict( raise ValueError("train_dataset must have a dataset_name") if cl_features is None: print("Loading cell line features ...") - cl_features = model.load_cell_line_features(data_path=path_data, dataset_name=train_dataset.dataset_name) + cl_features = _load_cell_line_features_cached(model, path_data, train_dataset.dataset_name) if drug_features is None: print("Loading drug features ...") - drug_features = model.load_drug_features(data_path=path_data, dataset_name=train_dataset.dataset_name) + drug_features = _load_drug_features_cached(model, path_data, train_dataset.dataset_name) cell_lines_to_keep = cl_features.identifiers if cl_features is not None else None drugs_to_keep = drug_features.identifiers if drug_features is not None else None @@ -1617,8 +1691,8 @@ def train_final_model( print(f"Best hyperparameters for final model: {best_hpams}") model.build_model(hyperparameters=best_hpams) - cl_features = model.load_cell_line_features(data_path=path_data, dataset_name=full_dataset.dataset_name) - drug_features = model.load_drug_features(data_path=path_data, dataset_name=full_dataset.dataset_name) + cl_features = _load_cell_line_features_cached(model, path_data, full_dataset.dataset_name) + drug_features = _load_drug_features_cached(model, path_data, full_dataset.dataset_name) cell_lines_to_keep = cl_features.identifiers drugs_to_keep = drug_features.identifiers if drug_features is not None else None diff --git a/tests/test_feature_cache.py b/tests/test_feature_cache.py new file mode 100644 index 00000000..43707ab7 --- /dev/null +++ b/tests/test_feature_cache.py @@ -0,0 +1,239 @@ +"""Tests for the feature matrix cache used by the experiment pipeline.""" + +import numpy as np +import pytest + +from drevalpy import experiment +from drevalpy.datasets.dataset import DrugResponseDataset, FeatureDataset +from drevalpy.experiment import _load_features_cached, clear_feature_cache, load_features +from drevalpy.models.drp_model import DRPModel + +#: Number of loader calls per feature kind, shared by all test models so the cache effect is visible +#: across instances and subclasses. +_CALLS = {"cell_line": 0, "drug": 0} + + +class _CountingModel(DRPModel): + """Model whose loaders count how often they were called and return a fresh object each time.""" + + cell_line_views = ["gene_expression"] + drug_views = ["fingerprints"] + + @classmethod + def get_model_name(cls) -> str: + """ + Returns the model name. + + :returns: name of this test model + """ + return "CountingModel" + + @classmethod + def get_hyperparameter_set(cls) -> list[dict]: + """ + Returns a single empty hyperparameter set, there is no hyperparameters.yaml for a test model. + + :returns: list with one empty hyperparameter dict + """ + return [{}] + + def build_model(self, hyperparameters: dict) -> None: + """ + Stores the hyperparameters, which are part of the cache key. + + :param hyperparameters: hyperparameters to use + """ + self.hyperparameters = hyperparameters + + def load_cell_line_features(self, data_path: str, dataset_name: str) -> FeatureDataset: + """ + Returns a new FeatureDataset and counts the call. + + :param data_path: unused + :param dataset_name: unused + :returns: cell line features + """ + _CALLS["cell_line"] += 1 + return FeatureDataset(features={"CL0": {"gene_expression": np.zeros(3)}}) + + def load_drug_features(self, data_path: str, dataset_name: str) -> FeatureDataset | None: + """ + Returns a new FeatureDataset and counts the call. + + :param data_path: unused + :param dataset_name: unused + :returns: drug features + """ + _CALLS["drug"] += 1 + return FeatureDataset(features={"D0": {"fingerprints": np.zeros(3)}}) + + def train( + self, + output: DrugResponseDataset, + cell_line_input: FeatureDataset, + drug_input: FeatureDataset | None = None, + output_earlystopping: DrugResponseDataset | None = None, + model_checkpoint_dir: str = "checkpoints", + ) -> None: + """ + Not used by these tests. + + :param output: unused + :param cell_line_input: unused + :param drug_input: unused + :param output_earlystopping: unused + :param model_checkpoint_dir: unused + """ + + def predict( + self, + cell_line_ids: np.ndarray, + drug_ids: np.ndarray, + cell_line_input: FeatureDataset, + drug_input: FeatureDataset | None = None, + ) -> np.ndarray: + """ + Returns zeros. + + :param cell_line_ids: cell line ids to predict for + :param drug_ids: unused + :param cell_line_input: unused + :param drug_input: unused + :returns: zero predictions + """ + return np.zeros(len(cell_line_ids)) + + +class _OtherCountingModel(_CountingModel): + """Same behaviour under a different class name, so it must get its own cache entry.""" + + @classmethod + def get_model_name(cls) -> str: + """ + Returns the model name. + + :returns: name of this test model + """ + return "OtherCountingModel" + + +@pytest.fixture(autouse=True) +def _reset_cache(): + """ + Start every test with an empty cache and zeroed call counters. + + :yields: nothing, this fixture only manages global state + """ + clear_feature_cache() + _CALLS["cell_line"] = 0 + _CALLS["drug"] = 0 + yield + clear_feature_cache() + + +def _build(hyperparameters: dict | None = None, model_class: type[_CountingModel] = _CountingModel) -> _CountingModel: + """ + Instantiate and build a counting model. + + :param hyperparameters: hyperparameters to build with + :param model_class: class to instantiate + :returns: the built model + """ + model = model_class() + model.build_model(hyperparameters=hyperparameters if hyperparameters is not None else {}) + return model + + +def test_repeated_loads_hit_the_cache() -> None: + """The same model, path and dataset must load each feature kind exactly once.""" + model = _build() + first = _load_features_cached(model, "data", "TOYv1", "cell_line") + second = _load_features_cached(model, "data", "TOYv1", "cell_line") + + assert _CALLS["cell_line"] == 1 + assert first is second + + +def test_cache_is_shared_across_model_instances() -> None: + """A fresh instance with the same hyperparameters must reuse the cached matrix. + + This is the case that matters in practice: single drug models re-instantiate the model + for every drug and every split. + """ + for _ in range(5): + _load_features_cached(_build(), "data", "TOYv1", "cell_line") + + assert _CALLS["cell_line"] == 1 + + +def test_cache_key_separates_kind_dataset_path_and_hyperparameters() -> None: + """Everything the matrices depend on has to be part of the cache key.""" + model = _build() + _load_features_cached(model, "data", "TOYv1", "cell_line") + _load_features_cached(model, "data", "TOYv1", "drug") + assert (_CALLS["cell_line"], _CALLS["drug"]) == (1, 1) + + _load_features_cached(model, "data", "TOYv2", "cell_line") + assert _CALLS["cell_line"] == 2 + + _load_features_cached(model, "other_data", "TOYv1", "cell_line") + assert _CALLS["cell_line"] == 3 + + _load_features_cached(_build({"gene_list": "landmark_genes_reduced"}), "data", "TOYv1", "cell_line") + assert _CALLS["cell_line"] == 4 + + _load_features_cached(_build(model_class=_OtherCountingModel), "data", "TOYv1", "cell_line") + assert _CALLS["cell_line"] == 5 + + +def test_hyperparameter_order_does_not_split_the_key() -> None: + """Equal hyperparameters written in a different order describe the same matrix.""" + _load_features_cached(_build({"a": 1, "b": 2}), "data", "TOYv1", "cell_line") + _load_features_cached(_build({"b": 2, "a": 1}), "data", "TOYv1", "cell_line") + + assert _CALLS["cell_line"] == 1 + + +def test_cache_evicts_oldest_entry_at_maxsize() -> None: + """The cache must not grow without bound, an entry is a full feature matrix.""" + for i in range(experiment._FEATURE_CACHE_MAXSIZE + 1): + _load_features_cached(_build(), "data", f"TOY{i}", "cell_line") + + assert len(experiment._FEATURE_CACHE) == experiment._FEATURE_CACHE_MAXSIZE + assert _CALLS["cell_line"] == experiment._FEATURE_CACHE_MAXSIZE + 1 + + # The first dataset was evicted and has to be loaded again. + _load_features_cached(_build(), "data", "TOY0", "cell_line") + assert _CALLS["cell_line"] == experiment._FEATURE_CACHE_MAXSIZE + 2 + + +def test_cache_can_be_disabled_via_environment(monkeypatch) -> None: + """DREVAL_FEATURE_CACHE=0 restores the previous load-every-time behaviour. + + :param monkeypatch: pytest monkeypatch fixture + """ + monkeypatch.setenv("DREVAL_FEATURE_CACHE", "0") + model = _build() + first = _load_features_cached(model, "data", "TOYv1", "cell_line") + second = _load_features_cached(model, "data", "TOYv1", "cell_line") + + assert _CALLS["cell_line"] == 2 + assert first is not second + assert experiment._FEATURE_CACHE == {} + + +def test_load_features_uses_the_cache() -> None: + """load_features is the entry point used by the pipeline and has to go through the cache.""" + dataset = DrugResponseDataset( + response=np.array([1.0]), + cell_line_ids=np.array(["CL0"]), + drug_ids=np.array(["D0"]), + dataset_name="TOYv1", + ) + model = _build() + cl_first, drug_first = load_features(model, "data", dataset) + cl_second, drug_second = load_features(model, "data", dataset) + + assert (_CALLS["cell_line"], _CALLS["drug"]) == (1, 1) + assert cl_first is cl_second + assert drug_first is drug_second From aaa1669ff36ffadfd53951e54c463eb4e3b83210 Mon Sep 17 00:00:00 2001 From: jfrog64 Date: Thu, 27 Aug 2026 09:59:47 +0200 Subject: [PATCH 2/2] fix: let models opt out of the feature cache when their loaders build state SparseGO's load_cell_line_features does not only return features, it also builds the ontology structure (layer_connections, gene2id_mapping_ont, ontology_gene_order, gene_dim_input) that train() and predict() need. The cache skips the loader for every instance after the first, so a fresh instance per CV split would have been left half built and train() would have raised "layer_connections or gene2id_mapping_ont are not set". DRPModel now carries supports_feature_caching (default True); SparseGO sets it to False and its features are loaded per instance again. A test walks MODEL_FACTORY and asserts that every model whose feature loaders assign to self has opted out, so the next such model cannot slip through. Co-Authored-By: Claude Opus 5 --- docs/runyourmodel.rst | 10 ++++ drevalpy/experiment.py | 4 +- drevalpy/models/SparseGO/sparsego.py | 4 ++ drevalpy/models/drp_model.py | 10 +++- tests/test_feature_cache.py | 81 ++++++++++++++++++++++++++++ 5 files changed, 107 insertions(+), 2 deletions(-) diff --git a/docs/runyourmodel.rst b/docs/runyourmodel.rst index 15a0f50c..1553a300 100644 --- a/docs/runyourmodel.rst +++ b/docs/runyourmodel.rst @@ -31,6 +31,7 @@ Additionally, you must define a unique model name to identify your model during is_single_drug_model = True / False # TODO: set to true if your model is a single drug model (i.e., it needs to be trained for each drug separately) early_stopping = True / False # TODO: set to true if you want to use a part of the validation set for early stopping + supports_feature_caching = True # TODO: set to false if your feature loading also initializes model state, see below cell_line_views = ["gene_expression", "methylation"] drug_views = ["fingerprints"] @@ -59,6 +60,15 @@ For our provided datasets, we have other loading methods implemented in the `dre * ``def load_drug_ids_from_csv`` * ``load_tissues_from_csv`` +Loading a feature matrix is often the most expensive part of a run, and the same matrix is needed again for +every CV split and, for single drug models, for every drug. The pipeline therefore caches what your loaders +return and hands the same object to later model instances, keyed by model class, data path, dataset name and +hyperparameters. Every consumer works on a ``copy()`` of it, so this is invisible to your model as long as your +loaders only return features. If they also initialize model state, e.g., an ontology structure that ``train`` +needs later, set ``supports_feature_caching = False`` on your class: the pipeline then calls your loaders for +every instance, as it did before the cache existed. Setting ``DREVAL_FEATURE_CACHE=0`` in the environment +disables the cache for all models. + .. code-block:: Python def load_drug_features(self, data_path: str, dataset_name: str) -> FeatureDataset: diff --git a/drevalpy/experiment.py b/drevalpy/experiment.py index d0dbcca7..b57da914 100644 --- a/drevalpy/experiment.py +++ b/drevalpy/experiment.py @@ -658,7 +658,9 @@ def _load_features_cached(model: DRPModel, path_data: str, dataset_name: str, ki :returns: the feature dataset, or None if the model does not use this kind of feature """ loader = model.load_cell_line_features if kind == "cell_line" else model.load_drug_features - if os.environ.get("DREVAL_FEATURE_CACHE", "1") == "0": + # Models whose loaders also initialize model state have to run them for every instance, and the cache + # can be switched off entirely to check that results are unchanged. + if not model.supports_feature_caching or os.environ.get("DREVAL_FEATURE_CACHE", "1") == "0": return loader(data_path=path_data, dataset_name=dataset_name) # The hyperparameters decide which views and which gene list are loaded, so they belong in the key. key = ( diff --git a/drevalpy/models/SparseGO/sparsego.py b/drevalpy/models/SparseGO/sparsego.py index 30c887b4..a9651d67 100644 --- a/drevalpy/models/SparseGO/sparsego.py +++ b/drevalpy/models/SparseGO/sparsego.py @@ -415,6 +415,10 @@ class SparseGOModel(DRPModel): cell_line_views = ["gene_expression", "mutations"] drug_views = ["fingerprints"] early_stopping = False + # load_cell_line_features also builds the ontology structure (layer_connections, gene2id_mapping_ont, + # ontology_gene_order, gene_dim_input) that train() and predict() need, so it has to run for every + # instance and the features must not be served from the pipeline's cache. + supports_feature_caching = False def __init__(self) -> None: """Initialize SparseGOModel.""" diff --git a/drevalpy/models/drp_model.py b/drevalpy/models/drp_model.py index 7599be2e..3ef3b02b 100644 --- a/drevalpy/models/drp_model.py +++ b/drevalpy/models/drp_model.py @@ -28,13 +28,21 @@ class DRPModel(ABC): The DRPModel class is an abstract wrapper class for drug response prediction models. It has a boolean attribute is_single_drug_model indicating whether it is a single drug model and a boolean - attribute early_stopping indicating whether early stopping is used. + attribute early_stopping indicating whether early stopping is used. The boolean attribute + supports_feature_caching tells the pipeline whether the loaded feature matrices may be reused across + model instances. """ # Used in the pipeline! early_stopping = False # Then, the model is trained per drug is_single_drug_model = False + # The pipeline caches the feature matrices a model loads and hands the same object to every later + # instance of that model, which is what makes repeated splits and single drug runs affordable. + # Set this to False if load_cell_line_features or load_drug_features do more than return features, + # i.e., if they also initialize model state. Such a loader has to run for every instance, so the + # pipeline then loads the features every time, as it did before the cache existed. See SparseGOModel. + supports_feature_caching = True def __init__(self): """Initialize the DRPModel instance.""" diff --git a/tests/test_feature_cache.py b/tests/test_feature_cache.py index 43707ab7..bb6cb45e 100644 --- a/tests/test_feature_cache.py +++ b/tests/test_feature_cache.py @@ -1,11 +1,16 @@ """Tests for the feature matrix cache used by the experiment pipeline.""" +import ast +import inspect +import textwrap + import numpy as np import pytest from drevalpy import experiment from drevalpy.datasets.dataset import DrugResponseDataset, FeatureDataset from drevalpy.experiment import _load_features_cached, clear_feature_cache, load_features +from drevalpy.models import MODEL_FACTORY from drevalpy.models.drp_model import DRPModel #: Number of loader calls per feature kind, shared by all test models so the cache effect is visible @@ -117,6 +122,52 @@ def get_model_name(cls) -> str: return "OtherCountingModel" +class _StatefulLoaderModel(_CountingModel): + """Model whose loader also initializes model state, like SparseGO builds its ontology there.""" + + supports_feature_caching = False + + @classmethod + def get_model_name(cls) -> str: + """ + Returns the model name. + + :returns: name of this test model + """ + return "StatefulLoaderModel" + + def load_cell_line_features(self, data_path: str, dataset_name: str) -> FeatureDataset: + """ + Returns features and, as a side effect, initializes state that train would need. + + :param data_path: unused + :param dataset_name: unused + :returns: cell line features + """ + features = super().load_cell_line_features(data_path=data_path, dataset_name=dataset_name) + self.ontology_loaded = True + return features + + +def _self_assignments(function) -> set[str]: + """ + Collect the attribute names a function assigns on ``self``. + + :param function: function to inspect + :returns: names assigned on self, empty if the function only returns something + """ + tree = ast.parse(textwrap.dedent(inspect.getsource(function))) + assigned = set() + for node in ast.walk(tree): + if isinstance(node, ast.Assign | ast.AugAssign | ast.AnnAssign): + targets = node.targets if isinstance(node, ast.Assign) else [node.target] + for target in targets: + if isinstance(target, ast.Attribute) and isinstance(target.value, ast.Name): + if target.value.id == "self": + assigned.add(target.attr) + return assigned + + @pytest.fixture(autouse=True) def _reset_cache(): """ @@ -237,3 +288,33 @@ def test_load_features_uses_the_cache() -> None: assert (_CALLS["cell_line"], _CALLS["drug"]) == (1, 1) assert cl_first is cl_second assert drug_first is drug_second + + +def test_models_that_opt_out_load_for_every_instance() -> None: + """A loader that initializes model state has to run for every instance, not once per process.""" + models = [_build(model_class=_StatefulLoaderModel) for _ in range(3)] + for model in models: + _load_features_cached(model, "data", "TOYv1", "cell_line") + + assert _CALLS["cell_line"] == 3 + assert experiment._FEATURE_CACHE == {} + assert all(getattr(model, "ontology_loaded", False) for model in models) + + +def test_stateful_loaders_are_marked_as_uncacheable() -> None: + """Every shipped model whose feature loaders assign to self must set supports_feature_caching False. + + The cache hands the same matrix to later instances and skips their loader, so a loader that also + initializes model state would leave those instances half built. + """ + stateful = { + name + for name, model_class in MODEL_FACTORY.items() + if _self_assignments(model_class.load_cell_line_features) or _self_assignments(model_class.load_drug_features) + } + not_opted_out = {name for name in stateful if MODEL_FACTORY[name].supports_feature_caching} + + assert not_opted_out == set(), ( + f"{sorted(not_opted_out)} initialize model state in their feature loaders and must set " + "supports_feature_caching = False" + )