LightningCLI fails with num_workers > 0 and custom dataset

[adosar]

Hello,

I have a custom dataset that looks like this:

class CustomDataset(Dataset):
    def __init__(
            self, pcd_names, path_to_X,
            path_to_Y=None, index_col=None, labels=None,
            transform_x=None, transform_y=None,
            ):

        if (labels is not None) and (type(labels) != list):
            raise ValueError('labels must be a list!')

        self._pcd_names = pcd_names
        self.path_to_X = path_to_X
        self.path_to_Y = path_to_Y
        self.labels = labels
        self.index_col = index_col
        self.transform_x = transform_x
        self.transform_y = transform_y

        self.X = None
        self.Y = None

    @property
    def pcd_names(self):
        return self._pcd_names

    def __len__(self):
        return len(self.pcd_names)

    def __getitem__(self, idx):
        # Account for np.load and multiprocessing.
        if self.X is None:
            self.X = np.load(self.path_to_X, mmap_mode='r')
        if self.Y is None and self.path_to_Y is not None:
            self.Y = pd.read_csv(self.path_to_Y, index_col=self.index_col)[self.labels]

        name = self.pcd_names[idx]
        sample_x = self.X[name]

        if self.transform_x is not None:
            sample_x = self.transform_x(sample_x)

        # Only for labeled datasets.
        if self.Y is not None:
            sample_y = self.Y.loc[name].values

            if self.transform_y is not None:
                sample_y = self.transform_y(sample_y)

            return (
                    torch.tensor(sample_x, dtype=torch.float),
                    torch.tensor(sample_y, dtype=torch.float)
                    )

        return torch.tensor(sample_x, dtype=torch.float)

Everything works fine even with num_workers > 0 when I train my model with a custom script. The problem appears when I train the model with the LightningCLI. I am getting the following error:

Seed set to 1
GPU available: True (cuda), used: True
TPU available: False, using: 0 TPU cores
IPU available: False, using: 0 IPUs
HPU available: False, using: 0 HPUs
LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0,1]

  | Name  | Type     | Params
-----------------------------------
0 | model | MyNet | 2.7 M 
-----------------------------------
2.7 M     Trainable params
0         Non-trainable params
2.7 M     Total params
10.652    Total estimated model params size (MB)
Epoch 0:   0%|                                                                                                                                                       | 0/1000 [00:00<?, ?it/s]terminate called after throwing an instance of 'c10::Error'
  what():  CUDA error: initialization error
CUDA kernel errors might be asynchronously reported at some other API call, so the stacktrace below might be incorrect.
For debugging consider passing CUDA_LAUNCH_BLOCKING=1.
Compile with `TORCH_USE_CUDA_DSA` to enable device-side assertions.

Exception raised from c10_cuda_check_implementation at ../c10/cuda/CUDAException.cpp:44 (most recent call first):
frame #0: c10::Error::Error(c10::SourceLocation, std::string) + 0x57 (0x7fdc7e4ced87 in /home/ansar/venvir/ml/lib64/python3.11/site-packages/torch/lib/libc10.so)
frame #1: c10::detail::torchCheckFail(char const*, char const*, unsigned int, std::string const&) + 0x64 (0x7fdc7e47f75f in /home/ansar/venvir/ml/lib64/python3.11/site-packages/torch/lib/libc10.so)
frame #2: c10::cuda::c10_cuda_check_implementation(int, char const*, char const*, int, bool) + 0x118 (0x7fdc7e59f8a8 in /home/ansar/venvir/ml/lib64/python3.11/site-packages/torch/lib/libc10_cuda.so)
frame #3: c10::cuda::ExchangeDevice(int) + 0x8a (0x7fdc7e59fd0a in /home/ansar/venvir/ml/lib64/python3.11/site-packages/torch/lib/libc10_cuda.so)
frame #4: <unknown function> + 0xfa560c (0x7fdc337a560c in /home/ansar/venvir/ml/lib64/python3.11/site-packages/torch/lib/libtorch_cuda.so)
frame #5: <unknown function> + 0x543010 (0x7fdc7cb43010 in /home/ansar/venvir/ml/lib64/python3.11/site-packages/torch/lib/libtorch_python.so)
frame #6: <unknown function> + 0x649bf (0x7fdc7e4b39bf in /home/ansar/venvir/ml/lib64/python3.11/site-packages/torch/lib/libc10.so)
frame #7: c10::TensorImpl::~TensorImpl() + 0x21b (0x7fdc7e4acc8b in /home/ansar/venvir/ml/lib64/python3.11/site-packages/torch/lib/libc10.so)
frame #8: c10::TensorImpl::~TensorImpl() + 0x9 (0x7fdc7e4ace39 in /home/ansar/venvir/ml/lib64/python3.11/site-packages/torch/lib/libc10.so)
frame #9: <unknown function> + 0x80b718 (0x7fdc7ce0b718 in /home/ansar/venvir/ml/lib64/python3.11/site-packages/torch/lib/libtorch_python.so)
frame #10: THPVariable_subclass_dealloc(_object*) + 0x2f6 (0x7fdc7ce0ba96 in /home/ansar/venvir/ml/lib64/python3.11/site-packages/torch/lib/libtorch_python.so)
<omitting python frames>

