Use swarming in run_lua CT task

ct/go/master_scripts/run_lua_on_workers/main.go

 // run_lua_on_workers is an application that runs the specified lua script on all
 // CT workers and uploads the results to Google Storage. The requester is emailed
 // when the task is done.
 package main
 
 import (
-        "bytes"
         "database/sql"
         "flag"
         "fmt"
         "io"
         "io/ioutil"
         "os"
         "path/filepath"
-        "strings"
-        "text/template"
+        "strconv"
         "time"
 
         "github.com/skia-dev/glog"
         "go.skia.org/infra/ct/go/ctfe/lua_scripts"
         "go.skia.org/infra/ct/go/frontend"
         "go.skia.org/infra/ct/go/master_scripts/master_common"
         "go.skia.org/infra/ct/go/util"
         "go.skia.org/infra/go/common"
         skutil "go.skia.org/infra/go/util"
 )
 
+const (
+        MAX_PAGES_PER_SWARMING_BOT = 100
+)
+
 var (
         emails        = flag.String("emails", "", "The comma separated email addresses to notify when the task is picked up and completes.")
         description   = flag.String("description", "", "The description of the run as entered by the requester.")
         gaeTaskID     = flag.Int64("gae_task_id", -1, "The key of the App Engine task. This task will be updated when the task is completed.")
         pagesetType   = flag.String("pageset_type", "", "The type of pagesets to use. Eg: 10k, Mobile10k, All.")
         chromiumBuild = flag.String("chromium_build", "", "The chromium build to use for this capture_archives run.")
         runID         = flag.String("run_id", "", "The unique run id (typically requester + timestamp).")
 
         taskCompletedSuccessfully     = false
         luaOutputRemoteLink           = ""
@@ skipping 54 matching lines @@
         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", *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")
@@ skipping 13 matching lines @@
 
         // 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
+        // Empty the remote dir before the workers upload to it.
+        gsBaseDir := filepath.Join(util.SWARMING_DIR_NAME, filepath.Join(util.LuaRunsDir, *runID), *pagesetType)

[dogben, 2016/05/19 17:02:41]

I would hope that runID is unique and therefore this dir wouldn't already
exist...

(Probably could/should have made this comment on other CLs.)

[rmistry, 2016/05/19 17:08:19]

The other CLs IIRC did not use a runID here.
But yes this dir would normally not already exist, but it is especially useful
for testing. I can keep using the same fake runID for local runs.

+        skutil.LogErr(gs.DeleteRemoteDir(gsBaseDir))
+
+        // Archive, trigger and collect swarming tasks.
+        isolateExtraArgs := map[string]string{
+                "CHROMIUM_BUILD": *chromiumBuild,
+                "RUN_ID":         *runID,
         }
-        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)
+        if err := util.TriggerSwarmingTask(*pagesetType, "run_lua", util.RUN_LUA_ISOLATE, 2*time.Hour, 1*time.Hour, MAX_PAGES_PER_SWARMING_BOT, isolateExtraArgs); err != nil {
+                glog.Errorf("Error encountered when swarming tasks: %s", err)
                 return
         }
 
         // Copy outputs from all slaves locally and combine it into one file.
         consolidatedFileName := "lua-output"
         consolidatedLuaOutput := filepath.Join(os.TempDir(), consolidatedFileName)
+        // If the file already exists it could be that there is another lua task running on this machine.

[dogben, 2016/05/19 17:02:41]

Please just use ioutil.TempDir instead of os.TempDir to avoid this issue.

[rmistry, 2016/05/19 17:08:19]

Unfortunately the file has to be in /tmp/lua-output. That is where all lua
aggregator scripts expect the output to reside. I could probably document that
lua aggregator scripts should instead include some unique string (eg:
__OUTPUT_FILE__) which can be replaced here. But for now it expects output in
/tmp/lua-output.

[dogben, 2016/05/19 17:15:38]

Oh, that is unfortunate.

+        // Wait for the file to be deleted within a deadline.
+        if err := waitForOutputFile(consolidatedLuaOutput); err != nil {
+                glog.Error(err)
+                return
+        }
         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")
+        for i := 1; i <= util.PagesetTypeToInfo[*pagesetType].NumPages/MAX_PAGES_PER_SWARMING_BOT; i++ {
+                startRange := strconv.Itoa(util.GetStartRange(i, MAX_PAGES_PER_SWARMING_BOT))
+                workerRemoteOutputPath := filepath.Join(util.LuaRunsDir, *runID, startRange, "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
@@ skipping 43 matching lines @@
                 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 waitForOutputFile(luaOutput string) error {
+        // Check every 10 secs and timeout after 10 mins.
+        ticker := time.NewTicker(10 * time.Second)
+        deadline := 10 * time.Minute
+        deadlineTicker := time.NewTicker(deadline)
+        defer ticker.Stop()
+        defer deadlineTicker.Stop()
+        for {
+                select {
+                case <-ticker.C:
+                        if _, err := os.Stat(luaOutput); os.IsNotExist(err) {
+                                return nil
+                        }
+                        glog.Infof("%s still exists. Waiting for the other lua task to complete.", luaOutput)
+                case <-deadlineTicker.C:
+                        return fmt.Errorf("%s still existed after %v secs", luaOutput, deadline.Seconds())
+                }
+        }
+}