如何将基于multiprocessing的并行代码迁移为PySpark实现？

一、先修复当前multiprocessing代码的问题 #

你的现有代码存在几个错误，导致无法充分利用多核：

1. executor.map(Simulation.simulation_steps(), iterations_generator()) 错误调用了simulation_steps()，应传入函数本身而非执行结果，正确写法是executor.map(Simulation.simulation_steps, iterations_generator())
2. 循环中experiment.simulations.append(results[v])逻辑错误，v就是迭代结果，直接使用v即可，无需取results[v]

修复后的单实验并行代码：

import concurrent.futures
import multiprocessing as mp

def process_simulation(experiment):
    number_of_workers = mp.cpu_count()
    with concurrent.futures.ProcessPoolExecutor(max_workers=number_of_workers) as executor:
        # 传入函数对象而非执行结果
        results = list(executor.map(Simulation.simulation_steps, iterations_generator()))
        experiment.simulations = list(results)

如果要并行处理多个实验，可将实验列表纳入并行逻辑：

def process_single_experiment(experiment):
    # 给每个实验分配合理的核心数，避免资源竞争
    worker_count = max(1, mp.cpu_count() // len(all_experiments))
    with concurrent.futures.ProcessPoolExecutor(max_workers=worker_count) as executor:
        results = list(executor.map(Simulation.simulation_steps, iterations_generator()))
        experiment.simulations = results
    return experiment

# 并行处理所有实验
with concurrent.futures.ProcessPoolExecutor() as executor:
    processed_experiments = list(executor.map(process_single_experiment, all_experiments))

二、迁移到PySpark的步骤 #

1. 核心思路 #

PySpark基于分布式数据集（RDD/DataFrame）实现并行计算，需将实验、模拟任务转化为可序列化的数据集，通过map/flatMap等操作分发到集群节点执行。

2. 代码改造要点 #

（1）确保类与函数可序列化

PySpark要求所有传递到集群的对象支持pickle序列化：

• 移除Experiment和Simulation类中无法序列化的属性（如打开的文件句柄、数据库连接）
• 纯Python编写的类方法默认可序列化，无需额外修改

（2）重构模拟逻辑为纯函数

将依赖类状态的模拟逻辑改为接收参数的纯函数，避免分布式执行时的状态传递问题：

def run_simulation(iteration_params):
    # 实例化Simulation并执行模拟，返回结果
    sim = Simulation()
    return sim.simulation_steps(iteration_params)

（3）PySpark分布式实现方案

方案一：基于RDD处理（灵活适配非结构化任务）

from pyspark.sql import SparkSession

# 初始化SparkSession
spark = SparkSession.builder.appName("SimulationExperiments").getOrCreate()
sc = spark.sparkContext

# 将所有实验转化为可分发的RDD
experiment_rdd = sc.parallelize(all_experiments)

def process_experiment_spark(experiment):
    # 生成当前实验的所有迭代参数
    iterations = list(iterations_generator(experiment))
    # 并行执行当前实验的所有模拟任务
    sim_results = sc.parallelize(iterations).map(run_simulation).collect()
    experiment.simulations = sim_results
    return experiment

# 分布式处理所有实验
processed_experiments_rdd = experiment_rdd.map(process_experiment_spark)
# 收集结果到本地（仅适用于结果量较小的场景）
processed_experiments = processed_experiments_rdd.collect()

# 关闭SparkSession
spark.stop()

方案二：基于DataFrame处理（推荐，适配结构化数据）

如果实验参数和模拟结果是结构化数据，用DataFrame更易管理：

from pyspark.sql import SparkSession
from pyspark.sql.functions import udf, collect_list
from pyspark.sql.types import StructType, StructField, FloatType

# 初始化SparkSession
spark = SparkSession.builder.appName("SimulationExperiments").getOrCreate()

# 1. 将实验列表转化为DataFrame（假设实验包含id和参数字段）
experiment_df = spark.createDataFrame(
    [(exp.id, exp.params) for exp in all_experiments],
    schema=["experiment_id", "params"]
)

# 2. 定义模拟UDF，指定返回结果的Schema
sim_result_schema = StructType([
    StructField("metric1", FloatType()),
    StructField("metric2", FloatType())
])

@udf(returnType=sim_result_schema)
def run_simulation_udf(experiment_params, iteration_id):
    sim = Simulation()
    return sim.simulation_steps(experiment_params, iteration_id)

# 3. 生成迭代参数DataFrame
iterations_df = spark.createDataFrame(
    [(i,) for i in range(num_iterations)],
    schema=["iteration_id"]
)

# 4. 生成所有实验-迭代组合任务
tasks_df = experiment_df.crossJoin(iterations_df)

# 5. 分布式执行模拟
results_df = tasks_df.withColumn(
    "sim_result",
    run_simulation_udf("params", "iteration_id")
)

# 6. 按实验ID聚合模拟结果
aggregated_results_df = results_df.groupBy("experiment_id").agg(
    collect_list("sim_result").alias("simulations")
)

# 7. 将聚合结果转回Experiment类（按需使用）
def row_to_experiment(row):
    exp = Experiment()
    exp.id = row.experiment_id
    exp.params = row.params
    exp.simulations = row.simulations
    return exp

processed_experiments = aggregated_results_df.join(
    experiment_df, on="experiment_id"
).rdd.map(row_to_experiment).collect()

# 关闭SparkSession
spark.stop()

3. 关键注意事项 #

• 确保集群所有节点已安装numpy、pandas等依赖库
• 避免在分布式函数中使用全局变量，所有参数需显式传递
• 若模拟结果数据量较大，不要用collect()拉取到本地，直接在Spark中完成后续分析或写入分布式存储（如HDFS、S3）

内容的提问来源于stack exchange，提问作者Loic Trobas