fm_bot.go 6.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281
  1. // Copyright 2019 Google LLC.
  2. // Use of this source code is governed by a BSD-style license that can be found in the LICENSE file.
  3. package main
  4. import (
  5. "bufio"
  6. "bytes"
  7. "flag"
  8. "fmt"
  9. "log"
  10. "math/rand"
  11. "os"
  12. "os/exec"
  13. "path/filepath"
  14. "runtime"
  15. "strings"
  16. "sync"
  17. "sync/atomic"
  18. "time"
  19. )
  20. // Too many GPU processes and we'll start to overwhelm your GPU,
  21. // even hanging your machine in the worst case. Here's a reasonable default.
  22. func defaultGpuLimit() int {
  23. limit := 8
  24. if n := runtime.NumCPU(); n < limit {
  25. return n
  26. }
  27. return limit
  28. }
  29. var script = flag.String("script", "", "A file with jobs to run, one per line. - for stdin.")
  30. var random = flag.Bool("random", true, "Assign sources into job batches randomly?")
  31. var quiet = flag.Bool("quiet", false, "Print only failures?")
  32. var exact = flag.Bool("exact", false, "Match GM names only exactly.")
  33. var cpuLimit = flag.Int("cpuLimit", runtime.NumCPU(),
  34. "Maximum number of concurrent processes for CPU-bound work.")
  35. var gpuLimit = flag.Int("gpuLimit", defaultGpuLimit(),
  36. "Maximum number of concurrent processes for GPU-bound work.")
  37. func init() {
  38. flag.StringVar(script, "s", *script, "Alias for --script.")
  39. flag.BoolVar(random, "r", *random, "Alias for --random.")
  40. flag.BoolVar(quiet, "q", *quiet, "Alias for --quiet.")
  41. flag.BoolVar(exact, "e", *exact, "Alias for --exact.")
  42. flag.IntVar(cpuLimit, "c", *cpuLimit, "Alias for --cpuLimit.")
  43. flag.IntVar(gpuLimit, "g", *gpuLimit, "Alias for --gpuLimit.")
  44. }
  45. func listAllGMs(fm string) (gms []string, err error) {
  46. // Query fm binary for list of all available GMs by running with no arguments.
  47. cmd := exec.Command(fm)
  48. stdout, err := cmd.Output()
  49. if err != nil {
  50. return
  51. }
  52. // GM names are listed line-by-line.
  53. scanner := bufio.NewScanner(bytes.NewReader(stdout))
  54. for scanner.Scan() {
  55. gms = append(gms, scanner.Text())
  56. }
  57. err = scanner.Err()
  58. return
  59. }
  60. type work struct {
  61. Sources []string
  62. Flags []string
  63. }
  64. func parseWork(args []string, gms []string) (*work, error) {
  65. w := &work{}
  66. for _, arg := range args {
  67. // I wish we could parse flags here too, but it's too late.
  68. if strings.HasPrefix(arg, "-") {
  69. msg := "Is '%s' an fm flag? If so please pass it using flag=value syntax."
  70. if flag.Lookup(arg[1:]) != nil {
  71. msg = "Please pass fm_bot flags like '%s' on the command line before the FM binary."
  72. }
  73. return nil, fmt.Errorf(msg, arg)
  74. }
  75. // Everything after a # is a comment.
  76. if strings.HasPrefix(arg, "#") {
  77. break
  78. }
  79. // Treat "gm" or "gms" as a shortcut for all known GMs.
  80. if arg == "gm" || arg == "gms" {
  81. w.Sources = append(w.Sources, gms...)
  82. continue
  83. }
  84. // Is this an option to pass through to fm?
  85. if parts := strings.Split(arg, "="); len(parts) == 2 {
  86. f := "-"
  87. if len(parts[0]) > 1 {
  88. f += "-"
  89. }
  90. f += parts[0]
  91. w.Flags = append(w.Flags, f, parts[1])
  92. continue
  93. }
  94. // Is this argument naming a GM?
  95. matchedAnyGM := false
  96. for _, gm := range gms {
  97. if (*exact && gm == arg) || (!*exact && strings.Contains(gm, arg)) {
  98. w.Sources = append(w.Sources, gm)
  99. matchedAnyGM = true
  100. }
  101. }
  102. if matchedAnyGM {
  103. continue
  104. }
  105. // Anything left ought to be on the file system: a file, a directory, or a glob.
  106. // Not all shells expand globs, so we'll do it here just in case.
  107. matches, err := filepath.Glob(arg)
  108. if err != nil {
  109. return nil, err
  110. }
  111. if len(matches) == 0 {
  112. return nil, fmt.Errorf("Don't understand '%s'.", arg)
  113. }
  114. for _, match := range matches {
  115. err := filepath.Walk(match, func(path string, info os.FileInfo, err error) error {
  116. if !info.IsDir() {
  117. w.Sources = append(w.Sources, path)
  118. }
  119. return err
  120. })
  121. if err != nil {
  122. return nil, err
  123. }
  124. }
  125. }
  126. return w, nil
  127. }
  128. func main() {
  129. flag.Parse()
  130. if flag.NArg() < 1 {
  131. log.Fatal("Please pass an fm binary as the first argument.")
  132. }
  133. fm := flag.Args()[0]
  134. gms, err := listAllGMs(fm)
  135. if err != nil {
  136. log.Fatalln("Could not query", fm, "for GMs:", err)
  137. }
  138. // One job can comes right on the command line,
  139. // and any number can come one per line from -script.
  140. jobs := [][]string{flag.Args()[1:]}
  141. if *script != "" {
  142. file := os.Stdin
  143. if *script != "-" {
  144. file, err = os.Open(*script)
  145. if err != nil {
  146. log.Fatal(err)
  147. }
  148. defer file.Close()
  149. }
  150. scanner := bufio.NewScanner(file)
  151. for scanner.Scan() {
  152. jobs = append(jobs, strings.Fields(scanner.Text()))
  153. }
  154. if err = scanner.Err(); err != nil {
  155. log.Fatal(err)
  156. }
  157. }
  158. wg := &sync.WaitGroup{}
  159. var failures int32 = 0
  160. worker := func(queue chan work) {
  161. for w := range queue {
  162. start := time.Now()
  163. args := w.Flags[:]
  164. args = append(args, "-s")
  165. args = append(args, w.Sources...)
  166. cmd := exec.Command(fm, args...)
  167. output, err := cmd.CombinedOutput()
  168. status := "#done"
  169. if err != nil {
  170. status = fmt.Sprintf("#failed (%v)", err)
  171. if len(w.Sources) == 1 {
  172. // If a source ran alone and failed, that's just a failure.
  173. atomic.AddInt32(&failures, 1)
  174. } else {
  175. // If a batch of sources ran and failed, split them up and try again.
  176. for _, source := range w.Sources {
  177. wg.Add(1)
  178. queue <- work{[]string{source}, w.Flags}
  179. }
  180. }
  181. }
  182. if !*quiet || (err != nil && len(w.Sources) == 1) {
  183. log.Printf("\n%v %v in %v:\n%s",
  184. strings.Join(cmd.Args, " "), status, time.Since(start), output)
  185. }
  186. wg.Done()
  187. }
  188. }
  189. cpu := make(chan work, *cpuLimit)
  190. for i := 0; i < *cpuLimit; i++ {
  191. go worker(cpu)
  192. }
  193. gpu := make(chan work, *gpuLimit)
  194. for i := 0; i < *gpuLimit; i++ {
  195. go worker(gpu)
  196. }
  197. for _, job := range jobs {
  198. // Skip blank lines, empty command lines.
  199. if len(job) == 0 {
  200. continue
  201. }
  202. w, err := parseWork(job, gms)
  203. if err != nil {
  204. log.Fatal(err)
  205. }
  206. // Determine if this is CPU-bound or GPU-bound work, conservatively assuming GPU.
  207. queue, limit := gpu, *gpuLimit
  208. backend := ""
  209. for i, flag := range w.Flags {
  210. if flag == "-b" || flag == "--backend" {
  211. backend = w.Flags[i+1]
  212. }
  213. }
  214. whitelisted := map[string]bool{
  215. "cpu": true,
  216. "skp": true,
  217. "pdf": true,
  218. }
  219. if whitelisted[backend] {
  220. queue, limit = cpu, *cpuLimit
  221. }
  222. if *random {
  223. rand.Shuffle(len(w.Sources), func(i, j int) {
  224. w.Sources[i], w.Sources[j] = w.Sources[j], w.Sources[i]
  225. })
  226. }
  227. // Round up so there's at least one source per batch.
  228. sourcesPerBatch := (len(w.Sources) + limit - 1) / limit
  229. for i := 0; i < len(w.Sources); i += sourcesPerBatch {
  230. end := i + sourcesPerBatch
  231. if end > len(w.Sources) {
  232. end = len(w.Sources)
  233. }
  234. batch := w.Sources[i:end]
  235. wg.Add(1)
  236. queue <- work{batch, w.Flags}
  237. }
  238. }
  239. wg.Wait()
  240. if failures > 0 {
  241. log.Fatalln(failures, "failures after retries")
  242. }
  243. }