Merge pull request #438 from jnwei/pl_upgrades

Upgrades pl_upgrades to match main branch changes.

.github/workflows/docker-image.yml

@@ -11,5 +11,7 @@ jobs:
     runs-on: ubuntu-latest
     steps:
     - uses: actions/checkout@v4
+    - name: Cleanup
+      run: sudo rm -rf /usr/share/dotnet && sudo rm -rf /opt/ghc && sudo rm -rf "/usr/local/share/boost" && sudo rm -rf "$AGENT_TOOLSDIRECTORY"
     - name: Build the Docker image
-      run: docker build . --file Dockerfile --tag openfold:$(date +%s)
+      run: docker build . --file Dockerfile --tag openfold:$(date +%s)

.readthedocs.yaml

@@ -0,0 +1,14 @@
+version: 2
+
+# Set the OS, Python version and other tools you might need
+build:
+  os: ubuntu-22.04
+  tools:
+    python: "mambaforge-4.10" 
+
+# Build documentation in the "docs/" directory with Sphinx
+sphinx:
+  configuration: docs/source/conf.py
+
+conda:
+  environment: docs/environment.yml

Dockerfile

@@ -1,17 +1,20 @@
-FROM nvidia/cuda:11.3.1-cudnn8-devel-ubuntu18.04
+FROM nvidia/cuda:12.1.1-cudnn8-devel-ubuntu22.04
 
 # metainformation
-LABEL org.opencontainers.image.version = "1.0.0"
-LABEL org.opencontainers.image.authors = "Gustaf Ahdritz"
+LABEL org.opencontainers.image.version = "2.0.0"
+LABEL org.opencontainers.image.authors = "OpenFold Team"
 LABEL org.opencontainers.image.source = "https://github.com/aqlaboratory/openfold"
 LABEL org.opencontainers.image.licenses = "Apache License 2.0"
-LABEL org.opencontainers.image.base.name="docker.io/nvidia/cuda:10.2-cudnn8-runtime-ubuntu18.04"
+LABEL org.opencontainers.image.base.name="docker.io/nvidia/cuda:12.4.1-devel-ubuntu22.04"
+
+RUN apt-get update && apt-get install -y wget
 
 RUN apt-key del 7fa2af80
-RUN apt-key adv --fetch-keys https://developer.download.nvidia.com/compute/cuda/repos/ubuntu1804/x86_64/7fa2af80.pub
-RUN apt-key adv --fetch-keys https://developer.download.nvidia.com/compute/cuda/repos/ubuntu1804/x86_64/3bf863cc.pub
+RUN wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/cuda-keyring_1.0-1_all.deb
+RUN dpkg -i cuda-keyring_1.0-1_all.deb
+
+RUN apt-get install -y libxml2 cuda-minimal-build-12-1 libcusparse-dev-12-1 libcublas-dev-12-1 libcusolver-dev-12-1 git
 
-RUN apt-get update && apt-get install -y wget libxml2 cuda-minimal-build-11-3 libcusparse-dev-11-3 libcublas-dev-11-3 libcusolver-dev-11-3 git
 RUN wget -P /tmp \
     "https://github.com/conda-forge/miniforge/releases/download/23.3.1-1/Miniforge3-Linux-x86_64.sh" \
     && bash /tmp/Miniforge3-Linux-x86_64.sh -b -p /opt/conda \

LICENSE

@@ -187,7 +187,7 @@
       same "printed page" as the copyright notice for easier
       identification within third-party archives.
 
-   Copyright [yyyy] [name of copyright owner]
+   Copyright 2024 AlQuraishi Laboratory 
 
    Licensed under the Apache License, Version 2.0 (the "License");
    you may not use this file except in compliance with the License.

docs/Makefile

