lib/lsf/dispatch.go

   1 // Copyright (C) The Arvados Authors. All rights reserved.
   2 //
   3 // SPDX-License-Identifier: AGPL-3.0
   4
   5 package lsf
   6
   7 import (
   8         "context"
   9         "errors"
  10         "fmt"
  11         "math"
  12         "net/http"
  13         "regexp"
  14         "strings"
  15         "sync"
  16         "time"
  17
  18         "git.arvados.org/arvados.git/lib/cmd"
  19         "git.arvados.org/arvados.git/lib/dispatchcloud"
  20         "git.arvados.org/arvados.git/lib/service"
  21         "git.arvados.org/arvados.git/sdk/go/arvados"
  22         "git.arvados.org/arvados.git/sdk/go/arvadosclient"
  23         "git.arvados.org/arvados.git/sdk/go/auth"
  24         "git.arvados.org/arvados.git/sdk/go/ctxlog"
  25         "git.arvados.org/arvados.git/sdk/go/dispatch"
  26         "git.arvados.org/arvados.git/sdk/go/health"
  27         "github.com/julienschmidt/httprouter"
  28         "github.com/prometheus/client_golang/prometheus"
  29         "github.com/prometheus/client_golang/prometheus/promhttp"
  30         "github.com/sirupsen/logrus"
  31 )
  32
  33 var DispatchCommand cmd.Handler = service.Command(arvados.ServiceNameDispatchLSF, newHandler)
  34
  35 func newHandler(ctx context.Context, cluster *arvados.Cluster, token string, reg *prometheus.Registry) service.Handler {
  36         ac, err := arvados.NewClientFromConfig(cluster)
  37         if err != nil {
  38                 return service.ErrorHandler(ctx, cluster, fmt.Errorf("error initializing client from cluster config: %s", err))
  39         }
  40         d := &dispatcher{
  41                 Cluster:   cluster,
  42                 Context:   ctx,
  43                 ArvClient: ac,
  44                 AuthToken: token,
  45                 Registry:  reg,
  46         }
  47         go d.Start()
  48         return d
  49 }
  50
  51 type dispatcher struct {
  52         Cluster   *arvados.Cluster
  53         Context   context.Context
  54         ArvClient *arvados.Client
  55         AuthToken string
  56         Registry  *prometheus.Registry
  57
  58         logger        logrus.FieldLogger
  59         lsfcli        lsfcli
  60         lsfqueue      lsfqueue
  61         arvDispatcher *dispatch.Dispatcher
  62         httpHandler   http.Handler
  63
  64         initOnce sync.Once
  65         stop     chan struct{}
  66         stopped  chan struct{}
  67 }
  68
  69 // Start starts the dispatcher. Start can be called multiple times
  70 // with no ill effect.
  71 func (disp *dispatcher) Start() {
  72         disp.initOnce.Do(func() {
  73                 disp.init()
  74                 go func() {
  75                         disp.checkLsfQueueForOrphans()
  76                         err := disp.arvDispatcher.Run(disp.Context)
  77                         if err != nil {
  78                                 disp.logger.Error(err)
  79                                 disp.Close()
  80                         }
  81                 }()
  82         })
  83 }
  84
  85 // ServeHTTP implements service.Handler.
  86 func (disp *dispatcher) ServeHTTP(w http.ResponseWriter, r *http.Request) {
  87         disp.Start()
  88         disp.httpHandler.ServeHTTP(w, r)
  89 }
  90
  91 // CheckHealth implements service.Handler.
  92 func (disp *dispatcher) CheckHealth() error {
  93         disp.Start()
  94         select {
  95         case <-disp.stopped:
  96                 return errors.New("stopped")
  97         default:
  98                 return nil
  99         }
 100 }
 101
 102 // Done implements service.Handler.
 103 func (disp *dispatcher) Done() <-chan struct{} {
 104         return disp.stopped
 105 }
 106
 107 // Stop dispatching containers and release resources. Used by tests.
 108 func (disp *dispatcher) Close() {
 109         disp.Start()
 110         select {
 111         case disp.stop <- struct{}{}:
 112         default:
 113         }
 114         <-disp.stopped
 115 }
 116
 117 func (disp *dispatcher) init() {
 118         disp.logger = ctxlog.FromContext(disp.Context)
 119         disp.lsfcli.logger = disp.logger
 120         disp.lsfqueue = lsfqueue{
 121                 logger: disp.logger,
 122                 period: time.Duration(disp.Cluster.Containers.CloudVMs.PollInterval),
 123                 lsfcli: &disp.lsfcli,
 124         }
 125         disp.ArvClient.AuthToken = disp.AuthToken
 126         disp.stop = make(chan struct{}, 1)
 127         disp.stopped = make(chan struct{})
 128
 129         arv, err := arvadosclient.New(disp.ArvClient)
 130         if err != nil {
 131                 disp.logger.Fatalf("Error making Arvados client: %v", err)
 132         }
 133         arv.Retries = 25
 134         arv.ApiToken = disp.AuthToken
 135         disp.arvDispatcher = &dispatch.Dispatcher{
 136                 Arv:            arv,
 137                 Logger:         disp.logger,
 138                 BatchSize:      disp.Cluster.API.MaxItemsPerResponse,
 139                 RunContainer:   disp.runContainer,
 140                 PollPeriod:     time.Duration(disp.Cluster.Containers.CloudVMs.PollInterval),
 141                 MinRetryPeriod: time.Duration(disp.Cluster.Containers.MinRetryPeriod),
 142         }
 143
 144         if disp.Cluster.ManagementToken == "" {
 145                 disp.httpHandler = http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 146                         http.Error(w, "Management API authentication is not configured", http.StatusForbidden)
 147                 })
 148         } else {
 149                 mux := httprouter.New()
 150                 metricsH := promhttp.HandlerFor(disp.Registry, promhttp.HandlerOpts{
 151                         ErrorLog: disp.logger,
 152                 })
 153                 mux.Handler("GET", "/metrics", metricsH)
 154                 mux.Handler("GET", "/metrics.json", metricsH)
 155                 mux.Handler("GET", "/_health/:check", &health.Handler{
 156                         Token:  disp.Cluster.ManagementToken,
 157                         Prefix: "/_health/",
 158                         Routes: health.Routes{"ping": disp.CheckHealth},
 159                 })
 160                 disp.httpHandler = auth.RequireLiteralToken(disp.Cluster.ManagementToken, mux)
 161         }
 162 }
 163
 164 func (disp *dispatcher) runContainer(_ *dispatch.Dispatcher, ctr arvados.Container, status <-chan arvados.Container) error {
 165         ctx, cancel := context.WithCancel(disp.Context)
 166         defer cancel()
 167
 168         if ctr.State != dispatch.Locked {
 169                 // already started by prior invocation
 170         } else if _, ok := disp.lsfqueue.JobID(ctr.UUID); !ok {
 171                 disp.logger.Printf("Submitting container %s to LSF", ctr.UUID)
 172                 cmd := []string{disp.Cluster.Containers.CrunchRunCommand}
 173                 cmd = append(cmd, "--runtime-engine="+disp.Cluster.Containers.RuntimeEngine)
 174                 cmd = append(cmd, disp.Cluster.Containers.CrunchRunArgumentsList...)
 175                 err := disp.submit(ctr, cmd)
 176                 if err != nil {
 177                         return err
 178                 }
 179         }
 180
 181         disp.logger.Printf("Start monitoring container %v in state %q", ctr.UUID, ctr.State)
 182         defer disp.logger.Printf("Done monitoring container %s", ctr.UUID)
 183
 184         // If the container disappears from the lsf queue, there is
 185         // no point in waiting for further dispatch updates: just
 186         // clean up and return.
 187         go func(uuid string) {
 188                 for ctx.Err() == nil {
 189                         if _, ok := disp.lsfqueue.JobID(uuid); !ok {
 190                                 disp.logger.Printf("container %s job disappeared from LSF queue", uuid)
 191                                 cancel()
 192                                 return
 193                         }
 194                 }
 195         }(ctr.UUID)
 196
 197         for done := false; !done; {
 198                 select {
 199                 case <-ctx.Done():
 200                         // Disappeared from lsf queue
 201                         if err := disp.arvDispatcher.Arv.Get("containers", ctr.UUID, nil, &ctr); err != nil {
 202                                 disp.logger.Printf("error getting final container state for %s: %s", ctr.UUID, err)
 203                         }
 204                         switch ctr.State {
 205                         case dispatch.Running:
 206                                 disp.arvDispatcher.UpdateState(ctr.UUID, dispatch.Cancelled)
 207                         case dispatch.Locked:
 208                                 disp.arvDispatcher.Unlock(ctr.UUID)
 209                         }
 210                         return nil
 211                 case updated, ok := <-status:
 212                         if !ok {
 213                                 // status channel is closed, which is
 214                                 // how arvDispatcher tells us to stop
 215                                 // touching the container record, kill
 216                                 // off any remaining LSF processes,
 217                                 // etc.
 218                                 done = true
 219                                 break
 220                         }
 221                         if updated.State != ctr.State {
 222                                 disp.logger.Infof("container %s changed state from %s to %s", ctr.UUID, ctr.State, updated.State)
 223                         }
 224                         ctr = updated
 225                         if ctr.Priority < 1 {
 226                                 disp.logger.Printf("container %s has state %s, priority %d: cancel lsf job", ctr.UUID, ctr.State, ctr.Priority)
 227                                 disp.bkill(ctr)
 228                         } else {
 229                                 disp.lsfqueue.SetPriority(ctr.UUID, int64(ctr.Priority))
 230                         }
 231                 }
 232         }
 233         disp.logger.Printf("container %s is done", ctr.UUID)
 234
 235         // Try "bkill" every few seconds until the LSF job disappears
 236         // from the queue.
 237         ticker := time.NewTicker(5 * time.Second)
 238         defer ticker.Stop()
 239         for jobid, ok := disp.lsfqueue.JobID(ctr.UUID); ok; _, ok = disp.lsfqueue.JobID(ctr.UUID) {
 240                 err := disp.lsfcli.Bkill(jobid)
 241                 if err != nil {
 242                         disp.logger.Warnf("%s: bkill(%d): %s", ctr.UUID, jobid, err)
 243                 }
 244                 <-ticker.C
 245         }
 246         return nil
 247 }
 248
 249 func (disp *dispatcher) submit(container arvados.Container, crunchRunCommand []string) error {
 250         // Start with an empty slice here to ensure append() doesn't
 251         // modify crunchRunCommand's underlying array
 252         var crArgs []string
 253         crArgs = append(crArgs, crunchRunCommand...)
 254         crArgs = append(crArgs, container.UUID)
 255         crScript := execScript(crArgs)
 256
 257         bsubArgs, err := disp.bsubArgs(container)
 258         if err != nil {
 259                 return err
 260         }
 261         return disp.lsfcli.Bsub(crScript, bsubArgs, disp.ArvClient)
 262 }
 263
 264 func (disp *dispatcher) bkill(ctr arvados.Container) {
 265         if jobid, ok := disp.lsfqueue.JobID(ctr.UUID); !ok {
 266                 disp.logger.Debugf("bkill(%s): redundant, job not in queue", ctr.UUID)
 267         } else if err := disp.lsfcli.Bkill(jobid); err != nil {
 268                 disp.logger.Warnf("%s: bkill(%d): %s", ctr.UUID, jobid, err)
 269         }
 270 }
 271
 272 func (disp *dispatcher) bsubArgs(container arvados.Container) ([]string, error) {
 273         tmpArgs := []string{}
 274         args := []string{"bsub"}
 275         tmpArgs = append(tmpArgs, disp.Cluster.Containers.LSF.BsubArgumentsList...)
 276
 277         tmp := int64(math.Ceil(float64(dispatchcloud.EstimateScratchSpace(&container)) / 1048576))
 278         vcpus := container.RuntimeConstraints.VCPUs
 279         mem := int64(math.Ceil(float64(container.RuntimeConstraints.RAM+
 280                 container.RuntimeConstraints.KeepCacheRAM+
 281                 int64(disp.Cluster.Containers.ReserveExtraRAM)) / 1048576))
 282
 283         r := regexp.MustCompile(`([^%]|^)%([^%])`)
 284         undoubleRE := regexp.MustCompile(`%%`)
 285         for _, a := range tmpArgs {
 286                 tmp := r.ReplaceAllStringFunc(a, func(m string) string {
 287                         parts := r.FindStringSubmatch(m)
 288                         return parts[1] + disp.substitute(parts[2], container.UUID, vcpus, mem, tmp)
 289                 })
 290                 // handle escaped literal % symbols
 291                 tmp = undoubleRE.ReplaceAllString(tmp, "%")
 292                 args = append(args, tmp)
 293         }
 294
 295         if u := disp.Cluster.Containers.LSF.BsubSudoUser; u != "" {
 296                 args = append([]string{"sudo", "-E", "-u", u}, args...)
 297         }
 298         return args, nil
 299 }
 300
 301 func (disp *dispatcher) substitute(l string, uuid string, vcpus int, mem, tmp int64) string {
 302         var arg string
 303         switch l {
 304         case "C":
 305                 arg = fmt.Sprintf("%d", vcpus)
 306         case "T":
 307                 arg = fmt.Sprintf("%d", tmp)
 308         case "M":
 309                 arg = fmt.Sprintf("%d", mem)
 310         case "U":
 311                 arg = uuid
 312         default:
 313                 arg = "%" + l
 314         }
 315         return arg
 316 }
 317
 318 // Check the next bjobs report, and invoke TrackContainer for all the
 319 // containers in the report. This gives us a chance to cancel existing
 320 // Arvados LSF jobs (started by a previous dispatch process) that
 321 // never released their LSF job allocations even though their
 322 // container states are Cancelled or Complete. See
 323 // https://dev.arvados.org/issues/10979
 324 func (disp *dispatcher) checkLsfQueueForOrphans() {
 325         containerUuidPattern := regexp.MustCompile(`^[a-z0-9]{5}-dz642-[a-z0-9]{15}$`)
 326         for _, uuid := range disp.lsfqueue.All() {
 327                 if !containerUuidPattern.MatchString(uuid) || !strings.HasPrefix(uuid, disp.Cluster.ClusterID) {
 328                         continue
 329                 }
 330                 err := disp.arvDispatcher.TrackContainer(uuid)
 331                 if err != nil {
 332                         disp.logger.Warnf("checkLsfQueueForOrphans: TrackContainer(%s): %s", uuid, err)
 333                 }
 334         }
 335 }
 336
 337 func execScript(args []string) []byte {
 338         s := "#!/bin/sh\nexec"
 339         for _, w := range args {
 340                 s += ` '`
 341                 s += strings.Replace(w, `'`, `'\''`, -1)
 342                 s += `'`
 343         }
 344         return []byte(s + "\n")
 345 }