Ich habe versucht, einen trainierten Faiss-Index in PySpark zu implementieren und eine verteilte Suche durchzuführen. Der gesamte Prozess umfasst also:
- Vor-Prozess
- Faiss-Index(~15G) laden und Faiss-Suche durchführen
- Post-Processing und Schreiben in HDFS
Ich stelle die CPUs pro Task auf 10 (spark.task.cpus=10) ein, um eine Multi-Thread-Suche durchzuführen. In Schritt 1 und Schritt 3 kann jedoch nur 1 CPU pro Task verwendet werden. Um alle CPUs auszunutzen, die ich spark.task.cpus=1vor Schritt 1 und 3 einstellen möchte. Ich habe die Set-Methode ausprobiert, RuntimeConfigaber es scheint, dass mein Programm dadurch hängen geblieben ist. Haben Sie einen Rat, wie Sie die Konfiguration zur Laufzeit ändern oder dieses Problem optimieren können?
Code-Beispiel:
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)
