scripts/run_on_swarming_bots/run_on_swarming_bots.go

package main

import (
	"flag"
	"fmt"
	"io/ioutil"
	"os"
	"path"
	"path/filepath"
	"sync"
	"time"

	"github.com/satori/go.uuid"
	"go.skia.org/infra/go/auth"
	"go.skia.org/infra/go/common"
	"go.skia.org/infra/go/isolate"
	"go.skia.org/infra/go/sklog"
	"go.skia.org/infra/go/swarming"
)

/*
	Run a specified command on all specified GCE instances.
*/

const (
	ISOLATE_DEFAULT_NAMESPACE = "default-gzip"
	TMP_ISOLATE_FILE_NAME     = "script.isolate"
	TMP_ISOLATE_FILE_CONTENTS = `{
  'variables': {
    'command': [
      'python', '%s',
    ],
    'files': [
      '%s',
    ],
  },
}`
)

var (
	dimensions = common.NewMultiStringFlag("dimension", nil, "Colon-separated key/value pair, eg: \"os:Linux\" Dimensions of the bots on which to run. Can specify multiple times.")
	pool       = flag.String("pool", swarming.DIMENSION_POOL_VALUE_SKIA, "Which Swarming pool to use.")
	script     = flag.String("script", "", "Path to a Python script to run.")
	taskName   = flag.String("task_name", "", "Name of the task to run.")
	workdir    = flag.String("workdir", os.TempDir(), "Working directory. Optional, but recommended not to use CWD.")
)

func main() {
	// Setup, parse args.
	defer common.LogPanic()
	common.Init()

	if *script == "" {
		sklog.Fatal("--script is required.")
	}
	scriptName := path.Base(*script)
	if *taskName == "" {
		*taskName = fmt.Sprintf("run_%s", scriptName)
	}

	dims, err := swarming.ParseDimensionFlags(dimensions)
	if err != nil {
		sklog.Fatalf("Problem parsing dimensions: %s", err)
	}
	dims["pool"] = *pool

	*workdir, err = filepath.Abs(*workdir)
	if err != nil {
		sklog.Fatal(err)
	}

	// Authenticated HTTP client.
	oauthCacheFile := path.Join(*workdir, "google_storage_token.data")
	httpClient, err := auth.NewClient(true, oauthCacheFile, swarming.AUTH_SCOPE)
	if err != nil {
		sklog.Fatal(err)
	}

	// Swarming API client.
	swarmApi, err := swarming.NewApiClient(httpClient, swarming.SWARMING_SERVER)
	if err != nil {
		sklog.Fatal(err)
	}

	// Obtain the list of bots.
	bots, err := swarmApi.ListBots(dims)
	if err != nil {
		sklog.Fatal(err)
	}

	swarming, err := swarming.NewSwarmingClient(*workdir, swarming.SWARMING_SERVER, isolate.ISOLATE_SERVER_URL)
	if err != nil {
		sklog.Fatal(err)
	}

	// Copy the script to the workdir.
	dstScript := path.Join(*workdir, scriptName)
	contents, err := ioutil.ReadFile(*script)
	if err != nil {
		sklog.Fatal(err)
	}
	if err := ioutil.WriteFile(dstScript, contents, 0644); err != nil {
		sklog.Fatal(err)
	}

	// Create an isolate file.
	isolateFile := path.Join(*workdir, TMP_ISOLATE_FILE_NAME)
	if err := ioutil.WriteFile(isolateFile, []byte(fmt.Sprintf(TMP_ISOLATE_FILE_CONTENTS, scriptName, scriptName)), 0644); err != nil {
		sklog.Fatal(err)
	}

	// Upload to isolate server.
	isolated, err := swarming.CreateIsolatedGenJSON(isolateFile, *workdir, "linux", *taskName, map[string]string{}, []string{})
	if err != nil {
		sklog.Fatal(err)
	}
	m, err := swarming.BatchArchiveTargets([]string{isolated}, 5*time.Minute)
	if err != nil {
		sklog.Fatal(err)
	}
	group := fmt.Sprintf("%s_%s", *taskName, uuid.NewV1())
	tags := map[string]string{
		"group": group,
	}

	// Trigger the task on each bot.
	var wg sync.WaitGroup
	for _, bot := range bots {
		wg.Add(1)
		go func(id string) {
			defer wg.Done()
			dims := map[string]string{
				"pool": *pool,
				"id":   id,
			}
			if _, err := swarming.TriggerSwarmingTasks(m, dims, tags, 100, 120*time.Minute, 120*time.Minute, 120*time.Minute, false, false); err != nil {
				sklog.Fatal(err)
			}
		}(bot.BotId)
	}
	wg.Wait()
	tasksLink := fmt.Sprintf("https://chromium-swarm.appspot.com/tasklist?f=group:%s", group)
	sklog.Infof("Triggered Swarming tasks. Visit this link to track progress:\n%s", tasksLink)
}