Python多进程与GPU并行计算全解：从Multiprocessing、Ray到PyTorch DDP实战

1
import multiprocessing as mp
2

3
if __name__ == "__main__":
4
    # 强制在程序入口处设置安全的进程派生机制
5
    mp.set_start_method("spawn", force=True)

1
import time
2
import math
3
from multiprocessing import Pool
4

5
def complex_simulation_task(param_seed):
6
    # 模拟高耗时复杂数值蒙特卡洛积分
7
    result = 0.0
8
    for i in range(1, 100000):
9
        result += math.sin(param_seed * i) * math.cos(i)
10
    return param_seed, result
11

12
def run_academic_pipeline():
13
    total_samples = 50000
14
    # 根据物理核心数配置合理的工作进程配额
15
    num_workers = mp.cpu_count()
16
    # 合理的分块尺寸能够显著降低进程间通信往返开销
17
    optimal_chunksize = max(1, total_samples // (num_workers * 16))
18

19
    print(f"启动科学计算流水线，核心数 {num_workers}，单块尺寸 {optimal_chunksize}")
20
    start_time = time.time()
21

22
    with Pool(processes=num_workers) as pool:
23
        task_inputs = range(total_samples)
24
        # 利用流式无序迭代器，计算完成一项立即返回一项，绝不堆积内存
25
        for idx, (seed, val) in enumerate(pool.imap_unordered(complex_simulation_task, task_inputs, chunksize=optimal_chunksize)):
26
            if idx % 5000 == 0:
27
                elapsed = time.time() - start_time
28
                print(f"已完成进度 [{idx}/{total_samples}]，耗时 {elapsed:.2f} 秒")
29

30
    print("科学仿真全部平稳计算完毕")
31

32
if __name__ == "__main__":
33
    mp.set_start_method("spawn", force=True)
34
    run_academic_pipeline()

1
import numpy as np
2
from multiprocessing import shared_memory, Process
3

4
def worker_read_matrix(shm_name, shape, dtype, row_slice):
5
    # 子进程接入已经开辟好的操作系统共享内存段
6
    existing_shm = shared_memory.SharedMemory(name=shm_name)
7
    # 基于共享物理内存构建轻量级 NumPy 视图，完全不存在数据拷贝
8
    shared_array = np.ndarray(shape, dtype=dtype, buffer=existing_shm.buf)
9

10
    # 局部数值处理
11
    local_mean = np.mean(shared_array[row_slice])
12
    print(f"进程处理切片 {row_slice}，局部均值计算结果为 {local_mean:.4f}")
13

14
    # 关闭当前子进程对共享内存的句柄引用
15
    existing_shm.close()
16

17
def main_shared_pipeline():
18
    matrix_shape = (10000, 10000)
19
    matrix_dtype = np.float64
20
    total_bytes = int(np.prod(matrix_shape) * np.dtype(matrix_dtype).itemsize)
21

22
    # 主进程申请一块专属物理共享内存空间
23
    shm = shared_memory.SharedMemory(create=True, size=total_bytes)
24
    # 将 NumPy 数组直接映射到该共享内存段
25
    master_array = np.ndarray(matrix_shape, dtype=matrix_dtype, buffer=shm.buf)
26
    master_array[:] = np.random.randn(*matrix_shape)
27

28
    processes = []
29
    # 切分任务并启动多个工作子进程
30
    slice_step = matrix_shape[0] // 4
31
    for i in range(4):
32
        r_slice = slice(i * slice_step, (i + 1) * slice_step)
33
        p = Process(target=worker_read_matrix, args=(shm.name, matrix_shape, matrix_dtype, r_slice))
34
        p.start()
35
        processes.append(p)
36

37
    for p in processes:
38
        p.join()
39

40
    # 主进程彻底释放并销毁操作系统共享内存段
41
    shm.close()
42
    shm.unlink()
43
    print("共享内存高维张量并行计算全流程顺利完成")
44

45
if __name__ == "__main__":
46
    mp.set_start_method("spawn", force=True)
47
    main_shared_pipeline()

1
import ray
2
import numpy as np
3
import time
4

5
# 启动本地 Ray 实例或自动挂接到现有的超算集群
6
ray.init(ignore_reinit_error=True)
7

8
# 定义一个常驻内存的科学仿真状态执行体
9
@ray.remote
10
class SimulationWorkerActor:
11
    def __init__(self, worker_id):
12
        self.worker_id = worker_id
13
        self.state_step = 0
14
        self.history_loss = []
15

16
    def step_simulation(self, shared_data_ref):
17
        # 零拷贝读取大型共享物理参数矩阵
18
        matrix = ray.get(shared_data_ref)
19
        self.state_step += 1
20
        loss_metric = float(np.sum(matrix) * 0.001 / self.state_step)
21
        self.history_loss.append(loss_metric)
22
        return self.worker_id, self.state_step, loss_metric
23

24
def run_ray_academic_cluster():
25
    # 生成大型科学输入数据并放入 Plasma 全局对象池
26
    large_matrix = np.random.randn(5000, 5000)
27
    data_ref = ray.put(large_matrix)
28

29
    # 实例化四个并行的状态执行体
30
    actors = [SimulationWorkerActor.remote(worker_id=i) for i in range(4)]
31

32
    # 向各个执行体异步分发多轮演化计算任务
33
    pending_refs = []
34
    for step in range(5):
35
        for actor in actors:
36
            pending_refs.append(actor.step_simulation.remote(data_ref))
37

38
    # 等待全部异步分布式任务收敛
39
    results = ray.get(pending_refs)
40
    for res in results[:8]:
41
        print(f"节点 {res[0]} 完成第 {res[1]} 步演化，损失评估值为 {res[2]:.6f}")
42

43
    ray.shutdown()
44
    print("Ray 分布式弹性科学调度圆满执行完毕")
45

46
if __name__ == "__main__":
47
    run_ray_academic_cluster()

1
import os
2
import torch
3
import torch.nn as nn
4
import torch.distributed as dist
5
from torch.nn.parallel import DistributedDataParallel as DDP
6
from torch.utils.data import DataLoader, Dataset
7
from torch.utils.data.distributed import DistributedSampler
8

9
class SyntheticPhysicsDataset(Dataset):
10
    def __init__(self, size=10000):
11
        self.data = torch.randn(size, 128)
12
        self.target = torch.randn(size, 1)
13
    def __len__(self):
14
        return len(self.data)
15
    def __getitem__(self, idx):
16
        return self.data[idx], self.target[idx]
17

18
def setup_distributed():
19
    # 从操作系统环境变量中读取由 torchrun 注入的分布式拓扑元数据
20
    dist.init_process_group(backend="nccl")
21
    local_rank = int(os.environ["LOCAL_RANK"])
22
    torch.cuda.set_device(local_rank)
23
    return local_rank
24

25
def cleanup_distributed():
26
    # 安全销毁分布式通信进程组
27
    dist.destroy_process_group()
28

29
def train_academic_ddp():
30
    local_rank = setup_distributed()
31
    global_rank = int(os.environ["RANK"])
32
    world_size = int(os.environ["WORLD_SIZE"])
33

34
    if global_rank == 0:
35
        print(f"成功激活 DDP 集群，总结点显卡规模为 {world_size}")
36

37
    # 构建模型并迁移至对应物理显卡
38
    model = nn.Sequential(
39
        nn.Linear(128, 256),
40
        nn.ReLU(),
41
        nn.Linear(256, 1)
42
    ).to(local_rank)
43

44
    # 包装为分布式数据并行模型
45
    model = DDP(model, device_ids=[local_rank], output_device=local_rank)
46
    criterion = nn.MSELoss()
47
    optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3)
