Cache repeated in-process predictions

Author: Veda Sheersh Boorla

catpred/inference/service.py

@@ -1,11 +1,14 @@
 from __future__ import annotations
 
+from collections import OrderedDict
 from pathlib import Path
 from functools import lru_cache
 import gzip
+import hashlib
 import json
 import os
 import subprocess
+import threading
 from typing import Tuple
 
 import numpy as np
@@ -22,6 +25,9 @@
 }
 _VALID_AAS = set("ACDEFGHIKLMNPQRSTVWY")
 _MODEL_CACHE_SIZE = max(int(os.environ.get("CATPRED_MODEL_CACHE_SIZE", "6")), 0)
+_PREDICTION_CACHE_SIZE = max(int(os.environ.get("CATPRED_PREDICTION_CACHE_SIZE", "128")), 0)
+_PREDICTION_CACHE: OrderedDict[tuple, str] = OrderedDict()
+_PREDICTION_CACHE_LOCK = threading.Lock()
 
 
 def _validate_parameter(parameter: str) -> str:
@@ -285,6 +291,89 @@ def _expand_deduplicated_prediction_output(
     expanded_df.to_csv(original_paths.output_csv, index=False)
 
 
+def _file_digest(path: str) -> str:
+    digest = hashlib.sha256()
+    with open(path, "rb") as handle:
+        for chunk in iter(lambda: handle.read(1024 * 1024), b""):
+            digest.update(chunk)
+    return digest.hexdigest()
+
+
+def _prediction_checkpoint_paths(checkpoint_dir: str, repo_root: Path) -> tuple[str, ...]:
+    checkpoint_path = Path(checkpoint_dir)
+    if not checkpoint_path.is_absolute():
+        checkpoint_path = (repo_root / checkpoint_path).resolve()
+    if checkpoint_path.is_file():
+        return (str(checkpoint_path),)
+    if not checkpoint_path.is_dir():
+        raise FileNotFoundError(f'Checkpoint directory not found: "{checkpoint_path}"')
+    model_paths = sorted(str(path.resolve()) for path in checkpoint_path.rglob("model.pt"))
+    if not model_paths:
+        raise FileNotFoundError(f'No model.pt checkpoints found in "{checkpoint_path}"')
+    return tuple(model_paths)
+
+
+def _prediction_cache_key(
+    parameter: str,
+    request: PredictionRequest,
+    paths: PreparedInputPaths,
+    repo_root: Path,
+) -> tuple:
+    protein_records_digest = None
+    if request.protein_records_file:
+        protein_records_path = _resolve_existing_path(
+            request.protein_records_file,
+            repo_root=repo_root,
+            purpose="Protein records file",
+        )
+        protein_records_digest = _file_digest(str(protein_records_path))
+
+    checkpoint_paths = _prediction_checkpoint_paths(request.checkpoint_dir, repo_root)
+    return (
+        parameter,
+        bool(request.use_gpu),
+        _file_digest(paths.input_csv),
+        protein_records_digest,
+        _checkpoint_fingerprint(checkpoint_paths),
+    )
+
+
+def _prediction_cache_get(cache_key: tuple) -> str | None:
+    if _PREDICTION_CACHE_SIZE <= 0:
+        return None
+    with _PREDICTION_CACHE_LOCK:
+        cached = _PREDICTION_CACHE.get(cache_key)
+        if cached is not None:
+            _PREDICTION_CACHE.move_to_end(cache_key)
+        return cached
+
+
+def _prediction_cache_put(cache_key: tuple, csv_text: str) -> None:
+    if _PREDICTION_CACHE_SIZE <= 0:
+        return
+    with _PREDICTION_CACHE_LOCK:
+        _PREDICTION_CACHE[cache_key] = csv_text
+        _PREDICTION_CACHE.move_to_end(cache_key)
+        while len(_PREDICTION_CACHE) > _PREDICTION_CACHE_SIZE:
+            _PREDICTION_CACHE.popitem(last=False)
+
+
+def _write_cached_prediction(
+    cached_csv: str,
+    paths: PreparedInputPaths,
+    repo_root: str | None,
+    results_dir: str,
+) -> str:
+    results_path = Path(results_dir)
+    if not results_path.is_absolute():
+        results_path = (_resolve_repo_root(repo_root) / results_path).resolve()
+    results_path.mkdir(parents=True, exist_ok=True)
+
+    final_output = results_path / Path(paths.output_csv).name
+    final_output.write_text(cached_csv, encoding="utf-8")
+    return str(final_output)
+
+
 @lru_cache(maxsize=_MODEL_CACHE_SIZE)
 def _load_cached_model_objects(
     checkpoint_paths: tuple[str, ...],
@@ -441,14 +530,22 @@ def run_inprocess_prediction_pipeline(
 ) -> str:
     parameter = _validate_parameter(request.parameter)
     paths = prepare_prediction_inputs(parameter, request.input_file, request.repo_root)
+    root = _resolve_repo_root(request.repo_root)
+    cache_key = _prediction_cache_key(parameter, request, paths, root)
+    cached_csv = _prediction_cache_get(cache_key)
+    if cached_csv is not None:
+        return _write_cached_prediction(cached_csv, paths, request.repo_root, results_dir)
+
     if request.protein_records_file:
         prediction_paths, was_deduplicated = paths, False
     else:
         prediction_paths, was_deduplicated = _deduplicate_prediction_input(paths)
     run_inprocess_prediction(request, prediction_paths)
     if was_deduplicated:
         _expand_deduplicated_prediction_output(paths, prediction_paths)
-    return _write_postprocessed_predictions(parameter, paths, request.repo_root, results_dir)
+    final_output = _write_postprocessed_predictions(parameter, paths, request.repo_root, results_dir)
+    _prediction_cache_put(cache_key, Path(final_output).read_text(encoding="utf-8"))
+    return final_output
 
 
 def _write_postprocessed_predictions(

tests/test_inference_fast_path.py

@@ -163,6 +163,7 @@ def test_duplicate_inputs_are_expanded_back_to_original_rows(self) -> None:
                 + "\n",
                 encoding="utf-8",
             )
+
             paths = service.PreparedInputPaths(
                 input_csv=str(input_csv),
                 records_file=str(records_file),
@@ -194,6 +195,70 @@ def test_duplicate_inputs_are_expanded_back_to_original_rows(self) -> None:
         self.assertIn("third,O,BBBB,seq2.pdb,2.0,0.2", expanded_lines)
 
 
+class PredictionResultCacheTests(unittest.TestCase):
+    def setUp(self) -> None:
+        service._PREDICTION_CACHE.clear()
+
+    def tearDown(self) -> None:
+        service._PREDICTION_CACHE.clear()
+
+    def test_pipeline_reuses_cached_prediction_for_identical_request(self) -> None:
+        with TemporaryDirectory() as tmp_dir:
+            repo_root = Path(tmp_dir)
+            input_csv = repo_root / "prepared.csv"
+            input_csv.write_text("SMILES,sequence\nC,AAAA\n", encoding="utf-8")
+            records_file = repo_root / "prepared.json.gz"
+            records_file.write_bytes(b"records")
+            output_csv = repo_root / "prepared_output.csv"
+            checkpoint = repo_root / "checkpoints" / "fold_0" / "model_0" / "model.pt"
+            checkpoint.parent.mkdir(parents=True)
+            checkpoint.write_bytes(b"checkpoint")
+
+            paths = service.PreparedInputPaths(
+                input_csv=str(input_csv),
+                records_file=str(records_file),
+                output_csv=str(output_csv),
+            )
+            request = PredictionRequest(
+                parameter="kcat",
+                input_file=str(input_csv),
+                checkpoint_dir=str(checkpoint.parent.parent.parent),
+                repo_root=str(repo_root),
+            )
+
+            def write_final(parameter, prepared_paths, repo_root_arg, results_dir):
+                final_output = Path(results_dir) / Path(prepared_paths.output_csv).name
+                final_output.parent.mkdir(parents=True, exist_ok=True)
+                final_output.write_text("SMILES,kcat\nC,1.23\n", encoding="utf-8")
+                return str(final_output)
+
+            with patch(
+                "catpred.inference.service.prepare_prediction_inputs",
+                return_value=paths,
+            ), patch(
+                "catpred.inference.service.run_inprocess_prediction",
+            ) as runner, patch(
+                "catpred.inference.service._write_postprocessed_predictions",
+                side_effect=write_final,
+            ) as postprocess:
+                first = service.run_inprocess_prediction_pipeline(
+                    request,
+                    results_dir=str(repo_root / "results" / "first"),
+                )
+                second = service.run_inprocess_prediction_pipeline(
+                    request,
+                    results_dir=str(repo_root / "results" / "second"),
+                )
+                first_text = Path(first).read_text(encoding="utf-8")
+                second_text = Path(second).read_text(encoding="utf-8")
+
+        runner.assert_called_once()
+        postprocess.assert_called_once()
+        self.assertEqual(first_text, "SMILES,kcat\nC,1.23\n")
+        self.assertEqual(second_text, "SMILES,kcat\nC,1.23\n")
+        self.assertNotEqual(first, second)
+
+
 class FastPredictArgsTests(unittest.TestCase):
     def test_fast_predict_args_save_components_without_individual_predictions(self) -> None:
         class FakePredictArgs: