Add multiprocess lock in dataset get item for ddp_spawn multigpu training.

[saedrna]

Hi,
For the PyTorch lightning dataset in this library, it only supports ddp_spawn for multi-GPU training. I have a dataset for large point clouds, in order to sample evenly distributed points during the training, a shared buffer recording already sampled parts are necessary. It is updated dynamically during getting the data. An example can be found in KPConf as below.

            with self.worker_lock:
                # Get potential minimum
                cloud_ind = int(torch.argmin(self.min_potentials))
                point_ind = int(self.argmin_potentials[cloud_ind])

                # Get potential points from tree structure
                pot_points = np.array(self.pot_trees[cloud_ind].data, copy=False)

                # Center point of input region
                center_point = pot_points[point_ind, :].reshape(1, -1)

                # Indices of points in input region
                pot_inds, dists = self.pot_trees[cloud_ind].query_radius(center_point,
                                                                         r=self.config.in_radius,
                                                                         return_distance=True)

                d2s = np.square(dists[0])
                pot_inds = pot_inds[0]

                # Update potentials (Tukey weights)
                if self.set != 'ERF':
                    tukeys = np.square(1 - d2s / np.square(self.config.in_radius))
                    tukeys[d2s > np.square(self.config.in_radius)] = 0
                    self.potentials[cloud_ind][pot_inds] += tukeys
                    min_ind = torch.argmin(self.potentials[cloud_ind])
                    self.min_potentials[[cloud_ind]] = self.potentials[cloud_ind][min_ind]
                    self.argmin_potentials[[cloud_ind]] = min_ind

I use the LightningDataset to wrap the dataloader and use DDPSpawnPlugin for strategy.
However, during training, the lock causes a segmentation fault ERROR: Unexpected segmentation fault encountered in worker.
I found in PyTorch, it says here.

Best practices and tips

Avoiding and fighting deadlocks

There are a lot of things that can go wrong when a new process is spawned, with the most common cause of deadlocks being background threads. If there’s any thread that holds a lock or imports a module, and fork is called, it’s very likely that the subprocess will be in a corrupted state and will deadlock or fail in a different way. Note that even if you don’t, Python built in libraries do - no need to look further than multiprocessing. multiprocessing.Queue is actually a very complex class, that spawns multiple threads used to serialize, send and receive objects, and they can cause aforementioned problems too. If you find yourself in such situation try using a multiprocessing.queues.SimpleQueue, that doesn’t use any additional threads.

We’re trying our best to make it easy for you and ensure these deadlocks don’t happen but some things are out of our control. If you have any issues you can’t cope with for a while, try reaching out on forums, and we’ll see if it’s an issue we can fix.

Is there a solution or workaround for this problem?

Thanks,
Han

[saedrna]

Another question is that: for example, setting num_workers = 0 in the dataloader for strategy=ddp_spawn with devices=4, how many threads will simultaneously access the dataloader to get the item?
Sorry, but I am new to PyTorch and can not find simple documentation for this.

[rusty1s]

If you use ddp_spawn on 4 devices with num_workers=0, there will be 4 processes requesting data from the data loading.

[rusty1s]

It is hard for me to give good advice from the snippet above. Is there a particular reason you are using a worker_lock?

[saedrna]

Because the dataset is shared between different processes. And I need to both read and write from the shared memory of the datasets. Commonly, the get operation from the dataset won't write anything to the dataset, but mine needs to update what is read (in order to keep balanced loading). The lock is necessary.
I have looked into the docs of both PyTorch and PyTorch lightning and found that ddp_spawn (in fact the mp.spawn) require all the stuff are pickleable. But unfortunately mp.Lock is not pickleable.
I have made some changes based on the LightningDataset, and changed to the ddp strategy. Currently, everything seems good enough for me. The additional memory costs may be inevitable.

[rusty1s]

Yes, this is a good workaround if the additional memory costs are negligible. However, keep in mind that each node/each GPU will only operate on its own isolated dataset. If you want to write to the dataset, these operations will only be visible within the same GPU.