func main() {
etcdUrlList := flag.String("etcd_urls", "", "ETCD server lists, sep by a comma.")
jobName := flag.String("job_name", "bwmf", "Job name in etcd path.")
jobType := flag.String("job_type", "c", "Job type, either 'c' for controller or 't' for task.")
numTasks := flag.Int("num_tasks", 1, "Num of tasks.")
taskConfigFile := flag.String("task_config", "", "Path to task config json file.")
flag.Parse()
if *jobName == "" {
log.Fatal("Job name is required.")
}
crd, oErr := filesystem.NewLocalFSClient().OpenReadCloser(*taskConfigFile)
if oErr != nil {
log.Fatalf("Failed opening task config file. %s", oErr)
}
confData, rdErr := ioutil.ReadAll(crd)
if rdErr != nil {
log.Fatalf("Failed reading task config. %s", rdErr)
}
log.Printf("conf data: %s", confData)
if *etcdUrlList == "" {
log.Fatal("Please specify the etcd server urls.")
}
etcdUrls := strings.Split(*etcdUrlList, ",")
log.Println("etcd urls: ", etcdUrls)
topoMaster := topo.NewFullTopologyOfMaster(uint64(*numTasks))
topoNeighbors := topo.NewFullTopologyOfNeighbor(uint64(*numTasks))
switch *jobType {
case "t":
bootstrap := framework.NewBootStrap(*jobName, etcdUrls, createListener(), nil)
taskBuilder := &bwmf.BWMFTaskBuilder{
NumOfTasks: uint64(*numTasks),
ConfBytes: confData,
}
bootstrap.SetTaskBuilder(taskBuilder)
bootstrap.AddLinkage("Master", topoMaster)
bootstrap.AddLinkage("Neighbors", topoNeighbors)
log.Println("Starting task..")
bootstrap.Start()
case "c":
controller := controller.New(*jobName, etcd.NewClient(etcdUrls), uint64(*numTasks), []string{"Master", "Neighbors"})
controller.Start()
log.Println("Controller started.")
controller.WaitForJobDone()
controller.Stop()
default:
log.Fatal("Please choose a type via '-jobtype': (c) controller, (t) task")
}
}
func TestBWMF(t *testing.T) {
etcdURLs := []string{"http://localhost:4001"}
job := "bwmf_basic_test"
numOfTasks := uint64(2)
generateTestData(t)
ctl := controller.New(job, etcd.NewClient(etcdURLs), numOfTasks, []string{"Neighbors", "Master"})
ctl.Start()
tb := &bwmf.BWMFTaskBuilder{
NumOfTasks: numOfTasks,
ConfBytes: []byte(`{
"OptConf": {
"Sigma":0.01,
"Alpha":1,
"Beta":0.1,
"GradTol":1e-06,
"FixedCnt": 200000,
"NumIters":4,
"DimLatent":2
},
"IOConf": {
"Fs":"local",
"IDPath":"../.tmp/row_shard.dat",
"ITPath":"../.tmp/column_shard.dat",
"ODPath":"../.tmp/dShard.dat",
"OTPath":"../.tmp/tShard.dat"
}
}`),
}
for i := uint64(0); i < numOfTasks; i++ {
go drive(
t,
job,
etcdURLs,
tb,
map[string]taskgraph.Topology{
"Master": topo.NewFullTopologyOfMaster(numOfTasks),
"Neighbors": topo.NewFullTopologyOfNeighbor(numOfTasks),
},
)
}
ctl.WaitForJobDone()
ctl.Stop()
}