|
1 | 1 | # tag::PRIMES_PROC_TOP[] |
2 | 2 | from time import perf_counter |
3 | | -from typing import Tuple, List, NamedTuple |
4 | | -from multiprocessing import Process, SimpleQueue # <1> |
| 3 | +from typing import Tuple, NamedTuple |
| 4 | +from multiprocessing import Process, SimpleQueue, cpu_count # <1> |
5 | 5 | from multiprocessing import queues # <2> |
| 6 | +import sys |
6 | 7 |
|
7 | 8 | from primes import is_prime, NUMBERS |
8 | 9 |
|
9 | 10 | class Result(NamedTuple): # <3> |
10 | 11 | flag: bool |
11 | 12 | elapsed: float |
12 | 13 |
|
13 | | -JobQueue = queues.SimpleQueue[Tuple[int, Result]] # <4> |
| 14 | +JobQueue = queues.SimpleQueue[int] # <4> |
| 15 | +ResultQueue = queues.SimpleQueue[Tuple[int, Result]] # <5> |
14 | 16 |
|
15 | | -def check(n: int) -> Result: # <5> |
| 17 | +def check(n: int) -> Result: # <6> |
16 | 18 | t0 = perf_counter() |
17 | 19 | res = is_prime(n) |
18 | 20 | return Result(res, perf_counter() - t0) |
19 | 21 |
|
20 | | -def job(n: int, results: JobQueue) -> None: # <6> |
21 | | - results.put((n, check(n))) # <7> |
| 22 | +def worker(jobs: JobQueue, results: ResultQueue) -> None: # <7> |
| 23 | + while n := jobs.get(): # <8> |
| 24 | + result = check(n) # <9> |
| 25 | + results.put((n, result)) # <10> |
22 | 26 | # end::PRIMES_PROC_TOP[] |
23 | 27 |
|
24 | 28 | # tag::PRIMES_PROC_MAIN[] |
25 | 29 | def main() -> None: |
| 30 | + if len(sys.argv) < 2: # <1> |
| 31 | + workers = cpu_count() |
| 32 | + else: |
| 33 | + workers = int(sys.argv[1]) |
| 34 | + |
26 | 35 | t0 = perf_counter() |
27 | | - results: JobQueue = SimpleQueue() # <1> |
28 | | - workers: List[Process] = [] # <2> |
| 36 | + jobs: JobQueue = SimpleQueue() # <2> |
| 37 | + results: ResultQueue = SimpleQueue() |
| 38 | + |
| 39 | + print(f'Checking {len(NUMBERS)} numbers with {workers} processes:') |
29 | 40 |
|
30 | | - for n in NUMBERS: |
31 | | - worker = Process(target=job, args=(n, results)) # <3> |
32 | | - worker.start() # <4> |
33 | | - workers.append(worker) # <5> |
| 41 | + for n in NUMBERS: # <3> |
| 42 | + jobs.put(n) |
34 | 43 |
|
35 | | - for _ in workers: # <6> |
| 44 | + for _ in range(workers): |
| 45 | + proc = Process(target=worker, args=(jobs, results)) # <4> |
| 46 | + proc.start() # <5> |
| 47 | + jobs.put(0) # <6> |
| 48 | + |
| 49 | + while True: |
36 | 50 | n, (prime, elapsed) = results.get() # <7> |
37 | 51 | label = 'P' if prime else ' ' |
38 | | - print(f'{n:16} {label} {elapsed:9.6f}s') |
39 | | - |
| 52 | + print(f'{n:16} {label} {elapsed:9.6f}s') # <8> |
| 53 | + if jobs.empty(): # <9> |
| 54 | + break |
40 | 55 |
|
41 | 56 | time = perf_counter() - t0 |
42 | 57 | print('Total time:', f'{time:0.2f}s') |
|
0 commit comments