使用 joblibspark 并行 Python：如何均匀分配作业？

Question

我有一个项目，其中 joblib 在一台计算机上运行良好，它有效地将功能发送到不同的内核。

现在我被分配在 Databricks 集群上做同样的事情。我今天尝试了很多方法，但结果的问题是作业并没有分散到每个计算节点上。我有 4 个执行器，我设置

n_jobs

=6，但是当我发送 4 个作业时，其中一些堆积在同一节点上，导致节点未使用。这是 Databricks Spark UI 的图片：

。有时，当我尝试这样做时，我会在一个节点上单独运行 1 个作业，而其余所有作业都堆积在一个节点上。

在

joblib

和

joblibspark

文档中，我看到参数

batch_size

指定发送到给定节点的任务数量。即使我将其设置为 1，我也会遇到同样的问题，节点未使用。

    from joblib import Parallel, delayed
    from joblibspark import register_spark
        
    register_spark()
        
    output = Parallel(backend="spark", n_jobs=6,
             verbose=config.JOBLIB_VERBOSE, batch_size=1)(
             delayed(fit_one)
             (x, model_data=model_data, dlmodel=dlmodel,
             outdir=outdir,  frac=sample_p,
             score_type=score_type,
             save=save,
             verbose=verbose) for x in ZZ)

我一整天都在研究这个问题，尝试各种后端和设置组合。我错过了什么？

Answer 1

spark 和 joblib Spark 的更新在 2022 年中期解决了这个问题。

https://github.com/joblib/joblib-spark/issues/34