Best practices for using multi-worker DataLoader (num_workers > 0) inside a ParallelEnv

[tlt18]

Hi community,

I'm working with a custom environment where each reset() call involves loading a significant amount of data from disk. To optimize this, I'd like to use a standard torch.utils.data.DataLoader with num_workers > 0 inside my environment's reset logic.

However, I've encountered a fundamental conflict when using this setup with torchrl.envs.ParallelEnv. I'm opening this discussion to seek advice on the recommended patterns for this scenario and to see if others have similar needs.

The Use Case: Data-loading on reset()
Our environment simulates scenarios based on real-world data logs. Each episode requires loading a different data file. The simplified logic looks like this:

Python

import torch
from torch.utils.data import Dataset, DataLoader
from torchrl.envs import EnvBase

class MyHeavyDataEnv(EnvBase):
    def __init__(self, file_paths, ...):
        super().__init__(...)
        self.file_paths = file_paths
        # This DataLoader will be re-initialized on each reset
        self.loader = None

    def _reset(self, tensordict):
        # Sample a new data file for the new episode
        file_path = random.choice(self.file_paths)
        
        # Create a dataset and a DataLoader to load the data for this episode
        dataset = MyDataset(file_path)
        
        # The ideal scenario is to parallelize this data loading
        self.loader = DataLoader(
            dataset,
            batch_size=None, # Loading the whole file
            num_workers=4   # <--- The desired optimization
        )
        
        # Consuming the loader to get the data
        episode_data = next(iter(self.loader))
        
        # ... use episode_data to set the initial state ...
        return initial_state_tensordict

When num_workers=0, this works perfectly but can be a performance bottleneck, as data loading and preprocessing are slow.

The Technical Challenge

The core issue is that ParallelEnv creates its worker processes as daemonic, while a DataLoader with num_workers > 0 attempts to spawn its own child processes. This leads to Python's well-known restriction:
AssertionError: daemonic processes are not allowed to have children

This creates a conflict between torchrl's process management for parallel environments and PyTorch's standard for parallel data loading.

What I've Considered

Approach 1: num_workers = 0 (The Safe-but-Slow Method)

This is the most straightforward workaround. It avoids the error but sacrifices performance at a critical point in our training loop. Given that we have many parallel environments, the total throughput is okay, but each reset is individually slow.

Approach 2: Custom Non-Daemonic ParallelEnv (The Powerful-but-Risky Method)

I've considered subclassing ParallelEnv and forcing its worker processes to be non-daemonic. This would permit them to have their own children (DataLoader workers).

However, the significant drawback is the need for manual and careful process lifecycle management to avoid creating zombie processes, which adds a lot of complexity and risk to the training framework.

Discussion Points & Questions
I'd love to hear the community's thoughts on this:

Is this a common scenario? Do other users have environments that require heavy, parallelizable I/O or computation during reset()?

Are there more elegant, recommended patterns to solve this? Perhaps there's a different way to structure the data loading pipeline with torchrl that I've missed.

Would the TorchRL team consider adding a feature to facilitate this? For example, a built-in option in ParallelEnv like daemon_workers=False could be a solution. This would, of course, come with a clear warning in the documentation about the user's responsibility to properly close the environment.

This seems like a potentially valuable capability for a broader set of use cases where environments are tightly coupled with large datasets.

Thank you for your time and for building such a great library! I look forward to the discussion.

[RenSolvyn]

Short version: I would not go down the non-daemonic road. Even if you get it working, it mostly will not buy you the speedup you are after, and there are two cleaner ways out.

Why the error exists. The AssertionError is not TorchRL policy, it is Python itself: the multiprocessing docs state that a daemonic process is not allowed to create child processes, precisely so nobody has to reap orphaned grandchildren when a worker dies. Subclassing ParallelEnv to make workers non-daemonic means you inherit that reaping job, which is the risk you already identified.

Why it would not help anyway. With N parallel envs each running num_workers=4, you have 4N processes queued against one storage device. Disk bandwidth is fixed, so past a handful of concurrent readers you are adding seek contention and IPC overhead, not throughput. Your ParallelEnv workers are already the parallelism. That is why num_workers=0 feels slow per reset but overall throughput is okay: you are likely near the I/O ceiling already.

Option 1, the real fix: preprocess your logs into memory-mapped TensorDicts once, offline.

from tensordict import TensorDict

# one-time conversion per log file
td = TensorDict(parse_log(path), batch_size=[n_steps])
td.memmap_(f"/data/prepped/{stem}")

Then reset stops being I/O:

def _reset(self, tensordict):
    stem = random.choice(self.stems)
    ep = TensorDict.load_memmap(f"/data/prepped/{stem}")
    return self._to_obs(ep)

load_memmap maps pages instead of reading bytes, so the cost moves to demand paging during the rollout, where the kernel overlaps it with your step compute. Quick synthetic check on my machine (400 MB float64 array, Linux, warm page cache, so the cold-disk gap is even larger): eager np.load took ~31 ms, opening the same array memory-mapped took ~1.5 ms, and pulling a random 1 MB slice through the map took ~0.2 ms. The reset-time cost basically disappears. The pattern is documented in the tensordict memory tutorial, and it is the same mechanism TorchRL already trusts internally: LazyMemmapStorage backs replay buffers with memmapped tensordicts for exactly this reason.

Option 2, if you must parse raw files at reset: prefetch on a thread. The daemon restriction applies to child processes, not threads, and CPython releases the GIL during blocking file reads, so a background thread genuinely overlaps I/O with your episode compute:

from concurrent.futures import ThreadPoolExecutor

def __init__(self, ...):
    self._pool = ThreadPoolExecutor(max_workers=1)
    self._next = self._pool.submit(self._load, random.choice(self.file_paths))

def _reset(self, tensordict):
    ep = self._next.result()
    self._next = self._pool.submit(self._load, random.choice(self.file_paths))
    return self._to_obs(ep)

You know at reset that another file will be needed, so you start loading it immediately and hide the latency behind the episode you are about to run.

One caveat on option 1: each ParallelEnv worker holds its own mapping, so if your prepped files are large relative to RAM and N is big, you can pressure the page cache. If you see that, cap the number of parallel envs before reaching back for DataLoader workers.

Happy to share the full conversion snippet if useful.

[tlt18]

Thanks so much for your great explanation! I ended up using a producer-consumer preloading scheme. It preloads data in the background, avoids the daemon process issue entirely, and resolves my reset I/O bottleneck.