title: Performance tuning LDA in spark
body: <p>I'm implementing an LDA model in spark (via the Scala API), and testing the model with varying numbers of topics. It seems to work fine in general, but encounters intermittent task failures that I'm pretty sure are related to memory issues. The relevant parts of my current code are below.</p>
  
  <p>Note that I'm loading my data from a a text dump of an RDD, in which each document is a sparse mllib vector. So an example line from my <code>LDA_vectors</code> file looks like this:</p>
  
  <pre><code>(7066346,(112312,[1,3,5,7,...],[2.0,57.0,10.0,2.0,...]))
  </code></pre>
  
  <p>This is the standard mllib sparse format, and can be read as</p>
  
  <pre><code>(document_id,(vocabulary_size,[term_id_1, term_id_2...], [term_1_freq,term_2, freq,...]))
  </code></pre>
  
  <p>So the <code>parse</code> function handles reading that into an RDD:</p>
  
  <pre><code>import org.apache.spark.mllib.clustering.{LDA, DistributedLDAModel, LocalLDAModel}
  import org.apache.spark.mllib.linalg.{Vector, Vectors}
  import org.apache.spark.rdd.RDD
  import java.io._
  
  def parse(rdd: RDD[String]): RDD[(Long, Vector)] = {
    val pattern: scala.util.matching.Regex = "\\(([0-9]+),(.*)\\)".r
    rdd .map{
    case pattern(k, v) =&gt; (k.toLong, Vectors.parse(v))
    }
   }
  
  val text = sc.textFile("/path/to/LDA_vectors")
  val docsWithFeatures = parse(text).repartition(192).cache()
  </code></pre>
  
  <p>And then I run a loop over different numbers of topics. Note that the code block for saving the word-document matrix to file follows a methods described <a href="https://stackoverflow.com/questions/29649904/save-spark-org-apache-spark-mllib-linalg-matrix-to-a-file">here</a>:</p>
  
  <pre><code>for (n_topics &lt;- Range(10,301,5) ){
  
    val s = n_topics.toString
    val lda = new LDA().setK(n_topics).setMaxIterations(20) 
    val ldaModel = lda.run(docsWithFeatures)
    val distLDAModel = ldaModel.asInstanceOf[DistributedLDAModel]
  
    // write some model info to file
    val outfile = new File(s"model_summary_$s")
    @transient val bw = new BufferedWriter(new FileWriter(outfile))
    bw.write("topicConcentration:"+"\t"+distLDAModel.topicConcentration+"\n")
    bw.write("docConcentration:"+"\t"+distLDAModel.docConcentration(0)+"\n")
    bw.write("LL:"+"\t"+distLDAModel.logLikelihood+"\n")
    bw.close()
  
    // convert to Distributed model so we can get our topic data out 
    val distLDAModel = ldaModel.asInstanceOf[DistributedLDAModel]
  
    // Save the document-topic matrix to file
    distLDAModel.topicDistributions.saveAsTextFile(s"doc_topic_$s")
  
    // this saves the word-topic matrix to file
    val topic_mat = distLDAModel.topicsMatrix
    val localMatrix: List[Array[Double]] = topic_mat.transpose.toArray.grouped(topic_mat.numCols).toList
    val lines: List[String] = localMatrix.map(line =&gt; line.mkString(" "))
    val outfile2 = new File(s"artist_topic_$s")
    @transient val bw2 = new BufferedWriter(new FileWriter(outfile2))
    for (line &lt;- lines) bw2.write(line+"\n")
    bw2.close()
  }
  </code></pre>
  
  <p>Ok. So this all works fine, but as I said I started running into task failures, with progressively greater likelihood I increase the number of topics. And I think these are caused by memory issues, leading me to wonder how I can performance tune LDA in spark.</p>
  
  <p>I'm running on Google Cloud Dataproc, so my resources are flexible, but I realize tgat I don't understand the internals of Spark's LDA enough to know how best to optimize performance here.</p>
  
  <p>My one attempt so far is what I do in this line:</p>
  
  <pre><code>val docsWithFeatures = parse(text).repartition(192).cache()
  </code></pre>
  
  <p>Here I repartition my document RDD into 192 partitions (for this example I was running spark on 48 cores, so used the rule of thumb of 4*n_cores) and cache it. This is reasonable if I was doing, say, repeated maps on the RDD, but I'm not certain if/how it will help performance here. What else can I do here?</p>
  
  <p>To facilitate any answers, here are some summary stats on my corpus:</p>
  
  <ul>
  <li>Documents: 166,784</li>
  <li>Vocabulary size (number of unique terms): 112,312</li>
  <li>Total tokens: 4,430,237,213</li>
  </ul>
  
  <p>Maybe my large number of tokens is the main issue here, and I just need to increase my per-task memory, i.e. by increasing available memory and reducing the number of executors. But of course this depends on how exact Spark LDA is working under the hood. See, e.g., my previous question <a href="https://stackoverflow.com/questions/34140667/google-cloud-dataproc-configuration-issues/34141695#34141695">here</a>.</p>