web-dev-qa-db-fra.com

Python multiprocessing avec générateur

J'essaie de traiter un fichier (chaque ligne est un document json). La taille du fichier peut aller jusqu'à 100 mbs en Go. J'ai donc écrit un code générateur pour récupérer chaque document ligne par ligne à partir d'un fichier.

def jl_file_iterator(file):
    with codecs.open(file, 'r', 'utf-8') as f:
        for line in f:
            document = json.loads(line)
            yield document

Mon système a 4 cœurs, je voudrais donc traiter 4 lignes du fichier en parallèle. Actuellement, j'ai ce code qui prend 4 lignes à la fois et appelle le code pour le traitement parallèle

threads = 4
files, i = [], 1
for jl in jl_file_iterator(input_path):
    files.append(jl)
    if i % (threads) == 0:
        # pool.map(processFile, files)
        parallelProcess(files, o)
        files = []
    i += 1

if files:
    parallelProcess(files, o)
    files = []

Ceci est mon code où le traitement réel se produit

def parallelProcess(files, outfile):
    processes = []
    for i in range(len(files)):
        p = Process(target=processFile, args=(files[i],))
        processes.append(p)
        p.start()
    for i in range(len(files)):
        processes[i].join()

def processFile(doc):
    extractors = {}
    ... do some processing on doc
    o.write(json.dumps(doc) + '\n')

Comme vous pouvez le voir, j'attends que les 4 lignes aient terminé le traitement avant d'envoyer les 4 prochains fichiers à traiter. Mais ce que je voudrais faire, c'est dès qu'un processus a terminé le traitement du fichier, je veux commencer la ligne suivante à affecter au processeur libéré. Comment je fais ça?

PS: Le problème est que c'est un générateur, je ne peux pas charger tous les fichiers et utiliser quelque chose comme map pour exécuter les processus.

Merci de votre aide

17
Muthu Rg

Comme @pvg l'a dit dans un commentaire, une file d'attente (limitée) est le moyen naturel de servir de médiateur entre un producteur et des consommateurs à des vitesses différentes, en veillant à ce qu'ils restent tous aussi occupés que possible mais sans laisser le producteur aller de l'avant.

Voici un exemple exécutable autonome. La file d'attente est limitée à une taille maximale égale au nombre de processus de travail. Si les consommateurs courent beaucoup plus vite que le producteur, il pourrait être judicieux de laisser la file d'attente s'allonger.

Dans votre cas spécifique, il serait probablement judicieux de transmettre des lignes aux consommateurs et de les laisser faire la partie document = json.loads(line) en parallèle.

import multiprocessing as mp

NCORE = 4

def process(q, iolock):
    from time import sleep
    while True:
        stuff = q.get()
        if stuff is None:
            break
        with iolock:
            print("processing", stuff)
        sleep(stuff)

if __name__ == '__main__':
    q = mp.Queue(maxsize=NCORE)
    iolock = mp.Lock()
    pool = mp.Pool(NCORE, initializer=process, initargs=(q, iolock))
    for stuff in range(20):
        q.put(stuff)  # blocks until q below its max size
        with iolock:
            print("queued", stuff)
    for _ in range(NCORE):  # tell workers we're done
        q.put(None)
    pool.close()
    pool.join()
18
Tim Peters

J'ai donc fini par gérer cela avec succès. En créant des morceaux de lignes à partir de mon fichier et en exécutant les lignes en parallèle. L'afficher ici afin qu'il puisse être utile à quelqu'un à l'avenir.

def run_parallel(self, processes=4):
    processes = int(processes)
    pool = mp.Pool(processes)
    try:
        pool = mp.Pool(processes)
        jobs = []
        # run for chunks of files
        for chunkStart,chunkSize in self.chunkify(input_path):
            jobs.append(pool.apply_async(self.process_wrapper,(chunkStart,chunkSize)))
        for job in jobs:
            job.get()
        pool.close()
    except Exception as e:
        print e

def process_wrapper(self, chunkStart, chunkSize):
    with open(self.input_file) as f:
        f.seek(chunkStart)
        lines = f.read(chunkSize).splitlines()
        for line in lines:
            document = json.loads(line)
            self.process_file(document)

# Splitting data into chunks for parallel processing
def chunkify(self, filename, size=1024*1024):
    fileEnd = os.path.getsize(filename)
    with open(filename,'r') as f:
        chunkEnd = f.tell()
        while True:
            chunkStart = chunkEnd
            f.seek(size,1)
            f.readline()
            chunkEnd = f.tell()
            yield chunkStart, chunkEnd - chunkStart
            if chunkEnd > fileEnd:
                break
8
Muthu Rg

La réponse de Tim Peters est super.
Mais mon cas spécifique était légèrement différent, et j'ai dû modifier sa réponse pour l'adapter à mes besoins. Référencement ici.
Cela répond à la question de @CpILL dans les commentaires.


Dans mon cas, j'ai utilisé une chaîne de générateur (pour créer un pipeline).
Parmi cette chaîne de générateurs, l'un d'eux effectuait de lourds calculs, ralentissant tout le pipeline.

Quelque chose comme ça :

def fast_generator1():
    for line in file:
        yield line

def slow_generator(lines):
    for line in lines:
        yield heavy_processing(line)

def fast_generator2():
    for line in lines:
        yield fast_func(line)

if __name__ == "__main__":
    lines = fast_generator1()
    lines = slow_generator(lines)
    lines = fast_generator2(lines)
    for line in lines:
        print(line)

Pour le rendre plus rapide, nous devons exécuter le générateur lent avec plusieurs processus.
Le code modifié ressemble à:

import multiprocessing as mp

NCORE = 4

def fast_generator1():
    for line in file:
        yield line

def slow_generator(lines):
    def gen_to_queue(input_q, lines):
        # This function simply consume our generator and write it to the input queue
        for line in lines:
            input_q.put(line)
        for _ in range(NCORE):    # Once generator is consumed, send end-signal
            input_q.put(None)

    def process(input_q, output_q):
        while True:
            line = input_q.get()
            if line is None:
                output_q.put(None)
                break
            output_q.put(heavy_processing(line))


    input_q = mp.Queue(maxsize=NCORE * 2)
    output_q = mp.Queue(maxsize=NCORE * 2)

    # Here we need 3 groups of worker :
    # * One that will consume the input generator and put it into a queue. It will be `gen_pool`. It's ok to have only 1 process doing this, since this is a very light task
    # * One that do the main processing. It will be `pool`.
    # * One that read the results and yield it back, to keep it as a generator. The main thread will do it.
    gen_pool = mp.Pool(1, initializer=gen_to_queue, initargs=(input_q, lines))
    pool = mp.Pool(NCORE, initializer=process, initargs=(input_q, output_q))

    finished_workers = 0
    while True:
        line = output_q.get()
        if line is None:
            finished_workers += 1
            if finished_workers == NCORE:
                break
        else:
            yield line

def fast_generator2():
    for line in lines:
        yield fast_func(line)

if __name__ == "__main__":
    lines = fast_generator1()
    lines = slow_generator(lines)
    lines = fast_generator2(lines)
    for line in lines:
        print(line)

Avec cette implémentation, nous avons un générateur multiprocessus: il est utilisé exactement comme les autres générateurs (comme dans le premier exemple de cette réponse), mais tous les calculs lourds se font en utilisant le multiprocessing, l'accélérant!

0
Astariul