Python并行计算multiprocessing

===== Python 并行计算 multiprocessing =====
===== 使用多进程绕过 GIL 实现真正的并行计算 =====

"""
multiprocessing 模块通过创建多个进程实现真正的并行计算。
每个进程拥有独立的 Python 解释器和内存空间，绕过了 GIL 限制。

对于 CPU 密集型任务，多进程可以充分利用多核 CPU。
"""

import multiprocessing as mp
import time
import random
import math
import os


# ===== 1. Pool.map：最简单的并行映射 =====

def cpu_intensive_task(x: float) -> float:
"""
一个 CPU 密集型任务，用于并行执行。
注意：函数必须在模块顶层定义，才能被子进程序列化。
"""
result = 0.0
for i in range(500000):
result += math.sin(x + i) * math.cos(x - i)
return result


def demo_pool_map():
"""演示 Pool.map 的基本用法"""
print("===== Pool.map 示例 =====")
data = [random.random() * 10 for _ in range(16)]

start = time.perf_counter()
with mp.Pool(processes=4) as pool:
results = pool.map(cpu_intensive_task, data)
elapsed = time.perf_counter() - start

print(f" 处理了 {len(data)} 个任务，耗时: {elapsed:.4f}s")
print(f" 结果 (前3): {results[:3]}")


# ===== 2. Pool.starmap：传递多个参数 =====

def multi_arg_task(a: float, b: float, c: float) -> float:
"""接受多个参数的函数"""
return math.sqrt(a * a + b * b) * math.sin(c)


def demo_starmap():
"""用 starmap 传递多参数"""
print("===== Pool.starmap 示例 =====")
args = [(1.0, 2.0, i) for i in range(10)]

with mp.Pool(processes=4) as pool:
results = pool.starmap(multi_arg_task, args)

print(f" starmap 结果: {results}")


# ===== 3. Pool.imap_unordered：无序返回 =====

def variable_time_task(x: int) -> int:
"""模拟耗时不同的任务"""
delay = random.uniform(0.1, 0.5)
time.sleep(delay)
return x * x


def demo_imap_unordered():
"""
使用 imap_unordered 无需等待前面的任务完成。
适合各任务执行时间差异很大的场景。
"""
print("===== Pool.imap_unordered 示例 =====")
tasks = range(20)

with mp.Pool(processes=4) as pool:
# 结果按完成顺序返回，而非提交顺序
for result in pool.imap_unordered(variable_time_task, tasks):
print(f" 完成: {result}", end="")
print()


# ===== 4. Process + Queue：进程间通信 =====

def producer(queue: mp.Queue, n: int):
"""
生产者进程：生成数据放入队列。
Queue 是进程安全的 FIFO 队列。
"""
for i in range(n):
item = random.randint(1, 100)
queue.put(item)
print(f"[生产者] 放入: {item}")
time.sleep(0.05)
queue.put(None) # 发送结束信号


def consumer(queue: mp.Queue, result_dict: dict):
"""
消费者进程：从队列取出数据并处理。
使用 Manager 共享字典来收集结果。
"""
total = 0
count = 0
while True:
item = queue.get()
if item is None: # 收到结束信号
break
total += item
count += 1
print(f"[消费者] 取出: {item}, 累计: {total}")
time.sleep(0.1)
result_dict['total'] = total
result_dict['count'] = count


def demo_process_queue():
"""演示 Process + Queue 的生产者消费者模式"""
print("===== Process + Queue 示例 =====")

queue = mp.Queue()
manager = mp.Manager()
result_dict = manager.dict() # 进程安全的共享字典

prod = mp.Process(target=producer, args=(queue, 10))
cons = mp.Process(target=consumer, args=(queue, result_dict))

prod.start()
cons.start()
prod.join()
cons.join()

print(f" 消费结果: 总数 {result_dict['total']}, 共 {result_dict['count']} 项")


# ===== 5. 共享内存：ctypes 共享变量 =====

def increment_counter(counter: mp.Value, lock: mp.Lock, n: int):
"""
使用共享内存中的计数器。
Lock 确保对共享变量的访问是线程安全的。
"""
for _ in range(n):
with lock:
counter.value += 1


def demo_shared_memory():
"""
演示进程间通过共享内存传递数据。
mp.Value 和 mp.Array 可以在进程间安全共享。
"""
print("===== 共享内存示例 =====")

counter = mp.Value('i', 0) # 'i' 表示 C int 类型
lock = mp.Lock()

processes = []
for _ in range(4):
p = mp.Process(target=increment_counter, args=(counter, lock, 1000))
processes.append(p)
p.start()

for p in processes:
p.join()

print(f" 最终计数器值: {counter.value} (预期: 4000)")


# ===== 6. apply_async 异步回调 =====

def async_task(x: int) -> int:
"""异步执行的耗时任务"""
time.sleep(random.uniform(0.1, 0.3))
return x * 10


def callback(result):
"""任务完成后的回调函数"""
print(f" [回调] 任务结果: {result}")


def error_callback(error):
"""任务失败时的错误回调"""
print(f" [错误] 任务失败: {error}")


def demo_apply_async():
"""
使用 apply_async 进行异步任务提交。
非阻塞方式提交任务，通过回调收集结果。
"""
print("===== apply_async 异步回调示例 =====")

with mp.Pool(processes=4) as pool:
results = []
for i in range(8):
result = pool.apply_async(
async_task,
args=(i,),
callback=callback,
error_callback=error_callback
)
results.append(result)

# 等待所有任务完成
for r in results:
r.wait()

# 或者用 get 获取所有结果
# final_results = [r.get() for r in results]
# print(f"最终结果: {final_results}")

print(" [完成] 所有异步任务已完成")


# ===== 7. 进程池上下文管理器与资源管理 =====

def demo_pool_context():
"""
使用上下文管理器确保资源释放。
推荐使用 with 语句管理 Pool 生命周期。
"""
print("===== Pool 资源管理 =====")

# chunksize 参数控制每个子进程一次获取的任务数
# 对于大量小任务，增大 chunksize 减少通信开销
with mp.Pool(processes=4) as pool:
tasks = range(100)
results = pool.map(
cpu_intensive_task,
tasks,
chunksize=10 # 每个进程一次获取 10 个任务
)

print(f" 共处理 {len(results)} 个任务")


def main():
"""主函数，依次运行所有示例"""
print(f"当前进程 PID: {os.getpid()}")
print(f"CPU 核心数: {mp.cpu_count()}")
print()

demo_pool_map()
print()
demo_starmap()
print()
demo_imap_unordered()
print()
demo_process_queue()
print()
demo_shared_memory()
print()
demo_apply_async()
print()
demo_pool_context()

print("\n===== 所有并行计算示例完成 =====")


if __name__ == '__main__':
# multiprocessing 在 Windows 上需要此保护
mp.freeze_support()
main()