Modernize conda environment

.gitattributes

@@ -0,0 +1,2 @@
+# SCM syntax highlighting & preventing 3-way merges
+pixi.lock merge=binary linguist-language=YAML linguist-generated=true

.gitignore

@@ -16,3 +16,10 @@ cutlass/
 
 # User-specific pre-commit settings
 .pre-commit-config.yaml
+
+# pixi environments
+.pixi/*
+!.pixi/config.toml
+
+# uv environments
+.venv/

docs/source/Installation.md

@@ -30,8 +30,34 @@ to install GPU accelerated {doc}`cuEquivariance attention kernels <kernels>`, us
 pip install openfold3[cuequivariance]
 ```
 
+### Modern conda environments with pixi
+
+OpenFold3 can now be installed in conda environments with [pixi](https://pixi.prefix.dev/latest/index.html).
+
+First [install pixi](https://pixi.prefix.dev/latest/installation/)
+
+```shell
+# Do this once and enjoy pixi for all your future projects!
+curl -fsSL https://pixi.sh/install.sh | sh
+# Then restart your shell and optionally install pixi completions
+```
+
+You can simply run openfold in one of the provided environments with pixi:
+```shell
+pixi run -e openfold3-cpu setup_openfold
+pixi run -e openfold3-cpu run_openfold
+```
+
+We provide the following environments:
+ - openfold3-cpu (linux-64, linux-aarch64, osx-64,osx-arm64)
+ - openfold3-cuda12 and openfold-cuda13 (linux-64, linux-aarch64) 
+
+For more information, including rationale, tips and tricks, see [Modern Conda Environments with Pixi](./modern-conda-environments-with-pixi.md).
+
 ### Environment variables
 
+> **Note:** This may need a revision given the pixi managed envs above (JD).
+
 OpenFold may need a few environment variables set so CUDA, compilation, and JIT-built extensions can be found correctly. 
 
 - `CUDA_HOME` should point to the CUDA installation. On many HPC clusters you will this can be set by loading the appropriate toolchain using environment modules, for example `module load cuda`.  If you do not set this you will likely get a `No such file or directory: '/usr/local/cuda/bin/nvcc'` error. 
@@ -54,7 +80,6 @@ OpenFold may need a few environment variables set so CUDA, compilation, and JIT-
     - Example: `export LIBRARY_PATH="$(echo "$CUDA_HOME" | sed 's|/cuda/|/math_libs/|')/targets/sbsa-linux/lib:${LIBRARY_PATH:-}"`
 
 
-
 ### OpenFold3 Docker Image
 
 #### Dockerhub

docs/source/modern-conda-environments-with-pixi.md

@@ -0,0 +1,3 @@
+# Modern OpenFold conda environments with pixi
+
+

examples/example_runner_yamls/osx.yaml

@@ -0,0 +1,14 @@
+model_update:
+  presets: 
+    - "predict"
+    - "pae_enabled"  # if using PAE enabled model
+    # - "low_mem"  # for lower memory systems
+  custom:
+    settings:
+      memory:
+        eval:
+          # otherwise the current default is wrongly to use deepspeed, even if not available
+          use_deepspeed_evo_attention: false
+
+data_module_args:
+  num_workers: 1

openfold3/__init__.py

@@ -19,8 +19,6 @@
 import gemmi
 from packaging import version
 
-from . import hacks  # noqa: F401
-
 if version.parse(gemmi.__version__) >= version.parse("0.7.3"):
     gemmi.set_leak_warnings(False)
 
@@ -32,3 +30,4 @@
     # This has weird effects with hanging if libaio is not installed and can
     # cause restart errors if run is preempted in the middle of autotuning
     deepspeed.HAS_TRITON = False
+    # FIXME: do we need this? it is really invasive with other potential users of DS

openfold3/core/data/framework/data_module.py

@@ -42,6 +42,9 @@
 import dataclasses
 import enum
 import logging
+import multiprocessing
+import platform
+import sys
 import warnings
 from functools import partial
 from typing import Any
@@ -148,14 +151,78 @@
         return self.get_subset(datasets_stage_mask)
 
 
+
 class DataModuleConfig(BaseModel):
     datasets: list[SerializeAsAny[BaseModel]]
     batch_size: int = 1
     num_workers: int = 0
     num_workers_validation: int = 0
+    multiprocessing_context: str = "openfold-default"
     data_seed: int = 42
     epoch_len: int = 1
 
+    @staticmethod
+    def safe_multiprocessing_context(
+            multiprocessing_context: str | None, num_workers: int
+    ) -> str | None:
+        """
+        Returns multiprocessing start methods with safer/sensible defaults:
+          - fork when using MPS
+          - forkserver for linux, matching the new 3.14 default
+          - default otherwise
+
+        For general info on risks and defaults across platformas and python versions see:
+          https://docs.pytorch.org/docs/stable/data.html#torch.utils.data.DataLoader
+          https://docs.pytorch.org/docs/stable/notes/multiprocessing.html#multiprocessing-poison-fork-note
+          https://docs.python.org/3/library/multiprocessing.html#contexts-and-start-methods
+        """
+
+        # Do not bother if not using multiprocessing
+        if num_workers == 0:
+            return None
+
+        # Set safe defaults
+        if multiprocessing_context == "openfold-default":
+
+            # Use fork to create processes when using MPS. See:
+            #  - https://github.com/pytorch/pytorch/issues/70344
+            #  - https://github.com/pytorch/pytorch/issues/87688
+            if platform.system() == "Darwin" and torch.backends.mps.is_available():
+                return "fork"
+
+            # Use forkserver in linux
+            # Backports the new python 3.14 default in previous python versions.
+            # An alternative for further safety would be "spawn". Avoid "fork".
+            # See: https://github.com/python/cpython/issues/84559
+            if platform.system() == "Linux":
+                return "forkserver"
+
+            # Use the platform default otherwise - "spawn" at the time of writing
+            return multiprocessing.get_start_method()
+
+        # Warn about unsafe defaults
+        else:
+            if platform.system() == "Darwin" and torch.backends.mps.is_available():
+                if multiprocessing_context != "fork":
+                    logger.warning(
+                        f"Using multiprocessing context {multiprocessing_context} on MPS may cause "
+                        "issues. Consider using 'fork' or 'openfold-default' (which resolves to 'fork' on MPS).",
+                        stacklevel=2,
+                    )
+            if platform.system() == "Linux":
+                dangerous_start_method = (
+                    multiprocessing_context == "fork" or
+                    multiprocessing_context is None and sys.version_info < (3, 14)
+                )
+                if dangerous_start_method:
+                    logger.warning(
+                        "Using 'fork' multiprocessing context in linux may cause issues. Consider using "
+                        "'spawn', 'forkserver' or 'openfold-default' (which resolves to 'forkserver' on linux).",
+                        stacklevel=2,
+                    )
+
+        return multiprocessing_context
+
 
 class DataModule(pl.LightningDataModule):
     """A LightningDataModule class for organizing Datasets and DataLoaders."""
@@ -167,6 +234,7 @@
         self.batch_size = data_module_config.batch_size
         self.num_workers = data_module_config.num_workers
         self.num_workers_validation = data_module_config.num_workers_validation
+        self.multiprocessing_context = data_module_config.multiprocessing_context
         self.data_seed = data_module_config.data_seed
         self.next_data_seed = data_module_config.data_seed
         self.epoch_len = data_module_config.epoch_len
@@ -433,8 +501,17 @@
         # instead of pl.seed_everything(workers=True), so this function is
         # passed explicitly here.
         worker_init_fn = partial(pl_worker_init_function, rank=self.global_rank)
+
+        # Set a sensible default for multiprocesssing start method
+        # depending on platform and python version.
+        multiprocessing_context = DataModuleConfig.safe_multiprocessing_context(
+            self.multiprocessing_context, num_workers
+        )
+
         logger.debug(
-            f"Creating {mode} dataloader: num_workers={num_workers}, "
+            f"Creating {mode} dataloader: "
+            f"num_workers={num_workers}, "
+            f"multiprocessing_context={multiprocessing_context}, "
             f"rank={self.global_rank}."
         )
         return DataLoader(
@@ -445,6 +522,7 @@
             collate_fn=openfold_batch_collator,
             generator=self.generators[mode],
             worker_init_fn=worker_init_fn,
+            multiprocessing_context=multiprocessing_context,
         )
 
     def train_dataloader(self) -> DataLoader:

openfold3/core/data/pipelines/featurization/conformer.py

@@ -108,6 +108,10 @@
             coords = conf.GetAtomPosition(atom.GetIdx())
             mol_ref_mask.append(int(atom.GetBoolProp("annot_used_atom_mask")))
             mol_ref_pos.append(coords)
+            # Some PyPI installations crash here due to ABI mismatch between RDKit and PyTorch
+            # Leaving a quick fix commented (beware, moving into slow python land)
+            # Remove if nobody else hits the problem
+            # mol_ref_pos.append([coords.x, coords.y, coords.z])
 
             # Atom elements (0-indexed)
             element_symbol = atom.GetSymbol()

openfold3/core/model/primitives/attention.py

@@ -507,11 +507,6 @@ def _deepspeed_evo_attn(
         biases:
             List of biases that broadcast to [*, H, Q, K]
     """
-    from openfold3 import hacks
-
-    hacks.prep_deepspeed()
-    hacks.prep_cutlass()
-
     if not ds4s_is_installed:
         raise ValueError(
             "_deepspeed_evo_attn requires that DeepSpeed be installed "

openfold3/tests/__init__.py

@@ -14,12 +14,11 @@
 
 import importlib
 
-from openfold3 import hacks  # noqa: F401
-
 if importlib.util.find_spec("deepspeed") is not None:
     import deepspeed
 
     # TODO: Resolve this
     # This is a hack to prevent deepspeed from doing the triton matmul autotuning
     # I'm not sure why it's doing this by default, but it's causing the tests to hang
     deepspeed.HAS_TRITON = False
+    # FIXME: potentially too invasive with other libraries, see also comments up about libaio

openfold3/tests/scripts/data_preprocessing/test_preparse_alignments_of3.py

@@ -62,7 +62,7 @@ def test_preparse_databases(self, cli_runner, tmp_path):
 
         # Check that npz files were created for both chains
         npz_files = list(tmp_path.glob("*.npz"))
-        assert [f.name for f in npz_files] == ["2q2k_B.npz", "2q2k_A.npz"]
+        assert set([f.name for f in npz_files]) == {"2q2k_B.npz", "2q2k_A.npz"}
 
         # Check contents of one npz file
         npz_data = np.load(tmp_path / "2q2k_A.npz", allow_pickle=True)

openfold3/tests/test_lmdb.py

@@ -66,7 +66,7 @@ def test_lmdb_roundtrip(self, tmp_path):
 
         # Create LMDB
         test_lmdb_dir = tmp_path / "test_lmdb"
-        map_size = 20 * 1024
+        map_size = 200 * 1024
         convert_datacache_to_lmdb(test_config_json, test_lmdb_dir, map_size)
 
         # read lmdb