Intentaba desplegar un índice Faiss entrenado en PySpark y hacer una búsqueda distribuida. Así que todo el proceso incluye:
- Pre-proceso
- Cargar el índice Faiss (~15G) y hacer la búsqueda Faiss
- Post-procesar y escribir en el HDFS
Establecí las CPUs por tarea como 10 (spark.task.cpus=10
) para hacer la búsqueda multi-hilo. Pero el paso 1 y el paso 3 sólo pueden utilizar 1 CPU por tarea. Para poder utilizar todas las CPUs quiero establecer spark.task.cpus=1
antes de los pasos 1 y 3. He probado el método de configuración RuntimeConfig
pero parece que hace que mi programa se atasque. ¿Algún consejo sobre cómo cambiar la configuración en tiempo de ejecución o cómo optimizar este problema?
Ejemplo de código:
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)