10868: Add fake uuid to test case for more realistic log messages.
[arvados.git] / services / crunch-dispatch-slurm / crunch-dispatch-slurm.go
1 package main
2
3 // Dispatcher service for Crunch that submits containers to the slurm queue.
4
5 import (
6         "flag"
7         "fmt"
8         "git.curoverse.com/arvados.git/sdk/go/arvados"
9         "git.curoverse.com/arvados.git/sdk/go/arvadosclient"
10         "git.curoverse.com/arvados.git/sdk/go/config"
11         "git.curoverse.com/arvados.git/sdk/go/dispatch"
12         "github.com/coreos/go-systemd/daemon"
13         "io"
14         "io/ioutil"
15         "log"
16         "math"
17         "os"
18         "os/exec"
19         "strings"
20         "time"
21 )
22
23 // Config used by crunch-dispatch-slurm
24 type Config struct {
25         Client arvados.Client
26
27         SbatchArguments []string
28         PollPeriod      arvados.Duration
29
30         // crunch-run command to invoke. The container UUID will be
31         // appended. If nil, []string{"crunch-run"} will be used.
32         //
33         // Example: []string{"crunch-run", "--cgroup-parent-subsystem=memory"}
34         CrunchRunCommand []string
35 }
36
37 func main() {
38         err := doMain()
39         if err != nil {
40                 log.Fatal(err)
41         }
42 }
43
44 var (
45         theConfig     Config
46         squeueUpdater Squeue
47 )
48
49 const defaultConfigPath = "/etc/arvados/crunch-dispatch-slurm/crunch-dispatch-slurm.yml"
50
51 func doMain() error {
52         flags := flag.NewFlagSet("crunch-dispatch-slurm", flag.ExitOnError)
53         flags.Usage = func() { usage(flags) }
54
55         configPath := flags.String(
56                 "config",
57                 defaultConfigPath,
58                 "`path` to JSON or YAML configuration file")
59
60         // Parse args; omit the first arg which is the command name
61         flags.Parse(os.Args[1:])
62
63         err := readConfig(&theConfig, *configPath)
64         if err != nil {
65                 return err
66         }
67
68         if theConfig.CrunchRunCommand == nil {
69                 theConfig.CrunchRunCommand = []string{"crunch-run"}
70         }
71
72         if theConfig.PollPeriod == 0 {
73                 theConfig.PollPeriod = arvados.Duration(10 * time.Second)
74         }
75
76         if theConfig.Client.APIHost != "" || theConfig.Client.AuthToken != "" {
77                 // Copy real configs into env vars so [a]
78                 // MakeArvadosClient() uses them, and [b] they get
79                 // propagated to crunch-run via SLURM.
80                 os.Setenv("ARVADOS_API_HOST", theConfig.Client.APIHost)
81                 os.Setenv("ARVADOS_API_TOKEN", theConfig.Client.AuthToken)
82                 os.Setenv("ARVADOS_API_HOST_INSECURE", "")
83                 if theConfig.Client.Insecure {
84                         os.Setenv("ARVADOS_API_HOST_INSECURE", "1")
85                 }
86                 os.Setenv("ARVADOS_KEEP_SERVICES", strings.Join(theConfig.Client.KeepServiceURIs, " "))
87                 os.Setenv("ARVADOS_EXTERNAL_CLIENT", "")
88         } else {
89                 log.Printf("warning: Client credentials missing from config, so falling back on environment variables (deprecated).")
90         }
91
92         arv, err := arvadosclient.MakeArvadosClient()
93         if err != nil {
94                 log.Printf("Error making Arvados client: %v", err)
95                 return err
96         }
97         arv.Retries = 25
98
99         squeueUpdater.StartMonitor(time.Duration(theConfig.PollPeriod))
100         defer squeueUpdater.Done()
101
102         dispatcher := dispatch.Dispatcher{
103                 Arv:            arv,
104                 RunContainer:   run,
105                 PollInterval:   time.Duration(theConfig.PollPeriod),
106                 DoneProcessing: make(chan struct{})}
107
108         if _, err := daemon.SdNotify(false, "READY=1"); err != nil {
109                 log.Printf("Error notifying init daemon: %v", err)
110         }
111
112         err = dispatcher.RunDispatcher()
113         if err != nil {
114                 return err
115         }
116
117         return nil
118 }
119
120 // sbatchCmd
121 func sbatchFunc(container arvados.Container) *exec.Cmd {
122         memPerCPU := math.Ceil(float64(container.RuntimeConstraints.RAM) / (float64(container.RuntimeConstraints.VCPUs) * 1048576))
123
124         var sbatchArgs []string
125         sbatchArgs = append(sbatchArgs, "--share")
126         sbatchArgs = append(sbatchArgs, theConfig.SbatchArguments...)
127         sbatchArgs = append(sbatchArgs, fmt.Sprintf("--job-name=%s", container.UUID))
128         sbatchArgs = append(sbatchArgs, fmt.Sprintf("--mem-per-cpu=%d", int(memPerCPU)))
129         sbatchArgs = append(sbatchArgs, fmt.Sprintf("--cpus-per-task=%d", container.RuntimeConstraints.VCPUs))
130         if container.SchedulingParameters.Partitions != nil {
131                 sbatchArgs = append(sbatchArgs, fmt.Sprintf("--partition=%s", strings.Join(container.SchedulingParameters.Partitions, ",")))
132         }
133
134         return exec.Command("sbatch", sbatchArgs...)
135 }
136
137 // scancelCmd
138 func scancelFunc(container arvados.Container) *exec.Cmd {
139         return exec.Command("scancel", "--name="+container.UUID)
140 }
141
142 // Wrap these so that they can be overridden by tests
143 var sbatchCmd = sbatchFunc
144 var scancelCmd = scancelFunc
145
146 // Submit job to slurm using sbatch.
147 func submit(dispatcher *dispatch.Dispatcher,
148         container arvados.Container, crunchRunCommand []string) (submitErr error) {
149         defer func() {
150                 // If we didn't get as far as submitting a slurm job,
151                 // unlock the container and return it to the queue.
152                 if submitErr == nil {
153                         // OK, no cleanup needed
154                         return
155                 }
156                 err := dispatcher.Unlock(container.UUID)
157                 if err != nil {
158                         log.Printf("Error unlocking container %s: %v", container.UUID, err)
159                 }
160         }()
161
162         // Create the command and attach to stdin/stdout
163         cmd := sbatchCmd(container)
164         stdinWriter, stdinerr := cmd.StdinPipe()
165         if stdinerr != nil {
166                 submitErr = fmt.Errorf("Error creating stdin pipe %v: %q", container.UUID, stdinerr)
167                 return
168         }
169
170         stdoutReader, stdoutErr := cmd.StdoutPipe()
171         if stdoutErr != nil {
172                 submitErr = fmt.Errorf("Error creating stdout pipe %v: %q", container.UUID, stdoutErr)
173                 return
174         }
175
176         stderrReader, stderrErr := cmd.StderrPipe()
177         if stderrErr != nil {
178                 submitErr = fmt.Errorf("Error creating stderr pipe %v: %q", container.UUID, stderrErr)
179                 return
180         }
181
182         // Mutex between squeue sync and running sbatch or scancel.
183         squeueUpdater.SlurmLock.Lock()
184         defer squeueUpdater.SlurmLock.Unlock()
185
186         log.Printf("sbatch starting: %+q", cmd.Args)
187         err := cmd.Start()
188         if err != nil {
189                 submitErr = fmt.Errorf("Error starting sbatch: %v", err)
190                 return
191         }
192
193         stdoutChan := make(chan []byte)
194         go func() {
195                 b, _ := ioutil.ReadAll(stdoutReader)
196                 stdoutReader.Close()
197                 stdoutChan <- b
198                 close(stdoutChan)
199         }()
200
201         stderrChan := make(chan []byte)
202         go func() {
203                 b, _ := ioutil.ReadAll(stderrReader)
204                 stderrReader.Close()
205                 stderrChan <- b
206                 close(stderrChan)
207         }()
208
209         // Send a tiny script on stdin to execute the crunch-run command
210         // slurm actually enforces that this must be a #! script
211         io.WriteString(stdinWriter, execScript(append(crunchRunCommand, container.UUID)))
212         stdinWriter.Close()
213
214         stdoutMsg := <-stdoutChan
215         stderrmsg := <-stderrChan
216
217         err = cmd.Wait()
218
219         if err != nil {
220                 submitErr = fmt.Errorf("Container submission failed: %v: %v (stderr: %q)", cmd.Args, err, stderrmsg)
221                 return
222         }
223
224         log.Printf("sbatch succeeded: %s", strings.TrimSpace(string(stdoutMsg)))
225         return
226 }
227
228 // If the container is marked as Locked, check if it is already in the slurm
229 // queue.  If not, submit it.
230 //
231 // If the container is marked as Running, check if it is in the slurm queue.
232 // If not, mark it as Cancelled.
233 func monitorSubmitOrCancel(dispatcher *dispatch.Dispatcher, container arvados.Container, monitorDone *bool) {
234         submitted := false
235         for !*monitorDone {
236                 if squeueUpdater.CheckSqueue(container.UUID) {
237                         // Found in the queue, so continue monitoring
238                         submitted = true
239                 } else if container.State == dispatch.Locked && !submitted {
240                         // Not in queue but in Locked state and we haven't
241                         // submitted it yet, so submit it.
242
243                         log.Printf("About to submit queued container %v", container.UUID)
244
245                         if err := submit(dispatcher, container, theConfig.CrunchRunCommand); err != nil {
246                                 log.Printf("Error submitting container %s to slurm: %v",
247                                         container.UUID, err)
248                                 // maybe sbatch is broken, put it back to queued
249                                 dispatcher.Unlock(container.UUID)
250                         }
251                         submitted = true
252                 } else {
253                         // Not in queue and we are not going to submit it.
254                         // Refresh the container state. If it is
255                         // Complete/Cancelled, do nothing, if it is Locked then
256                         // release it back to the Queue, if it is Running then
257                         // clean up the record.
258
259                         var con arvados.Container
260                         err := dispatcher.Arv.Get("containers", container.UUID, nil, &con)
261                         if err != nil {
262                                 log.Printf("Error getting final container state: %v", err)
263                         }
264
265                         switch con.State {
266                         case dispatch.Locked:
267                                 log.Printf("Container %s in state %v but missing from slurm queue, changing to %v.",
268                                         container.UUID, con.State, dispatch.Queued)
269                                 dispatcher.Unlock(container.UUID)
270                         case dispatch.Running:
271                                 st := dispatch.Cancelled
272                                 log.Printf("Container %s in state %v but missing from slurm queue, changing to %v.",
273                                         container.UUID, con.State, st)
274                                 dispatcher.UpdateState(container.UUID, st)
275                         default:
276                                 // Container state is Queued, Complete or Cancelled so stop monitoring it.
277                                 return
278                         }
279                 }
280         }
281 }
282
283 // Run or monitor a container.
284 //
285 // Monitor status updates.  If the priority changes to zero, cancel the
286 // container using scancel.
287 func run(dispatcher *dispatch.Dispatcher,
288         container arvados.Container,
289         status chan arvados.Container) {
290
291         log.Printf("Monitoring container %v started", container.UUID)
292         defer log.Printf("Monitoring container %v finished", container.UUID)
293
294         monitorDone := false
295         go monitorSubmitOrCancel(dispatcher, container, &monitorDone)
296
297         for container = range status {
298                 if container.State == dispatch.Locked || container.State == dispatch.Running {
299                         if container.Priority == 0 {
300                                 log.Printf("Canceling container %s", container.UUID)
301
302                                 // Mutex between squeue sync and running sbatch or scancel.
303                                 squeueUpdater.SlurmLock.Lock()
304                                 cmd := scancelCmd(container)
305                                 msg, err := cmd.CombinedOutput()
306                                 squeueUpdater.SlurmLock.Unlock()
307
308                                 if err != nil {
309                                         log.Printf("Error stopping container %s with %v %v: %v %v",
310                                                 container.UUID, cmd.Path, cmd.Args, err, string(msg))
311                                         if squeueUpdater.CheckSqueue(container.UUID) {
312                                                 log.Printf("Container %s is still in squeue after scancel.",
313                                                         container.UUID)
314                                                 continue
315                                         }
316                                 }
317
318                                 err = dispatcher.UpdateState(container.UUID, dispatch.Cancelled)
319                         }
320                 }
321         }
322         monitorDone = true
323 }
324
325 func readConfig(dst interface{}, path string) error {
326         err := config.LoadFile(dst, path)
327         if err != nil && os.IsNotExist(err) && path == defaultConfigPath {
328                 log.Printf("Config not specified. Continue with default configuration.")
329                 err = nil
330         }
331         return err
332 }