terminate called after throwing an instance of 'c10::Error'
  what():  CUDA error: initialization error
CUDA kernel errors might be asynchronously reported at some other API call, so the stacktrace below might be incorrect.
For debugging consider passing CUDA_LAUNCH_BLOCKING=1.
Compile with `TORCH_USE_CUDA_DSA` to enable device-side assertions.

Exception raised from c10_cuda_check_implementation at ../c10/cuda/CUDAException.cpp:44 (most recent call first):
frame #0: c10::Error::Error(c10::SourceLocation, std::string) + 0x57 (0x7fdc7e4ced87 in /home/ansar/venvir/ml/lib64/python3.11/site-packages/torch/lib/libc10.so)
frame #1: c10::detail::torchCheckFail(char const*, char const*, unsigned int, std::string const&) + 0x64 (0x7fdc7e47f75f in /home/ansar/venvir/ml/lib64/python3.11/site-packages/torch/lib/libc10.so)
frame #2: c10::cuda::c10_cuda_check_implementation(int, char const*, char const*, int, bool) + 0x118 (0x7fdc7e59f8a8 in /home/ansar/venvir/ml/lib64/python3.11/site-packages/torch/lib/libc10_cuda.so)
frame #3: c10::cuda::ExchangeDevice(int) + 0x8a (0x7fdc7e59fd0a in /home/ansar/venvir/ml/lib64/python3.11/site-packages/torch/lib/libc10_cuda.so)
frame #4: <unknown function> + 0xfa560c (0x7fdc337a560c in /home/ansar/venvir/ml/lib64/python3.11/site-packages/torch/lib/libtorch_cuda.so)
frame #5: <unknown function> + 0x543010 (0x7fdc7cb43010 in /home/ansar/venvir/ml/lib64/python3.11/site-packages/torch/lib/libtorch_python.so)
frame #6: <unknown function> + 0x649bf (0x7fdc7e4b39bf in /home/ansar/venvir/ml/lib64/python3.11/site-packages/torch/lib/libc10.so)
frame #7: c10::TensorImpl::~TensorImpl() + 0x21b (0x7fdc7e4acc8b in /home/ansar/venvir/ml/lib64/python3.11/site-packages/torch/lib/libc10.so)
frame #8: c10::TensorImpl::~TensorImpl() + 0x9 (0x7fdc7e4ace39 in /home/ansar/venvir/ml/lib64/python3.11/site-packages/torch/lib/libc10.so)
frame #9: <unknown function> + 0x80b718 (0x7fdc7ce0b718 in /home/ansar/venvir/ml/lib64/python3.11/site-packages/torch/lib/libtorch_python.so)
frame #10: THPVariable_subclass_dealloc(_object*) + 0x2f6 (0x7fdc7ce0ba96 in /home/ansar/venvir/ml/lib64/python3.11/site-packages/torch/lib/libtorch_python.so)
<omitting python frames>

terminate called after throwing an instance of 'c10::Error'
  what():  CUDA error: initialization error
CUDA kernel errors might be asynchronously reported at some other API call, so the stacktrace below might be incorrect.
For debugging consider passing CUDA_LAUNCH_BLOCKING=1.
Compile with `TORCH_USE_CUDA_DSA` to enable device-side assertions.

