Problem with multigpu parallel External Source operator

So I was following the tutorials for parallel External Source operator and have chosen to setup py_start_method to "spawn". Unfortunately while running it with pytorch lightning I've only recieved errors similar to ones attached here 4gpus_error.log.

I'm not sure whether it can be attributed to DALI or pytorch-lightning. I've tried all of the pickling methods mentioned in tutorials but none resolved this issue. The example below uses DDP but the error occured also in 1gpu setup. Locally on my laptop it runs smoothly, but I cannot reproduce the setup with multigpu. In my normal setup I'm running it on GCP NVIDIA 4xA100 machine.

Self-contained example reproducing the issue for 4-gpus based machine:

import random

from pathlib import Path
from typing import Union, Callable, List, Optional, Dict
import numpy as np
from PIL import Image
import cv2

from nvidia import dali
from nvidia.dali import pipeline_def, Pipeline
from nvidia.dali import fn
from nvidia.dali.plugin.pytorch import DALIGenericIterator
from nvidia.dali import pickling as dali_pickle
from pytorch_lightning import LightningDataModule, LightningModule, Trainer
import torch
from torch import nn


@dali_pickle.pickle_by_value
def SegmentationCallback(
    imgs_list: List[str],
    depths_list: List[str],
    labels_list: List[str],
    batch_size: int,
    seed: int,
    shard_id: int = 0,
    num_shards: int = 1,
    **kwargs,
):
    indices = np.arange(len(imgs_list))
    np.random.shuffle(indices)

    dataset_size = len(imgs_list)
    shard_size = dataset_size // num_shards
    shard_offset = shard_size * shard_id
    n_iterations = shard_size // batch_size
    perm = None
    last_seen_epoch = None

    def callback(sample_info):
        nonlocal perm, last_seen_epoch
        if sample_info.iteration >= n_iterations:
            raise StopIteration
        if last_seen_epoch != sample_info.epoch_idx:
            last_seen_epoch = sample_info.epoch_idx
            perm = np.random.default_rng(seed=seed + sample_info.epoch_idx).permutation(dataset_size)

        index = perm[sample_info.idx_in_epoch + shard_offset]
        img = np.frombuffer(open(imgs_list[index], "rb").read(), dtype=np.uint8)
        dsm = cv2.imread(depths_list[index], cv2.IMREAD_ANYDEPTH)[..., np.newaxis]
        label = np.frombuffer(open(labels_list[index], "rb").read(), dtype=np.uint8)
        return img, dsm, label

    return callback


@pipeline_def
def EnhancedPipeline(
    data_iterator: Callable,
    precision: int = 32,
):

    use_gpu = Pipeline.current().device_id is not None
    decoder_device, device = ("mixed", "gpu") if use_gpu else ("cpu", "cpu")

    img_type = dali.types.FLOAT16 if precision == 16 else dali.types.FLOAT
    imgs, depth, labels = fn.external_source(
        source=data_iterator,
        num_outputs=3,
        batch=False,
        parallel=True,
        dtype=[dali.types.UINT8, dali.types.FLOAT, dali.types.UINT8],
    )
    if use_gpu:
        depth = depth.gpu()

    imgs = fn.decoders.image(imgs, device=decoder_device)
    labels = fn.decoders.image(labels, device=decoder_device, output_type=dali.types.GRAY)

    imgs /= 255

    labels = fn.cast(labels, dtype=dali.types.INT64)
    imgs = fn.cast(imgs, dtype=img_type)
    depth = fn.cast(depth, dtype=img_type)

    imgs = fn.cat(imgs, depth, axis=2)
    imgs, labels = fn.transpose(
        [imgs, labels],
        perm=[2, 0, 1],
    )

    return imgs, labels


class DALIGenericIteratorWrapper(DALIGenericIterator):
    """
    Class wrapping DALIGenericIterator in order to obtain the batches in the pytorches iterable format (input, label).
    """

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

    def __next__(self):
        out = super().__next__()
        out = [out[0][output_name] for output_name in self.output_map]
        return out


