Not working within multiprocessing

[DirkRemmers]

Hi all,

First of all I want to say what a super nice package this is, thanks for the hard work! Second, I am not sure if this is an issue on the pyclesperanto side, or on the multiprocessing package. However, I thought it would be good to report the issue here for sure.

At the moment I'm trying to setup a high-throughput image analysis pipeline. For some of the functionality I need CPU processes, while for other processes I want to use the GPU. To speed up the CPU processes, I would like to use multiprocessing to optimize my workstation. However, I might have found an issue while trying to design a script that uses both at the same time.

When I have the GPU processes outside of the multiprocessing functionality, all works nicely. A simplified example:

from multiprocessing import Pool
from skimage import filters
import pyclesperanto_prototype as cle

def image_analysis_CPU(image):
     # do something with the image
     blurred_image = filters.gaussian(image, sigma = 2)
     return blurred_image
     
# get a list of images called image_list and continue with the CPU processes

n_cores = 4 # example
pool = multiprocessing.Pool(n_cores)
result_list = pool.map(image_analysis_CPU, image_list)
pool.close()
pool.join()

# then loop over the result_list and do all GPU processes
# notice that if there are many images, the memory requirements of result_list are becomming very high
cle.select_device('RTX') 
gpu_result_list = []
for result in result_list:
     gpu_input = cle.push(result)
     # do my image analysis to get to gpu_result
     analysis_result = cle.pull(gpu_result)
     gpu_result_list.append(analysis_result)

# continue with the rest of the script

However, due to the file sizes I'm going to process I am affraid this solution is not scalable since it would require to store a lot of data before processing using the GPU. The data that is required for the GPU process is relatively large, but can be removed after the GPU process is finished (I'm only interested in the result of this). After the GPU process has been performed, I am able to efficiently save my data, but saving the GPU input is not efficient and will require some unelegant data handling afterwards. Ideally, I would like to streamline my script so that I could make use of the GPU functionality within the process of a Pool worker. This way, I can prevent saving lots of data that is required to do the GPU processing. I included a lock to prevent situations where multiple workers are trying to access the GPU at the same time. A simplified example code would be:

from multiprocessing import Pool, Lock
from skimage import filters
import pyclesperanto_prototype as cle

def image_analysis(image):
    blurred_image = filters.gaussian(image, sigma = 2)
    lock.acquire() # to prevent multiple processes trying to access the GPU at the same time
    cle.select_device('RTX')
    gpu_input = cle.push(blurred_image)
    # do my image analysis to get to gpu_result
    analysis_result = cle.pull(gpu_result)
    lock.release()
    return analysis_result
    
def init(l):
    global lock
    lock = l
    
def main():
     # get a list of images called image_list and continue
     l = Lock() # define the lock outside the Pool so that they all have access to the same variable
     n_cores = 4 # example
     pool  = Pool(n_cores, initializer = init, initargs = (l,))
     result_list = pool.map(image_analysis, image_list)
     pool.close()
     pool.join()

# continue with the rest of the script

The issue that I get with this approach though is that all pyclesperanto functionality causes the Pool worker to freeze, it cannot pass any of the cle functions. Even selecting a GPU via the cle.select_device('RTX') line is not working. There are no error messages for me to include in this issue, the script just seems to get stuck while trying to perform a pyclesperanto related function.

Obviously the example code is simplified drastically, otherwise I could do the entire thing on the GPU. When a full script is required, please let me know and I can send all required example files.

Thanks in advance, I'm looking forward to your response.