Eu estava a tentar implementar um índice Faiss treinado para o PySpark e fazer uma pesquisa distribuída. Então todo o processo inclui:
- Pré-processo
- Carregar Índice de Faiss(~15G) e fazer Pesquisa de Faiss
- Pós-processar e escrever para o HDFS
Eu defino CPUs por tarefa como 10 (spark.task.cpus=10
) a fim de fazer uma pesquisa multi-tarefa. Mas o passo 1 e o passo 3 só podem utilizar 1 CPU por tarefa. Para utilizar todas as CPUs que eu quero definir spark.task.cpus=1
antes dos passos 1 e 3. Eu tentei definir o método, RuntimeConfig
mas parece que isso fez meu programa ficar preso. Algum conselho sobre como alterar a configuração em tempo de execução ou como otimizar este problema?
Exemplo 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)