From a36778d4f81cd0af62004f7e638691899f077cc5 Mon Sep 17 00:00:00 2001 From: Jesse Rusak Date: Wed, 17 Jun 2026 14:42:39 -0400 Subject: [PATCH] Always use spawn for data loader workers --- src/samudra/config.py | 30 +++++++++--------------------- src/samudra/datasets.py | 4 ++-- src/samudra/train.py | 5 +---- src/samudra/utils/data.py | 1 - src/samudra/utils/location.py | 10 ---------- tests/test_trainer.py | 3 +++ 6 files changed, 15 insertions(+), 38 deletions(-) diff --git a/src/samudra/config.py b/src/samudra/config.py index 4f5c1b27..8beae571 100644 --- a/src/samudra/config.py +++ b/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, ) diff --git a/src/samudra/datasets.py b/src/samudra/datasets.py index 3ffeb6c3..c838feb9 100644 --- a/src/samudra/datasets.py +++ b/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 diff --git a/src/samudra/train.py b/src/samudra/train.py index 33c0f2d0..8156164f 100644 --- a/src/samudra/train.py +++ b/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) diff --git a/src/samudra/utils/data.py b/src/samudra/utils/data.py index 7a869a2a..9290dbe1 100644 --- a/src/samudra/utils/data.py +++ b/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. diff --git a/src/samudra/utils/location.py b/src/samudra/utils/location.py index 3b3b5caf..11d681f8 100644 --- a/src/samudra/utils/location.py +++ b/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) diff --git a/tests/test_trainer.py b/tests/test_trainer.py index 1e22884f..8e6e66c7 100644 --- a/tests/test_trainer.py +++ b/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