<reponame>huyutuo/horovod
from bayes_opt import BayesianOptimization
from bayes_opt import UtilityFunction
from bayes_opt.logger import JSONLogger
from bayes_opt.event import Events
from bayes_opt.util import load_logs
import datetime

from numpy.core.fromnumeric import take
import build_command
import json
import read_config
import lhs
import numpy as np
import time
import os



def train_function(reduce, gather, fusion_buffer, cycle_time, cachen, nccl_thread, nccl_socks_pthread):
  run_command = build_command.get_run_command(reduce, gather, fusion_buffer, cycle_time, cachen, nccl_thread, nccl_socks_pthread,
      train_file_name  = " python3 ~/hyt/dis-train/vgg16-dis-train.py --batch-size=32  --epochs=10 ")
  print(run_command)

  start = datetime.datetime.now()
  os.system(run_command)
  # time.sleep(cycle_time/100)
  end = datetime.datetime.now()
  dur = (end - start).total_seconds()
  return -dur


def get_test_point_lhs():
  filename = "parameter.json"
  parameter_count, names, bounds, layers = read_config.get_config(filename)
  lhs_sample = lhs.LHSample(parameter_count, names, bounds, layers)
  samples = lhs_sample.getSample()
  samples = np.round(samples, decimals = 2)
  print(samples)
  return samples
  read_config.out_to_config_file(samples.tolist())


def get_pbounds():
  # Bounded region of parameter space
  filename = "parameter.json"
  _, _, bounds, _ = read_config.get_config(filename)
  pbounds = {}
  pbounds['reduce'] = (bounds[2][0], bounds[2][1])
  pbounds['gather'] = (bounds[3][0], bounds[3][1])
  pbounds['fusion_buffer'] = (bounds[0][0], bounds[0][1])
  pbounds['cycle_time'] = (bounds[1][0], bounds[1][1])
  pbounds['cachen'] = (bounds[4][0], bounds[4][1])
  pbounds['nccl_thread'] = (bounds[5][0], bounds[5][1])
  pbounds['nccl_socks_pthread'] = (bounds[6][0], bounds[6][1])
  return pbounds

def get_one_sample(one_point):
  params = {}
  params['reduce'] = one_point[2] 
  params['gather'] = one_point[3]
  params['fusion_buffer'] = one_point[0]
  params['cycle_time'] = one_point[1]
  params['cachen'] = one_point[4]
  params['nccl_thread'] = one_point[5]
  params['nccl_socks_pthread'] = one_point[6]
  return params


def run_test_point(optimizer):
  test_points = get_test_point_lhs()
  for point in test_points:
    params = get_one_sample(point)
    # print(params)
    target = train_function(point[2], point[3], point[0], point[1], point[4], point[5], point[6])
    # print(target)
    optimizer.register(
      params = params,
      target = target,
    )

if __name__ == "__main__":
  pbounds = get_pbounds()
  print(pbounds)
  optimizer = BayesianOptimization(
    f = train_function,
    pbounds = pbounds,
  )

  logger = JSONLogger(path="./vgg16-bayes-log.json")
  optimizer.subscribe(Events.OPTIMIZATION_STEP, logger) 

  run_test_point(optimizer)

  optimizer.maximize(
    init_points=0,
    n_iter=30,
    acq='ei'
  )

  print(optimizer.max)


  '''
  load_logs(new_optimizer, logs=["./logs.json"])
  
  next_point = new_optimizer.suggest(utility) 
  print(next_point)
  '''
from bayes_opt import BayesianOptimization
from bayes_opt import UtilityFunction
from bayes_opt.logger import JSONLogger
from bayes_opt.event import Events
from bayes_opt.util import load_logs
import datetime

from numpy.core.fromnumeric import take
import build_command
import json
import read_config
import lhs
import numpy as np
import time
import os



def train_function(reduce, gather, fusion_buffer, cycle_time, cachen, nccl_thread, nccl_socks_pthread):
  run_command = build_command.get_run_command(reduce, gather, fusion_buffer, cycle_time, cachen, nccl_thread, nccl_socks_pthread,
      train_file_name  = " python3 ~/hyt/dis-train/vgg16-dis-train.py --batch-size=32  --epochs=10 ")
  print(run_command)

  start = datetime.datetime.now()
  os.system(run_command)
  # time.sleep(cycle_time/100)
  end = datetime.datetime.now()
  dur = (end - start).total_seconds()
  return -dur


def get_test_point_lhs():
  filename = "parameter.json"
  parameter_count, names, bounds, layers = read_config.get_config(filename)
  lhs_sample = lhs.LHSample(parameter_count, names, bounds, layers)
  samples = lhs_sample.getSample()
  samples = np.round(samples, decimals = 2)
  print(samples)
  return samples
  read_config.out_to_config_file(samples.tolist())


def get_pbounds():
  # Bounded region of parameter space
  filename = "parameter.json"
  _, _, bounds, _ = read_config.get_config(filename)
  pbounds = {}
  pbounds['reduce'] = (bounds[2][0], bounds[2][1])
  pbounds['gather'] = (bounds[3][0], bounds[3][1])
  pbounds['fusion_buffer'] = (bounds[0][0], bounds[0][1])
  pbounds['cycle_time'] = (bounds[1][0], bounds[1][1])
  pbounds['cachen'] = (bounds[4][0], bounds[4][1])
  pbounds['nccl_thread'] = (bounds[5][0], bounds[5][1])
  pbounds['nccl_socks_pthread'] = (bounds[6][0], bounds[6][1])
  return pbounds

def get_one_sample(one_point):
  params = {}
  params['reduce'] = one_point[2] 
  params['gather'] = one_point[3]
  params['fusion_buffer'] = one_point[0]
  params['cycle_time'] = one_point[1]
  params['cachen'] = one_point[4]
  params['nccl_thread'] = one_point[5]
  params['nccl_socks_pthread'] = one_point[6]
  return params


def run_test_point(optimizer):
  test_points = get_test_point_lhs()
  for point in test_points:
    params = get_one_sample(point)
    # print(params)
    target = train_function(point[2], point[3], point[0], point[1], point[4], point[5], point[6])
    # print(target)
    optimizer.register(
      params = params,
      target = target,
    )

if __name__ == "__main__":
  pbounds = get_pbounds()
  print(pbounds)
  optimizer = BayesianOptimization(
    f = train_function,
    pbounds = pbounds,
  )

  logger = JSONLogger(path="./vgg16-bayes-log.json")
  optimizer.subscribe(Events.OPTIMIZATION_STEP, logger) 

  run_test_point(optimizer)

  optimizer.maximize(
    init_points=0,
    n_iter=30,
    acq='ei'
  )

  print(optimizer.max)


  '''
  load_logs(new_optimizer, logs=["./logs.json"])
  
  next_point = new_optimizer.suggest(utility) 
  print(next_point)
  '''
# time.sleep(cycle_time/100) # Bounded region of parameter space # print(params) # print(target) load_logs(new_optimizer, logs=["./logs.json"])      next_point = new_optimizer.suggest(utility)    print(next_point)