Use swarming in capture_skps and capture_skps_from_pdfs CT tasks

ct/go/master_scripts/capture_skps_on_workers/main.go

 // capture_skps_on_workers is an application that captures SKPs of the
 // specified patchset type on all CT workers and uploads the results to Google
 // Storage. The requester is emailed when the task is done.
 package main
 
 import (
-        "bytes"
         "flag"
         "fmt"
+        "io/ioutil"
+        "path"
+        "path/filepath"
+        "runtime"
+        "strconv"
         "strings"
-        "text/template"
         "time"
 
         "github.com/skia-dev/glog"
         "go.skia.org/infra/ct/go/ctfe/capture_skps"
         "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"
+        "go.skia.org/infra/go/swarming"
         skutil "go.skia.org/infra/go/util"
 )
 
+const (
+        MAX_PAGES_PER_SWARMING_BOT_CAPTURE_SKPS           = 100
+        MAX_PAGES_PER_SWARMING_BOT_CAPTURE_SKPS_FROM_PDFS = 1000
+)
+
 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 SKPs run.")
         targetPlatform = flag.String("target_platform", util.PLATFORM_LINUX, "The platform the benchmark will run on (Android / Linux).")
         runID          = flag.String("run_id", "", "The unique run id (typically requester + timestamp).")
 
         taskCompletedSuccessfully = false
@@ skipping 39 matching lines @@
                 glog.Error("At least one email address must be specified")
                 return
         }
         skutil.LogErr(frontend.UpdateWebappTaskSetStarted(&capture_skps.UpdateVars{}, *gaeTaskID))
         skutil.LogErr(util.SendTaskStartEmail(emailsArr, "Capture SKPs", *runID, *description))
         // Ensure webapp is updated and completion email is sent even if task
         // fails.
         defer updateWebappTask()
         defer sendEmail(emailsArr)
 
-        if !*master_common.Local {
-                // Cleanup tmp files after the run.
-                defer util.CleanTmpDir()
-        }
         // Finish with glog flush and how long the task took.
         defer util.TimeTrack(time.Now(), "Running capture skps task 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
         }
 
-        workerScript := "capture_skps"
+        isolateFile := util.CAPTURE_SKPS_ISOLATE
+        maxPages := MAX_PAGES_PER_SWARMING_BOT_CAPTURE_SKPS
         if strings.Contains(strings.ToUpper(*pagesetType), "PDF") {
                 // For PDF pagesets use the capture_skps_from_pdfs worker script.
-                workerScript = "capture_skps_from_pdfs"
+                isolateFile = util.CAPTURE_SKPS_FROM_PDFS_ISOLATE
+                maxPages = MAX_PAGES_PER_SWARMING_BOT_CAPTURE_SKPS_FROM_PDFS
                 // TODO(rmistry): Uncomment when ready to capture SKPs.
                 //// Sync PDFium and build pdfium_test binary which will be used by the worker script.
                 //if err := util.SyncDir(util.PDFiumTreeDir); err != nil {
                 //      glog.Errorf("Could not sync PDFium: %s", err)
                 //      return
                 //}
                 //if err := util.BuildPDFium(); err != nil {
                 //      glog.Errorf("Could not build PDFium: %s", err)
                 //      return
                 //}
                 //// Copy pdfium_test to Google Storage.
                 //pdfiumLocalDir := path.Join(util.PDFiumTreeDir, "out", "Debug")
                 //pdfiumRemoteDir := path.Join(util.BINARIES_DIR_NAME, *chromiumBuild)
                 //// Instantiate GsUtil object.
                 //gs, err := util.NewGsUtil(nil)
                 //if err != nil {
                 //      glog.Error(err)
                 //      return
                 //}
                 //if err := gs.UploadFile(util.BINARY_PDFIUM_TEST, pdfiumLocalDir, pdfiumRemoteDir); err != nil {
                 //      glog.Errorf("Could not upload %s to %s: %s", util.BINARY_PDFIUM_TEST, pdfiumRemoteDir, err)
                 //      return
                 //}
         }
 
