Always use spawn for data loader workers

src/samudra/config.py

@@ -268,11 +268,11 @@ def make_source(
             means_location: Location,
             stds_location: Location,
             turn_on_dask: bool = use_dask,
-        ) -> tuple[DataSource, bool]:
+        ) -> DataSource:
             resolved_data_location = data_root.resolve(data_location)
             resolved_means_location = data_root.resolve(means_location)
             resolved_stds_location = data_root.resolve(stds_location)
-            data_source = DataSource.from_locations(
+            return DataSource.from_locations(
                 data_location=resolved_data_location,
                 means_location=resolved_means_location,
                 stds_location=resolved_stds_location,
@@ -283,26 +283,15 @@ def make_source(
                 use_dask=turn_on_dask,
             )
 
-            return data_source, all(
-                loc.supports_fork
-                for loc in [
-                    resolved_data_location,
-                    resolved_means_location,
-                    resolved_stds_location,
-                ]
-            )
-
         sources = []
-        supports_forks = []
         for source_cfg in self.sources:
-            src, fork = make_source(
-                source_cfg.data_location,
-                source_cfg.data_means_location,
-                source_cfg.data_stds_location,
+            sources.append(
+                make_source(
+                    source_cfg.data_location,
+                    source_cfg.data_means_location,
+                    source_cfg.data_stds_location,
+                )
             )
-            sources.append(src)
-            supports_forks.append(fork)
-        supports_fork = all(supports_forks)
 
         primary_source = sources[0]
         if use_dask:
@@ -311,7 +300,7 @@ def make_source(
         else:
             # If we're not using dask for the main source, create a separate one
             primary = self.sources[0]
-            inference_source, _ = make_source(
+            inference_source = make_source(
                 primary.data_location,
                 primary.data_means_location,
                 primary.data_stds_location,
@@ -328,7 +317,6 @@ def make_source(
             sources=sources,
             inference_source=inference_source,
             loader_version=loader_version,
-            supports_fork=supports_fork,
             dataset_spec=dataset_spec,
             static_data=static_data,
         )

src/samudra/datasets.py

@@ -443,8 +443,8 @@ class TorchTrainDataset(Dataset[RawTrainData]):
     FLAG = LoaderVersion.OM4_TORCH
 
     # Shared across all instances within a process. Created lazily on first
-    # __getitem__ call so that each forked DataLoader worker gets its own
-    # clean executor — avoids inheriting fork-corrupted locks from the parent.
+    # __getitem__ call so that each DataLoader worker gets its own clean
+    # executor.
     _shared_executor: ClassVar[ThreadPoolExecutor | None] = None
 
     @classmethod

src/samudra/train.py

@@ -165,10 +165,7 @@ def __init__(self, cfg: TrainConfig) -> None:
 
         self.mp_context: BaseContext | None = None
         if data_num_workers > 0:
-            if self.data_container.supports_fork:
-                self.mp_context = multiprocessing.get_context("fork")
-            else:
-                self.mp_context = multiprocessing.get_context("spawn")
+            self.mp_context = multiprocessing.get_context("spawn")
 
         self.num_prog_in = int((cfg.data.hist + 1) * self.N_prog)
         self.num_boundary_in = int((cfg.data.hist + 1) * self.N_bound)

src/samudra/utils/data.py

@@ -412,7 +412,6 @@ class DataContainer:
     sources: list[DataSource]
     inference_source: DataSource
     loader_version: LoaderVersion
-    supports_fork: bool
     dataset_spec: DatasetSpec
     # TODO(559): static_data should belong to the DataSource, since we now
     #  deal with multiple resolutions.

src/samudra/utils/location.py

@@ -52,10 +52,6 @@ def open(self, chunks: dict[str, int] | None = None) -> xr.Dataset:
     def resolve(self, location: "Location") -> "ResolvedLocation":
         pass
 
-    @abstractmethod
-    def supports_fork(self) -> bool:
-        pass
-
     def __truediv__(self, other: "Location") -> "ResolvedLocation":
         return self.resolve(other)
 
@@ -102,9 +98,6 @@ def resolve(self, location: "Location") -> "ResolvedLocation":
             )
         return location
 
-    def supports_fork(self) -> bool:
-        return False  # s3fs does not support forking
-
     def __str__(self) -> str:
         return self.url()
 
@@ -143,9 +136,6 @@ def resolve(self, location: "Location") -> "ResolvedLocation":
             return LocalLocation(path=self.path / location.path)
         return location
 
-    def supports_fork(self) -> bool:
-        return True
-
     def __str__(self) -> str:
         return str(self.path)
 

tests/test_trainer.py

@@ -199,6 +199,8 @@ def test_data_loaders_enable_persistent_workers_on_positive_num_workers(
 ):
     _, trainer = trainer_pair
 
+    assert trainer.mp_context is not None
+    assert trainer.mp_context.get_start_method() == "spawn"
     assert trainer.train_loader._dataloader.persistent_workers is True
     assert trainer.val_loader._dataloader.persistent_workers is True
 
@@ -220,5 +222,6 @@ def test_data_loaders_disable_persistent_workers_when_num_workers_is_zero(
         trainer = Trainer(train_config)
         trainer.init_data_loaders(cur_step=train_config.steps[0])
 
+    assert trainer.mp_context is None
     assert trainer.train_loader._dataloader.persistent_workers is False
     assert trainer.val_loader._dataloader.persistent_workers is False