Merge pull request #33 from theproteinbot/codex/e2e-fixes

Harden inference API and add secure backend/deserialization architecture

Author: Vedasheersh

.gitignore

@@ -32,3 +32,16 @@ catpred/data/__pycache__/esm_utils.cpython-312.pyc
 catpred/data/__pycache__/data.cpython-312.pyc
 catpred/data/__pycache__/cache_utils.cpython-312.pyc
 catpred/data/__pycache__/__init__.cpython-312.pyc
+
+# Generic Python/OS artifacts
+__pycache__/
+*.pyo
+*.pyd
+.DS_Store
+*.egg-info/
+.venv/
+.ipynb_checkpoints/
+
+# Local validation artifacts
+.e2e-assets/
+.e2e-tests/

README.md

@@ -23,6 +23,7 @@
    - [System Requirements](#requirements)
    - [Installation](#installing)
    - [Prediction](#predict)
+   - [Web API (Optional)](#web-api-optional)
    - [Reproducibility](#reproduce)
 - [Acknowledgements](#acknw)
 - [License](#license)
@@ -65,7 +66,8 @@ Then proceed to either option below to complete the installation. If installing
 ```bash
 mkdir catpred_pipeline catpred_pipeline/results
 cd catpred_pipeline
-wget https://catpred.s3.us-east-1.amazonaws.com/capsule_data_update.tar.gz
+wget -c --tries=5 --timeout=30 https://catpred.s3.us-east-1.amazonaws.com/capsule_data_update.tar.gz || \
+wget -c --tries=5 --timeout=30 https://catpred.s3.amazonaws.com/capsule_data_update.tar.gz
 tar -xzf capsule_data_update.tar.gz
 git clone https://github.com/maranasgroup/catpred.git
 cd catpred
@@ -74,10 +76,147 @@ conda activate catpred
 pip install -e .
 ````
 
+`stride` is Linux-only and optional for the default demos. If needed for your workflow, install it separately on Linux:
+
+```bash
+conda install -c kimlab stride
+```
+
 ### 🔮 Prediction <a name="predict"></a>
 
 The Jupyter Notebook `batch_demo.ipynb` and the Python script `demo_run.py` show the usage of pre-trained models for prediction.
 
+Input CSV requirements for `demo_run.py` and batch prediction:
+- Required columns: `SMILES`, `sequence`, `pdbpath`.
+- `pdbpath` must be unique per unique sequence. Reusing the same `pdbpath` for different sequences can produce incorrect cached embeddings.
+- Reusing the same `pdbpath` for repeated measurements of the same sequence is supported.
+
+The helper script used to build protein records is:
+
+```bash
+python ./scripts/create_pdbrecords.py --data_file <input.csv> --out_file <input.json.gz>
+```
+
+CatPred currently expects one sequence per row. Multi-protein complexes (e.g., heteromers/homodimers) are not explicitly modeled as separate chains in the default prediction workflow.
+
+For released benchmark datasets, the number of entries with 3D structure can be smaller than the total sequence/substrate pairs; 3D-derived artifacts are available only for the subset with valid structure mapping.
+
+### 🌍 Web API (Optional)
+
+CatPred also provides an optional FastAPI service for prediction workflows.
+
+Install web dependencies:
+
+```bash
+pip install -e ".[web]"
+```
+
+Run the API:
+
+```bash
+catpred_web --host 0.0.0.0 --port 8000
+```
+
+Endpoints:
+- `GET /health` — liveness check.
+- `GET /ready` — backend configuration/readiness.
+- `POST /predict` — run inference.
+
+By default, the API is hardened for service use:
+- `input_file` requests are disabled (use `input_rows` instead).
+- request-time overrides of `repo_root` / `python_executable` are disabled.
+- `results_dir` is constrained under `CATPRED_API_RESULTS_ROOT`.
+
+Minimal `POST /predict` example for local inference using `input_rows`:
+
+```bash
+curl -X POST http://127.0.0.1:8000/predict \
+  -H "Content-Type: application/json" \
+  -d '{
+    "parameter": "kcat",
+    "checkpoint_dir": "../data/pretrained/reproduce_checkpoints/kcat",
+    "input_rows": [
+      {"SMILES": "CCO", "sequence": "ACDEFGHIK", "pdbpath": "seq_a"},
+      {"SMILES": "CCN", "sequence": "LMNPQRSTV", "pdbpath": "seq_b"}
+    ],
+    "results_dir": "batch1",
+    "backend": "local"
+  }'
+```
+
+You can keep local inference as default and optionally enable Modal as another backend:
+
+```bash
+export CATPRED_DEFAULT_BACKEND=local
+export CATPRED_MODAL_ENDPOINT="https://<your-modal-endpoint>"
+export CATPRED_MODAL_TOKEN="<optional-token>"
+export CATPRED_MODAL_FALLBACK_TO_LOCAL=1
+```
+
+Use `"backend": "modal"` in `/predict` requests to route through Modal. If fallback is enabled (env var above or request field `fallback_to_local`), failed modal requests can automatically reroute to local inference.
+
+Optional API environment variables:
+
+```bash
+# Root directories used by API path constraints
+export CATPRED_API_INPUT_ROOT="/absolute/path/for/input-csvs"
+export CATPRED_API_RESULTS_ROOT="/absolute/path/for/results"
+export CATPRED_API_CHECKPOINT_ROOT="/absolute/path/for/checkpoints"
+
+# Enable only for trusted local workflows (not recommended for public deployments)
+export CATPRED_API_ALLOW_INPUT_FILE=1
+export CATPRED_API_ALLOW_UNSAFE_OVERRIDES=1
+
+# Request limits
+export CATPRED_API_MAX_INPUT_ROWS=1000
+export CATPRED_API_MAX_INPUT_FILE_BYTES=5000000
+```
+
+Deserialization hardening controls:
+
+```bash
+# Trusted roots used by secure loaders (colon-separated list on Unix)
+export CATPRED_TRUSTED_DESERIALIZATION_ROOTS="/srv/catpred:/srv/catpred-data"
+
+# Backward-compatible default is enabled (1). Set to 0 to block unsafe pickle-based loading.
+# Use 0 only after validating your artifacts are safe-load compatible.
+export CATPRED_ALLOW_UNSAFE_DESERIALIZATION=1
+```
+
+### 🧪 Fine-Tuning On Custom Data
+
+You can fine-tune CatPred on your own regression targets using `train.py`.
+
+1. Prepare train/val/test CSVs with at least:
+- `SMILES`
+- `sequence`
+- `pdbpath` (unique per unique sequence)
+- one numeric target column (for example: `log10kcat_max`)
+
+2. Build a protein-records file that covers all `pdbpath` values in your splits:
+
+```bash
+python ./scripts/create_pdbrecords.py --data_file <combined_or_train_csv> --out_file <protein_records.json.gz>
+```
+
+3. Train:
+
+```bash
+python train.py \
+  --protein_records_path <protein_records.json.gz> \
+  --data_path <train.csv> \
+  --separate_val_path <val.csv> \
+  --separate_test_path <test.csv> \
+  --dataset_type regression \
+  --smiles_columns SMILES \
+  --target_columns <target_column_name> \
+  --add_esm_feats \
+  --loss_function mve \
+  --save_dir <output_checkpoint_dir>
+```
+
+For working end-to-end examples, see the training commands in scripts such as `scripts/reproduce_figS10_catpred.sh`.
+
 ### 🔄 Reproducing Publication Results <a name="reproduce"></a>
 
 We provide three separate ways for reproducing the results of the publication.

catpred/__init__.py

@@ -1,13 +1,30 @@
-import catpred.data
-import catpred.features
-import catpred.models
-import catpred.train
-import catpred.uncertainty
-
-import catpred.args
-import catpred.constants
-import catpred.nn_utils
-import catpred.utils
-import catpred.rdkit
+from __future__ import annotations
+
+import importlib
 
 __version__ = "0.0.1"
+
+_LAZY_SUBMODULES = {
+    "args",
+    "constants",
+    "data",
+    "features",
+    "inference",
+    "models",
+    "nn_utils",
+    "rdkit",
+    "security",
+    "train",
+    "uncertainty",
+    "utils",
+}
+
+__all__ = sorted(_LAZY_SUBMODULES) + ["__version__"]
+
+
+def __getattr__(name: str):
+    if name in _LAZY_SUBMODULES:
+        module = importlib.import_module(f"catpred.{name}")
+        globals()[name] = module
+        return module
+    raise AttributeError(f"module 'catpred' has no attribute '{name}'")

catpred/args.py

@@ -1,7 +1,6 @@
 import json
 import os
 from tempfile import TemporaryDirectory
-import pickle
 from typing import List, Optional
 from typing_extensions import Literal
 from packaging import version
@@ -15,6 +14,7 @@
 import catpred.data.utils
 from catpred.data import set_cache_mol, empty_cache
 from catpred.features import get_available_features_generators
+from catpred.security import load_index_artifact
 
 
 Metric = Literal['auc', 'prc-auc', 'rmse', 'mae', 'mse', 'r2', 'accuracy', 'cross_entropy', 'binary_cross_entropy', 'sid', 'wasserstein', 'f1', 'mcc', 'bounded_rmse', 'bounded_mae', 'bounded_mse']
@@ -815,8 +815,10 @@ def process_args(self) -> None:
             raise ValueError('When using crossval or index_predetermined split type, must provide crossval_index_file.')
 
         if self.split_type in ['crossval', 'index_predetermined']:
-            with open(self.crossval_index_file, 'rb') as rf:
-                self._crossval_index_sets = pickle.load(rf)
+            self._crossval_index_sets = load_index_artifact(
+                self.crossval_index_file,
+                purpose="cross-validation index file",
+            )
             self.num_folds = len(self.crossval_index_sets)
             self.seed = 0
 

catpred/data/cache_utils.py

@@ -4,6 +4,7 @@
 import hashlib
 from functools import wraps
 from pathlib import Path
+from catpred.security import load_torch_artifact
 
 def exists(val):
     return val is not None
@@ -92,7 +93,11 @@ def inner(t, *args, __cache_key = None, **kwargs):
 
         if entry_path.exists():
             log(f'cache hit: fetching {t} from {str(entry_path)}')
-            return torch.load(str(entry_path))
+            return load_torch_artifact(
+                str(entry_path),
+                purpose="esm cache entry",
+                roots=[CACHE_PATH],
+            )
 
         out = fn(t, *args, **kwargs)