arvados.go

   1 package main
   2
   3 import (
   4         "context"
   5         "encoding/json"
   6         "errors"
   7         "fmt"
   8         "io/ioutil"
   9         "net/url"
  10         "os"
  11         "regexp"
  12         "strings"
  13         "sync"
  14         "time"
  15
  16         "git.arvados.org/arvados.git/sdk/go/arvados"
  17         "git.arvados.org/arvados.git/sdk/go/arvadosclient"
  18         "git.arvados.org/arvados.git/sdk/go/keepclient"
  19         log "github.com/sirupsen/logrus"
  20         "golang.org/x/crypto/blake2b"
  21         "golang.org/x/net/websocket"
  22 )
  23
  24 type eventMessage struct {
  25         Status     int
  26         ObjectUUID string `json:"object_uuid"`
  27         EventType  string `json:"event_type"`
  28         Properties struct {
  29                 Text string
  30         }
  31 }
  32
  33 type arvadosClient struct {
  34         *arvados.Client
  35         notifying map[string]map[chan<- eventMessage]int
  36         wantClose chan struct{}
  37         wsconn    *websocket.Conn
  38         mtx       sync.Mutex
  39 }
  40
  41 // Listen for events concerning the given uuids. When an event occurs
  42 // (and after connecting/reconnecting to the event stream), send each
  43 // uuid to ch. If a {ch, uuid} pair is subscribed twice, the uuid will
  44 // be sent only once for each update, but two Unsubscribe calls will
  45 // be needed to stop sending them.
  46 func (client *arvadosClient) Subscribe(ch chan<- eventMessage, uuid string) {
  47         client.mtx.Lock()
  48         defer client.mtx.Unlock()
  49         if client.notifying == nil {
  50                 client.notifying = map[string]map[chan<- eventMessage]int{}
  51                 client.wantClose = make(chan struct{})
  52                 go client.runNotifier()
  53         }
  54         chmap := client.notifying[uuid]
  55         if chmap == nil {
  56                 chmap = map[chan<- eventMessage]int{}
  57                 client.notifying[uuid] = chmap
  58         }
  59         needSub := true
  60         for _, nch := range chmap {
  61                 if nch > 0 {
  62                         needSub = false
  63                         break
  64                 }
  65         }
  66         chmap[ch]++
  67         if needSub && client.wsconn != nil {
  68                 go json.NewEncoder(client.wsconn).Encode(map[string]interface{}{
  69                         "method": "subscribe",
  70                         "filters": [][]interface{}{
  71                                 {"object_uuid", "=", uuid},
  72                                 {"event_type", "in", []string{"stderr", "crunch-run", "update"}},
  73                         },
  74                 })
  75         }
  76 }
  77
  78 func (client *arvadosClient) Unsubscribe(ch chan<- eventMessage, uuid string) {
  79         client.mtx.Lock()
  80         defer client.mtx.Unlock()
  81         chmap := client.notifying[uuid]
  82         if n := chmap[ch] - 1; n == 0 {
  83                 delete(chmap, ch)
  84                 if len(chmap) == 0 {
  85                         delete(client.notifying, uuid)
  86                 }
  87                 if client.wsconn != nil {
  88                         go json.NewEncoder(client.wsconn).Encode(map[string]interface{}{
  89                                 "method": "unsubscribe",
  90                                 "filters": [][]interface{}{
  91                                         {"object_uuid", "=", uuid},
  92                                         {"event_type", "in", []string{"stderr", "crunch-run", "update"}},
  93                                 },
  94                         })
  95                 }
  96         } else if n > 0 {
  97                 chmap[ch] = n
  98         }
  99 }
 100
 101 func (client *arvadosClient) Close() {
 102         client.mtx.Lock()
 103         defer client.mtx.Unlock()
 104         if client.notifying != nil {
 105                 client.notifying = nil
 106                 close(client.wantClose)
 107         }
 108 }
 109
 110 func (client *arvadosClient) runNotifier() {
 111 reconnect:
 112         for {
 113                 var cluster arvados.Cluster
 114                 err := client.RequestAndDecode(&cluster, "GET", arvados.EndpointConfigGet.Path, nil, nil)
 115                 if err != nil {
 116                         log.Warnf("error getting cluster config: %s", err)
 117                         time.Sleep(5 * time.Second)
 118                         continue reconnect
 119                 }
 120                 wsURL := cluster.Services.Websocket.ExternalURL
 121                 wsURL.Scheme = strings.Replace(wsURL.Scheme, "http", "ws", 1)
 122                 wsURL.Path = "/websocket"
 123                 wsURLNoToken := wsURL.String()
 124                 wsURL.RawQuery = url.Values{"api_token": []string{client.AuthToken}}.Encode()
 125                 conn, err := websocket.Dial(wsURL.String(), "", cluster.Services.Controller.ExternalURL.String())
 126                 if err != nil {
 127                         log.Warnf("websocket connection error: %s", err)
 128                         time.Sleep(5 * time.Second)
 129                         continue reconnect
 130                 }
 131                 log.Printf("connected to websocket at %s", wsURLNoToken)
 132
 133                 client.mtx.Lock()
 134                 client.wsconn = conn
 135                 resubscribe := make([]string, 0, len(client.notifying))
 136                 for uuid := range client.notifying {
 137                         resubscribe = append(resubscribe, uuid)
 138                 }
 139                 client.mtx.Unlock()
 140
 141                 go func() {
 142                         w := json.NewEncoder(conn)
 143                         for _, uuid := range resubscribe {
 144                                 w.Encode(map[string]interface{}{
 145                                         "method": "subscribe",
 146                                         "filters": [][]interface{}{
 147                                                 {"object_uuid", "=", uuid},
 148                                                 {"event_type", "in", []string{"stderr", "crunch-run", "crunchstat", "update"}},
 149                                         },
 150                                 })
 151                         }
 152                 }()
 153
 154                 r := json.NewDecoder(conn)
 155                 for {
 156                         var msg eventMessage
 157                         err := r.Decode(&msg)
 158                         select {
 159                         case <-client.wantClose:
 160                                 return
 161                         default:
 162                                 if err != nil {
 163                                         log.Printf("error decoding websocket message: %s", err)
 164                                         client.mtx.Lock()
 165                                         client.wsconn = nil
 166                                         client.mtx.Unlock()
 167                                         go conn.Close()
 168                                         continue reconnect
 169                                 }
 170                                 client.mtx.Lock()
 171                                 for ch := range client.notifying[msg.ObjectUUID] {
 172                                         ch <- msg
 173                                 }
 174                                 client.mtx.Unlock()
 175                         }
 176                 }
 177         }
 178 }
 179
 180 type arvadosContainerRunner struct {
 181         Client      *arvados.Client
 182         Name        string
 183         OutputName  string
 184         ProjectUUID string
 185         VCPUs       int
 186         RAM         int64
 187         Prog        string // if empty, run /proc/self/exe
 188         Args        []string
 189         Mounts      map[string]map[string]interface{}
 190         Priority    int
 191 }
 192
 193 func (runner *arvadosContainerRunner) Run() (string, error) {
 194         return runner.RunContext(context.Background())
 195 }
 196
 197 func (runner *arvadosContainerRunner) RunContext(ctx context.Context) (string, error) {
 198         if runner.ProjectUUID == "" {
 199                 return "", errors.New("cannot run arvados container: ProjectUUID not provided")
 200         }
 201
 202         mounts := map[string]map[string]interface{}{
 203                 "/mnt/output": {
 204                         "kind":     "collection",
 205                         "writable": true,
 206                 },
 207         }
 208         for path, mnt := range runner.Mounts {
 209                 mounts[path] = mnt
 210         }
 211
 212         prog := runner.Prog
 213         if prog == "" {
 214                 prog = "/mnt/cmd/lightning"
 215                 cmdUUID, err := runner.makeCommandCollection()
 216                 if err != nil {
 217                         return "", err
 218                 }
 219                 mounts["/mnt/cmd"] = map[string]interface{}{
 220                         "kind": "collection",
 221                         "uuid": cmdUUID,
 222                 }
 223         }
 224         command := append([]string{prog}, runner.Args...)
 225
 226         priority := runner.Priority
 227         if priority < 1 {
 228                 priority = 500
 229         }
 230         rc := arvados.RuntimeConstraints{
 231                 VCPUs:        runner.VCPUs,
 232                 RAM:          runner.RAM,
 233                 KeepCacheRAM: (1 << 26) * 2 * int64(runner.VCPUs),
 234         }
 235         outname := &runner.OutputName
 236         if *outname == "" {
 237                 outname = nil
 238         }
 239         var cr arvados.ContainerRequest
 240         err := runner.Client.RequestAndDecode(&cr, "POST", "arvados/v1/container_requests", nil, map[string]interface{}{
 241                 "container_request": map[string]interface{}{
 242                         "owner_uuid":          runner.ProjectUUID,
 243                         "name":                runner.Name,
 244                         "container_image":     "lightning-runtime",
 245                         "command":             command,
 246                         "mounts":              mounts,
 247                         "use_existing":        true,
 248                         "output_path":         "/mnt/output",
 249                         "output_name":         outname,
 250                         "runtime_constraints": rc,
 251                         "priority":            runner.Priority,
 252                         "state":               arvados.ContainerRequestStateCommitted,
 253                 },
 254         })
 255         if err != nil {
 256                 return "", err
 257         }
 258         log.Printf("container request UUID: %s", cr.UUID)
 259         log.Printf("container UUID: %s", cr.ContainerUUID)
 260
 261         logch := make(chan eventMessage)
 262         client := arvadosClient{Client: runner.Client}
 263         defer client.Close()
 264         subscribedUUID := ""
 265         defer func() {
 266                 if subscribedUUID != "" {
 267                         client.Unsubscribe(logch, subscribedUUID)
 268                 }
 269         }()
 270
 271         ticker := time.NewTicker(5 * time.Second)
 272         defer ticker.Stop()
 273
 274         neednewline := ""
 275
 276         lastState := cr.State
 277         refreshCR := func() {
 278                 err = runner.Client.RequestAndDecode(&cr, "GET", "arvados/v1/container_requests/"+cr.UUID, nil, nil)
 279                 if err != nil {
 280                         fmt.Fprint(os.Stderr, neednewline)
 281                         log.Printf("error getting container request: %s", err)
 282                         return
 283                 }
 284                 if lastState != cr.State {
 285                         fmt.Fprint(os.Stderr, neednewline)
 286                         log.Printf("container request state: %s", cr.State)
 287                         lastState = cr.State
 288                 }
 289                 if subscribedUUID != cr.ContainerUUID {
 290                         fmt.Fprint(os.Stderr, neednewline)
 291                         neednewline = ""
 292                         if subscribedUUID != "" {
 293                                 client.Unsubscribe(logch, subscribedUUID)
 294                         }
 295                         client.Subscribe(logch, cr.ContainerUUID)
 296                         subscribedUUID = cr.ContainerUUID
 297                 }
 298         }
 299
 300         var reCrunchstat = regexp.MustCompile(`mem .* rss`)
 301 waitctr:
 302         for cr.State != arvados.ContainerRequestStateFinal {
 303                 select {
 304                 case <-ctx.Done():
 305                         err := runner.Client.RequestAndDecode(&cr, "PATCH", "arvados/v1/container_requests/"+cr.UUID, nil, map[string]interface{}{
 306                                 "container_request": map[string]interface{}{
 307                                         "priority": 0,
 308                                 },
 309                         })
 310                         if err != nil {
 311                                 log.Errorf("error while trying to cancel container request %s: %s", cr.UUID, err)
 312                         }
 313                         break waitctr
 314                 case <-ticker.C:
 315                         refreshCR()
 316                 case msg := <-logch:
 317                         switch msg.EventType {
 318                         case "update":
 319                                 refreshCR()
 320                         case "stderr":
 321                                 for _, line := range strings.Split(msg.Properties.Text, "\n") {
 322                                         if line != "" {
 323                                                 fmt.Fprint(os.Stderr, neednewline)
 324                                                 neednewline = ""
 325                                                 log.Print(line)
 326                                         }
 327                                 }
 328                         case "crunchstat":
 329                                 for _, line := range strings.Split(msg.Properties.Text, "\n") {
 330                                         mem := reCrunchstat.FindString(line)
 331                                         if mem != "" {
 332                                                 fmt.Fprintf(os.Stderr, "%s               \r", mem)
 333                                                 neednewline = "\n"
 334                                         }
 335                                 }
 336                         }
 337                 }
 338         }
 339         fmt.Fprint(os.Stderr, neednewline)
 340
 341         if err := ctx.Err(); err != nil {
 342                 return "", err
 343         }
 344
 345         var c arvados.Container
 346         err = runner.Client.RequestAndDecode(&c, "GET", "arvados/v1/containers/"+cr.ContainerUUID, nil, nil)
 347         if err != nil {
 348                 return "", err
 349         } else if c.State != arvados.ContainerStateComplete {
 350                 return "", fmt.Errorf("container did not complete: %s", c.State)
 351         } else if c.ExitCode != 0 {
 352                 return "", fmt.Errorf("container exited %d", c.ExitCode)
 353         }
 354         return cr.OutputUUID, err
 355 }
 356
 357 var collectionInPathRe = regexp.MustCompile(`^(.*/)?([0-9a-f]{32}\+[0-9]+|[0-9a-z]{5}-[0-9a-z]{5}-[0-9a-z]{15})(/.*)?$`)
 358
 359 func (runner *arvadosContainerRunner) TranslatePaths(paths ...*string) error {
 360         if runner.Mounts == nil {
 361                 runner.Mounts = make(map[string]map[string]interface{})
 362         }
 363         for _, path := range paths {
 364                 if *path == "" || *path == "-" {
 365                         continue
 366                 }
 367                 m := collectionInPathRe.FindStringSubmatch(*path)
 368                 if m == nil {
 369                         return fmt.Errorf("cannot find uuid in path: %q", *path)
 370                 }
 371                 uuid := m[2]
 372                 mnt, ok := runner.Mounts["/mnt/"+uuid]
 373                 if !ok {
 374                         mnt = map[string]interface{}{
 375                                 "kind": "collection",
 376                                 "uuid": uuid,
 377                         }
 378                         runner.Mounts["/mnt/"+uuid] = mnt
 379                 }
 380                 *path = "/mnt/" + uuid + m[3]
 381         }
 382         return nil
 383 }
 384
 385 var mtxMakeCommandCollection sync.Mutex
 386
 387 func (runner *arvadosContainerRunner) makeCommandCollection() (string, error) {
 388         mtxMakeCommandCollection.Lock()
 389         defer mtxMakeCommandCollection.Unlock()
 390         exe, err := ioutil.ReadFile("/proc/self/exe")
 391         if err != nil {
 392                 return "", err
 393         }
 394         b2 := blake2b.Sum256(exe)
 395         cname := fmt.Sprintf("lightning-%x", b2)
 396         var existing arvados.CollectionList
 397         err = runner.Client.RequestAndDecode(&existing, "GET", "arvados/v1/collections", nil, arvados.ListOptions{
 398                 Limit: 1,
 399                 Count: "none",
 400                 Filters: []arvados.Filter{
 401                         {Attr: "name", Operator: "=", Operand: cname},
 402                         {Attr: "owner_uuid", Operator: "=", Operand: runner.ProjectUUID},
 403                 },
 404         })
 405         if err != nil {
 406                 return "", err
 407         }
 408         if len(existing.Items) > 0 {
 409                 uuid := existing.Items[0].UUID
 410                 log.Printf("using lightning binary in existing collection %s (name is %q; did not verify whether content matches)", uuid, cname)
 411                 return uuid, nil
 412         }
 413         log.Printf("writing lightning binary to new collection %q", cname)
 414         ac, err := arvadosclient.New(runner.Client)
 415         if err != nil {
 416                 return "", err
 417         }
 418         kc := keepclient.New(ac)
 419         var coll arvados.Collection
 420         fs, err := coll.FileSystem(runner.Client, kc)
 421         if err != nil {
 422                 return "", err
 423         }
 424         f, err := fs.OpenFile("lightning", os.O_CREATE|os.O_WRONLY, 0777)
 425         if err != nil {
 426                 return "", err
 427         }
 428         _, err = f.Write(exe)
 429         if err != nil {
 430                 return "", err
 431         }
 432         err = f.Close()
 433         if err != nil {
 434                 return "", err
 435         }
 436         mtxt, err := fs.MarshalManifest(".")
 437         if err != nil {
 438                 return "", err
 439         }
 440         err = runner.Client.RequestAndDecode(&coll, "POST", "arvados/v1/collections", nil, map[string]interface{}{
 441                 "collection": map[string]interface{}{
 442                         "owner_uuid":    runner.ProjectUUID,
 443                         "manifest_text": mtxt,
 444                         "name":          cname,
 445                 },
 446         })
 447         if err != nil {
 448                 return "", err
 449         }
 450         log.Printf("stored lightning binary in new collection %s", coll.UUID)
 451         return coll.UUID, nil
 452 }