mp.set_start_method('spawn') triggered an error saying the context is already been set

Question

Here is my full code

I have succeeded to reproduce the behavior of my main code with a little snippet.

In a Google Colab Env, suppose I setup hardware accelerator to GPU.

Here is the little snippet :

import multiprocessing as mp

def foo(q):
    q.put('hello')

if __name__ == '__main__':
    mp.set_start_method('spawn')
    q = mp.Queue()
    p = mp.Process(target=foo, args=(q,))
    p.start()
    print(q.get())
    p.join()

With mp.set_start_method('spawn'), I got the error RuntimeError: context has already been set and if I use mp.set_start_method('spawn', force=True) it got stuck in an infinite loop if I can say.

Is there any way I can prevent that error from occurring in the Colab Env?

P.S. Be aware mp.set_start_method('spawn') this line of code is necessary in my code. Otherwise, I got the error RuntimeError: Cannot re-initialize CUDA in forked subprocess. To use CUDA with multiprocessing, you must use the 'spawn' start method

Overkilled Solution

The problem here is that the spawned subprocess can't find __main__.foo.

One (less than perfect) option is to put the code in a separate file, e.g. make a new directory and add it to the path.

I don't want to use that solution because it is way overkilled for my main code. Maybe it will guide you to a more elegant solution.

UPDATE

This seems to be a more reasonable solution, but still looking for an answer more elegant.

Answer 1

You should set start method before a function call(except main) as spawn.

Example usage:

import multiprocessing as mp

try:
   mp.set_start_method('spawn', force=True)
   print("spawned")
except RuntimeError:
   pass

I usually use this block for inference with multiprocessing in PyTorch