<commit_before>package xerial.silk.framework

import xerial.silk._
import xerial.core.log.{LoggerFactory, Logger}
import java.util.UUID
import scala.util.Random
import scala.collection.parallel.ParSeq
import scala.collection.{GenTraversable, GenSeq}
import xerial.silk.framework.ops.FlatMapOp
import xerial.silk.framework.ops.FlatMapSeqOp
import xerial.silk.framework.ops.ShuffleReduceOp
import xerial.silk.framework.ops.ShuffleOp
import xerial.silk.framework.ops.SortOp
import xerial.silk.framework.ops.ConcatOp
import xerial.silk.framework.ops.RawSeq
import xerial.silk.framework.ops.ReduceOp
import xerial.silk.framework.ops.FilterOp
import xerial.silk.framework.ops.SizeOp
import xerial.silk.framework.ops.SamplingOp
import xerial.silk.framework.ops.MapOp


trait ClassBoxAPI {
  def classLoader : ClassLoader
}


/**
 * ClassBoxComponent has a role to provide the current ClassBoxID and distribute
 * the ClassBox to cluster nodes.
 */
trait ClassBoxComponent {
  self: SilkFramework =>

  type ClassBoxType <: ClassBoxAPI

  /**
   * Get the current class box id
   */
  def classBoxID : UUID

  /**
   * Retrieve the class box having the specified id
   * @param classBoxID
   * @return
   */
  def getClassBox(classBoxID:UUID) : ClassBoxAPI

}


trait ExecutorAPI {
  def getSlices[A](op: Silk[A]) : GenSeq[SilkFuture[Slice]]
}


trait DefaultExecutor extends ExecutorComponent {
  self : SilkFramework
    with LocalTaskManagerComponent
    with ClassBoxComponent
    with SliceStorageComponent =>

  type Executor = ExecutorImpl
  def executor : Executor = new ExecutorImpl

  class ExecutorImpl extends ExecutorAPI with ExecutorBase {}
}



/**
 * Executor receives a silk, optimize a plan, then submit evaluation tasks to the local task manager
 */
trait ExecutorComponent {
  self : SilkFramework
    with LocalTaskManagerComponent
    with ClassBoxComponent
    with SliceStorageComponent =>

  type Executor <: ExecutorAPI with ExecutorBase
  def executor : Executor

  trait ExecutorBase extends Logger {

    implicit class toGenFun[A,B](f:A=>B) {
      def toF1 : Any=>Any = f.asInstanceOf[Any=>Any]

      def toFlatMap : Any=>SilkSeq[Any] = f.asInstanceOf[Any=>SilkSeq[Any]]
      def toFilter : Any=>Boolean = f.asInstanceOf[Any=>Boolean]
    }

    implicit class toGenFMap[A, B](f:A=>GenTraversable[B]) {
      def toFmap = f.asInstanceOf[Any=>GenTraversable[Any]]
    }

    def run[A](session:Session, silk: Silk[A]): Result[A] = {

      val dataSeq : ParSeq[Seq[A]] = for{
        future <- getSlices(silk)
      }
      yield {
        val slice = future.get
        //debug(s"Fetch slice: $slice in $silk")
        sliceStorage.retrieve(silk.id, slice).asInstanceOf[Seq[A]]
      }

      val result = dataSeq.seq.flatten
      result
    }

    /**
     * Get Stage information
     * @param op
     * @tparam A
     * @return
     */
    def getStage[A](op:Silk[A]) : StageInfo = {
      sliceStorage.getStageInfo(op).map { si =>
        si.status match {
          case StageStarted(ts) => si
          case StageFinished(ts) => si
          case StageAborted(cause, ts) => // TODO restart
            SilkException.error(s"stage of [${op.idPrefix}] has been aborted: $cause")
        }
      } getOrElse {
        // Start a stage for evaluating the Silk
        executor.startStage(op)
      }
    }

    private def startStage[A,Out](op:Silk[A], in:Silk[A], f:Seq[_]=>Out) = {
      val inputStage = getStage(in)
      val N = inputStage.numSlices
      val stageInfo = StageInfo(0, N, StageStarted(System.currentTimeMillis()))
      sliceStorage.setStageInfo(op, stageInfo)
      // TODO append par
      for(i <- (0 until N)) {
        // Get an input slice
        val inputSlice = sliceStorage.get(in.id, i).get
        // Send the slice processing task to a node close to the inputSlice location
        localTaskManager.submit(EvalSliceTask(s"eval slice of ${op}", UUID.randomUUID, classBoxID, op.id, in.id, inputSlice, f.toF1, Seq(inputSlice.nodeName)))
      }
      stageInfo
    }


