如何通过多处理实现生产者消费者？

Question

I have a program where I need to download files from some source and upload them . But I need to make sure that there are at max 10 files in the download location. Is there way to use Managers() as well?

It sounded like a typical Producer - Consumer problem. Below is my program.

Below is my implementation

from multiprocessing import Process, Queue, Lock
import requests
import json
import shutil
import os
import time
import random
import warnings
warnings.filterwarnings("ignore")

sha_list = [line.strip() for line in open("ShaList")]


def save_file_from_sofa(sha1):
    r = requests.get("https://DOWNLOAD_URL/{}".format(sha1), verify=False, stream=True)
    with open(sha1, 'wb') as handle:
        shutil.copyfileobj(r.raw, handle)


def mock_upload():
    time.sleep(random.randint(10,16))


def producer(queue, lock):
    with lock:
        print("Starting Producer {}".format(os.getpid()))

    while sha_list:
        if not queue.full():
            sha1 = sha_list.pop()
            save_file_from_sofa(sha1)
            queue.put(sha1)


def consumer(queue, lock):
    with lock:
        print("Starting Consumer {}".format(os.getpid()))

    while True:
        sha1 = queue.get()
        mock_upload()
        with lock:
            print("{} GOT {}".format(os.getpid(), sha1))

if __name__ == "__main__":
    queue = Queue(5)
    lock = Lock()

    producers = [Process(target=producer, args=(queue, lock)) for _ in range(2)]
    consumers = []

    for _ in range(3):
        p = Process(target=consumer, args=(queue, lock))
        p.daemon = True #Do not forget to set it to true
        consumers.append(p)

    for p in producers:
        p.start()
    for c in consumers:
        c.start()

    for p in producers:
        p.join()

    print("DONE")

But It does not do what is expected, as you can see from the output below

Starting Producer 623

Starting Producer 624

Starting Consumer 626

Starting Consumer 625

Starting Consumer 627

626 GOT 4ff551490d6b2eec7c6c0470f4b092fdc34cd521

625 GOT 83a53a3400fc83f2b02135ba0cc6c8625ecc7dc4

627 GOT 4ff551490d6b2eec7c6c0470f4b092fdc34cd521

626 GOT 83a53a3400fc83f2b02135ba0cc6c8625ecc7dc4

625 GOT 4e7132301ce9d61445db07910ff90a64474e6a88

626 GOT 0efbd413d733b3903e6dee777ace5ef47a2ec144

627 GOT 4e7132301ce9d61445db07910ff90a64474e6a88

625 GOT 0efbd413d733b3903e6dee777ace5ef47a2ec144

626 GOT 0a3fc4bdd56fa2bf52f5f43277f3b4ee0f040937

625 GOT eb9c07329a8b5cb66e47f0dd8e56894707a84d94

627 GOT 0a3fc4bdd56fa2bf52f5f43277f3b4ee0f040937

626 GOT eb9c07329a8b5cb66e47f0dd8e56894707a84d94

DONE

As you can see consumer picks up same SHA1s multiple times. So, I need a program to make sure that all the SHA1s put in the queue by producer is picked up by only 1 consumer.

PS I had also thought to make it work using pool. For producer it can work fine as I already have list of SHA1s to be put in the queue, But in case of consumer how would I use any list to make sure that consumer is actually stopping.

Answer 1

Just use a pool from either multiprocessing.Pool or concurrent.futures . The pool allows you to set how many workers you want running at the same time. This means you will have maximum max_workers files downloaded at the same time.

As the download/upload is sequential (you cannot start an upload until the download is complete), you gain no value from running them in two separated threads/processes. Just join the two operations in a single job unit and then run multiple jobs concurrently.

Moreover, as long as you just need to download/upload files (IO bound operations) you'd better use threads instead of processes as they are more lightweight.

from concurrent.futures import ThreadPoolExecutor

list_of_sha1s = ['foobar', 'foobaz']

def worker(sha1):
    path = save_file_from_sofa(sha1)
    upload_file(path)

    return sha1

with ThreadPoolExecutor(max_workers=10) as pool:
    for sha1 in pool.map(worker, list_of_sha1s):
        print("Done SHA1: %s" % sha1)