Stavo cercando di implementare un indice Faiss addestrato su PySpark e di fare una ricerca distribuita. Quindi l'intero processo include:
- Pre-processo
- Carica l'indice Faiss(~15G) e fai la ricerca Faiss
- 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=1
prima dei passi 1 e 3. Ho provato a impostare il metodo di RuntimeConfig
ma 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)