在运行时更改 PySpark 的配置

Question

我试图将经过训练的 Faiss 索引部署到 PySpark 并进行分布式搜索。所以整个过程包括：

1. 预处理
2. 加载 Faiss 索引（~15G）并进行 Faiss 搜索
3. 后处理并写入 HDFS

我将每个任务的 CPU 设置为 10 (spark.task.cpus=10) 以便进行多线程搜索。但是步骤 1 和步骤 3 每个任务只能使用 1 个 CPU。为了利用所有 CPU，我想在第 1 步和第 3 步之前设置spark.task.cpus=1。我尝试了RuntimeConfig 的设置方法，但它似乎让我的程序卡住了。关于如何在运行时更改配置或如何优化此问题的任何建议？

代码示例：

def load_and_search(x, model_path):
    faiss_idx = faiss.read_index(model_path)
    q_vec = np.concatenate(x)
    _, idx_array = faiss_idx.search(q_vec, k=10)
    return idx_array


data = sc.textFile(input_path)

# preprocess, only used one cpu per task
data = data.map(lambda x: x)

# load faiss index and search, used multiple cpus per task
data = data.mapPartitioins(lambda x: load_and_search(x, model_path))

# postprocess and write, one cpu per task
data = data.map(lambda x: x).saveAsTextFile(result_path)

Answer 1

另一种想法：在第 1 步和第 3 步中使用 mapPartitions。然后，在每个 worker 中使用多处理池来并行映射分区中的项目。这样，您可以使用分配给工作人员的所有 CPU 而无需更改配置（我不知道这是否可能）。

伪代码：

def item_mapper(item):
    return ...

def partition_mapper(partition):
    with mp.Pool(processes=10) as pool:
        yield from pool.imap(item_mapper, partition)

rdd.mapPartitions(partition_mapper)

Answer 2

您可以通过以下方式更改 sparkContext 属性：

conf = sc._conf.setAll([('spark.task.cpus','1')])
sc._conf.getAll()
data = data.map(lambda x: x)

conf = sc._conf.setAll([('spark.task.cpus','10')])
sc._conf.getAll()
# load faiss index and search, used multiple cpus per task
data = data.mapPartitioins(lambda x: load_and_search(x, model_path))

conf = sc._conf.setAll([('spark.task.cpus','1')])
sc._conf.getAll()
# postprocess and write, one cpu per task
data = data.map(lambda x: x).saveAsTextFile(result_path)

getAll() 可以去掉，添加只是为了检查当前配置。