| package main |
| |
| import ( |
| "context" |
| "flag" |
| "fmt" |
| "io/ioutil" |
| "os" |
| "path" |
| "path/filepath" |
| "regexp" |
| "strings" |
| "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/cas/rbe" |
| "go.skia.org/infra/go/cipd" |
| "go.skia.org/infra/go/common" |
| "go.skia.org/infra/go/httputils" |
| "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. |
| */ |
| |
| var ( |
| dev = flag.Bool("dev", false, "Run against dev swarming instance.") |
| 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, "Include these bots, regardless of whether they match the requested dimensions. Calculated AFTER the dimensions are computed. Can be simple strings or regexes.") |
| excludeBots = common.NewMultiStringFlag("exclude_bot", nil, "Exclude these bots, regardless of whether they match the requested dimensions. 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 instance.") |
| 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) |
| } |
| |
| casInstance := rbe.InstanceChromiumSwarm |
| swarmingServer := swarming.SWARMING_SERVER |
| if *internal { |
| casInstance = rbe.InstanceChromeSwarming |
| swarmingServer = swarming.SWARMING_SERVER_PRIVATE |
| } else if *dev { |
| casInstance = rbe.InstanceChromiumSwarmDev |
| 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 casInput *swarming_api.SwarmingRpcsCASReference |
| if !*dryRun { |
| if *script == "" { |
| sklog.Fatal("--script is required if not running in dry run mode.") |
| } |
| |
| // Copy the script to the workdir. |
| casRoot, err := ioutil.TempDir(*workdir, "run_on_swarming_bots") |
| if err != nil { |
| sklog.Fatal(err) |
| } |
| defer util.RemoveAll(casRoot) |
| dstScript := path.Join(casRoot, scriptName) |
| if err := util.CopyFile(*script, dstScript); err != nil { |
| sklog.Fatal(err) |
| } |
| |
| // Upload to CAS. |
| casClient, err := rbe.NewClient(ctx, casInstance, ts) |
| if err != nil { |
| sklog.Fatal(err) |
| } |
| digest, err := casClient.Upload(ctx, casRoot, []string{scriptName}, nil) |
| if err != nil { |
| sklog.Fatal(err) |
| } |
| casInput, err = swarming.MakeCASReference(digest, casInstance) |
| 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 |
| } |
| var os string |
| var arch string |
| var platform string |
| for _, dim := range bot.Dimensions { |
| if dim.Key == "os" { |
| for _, value := range dim.Value { |
| if strings.Contains(value, "Mac") { |
| os = "Mac" |
| break |
| } else if strings.Contains(value, "Linux") { |
| os = "Linux" |
| break |
| } else if strings.Contains(value, "Windows") { |
| os = "Windows" |
| break |
| } |
| } |
| } else if dim.Key == "cpu" { |
| for _, value := range dim.Value { |
| if strings.Contains(value, "arm") { |
| arch = "arm" |
| break |
| } else if strings.Contains(value, "x86-64") { |
| arch = "x86-64" |
| // Don't break, in case a bot has both "x86" and |
| // "x86-64" dimensions. |
| } else if strings.Contains(value, "x86-32") { |
| arch = "386" |
| // Don't break, in case a bot has both "x86" and |
| // "x86-64" dimensions. |
| } |
| } |
| } |
| } |
| if os == "Linux" { |
| if arch == "arm" { |
| platform = cipd.PlatformLinuxArm64 |
| } else { |
| platform = cipd.PlatformLinuxAmd64 |
| } |
| } else if os == "Windows" { |
| if arch == "386" { |
| platform = cipd.PlatformWindows386 |
| } else { |
| platform = cipd.PlatformWindowsAmd64 |
| } |
| } else if os == "Mac" { |
| platform = cipd.PlatformMacAmd64 |
| } |
| if platform == "" { |
| sklog.Fatal("Unable to find platform for %s", bot.BotId) |
| } |
| |
| 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", |
| }, |
| }, |
| CasInputRoot: casInput, |
| CipdInput: swarming.ConvertCIPDInput(cipd.PkgsPython[platform]), |
| 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, |
| 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 |
| } |