Tensorflow-gpu and multiprocessing

Question

I have finished implementing an Asynchronous Advantage Actor-Critic (A3C) agent for TensorFlow (gpu). By using a single RMSprop optimizer with shared statistics. To do so, a central controller holds both the Global Network (ActorCriticModel) and the Optimizer. Workers, to update the global network, need to communicate gradients to the Controller, while receiving back the updated model weights. It works decently: despite there might be some errors in the code, it still achieves the results pretty fast. Each Worker is a thread on its own.

class Worker(Thread):

Problem is Workers need to extend the Thread class while I would like to use heavy Processes (from multiprocessing) instead (Queues used in the following snippet of code are imported from the multiprocessing class).

class Controller:

    def __init__(self, **kwargs):

        self.terminated_workers = 0
        self.env_name = kwargs["env_name"]
        self.available_actions = kwargs["available_actions"]
        self.action_size = len(self.available_actions)
        self.env = make_env(kwargs, env_name=self.env_name)
        self.state_shape = self.env.observation_space.shape

        self.number_of_workers = kwargs["number_of_workers"]
        self.workers = []
        self.controller_queue = Queue()
        self.worker_queues = []

        # Initialization of threads
        for i in range(self.number_of_workers):
            worker_queue = Queue()
            self.worker_queues.append(worker_queue)

            updater = GlobalUpdater(self.controller_queue, worker_queue)
            worker = Worker(i, updater, self.state_shape, kwargs)
            self.workers.append(worker)

        # Only import tensorflow after heavy Process initialization
        import tensorflow as tf

        self.global_model = ActorCriticModel(self.state_shape, self.action_size)
        self.opt = tf.keras.optimizers.RMSprop(
            learning_rate=kwargs["lr_schedule"],
            rho=kwargs["rho"],
            centered=kwargs["centered"],
            clipvalue=kwargs["clipvalue"]
        ) 

As soon as i change the Worker to extend Processes instead of threads, I get the following message for every worker instantiated:

F tensorflow/stream_executor/cuda/cuda_driver.cc:152] Failed setting context: CUDA_ERROR_NOT_INITIALIZED: initialization error

Technically, since multiple methods of the Worker class need to use tensorflow, the library is imported statically, before declaring the Worker class, as follows:

import tensorflow as tf

class Worker(Process):
   ...

I'm using Linux and the libraries' versions are:

• TensorFlow: 2.9.1
• CUDA: 11.2.2-1
• cuDNN: 8.1.1.33-1+cuda11.2

I tried every other solution I found online, including setting mp.set_start_method('spawn') and trying to postpone as much as i can the import of the tensorflow library, so to postpone as much as possible the heavy process initialization but nothing seems to work. How would you approach this problem? Do you have any idea on how to solve it? Thanks in advance.