Exception raised from c10_cuda_check_implementation at ../c10/cuda/CUDAException.cpp:44 (most recent call first):
frame #0: c10::Error::Error(c10::SourceLocation, std::string) + 0x57 (0x7fdc7e4ced87 in /home/ansar/venvir/ml/lib64/python3.11/site-packages/torch/lib/libc10.so)
frame #1: c10::detail::torchCheckFail(char const*, char const*, unsigned int, std::string const&) + 0x64 (0x7fdc7e47f75f in /home/ansar/venvir/ml/lib64/python3.11/site-packages/torch/lib/libc10.so)
frame #2: c10::cuda::c10_cuda_check_implementation(int, char const*, char const*, int, bool) + 0x118 (0x7fdc7e59f8a8 in /home/ansar/venvir/ml/lib64/python3.11/site-packages/torch/lib/libc10_cuda.so)
frame #3: c10::cuda::ExchangeDevice(int) + 0x8a (0x7fdc7e59fd0a in /home/ansar/venvir/ml/lib64/python3.11/site-packages/torch/lib/libc10_cuda.so)
frame #4: <unknown function> + 0xfa560c (0x7fdc337a560c in /home/ansar/venvir/ml/lib64/python3.11/site-packages/torch/lib/libtorch_cuda.so)
frame #5: <unknown function> + 0x543010 (0x7fdc7cb43010 in /home/ansar/venvir/ml/lib64/python3.11/site-packages/torch/lib/libtorch_python.so)
frame #6: <unknown function> + 0x649bf (0x7fdc7e4b39bf in /home/ansar/venvir/ml/lib64/python3.11/site-packages/torch/lib/libc10.so)
frame #7: c10::TensorImpl::~TensorImpl() + 0x21b (0x7fdc7e4acc8b in /home/ansar/venvir/ml/lib64/python3.11/site-packages/torch/lib/libc10.so)
frame #8: c10::TensorImpl::~TensorImpl() + 0x9 (0x7fdc7e4ace39 in /home/ansar/venvir/ml/lib64/python3.11/site-packages/torch/lib/libc10.so)
frame #9: <unknown function> + 0x80b718 (0x7fdc7ce0b718 in /home/ansar/venvir/ml/lib64/python3.11/site-packages/torch/lib/libtorch_python.so)
frame #10: THPVariable_subclass_dealloc(_object*) + 0x2f6 (0x7fdc7ce0ba96 in /home/ansar/venvir/ml/lib64/python3.11/site-packages/torch/lib/libtorch_python.so)
<omitting python frames>

terminate called after throwing an instance of 'c10::Error'
  what():  CUDA error: initialization error
CUDA kernel errors might be asynchronously reported at some other API call, so the stacktrace below might be incorrect.
For debugging consider passing CUDA_LAUNCH_BLOCKING=1.
Compile with `TORCH_USE_CUDA_DSA` to enable device-side assertions.

Exception raised from c10_cuda_check_implementation at ../c10/cuda/CUDAException.cpp:44 (most recent call first):
frame #0: c10::Error::Error(c10::SourceLocation, std::string) + 0x57 (0x7fdc7e4ced87 in /home/ansar/venvir/ml/lib64/python3.11/site-packages/torch/lib/libc10.so)
frame #1: c10::detail::torchCheckFail(char const*, char const*, unsigned int, std::string const&) + 0x64 (0x7fdc7e47f75f in /home/ansar/venvir/ml/lib64/python3.11/site-packages/torch/lib/libc10.so)
frame #2: c10::cuda::c10_cuda_check_implementation(int, char const*, char const*, int, bool) + 0x118 (0x7fdc7e59f8a8 in /home/ansar/venvir/ml/lib64/python3.11/site-packages/torch/lib/libc10_cuda.so)
frame #3: c10::cuda::ExchangeDevice(int) + 0x8a (0x7fdc7e59fd0a in /home/ansar/venvir/ml/lib64/python3.11/site-packages/torch/lib/libc10_cuda.so)
frame #4: <unknown function> + 0xfa560c (0x7fdc337a560c in /home/ansar/venvir/ml/lib64/python3.11/site-packages/torch/lib/libtorch_cuda.so)
frame #5: <unknown function> + 0x543010 (0x7fdc7cb43010 in /home/ansar/venvir/ml/lib64/python3.11/site-packages/torch/lib/libtorch_python.so)
frame #6: <unknown function> + 0x649bf (0x7fdc7e4b39bf in /home/ansar/venvir/ml/lib64/python3.11/site-packages/torch/lib/libc10.so)
frame #7: c10::TensorImpl::~TensorImpl() + 0x21b (0x7fdc7e4acc8b in /home/ansar/venvir/ml/lib64/python3.11/site-packages/torch/lib/libc10.so)
frame #8: c10::TensorImpl::~TensorImpl() + 0x9 (0x7fdc7e4ace39 in /home/ansar/venvir/ml/lib64/python3.11/site-packages/torch/lib/libc10.so)
frame #9: <unknown function> + 0x80b718 (0x7fdc7ce0b718 in /home/ansar/venvir/ml/lib64/python3.11/site-packages/torch/lib/libtorch_python.so)
frame #10: THPVariable_subclass_dealloc(_object*) + 0x2f6 (0x7fdc7ce0ba96 in /home/ansar/venvir/ml/lib64/python3.11/site-packages/torch/lib/libtorch_python.so)
<omitting python frames>

Traceback (most recent call last):
  File "/home/ansar/venvir/ml/lib64/python3.11/site-packages/torch/utils/data/dataloader.py", line 1133, in _try_get_data
    data = self._data_queue.get(timeout=timeout)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/lib64/python3.11/queue.py", line 180, in get
    self.not_empty.wait(remaining)
  File "/usr/lib64/python3.11/threading.py", line 324, in wait
    gotit = waiter.acquire(True, timeout)
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/ansar/venvir/ml/lib64/python3.11/site-packages/torch/utils/data/_utils/signal_handling.py", line 66, in handler
    _error_if_any_worker_fails()
