func mergeUploadCSVFiles(runID string, gs *util.GsUtil) ([]string, error) {
localOutputDir := filepath.Join(util.StorageDir, util.BenchmarkRunsDir, runID)
skutil.MkdirAll(localOutputDir, 0700)
noOutputSlaves := []string{}
// Copy outputs from all slaves locally.
for i := 0; i < util.NumWorkers(); i++ {
workerNum := i + 1
workerLocalOutputPath := filepath.Join(localOutputDir, fmt.Sprintf("slave%d", workerNum)+".csv")
workerRemoteOutputPath := filepath.Join(util.BenchmarkRunsDir, runID, fmt.Sprintf("slave%d", workerNum), "outputs", runID+".output")
respBody, err := gs.GetRemoteFileContents(workerRemoteOutputPath)
if err != nil {
glog.Errorf("Could not fetch %s: %s", workerRemoteOutputPath, err)
noOutputSlaves = append(noOutputSlaves, fmt.Sprintf(util.WORKER_NAME_TEMPLATE, workerNum))
continue
}
defer skutil.Close(respBody)
out, err := os.Create(workerLocalOutputPath)
if err != nil {
return noOutputSlaves, fmt.Errorf("Unable to create file %s: %s", workerLocalOutputPath, err)
}
defer skutil.Close(out)
defer skutil.Remove(workerLocalOutputPath)
if _, err = io.Copy(out, respBody); err != nil {
return noOutputSlaves, fmt.Errorf("Unable to copy to file %s: %s", workerLocalOutputPath, err)
}
}
// Call csv_merger.py to merge all results into a single results CSV.
_, currentFile, _, _ := runtime.Caller(0)
pathToPyFiles := filepath.Join(
filepath.Dir((filepath.Dir(filepath.Dir(filepath.Dir(currentFile))))),
"py")
pathToCsvMerger := filepath.Join(pathToPyFiles, "csv_merger.py")
outputFileName := runID + ".output"
args := []string{
pathToCsvMerger,
"--csv_dir=" + localOutputDir,
"--output_csv_name=" + filepath.Join(localOutputDir, outputFileName),
}
err := util.ExecuteCmd("python", args, []string{}, util.CSV_MERGER_TIMEOUT, nil, nil)
if err != nil {
return noOutputSlaves, fmt.Errorf("Error running csv_merger.py: %s", err)
}
// Copy the output file to Google Storage.
remoteOutputDir := filepath.Join(util.BenchmarkRunsDir, runID, "consolidated_outputs")
if err := gs.UploadFile(outputFileName, localOutputDir, remoteOutputDir); err != nil {
return noOutputSlaves, fmt.Errorf("Unable to upload %s to %s: %s", outputFileName, remoteOutputDir, err)
}
return noOutputSlaves, nil
}
func main() {
defer common.LogPanic()
master_common.Init()
// Send start email.
emailsArr := util.ParseEmails(*emails)
emailsArr = append(emailsArr, util.CtAdmins...)
if len(emailsArr) == 0 {
glog.Error("At least one email address must be specified")
return
}
skutil.LogErr(frontend.UpdateWebappTaskSetStarted(&lua_scripts.UpdateVars{}, *gaeTaskID))
skutil.LogErr(util.SendTaskStartEmail(emailsArr, "Lua script", util.GetMasterLogLink(*runID), *description))
// Ensure webapp is updated and email is sent even if task fails.
defer updateWebappTask()
defer sendEmail(emailsArr)
// Cleanup tmp files after the run.
if !*master_common.Local {
defer util.CleanTmpDir()
}
// Finish with glog flush and how long the task took.
defer util.TimeTrack(time.Now(), "Running Lua script on workers")
defer glog.Flush()
if *pagesetType == "" {
glog.Error("Must specify --pageset_type")
return
}
if *chromiumBuild == "" {
glog.Error("Must specify --chromium_build")
return
}
if *runID == "" {
glog.Error("Must specify --run_id")
return
}
// Instantiate GsUtil object.
gs, err := util.NewGsUtil(nil)
if err != nil {
glog.Error(err)
return
}
// Upload the lua script for this run to Google storage.
luaScriptName := *runID + ".lua"
defer skutil.Remove(filepath.Join(os.TempDir(), luaScriptName))
luaScriptRemoteDir := filepath.Join(util.LuaRunsDir, *runID, "scripts")
if err := gs.UploadFile(luaScriptName, os.TempDir(), luaScriptRemoteDir); err != nil {
glog.Errorf("Could not upload %s to %s: %s", luaScriptName, luaScriptRemoteDir, err)
return
}
// Run the run_lua script on all workers.
runLuaCmdTemplate := "DISPLAY=:0 run_lua --worker_num={{.WorkerNum}} --log_dir={{.LogDir}} --log_id={{.RunID}} --pageset_type={{.PagesetType}} --chromium_build={{.ChromiumBuild}} --run_id={{.RunID}} --local={{.Local}};"
runLuaTemplateParsed := template.Must(template.New("run_lua_cmd").Parse(runLuaCmdTemplate))
luaCmdBytes := new(bytes.Buffer)
if err := runLuaTemplateParsed.Execute(luaCmdBytes, struct {
WorkerNum string
LogDir string
PagesetType string
ChromiumBuild string
RunID string
Local bool
}{
WorkerNum: util.WORKER_NUM_KEYWORD,
LogDir: util.GLogDir,
PagesetType: *pagesetType,
ChromiumBuild: *chromiumBuild,
RunID: *runID,
Local: *master_common.Local,
}); err != nil {
glog.Errorf("Failed to execute template: %s", err)
return
}
cmd := append(master_common.WorkerSetupCmds(),
// The main command that runs run_lua on all workers.
luaCmdBytes.String())
_, err = util.SSH(strings.Join(cmd, " "), util.Slaves, util.RUN_LUA_TIMEOUT)
if err != nil {
glog.Errorf("Error while running cmd %s: %s", cmd, err)
return
}
// Copy outputs from all slaves locally and combine it into one file.
consolidatedFileName := "lua-output"
consolidatedLuaOutput := filepath.Join(os.TempDir(), consolidatedFileName)
defer skutil.Remove(consolidatedLuaOutput)
if err := ioutil.WriteFile(consolidatedLuaOutput, []byte{}, 0660); err != nil {
glog.Errorf("Could not create %s: %s", consolidatedLuaOutput, err)
return
}
for i := 0; i < util.NumWorkers(); i++ {
workerNum := i + 1
workerRemoteOutputPath := filepath.Join(util.LuaRunsDir, *runID, fmt.Sprintf("slave%d", workerNum), "outputs", *runID+".output")
respBody, err := gs.GetRemoteFileContents(workerRemoteOutputPath)
if err != nil {
glog.Errorf("Could not fetch %s: %s", workerRemoteOutputPath, err)
continue
}
defer skutil.Close(respBody)
out, err := os.OpenFile(consolidatedLuaOutput, os.O_RDWR|os.O_APPEND, 0660)
if err != nil {
glog.Errorf("Unable to open file %s: %s", consolidatedLuaOutput, err)
return
}
defer skutil.Close(out)
if _, err = io.Copy(out, respBody); err != nil {
glog.Errorf("Unable to write out %s to %s: %s", workerRemoteOutputPath, consolidatedLuaOutput, err)
return
}
}
// Copy the consolidated file into Google Storage.
consolidatedOutputRemoteDir := filepath.Join(util.LuaRunsDir, *runID, "consolidated_outputs")
luaOutputRemoteLink = util.GS_HTTP_LINK + filepath.Join(util.GSBucketName, consolidatedOutputRemoteDir, consolidatedFileName)
if err := gs.UploadFile(consolidatedFileName, os.TempDir(), consolidatedOutputRemoteDir); err != nil {
glog.Errorf("Unable to upload %s to %s: %s", consolidatedLuaOutput, consolidatedOutputRemoteDir, err)
return
}
// Upload the lua aggregator (if specified) for this run to Google storage.
luaAggregatorName := *runID + ".aggregator"
luaAggregatorPath := filepath.Join(os.TempDir(), luaAggregatorName)
defer skutil.Remove(luaAggregatorPath)
luaAggregatorFileInfo, err := os.Stat(luaAggregatorPath)
if !os.IsNotExist(err) && luaAggregatorFileInfo.Size() > 10 {
if err := gs.UploadFile(luaAggregatorName, os.TempDir(), luaScriptRemoteDir); err != nil {
glog.Errorf("Could not upload %s to %s: %s", luaAggregatorName, luaScriptRemoteDir, err)
return
}
// Run the aggregator and save stdout.
luaAggregatorOutputFileName := *runID + ".agg.output"
luaAggregatorOutputFilePath := filepath.Join(os.TempDir(), luaAggregatorOutputFileName)
luaAggregatorOutputFile, err := os.Create(luaAggregatorOutputFilePath)
defer skutil.Close(luaAggregatorOutputFile)
defer skutil.Remove(luaAggregatorOutputFilePath)
if err != nil {
glog.Errorf("Could not create %s: %s", luaAggregatorOutputFilePath, err)
return
}
err = util.ExecuteCmd(util.BINARY_LUA, []string{luaAggregatorPath}, []string{},
util.LUA_AGGREGATOR_TIMEOUT, luaAggregatorOutputFile, nil)
if err != nil {
glog.Errorf("Could not execute the lua aggregator %s: %s", luaAggregatorPath, err)
return
}
// Copy the aggregator output into Google Storage.
luaAggregatorOutputRemoteLink = util.GS_HTTP_LINK + filepath.Join(util.GSBucketName, consolidatedOutputRemoteDir, luaAggregatorOutputFileName)
if err := gs.UploadFile(luaAggregatorOutputFileName, os.TempDir(), consolidatedOutputRemoteDir); err != nil {
glog.Errorf("Unable to upload %s to %s: %s", luaAggregatorOutputFileName, consolidatedOutputRemoteDir, err)
return
}
} else {
glog.Info("A lua aggregator has not been specified.")
}
taskCompletedSuccessfully = true
}
func main() {
defer common.LogPanic()
worker_common.Init()
defer util.TimeTrack(time.Now(), "Creating Pagesets")
defer glog.Flush()
// Create the task file so that the master knows this worker is still busy.
skutil.LogErr(util.CreateTaskFile(util.ACTIVITY_CREATING_PAGESETS))
defer util.DeleteTaskFile(util.ACTIVITY_CREATING_PAGESETS)
// Delete and remake the local pagesets directory.
pathToPagesets := filepath.Join(util.PagesetsDir, *pagesetType)
skutil.RemoveAll(pathToPagesets)
skutil.MkdirAll(pathToPagesets, 0700)
// Get info about the specified pageset type.
pagesetTypeInfo := util.PagesetTypeToInfo[*pagesetType]
csvSource := pagesetTypeInfo.CSVSource
numPages := pagesetTypeInfo.NumPages
userAgent := pagesetTypeInfo.UserAgent
// Download the CSV file from Google Storage to a tmp location.
gs, err := util.NewGsUtil(nil)
if err != nil {
glog.Error(err)
return
}
respBody, err := gs.GetRemoteFileContents(csvSource)
if err != nil {
glog.Error(err)
return
}
defer skutil.Close(respBody)
csvFile := filepath.Join(os.TempDir(), filepath.Base(csvSource))
out, err := os.Create(csvFile)
if err != nil {
glog.Errorf("Unable to create file %s: %s", csvFile, err)
return
}
defer skutil.Close(out)
defer skutil.Remove(csvFile)
if _, err = io.Copy(out, respBody); err != nil {
glog.Error(err)
return
}
// Figure out which pagesets this worker should generate.
numPagesPerSlave := numPages / util.NumWorkers()
if *worker_common.Local {
// When running locally, just do 10 pagesets to make things fast.
numPagesPerSlave = 10
}
startNum := (*workerNum-1)*numPagesPerSlave + 1
endNum := *workerNum * numPagesPerSlave
// Construct path to the create_page_set.py python script.
_, currentFile, _, _ := runtime.Caller(0)
createPageSetScript := filepath.Join(
filepath.Dir((filepath.Dir(filepath.Dir(filepath.Dir(currentFile))))),
"py", "create_page_set.py")
// Execute the create_page_set.py python script.
timeoutSecs := util.PagesetTypeToInfo[*pagesetType].CreatePagesetsTimeoutSecs
for currNum := startNum; currNum <= endNum; currNum++ {
args := []string{
createPageSetScript,
"-s", strconv.Itoa(currNum),
"-e", strconv.Itoa(currNum),
"-c", csvFile,
"-p", *pagesetType,
"-u", userAgent,
"-o", pathToPagesets,
}
if err := util.ExecuteCmd("python", args, []string{}, time.Duration(timeoutSecs)*time.Second, nil, nil); err != nil {
glog.Error(err)
return
}
}
// Write timestamp to the pagesets dir.
skutil.LogErr(util.CreateTimestampFile(pathToPagesets))
// Upload pagesets dir to Google Storage.
if err := gs.UploadWorkerArtifacts(util.PAGESETS_DIR_NAME, *pagesetType, *workerNum); err != nil {
glog.Error(err)
return
}
}