Compare commits

...

4 Commits

Author SHA1 Message Date
Steven Palma f71a4b762a chore(config): add libs note 2026-07-24 13:17:42 +02:00
Steven Palma 0f174bd0cc chore(scripts): add multiprocessing_context safeguards 2026-07-24 13:14:00 +02:00
0o8o0-blip acbab791f1 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) <noreply@anthropic.com>
2026-07-24 12:21:07 +02:00
0o8o0-blip f8fc30bfda 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.
2026-07-24 12:20:41 +02:00
2 changed files with 30 additions and 4 deletions
+18
View File
@@ -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:
+12 -4
View File
@@ -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