    private def startReduceStage[A](op:Silk[A], in:Silk[A], reducer:Seq[_] => Any, aggregator:Seq[_]=>Any) = {
      val inputStage = getStage(in)
      val N = inputStage.numSlices

      // The outer reduce task produces only 1 slice
      val stageInfo = StageInfo(0, 1, StageStarted(System.currentTimeMillis()))

      // Evaluate reduce at each slice
      val subStageID = SilkUtil.newUUID
      for(i <- 0 until N) {
        // Get an input slice
        val inputSlice = sliceStorage.get(in.id, i).get
        // Locality-aware job submission
        localTaskManager.submit(ReduceTask(s"reduce each ${op}", UUID.randomUUID, classBoxID, subStageID, in.id, Seq(i), i, reducer, aggregator,  Seq(inputSlice.nodeName)))
      }

      // Determine the number of reducers to use. The default is 1/3 of the number of the input slices
      val R = ((N + (3-1))/ 3.0).toInt
      val W = (N + (R-1)) / R
      info(s"num reducers:$R, W:$W")

      // Reduce the previous sub stage using R reducers
      val aggregateStageID = SilkUtil.newUUID
      for((sliceRange, i) <- (0 until N).sliding(W, W).zipWithIndex) {
        val sliceIndexSet = sliceRange.toIndexedSeq
        localTaskManager.submit(ReduceTask(s"reduce aggregate ${op}", UUID.randomUUID, classBoxID, aggregateStageID, subStageID, sliceIndexSet, i, aggregator, aggregator, Seq.empty))
      }

      // The final aggregate task
      localTaskManager.submit(ReduceTask(s"reduce final of ${op}", UUID.randomUUID, classBoxID, op.id, aggregateStageID, (0 until R).toIndexedSeq, 0, aggregator, aggregator, Seq.empty))

      stageInfo
    }


