X-Git-Url: https://git.arvados.org/arvados.git/blobdiff_plain/d77c4cc58d393c48ce46b987f6eada7c7cc381c6..4dde672661ea9dca680ec9eb1cdba7bd7d87fca7:/services/crunch-dispatch-slurm/squeue.go

diff --git a/services/crunch-dispatch-slurm/squeue.go b/services/crunch-dispatch-slurm/squeue.go
index b86a4d9568..b8e3108c7c 100644
--- a/services/crunch-dispatch-slurm/squeue.go
+++ b/services/crunch-dispatch-slurm/squeue.go
@@ -1,120 +1,181 @@
+// Copyright (C) The Arvados Authors. All rights reserved.
+//
+// SPDX-License-Identifier: AGPL-3.0
+
 package main
 
 import (
-	"bufio"
+	"bytes"
+	"fmt"
 	"log"
-	"os/exec"
+	"sort"
+	"strings"
 	"sync"
 	"time"
 )
 
-type Squeue struct {
-	squeueContents []string
-	squeueDone     chan struct{}
-	squeueError    error
-	squeueCond     *sync.Cond
-	SlurmLock      sync.Mutex
+type slurmJob struct {
+	uuid         string
+	wantPriority int64
+	priority     int64 // current slurm priority (incorporates nice value)
+	nice         int64 // current slurm nice value
 }
 
-// squeueFunc
-func squeueFunc() *exec.Cmd {
-	return exec.Command("squeue", "--format=%j")
+// Squeue implements asynchronous polling monitor of the SLURM queue using the
+// command 'squeue'.
+type SqueueChecker struct {
+	Period         time.Duration
+	PrioritySpread int64
+	Slurm          Slurm
+	queue          map[string]*slurmJob
+	startOnce      sync.Once
+	done           chan struct{}
+	sync.Cond
 }
 
-var squeueCmd = squeueFunc
+// HasUUID checks if a given container UUID is in the slurm queue.
+// This does not run squeue directly, but instead blocks until woken
+// up by next successful update of squeue.
+func (sqc *SqueueChecker) HasUUID(uuid string) bool {
+	sqc.startOnce.Do(sqc.start)
 
-func (squeue *Squeue) RunSqueue() error {
-	var newSqueueContents []string
+	sqc.L.Lock()
+	defer sqc.L.Unlock()
 
-	// Mutex between squeue sync and running sbatch or scancel.  This
-	// establishes a sequence so that squeue doesn't run concurrently with
-	// sbatch or scancel; the next update of squeue will occur only after
-	// sbatch or scancel has completed.
-	squeueUpdater.SlurmLock.Lock()
-	defer squeueUpdater.SlurmLock.Unlock()
-
-	// Also ensure unlock on all return paths
-	defer squeueUpdater.squeueCond.L.Unlock()
-
-	cmd := squeueCmd()
-	sq, err := cmd.StdoutPipe()
-	if err != nil {
-		log.Printf("Error creating stdout pipe for squeue: %v", err)
-		squeueUpdater.squeueCond.L.Lock()
-		squeueUpdater.squeueError = err
-		return err
-	}
-	cmd.Start()
-	scanner := bufio.NewScanner(sq)
-	for scanner.Scan() {
-		newSqueueContents = append(newSqueueContents, scanner.Text())
-	}
-	if err := scanner.Err(); err != nil {
-		cmd.Wait()
-		log.Printf("Error reading from squeue pipe: %v", err)
-		squeueUpdater.squeueCond.L.Lock()
-		squeueUpdater.squeueError = err
-		return err
-	}
+	// block until next squeue broadcast signaling an update.
+	sqc.Wait()
+	_, exists := sqc.queue[uuid]
+	return exists
+}
 
-	err = cmd.Wait()
-	if err != nil {
-		log.Printf("Error running squeue: %v", err)
-		squeueUpdater.squeueCond.L.Lock()
-		squeueUpdater.squeueError = err
-		return err
+// SetPriority sets or updates the desired (Arvados) priority for a
+// container.
+func (sqc *SqueueChecker) SetPriority(uuid string, want int64) {
+	sqc.startOnce.Do(sqc.start)
+	sqc.L.Lock()
+	defer sqc.L.Unlock()
+	job, ok := sqc.queue[uuid]
+	if !ok {
+		// Wait in case the slurm job was just submitted and
+		// will appear in the next squeue update.
+		sqc.Wait()
+		if job, ok = sqc.queue[uuid]; !ok {
+			return
+		}
 	}
-
-	squeueUpdater.squeueCond.L.Lock()
-	squeueUpdater.squeueError = nil
-	squeueUpdater.squeueContents = newSqueueContents
-	squeueUpdater.squeueCond.Broadcast()
-
-	return nil
+	job.wantPriority = want
 }
 
-// Check if a container UUID is in the slurm queue.  This will block until the
-// next successful update from SLURM.
-func (squeue *Squeue) CheckSqueue(uuid string) (bool, error) {
-	squeueUpdater.squeueCond.L.Lock()
-	// block until next squeue broadcast signaling an update.
-	squeueUpdater.squeueCond.Wait()
-	if squeueUpdater.squeueError != nil {
-		e := squeueUpdater.squeueError
-		squeueUpdater.squeueCond.L.Unlock()
-		return false, e
+// adjust slurm job nice values as needed to ensure slurm priority
+// order matches Arvados priority order.
+func (sqc *SqueueChecker) reniceAll() {
+	sqc.L.Lock()
+	defer sqc.L.Unlock()
+
+	jobs := make([]*slurmJob, 0, len(sqc.queue))
+	for _, j := range sqc.queue {
+		if j.wantPriority == 0 {
+			// SLURM job with unknown Arvados priority
+			// (perhaps it's not an Arvados job)
+			continue
+		}
+		jobs = append(jobs, j)
 	}
-	contents := squeueUpdater.squeueContents
-	squeueUpdater.squeueCond.L.Unlock()
 
-	for _, k := range contents {
-		if k == uuid {
-			return true, nil
+	sort.Slice(jobs, func(i, j int) bool {
+		return jobs[i].wantPriority > jobs[j].wantPriority
+	})
+	renice := wantNice(jobs, sqc.PrioritySpread)
+	for i, job := range jobs {
+		if renice[i] == job.nice {
+			continue
 		}
+		log.Printf("updating slurm priority for %q: nice %d => %d", job.uuid, job.nice, renice[i])
+		sqc.Slurm.Renice(job.uuid, renice[i])
 	}
-	return false, nil
 }
 
-func (squeue *Squeue) StartMonitor(pollInterval time.Duration) {
-	squeueUpdater.squeueCond = sync.NewCond(&sync.Mutex{})
-	squeueUpdater.squeueDone = make(chan struct{})
-	squeueUpdater.RunSqueue()
-	go squeueUpdater.SyncSqueue(pollInterval)
+// Stop stops the squeue monitoring goroutine. Do not call HasUUID
+// after calling Stop.
+func (sqc *SqueueChecker) Stop() {
+	if sqc.done != nil {
+		close(sqc.done)
+	}
 }
 
-func (squeue *Squeue) Done() {
-	squeueUpdater.squeueDone <- struct{}{}
-	close(squeueUpdater.squeueDone)
+// check gets the names of jobs in the SLURM queue (running and
+// queued). If it succeeds, it updates sqc.queue and wakes up any
+// goroutines that are waiting in HasUUID() or All().
+func (sqc *SqueueChecker) check() {
+	// Mutex between squeue sync and running sbatch or scancel.  This
+	// establishes a sequence so that squeue doesn't run concurrently with
+	// sbatch or scancel; the next update of squeue will occur only after
+	// sbatch or scancel has completed.
+	sqc.L.Lock()
+	defer sqc.L.Unlock()
+
+	cmd := sqc.Slurm.QueueCommand([]string{"--all", "--format=%j %y %Q"})
+	stdout, stderr := &bytes.Buffer{}, &bytes.Buffer{}
+	cmd.Stdout, cmd.Stderr = stdout, stderr
+	if err := cmd.Run(); err != nil {
+		log.Printf("Error running %q %q: %s %q", cmd.Path, cmd.Args, err, stderr.String())
+		return
+	}
+
+	lines := strings.Split(stdout.String(), "\n")
+	newq := make(map[string]*slurmJob, len(lines))
+	for _, line := range lines {
+		if line == "" {
+			continue
+		}
+		var uuid string
+		var n, p int64
+		if _, err := fmt.Sscan(line, &uuid, &n, &p); err != nil {
+			log.Printf("warning: ignoring unparsed line in squeue output: %q", line)
+			continue
+		}
+		replacing, ok := sqc.queue[uuid]
+		if !ok {
+			replacing = &slurmJob{uuid: uuid}
+		}
+		replacing.priority = p
+		replacing.nice = n
+		newq[uuid] = replacing
+	}
+	sqc.queue = newq
+	sqc.Broadcast()
 }
 
-func (squeue *Squeue) SyncSqueue(pollInterval time.Duration) {
-	ticker := time.NewTicker(pollInterval)
-	for {
-		select {
-		case <-squeueUpdater.squeueDone:
-			return
-		case <-ticker.C:
-			squeueUpdater.RunSqueue()
+// Initialize, and start a goroutine to call check() once per
+// squeue.Period until terminated by calling Stop().
+func (sqc *SqueueChecker) start() {
+	sqc.L = &sync.Mutex{}
+	sqc.done = make(chan struct{})
+	go func() {
+		ticker := time.NewTicker(sqc.Period)
+		for {
+			select {
+			case <-sqc.done:
+				ticker.Stop()
+				return
+			case <-ticker.C:
+				sqc.check()
+				sqc.reniceAll()
+			}
 		}
+	}()
+}
+
+// All waits for the next squeue invocation, and returns all job
+// names reported by squeue.
+func (sqc *SqueueChecker) All() []string {
+	sqc.startOnce.Do(sqc.start)
+	sqc.L.Lock()
+	defer sqc.L.Unlock()
+	sqc.Wait()
+	var uuids []string
+	for u := range sqc.queue {
+		uuids = append(uuids, u)
 	}
+	return uuids
 }