<commit_before>/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance with
 * the License.  You may obtain a copy of the License at
 *
 *    http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package org.apache.spark.deploy.control

import org.apache.spark.deploy.DeployMessages.{KillDriverResponse, KillExecutors, RequestKillDriver}
import org.apache.spark.deploy.master.Master
import org.apache.spark.internal.Logging
import org.apache.spark.rpc.{RpcAddress, RpcEnv, ThreadSafeRpcEndpoint}
import org.apache.spark.scheduler.StageInfo
import org.apache.spark.scheduler.cluster.CoarseGrainedClusterMessages._
import org.apache.spark.{SecurityManager, SparkConf}

import scala.collection.mutable.{HashMap, ListBuffer}


class ControllerJob(conf: SparkConf, deadlineJobMillisecond: Long) extends Logging {

  var NOMINAL_RATE_RECORD_S: Double = conf.getDouble("spark.control.nominalrate", 1000.0)
  val OVERSCALE: Int = conf.getInt("spark.control.overscale", 2)
  val LOCALITY_FACTOR: Int = 2

  val NOMINAL_RATE_DATA_S: Double = conf.getDouble(
    "spark.control.nominalratedata", 48000000.0)

  val DEADLINE_TIMEOUT = 300000

  val numMaxExecutor: Int = conf.getInt("spark.control.maxexecutor", 4)
  val coreForVM: Int = conf.getInt("spark.control.coreforvm", 8)

  var numExecutor = 0
  var totalCore = 0
  var coreForExecutor = new HashMap[Int, Int]

  val securityMgr = new SecurityManager(conf)
  val rpcEnv = RpcEnv.create("ControllEnv", "localhost", 6666, conf, securityMgr, clientMode = true)
  val controllerEndpoint = rpcEnv.setupEndpoint("ControllJob",
    new ControllerJob(rpcEnv, "ControllEnv", "ControllJob", conf, securityMgr))
  // rpcEnv.awaitTermination()


  def stop(): Unit = {
    rpcEnv.stop(controllerEndpoint)
  }

  def computeDeadlineStage(stage: StageInfo, weight: Long, startTime: Long,
                           alpha: Double, deadline: Long): Long = {
    var deadline = (deadlineJobMillisecond - startTime) / (weight + 1)
    if (deadline < 0) {
      logError("ALPHA DEADLINE NEGATIVE -> ALPHA DEADLINE NOT SATISFIED")
      deadline = (((deadlineJobMillisecond + (1 - alpha) * deadline) - startTime)
        / (weight + 1)).toLong
    }
    if (deadline < 0) {
      logError("DEADLINE NEGATIVE -> DEADLINE NOT SATISFIED")
      deadline = DEADLINE_TIMEOUT
    }
    deadline
  }

  def computeNominalRecord(stage: StageInfo, duration: Long, recordsRead: Double): Unit = {
    // val duration = (stage.completionTime.get - stage.submissionTime.get) / 1000.0
    NOMINAL_RATE_RECORD_S = recordsRead / (duration / 1000.0)
    logInfo("DURATION STAGE ID " + stage.stageId + " : " + duration)
    logInfo("NOMINAL RECORD/S STAGE ID " + stage.stageId + " : " + NOMINAL_RATE_RECORD_S)
    conf.set("spark.control.nominalrate", NOMINAL_RATE_RECORD_S.toString)
  }

  def computeCoreStage(deadlineStage: Long, numRecord: Long): Int = {
    logInfo("NumRecords: " + numRecord.toString +
      " DeadlineStage : " + deadlineStage.toString +
      " NominalRate: " + NOMINAL_RATE_RECORD_S.toString)
    OVERSCALE * math.ceil((numRecord / (deadlineStage / 1000.0)) / NOMINAL_RATE_RECORD_S).toInt
  }

  def computeDeadlineFirstStage(stage: StageInfo, weight: Long): Long = {
    val deadline = (deadlineJobMillisecond - stage.submissionTime.get) / (weight + 1)
    if (deadline < 0) {
      logError("DEADLINE NEGATIVE -> DEADLINE NOT SATISFIED")
    }
    deadline
  }

  def computeCoreFirstStage(stage: StageInfo): Int = {
    numMaxExecutor * coreForVM
  }

  def computeCoreStageFromSize(deadlineStage: Long, totalSize: Long): Int = {
    logInfo("TotalSize RDD First Stage: " + totalSize.toString)
    if (deadlineStage > 0) {
      OVERSCALE * math.ceil(totalSize / (deadlineStage / 1000.0) / NOMINAL_RATE_DATA_S).toInt
    } else {
      numMaxExecutor * coreForVM
    }
  }

