123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281 |
- // Copyright 2019 Google LLC.
- // Use of this source code is governed by a BSD-style license that can be found in the LICENSE file.
- package main
- import (
- "bufio"
- "bytes"
- "flag"
- "fmt"
- "log"
- "math/rand"
- "os"
- "os/exec"
- "path/filepath"
- "runtime"
- "strings"
- "sync"
- "sync/atomic"
- "time"
- )
- // Too many GPU processes and we'll start to overwhelm your GPU,
- // even hanging your machine in the worst case. Here's a reasonable default.
- func defaultGpuLimit() int {
- limit := 8
- if n := runtime.NumCPU(); n < limit {
- return n
- }
- return limit
- }
- var script = flag.String("script", "", "A file with jobs to run, one per line. - for stdin.")
- var random = flag.Bool("random", true, "Assign sources into job batches randomly?")
- var quiet = flag.Bool("quiet", false, "Print only failures?")
- var exact = flag.Bool("exact", false, "Match GM names only exactly.")
- var cpuLimit = flag.Int("cpuLimit", runtime.NumCPU(),
- "Maximum number of concurrent processes for CPU-bound work.")
- var gpuLimit = flag.Int("gpuLimit", defaultGpuLimit(),
- "Maximum number of concurrent processes for GPU-bound work.")
- func init() {
- flag.StringVar(script, "s", *script, "Alias for --script.")
- flag.BoolVar(random, "r", *random, "Alias for --random.")
- flag.BoolVar(quiet, "q", *quiet, "Alias for --quiet.")
- flag.BoolVar(exact, "e", *exact, "Alias for --exact.")
- flag.IntVar(cpuLimit, "c", *cpuLimit, "Alias for --cpuLimit.")
- flag.IntVar(gpuLimit, "g", *gpuLimit, "Alias for --gpuLimit.")
- }
- func listAllGMs(fm string) (gms []string, err error) {
- // Query fm binary for list of all available GMs by running with no arguments.
- cmd := exec.Command(fm)
- stdout, err := cmd.Output()
- if err != nil {
- return
- }
- // GM names are listed line-by-line.
- scanner := bufio.NewScanner(bytes.NewReader(stdout))
- for scanner.Scan() {
- gms = append(gms, scanner.Text())
- }
- err = scanner.Err()
- return
- }
- type work struct {
- Sources []string
- Flags []string
- }
- func parseWork(args []string, gms []string) (*work, error) {
- w := &work{}
- for _, arg := range args {
- // I wish we could parse flags here too, but it's too late.
- if strings.HasPrefix(arg, "-") {
- msg := "Is '%s' an fm flag? If so please pass it using flag=value syntax."
- if flag.Lookup(arg[1:]) != nil {
- msg = "Please pass fm_bot flags like '%s' on the command line before the FM binary."
- }
- return nil, fmt.Errorf(msg, arg)
- }
- // Everything after a # is a comment.
- if strings.HasPrefix(arg, "#") {
- break
- }
- // Treat "gm" or "gms" as a shortcut for all known GMs.
- if arg == "gm" || arg == "gms" {
- w.Sources = append(w.Sources, gms...)
- continue
- }
- // Is this an option to pass through to fm?
- if parts := strings.Split(arg, "="); len(parts) == 2 {
- f := "-"
- if len(parts[0]) > 1 {
- f += "-"
- }
- f += parts[0]
- w.Flags = append(w.Flags, f, parts[1])
- continue
- }
- // Is this argument naming a GM?
- matchedAnyGM := false
- for _, gm := range gms {
- if (*exact && gm == arg) || (!*exact && strings.Contains(gm, arg)) {
- w.Sources = append(w.Sources, gm)
- matchedAnyGM = true
- }
- }
- if matchedAnyGM {
- continue
- }
- // Anything left ought to be on the file system: a file, a directory, or a glob.
- // Not all shells expand globs, so we'll do it here just in case.
- matches, err := filepath.Glob(arg)
- if err != nil {
- return nil, err
- }
- if len(matches) == 0 {
- return nil, fmt.Errorf("Don't understand '%s'.", arg)
- }
- for _, match := range matches {
- err := filepath.Walk(match, func(path string, info os.FileInfo, err error) error {
- if !info.IsDir() {
- w.Sources = append(w.Sources, path)
- }
- return err
- })
- if err != nil {
- return nil, err
- }
- }
- }
- return w, nil
- }
- func main() {
- flag.Parse()
- if flag.NArg() < 1 {
- log.Fatal("Please pass an fm binary as the first argument.")
- }
- fm := flag.Args()[0]
- gms, err := listAllGMs(fm)
- if err != nil {
- log.Fatalln("Could not query", fm, "for GMs:", err)
- }
- // One job can comes right on the command line,
- // and any number can come one per line from -script.
- jobs := [][]string{flag.Args()[1:]}
- if *script != "" {
- file := os.Stdin
- if *script != "-" {
- file, err = os.Open(*script)
- if err != nil {
- log.Fatal(err)
- }
- defer file.Close()
- }
- scanner := bufio.NewScanner(file)
- for scanner.Scan() {
- jobs = append(jobs, strings.Fields(scanner.Text()))
- }
- if err = scanner.Err(); err != nil {
- log.Fatal(err)
- }
- }
- wg := &sync.WaitGroup{}
- var failures int32 = 0
- worker := func(queue chan work) {
- for w := range queue {
- start := time.Now()
- args := w.Flags[:]
- args = append(args, "-s")
- args = append(args, w.Sources...)
- cmd := exec.Command(fm, args...)
- output, err := cmd.CombinedOutput()
- status := "#done"
- if err != nil {
- status = fmt.Sprintf("#failed (%v)", err)
- if len(w.Sources) == 1 {
- // If a source ran alone and failed, that's just a failure.
- atomic.AddInt32(&failures, 1)
- } else {
- // If a batch of sources ran and failed, split them up and try again.
- for _, source := range w.Sources {
- wg.Add(1)
- queue <- work{[]string{source}, w.Flags}
- }
- }
- }
- if !*quiet || (err != nil && len(w.Sources) == 1) {
- log.Printf("\n%v %v in %v:\n%s",
- strings.Join(cmd.Args, " "), status, time.Since(start), output)
- }
- wg.Done()
- }
- }
- cpu := make(chan work, *cpuLimit)
- for i := 0; i < *cpuLimit; i++ {
- go worker(cpu)
- }
- gpu := make(chan work, *gpuLimit)
- for i := 0; i < *gpuLimit; i++ {
- go worker(gpu)
- }
- for _, job := range jobs {
- // Skip blank lines, empty command lines.
- if len(job) == 0 {
- continue
- }
- w, err := parseWork(job, gms)
- if err != nil {
- log.Fatal(err)
- }
- // Determine if this is CPU-bound or GPU-bound work, conservatively assuming GPU.
- queue, limit := gpu, *gpuLimit
- backend := ""
- for i, flag := range w.Flags {
- if flag == "-b" || flag == "--backend" {
- backend = w.Flags[i+1]
- }
- }
- whitelisted := map[string]bool{
- "cpu": true,
- "skp": true,
- "pdf": true,
- }
- if whitelisted[backend] {
- queue, limit = cpu, *cpuLimit
- }
- if *random {
- rand.Shuffle(len(w.Sources), func(i, j int) {
- w.Sources[i], w.Sources[j] = w.Sources[j], w.Sources[i]
- })
- }
- // Round up so there's at least one source per batch.
- sourcesPerBatch := (len(w.Sources) + limit - 1) / limit
- for i := 0; i < len(w.Sources); i += sourcesPerBatch {
- end := i + sourcesPerBatch
- if end > len(w.Sources) {
- end = len(w.Sources)
- }
- batch := w.Sources[i:end]
- wg.Add(1)
- queue <- work{batch, w.Flags}
- }
- }
- wg.Wait()
- if failures > 0 {
- log.Fatalln(failures, "failures after retries")
- }
- }
|