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                                         go func() { 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         Preemptible bool
 205 }
 206
 207 func (runner *arvadosContainerRunner) Run() (string, error) {
 208         return runner.RunContext(context.Background())
 209 }
 210
 211 func (runner *arvadosContainerRunner) RunContext(ctx context.Context) (string, error) {
 212         if runner.ProjectUUID == "" {
 213                 return "", errors.New("cannot run arvados container: ProjectUUID not provided")
 214         }
 215
 216         mounts := map[string]map[string]interface{}{
 217                 "/mnt/output": {
 218                         "kind":     "collection",
 219                         "writable": true,
 220                 },
 221         }
 222         for path, mnt := range runner.Mounts {
 223                 mounts[path] = mnt
 224         }
 225
 226         prog := runner.Prog
 227         if prog == "" {
 228                 prog = "/mnt/cmd/lightning"
 229                 cmdUUID, err := runner.makeCommandCollection()
 230                 if err != nil {
 231                         return "", err
 232                 }
 233                 mounts["/mnt/cmd"] = map[string]interface{}{
 234                         "kind": "collection",
 235                         "uuid": cmdUUID,
 236                 }
 237         }
 238         command := append([]string{prog}, runner.Args...)
 239
 240         priority := runner.Priority
 241         if priority < 1 {
 242                 priority = 500
 243         }
 244         keepCache := runner.KeepCache
 245         if keepCache < 1 {
 246                 keepCache = 2
 247         }
 248         rc := arvados.RuntimeConstraints{
 249                 API:          runner.APIAccess,
 250                 VCPUs:        runner.VCPUs,
 251                 RAM:          runner.RAM,
 252                 KeepCacheRAM: (1 << 26) * int64(keepCache) * int64(runner.VCPUs),
 253         }
 254         outname := &runner.OutputName
 255         if *outname == "" {
 256                 outname = nil
 257         }
 258         var cr arvados.ContainerRequest
 259         err := runner.Client.RequestAndDecode(&cr, "POST", "arvados/v1/container_requests", nil, map[string]interface{}{
 260                 "container_request": map[string]interface{}{
 261                         "owner_uuid":          runner.ProjectUUID,
 262                         "name":                runner.Name,
 263                         "container_image":     "lightning-runtime",
 264                         "command":             command,
 265                         "mounts":              mounts,
 266                         "use_existing":        true,
 267                         "output_path":         "/mnt/output",
 268                         "output_name":         outname,
 269                         "runtime_constraints": rc,
 270                         "priority":            runner.Priority,
 271                         "state":               arvados.ContainerRequestStateCommitted,
 272                         "scheduling_parameters": arvados.SchedulingParameters{
 273                                 Preemptible: runner.Preemptible,
 274                                 Partitions:  []string{},
 275                         },
 276                         "environment": map[string]string{
 277                                 "GOMAXPROCS": fmt.Sprintf("%d", rc.VCPUs),
 278                         },
 279                 },
 280         })
 281         if err != nil {
 282                 return "", err
 283         }
 284         log.Printf("container request UUID: %s", cr.UUID)
 285         log.Printf("container UUID: %s", cr.ContainerUUID)
 286
 287         logch := make(chan eventMessage)
 288         client := arvadosClient{Client: runner.Client}
 289         defer client.Close()
 290         subscribedUUID := ""
 291         defer func() {
 292                 if subscribedUUID != "" {
 293                         log.Printf("unsubscribe container UUID: %s", subscribedUUID)
 294                         client.Unsubscribe(logch, subscribedUUID)
 295                 }
 296         }()
 297
 298         neednewline := ""
 299
 300         lastState := cr.State
 301         refreshCR := func() {
 302                 ctx, cancel := context.WithDeadline(ctx, time.Now().Add(time.Minute))
 303                 defer cancel()
 304                 err = runner.Client.RequestAndDecodeContext(ctx, &cr, "GET", "arvados/v1/container_requests/"+cr.UUID, nil, nil)
 305                 if err != nil {
 306                         fmt.Fprint(os.Stderr, neednewline)
 307                         neednewline = ""
 308                         log.Printf("error getting container request: %s", err)
 309                         return
 310                 }
 311                 if lastState != cr.State {
 312                         fmt.Fprint(os.Stderr, neednewline)
 313                         neednewline = ""
 314                         log.Printf("container request state: %s", cr.State)
 315                         lastState = cr.State
 316                 }
 317                 if subscribedUUID != cr.ContainerUUID {
 318                         fmt.Fprint(os.Stderr, neednewline)
 319                         neednewline = ""
 320                         if subscribedUUID != "" {
 321                                 log.Printf("unsubscribe container UUID: %s", subscribedUUID)
 322                                 client.Unsubscribe(logch, subscribedUUID)
 323                         }
 324                         log.Printf("subscribe container UUID: %s", cr.ContainerUUID)
 325                         client.Subscribe(logch, cr.ContainerUUID)
 326                         subscribedUUID = cr.ContainerUUID
 327                 }
 328         }
 329
 330         var reCrunchstat = regexp.MustCompile(`mem .* (\d+) rss`)
 331 waitctr:
 332         for cr.State != arvados.ContainerRequestStateFinal {
 333                 select {
 334                 case <-ctx.Done():
 335                         err := runner.Client.RequestAndDecode(&cr, "PATCH", "arvados/v1/container_requests/"+cr.UUID, nil, map[string]interface{}{
 336                                 "container_request": map[string]interface{}{
 337                                         "priority": 0,
 338                                 },
 339                         })
 340                         if err != nil {
 341                                 log.Errorf("error while trying to cancel container request %s: %s", cr.UUID, err)
 342                         }
 343                         break waitctr
 344                 case <-refreshTicker.C:
 345                         refreshCR()
 346                 case msg := <-logch:
 347                         switch msg.EventType {
 348                         case "update":
 349                                 refreshCR()
 350                         case "stderr":
 351                                 for _, line := range strings.Split(msg.Properties.Text, "\n") {
 352                                         if line != "" {
 353                                                 fmt.Fprint(os.Stderr, neednewline)
 354                                                 neednewline = ""
 355                                                 log.Print(line)
 356                                         }
 357                                 }
 358                         case "crunchstat":
 359                                 for _, line := range strings.Split(msg.Properties.Text, "\n") {
 360                                         m := reCrunchstat.FindStringSubmatch(line)
 361                                         if m != nil {
 362                                                 rss, _ := strconv.ParseInt(m[1], 10, 64)
 363                                                 fmt.Fprintf(os.Stderr, "%s rss %.3f GB           \r", cr.UUID, float64(rss)/1e9)
 364                                                 neednewline = "\n"
 365                                         }
 366                                 }
 367                         }
 368                 }
 369         }
 370         fmt.Fprint(os.Stderr, neednewline)
 371
 372         if err := ctx.Err(); err != nil {
 373                 return "", err
 374         }
 375
 376         var c arvados.Container
 377         err = runner.Client.RequestAndDecode(&c, "GET", "arvados/v1/containers/"+cr.ContainerUUID, nil, nil)
 378         if err != nil {
 379                 return "", err
 380         } else if c.State != arvados.ContainerStateComplete {
 381                 return "", fmt.Errorf("container did not complete: %s", c.State)
 382         } else if c.ExitCode != 0 {
 383                 return "", fmt.Errorf("container exited %d", c.ExitCode)
 384         }
 385         return cr.OutputUUID, err
 386 }
 387
 388 var collectionInPathRe = regexp.MustCompile(`^(.*/)?([0-9a-f]{32}\+[0-9]+|[0-9a-z]{5}-[0-9a-z]{5}-[0-9a-z]{15})(/.*)?$`)
 389
 390 func (runner *arvadosContainerRunner) TranslatePaths(paths ...*string) error {
 391         if runner.Mounts == nil {
 392                 runner.Mounts = make(map[string]map[string]interface{})
 393         }
 394         for _, path := range paths {
 395                 if *path == "" || *path == "-" {
 396                         continue
 397                 }
 398                 m := collectionInPathRe.FindStringSubmatch(*path)
 399                 if m == nil {
 400                         return fmt.Errorf("cannot find uuid in path: %q", *path)
 401                 }
 402                 collID := m[2]
 403                 mnt, ok := runner.Mounts["/mnt/"+collID]
 404                 if !ok {
 405                         mnt = map[string]interface{}{
 406                                 "kind": "collection",
 407                         }
 408                         if len(collID) == 27 {
 409                                 mnt["uuid"] = collID
 410                         } else {
 411                                 mnt["portable_data_hash"] = collID
 412                         }
 413                         runner.Mounts["/mnt/"+collID] = mnt
 414                 }
 415                 *path = "/mnt/" + collID + m[3]
 416         }
 417         return nil
 418 }
 419
 420 var mtxMakeCommandCollection sync.Mutex
 421
 422 func (runner *arvadosContainerRunner) makeCommandCollection() (string, error) {
 423         mtxMakeCommandCollection.Lock()
 424         defer mtxMakeCommandCollection.Unlock()
 425         exe, err := ioutil.ReadFile("/proc/self/exe")
 426         if err != nil {
 427                 return "", err
 428         }
 429         b2 := blake2b.Sum256(exe)
 430         cname := "lightning " + cmd.Version.String() // must build with "make", not just "go install"
 431         var existing arvados.CollectionList
 432         err = runner.Client.RequestAndDecode(&existing, "GET", "arvados/v1/collections", nil, arvados.ListOptions{
 433                 Limit: 1,
 434                 Count: "none",
 435                 Filters: []arvados.Filter{
 436                         {Attr: "name", Operator: "=", Operand: cname},
 437                         {Attr: "owner_uuid", Operator: "=", Operand: runner.ProjectUUID},
 438                         {Attr: "properties.blake2b", Operator: "=", Operand: fmt.Sprintf("%x", b2)},
 439                 },
 440         })
 441         if err != nil {
 442                 return "", err
 443         }
 444         if len(existing.Items) > 0 {
 445                 coll := existing.Items[0]
 446                 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"])
 447                 return coll.UUID, nil
 448         }
 449         log.Printf("writing lightning binary to new collection %q", cname)
 450         ac, err := arvadosclient.New(runner.Client)
 451         if err != nil {
 452                 return "", err
 453         }
 454         kc := keepclient.New(ac)
 455         var coll arvados.Collection
 456         fs, err := coll.FileSystem(runner.Client, kc)
 457         if err != nil {
 458                 return "", err
 459         }
 460         f, err := fs.OpenFile("lightning", os.O_CREATE|os.O_WRONLY, 0777)
 461         if err != nil {
 462                 return "", err
 463         }
 464         _, err = f.Write(exe)
 465         if err != nil {
 466                 return "", err
 467         }
 468         err = f.Close()
 469         if err != nil {
 470                 return "", err
 471         }
 472         mtxt, err := fs.MarshalManifest(".")
 473         if err != nil {
 474                 return "", err
 475         }
 476         err = runner.Client.RequestAndDecode(&coll, "POST", "arvados/v1/collections", nil, map[string]interface{}{
 477                 "collection": map[string]interface{}{
 478                         "owner_uuid":    runner.ProjectUUID,
 479                         "manifest_text": mtxt,
 480                         "name":          cname,
 481                         "properties": map[string]interface{}{
 482                                 "blake2b": fmt.Sprintf("%x", b2),
 483                         },
 484                 },
 485         })
 486         if err != nil {
 487                 return "", err
 488         }
 489         log.Printf("stored lightning binary in new collection %s", coll.UUID)
 490         return coll.UUID, nil
 491 }
 492
 493 // zopen returns a reader for the given file, using the arvados API
 494 // instead of arv-mount/fuse where applicable, and transparently
 495 // decompressing the input if fnm ends with ".gz".
 496 func zopen(fnm string) (io.ReadCloser, error) {
 497         f, err := open(fnm)
 498         if err != nil || !strings.HasSuffix(fnm, ".gz") {
 499                 return f, err
 500         }
 501         rdr, err := pgzip.NewReader(bufio.NewReaderSize(f, 4*1024*1024))
 502         if err != nil {
 503                 f.Close()
 504                 return nil, err
 505         }
 506         return gzipr{rdr, f}, nil
 507 }
 508
 509 // gzipr wraps a ReadCloser and a Closer, presenting a single Close()
 510 // method that closes both wrapped objects.
 511 type gzipr struct {
 512         io.ReadCloser
 513         io.Closer
 514 }
 515
 516 func (gr gzipr) Close() error {
 517         e1 := gr.ReadCloser.Close()
 518         e2 := gr.Closer.Close()
 519         if e1 != nil {
 520                 return e1
 521         }
 522         return e2
 523 }
 524
 525 var (
 526         arvadosClientFromEnv = arvados.NewClientFromEnv()
 527         keepClient           *keepclient.KeepClient
 528         siteFS               arvados.CustomFileSystem
 529         siteFSMtx            sync.Mutex
 530 )
 531
 532 type file interface {
 533         io.ReadCloser
 534         io.Seeker
 535         Readdir(n int) ([]os.FileInfo, error)
 536 }
 537
 538 func open(fnm string) (file, error) {
 539         if os.Getenv("ARVADOS_API_HOST") == "" {
 540                 return os.Open(fnm)
 541         }
 542         m := collectionInPathRe.FindStringSubmatch(fnm)
 543         if m == nil {
 544                 return os.Open(fnm)
 545         }
 546         collectionUUID := m[2]
 547         collectionPath := m[3]
 548
 549         siteFSMtx.Lock()
 550         defer siteFSMtx.Unlock()
 551         if siteFS == nil {
 552                 log.Info("setting up Arvados client")
 553                 ac, err := arvadosclient.New(arvadosClientFromEnv)
 554                 if err != nil {
 555                         return nil, err
 556                 }
 557                 ac.Client = arvados.DefaultSecureClient
 558                 keepClient = keepclient.New(ac)
 559                 // Don't use keepclient's default short timeouts.
 560                 keepClient.HTTPClient = arvados.DefaultSecureClient
 561                 keepClient.BlockCache = &keepclient.BlockCache{MaxBlocks: 4}
 562                 siteFS = arvadosClientFromEnv.SiteFileSystem(keepClient)
 563         } else {
 564                 keepClient.BlockCache.MaxBlocks += 2
 565         }
 566
 567         log.Infof("reading %q from %s using Arvados client", collectionPath, collectionUUID)
 568         f, err := siteFS.Open("by_id/" + collectionUUID + collectionPath)
 569         if err != nil {
 570                 return nil, err
 571         }
 572         return &reduceCacheOnClose{file: f}, nil
 573 }
 574
 575 type reduceCacheOnClose struct {
 576         file
 577         once sync.Once
 578 }
 579
 580 func (rc *reduceCacheOnClose) Close() error {
 581         rc.once.Do(func() { keepClient.BlockCache.MaxBlocks -= 2 })
 582         return rc.file.Close()
 583 }