    private def startShuffleStage[A, K](shuffleOp:ShuffleOp[A, K]) = {
      val inputStage = getStage(shuffleOp.in)<commit_msg>Parallel job submission<commit_after>package xerial.silk.framework

import xerial.silk._
import xerial.core.log.{LoggerFactory, Logger}
import java.util.UUID
import scala.util.Random
import scala.collection.parallel.ParSeq
import scala.collection.{GenTraversable, GenSeq}
import xerial.silk.framework.ops.FlatMapOp
import xerial.silk.framework.ops.FlatMapSeqOp
import xerial.silk.framework.ops.ShuffleReduceOp
import xerial.silk.framework.ops.ShuffleOp
import xerial.silk.framework.ops.SortOp
import xerial.silk.framework.ops.ConcatOp
import xerial.silk.framework.ops.RawSeq
import xerial.silk.framework.ops.ReduceOp
import xerial.silk.framework.ops.FilterOp
import xerial.silk.framework.ops.SizeOp
import xerial.silk.framework.ops.SamplingOp
import xerial.silk.framework.ops.MapOp


trait ClassBoxAPI {
  def classLoader : ClassLoader
}


/**
 * ClassBoxComponent has a role to provide the current ClassBoxID and distribute
 * the ClassBox to cluster nodes.
 */
trait ClassBoxComponent {
  self: SilkFramework =>

  type ClassBoxType <: ClassBoxAPI

  /**
   * Get the current class box id
   */
  def classBoxID : UUID

  /**
   * Retrieve the class box having the specified id
   * @param classBoxID
   * @return
   */
  def getClassBox(classBoxID:UUID) : ClassBoxAPI

}


trait ExecutorAPI {
  def getSlices[A](op: Silk[A]) : GenSeq[SilkFuture[Slice]]
}


trait DefaultExecutor extends ExecutorComponent {
  self : SilkFramework
    with LocalTaskManagerComponent
    with ClassBoxComponent
    with SliceStorageComponent =>

  type Executor = ExecutorImpl
  def executor : Executor = new ExecutorImpl

  class ExecutorImpl extends ExecutorAPI with ExecutorBase {}
}



/**
 * Executor receives a silk, optimize a plan, then submit evaluation tasks to the local task manager
 */
trait ExecutorComponent {
  self : SilkFramework
    with LocalTaskManagerComponent
    with ClassBoxComponent
    with SliceStorageComponent =>

  type Executor <: ExecutorAPI with ExecutorBase
  def executor : Executor

  trait ExecutorBase extends Logger {

    implicit class toGenFun[A,B](f:A=>B) {
      def toF1 : Any=>Any = f.asInstanceOf[Any=>Any]

      def toFlatMap : Any=>SilkSeq[Any] = f.asInstanceOf[Any=>SilkSeq[Any]]
      def toFilter : Any=>Boolean = f.asInstanceOf[Any=>Boolean]
    }

    implicit class toGenFMap[A, B](f:A=>GenTraversable[B]) {
      def toFmap = f.asInstanceOf[Any=>GenTraversable[Any]]
    }

    def run[A](session:Session, silk: Silk[A]): Result[A] = {

      val dataSeq : ParSeq[Seq[A]] = for{
        future <- getSlices(silk)
      }
      yield {
        val slice = future.get
        //debug(s"Fetch slice: $slice in $silk")
        sliceStorage.retrieve(silk.id, slice).asInstanceOf[Seq[A]]
      }

      val result = dataSeq.seq.flatten
      result
    }

    /**
     * Get Stage information
     * @param op
     * @tparam A
     * @return
     */
    def getStage[A](op:Silk[A]) : StageInfo = {
      sliceStorage.getStageInfo(op).map { si =>
        si.status match {
          case StageStarted(ts) => si
          case StageFinished(ts) => si
          case StageAborted(cause, ts) => // TODO restart
            SilkException.error(s"stage of [${op.idPrefix}] has been aborted: $cause")
        }
      } getOrElse {
        // Start a stage for evaluating the Silk
        executor.startStage(op)
      }
    }

    private def startStage[A,Out](op:Silk[A], in:Silk[A], f:Seq[_]=>Out) = {
      val inputStage = getStage(in)
      val N = inputStage.numSlices
      val stageInfo = StageInfo(0, N, StageStarted(System.currentTimeMillis()))
      sliceStorage.setStageInfo(op, stageInfo)
      // TODO append par
      for(i <- (0 until N).par) {
        // Get an input slice
        val inputSlice = sliceStorage.get(in.id, i).get
        // Send the slice processing task to a node close to the inputSlice location
        localTaskManager.submit(EvalSliceTask(s"eval slice of ${op}", UUID.randomUUID, classBoxID, op.id, in.id, inputSlice, f.toF1, Seq(inputSlice.nodeName)))
      }
      stageInfo
    }


    private def startReduceStage[A](op:Silk[A], in:Silk[A], reducer:Seq[_] => Any, aggregator:Seq[_]=>Any) = {
      val inputStage = getStage(in)
      val N = inputStage.numSlices

      // The outer reduce task produces only 1 slice
      val stageInfo = StageInfo(0, 1, StageStarted(System.currentTimeMillis()))

      // Evaluate reduce at each slice
      val subStageID = SilkUtil.newUUID
      for(i <- (0 until N).par) {
        // Get an input slice
        val inputSlice = sliceStorage.get(in.id, i).get
        // Locality-aware job submission
        localTaskManager.submit(ReduceTask(s"reduce each ${op}", UUID.randomUUID, classBoxID, subStageID, in.id, Seq(i), i, reducer, aggregator,  Seq(inputSlice.nodeName)))
      }

      // Determine the number of reducers to use. The default is 1/3 of the number of the input slices
      val R = ((N + (3-1))/ 3.0).toInt
      val W = (N + (R-1)) / R
      info(s"num reducers:$R, W:$W")

      // Reduce the previous sub stage using R reducers
      val aggregateStageID = SilkUtil.newUUID
      for((sliceRange, i) <- (0 until N).sliding(W, W).zipWithIndex) {
        val sliceIndexSet = sliceRange.toIndexedSeq
        localTaskManager.submit(ReduceTask(s"reduce aggregate ${op}", UUID.randomUUID, classBoxID, aggregateStageID, subStageID, sliceIndexSet, i, aggregator, aggregator, Seq.empty))
      }

      // The final aggregate task
      localTaskManager.submit(ReduceTask(s"reduce final of ${op}", UUID.randomUUID, classBoxID, op.id, aggregateStageID, (0 until R).toIndexedSeq, 0, aggregator, aggregator, Seq.empty))

      stageInfo
    }


    private def startShuffleStage[A, K](shuffleOp:ShuffleOp[A, K]) = {
      val inputStage = getStage(shuffleOp.in)
{'fraction_non_alphanumeric': 0.10933980246057876, 'fraction_numerical': 0.0031190434933287127, 'mean_word_length': 3.156643860280879, 'pattern_counts': {'":': 0, '<': 15, '<?xml version=': 0, '>': 41, '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': '4628307', 'n_tokens_mistral': 3767, 'n_tokens_neox': 3697, 'n_words': 923}