48

49
    dataset = SyntheticPhysicsDataset()
50
    # 极为关键的分布式采样器，确保每张显卡分流互不重叠的数据子集
51
    sampler = DistributedSampler(dataset, num_replicas=world_size, rank=global_rank, shuffle=True)
52
    dataloader = DataLoader(dataset, batch_size=64, sampler=sampler, num_workers=2, pin_memory=True)
53

54
    for epoch in range(3):
55
        # 每一个周期必须显式设置 epoch，以保证数据打乱随机种子的动态演化
56
        sampler.set_epoch(epoch)
57
        model.train()
58
        total_epoch_loss = 0.0
59

60
        for batch_x, batch_y in dataloader:
61
            batch_x = batch_x.to(local_rank, non_blocking=True)
62
            batch_y = batch_y.to(local_rank, non_blocking=True)
63

64
            optimizer.zero_grad()
65
            pred = model(batch_x)
66
            loss = criterion(pred, batch_y)
67
            loss.backward()
68
            optimizer.step()
69
            total_epoch_loss += loss.item()
70

71
        if global_rank == 0:
72
            avg_loss = total_epoch_loss / len(dataloader)
73
            print(f"训练周期 [{epoch + 1}/3] 完成，主节点监控损失值为 {avg_loss:.6f}")
74

75
    cleanup_distributed()
76

77
if __name__ == "__main__":
78
    train_academic_ddp()

