[Bug] Deadlock in compute_standardization_stats due to Dask + PyTorch DataLoader

[Sharkyii]

Description

The compute_standardization_stats.py script hangs indefinitely when run with default settings on Linux. The script gets stuck at 0% progress during the data loading phase and never completes.

Steps to Reproduce

1. Install neural-lam from main branch on Linux
2. Download and extract the MEPS example dataset
3. Place the datastore config in the correct location:

   cp data/meps.datastore.yaml

4. Run the standardization stats computation:

   python -m neural_lam.datastore.npyfilesmeps.compute_standardization_stats \
       --datastore_config_path data/meps_example/meps.datastore.yaml

Expected Behavior

The script should compute mean and standard deviation statistics for the dataset and complete within a few minutes for the small example dataset.

Actual Behavior

The script hangs indefinitely at:

Computing mean and std.-dev. for parameters...
0%|                                                              | 0/1 [00:00<?, ?it/s]

• No progress is made
• CPU usage drops to near zero
• Process must be killed manually (Ctrl+C)
• Tested for over 1 hour with no completion

Root Cause

This is a fork-safety deadlock specific to Linux, caused by PyTorch DataLoader using fork() after Dask has initialized thread pools.

It is NOT caused by Dask spawning processes inside workers.

What Actually Happens:

1. The main process initializes Dask arrays
2. Dask creates internal thread pools
3. PyTorch DataLoader with num_workers > 0 uses fork() on Linux (default multiprocessing start method)
4. fork() copies the entire memory space and lock states, but does NOT copy running threads
5. Worker processes inherit locks that appear "held" by threads that don't exist in the child
6. When Dask tries to acquire those locks in the worker, it blocks forever waiting for non-existent threads
7. Result: Workers freeze, progress bar stays at 0%

Why This Only Happens on Linux:

Platform	Default Start Method	Behavior
Linux	fork	💀 Deadlock (inherits broken locks)
Windows	spawn	✅ Works (fresh interpreter per worker)
macOS	spawn	✅ Works (fresh interpreter per worker)

Environment

• OS: Linux (Ubuntu)
• Python: 3.11

[Sharkyii]

Script that hangs in Linux

import torch
from torch.utils.data import Dataset, DataLoader
import dask.array as da
import dask
import time

x = da.random.random((5000, 5000), chunks=(1000, 1000))
print("Running pre-fork compute...")
x.mean().compute()  # this initializes threadpools
print("Pre-fork compute done.")

class MyDataset(Dataset):
    def __len__(self):
        return 10

    def __getitem__(self, idx):
        # This compute runs inside DataLoader worker
        return x.std().compute()

print("Creating DataLoader...")
loader = DataLoader(MyDataset(), num_workers=4)
print("Starting iteration...")

for batch in loader:
    print("Got batch")

[Sharkyii]

What could be a better approach for this?
Proposed Solutions

1. Reinitialize the Dask thread pool inside each DataLoader worker
   This could prevent inherited thread state conflicts across forked processes. (Still thinking how to impliment...)
2. Force the multiprocessing start method to spawn
   This may avoid issues caused by fork inheriting active threads.
3. Set num_workers=0 by default
   @joeloskarsson @leifdenby @khintz @j6k4m8

[observingClouds]

Hi @Sharkyii 👋,

Thank you for raising this issue and doing some in-depth analysis. While this is indeed not ideal, the MEPS data-loader is very specific to the MEPS dataset. A more general dataset loader is the MDP dataloader. MDP dataset can be created for example with mllam-data-prep. If you encounter this issue there as well, it would definitely good to fix it. While a more robust MEPS loader would be great, a complicated fix would not worth the effort as it is rather deprecated.

[Sharkyii]

Hello @observingClouds
I checked it and there is no Deadlock situation in MDP dataloader

Why MDP is safe:

• No PyTorch DataLoader usage
• No fork-based multiprocessing (multiprocessing.Pool, ProcessPoolExecutor, etc.)
• Only uses Dask, which handles process spawning correctly via LocalCluster

for the MEPS dataset a easy solution would be adding this (Tested ✅)

import torch.multiprocessing as mp
mp.set_start_method("spawn", force=True)

if this is fine i could add a PR and close this issue

[observingClouds]

I ran your example but could not reproduce the deadlock.

I ran this on a Linux system with torch 2.9.1+cu126 and dask 2025.9.1. What versions are you running?

[Sharkyii]

2.10.0+cu128 torch
2026.1.2 dask
i downgraded dask still didn't worked for me

[Sharkyii]

@observingClouds i tried running on collab but still it failed!

[Sharkyii]

@observingClouds just wanted to confirm are you able to see Got batch in the output?

[observingClouds]

Yes, I will checkout your versions next.

In [1]: import torch
   ...: from torch.utils.data import Dataset, DataLoader
   ...: import dask.array as da
   ...: import dask
   ...: import time
   ...: 
   ...: x = da.random.random((5000, 5000), chunks=(1000, 1000))
   ...: print("Running pre-fork compute...")
   ...: x.mean().compute()  # this initializes threadpools
   ...: print("Pre-fork compute done.")
   ...: 
   ...: class MyDataset(Dataset):
   ...:     def __len__(self):
   ...:         return 10
   ...: 
   ...:     def __getitem__(self, idx):
   ...:         # This compute runs inside DataLoader worker
   ...:         return x.std().compute()
   ...: 
   ...: print("Creating DataLoader...")
   ...: loader = DataLoader(MyDataset(), num_workers=4)
   ...: print("Starting iteration...")
   ...: 
   ...: for batch in loader:
   ...:     print("Got batch")
   ...: 
Running pre-fork compute...
Pre-fork compute done.
Creating DataLoader...
Starting iteration...
Got batch
Got batch
Got batch
Got batch
Got batch
Got batch
Got batch
Got batch
Got batch
Got batch

[Shalinis19137]

I would like to work on this issue. Could you please assign it to me?

[sadamov]

Closing as the test/CI/packaging stack has been substantially modernised since this was reported: uv with CPU/GPU extras (#604), torch wheel index resolution (#639), Ruff and pre-commit modernisation (#614), CI dependency caching (#606), plus xarray/zarr have moved on too. If the deadlock still reproduces on current main with the updated stack, please reopen with a fresh reproducer.