<commit_before>package edu.mit.csail.db.ml.modeldb.client

import java.util.UUID

import com.twitter.finagle.Thrift
import com.twitter.util.Await
import edu.mit.csail.db.ml.modeldb.client.SyncingStrategy.SyncingStrategy
import edu.mit.csail.db.ml.modeldb.client.event.{ExperimentEvent, ExperimentRunEvent, ModelDbEvent, ProjectEvent}
import org.apache.spark.ml.FeatureTracker
import edu.mit.csail.db.ml.modeldb.evaluation.Timer
import modeldb.ModelDBService.FutureIface
import modeldb._
import org.apache.spark.ml.classification.LogisticRegressionModel
import org.apache.spark.ml.regression.LinearRegressionModel
import org.apache.spark.ml.util.MLReader
import org.apache.spark.ml.{SyncableCrossValidator, SyncableEstimator, SyncablePipeline, Transformer}
import org.apache.spark.sql.DataFrame

import scala.collection.mutable
import scala.collection.mutable.ArrayBuffer

/**
  * Represents the strategy used by the ModelDb when deciding how to sync.
  * Eager means that it will sync right after an event is buffered.
  * Manual means it will not sync until the user calls the sync() function.
  */
object SyncingStrategy extends Enumeration {
  type SyncingStrategy = Value
  val Eager, Manual = Value
}

abstract class ProjectConfig(val id: Int)
case class ExistingProject(override val id: Int) extends ProjectConfig(id)
case class NewOrExistingProject(name: String, author: String, description: String) extends ProjectConfig(-1)

abstract class ExperimentConfig(val id: Int)
case class ExistingExperiment(override val id: Int) extends ExperimentConfig(-1)
case class DefaultExperiment() extends ExperimentConfig(-1)
case class NewOrExistingExperiment(name: String, description: String) extends ExperimentConfig(-1)

abstract class ExperimentRunConfig(val id: Int)
case class ExistingExperimentRun(override val id: Int) extends ExperimentRunConfig(id)
case class NewExperimentRun(description: String="") extends ExperimentRunConfig(-1)

/**
  * This is the Syncer that is responsible for storing events in the ModelDB.
  *
  * The shouldCountRows parameter is a boolean indicating whether ModelDB should count the number of rows in each
  * DataFrame and store the count in the database. Counting the number of rows requires a full sequential scan of the
  * DataFrame, which is a performance intensive operation. If shouldCountRows is set to true, then the rows will
  * be counted and stored in the database. if shouldCountRows is set to false, then ModelDB will simply store -1 as
  * the number of rows in each DataFrame. By default, we count the number of rows.
  */
