package resa.evaluation.scheduler.plan;
import backtype.storm.Config;
import backtype.storm.generated.Nimbus;
import backtype.storm.generated.RebalanceOptions;
import backtype.storm.scheduler.ExecutorDetails;
import backtype.storm.utils.NimbusClient;
import backtype.storm.utils.Utils;
import org.apache.curator.framework.CuratorFramework;
import org.junit.Before;
import org.junit.Test;
import resa.migrate.plan.KuhnMunkres;
import resa.migrate.plan.PackCalculator;
import resa.util.TopologyHelper;
import java.nio.file.Files;
import java.nio.file.Paths;
import java.util.*;
import java.util.stream.Collectors;
import java.util.stream.IntStream;
import java.util.stream.Stream;
/**
* Created by ding on 14-8-12.
*/
public class MigrateSimulate {
private CuratorFramework zk;
private Map<String, Object> conf;
private List<int[]> allocations;
private List<String> slots;
private double[] dataSizes;
private KuhnMunkres km;
private ExecutorDetails[] executors;
private String topoName = "fpt";
@Before
public void init() throws Exception {
conf = Utils.readStormConfig();
conf.put(Config.NIMBUS_HOST, "192.168.0.30");
conf.put(Config.STORM_ZOOKEEPER_SERVERS, Arrays.asList("192.168.0.30"));
zk = Utils.newCuratorStarted(conf, (List<String>) conf.get(Config.STORM_ZOOKEEPER_SERVERS),
conf.get(Config.STORM_ZOOKEEPER_PORT));
allocations = Files.readAllLines(Paths.get("/Volumes/Data/work/doctor/resa/exp/local-migrate.txt")).stream()
.map(String::trim).filter(s -> !s.isEmpty())
.map(s -> Stream.of(s.split(",")).mapToInt(Integer::parseInt).toArray())
.collect(Collectors.toList());
dataSizes = Files.readAllLines(Paths.get("/Volumes/Data/work/doctor/resa/exp/fp-data-sizes-064-70.txt")).stream()
.map(String::trim).filter(s -> !s.isEmpty()).mapToDouble(Double::valueOf).toArray();
slots = getSlots();
km = new KuhnMunkres(dataSizes.length);
String topoId = TopologyHelper.getTopologyId(topoName, conf);
Map<String, Map<ExecutorDetails, String>> assignment = TopologyHelper.getAssignment(conf, topoId);
Map<ExecutorDetails, String> compAssignment = assignment.get("detector");
executors = compAssignment.keySet().stream().sorted(Comparator.comparing(ExecutorDetails::getStartTask))
.toArray(ExecutorDetails[]::new);
}
@Test
public void deleteZkNode() throws Exception {
// ZKUtil.deleteRecursive(zk.getZookeeperClient().getZooKeeper(), "/resa");
}
@Test
public void run() throws Exception {
int state = 2;
String[] workers = slots.stream().limit(allocations.get(state).length).toArray(String[]::new);
assignment2State(allocations.get(state), workers);
doRebalance(topoName, allocations.get(state).length);
System.out.println("Init state, curr alloc=" + Arrays.toString(allocations.get(state)));
Scanner scanner = new Scanner(System.in);
String line;
int count = 0;
while ((line = scanner.nextLine()) != null) {
if (line.equals("c")) {
count++;
System.out.println("run " + count + ", curr alloc=" + Arrays.toString(allocations.get(state)) +
", target alloc=" + Arrays.toString(allocations.get(state + 1)));
System.out.println("old workers:" + Arrays.toString(workers));
workers = migrateTo(allocations.get(state), workers, allocations.get(state + 1));
System.out.println("new workers:" + Arrays.toString(workers));
assignment2State(allocations.get(state + 1), workers);
doRebalance(topoName, allocations.get(state + 1).length);
state++;
}
}
}
private String[] migrateTo(int[] from, String[] workers, int[] to) {
PackCalculator.Range[] pack1 = PackCalculator.convertPack(from);
PackCalculator.Range[] pack2 = PackCalculator.convertPack(to);
double[][] weights = new double[pack1.length][pack2.length];
for (int i = 0; i < pack1.length; i++) {
for (int j = 0; j < pack2.length; j++) {
weights[i][j] = overlap(pack1[i], pack2[j]);
}
}
int[][] map = km.getMaxBipartie(weights, new double[1]);
String[] newWorkers = new String[to.length];
Set<String> selected = new HashSet<>();
for (int i = 0; i < map.length; i++) {
newWorkers[map[i][1]] = workers[map[i][0]];
selected.add(newWorkers[map[i][1]]);
}
for (int i = 0; i < newWorkers.length; i++) {
if (newWorkers[i] == null) {
for (String w : slots) {
if (!selected.contains(w)) {
selected.add(w);
newWorkers[i] = w;
break;
}
}
}
}
return newWorkers;
}
private double overlap(PackCalculator.Range r1, PackCalculator.Range r2) {
if (r1.end < r2.start || r1.start > r2.end) {
return 0;
} else if (r1.start <= r2.start && r1.end >= r2.start) {
return IntStream.rangeClosed(r2.start, Math.min(r2.end, r1.end)).mapToDouble(i -> dataSizes[i]).sum();
} else if (r1.start >= r2.start && r1.start <= r2.end) {
return IntStream.rangeClosed(r1.start, Math.min(r2.end, r1.end)).mapToDouble(i -> dataSizes[i]).sum();
}
return 0;
}
public void assignment2State(int[] alloc, String[] workers) throws Exception {
Map<String, String> exe2WorkerSlot = new HashMap<>();
int j = 0;
for (int i = 0; i < alloc.length; i++) {
for (int k = 0; k < alloc[i]; k++) {
exe2WorkerSlot.put(executors[j].getStartTask() + "-" + executors[j].getEndTask(), workers[i]);
j++;
}
}
// System.out.println(exe2WorkerSlot);
String path = "/resa/" + TopologyHelper.getTopologyId(topoName, conf);
if (zk.checkExists().forPath(path) != null) {
zk.setData().forPath(path, Utils.serialize(new HashMap<>(Collections.singletonMap("assignment",
exe2WorkerSlot))));
} else {
zk.create().creatingParentsIfNeeded().forPath(path, Utils.serialize(
new HashMap<>(Collections.singletonMap("assignment", exe2WorkerSlot))));
}
}
private List<String> getSlots() throws Exception {
NimbusClient nimbusClient = NimbusClient.getConfiguredClient(conf);
Nimbus.Client nimbus = nimbusClient.getClient();
List<String> slots = nimbus.getClusterInfo().get_supervisors().stream()
.flatMap(s -> Arrays.asList(s.get_host() + ":6700", s.get_host() + ":6701").stream())
.collect(Collectors.toList());
Collections.shuffle(slots);
System.out.println(slots.size());
return slots;
}
private void doRebalance(String topoName, int numWorkers) throws Exception {
NimbusClient nimbusClient = NimbusClient.getConfiguredClient(conf);
Nimbus.Client nimbus = nimbusClient.getClient();
RebalanceOptions options = new RebalanceOptions();
options.set_num_workers(numWorkers);
options.set_wait_secs(1);
System.out.println("Reassigning to " + numWorkers + " workers");
nimbus.rebalance(topoName, options);
}
public static void main(String[] args) throws Exception {
MigrateSimulate simulate = new MigrateSimulate();
simulate.init();
simulate.run();
}
}