Python使用apply_async和imap时遇到的区别

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99

# import multiprocessing as mp
# this module is a fork of official multiprocessing and I believe is faster
import multiprocess as mp

# pool
# pool = mp.Pool(mp.cpu_count())
pool = mp.Pool(16)

'''
multiprocessing a for-range function with tqdm showing progess bar, https://github.com/tqdm/tqdm/issues/484
using imap_unordered to start process in a pool and get returned value asynchronously
- imap_unordered only takes one argument(the iterator), if you need multiple arguments, you have to wrap all arguments into one new iterator
# myRange: range of for loop, e.g. range(2001,2020)
# MyFunc: function to by multiprocessed, takes returned value of myRange as the only argument
''' 
def forloop_mp(myRange, myFunc):
    if __name__ == '__main__':
        # start a pool
        # with mp.Pool(processes=mp.cpu_count()) as pool:
        with mp.Pool(processes=16) as pool:
         
            # to store returned values
            results = []
            # run iteration with tqdm
            # expecting a message to print and the returned result of each iteration
            for (msg, result) in tqdm(pool.imap_unordered(myFunc, myRange), total=len(myRange)):
                # if you want to print returned message
                print(msg)
                results.append(result)
            # stop receiving new tasks
            pool.close()
            # block the method until the process whose join() method is called terminates
            pool.join()
            # exiting the 'with'-block has stopped the pool, dont terminate it manually
            # pool.terminate()
            
        return results
    else:
        raise "Not in Jupyter Notebook"

'''
multiprocessing a for-range function with tqdm showing progess bar, https://github.com/tqdm/tqdm/issues/484
using apply_async to start process in a pool (another way is to use imap_unordered)
# myRange: range of for loop, e.g. range(2001,2020)
# MyFunc: function to by multiprocessed
# params: apply_async can receive multiple args as a list, or you can simple use args=(i, ), which will be the same as imap_unordered
''' 
def forloop_mp_applyasync(myRange, myFunc, params):
    if __name__ == '__main__':
        # start a pool
        # with mp.Pool(processes=mp.cpu_count()) as pool:
        with mp.Pool(processes=16) as pool:
            
            # progress bar and result list
            pbar = tqdm(total=len(myRange))
            results = [None] * len(myRange)  # result list of correct size

            # callback function to update pbar
            # expecting a message to print and the returned result of each iteration    
            def update(msg, result):
                pbar.update()
                # if you want to print returned value
                print(msg)
                results.append(result)

            for i in myRange:
                pool.apply_async(myFunc, args=(i, params) , callback=update)

            # stop receiving new tasks
            pool.close()
            # block the method until the process whose join() method is called terminates
            pool.join()
            # stop updating pbar
            pbar.close()
            # exiting the 'with'-block has stopped the pool, dont terminate it manually
            # pool.terminate()
            
        return results
    else:
        raise "Not in Jupyter Notebook"

''' 
clip raster by mask layer
'''
def clip_raster_by_mask_layer(inputfile, masklayer, outputfile):
    # warp
    gdal.Warp(outputfile, inputfile, format = 'GTiff', cutlineDSName = masklayer)
    print('clipping ' + inputfile)
    return ('clipped ' + inputfile)
'''
the function to be multi processed
'''
def clip_LC_rasters(i):
    input_file = LC_PATH + 'Aligned_LandCover' + str(i) + '_4326.tif'
    output_file = LC_PATH + 'Clipped_LandCover' + str(i) + '_4326.tif'
    mask_layer = BOUND_PATH + 'ChinaBound-FromGov-1000000.shp'
    return clip_raster_by_mask_layer(input_file, mask_layer, output_file)
# call
forloop_mp(range(2001,2020), clip_LC_rasters)