class Model(LightningModule):
    def __init__(self, num_classes=int, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.model = nn.Conv2d(4, num_classes, 1)
        self.criterion = nn.CrossEntropyLoss()

    def forward(self, x):
        return self.model(x)

    def training_step(self, batch, batch_idx):
        return self._step(batch)

    def validation_step(self, batch, batch_idx):
        return self._step(batch)

    def test_step(self, batch, batch_idx, dataloader_idx: Optional[int] = 0):
        return self._step(batch, dataloader_idx=dataloader_idx)

    def _step(self, batch, dataloader_idx=None):
        inputs, masks = batch
        outputs = self.forward(inputs)
        loss = self.criterion(outputs, masks.squeeze(1))
        return loss

    def configure_optimizers(self):
        optimizer = torch.optim.Adam(self.parameters(), lr=1e-3)
        return optimizer


class SegmentationDataModule(LightningDataModule):
    def __init__(
        self,
        imgs_list: str,
        depth_list: str,
        labels_list: str,
        batch_size: int,
        num_workers: int,
        gpus: Optional[Union[List[int], str, int]] = None,
        accelerator: str = None,
        precision: int = 32,
        seed: int = 42,
        **kwargs,
    ):
        super().__init__()
        self.imgs_list = imgs_list
        self.depth_list = depth_list
        self.labels_list = labels_list

        self.pipeline_config = {
            "num_threads": num_workers if num_workers != 0 else 1,
            "device_id": None if not gpus else 0,
            "seed": seed,
            "batch_size": batch_size,
            "precision": precision,
            "py_num_workers": num_workers if num_workers != 0 else 1,
            "py_start_method": "spawn",
        }
        self.accelerator = accelerator

    def setup(self, stage: Optional[str] = None):
        paths = {"imgs": self.imgs_list, "depths": self.depth_list, "labels": self.labels_list}
        self.train_dataset = {**paths}
        self.val_dataset = {**paths}
        self.test_dataset = {**paths}

    def train_dataloader(self):
        return self._get_iterator(
            self.train_dataset["imgs"],
            self.train_dataset["depths"],
            self.train_dataset["labels"],
            self.pipeline_config,
        )

    def val_dataloader(self):
        return self._get_iterator(
            self.val_dataset["imgs"],
            self.val_dataset["depths"],
            self.val_dataset["labels"],
            self.pipeline_config,
        )

    def test_dataloader(self):
        return self._get_iterator(
            self.test_dataset["imgs"],
            self.test_dataset["depths"],
            self.test_dataset["labels"],
            self.pipeline_config,
        )

    def _get_iterator(self, imgs_list: List, depths_list: List, labels_list: List, pipeline_config: Dict):
        assert len(imgs_list) == len(labels_list)
        if (
            self.accelerator
        ):  # has to be assigned for DDP, is available only after pytorch-lightning trainer initialization
            pipeline_config = {
                **pipeline_config,
                "device_id": self.trainer.local_rank,
            }

        data_callable = SegmentationCallback(
            imgs_list,
            depths_list,
            labels_list,
            shard_id=self.trainer.global_rank,
            num_shards=self.trainer.world_size,
            **pipeline_config,
        )
        pipe = EnhancedPipeline(data_callable, **pipeline_config)
        pipe.build()

        return DALIGenericIteratorWrapper(
            pipelines=pipe, output_map=["imgs", "labels"], size=len(imgs_list) // (self.trainer.world_size)
        )


if __name__ == "__main__":
    seed = 42
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)

    img = np.zeros((512, 512, 3), dtype=np.uint8)
    img[256:, 256:, 0] = 255

    depth = np.zeros((512, 512))
    depth[256:, 256:] = 100

    label = np.zeros((512, 512), dtype=np.uint8)
    label[256:, 256:] = 1

    data_root = Path("imgs")
    data_root.mkdir(parents=True, exist_ok=True)
    Image.fromarray(img).save(data_root / "img.png")
    Image.fromarray(depth, mode="F").save(data_root / "depth.tif", compression="tiff_lzw")
    Image.fromarray(label).save(data_root / "label.png")

    imgs_list = [str(data_root / "img.png")] * 100
    depths_list = [str(data_root / "depth.tif")] * 100
    labels_list = [str(data_root / "label.png")] * 100

    batch_size = 4
    num_workers = 6
    gpus = 4
    accelerator = "ddp"
    precision = 16
    model = Model(2)
    datamodule = SegmentationDataModule(
        imgs_list, depths_list, labels_list, batch_size, num_workers, gpus, accelerator, precision
    )

    trainer_config = {"gpus": gpus, "max_epochs": 10, "accelerator": accelerator, "precision": precision}
    trainer = Trainer(**trainer_config)
    trainer.fit(model=model, datamodule=datamodule)

With this one solved, it would be the last of one from my odyssey of DALI issues, thanks a lot for your help!