Merge pull request #141 from aqlaboratory/hotfix/setup-download-all-parameters

Fix setup-openfold for download all parameters option

Author: jnwei

openfold3/entry_points/parameters.py

@@ -36,19 +36,23 @@ class CheckpointEntry:
 
 
 OPENFOLD_MODEL_CHECKPOINT_REGISTRY = {
-    "openfold3-p1": CheckpointEntry(
-        file_name="of3_ft3_v1.pt", version_compatibility="<0.4"
-    ),
     "openfold3-p2-145k": CheckpointEntry(
         file_name="of3-p2-145k.pt", version_compatibility=">=0.4"
     ),
     "openfold3-p2-155k": CheckpointEntry(
         file_name="of3-p2-155k.pt", version_compatibility=">=0.4"
     ),
+    "openfold3-p1": CheckpointEntry(  # legacy
+        file_name="of3_ft3_v1.pt", version_compatibility="<0.4"
+    ),
 }
 
 DEFAULT_CHECKPOINT_NAME = "openfold3-p2-155k"
 
+# These checkpoints are not supported for download and use in the current version,
+# but are left in the registry for record-keeping and compatibility checks.
+LEGACY_CHECKPOINTS = ["openfold3-p1"]
+
 
 def download_model_parameters(
     download_dir: Path,

openfold3/setup_openfold.py

@@ -29,6 +29,7 @@
 from openfold3.core.utils.s3 import download_s3_file, s3_file_matches_local
 from openfold3.entry_points.parameters import (
     DEFAULT_CHECKPOINT_NAME,
+    LEGACY_CHECKPOINTS,
     OPENFOLD_MODEL_CHECKPOINT_REGISTRY,
     download_model_parameters,
 )
@@ -131,7 +132,12 @@ def setup_param_directory(
 
 def download_parameters(param_dir) -> None:
     """Perform the parameter download."""
-    all_checkpoints = list(OPENFOLD_MODEL_CHECKPOINT_REGISTRY.keys())
+    # Exclude incompatible checkpoints:
+    all_checkpoints = [
+        name
+        for name in OPENFOLD_MODEL_CHECKPOINT_REGISTRY
+        if name not in LEGACY_CHECKPOINTS
+    ]
 
     logger.info("Select parameters to download:")
     logger.info(f"1) Download only the default checkpoint ({DEFAULT_CHECKPOINT_NAME})")

openfold3/tests/test_entry_points.py

@@ -36,6 +36,7 @@
 from openfold3.entry_points.parameters import (
     CHECKPOINT_ROOT_FILENAME,
     DEFAULT_CHECKPOINT_NAME,
+    LEGACY_CHECKPOINTS,
     OPENFOLD_MODEL_CHECKPOINT_REGISTRY,
     CheckpointEntry,
 )
@@ -699,7 +700,7 @@ def test_remove_duplicates(self, dummy_ckpt_file, dummy_output_path):
 
 
 class TestSetupOpenFold:
-    def test_fresh_parameter_download(self, tmp_path):
+    def test_fresh_parameter_default_download(self, tmp_path):
         inputs = iter(
             [
                 str(tmp_path),  # Set cache directory
@@ -726,3 +727,34 @@ def test_fresh_parameter_download(self, tmp_path):
             tmp_path
             / OPENFOLD_MODEL_CHECKPOINT_REGISTRY[DEFAULT_CHECKPOINT_NAME].file_name
         ).exists()
+
+    def test_fresh_parameter_download_all(self, tmp_path):
+        inputs = iter(
+            [
+                str(tmp_path),  # Set cache directory
+                "",  # Use default (cache) directory for params directory
+                "2",  # download choice: all parameters
+                "no",  # skip integration tests
+            ]
+        )
+
+        with (
+            patch("builtins.input", side_effect=inputs),
+            patch(
+                "openfold3.setup_openfold.download_s3_file",
+                side_effect=_fake_download_s3_file,
+            ),
+        ):
+            setup_openfold.main()
+
+        # Check that the checkpoint root file exists and has the expected path
+        assert (tmp_path / CHECKPOINT_ROOT_FILENAME).exists()
+        assert (tmp_path / CHECKPOINT_ROOT_FILENAME).read_text() == str(tmp_path)
+
+        expected_checkpoints = list(
+            set(OPENFOLD_MODEL_CHECKPOINT_REGISTRY.keys()) - set(LEGACY_CHECKPOINTS)
+        )
+        for ckpt_name in expected_checkpoints:
+            assert (
+                tmp_path / OPENFOLD_MODEL_CHECKPOINT_REGISTRY[ckpt_name].file_name
+            ).exists()