Better defaults and logging for multiprocessing context

We would still need to document this somewhere

Author: sdvillal

openfold3/core/data/framework/data_module.py

@@ -44,6 +44,7 @@
 import logging
 import multiprocessing
 import platform
+import sys
 import warnings
 from functools import partial
 from typing import Any
@@ -150,15 +151,77 @@ def get_config_for_mode(self, mode: DatasetMode) -> "MultiDatasetConfig":
         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 = None
+    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:
+
+            # 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."""
@@ -170,7 +233,7 @@ def __init__(self, data_module_config: DataModuleConfig) -> None:
         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.multiprocessing_context = data_module_config.safe_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
@@ -438,22 +501,11 @@ def generate_dataloader(self, mode: DatasetMode, sampler: Sampler | None = None)
         # passed explicitly here.
         worker_init_fn = partial(pl_worker_init_function, rank=self.global_rank)
 
-        # Configure multiprocessing_context with sensible defaults
-        # For general info on risks 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
-        multiprocessing_context = self.multiprocessing_context
-        # Use known/safer working defaults
-        if multiprocessing_context is None and num_workers > 0:
-            # 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():
-                multiprocessing_context = "fork"
-            # Use spawn by default in aarch64 as it is the safer bet (we observed failures with default)
-            elif platform.system() == "linux" and platform.machine() == "aarch64":
-                multiprocessing_context = "spawn"
+        # 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: "