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 9fee1a6f..b57da914 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,82 @@ 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 + # 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 = ( + 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 +712,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 +1181,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 +1693,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/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 new file mode 100644 index 00000000..bb6cb45e --- /dev/null +++ b/tests/test_feature_cache.py @@ -0,0 +1,320 @@ +"""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 +#: 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" + + +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(): + """ + 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 + + +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" + )