web-dev-qa-db-fra.com

comment utiliser l'initialiseur pour configurer mon pool multiprocessus?

J'essaie d'utiliser l'objet Pool multiprocessus. J'aimerais que chaque processus ouvre une connexion à la base de données au démarrage, puis utilise cette connexion pour traiter les données transmises. (Plutôt que d'ouvrir et de fermer la connexion pour chaque bit de données.) Cela ressemble à ce qu'est l'initialiseur mais je ne peux pas comprendre comment le travailleur et l'initialiseur communiquent. J'ai donc quelque chose comme ça:

def get_cursor():
  return psycopg2.connect(...).cursor()

def process_data(data):
   # here I'd like to have the cursor so that I can do things with the data

if __name__ == "__main__":
  pool = Pool(initializer=get_cursor, initargs=())
  pool.map(process_data, get_some_data_iterator())

comment puis-je (ou dois-je) récupérer le curseur de get_cursor () dans process_data ()?

pythonmultiprocessing
54
Chris Curvey

La fonction d'initialisation est appelée ainsi:

def worker(...):
    ...
    if initializer is not None:
        initializer(*args)

il n'y a donc aucune valeur de retour enregistrée nulle part. Vous pensez peut-être que cela vous condamne, mais non! Chaque travailleur est dans un processus distinct. Ainsi, vous pouvez utiliser une variable global ordinaire.

Ce n'est pas exactement joli, mais ça marche:

cursor = None
def set_global_cursor(...):
    global cursor
    cursor = ...

Maintenant, vous pouvez simplement utiliser cursor dans votre process_data fonction. La variable cursor à l'intérieur de chaque processus distinct est distincte de tous les autres processus, de sorte qu'ils ne marchent pas les uns sur les autres.

(Je ne sais pas si psycopg2 a une manière différente de gérer cela qui n'implique pas d'utiliser multiprocessing en premier lieu; il s'agit d'une réponse générale à un problème général avec le module multiprocessing.)

81
torek

torek a déjà bien expliqué pourquoi l'initialiseur ne fonctionne pas dans ce cas. Cependant, je ne suis pas fan de variable globale personnellement, je voudrais donc coller une autre solution ici.

L'idée est d'utiliser une classe pour encapsuler la fonction et initialiser la classe avec la variable "globale".

class Processor(object):
  """Process the data and save it to database."""

  def __init__(self, credentials):
    """Initialize the class with 'global' variables"""
    self.cursor = psycopg2.connect(credentials).cursor()

  def __call__(self, data):
    """Do something with the cursor and data"""
    self.cursor.find(data.key)

Et puis appeler avec

p = Pool(5)
p.map(Processor(credentials), list_of_data)

Ainsi, le premier paramètre a initialisé la classe avec des informations d'identification, renvoie une instance de la classe et la carte appelle l'instance avec des données.

Bien que ce ne soit pas aussi simple que la solution de variable globale, je suggère fortement d'éviter la variable globale et d'encapsuler les variables de manière sûre. (Et je souhaite vraiment qu'ils puissent soutenir l'expression lambda un jour, cela rendra les choses beaucoup plus faciles ...)

11
yeelan

Vous pouvez également envoyer la fonction à l'initialiseur et y créer une connexion. Ensuite, vous ajoutez le curseur à la fonction.

def init_worker(function):
    function.cursor = db.conn()

Vous pouvez maintenant accéder à la base de données via function.cursor sans utiliser de globaux, par exemple:

def use_db(i):
    print(use_db.cursor) #process local
pool = Pool(initializer=init_worker, initargs=(use_db,))
pool.map(use_db, range(10))
9
The Unfun Cat

Étant donné que la définition de variables globales dans l'initialiseur n'est généralement pas souhaitable, nous pouvons éviter leur utilisation et également éviter de répéter une initialisation coûteuse dans chaque appel avec une mise en cache simple dans chaque sous-processus:

from functools import lru_cache
from multiprocessing.pool import Pool
from time import sleep


@lru_cache(maxsize=None)
def _initializer(a, b):
    print(f'Initialized with {a}, {b}')


def _pool_func(a, b, i):
    _initializer(a, b)
    sleep(1)
    print(f'got {i}')


arg_a = 1
arg_b = 2

with Pool(processes=5) as pool:
    pool.starmap(_pool_func, ((arg_a, arg_b, i) for i in range(0, 20)))

Production:

Initialized with 1, 2
Initialized with 1, 2
Initialized with 1, 2
Initialized with 1, 2
Initialized with 1, 2
got 1
got 0
got 4
got 2
got 3
got 5
got 7
got 8
got 6
got 9
got 10
got 11
got 12
got 14
got 13
got 15
got 16
got 17
got 18
got 19
5
mcguip