1
#!/usr/bin/env python3
2
# ==============================================================================
3
# 学术 GPU 集群算力状态巡检与死锁进程识别工具
4
# 依赖库：pynvml (pip install nvidia-ml-py3) 与 psutil (pip install psutil)
5
# ==============================================================================
6

7
import os
8
import sys
9
import time
10
import psutil
11
try:
12
    import pynvml
13
    HAS_NVML = True
14
except ImportError:
15
    HAS_NVML = False
16

17
def inspect_system_compute_health():
18
    print("========================================================")
19
    print("开始执行学术多核 CPU 与 GPU 并行计算集群环境健康巡检")
20
    print("========================================================")
21

22
    # 1. 检查物理 CPU 与逻辑核心利用率
23
    cpu_cores = psutil.cpu_count(logical=False)
24
    logical_cores = psutil.cpu_count(logical=True)
25
    overall_cpu_util = psutil.cpu_percent(interval=0.5)
26
    vm = psutil.virtual_memory()
27

28
    print(f"物理核心数量 {cpu_cores}，逻辑线程总数 {logical_cores}")
29
    print(f"当前全机 CPU 整体负载 {overall_cpu_util:.1f}%")
30
    print(f"系统物理内存总量 {vm.total / (1024**3):.1f} GB，当前已使用 {vm.percent}%")
31

32
    if not HAS_NVML:
33
        print("[警告] 未检测到 pynvml 库，跳过底层 GPU 显存与通信状态深度巡检")
34
        return
35

36
    pynvml.nvmlInit()
37
    device_count = pynvml.nvmlDeviceGetCount()
38
    print(f"检测到系统中已挂载的独立加速显卡总数为 {device_count} 块")
39

40
    flagged_zombies = []
41

42
    for i in range(device_count):
43
        handle = pynvml.nvmlDeviceGetHandleByIndex(i)
44
        name = pynvml.nvmlDeviceGetName(handle)
45
        mem_info = pynvml.nvmlDeviceGetMemoryInfo(handle)
46
        util = pynvml.nvmlDeviceGetUtilizationRates(handle)
47

48
        used_gb = mem_info.used / (1024**3)
49
        total_gb = mem_info.total / (1024**3)
50
        print("--------------------------------------------------------")
51
        print(f"显卡编号 [{i}] 设备型号 {name}")
52
        print(f"显存使用情况 {used_gb:.2f} GB / {total_gb:.2f} GB ({mem_info.used / mem_info.total * 100:.1f}%)")
53
        print(f"当前 GPU 计算核心利用率 {util.gpu}%，显存读写带宽利用率 {util.memory}%")
54

55
        # 深度检索当前正在霸占显卡显存的所有底层进程句柄
56
        processes = pynvml.nvmlDeviceGetComputeRunningProcesses(handle)
57
        for p in processes:
58
            pid = p.pid
59
            vram_mb = p.usedGpuMemory / (1024**2)
60
            try:
61
                proc = psutil.Process(pid)
62
                p_user = proc.username()
63
                p_name = proc.name()
64
                p_status = proc.status()
65
                # 识别典型的僵尸进程：占用大量显存但计算利用率为零且处于休眠停滞状态
66
                if util.gpu == 0 and vram_mb > 2048 and p_status == psutil.STATUS_SLEEPING:
67
                    flagged_zombies.append((i, pid, p_user, vram_mb))
68
                print(f"  -> 进程 PID {pid} | 属主 {p_user} | 进程名 {p_name} | 显存占用 {vram_mb:.1f} MB | 状态 {p_status}")
69
            except psutil.NoSuchProcess:
70
                print(f"  -> 进程 PID {pid} [句柄残留孤儿，已无法在系统进程表中检索]")
71
                flagged_zombies.append((i, pid, "Unknown", vram_mb))
72

73
    pynvml.nvmlShutdown()
74

75
    print("========================================================")
76
    if flagged_zombies:
77
        print(f"[潜在风险告警] 巡检发现 {len(flagged_zombies)} 个疑似陷入死锁或挂起的显存僵尸进程")
78
        for z in flagged_zombies:
79
            print(f"  显卡 [{z[0]}] 进程 PID {z[1]}，属主 {z[2]}，持续霸占显存 {z[3]:.1f} MB")
80
        print("建议通过 kill -9 指令精准释放被死锁的宝贵科研显存资源")
81
    else:
82
        print("[状态优良] 未发现异常死锁或显存泄漏隐患，计算集群运转健康")
83
    print("========================================================")
84

85
if __name__ == "__main__":
86
    inspect_system_compute_health()