Modifica della configurazione a runtime per PySpark

Stavo cercando di implementare un indice Faiss addestrato su PySpark e di fare una ricerca distribuita. Quindi l'intero processo include:

1. Pre-processo
2. Carica l'indice Faiss(~15G) e fai la ricerca Faiss
3. Post-processo e scrivere a HDFS

Ho impostato le CPU per ogni attività come 10 (spark.task.cpus=10) per poter effettuare la ricerca multithread. Ma il passo 1 e il passo 3 possono utilizzare solo 1 CPU per ogni attività. Al fine di utilizzare tutte le CPU che voglio impostare spark.task.cpus=1prima dei passi 1 e 3. Ho provato a impostare il metodo di RuntimeConfigma sembra che abbia bloccato il mio programma. Qualche consiglio su come cambiare la configurazione a runtime o come ottimizzare questo problema?

Esempio di codice:

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)