From 95211b98f1cd6b638bda84a8d28f9e41323229dd Mon Sep 17 00:00:00 2001 From: Steven Palma Date: Tue, 28 Jul 2026 00:42:55 +0200 Subject: [PATCH] feat(config): add `multiprocessing` option to `DataLoader` context and sets `spawn` as default (#4139) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * Add dataloader_multiprocessing_context, default to spawn Make the DataLoader multiprocessing start method configurable on TrainPipelineConfig and default it to 'spawn'. The previous default (fork on Linux) is unsafe with libraries that hold non-fork-safe state in the parent process — common ones in this codebase are PyAV, torchcodec, and the ffmpeg shared libs they wrap. Symptoms reported in #2488, #2209, and observed locally include: - multiprocessing.context.AuthenticationError: digest received was wrong - RuntimeError: Pin memory thread exited unexpectedly - RuntimeError: DataLoader worker exited unexpectedly - Random SIGSEGV inside worker processes during video decode Switching to spawn re-imports modules cleanly in each worker and eliminates these failure modes. Added the setting as a config field rather than hard-coding so users on platforms where fork is preferred can opt back in via --dataloader-multiprocessing-context=fork. * Address review: shorten config comment, note spawn startup tradeoff Per @jashshah999, mention that spawn workers re-import modules and so add some startup time vs fork. Also trim the failure-mode dump from the inline comment — the linked issue covers the symptoms in detail. Co-Authored-By: Claude Opus 4.7 (1M context) * chore(scripts): add multiprocessing_context safeguards * chore(config): add libs note --------- Co-authored-by: 0o8o0-blip <0o8o0-blip@users.noreply.github.com> --- src/lerobot/configs/train.py | 18 ++++++++++++++++++ src/lerobot/scripts/lerobot_train.py | 16 ++++++++++++---- 2 files changed, 30 insertions(+), 4 deletions(-) diff --git a/src/lerobot/configs/train.py b/src/lerobot/configs/train.py index e3d354691..e92247188 100644 --- a/src/lerobot/configs/train.py +++ b/src/lerobot/configs/train.py @@ -14,6 +14,7 @@ import builtins import datetime as dt import json +import multiprocessing import os import tempfile from dataclasses import dataclass, field @@ -101,6 +102,12 @@ class TrainPipelineConfig(HubMixin): batch_size: int = 8 prefetch_factor: int = 4 persistent_workers: bool = True + # DataLoader worker start method. "spawn" is safer than "fork" with + # non-fork-safe libs (PyAV / torchcodec / ffmpeg), but adds some + # worker-startup time per run since workers re-import modules instead + # of inheriting parent state. Override with `--dataloader_multiprocessing_context=fork` + # when appropriate, or set it to `null` to use Python's platform default. + dataloader_multiprocessing_context: str | None = "spawn" steps: int = 100_000 # Run policy in the simulation environment every N steps to measure reward/success (0 = disabled). env_eval_freq: int = 20_000 @@ -212,6 +219,17 @@ class TrainPipelineConfig(HubMixin): self.reward_model.pretrained_path = str(policy_dir) def validate(self) -> None: + available_contexts = multiprocessing.get_all_start_methods() + if ( + self.dataloader_multiprocessing_context is not None + and self.dataloader_multiprocessing_context not in available_contexts + ): + raise ValueError( + "`dataloader_multiprocessing_context` must be None or one of " + f"{available_contexts} on this platform, got " + f"{self.dataloader_multiprocessing_context!r}." + ) + self._resolve_pretrained_from_cli() if self.policy is None and self.reward_model is None: diff --git a/src/lerobot/scripts/lerobot_train.py b/src/lerobot/scripts/lerobot_train.py index ec5565cf4..8bfa16a98 100644 --- a/src/lerobot/scripts/lerobot_train.py +++ b/src/lerobot/scripts/lerobot_train.py @@ -71,6 +71,16 @@ from lerobot.utils.utils import ( from .lerobot_eval import eval_policy_all +def _dataloader_worker_kwargs(cfg: TrainPipelineConfig) -> dict[str, Any]: + """Return worker-only DataLoader options, disabling them for single-process loading.""" + workers_enabled = cfg.num_workers > 0 + return { + "prefetch_factor": cfg.prefetch_factor if workers_enabled else None, + "persistent_workers": cfg.persistent_workers and workers_enabled, + "multiprocessing_context": cfg.dataloader_multiprocessing_context if workers_enabled else None, + } + + def update_policy( train_metrics: MetricsTracker, policy: PreTrainedPolicy, @@ -473,8 +483,7 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None): pin_memory=device.type == "cuda", drop_last=False, collate_fn=collate_fn, - prefetch_factor=cfg.prefetch_factor if cfg.num_workers > 0 else None, - persistent_workers=cfg.persistent_workers and cfg.num_workers > 0, + **_dataloader_worker_kwargs(cfg), ) # Build eval dataloader if a held-out split exists @@ -500,8 +509,7 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None): pin_memory=device.type == "cuda", drop_last=False, collate_fn=eval_collate_fn, - prefetch_factor=cfg.prefetch_factor if cfg.num_workers > 0 else None, - persistent_workers=cfg.persistent_workers and cfg.num_workers > 0, + **_dataloader_worker_kwargs(cfg), ) # Prepare everything with accelerator