scripts/run_on_swarming_bots/run_on_swarming_bots.go

package main

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

	"github.com/google/uuid"
	swarming_api "go.chromium.org/luci/common/api/swarming/swarming/v1"

	"go.skia.org/infra/go/auth"
	"go.skia.org/infra/go/cipd"
	"go.skia.org/infra/go/common"
	"go.skia.org/infra/go/exec"
	"go.skia.org/infra/go/httputils"
	"go.skia.org/infra/go/isolate"
	"go.skia.org/infra/go/sklog"
	"go.skia.org/infra/go/swarming"
	"go.skia.org/infra/go/util"
)

/*
	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': {
    'files': [
      '%s',
    ],
  },
}`
)

var (
	dev         = flag.Bool("dev", false, "Run against dev swarming and isolate instances.")
	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.")
	dryRun      = flag.Bool("dry_run", false, "List the bots, don't actually run any tasks")
	includeBots = common.NewMultiStringFlag("include_bot", nil, "If specified, treated as a white list of bots which will be affected, calculated AFTER the dimensions is computed. Can be simple strings or regexes")
	excludeBots = common.NewMultiStringFlag("exclude_bot", nil, "If specified, treated as a blacklist of bots which will not run the task, calculated AFTER the dimensions are computed and after --include_bot is applied. Can be simple strings or regexes.")
	internal    = flag.Bool("internal", false, "Run against internal swarming and isolate instances.")
	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.
	common.Init()

	ctx := context.Background()
	if *internal && *dev {
		sklog.Fatal("Both --internal and --dev cannot be specified.")
	}
	scriptName := path.Base(*script)
	if *taskName == "" {
		*taskName = fmt.Sprintf("run_%s", scriptName)
	}

	dims, err := swarming.ParseDimensionsSingleValue(*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)
	}

	includeRegs, err := parseRegex(*includeBots)
	if err != nil {
		sklog.Fatal(err)
	}
	excludeRegs, err := parseRegex(*excludeBots)
	if err != nil {
		sklog.Fatal(err)
	}

	// validate isolate is on PATH
	if err := exec.Run(context.Background(), &exec.Command{
		Name: "isolate",
		Args: []string{"version"},
	}); err != nil {
		sklog.Fatalf(`isolated not found on PATH. Checkout the README for installation details.`)
	}
	sklog.Info("isolate detected on PATH")

	// validate isolated is on PATH
	if err := exec.Run(context.Background(), &exec.Command{
		Name: "isolated",
		Args: []string{"version"},
	}); err != nil {
		sklog.Fatalf(`isolated not found on PATH. Checkout the README for installation details.`)
	}
	sklog.Info("isolated detected on PATH")

	isolateServer := isolate.ISOLATE_SERVER_URL
	swarmingServer := swarming.SWARMING_SERVER
	if *internal {
		isolateServer = isolate.ISOLATE_SERVER_URL_PRIVATE
		swarmingServer = swarming.SWARMING_SERVER_PRIVATE
	} else if *dev {
		isolateServer = isolate.ISOLATE_SERVER_URL_DEV
		swarmingServer = swarming.SWARMING_SERVER_DEV
	}

	// Authenticated HTTP client.
	ts, err := auth.NewDefaultTokenSource(true, swarming.AUTH_SCOPE)
	if err != nil {
		sklog.Fatal(err)
	}
	httpClient := httputils.DefaultClientConfig().WithTokenSource(ts).With2xxOnly().Client()

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

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

	var hashes []string
	if !*dryRun {
		if *script == "" {
			sklog.Fatal("--script is required if not running in dry run mode.")
		}

		// Copy the script to the workdir.
		isolateDir, err := ioutil.TempDir(*workdir, "run_on_swarming_bots")
		if err != nil {
			sklog.Fatal(err)
		}
		defer util.RemoveAll(isolateDir)
		dstScript := path.Join(isolateDir, scriptName)
		if err := util.CopyFile(*script, dstScript); err != nil {
			sklog.Fatal(err)
		}

		// Create an isolate file.
		isolateFile := path.Join(isolateDir, TMP_ISOLATE_FILE_NAME)
		if err := util.WithWriteFile(isolateFile, func(w io.Writer) error {
			_, err := w.Write([]byte(fmt.Sprintf(TMP_ISOLATE_FILE_CONTENTS, scriptName)))
			return err
		}); err != nil {
			sklog.Fatal(err)
		}

		// Upload to isolate server.
		isolateClient, err := isolate.NewClient(*workdir, isolateServer)
		if err != nil {
			sklog.Fatal(err)
		}
		isolateTask := &isolate.Task{
			BaseDir:     isolateDir,
			IsolateFile: isolateFile,
		}
		hashes, _, err = isolateClient.IsolateTasks(ctx, []*isolate.Task{isolateTask})
		if err != nil {
			sklog.Fatal(err)
		}
	}

	// Trigger the task on each bot.
	cmd := []string{"python", "-u", scriptName}
	group := fmt.Sprintf("%s_%s", *taskName, uuid.New())
	tags := []string{
		fmt.Sprintf("group:%s", group),
	}
	if *dryRun {
		sklog.Info("Dry run mode.  Would run on following bots:")
	}
	var wg sync.WaitGroup
	for _, bot := range bots {
		if len(includeRegs) > 0 && !matchesAny(bot.BotId, includeRegs) {
			sklog.Debugf("Skipping %s because it does not match --include_bot", bot.BotId)
			continue
		}
		if matchesAny(bot.BotId, excludeRegs) {
			sklog.Debugf("Skipping %s because it matches --exclude_bot", bot.BotId)
			continue
		}
		if *dryRun {
			sklog.Info(bot.BotId)
			continue
		}
		wg.Add(1)
		go func(id string) {
			defer wg.Done()
			dims := []*swarming_api.SwarmingRpcsStringPair{
				{
					Key:   "pool",
					Value: *pool,
				},
				{
					Key:   "id",
					Value: id,
				},
			}
			sklog.Infof("Triggering on %s", id)
			req := &swarming_api.SwarmingRpcsNewTaskRequest{
				Name:     *taskName,
				Priority: swarming.HIGHEST_PRIORITY,
				TaskSlices: []*swarming_api.SwarmingRpcsTaskSlice{
					{
						ExpirationSecs: int64((120 * time.Minute).Seconds()),
						Properties: &swarming_api.SwarmingRpcsTaskProperties{
							Caches: []*swarming_api.SwarmingRpcsCacheEntry{
								{
									Name: "vpython",
									Path: "cache/vpython",
								},
							},
							CipdInput:  swarming.ConvertCIPDInput(cipd.PkgsPython),
							Command:    cmd,
							Dimensions: dims,
							EnvPrefixes: []*swarming_api.SwarmingRpcsStringListPair{
								{
									Key:   "PATH",
									Value: []string{"cipd_bin_packages", "cipd_bin_packages/bin"},
								},
								{
									Key:   "VPYTHON_VIRTUALENV_ROOT",
									Value: []string{"cache/vpython"},
								},
							},
							ExecutionTimeoutSecs: int64((120 * time.Minute).Seconds()),
							Idempotent:           false,
							InputsRef: &swarming_api.SwarmingRpcsFilesRef{
								Isolated:       hashes[0],
								Isolatedserver: isolateServer,
								Namespace:      isolate.DEFAULT_NAMESPACE,
							},
							IoTimeoutSecs: int64((120 * time.Minute).Seconds()),
						},
					},
				},
				Tags: tags,
			}
			if _, err := swarmApi.TriggerTask(req); err != nil {
				sklog.Fatal(err)
			}
		}(bot.BotId)
	}

	wg.Wait()
	if !*dryRun {
		tasksLink := fmt.Sprintf("https://%s/tasklist?f=group:%s", swarmingServer, group)
		sklog.Infof("Triggered Swarming tasks. Visit this link to track progress:\n%s", tasksLink)
	}
}

func parseRegex(flags []string) (retval []*regexp.Regexp, e error) {
	for _, s := range flags {
		r, err := regexp.Compile(s)
		if err != nil {
			return nil, err
		}
		retval = append(retval, r)
	}
	return retval, nil
}

func matchesAny(s string, xr []*regexp.Regexp) bool {
	for _, r := range xr {
		if r.MatchString(s) {
			return true
		}
	}
	return false
}