class ModelDbSyncer(hostPortPair: Option[(String, Int)] = Some("localhost", 6543),
                    syncingStrategy: SyncingStrategy = SyncingStrategy.Eager,
                    projectConfig: ProjectConfig,
                    experimentConfig: ExperimentConfig = new DefaultExperiment,
                    experimentRunConfig: ExperimentRunConfig,
                    val shouldCountRows: Boolean = true,
                    val shouldStoreGSCVE: Boolean = true,
                    val shouldStoreSpecificModels: Boolean = true) {
  /**
    * This is a helper class that will constitute the entries in the buffer.
    * @param event - The event in the buffer.
    * @param postSync - The function to execute after the event has been sycned to the ModelDB.
    */
  case class BufferEntry(event: ModelDbEvent, postSync: (ModelDbSyncer, ModelDbEvent) => Unit)

  /**
    * As we get events, we buffer them in this ArrayBuffer.
    */
  var buffered: ArrayBuffer[BufferEntry] = ArrayBuffer()

  /**
    * This is the Thrift client that is responsible for talking to the ModelDB server.
    */
  private val client: Option[FutureIface] = hostPortPair match {
    case Some((host, port)) => Some(Thrift.client.newIface[ModelDBService.FutureIface](s"$host:$port"))
    case None => None
  }<commit_msg>Default should in Spark client to false<commit_after>package edu.mit.csail.db.ml.modeldb.client

import java.util.UUID

import com.twitter.finagle.Thrift
import com.twitter.util.Await
import edu.mit.csail.db.ml.modeldb.client.SyncingStrategy.SyncingStrategy
import edu.mit.csail.db.ml.modeldb.client.event.{ExperimentEvent, ExperimentRunEvent, ModelDbEvent, ProjectEvent}
import org.apache.spark.ml.FeatureTracker
import edu.mit.csail.db.ml.modeldb.evaluation.Timer
import modeldb.ModelDBService.FutureIface
import modeldb._
import org.apache.spark.ml.classification.LogisticRegressionModel
import org.apache.spark.ml.regression.LinearRegressionModel
import org.apache.spark.ml.util.MLReader
import org.apache.spark.ml.{SyncableCrossValidator, SyncableEstimator, SyncablePipeline, Transformer}
import org.apache.spark.sql.DataFrame

import scala.collection.mutable
import scala.collection.mutable.ArrayBuffer

/**
  * Represents the strategy used by the ModelDb when deciding how to sync.
  * Eager means that it will sync right after an event is buffered.
  * Manual means it will not sync until the user calls the sync() function.
  */
object SyncingStrategy extends Enumeration {
  type SyncingStrategy = Value
  val Eager, Manual = Value
}

abstract class ProjectConfig(val id: Int)
case class ExistingProject(override val id: Int) extends ProjectConfig(id)
case class NewOrExistingProject(name: String, author: String, description: String) extends ProjectConfig(-1)

abstract class ExperimentConfig(val id: Int)
case class ExistingExperiment(override val id: Int) extends ExperimentConfig(-1)
case class DefaultExperiment() extends ExperimentConfig(-1)
case class NewOrExistingExperiment(name: String, description: String) extends ExperimentConfig(-1)

abstract class ExperimentRunConfig(val id: Int)
case class ExistingExperimentRun(override val id: Int) extends ExperimentRunConfig(id)
case class NewExperimentRun(description: String="") extends ExperimentRunConfig(-1)

/**
  * This is the Syncer that is responsible for storing events in the ModelDB.
  *
  * The shouldCountRows parameter is a boolean indicating whether ModelDB should count the number of rows in each
  * DataFrame and store the count in the database. Counting the number of rows requires a full sequential scan of the
  * DataFrame, which is a performance intensive operation. If shouldCountRows is set to true, then the rows will
  * be counted and stored in the database. if shouldCountRows is set to false, then ModelDB will simply store -1 as
  * the number of rows in each DataFrame. By default, we count the number of rows.
  */
class ModelDbSyncer(hostPortPair: Option[(String, Int)] = Some("localhost", 6543),
                    syncingStrategy: SyncingStrategy = SyncingStrategy.Eager,
                    projectConfig: ProjectConfig,
                    experimentConfig: ExperimentConfig = new DefaultExperiment,
                    experimentRunConfig: ExperimentRunConfig,
                    val shouldCountRows: Boolean = false,
                    val shouldStoreGSCVE: Boolean = false,
                    val shouldStoreSpecificModels: Boolean = false) {
  /**
    * This is a helper class that will constitute the entries in the buffer.
    * @param event - The event in the buffer.
    * @param postSync - The function to execute after the event has been sycned to the ModelDB.
    */
  case class BufferEntry(event: ModelDbEvent, postSync: (ModelDbSyncer, ModelDbEvent) => Unit)

  /**
    * As we get events, we buffer them in this ArrayBuffer.
    */
  var buffered: ArrayBuffer[BufferEntry] = ArrayBuffer()

  /**
    * This is the Thrift client that is responsible for talking to the ModelDB server.
    */
  private val client: Option[FutureIface] = hostPortPair match {
    case Some((host, port)) => Some(Thrift.client.newIface[ModelDBService.FutureIface](s"$host:$port"))
    case None => None
  }
{'fraction_non_alphanumeric': 0.07529230373891815, 'fraction_numerical': 0.002569703199280483, 'mean_word_length': 4.131179960448253, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 9, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1379301', 'n_tokens_mistral': 2154, 'n_tokens_neox': 2073, 'n_words': 748}