X-Git-Url: https://git.arvados.org/arvados.git/blobdiff_plain/83a4494e66f4f7447091779f25e6f202b2379de9..dfe0ec7bfec3fd72cd40d3962e5c8af08d2413d2:/services/crunch-dispatch-slurm/crunch-dispatch-slurm.go

diff --git a/services/crunch-dispatch-slurm/crunch-dispatch-slurm.go b/services/crunch-dispatch-slurm/crunch-dispatch-slurm.go
index 740df55ecd..0c1ce49592 100644
--- a/services/crunch-dispatch-slurm/crunch-dispatch-slurm.go
+++ b/services/crunch-dispatch-slurm/crunch-dispatch-slurm.go
@@ -3,12 +3,14 @@ package main
 // Dispatcher service for Crunch that submits containers to the slurm queue.
 
 import (
-	"encoding/json"
 	"flag"
 	"fmt"
 	"git.curoverse.com/arvados.git/sdk/go/arvados"
 	"git.curoverse.com/arvados.git/sdk/go/arvadosclient"
+	"git.curoverse.com/arvados.git/sdk/go/config"
 	"git.curoverse.com/arvados.git/sdk/go/dispatch"
+	"github.com/coreos/go-systemd/daemon"
+	"io"
 	"io/ioutil"
 	"log"
 	"math"
@@ -20,52 +22,73 @@ import (
 
 // Config used by crunch-dispatch-slurm
 type Config struct {
-	SbatchArguments  []string
-	PollPeriod       *time.Duration
-	CrunchRunCommand *string
+	Client arvados.Client
+
+	SbatchArguments []string
+	PollPeriod      arvados.Duration
+
+	// crunch-run command to invoke. The container UUID will be
+	// appended. If nil, []string{"crunch-run"} will be used.
+	//
+	// Example: []string{"crunch-run", "--cgroup-parent-subsystem=memory"}
+	CrunchRunCommand []string
 }
 
 func main() {
 	err := doMain()
 	if err != nil {
-		log.Fatalf("%q", err)
+		log.Fatal(err)
 	}
 }
 
 var (
-	config        Config
+	theConfig     Config
 	squeueUpdater Squeue
 )
 
-const defaultConfigPath = "/etc/arvados/crunch-dispatch-slurm/config.json"
+const defaultConfigPath = "/etc/arvados/crunch-dispatch-slurm/crunch-dispatch-slurm.yml"
 
 func doMain() error {
 	flags := flag.NewFlagSet("crunch-dispatch-slurm", flag.ExitOnError)
+	flags.Usage = func() { usage(flags) }
 
 	configPath := flags.String(
 		"config",
 		defaultConfigPath,
-		"`path` to json configuration file")
-
-	config.PollPeriod = flags.Duration(
-		"poll-interval",
-		10*time.Second,
-		"Time duration to poll for queued containers")
-
-	config.CrunchRunCommand = flags.String(
-		"crunch-run-command",
-		"/usr/bin/crunch-run",
-		"Crunch command to run container")
+		"`path` to JSON or YAML configuration file")
 
 	// Parse args; omit the first arg which is the command name
 	flags.Parse(os.Args[1:])
 
-	err := readConfig(&config, *configPath)
+	err := readConfig(&theConfig, *configPath)
 	if err != nil {
-		log.Printf("Error reading configuration: %v", err)
 		return err
 	}
 
+	if theConfig.CrunchRunCommand == nil {
+		theConfig.CrunchRunCommand = []string{"crunch-run"}
+	}
+
+	if theConfig.PollPeriod == 0 {
+		theConfig.PollPeriod = arvados.Duration(10 * time.Second)
+	}
+
+	if theConfig.Client.APIHost != "" || theConfig.Client.AuthToken != "" {
+		// Copy real configs into env vars so [a]
+		// MakeArvadosClient() uses them, and [b] they get
+		// propagated to crunch-run via SLURM.
+		os.Setenv("ARVADOS_API_HOST", theConfig.Client.APIHost)
+		os.Setenv("ARVADOS_API_TOKEN", theConfig.Client.AuthToken)
+		os.Setenv("ARVADOS_API_HOST_INSECURE", "")
+		if theConfig.Client.Insecure {
+			os.Setenv("ARVADOS_API_HOST_INSECURE", "1")
+		}
+		os.Setenv("ARVADOS_KEEP_SERVICES", strings.Join(theConfig.Client.KeepServiceURIs, " "))
+		os.Setenv("ARVADOS_EXTERNAL_CLIENT", "")
+	} else {
+		log.Printf("warning: Client credentials missing from config, so falling back on environment variables (deprecated).")
+	}
+
 	arv, err := arvadosclient.MakeArvadosClient()
 	if err != nil {
 		log.Printf("Error making Arvados client: %v", err)
@@ -73,15 +96,19 @@ func doMain() error {
 	}
 	arv.Retries = 25
 
-	squeueUpdater.StartMonitor(*config.PollPeriod)
+	squeueUpdater.StartMonitor(time.Duration(theConfig.PollPeriod))
 	defer squeueUpdater.Done()
 
 	dispatcher := dispatch.Dispatcher{
 		Arv:            arv,
 		RunContainer:   run,
-		PollInterval:   *config.PollPeriod,
+		PollInterval:   time.Duration(theConfig.PollPeriod),
 		DoneProcessing: make(chan struct{})}
 
+	if _, err := daemon.SdNotify("READY=1"); err != nil {
+		log.Printf("Error notifying init daemon: %v", err)
+	}
+
 	err = dispatcher.RunDispatcher()
 	if err != nil {
 		return err
@@ -96,10 +123,13 @@ func sbatchFunc(container arvados.Container) *exec.Cmd {
 
 	var sbatchArgs []string
 	sbatchArgs = append(sbatchArgs, "--share")
-	sbatchArgs = append(sbatchArgs, config.SbatchArguments...)
+	sbatchArgs = append(sbatchArgs, theConfig.SbatchArguments...)
 	sbatchArgs = append(sbatchArgs, fmt.Sprintf("--job-name=%s", container.UUID))
 	sbatchArgs = append(sbatchArgs, fmt.Sprintf("--mem-per-cpu=%d", int(memPerCPU)))
 	sbatchArgs = append(sbatchArgs, fmt.Sprintf("--cpus-per-task=%d", container.RuntimeConstraints.VCPUs))
+	if container.RuntimeConstraints.Partition != nil {
+		sbatchArgs = append(sbatchArgs, fmt.Sprintf("--partition=%s", strings.Join(container.RuntimeConstraints.Partition, ",")))
+	}
 
 	return exec.Command("sbatch", sbatchArgs...)
 }
@@ -115,7 +145,7 @@ var scancelCmd = scancelFunc
 
 // Submit job to slurm using sbatch.
 func submit(dispatcher *dispatch.Dispatcher,
-	container arvados.Container, crunchRunCommand string) (submitErr error) {
+	container arvados.Container, crunchRunCommand []string) (submitErr error) {
 	defer func() {
 		// If we didn't get as far as submitting a slurm job,
 		// unlock the container and return it to the queue.
@@ -123,10 +153,7 @@ func submit(dispatcher *dispatch.Dispatcher,
 			// OK, no cleanup needed
 			return
 		}
-		err := dispatcher.Arv.Update("containers", container.UUID,
-			arvadosclient.Dict{
-				"container": arvadosclient.Dict{"state": "Queued"}},
-			nil)
+		err := dispatcher.Unlock(container.UUID)
 		if err != nil {
 			log.Printf("Error unlocking container %s: %v", container.UUID, err)
 		}
@@ -156,9 +183,10 @@ func submit(dispatcher *dispatch.Dispatcher,
 	squeueUpdater.SlurmLock.Lock()
 	defer squeueUpdater.SlurmLock.Unlock()
 
+	log.Printf("sbatch starting: %+q", cmd.Args)
 	err := cmd.Start()
 	if err != nil {
-		submitErr = fmt.Errorf("Error starting %v: %v", cmd.Args, err)
+		submitErr = fmt.Errorf("Error starting sbatch: %v", err)
 		return
 	}
 
@@ -178,7 +206,7 @@ func submit(dispatcher *dispatch.Dispatcher,
 
 	// Send a tiny script on stdin to execute the crunch-run command
 	// slurm actually enforces that this must be a #! script
-	fmt.Fprintf(stdinWriter, "#!/bin/sh\nexec '%s' '%s'\n", crunchRunCommand, container.UUID)
+	io.WriteString(stdinWriter, execScript(append(crunchRunCommand, container.UUID)))
 	stdinWriter.Close()
 
 	err = cmd.Wait()
@@ -215,11 +243,11 @@ func monitorSubmitOrCancel(dispatcher *dispatch.Dispatcher, container arvados.Co
 
 			log.Printf("About to submit queued container %v", container.UUID)
 
-			if err := submit(dispatcher, container, *config.CrunchRunCommand); err != nil {
+			if err := submit(dispatcher, container, theConfig.CrunchRunCommand); err != nil {
 				log.Printf("Error submitting container %s to slurm: %v",
 					container.UUID, err)
 				// maybe sbatch is broken, put it back to queued
-				dispatcher.UpdateState(container.UUID, dispatch.Queued)
+				dispatcher.Unlock(container.UUID)
 			}
 			submitted = true
 		} else {
@@ -235,20 +263,20 @@ func monitorSubmitOrCancel(dispatcher *dispatch.Dispatcher, container arvados.Co
 				log.Printf("Error getting final container state: %v", err)
 			}
 
-			var st arvados.ContainerState
 			switch con.State {
 			case dispatch.Locked:
-				st = dispatch.Queued
+				log.Printf("Container %s in state %v but missing from slurm queue, changing to %v.",
+					container.UUID, con.State, dispatch.Queued)
+				dispatcher.Unlock(container.UUID)
 			case dispatch.Running:
-				st = dispatch.Cancelled
+				st := dispatch.Cancelled
+				log.Printf("Container %s in state %v but missing from slurm queue, changing to %v.",
+					container.UUID, con.State, st)
+				dispatcher.UpdateState(container.UUID, st)
 			default:
 				// Container state is Queued, Complete or Cancelled so stop monitoring it.
 				return
 			}
-
-			log.Printf("Container %s in state %v but missing from slurm queue, changing to %v.",
-				container.UUID, con.State, st)
-			dispatcher.UpdateState(container.UUID, st)
 		}
 	}
 }
@@ -295,16 +323,10 @@ func run(dispatcher *dispatch.Dispatcher,
 }
 
 func readConfig(dst interface{}, path string) error {
-	if buf, err := ioutil.ReadFile(path); err != nil && os.IsNotExist(err) {
-		if path == defaultConfigPath {
-			log.Printf("Config not specified. Continue with default configuration.")
-		} else {
-			return fmt.Errorf("Config file not found %q: %v", path, err)
-		}
-	} else if err != nil {
-		return fmt.Errorf("Error reading config %q: %v", path, err)
-	} else if err = json.Unmarshal(buf, dst); err != nil {
-		return fmt.Errorf("Error decoding config %q: %v", path, err)
+	err := config.LoadFile(dst, path)
+	if err != nil && os.IsNotExist(err) && path == defaultConfigPath {
+		log.Printf("Config not specified. Continue with default configuration.")
+		err = nil
 	}
-	return nil
+	return err
 }