RuntimeError: DataLoader worker (pid 1274565) is killed by signal: Aborted. 

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "/home/ansar/point_clouds/cli.py", line 15, in <module>
    cli_main()
  File "/home/ansar/point_clouds/cli.py", line 11, in cli_main
    cli = LightningCLI(PointNetLit, PCDDataModule)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/ansar/venvir/ml/lib64/python3.11/site-packages/lightning/pytorch/cli.py", line 388, in __init__
    self._run_subcommand(self.subcommand)
  File "/home/ansar/venvir/ml/lib64/python3.11/site-packages/lightning/pytorch/cli.py", line 679, in _run_subcommand
    fn(**fn_kwargs)
  File "/home/ansar/venvir/ml/lib64/python3.11/site-packages/lightning/pytorch/trainer/trainer.py", line 544, in fit
    call._call_and_handle_interrupt(
  File "/home/ansar/venvir/ml/lib64/python3.11/site-packages/lightning/pytorch/trainer/call.py", line 44, in _call_and_handle_interrupt
    return trainer_fn(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/ansar/venvir/ml/lib64/python3.11/site-packages/lightning/pytorch/trainer/trainer.py", line 580, in _fit_impl
    self._run(model, ckpt_path=ckpt_path)
  File "/home/ansar/venvir/ml/lib64/python3.11/site-packages/lightning/pytorch/trainer/trainer.py", line 987, in _run
    results = self._run_stage()
              ^^^^^^^^^^^^^^^^^
  File "/home/ansar/venvir/ml/lib64/python3.11/site-packages/lightning/pytorch/trainer/trainer.py", line 1033, in _run_stage
    self.fit_loop.run()
  File "/home/ansar/venvir/ml/lib64/python3.11/site-packages/lightning/pytorch/loops/fit_loop.py", line 205, in run
    self.advance()
  File "/home/ansar/venvir/ml/lib64/python3.11/site-packages/lightning/pytorch/loops/fit_loop.py", line 363, in advance
    self.epoch_loop.run(self._data_fetcher)
  File "/home/ansar/venvir/ml/lib64/python3.11/site-packages/lightning/pytorch/loops/training_epoch_loop.py", line 140, in run
    self.advance(data_fetcher)
  File "/home/ansar/venvir/ml/lib64/python3.11/site-packages/lightning/pytorch/loops/training_epoch_loop.py", line 212, in advance
    batch, _, __ = next(data_fetcher)
                   ^^^^^^^^^^^^^^^^^^
  File "/home/ansar/venvir/ml/lib64/python3.11/site-packages/lightning/pytorch/loops/fetchers.py", line 133, in __next__
    batch = super().__next__()
            ^^^^^^^^^^^^^^^^^^
  File "/home/ansar/venvir/ml/lib64/python3.11/site-packages/lightning/pytorch/loops/fetchers.py", line 60, in __next__
    batch = next(self.iterator)
            ^^^^^^^^^^^^^^^^^^^
  File "/home/ansar/venvir/ml/lib64/python3.11/site-packages/lightning/pytorch/utilities/combined_loader.py", line 341, in __next__
    out = next(self._iterator)
          ^^^^^^^^^^^^^^^^^^^^
  File "/home/ansar/venvir/ml/lib64/python3.11/site-packages/lightning/pytorch/utilities/combined_loader.py", line 78, in __next__
    out[i] = next(self.iterators[i])
             ^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/ansar/venvir/ml/lib64/python3.11/site-packages/torch/utils/data/dataloader.py", line 631, in __next__
    data = self._next_data()
           ^^^^^^^^^^^^^^^^^
  File "/home/ansar/venvir/ml/lib64/python3.11/site-packages/torch/utils/data/dataloader.py", line 1329, in _next_data
    idx, data = self._get_data()
                ^^^^^^^^^^^^^^^^
  File "/home/ansar/venvir/ml/lib64/python3.11/site-packages/torch/utils/data/dataloader.py", line 1285, in _get_data
    success, data = self._try_get_data()
                    ^^^^^^^^^^^^^^^^^^^^
  File "/home/ansar/venvir/ml/lib64/python3.11/site-packages/torch/utils/data/dataloader.py", line 1146, in _try_get_data
    raise RuntimeError(f'DataLoader worker (pid(s) {pids_str}) exited unexpectedly') from e
RuntimeError: DataLoader worker (pid(s) 1274565, 1274568) exited unexpectedly
Epoch 0:   0%|          | 0/1000 [00:00<?, ?it/s]

It passes the sanity checks for the validation dataloader but then fails. If I just use num_workers == 0 it runs without problems also with LightningCLI. Any ideas?

Note

Adding fast_dev_run == True also runs without problems with num_workers > 0 and LightningCLI.