  def computeTaskForExecutors(coresToBeAllocated: Int, totalTasksStage: Int): IndexedSeq[Int] = {
    numExecutor = math.ceil(coresToBeAllocated.toDouble / coreForVM.toDouble).toInt

    if (numExecutor > numMaxExecutor) {
      logError("NUM EXECUTORS TOO HIGH: %d > NUM MAX EXECUTORS %d".format(
        numExecutor, numMaxExecutor
      ))
      IndexedSeq(-1)
    }
    numExecutor = math.min(math.ceil(coresToBeAllocated.toDouble / coreForVM.toDouble).toInt
      * LOCALITY_FACTOR,
      numMaxExecutor)
    if ((coresToBeAllocated / numExecutor) <= 1) {
      numExecutor = math.ceil(coresToBeAllocated.toDouble / coreForVM.toDouble).toInt
    }
    var coresToStart = coresToBeAllocated.toInt
    var n = numExecutor
    var coresForExecutor = new ListBuffer[Int]()
    while (coresToStart > 0 && n > 0) {
      val a = math.floor(coresToStart / n).toInt
      coresToStart -= a
      n -= 1
      coresForExecutor += a
    }
    val coresPerExecutor = coresForExecutor.toList
//    val coresPerExecutor = (1 to numExecutor).map {
//      i => if (coresToBeAllocated % numExecutor >= i) {
//        1 + (coresToBeAllocated / numExecutor)
//      } else coresToBeAllocated / numExecutor
//    }

    val remainingTasks = totalTasksStage - coresPerExecutor.foldLeft(0) {
      (agg, x) => totalTasksStage * x / coresToBeAllocated + agg
    }

    val taskPerExecutor = (0 until numExecutor).map { i =>
      if (i < remainingTasks) {
        totalTasksStage * coresPerExecutor(i) / coresToBeAllocated + 1
      }
      else {
        totalTasksStage * coresPerExecutor(i) / coresToBeAllocated
      }
    }

    // val taskPerExecutor = scala.collection.mutable.IndexedSeq((0 until numExecutor).map {
    //  tasks * coresPerExecutor(_) / coresToBeAllocated
    // }: _*)

    // val remainingTasks = tasks - taskPerExecutor.sum

    // (0 until remainingTasks).foreach { i =>
    //  taskPerExecutor(i % numExecutor) = taskPerExecutor(i % numExecutor) + 1
    // }

    taskPerExecutor<commit_msg>[CONTROL] Fix task for core<commit_after>/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance with
 * the License.  You may obtain a copy of the License at
 *
 *    http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package org.apache.spark.deploy.control

import org.apache.spark.deploy.DeployMessages.{KillDriverResponse, KillExecutors, RequestKillDriver}
import org.apache.spark.deploy.master.Master
import org.apache.spark.internal.Logging
import org.apache.spark.rpc.{RpcAddress, RpcEnv, ThreadSafeRpcEndpoint}
import org.apache.spark.scheduler.StageInfo
import org.apache.spark.scheduler.cluster.CoarseGrainedClusterMessages._
import org.apache.spark.{SecurityManager, SparkConf}

import scala.collection.mutable.{HashMap, ListBuffer}


class ControllerJob(conf: SparkConf, deadlineJobMillisecond: Long) extends Logging {

  var NOMINAL_RATE_RECORD_S: Double = conf.getDouble("spark.control.nominalrate", 1000.0)
  val OVERSCALE: Int = conf.getInt("spark.control.overscale", 2)
  val LOCALITY_FACTOR: Int = 2

  val NOMINAL_RATE_DATA_S: Double = conf.getDouble(
    "spark.control.nominalratedata", 48000000.0)

  val DEADLINE_TIMEOUT = 300000

  val numMaxExecutor: Int = conf.getInt("spark.control.maxexecutor", 4)
  val coreForVM: Int = conf.getInt("spark.control.coreforvm", 8)

  var numExecutor = 0
  var totalCore = 0
  var coreForExecutor = new HashMap[Int, Int]

  val securityMgr = new SecurityManager(conf)
  val rpcEnv = RpcEnv.create("ControllEnv", "localhost", 6666, conf, securityMgr, clientMode = true)
  val controllerEndpoint = rpcEnv.setupEndpoint("ControllJob",
    new ControllerJob(rpcEnv, "ControllEnv", "ControllJob", conf, securityMgr))
  // rpcEnv.awaitTermination()


  def stop(): Unit = {
    rpcEnv.stop(controllerEndpoint)
  }

  def computeDeadlineStage(stage: StageInfo, weight: Long, startTime: Long,
                           alpha: Double, deadline: Long): Long = {
    var deadline = (deadlineJobMillisecond - startTime) / (weight + 1)
    if (deadline < 0) {
      logError("ALPHA DEADLINE NEGATIVE -> ALPHA DEADLINE NOT SATISFIED")
      deadline = (((deadlineJobMillisecond + (1 - alpha) * deadline) - startTime)
        / (weight + 1)).toLong
    }
    if (deadline < 0) {
      logError("DEADLINE NEGATIVE -> DEADLINE NOT SATISFIED")
      deadline = DEADLINE_TIMEOUT
    }
    deadline
  }

  def computeNominalRecord(stage: StageInfo, duration: Long, recordsRead: Double): Unit = {
    // val duration = (stage.completionTime.get - stage.submissionTime.get) / 1000.0
    NOMINAL_RATE_RECORD_S = recordsRead / (duration / 1000.0)
    logInfo("DURATION STAGE ID " + stage.stageId + " : " + duration)
    logInfo("NOMINAL RECORD/S STAGE ID " + stage.stageId + " : " + NOMINAL_RATE_RECORD_S)
    conf.set("spark.control.nominalrate", NOMINAL_RATE_RECORD_S.toString)
  }

  def computeCoreStage(deadlineStage: Long, numRecord: Long): Int = {
    logInfo("NumRecords: " + numRecord.toString +
      " DeadlineStage : " + deadlineStage.toString +
      " NominalRate: " + NOMINAL_RATE_RECORD_S.toString)
    OVERSCALE * math.ceil((numRecord / (deadlineStage / 1000.0)) / NOMINAL_RATE_RECORD_S).toInt
  }

  def computeDeadlineFirstStage(stage: StageInfo, weight: Long): Long = {
    val deadline = (deadlineJobMillisecond - stage.submissionTime.get) / (weight + 1)
    if (deadline < 0) {
      logError("DEADLINE NEGATIVE -> DEADLINE NOT SATISFIED")
    }
    deadline
  }

  def computeCoreFirstStage(stage: StageInfo): Int = {
    numMaxExecutor * coreForVM
  }

  def computeCoreStageFromSize(deadlineStage: Long, totalSize: Long): Int = {
    logInfo("TotalSize RDD First Stage: " + totalSize.toString)
    if (deadlineStage > 0) {
      OVERSCALE * math.ceil(totalSize / (deadlineStage / 1000.0) / NOMINAL_RATE_DATA_S).toInt
    } else {
      numMaxExecutor * coreForVM
    }
  }

  def computeTaskForExecutors(coresToBeAllocated: Int, totalTasksStage: Int): List[Int] = {
    numExecutor = math.ceil(coresToBeAllocated.toDouble / coreForVM.toDouble).toInt

    if (numExecutor > numMaxExecutor) {
      logError("NUM EXECUTORS TOO HIGH: %d > NUM MAX EXECUTORS %d".format(
        numExecutor, numMaxExecutor
      ))
      List(-1)
    }
    numExecutor = math.min(math.ceil(coresToBeAllocated.toDouble / coreForVM.toDouble).toInt
      * LOCALITY_FACTOR,
      numMaxExecutor)
    if ((coresToBeAllocated / numExecutor) <= 1) {
      numExecutor = math.ceil(coresToBeAllocated.toDouble / coreForVM.toDouble).toInt
    }
    var coresToStart = math.ceil(coresToBeAllocated.toDouble / OVERSCALE).toInt
    var n = numExecutor
    var coresForExecutor = new ListBuffer[Int]()
    while (coresToStart > 0 && n > 0) {
      val a = math.floor(coresToStart / n).toInt
      coresToStart -= a
      n -= 1
      coresForExecutor += a
    }
    val coresPerExecutor = coresForExecutor.toList
    //    val coresPerExecutor = (1 to numExecutor).map {
    //      i => if (coresToBeAllocated % numExecutor >= i) {
    //        1 + (coresToBeAllocated / numExecutor)
    //      } else coresToBeAllocated / numExecutor
    //    }
    val taskForOneCore = math.floor(totalTasksStage / (coresToBeAllocated / OVERSCALE)).toInt
    var remainingTasks = totalTasksStage.toInt
    val taskPerExecutor = (0 until numExecutor).map { i =>
      if (remainingTasks > taskForOneCore * coresPerExecutor(i)) {
        remainingTasks -= taskForOneCore * coresPerExecutor(i)
        taskForOneCore * coresPerExecutor(i)
      } else {
        remainingTasks = 0
        remainingTasks
      }
    }
    val taskForExecutor = scala.collection.mutable.IndexedSeq(taskPerExecutor: _*)
    var j = taskForExecutor.size - 1
    while (remainingTasks > 0) {
      taskForExecutor(j) += 1
      remainingTasks -= 1
      j -= 1
    }
    taskForExecutor.toList
{'fraction_non_alphanumeric': 0.10306710260003114, 'fraction_numerical': 0.011754631791997509, 'mean_word_length': 3.5014015416958655, 'pattern_counts': {'":': 0, '<': 12, '<?xml version=': 0, '>': 29, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, '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': '1919879', 'n_tokens_mistral': 4224, 'n_tokens_neox': 3977, 'n_words': 973}