PySpark cluster and session sizing

[cjer]

We are using autofaiss to build a pipeline for creating indexes of different sizes. We have embedding datasets ranging from 1M to 1B examples. Currently all the embeddings we use are 768 dims.

We are running on a GCP dataproc cluster. What are the recommended cluster machine sizes and session configs for the master and worker nodes to best utilize autofaiss' parallelization code for all the different stages (training, adding embeddings and merging)?