@@ -0,0 +1,20 @@
+# Minimal makefile for Sphinx documentation
+#
+
+# You can set these variables from the command line, and also
+# from the environment for the first two.
+SPHINXOPTS    ?=
+SPHINXBUILD   ?= sphinx-build
+SOURCEDIR     = source
+BUILDDIR      = build
+
+# Put it first so that "make" without argument is like "make help".
+help:
+	@$(SPHINXBUILD) -M help "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O)
+
+.PHONY: help Makefile
+
+# Catch-all target: route all unknown targets to Sphinx using the new
+# "make mode" option.  $(O) is meant as a shortcut for $(SPHINXOPTS).
+%: Makefile
+	@$(SPHINXBUILD) -M $@ "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O)

docs/environment.yml

@@ -0,0 +1,7 @@
+name: openfold-docs
+channels:
+  - conda-forge
+dependencies:
+  - sphinx=7
+  - myst-parser=3
+  - furo

docs/source/Aux_seq_files.md

@@ -0,0 +1,203 @@
+# Auxiliary Sequence Files for OpenFold Training
+
+The training dataset of OpenFold is very large. The `pdb` directory alone contains 185,000 mmcifs;  each chain for has multiple sequence alignment files and mmcif files. 
+
+OpenFold introduces a few new file structures for faster access to alignments and mmcif data.
+
+This documentation will explain the benefits of having the condensed file structure, and explain the contents of each of the files.
+
+## Default alignment file structure 
+
+One way to store mmcifs and alignments files would be to have a directory for each mmcif chain. 
+
+For example, consider two protein as a case study
+```
+- OpenProteinSet
+  └── mmcifs 
+	 ├── 3lrm.cif
+	 └── 6kwc.cif
+	 ...
+```
+
+In the `alignments` directory, [PDB:6KWC](https://www.rcsb.org/structure/6KWC) is a monomer with one chain, and thus would have one alignment direcotry. [PDB:3LRM](https://www.rcsb.org/structure/3lrm), a homotetramer, would have one alignment directory for each of its four chains.
+```
+- OpenProteinSet 
+  └── alignments 
+	  └── 3lrm_A 
+		  ├── bfd_uniclust_hits.a3m
+		  ├── mgnify_hits.a3m
+		  ├── pdb70_hits.hhr
+		  └── uniref90_hits.a3m 
+	  └── 3lrm_B 
+		  ├── bfd_uniclust_hits.a3m
+		  ├── mgnify_hits.a3m
+		  ├── pdb70_hits.hhr
+		  └── uniref90_hits.a3m 
+	  └── 3lrm_C 
+		  ├── bfd_uniclust_hits.a3m
+		  ├── mgnify_hits.a3m
+		  ├── pdb70_hits.hhr
+		  └── uniref90_hits.a3m 
+	  └── 3lrm_D 
+		  ├── bfd_uniclust_hits.a3m
+		  ├── mgnify_hits.a3m
+		  ├── pdb70_hits.hhr
+		  └── uniref90_hits.a3m 
+	  └── 6kwc_A 
+		  ├── bfd_uniclust_hits.a3m
+		  ├── mgnify_hits.a3m
+		  ├── pdb70_hits.hhr
+		  └── uniref90_hits.a3m 
+		...
+```
+
+In practice, the IO overhead of having one directory per protein chain makes accessing the alignments slow. 
+
+## OpenFold DB file structure 
+
+ Here we describe a new filesystem that can be used by OpenFold for more efficient access of alignment file and index file contents
+
+All together, the file directory would look like:
+```
+- OpenProteinSet 
+  ├── duplicate_pdb_chains.txt
+  └── pdb
+	  ├── mmcif_cache.json 
+	  └── mmcifs 
+		  ├── 3lrm.cif
+		  └── 6kwc.cif
+	  └── alignment_db
+		  ├── alignment_db_0.db 
+          ├── alignment_db_1.db
+          ...
+          ├── alignment_db_9.db
+		  └── alignment_db.index 
+```
+
+We will describe each of the file types here.
+
+### Alignments db files and index files
+
+To speed up access of MSAs, OpenFold has an alternate alignments storage procedure. Instead of storing dedicated files for each single alignment, we consolidate large sets of alignments to single files referred to as _alignments_db's_. This can reduce I/O overhead and in practice we recommend using around 10 `alignments_db_x.db` files to store the total training set of OpenFold. During training, OpenFold can access each alignment using byte index pointers that are stored in a separate index file (`alignments_db.index`). The alignments for the `3LRM` and `6KWC` examples would be recorded in the index file as follows:
+
+```alignments_db.index
+{
+	...
+	"3lrm_A": {
+        "db": "alignment_db_0.db",
+        "files": [
+            ["bfd_uniclust_hits.a3m", 212896478938, 1680200],
+            ["mgnify_hits.a3m", 212893696883, 2782055],
+            ["pdb70_hits.hhr", 212898159138, 614978],
+            ["uniref90_hits.a3m", 212898774116, 6165789]
+        ]
+    },
+    "6kwc_A": {
+        "db": "alignment_db_1.db",
+        "files": [
+            ["bfd_uniclust_hits.a3m", 415618723280, 380289],
+            ["mgnify_hits.a3m", 415618556077, 167203],
+            ["pdb70_hits.hhr", 415619103569, 148672],
+            ["uniref90_hits.a3m", 415617547852, 1008225]
+        ]
+    }
+	...
+}
+```
+
+For each entry, the corresponding `alignment_db` file and the byte start location and number of bytes to read the respective alignments are given. For example, the alignment information in `bfd_uniclust_hits.a3m` for chain `3lrm_A` can be found in the database file `alignment_db_0.db`, starting at byte location `212896478938` and reading in the next `1680200` bytes.
+
+### Chain cache files and mmCIF cache files
+
+Information from the mmcif files can be parsed in advance to create a `chain_cache.json` or a `mmcif_cache.json`. For OpenFold, the `chain_cache.json` is used to sample chains for training, and the `mmcif_cache.json` is used to prefilter templates. 
+
+Here's what the chain_cache.json entry looks like for our examples:
+
+```chain_cache.json
+{
+	...
+	    "3lrm_A": {
+        "release_date": "2010-06-30",
+        "seq": "MFAFYFLTACISLKGVFGVSPSYNGLGLTPQMGWDNWNTFACDVSEQLLLDTADRISDLGLKDMGYKYIILDDCWSSGRDSDGFLVADEQKFPNGMGHVADHLHNNSFLFGMYSSAGEYTCAGYPGSLGREEEDAQFFANNRVDYLKYANCYNKGQFGTPEISYHRYKAMSDALNKTGRPVFYSLCNWGQDLTFYWGSGIANSWRMSGDVTAEFTRPDSRCPCDGDEYDCKYAGFHCSIMNILNKAAPMGQNAGVGGWNDLDNLEVGVGNLTDDEEKAHFSMWAMVKSPLIIGANVNNLKASSYSIYSQASVIAINQDSNGIPATRVWRYYVSDTDEYGQGEIQMWSGPLDNGDQVVALLNGGSVSRPMNTTLEEIFFDSNLGSKKLTSTWDIYDLWANRVDNSTASAILGRNKTATGILYNATEQSYKDGLSKNDTRLFGQKIGSLSPNAILNTTVPAHGIAFYRLRPSSDYKDDDDK",
+        "resolution": 2.7,
+        "cluster_size": 6
+    },
+    "3lrm_B": {
+        "release_date": "2010-06-30",
+        "seq": "MFAFYFLTACISLKGVFGVSPSYNGLGLTPQMGWDNWNTFACDVSEQLLLDTADRISDLGLKDMGYKYIILDDCWSSGRDSDGFLVADEQKFPNGMGHVADHLHNNSFLFGMYSSAGEYTCAGYPGSLGREEEDAQFFANNRVDYLKYANCYNKGQFGTPEISYHRYKAMSDALNKTGRPVFYSLCNWGQDLTFYWGSGIANSWRMSGDVTAEFTRPDSRCPCDGDEYDCKYAGFHCSIMNILNKAAPMGQNAGVGGWNDLDNLEVGVGNLTDDEEKAHFSMWAMVKSPLIIGANVNNLKASSYSIYSQASVIAINQDSNGIPATRVWRYYVSDTDEYGQGEIQMWSGPLDNGDQVVALLNGGSVSRPMNTTLEEIFFDSNLGSKKLTSTWDIYDLWANRVDNSTASAILGRNKTATGILYNATEQSYKDGLSKNDTRLFGQKIGSLSPNAILNTTVPAHGIAFYRLRPSSDYKDDDDK",
+        "resolution": 2.7,
+        "cluster_size": 6
+    },
+    "3lrm_C": {
+        "release_date": "2010-06-30",
+        "seq": "MFAFYFLTACISLKGVFGVSPSYNGLGLTPQMGWDNWNTFACDVSEQLLLDTADRISDLGLKDMGYKYIILDDCWSSGRDSDGFLVADEQKFPNGMGHVADHLHNNSFLFGMYSSAGEYTCAGYPGSLGREEEDAQFFANNRVDYLKYANCYNKGQFGTPEISYHRYKAMSDALNKTGRPVFYSLCNWGQDLTFYWGSGIANSWRMSGDVTAEFTRPDSRCPCDGDEYDCKYAGFHCSIMNILNKAAPMGQNAGVGGWNDLDNLEVGVGNLTDDEEKAHFSMWAMVKSPLIIGANVNNLKASSYSIYSQASVIAINQDSNGIPATRVWRYYVSDTDEYGQGEIQMWSGPLDNGDQVVALLNGGSVSRPMNTTLEEIFFDSNLGSKKLTSTWDIYDLWANRVDNSTASAILGRNKTATGILYNATEQSYKDGLSKNDTRLFGQKIGSLSPNAILNTTVPAHGIAFYRLRPSSDYKDDDDK",
+        "resolution": 2.7,
+        "cluster_size": 6
+    },
+    "3lrm_D": {
+        "release_date": "2010-06-30",
+        "seq": "MFAFYFLTACISLKGVFGVSPSYNGLGLTPQMGWDNWNTFACDVSEQLLLDTADRISDLGLKDMGYKYIILDDCWSSGRDSDGFLVADEQKFPNGMGHVADHLHNNSFLFGMYSSAGEYTCAGYPGSLGREEEDAQFFANNRVDYLKYANCYNKGQFGTPEISYHRYKAMSDALNKTGRPVFYSLCNWGQDLTFYWGSGIANSWRMSGDVTAEFTRPDSRCPCDGDEYDCKYAGFHCSIMNILNKAAPMGQNAGVGGWNDLDNLEVGVGNLTDDEEKAHFSMWAMVKSPLIIGANVNNLKASSYSIYSQASVIAINQDSNGIPATRVWRYYVSDTDEYGQGEIQMWSGPLDNGDQVVALLNGGSVSRPMNTTLEEIFFDSNLGSKKLTSTWDIYDLWANRVDNSTASAILGRNKTATGILYNATEQSYKDGLSKNDTRLFGQKIGSLSPNAILNTTVPAHGIAFYRLRPSSDYKDDDDK",
+        "resolution": 2.7,
+        "cluster_size": 6
+    },
+	"6kwc_A": {
+        "release_date": "2021-01-27",
+        "seq": "GSTIQPGTGYNNGYFYSYWNDGHGGVTYTNGPGGQFSVNWSNSGEFVGGKGWQPGTKNKVINFSGSYNPNGNSYLSVYGWSRNPLIEYYIVENFGTYNPSTGATKLGEVTSDGSVYDIYRTQRVNQPSIIGTATFYQYWSVRRNHRSSGSVNTANHFNAWAQQGLTLGTMDYQIVAVQGYFSSGSASITVS",
+        "resolution": 1.297,
+        "cluster_size": 195
+    },
+	...
+}
+```
+
+The mmcif_cache.json file would contain similar information, but condensed by mmcif id, e.g.
+
+```mmcif_cache.json
+{
+    "3lrm": {
+        "release_date": "2010-06-30",
+        "chain_ids": [
+            "A",
+            "B",
+            "C",
+            "D"
+        ],
+        "seqs": [
+            "MFAFYFLTACISLKGVFGVSPSYNGLGLTPQMGWDNWNTFACDVSEQLLLDTADRISDLGLKDMGYKYIILDDCWSSGRDSDGFLVADEQKFPNGMGHVADHLHNNSFLFGMYSSAGEYTCAGYPGSLGREEEDAQFFANNRVDYLKYANCYNKGQFGTPEISYHRYKAMSDALNKTGRPVFYSLCNWGQDLTFYWGSGIANSWRMSGDVTAEFTRPDSRCPCDGDEYDCKYAGFHCSIMNILNKAAPMGQNAGVGGWNDLDNLEVGVGNLTDDEEKAHFSMWAMVKSPLIIGANVNNLKASSYSIYSQASVIAINQDSNGIPATRVWRYYVSDTDEYGQGEIQMWSGPLDNGDQVVALLNGGSVSRPMNTTLEEIFFDSNLGSKKLTSTWDIYDLWANRVDNSTASAILGRNKTATGILYNATEQSYKDGLSKNDTRLFGQKIGSLSPNAILNTTVPAHGIAFYRLRPSSDYKDDDDK",
+            "MFAFYFLTACISLKGVFGVSPSYNGLGLTPQMGWDNWNTFACDVSEQLLLDTADRISDLGLKDMGYKYIILDDCWSSGRDSDGFLVADEQKFPNGMGHVADHLHNNSFLFGMYSSAGEYTCAGYPGSLGREEEDAQFFANNRVDYLKYANCYNKGQFGTPEISYHRYKAMSDALNKTGRPVFYSLCNWGQDLTFYWGSGIANSWRMSGDVTAEFTRPDSRCPCDGDEYDCKYAGFHCSIMNILNKAAPMGQNAGVGGWNDLDNLEVGVGNLTDDEEKAHFSMWAMVKSPLIIGANVNNLKASSYSIYSQASVIAINQDSNGIPATRVWRYYVSDTDEYGQGEIQMWSGPLDNGDQVVALLNGGSVSRPMNTTLEEIFFDSNLGSKKLTSTWDIYDLWANRVDNSTASAILGRNKTATGILYNATEQSYKDGLSKNDTRLFGQKIGSLSPNAILNTTVPAHGIAFYRLRPSSDYKDDDDK",
+            "MFAFYFLTACISLKGVFGVSPSYNGLGLTPQMGWDNWNTFACDVSEQLLLDTADRISDLGLKDMGYKYIILDDCWSSGRDSDGFLVADEQKFPNGMGHVADHLHNNSFLFGMYSSAGEYTCAGYPGSLGREEEDAQFFANNRVDYLKYANCYNKGQFGTPEISYHRYKAMSDALNKTGRPVFYSLCNWGQDLTFYWGSGIANSWRMSGDVTAEFTRPDSRCPCDGDEYDCKYAGFHCSIMNILNKAAPMGQNAGVGGWNDLDNLEVGVGNLTDDEEKAHFSMWAMVKSPLIIGANVNNLKASSYSIYSQASVIAINQDSNGIPATRVWRYYVSDTDEYGQGEIQMWSGPLDNGDQVVALLNGGSVSRPMNTTLEEIFFDSNLGSKKLTSTWDIYDLWANRVDNSTASAILGRNKTATGILYNATEQSYKDGLSKNDTRLFGQKIGSLSPNAILNTTVPAHGIAFYRLRPSSDYKDDDDK",
+            "MFAFYFLTACISLKGVFGVSPSYNGLGLTPQMGWDNWNTFACDVSEQLLLDTADRISDLGLKDMGYKYIILDDCWSSGRDSDGFLVADEQKFPNGMGHVADHLHNNSFLFGMYSSAGEYTCAGYPGSLGREEEDAQFFANNRVDYLKYANCYNKGQFGTPEISYHRYKAMSDALNKTGRPVFYSLCNWGQDLTFYWGSGIANSWRMSGDVTAEFTRPDSRCPCDGDEYDCKYAGFHCSIMNILNKAAPMGQNAGVGGWNDLDNLEVGVGNLTDDEEKAHFSMWAMVKSPLIIGANVNNLKASSYSIYSQASVIAINQDSNGIPATRVWRYYVSDTDEYGQGEIQMWSGPLDNGDQVVALLNGGSVSRPMNTTLEEIFFDSNLGSKKLTSTWDIYDLWANRVDNSTASAILGRNKTATGILYNATEQSYKDGLSKNDTRLFGQKIGSLSPNAILNTTVPAHGIAFYRLRPSSDYKDDDDK"
+        ],
+        "no_chains": 4,
+        "resolution": 2.7
+    },
+	    "6kwc": {
+        "release_date": "2021-01-27",
+        "chain_ids": [
+            "A"
+        ],
+        "seqs": [
+            "GSTIQPGTGYNNGYFYSYWNDGHGGVTYTNGPGGQFSVNWSNSGEFVGGKGWQPGTKNKVINFSGSYNPNGNSYLSVYGWSRNPLIEYYIVENFGTYNPSTGATKLGEVTSDGSVYDIYRTQRVNQPSIIGTATFYQYWSVRRNHRSSGSVNTANHFNAWAQQGLTLGTMDYQIVAVQGYFSSGSASITVS"
+        ],
+        "no_chains": 1,
+        "resolution": 1.297
+    },
+    ...
+}
+```
+
+
+### Duplicate pdb chain files 
+
+Duplicate chains occur across pdb entries. Some of these chains are the homomeric units of a multimer, others are subunits that are shared across different protein.
+
+To reduce storage overhead of creating / storing identical data for duplicate entries, we have a duplicate chain file. Each line stores the all chains that are identical. Our `6kwc` and `3lrm` examples would be stored as follows.
+
+```duplicate_pdb_chains.txt
+...
+6kwc_A
+3lrm_A 3lrm_B 3lrm_C 3lrm_D
+...
+```
+
+

docs/source/FAQ.md

@@ -0,0 +1,34 @@
+# FAQ
+
+Frequently asked questions or encountered issues when running OpenFold. 
+
+## Setup
+
+- When running unit tests (e.g. [`./scripts/run_unit_tests.sh`](https://github.com/aqlaboratory/openfold/blob/main/scripts/run_unit_tests.sh)), I see an error such as  
+	```
+	ImportError: version GLIBCXX_3.4.30 not found
+	```
+
+	> Solution: Make sure that the `$LD_LIBRARY_PATH` environment has been set to include the conda path, e.g. `export $LD_LIBRARY_PATH=$CONDA_PREFIX/lib:$LD_LIBRARY_PATH`
+
+- I see a CUDA mismatch error, eg. 
+```
+The detected CUDA version (11.8) mismatches the version that was used to compile
+PyTorch (12.1). Please make sure to use the same CUDA versions.
+```
+
+ > 	Solution: Ensure that your system's CUDA driver and toolkit match your intended OpenFold installation (CUDA 11 by default).  You can check the CUDA driver version with a command such as `nvidia-smi`
+
+- I get some error involving `fatal error: cuda_runtime.h: No such file or directory` and or `ninja: build stopped: subcommand failed.`. 
+
+> Solution: Something went wrong with setting up some of the custom kernels. Try running `install_third_party_dependencies.sh` again or try `python3 setup.py install` from inside the OpenFold folder. Make sure to prepend the conda environment as described above before running this.
+
+## Training
+
+- My model training is hanging on the data loading step:
+	 > Solution: While each system is different, a few general suggestions:
+		 - Check your `$KMP_AFFINITY` environment setting and see if it is suitable for your system.
+  		 - Adjust the number of data workers used to prepare data with the `--num_workers` setting. Increasing the number could help with dataset processing speed. However, to many workers could cause an OOM issue. 
+
+- When I reload my pretrained model weights or checkpoints, I get `RuntimeError: Error(s) in loading state_dict for OpenFoldWrapper: Unexpected key(s) in state_dict:`
+	> Solution: This suggests that your checkpoint / model weights are in OpenFold v1 format with outdated model layer names. Convert your weights/checkpoints following [this guide](convert_of_v1_weights.md).