Use swarming in capture_skps and capture_skps_from_pdfs CT tasks

ct/go/master_scripts/capture_skps_on_workers/main.go

Index: ct/go/master_scripts/capture_skps_on_workers/main.go
diff --git a/ct/go/master_scripts/capture_skps_on_workers/main.go b/ct/go/master_scripts/capture_skps_on_workers/main.go
index 34a2c38a5ed80f88a1c45fb19f5fa96e5f95212d..11659adc8e406113737ea76d56c607f41c428a19 100644
--- a/ct/go/master_scripts/capture_skps_on_workers/main.go
+++ b/ct/go/master_scripts/capture_skps_on_workers/main.go
@@ -4,11 +4,14 @@
 package main
 
 import (
-	"bytes"
 	"flag"
 	"fmt"
+	"io/ioutil"
+	"path"
+	"path/filepath"
+	"runtime"
+	"strconv"
 	"strings"
-	"text/template"
 	"time"
 
 	"github.com/skia-dev/glog"
@@ -17,9 +20,15 @@ import (
 	"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.")
@@ -79,10 +88,6 @@ func main() {
 	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()
@@ -100,10 +105,12 @@ func main() {
 		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 {
@@ -129,41 +136,58 @@ func main() {
 		//}
 	}
 
-	// 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")
 
-	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)
+	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.Errorf("Error while running cmd %s: %s", cmd, err)
+		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
 	}