Add a friendly deprecation warning for pae_enabled model preset

Author: jnwei

openfold3/projects/of3_all_atom/project_entry.py

@@ -18,19 +18,32 @@
 from dataclasses import dataclass
 
 from ml_collections import ConfigDict
-from pydantic import BaseModel
+from pydantic import BaseModel, field_validator
 from pydantic import ConfigDict as PydanticConfigDict
 
 from openfold3.core.config.config_utils import load_yaml
 from openfold3.projects.of3_all_atom.config.model_config import model_config
 from openfold3.projects.of3_all_atom.runner import OpenFold3AllAtom
 
+logger = logging.getLogger(__name__)
+
 
 class ModelUpdate(BaseModel):
     model_config = PydanticConfigDict(extra="forbid")
     presets: list[str] = []
     custom: dict = {}
 
+    @field_validator("presets", mode="before")
+    def warn_if_pae_enabled_included(cls, v):
+        if "pae_enabled" in v:
+            logger.warning(
+                "The `pae_enabled` model preset is deprecated and will be removed."
+                " from the ModelUpdate.presets list. Please remove `pae_enabled`"
+                " from your model update section"
+            )
+            v.remove("pae_enabled")
+        return v
+
 
 @dataclass
 class OF3ProjectEntry:

openfold3/tests/test_entry_points.py

@@ -316,6 +316,14 @@ def test_low_mem_model_config_preset(self, tmp_path, dummy_ckpt_file):
         # test existing setting in experiment runner is not overwritten
         assert not model_cfg.settings.memory.eval.use_lma
 
+    def test_model_update_with_pae_enabled_triggers_warning(self):
+        with patch(
+            "openfold3.projects.of3_all_atom.project_entry.logger"
+        ) as mock_logger:
+            ModelUpdate.model_validate({"presets": ["predict", "pae_enabled"]})
+        warning_messages = [call.args[0] for call in mock_logger.warning.call_args_list]
+        assert any("model preset is deprecated" in msg for msg in warning_messages)
+
 
 class DummyWandbExperiment:
     def __init__(self, directory):