private def dataflowWorkerMemory()

in ratatool-sampling/src/main/scala/com/spotify/ratatool/samplers/BigSampler.scala [155:176]


  private def dataflowWorkerMemory(options: PipelineOptions): Option[Int] = Try {
    val dataflowPipelineWorkerPoolOptions = Class
      .forName("org.apache.beam.runners.dataflow.options.DataflowPipelineWorkerPoolOptions")
      .asInstanceOf[Class[PipelineOptions]]
    val machineType = dataflowPipelineWorkerPoolOptions
      .getMethod("getWorkerMachineType")
      .invoke(options.as(dataflowPipelineWorkerPoolOptions))
      .asInstanceOf[String]
    val machineTypeParts = machineType.split("-")
    // approximation. does not work for all machine types families
    val cpuMemFactor = machineTypeParts(1) match {
      case "standard" => 4
      case "highmem"  => 8
      case "megamem"  => 14
      case "hypermem" => 20
      case "ultramem" => 28
      case "highcpu"  => 2
      case _          => 1
    }
    val machineCpus = machineTypeParts(2).toInt
    machineCpus * cpuMemFactor
  }.toOption