-        // Run the capture SKPs script on all workers.
-        captureSKPsCmdTemplate := "DISPLAY=:0 {{.WorkerScript}} --worker_num={{.WorkerNum}} --log_dir={{.LogDir}} --log_id={{.RunID}} " +
-                "--pageset_type={{.PagesetType}} --chromium_build={{.ChromiumBuild}} --run_id={{.RunID}} " +
-                "--target_platform={{.TargetPlatform}} --local={{.Local}};"
-        captureSKPsTemplateParsed := template.Must(template.New("capture_skps_cmd").Parse(captureSKPsCmdTemplate))
-        captureSKPsCmdBytes := new(bytes.Buffer)
-        if err := captureSKPsTemplateParsed.Execute(captureSKPsCmdBytes, struct {
-                WorkerScript   string
-                WorkerNum      string
-                LogDir         string
-                PagesetType    string
-                ChromiumBuild  string
-                RunID          string
-                TargetPlatform string
-                Local          bool
-        }{
-                WorkerScript:   workerScript,
-                WorkerNum:      util.WORKER_NUM_KEYWORD,
-                LogDir:         util.GLogDir,
-                PagesetType:    *pagesetType,
-                ChromiumBuild:  *chromiumBuild,
-                RunID:          *runID,
-                TargetPlatform: *targetPlatform,
-                Local:          *master_common.Local,
-        }); err != nil {
-                glog.Errorf("Failed to execute template: %s", err)
+        // Instantiate the swarming client.
+        workDir, err := ioutil.TempDir("", "swarming_work_")
+        if err != nil {
+                glog.Errorf("Could not get temp dir: %s", err)
+                return
+        }
+        s, err := swarming.NewSwarmingClient(workDir)
+        if err != nil {
+                glog.Errorf("Could not instantiate swarming client: %s", err)
+                return
+        }
+        defer s.Cleanup()
+        // Create isolated.gen.json files from tasks.
+        taskNames := []string{}
+        for i := 1; i <= util.PagesetTypeToInfo[*pagesetType].NumPages/maxPages; i++ {
+                taskNames = append(taskNames, fmt.Sprintf("capture_skps_%d", i))
+        }
+        genJSONs := []string{}
+        // Get path to isolate files.
+        _, currentFile, _, _ := runtime.Caller(0)
+        pathToIsolates := filepath.Join(filepath.Dir((filepath.Dir(filepath.Dir(filepath.Dir(currentFile))))), "isolates")
+
+        for i, taskName := range taskNames {
+                extraArgs := map[string]string{
+                        "START_RANGE":    strconv.Itoa(util.GetStartRange(i+1, maxPages)),
+                        "NUM":            strconv.Itoa(maxPages),
+                        "PAGESET_TYPE":   *pagesetType,
+                        "CHROMIUM_BUILD": *chromiumBuild,
+                        "RUN_ID":         *runID,
+                }
+                genJSON, err := s.CreateIsolatedGenJSON(path.Join(pathToIsolates, isolateFile), s.WorkDir, "linux", taskName, extraArgs, []string{})
+                if err != nil {
+                        glog.Errorf("Could not create isolated.gen.json for task %s: %s", taskName, err)
+                        return
+                }
+                genJSONs = append(genJSONs, genJSON)
+        }
+        // Empty the remote dir before the workers upload to it.
+        gs, err := util.NewGsUtil(nil)
+        if err != nil {
+                glog.Error(err)
+                return
+        }
+        skpGSBaseDir := filepath.Join(util.SWARMING_DIR_NAME, util.SKPS_DIR_NAME, *pagesetType)
+        skutil.LogErr(gs.DeleteRemoteDir(skpGSBaseDir))
+        if strings.Contains(strings.ToUpper(*pagesetType), "PDF") {
+                pdfGSBaseDir := filepath.Join(util.SWARMING_DIR_NAME, util.PDFS_DIR_NAME, *pagesetType)
+                skutil.LogErr(gs.DeleteRemoteDir(pdfGSBaseDir))
+        }
+        // Archive, trigger and collect swarming tasks.
+        if err := util.ArchiveTriggerCollectSwarmingTask(s, taskNames, genJSONs, 2*time.Hour, 1*time.Hour); err != nil {
+                glog.Errorf("Error encountered when swarming tasks: %s", err)
                 return
         }
 
-        cmd := append(master_common.WorkerSetupCmds(),
-                // The main command that captures SKPs on all workers.
-                captureSKPsCmdBytes.String())
-        _, err := util.SSH(strings.Join(cmd, " "), util.Slaves, util.CAPTURE_SKPS_TIMEOUT)
-        if err != nil {
-                glog.Errorf("Error while running cmd %s: %s", cmd, err)
-                return
-        }
-
         taskCompletedSuccessfully = true
 }