arvados.go

   1 // Copyright (C) The Lightning Authors. All rights reserved.
   2 //
   3 // SPDX-License-Identifier: AGPL-3.0
   4
   5 package lightning
   6
   7 import (
   8         "bufio"
   9         "context"
  10         "encoding/json"
  11         "errors"
  12         "fmt"
  13         "io"
  14         "io/ioutil"
  15         "net/url"
  16         "os"
  17         "regexp"
  18         "strconv"
  19         "strings"
  20         "sync"
  21         "time"
  22
  23         "git.arvados.org/arvados.git/lib/cmd"
  24         "git.arvados.org/arvados.git/sdk/go/arvados"
  25         "git.arvados.org/arvados.git/sdk/go/arvadosclient"
  26         "git.arvados.org/arvados.git/sdk/go/keepclient"
  27         "github.com/klauspost/pgzip"
  28         log "github.com/sirupsen/logrus"
  29         "golang.org/x/crypto/blake2b"
  30         "golang.org/x/net/websocket"
  31 )
  32
  33 type eventMessage struct {
  34         Status     int
  35         ObjectUUID string `json:"object_uuid"`
  36         EventType  string `json:"event_type"`
  37         Properties struct {
  38                 Text string
  39         }
  40 }
  41
  42 type arvadosClient struct {
  43         *arvados.Client
  44         notifying map[string]map[chan<- eventMessage]int
  45         wantClose chan struct{}
  46         wsconn    *websocket.Conn
  47         mtx       sync.Mutex
  48 }
  49
  50 // Listen for events concerning the given uuids. When an event occurs
  51 // (and after connecting/reconnecting to the event stream), send each
  52 // uuid to ch. If a {ch, uuid} pair is subscribed twice, the uuid will
  53 // be sent only once for each update, but two Unsubscribe calls will
  54 // be needed to stop sending them.
  55 func (client *arvadosClient) Subscribe(ch chan<- eventMessage, uuid string) {
  56         client.mtx.Lock()
  57         defer client.mtx.Unlock()
  58         if client.notifying == nil {
  59                 client.notifying = map[string]map[chan<- eventMessage]int{}
  60                 client.wantClose = make(chan struct{})
  61                 go client.runNotifier()
  62         }
  63         chmap := client.notifying[uuid]
  64         if chmap == nil {
  65                 chmap = map[chan<- eventMessage]int{}
  66                 client.notifying[uuid] = chmap
  67         }
  68         needSub := true
  69         for _, nch := range chmap {
  70                 if nch > 0 {
  71                         needSub = false
  72                         break
  73                 }
  74         }
  75         chmap[ch]++
  76         if needSub && client.wsconn != nil {
  77                 go json.NewEncoder(client.wsconn).Encode(map[string]interface{}{
  78                         "method": "subscribe",
  79                         "filters": [][]interface{}{
  80                                 {"object_uuid", "=", uuid},
  81                                 {"event_type", "in", []string{"stderr", "crunch-run", "update"}},
  82                         },
  83                 })
  84         }
  85 }
  86
  87 func (client *arvadosClient) Unsubscribe(ch chan<- eventMessage, uuid string) {
  88         client.mtx.Lock()
  89         defer client.mtx.Unlock()
  90         chmap := client.notifying[uuid]
  91         if n := chmap[ch] - 1; n == 0 {
  92                 delete(chmap, ch)
  93                 if len(chmap) == 0 {
  94                         delete(client.notifying, uuid)
  95                 }
  96                 if client.wsconn != nil {
  97                         go json.NewEncoder(client.wsconn).Encode(map[string]interface{}{
  98                                 "method": "unsubscribe",
  99                                 "filters": [][]interface{}{
 100                                         {"object_uuid", "=", uuid},
 101                                         {"event_type", "in", []string{"stderr", "crunch-run", "update"}},
 102                                 },
 103                         })
 104                 }
 105         } else if n > 0 {
 106                 chmap[ch] = n
 107         }
 108 }
 109
 110 func (client *arvadosClient) Close() {
 111         client.mtx.Lock()
 112         defer client.mtx.Unlock()
 113         if client.notifying != nil {
 114                 client.notifying = nil
 115                 close(client.wantClose)
 116         }
 117 }
 118
 119 func (client *arvadosClient) runNotifier() {
 120 reconnect:
 121         for {
 122                 var cluster arvados.Cluster
 123                 err := client.RequestAndDecode(&cluster, "GET", arvados.EndpointConfigGet.Path, nil, nil)
 124                 if err != nil {
 125                         log.Warnf("error getting cluster config: %s", err)
 126                         time.Sleep(5 * time.Second)
 127                         continue reconnect
 128                 }
 129                 wsURL := cluster.Services.Websocket.ExternalURL
 130                 wsURL.Scheme = strings.Replace(wsURL.Scheme, "http", "ws", 1)
 131                 wsURL.Path = "/websocket"
 132                 wsURLNoToken := wsURL.String()
 133                 wsURL.RawQuery = url.Values{"api_token": []string{client.AuthToken}}.Encode()
 134                 conn, err := websocket.Dial(wsURL.String(), "", cluster.Services.Controller.ExternalURL.String())
 135                 if err != nil {
 136                         log.Warnf("websocket connection error: %s", err)
 137                         time.Sleep(5 * time.Second)
 138                         continue reconnect
 139                 }
 140                 log.Printf("connected to websocket at %s", wsURLNoToken)
 141
 142                 client.mtx.Lock()
 143                 client.wsconn = conn
 144                 resubscribe := make([]string, 0, len(client.notifying))
 145                 for uuid := range client.notifying {
 146                         resubscribe = append(resubscribe, uuid)
 147                 }
 148                 client.mtx.Unlock()
 149
 150                 go func() {
 151                         w := json.NewEncoder(conn)
 152                         for _, uuid := range resubscribe {
 153                                 w.Encode(map[string]interface{}{
 154                                         "method": "subscribe",
 155                                         "filters": [][]interface{}{
 156                                                 {"object_uuid", "=", uuid},
 157                                                 {"event_type", "in", []string{"stderr", "crunch-run", "crunchstat", "update"}},
 158                                         },
 159                                 })
 160                         }
 161                 }()
 162
 163                 r := json.NewDecoder(conn)
 164                 for {
 165                         var msg eventMessage
 166                         err := r.Decode(&msg)
 167                         select {
 168                         case <-client.wantClose:
 169                                 return
 170                         default:
 171                                 if err != nil {
 172                                         log.Printf("error decoding websocket message: %s", err)
 173                                         client.mtx.Lock()
 174                                         client.wsconn = nil
 175                                         client.mtx.Unlock()
 176                                         go conn.Close()
 177                                         continue reconnect
 178                                 }
 179                                 client.mtx.Lock()
 180                                 for ch := range client.notifying[msg.ObjectUUID] {
 181                                         ch <- msg
 182                                 }
 183                                 client.mtx.Unlock()
 184                         }
 185                 }
 186         }
 187 }
 188
 189 var refreshTicker = time.NewTicker(5 * time.Second)
 190
 191 type arvadosContainerRunner struct {
 192         Client      *arvados.Client
 193         Name        string
 194         OutputName  string
 195         ProjectUUID string
 196         APIAccess   bool
 197         VCPUs       int
 198         RAM         int64
 199         Prog        string // if empty, run /proc/self/exe
 200         Args        []string
 201         Mounts      map[string]map[string]interface{}
 202         Priority    int
 203         KeepCache   int // cache buffers per VCPU (0 for default)
 204 }
 205
 206 func (runner *arvadosContainerRunner) Run() (string, error) {
 207         return runner.RunContext(context.Background())
 208 }
 209
 210 func (runner *arvadosContainerRunner) RunContext(ctx context.Context) (string, error) {
 211         if runner.ProjectUUID == "" {
 212                 return "", errors.New("cannot run arvados container: ProjectUUID not provided")
 213         }
 214
 215         mounts := map[string]map[string]interface{}{
 216                 "/mnt/output": {
 217                         "kind":     "collection",
 218                         "writable": true,
 219                 },
 220         }
 221         for path, mnt := range runner.Mounts {
 222                 mounts[path] = mnt
 223         }
 224
 225         prog := runner.Prog
 226         if prog == "" {
 227                 prog = "/mnt/cmd/lightning"
 228                 cmdUUID, err := runner.makeCommandCollection()
 229                 if err != nil {
 230                         return "", err
 231                 }
 232                 mounts["/mnt/cmd"] = map[string]interface{}{
 233                         "kind": "collection",
 234                         "uuid": cmdUUID,
 235                 }
 236         }
 237         command := append([]string{prog}, runner.Args...)
 238
 239         priority := runner.Priority
 240         if priority < 1 {
 241                 priority = 500
 242         }
 243         keepCache := runner.KeepCache
 244         if keepCache < 1 {
 245                 keepCache = 2
 246         }
 247         rc := arvados.RuntimeConstraints{
 248                 API:          &runner.APIAccess,
 249                 VCPUs:        runner.VCPUs,
 250                 RAM:          runner.RAM,
 251                 KeepCacheRAM: (1 << 26) * int64(keepCache) * int64(runner.VCPUs),
 252         }
 253         outname := &runner.OutputName
 254         if *outname == "" {
 255                 outname = nil
 256         }
 257         var cr arvados.ContainerRequest
 258         err := runner.Client.RequestAndDecode(&cr, "POST", "arvados/v1/container_requests", nil, map[string]interface{}{
 259                 "container_request": map[string]interface{}{
 260                         "owner_uuid":          runner.ProjectUUID,
 261                         "name":                runner.Name,
 262                         "container_image":     "lightning-runtime",
 263                         "command":             command,
 264                         "mounts":              mounts,
 265                         "use_existing":        true,
 266                         "output_path":         "/mnt/output",
 267                         "output_name":         outname,
 268                         "runtime_constraints": rc,
 269                         "priority":            runner.Priority,
 270                         "state":               arvados.ContainerRequestStateCommitted,
 271                         "scheduling_parameters": arvados.SchedulingParameters{
 272                                 Preemptible: false,
 273                                 Partitions:  []string{},
 274                         },
 275                         "environment": map[string]string{
 276                                 "GOMAXPROCS": fmt.Sprintf("%d", rc.VCPUs),
 277                         },
 278                 },
 279         })
 280         if err != nil {
 281                 return "", err
 282         }
 283         log.Printf("container request UUID: %s", cr.UUID)
 284         log.Printf("container UUID: %s", cr.ContainerUUID)
 285
 286         logch := make(chan eventMessage)
 287         client := arvadosClient{Client: runner.Client}
 288         defer client.Close()
 289         subscribedUUID := ""
 290         defer func() {
 291                 if subscribedUUID != "" {
 292                         log.Printf("unsubscribe container UUID: %s", subscribedUUID)
 293                         client.Unsubscribe(logch, subscribedUUID)
 294                 }
 295         }()
 296
 297         neednewline := ""
 298
 299         lastState := cr.State
 300         refreshCR := func() {
 301                 ctx, cancel := context.WithDeadline(ctx, time.Now().Add(time.Minute))
 302                 defer cancel()
 303                 err = runner.Client.RequestAndDecodeContext(ctx, &cr, "GET", "arvados/v1/container_requests/"+cr.UUID, nil, nil)
 304                 if err != nil {
 305                         fmt.Fprint(os.Stderr, neednewline)
 306                         neednewline = ""
 307                         log.Printf("error getting container request: %s", err)
 308                         return
 309                 }
 310                 if lastState != cr.State {
 311                         fmt.Fprint(os.Stderr, neednewline)
 312                         neednewline = ""
 313                         log.Printf("container request state: %s", cr.State)
 314                         lastState = cr.State
 315                 }
 316                 if subscribedUUID != cr.ContainerUUID {
 317                         fmt.Fprint(os.Stderr, neednewline)
 318                         neednewline = ""
 319                         if subscribedUUID != "" {
 320                                 log.Printf("unsubscribe container UUID: %s", subscribedUUID)
 321                                 client.Unsubscribe(logch, subscribedUUID)
 322                         }
 323                         log.Printf("subscribe container UUID: %s", cr.ContainerUUID)
 324                         client.Subscribe(logch, cr.ContainerUUID)
 325                         subscribedUUID = cr.ContainerUUID
 326                 }
 327         }
 328
 329         var reCrunchstat = regexp.MustCompile(`mem .* (\d+) rss`)
 330 waitctr:
 331         for cr.State != arvados.ContainerRequestStateFinal {
 332                 select {
 333                 case <-ctx.Done():
 334                         err := runner.Client.RequestAndDecode(&cr, "PATCH", "arvados/v1/container_requests/"+cr.UUID, nil, map[string]interface{}{
 335                                 "container_request": map[string]interface{}{
 336                                         "priority": 0,
 337                                 },
 338                         })
 339                         if err != nil {
 340                                 log.Errorf("error while trying to cancel container request %s: %s", cr.UUID, err)
 341                         }
 342                         break waitctr
 343                 case <-refreshTicker.C:
 344                         refreshCR()
 345                 case msg := <-logch:
 346                         switch msg.EventType {
 347                         case "update":
 348                                 refreshCR()
 349                         case "stderr":
 350                                 for _, line := range strings.Split(msg.Properties.Text, "\n") {
 351                                         if line != "" {
 352                                                 fmt.Fprint(os.Stderr, neednewline)
 353                                                 neednewline = ""
 354                                                 log.Print(line)
 355                                         }
 356                                 }
 357                         case "crunchstat":
 358                                 for _, line := range strings.Split(msg.Properties.Text, "\n") {
 359                                         m := reCrunchstat.FindStringSubmatch(line)
 360                                         if m != nil {
 361                                                 rss, _ := strconv.ParseInt(m[1], 10, 64)
 362                                                 fmt.Fprintf(os.Stderr, "%s rss %.3f GB           \r", cr.UUID, float64(rss)/1e9)
 363                                                 neednewline = "\n"
 364                                         }
 365                                 }
 366                         }
 367                 }
 368         }
 369         fmt.Fprint(os.Stderr, neednewline)
 370
 371         if err := ctx.Err(); err != nil {
 372                 return "", err
 373         }
 374
 375         var c arvados.Container
 376         err = runner.Client.RequestAndDecode(&c, "GET", "arvados/v1/containers/"+cr.ContainerUUID, nil, nil)
 377         if err != nil {
 378                 return "", err
 379         } else if c.State != arvados.ContainerStateComplete {
 380                 return "", fmt.Errorf("container did not complete: %s", c.State)
 381         } else if c.ExitCode != 0 {
 382                 return "", fmt.Errorf("container exited %d", c.ExitCode)
 383         }
 384         return cr.OutputUUID, err
 385 }
 386
 387 var collectionInPathRe = regexp.MustCompile(`^(.*/)?([0-9a-f]{32}\+[0-9]+|[0-9a-z]{5}-[0-9a-z]{5}-[0-9a-z]{15})(/.*)?$`)
 388
 389 func (runner *arvadosContainerRunner) TranslatePaths(paths ...*string) error {
 390         if runner.Mounts == nil {
 391                 runner.Mounts = make(map[string]map[string]interface{})
 392         }
 393         for _, path := range paths {
 394                 if *path == "" || *path == "-" {
 395                         continue
 396                 }
 397                 m := collectionInPathRe.FindStringSubmatch(*path)
 398                 if m == nil {
 399                         return fmt.Errorf("cannot find uuid in path: %q", *path)
 400                 }
 401                 collID := m[2]
 402                 mnt, ok := runner.Mounts["/mnt/"+collID]
 403                 if !ok {
 404                         mnt = map[string]interface{}{
 405                                 "kind": "collection",
 406                         }
 407                         if len(collID) == 27 {
 408                                 mnt["uuid"] = collID
 409                         } else {
 410                                 mnt["portable_data_hash"] = collID
 411                         }
 412                         runner.Mounts["/mnt/"+collID] = mnt
 413                 }
 414                 *path = "/mnt/" + collID + m[3]
 415         }
 416         return nil
 417 }
 418
 419 var mtxMakeCommandCollection sync.Mutex
 420
 421 func (runner *arvadosContainerRunner) makeCommandCollection() (string, error) {
 422         mtxMakeCommandCollection.Lock()
 423         defer mtxMakeCommandCollection.Unlock()
 424         exe, err := ioutil.ReadFile("/proc/self/exe")
 425         if err != nil {
 426                 return "", err
 427         }
 428         b2 := blake2b.Sum256(exe)
 429         cname := "lightning " + cmd.Version.String() // must build with "make", not just "go install"
 430         var existing arvados.CollectionList
 431         err = runner.Client.RequestAndDecode(&existing, "GET", "arvados/v1/collections", nil, arvados.ListOptions{
 432                 Limit: 1,
 433                 Count: "none",
 434                 Filters: []arvados.Filter{
 435                         {Attr: "name", Operator: "=", Operand: cname},
 436                         {Attr: "owner_uuid", Operator: "=", Operand: runner.ProjectUUID},
 437                         {Attr: "properties.blake2b", Operator: "=", Operand: fmt.Sprintf("%x", b2)},
 438                 },
 439         })
 440         if err != nil {
 441                 return "", err
 442         }
 443         if len(existing.Items) > 0 {
 444                 coll := existing.Items[0]
 445                 log.Printf("using lightning binary in existing collection %s (name is %q, hash is %q; did not verify whether content matches)", coll.UUID, cname, coll.Properties["blake2b"])
 446                 return coll.UUID, nil
 447         }
 448         log.Printf("writing lightning binary to new collection %q", cname)
 449         ac, err := arvadosclient.New(runner.Client)
 450         if err != nil {
 451                 return "", err
 452         }
 453         kc := keepclient.New(ac)
 454         var coll arvados.Collection
 455         fs, err := coll.FileSystem(runner.Client, kc)
 456         if err != nil {
 457                 return "", err
 458         }
 459         f, err := fs.OpenFile("lightning", os.O_CREATE|os.O_WRONLY, 0777)
 460         if err != nil {
 461                 return "", err
 462         }
 463         _, err = f.Write(exe)
 464         if err != nil {
 465                 return "", err
 466         }
 467         err = f.Close()
 468         if err != nil {
 469                 return "", err
 470         }
 471         mtxt, err := fs.MarshalManifest(".")
 472         if err != nil {
 473                 return "", err
 474         }
 475         err = runner.Client.RequestAndDecode(&coll, "POST", "arvados/v1/collections", nil, map[string]interface{}{
 476                 "collection": map[string]interface{}{
 477                         "owner_uuid":    runner.ProjectUUID,
 478                         "manifest_text": mtxt,
 479                         "name":          cname,
 480                         "properties": map[string]interface{}{
 481                                 "blake2b": fmt.Sprintf("%x", b2),
 482                         },
 483                 },
 484         })
 485         if err != nil {
 486                 return "", err
 487         }
 488         log.Printf("stored lightning binary in new collection %s", coll.UUID)
 489         return coll.UUID, nil
 490 }
 491
 492 // zopen returns a reader for the given file, using the arvados API
 493 // instead of arv-mount/fuse where applicable, and transparently
 494 // decompressing the input if fnm ends with ".gz".
 495 func zopen(fnm string) (io.ReadCloser, error) {
 496         f, err := open(fnm)
 497         if err != nil || !strings.HasSuffix(fnm, ".gz") {
 498                 return f, err
 499         }
 500         rdr, err := pgzip.NewReader(bufio.NewReaderSize(f, 4*1024*1024))
 501         if err != nil {
 502                 f.Close()
 503                 return nil, err
 504         }
 505         return gzipr{rdr, f}, nil
 506 }
 507
 508 // gzipr wraps a ReadCloser and a Closer, presenting a single Close()
 509 // method that closes both wrapped objects.
 510 type gzipr struct {
 511         io.ReadCloser
 512         io.Closer
 513 }
 514
 515 func (gr gzipr) Close() error {
 516         e1 := gr.ReadCloser.Close()
 517         e2 := gr.Closer.Close()
 518         if e1 != nil {
 519                 return e1
 520         }
 521         return e2
 522 }
 523
 524 var (
 525         arvadosClientFromEnv = arvados.NewClientFromEnv()
 526         keepClient           *keepclient.KeepClient
 527         siteFS               arvados.CustomFileSystem
 528         siteFSMtx            sync.Mutex
 529 )
 530
 531 type file interface {
 532         io.ReadCloser
 533         io.Seeker
 534         Readdir(n int) ([]os.FileInfo, error)
 535 }
 536
 537 func open(fnm string) (file, error) {
 538         if os.Getenv("ARVADOS_API_HOST") == "" {
 539                 return os.Open(fnm)
 540         }
 541         m := collectionInPathRe.FindStringSubmatch(fnm)
 542         if m == nil {
 543                 return os.Open(fnm)
 544         }
 545         collectionUUID := m[2]
 546         collectionPath := m[3]
 547
 548         siteFSMtx.Lock()
 549         defer siteFSMtx.Unlock()
 550         if siteFS == nil {
 551                 log.Info("setting up Arvados client")
 552                 ac, err := arvadosclient.New(arvadosClientFromEnv)
 553                 if err != nil {
 554                         return nil, err
 555                 }
 556                 ac.Client = arvados.DefaultSecureClient
 557                 keepClient = keepclient.New(ac)
 558                 // Don't use keepclient's default short timeouts.
 559                 keepClient.HTTPClient = arvados.DefaultSecureClient
 560                 keepClient.BlockCache = &keepclient.BlockCache{MaxBlocks: 4}
 561                 siteFS = arvadosClientFromEnv.SiteFileSystem(keepClient)
 562         } else {
 563                 keepClient.BlockCache.MaxBlocks += 2
 564         }
 565
 566         log.Infof("reading %q from %s using Arvados client", collectionPath, collectionUUID)
 567         f, err := siteFS.Open("by_id/" + collectionUUID + collectionPath)
 568         if err != nil {
 569                 return nil, err
 570         }
 571         return &reduceCacheOnClose{file: f}, nil
 572 }
 573
 574 type reduceCacheOnClose struct {
 575         file
 576         once sync.Once
 577 }
 578
 579 func (rc *reduceCacheOnClose) Close() error {
 580         rc.once.Do(func() { keepClient.BlockCache.MaxBlocks -= 2 })
 581         return rc.file.Close()
 582 }