19524: Limit size of PCA input matrix.
[lightning.git] / 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 }
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 }