arvados.go

   1 package main
   2
   3 import (
   4         "errors"
   5         "fmt"
   6         "io/ioutil"
   7         "log"
   8         "os"
   9         "regexp"
  10
  11         "git.arvados.org/arvados.git/sdk/go/arvados"
  12         "git.arvados.org/arvados.git/sdk/go/arvadosclient"
  13         "git.arvados.org/arvados.git/sdk/go/keepclient"
  14         "golang.org/x/crypto/blake2b"
  15 )
  16
  17 type arvadosContainerRunner struct {
  18         Client      *arvados.Client
  19         Name        string
  20         ProjectUUID string
  21         VCPUs       int
  22         RAM         int64
  23         Prog        string // if empty, run /proc/self/exe
  24         Args        []string
  25         Mounts      map[string]string
  26 }
  27
  28 var (
  29         collectionInPathRe = regexp.MustCompile(`^(.*/)?([0-9a-f]{32}\+[0-9]+|[0-9a-z]{5}-[0-9a-z]{5}-[0-9a-z]{15})(/.*)?$`)
  30 )
  31
  32 func (runner *arvadosContainerRunner) Run() error {
  33         if runner.ProjectUUID == "" {
  34                 return errors.New("cannot run arvados container: ProjectUUID not provided")
  35         }
  36         mounts := map[string]map[string]interface{}{
  37                 "/mnt/output": {
  38                         "kind":     "tmp",
  39                         "writable": true,
  40                         "capacity": 100000000000,
  41                 },
  42         }
  43
  44         prog := runner.Prog
  45         if prog == "" {
  46                 prog = "/mnt/cmd/lightning"
  47                 cmdUUID, err := runner.makeCommandCollection()
  48                 if err != nil {
  49                         return err
  50                 }
  51                 mounts["/mnt/cmd"] = map[string]interface{}{
  52                         "kind": "collection",
  53                         "uuid": cmdUUID,
  54                 }
  55         }
  56         command := append([]string{prog}, runner.Args...)
  57
  58         for uuid, mnt := range runner.Mounts {
  59                 mounts[mnt] = map[string]interface{}{
  60                         "kind": "collection",
  61                         "uuid": uuid,
  62                 }
  63         }
  64         rc := arvados.RuntimeConstraints{
  65                 VCPUs:        runner.VCPUs,
  66                 RAM:          runner.RAM,
  67                 KeepCacheRAM: (1 << 26) * 2 * int64(runner.VCPUs),
  68         }
  69         var cr arvados.ContainerRequest
  70         err := runner.Client.RequestAndDecode(&cr, "POST", "arvados/v1/container_requests", nil, map[string]interface{}{
  71                 "container_request": map[string]interface{}{
  72                         "owner_uuid":          runner.ProjectUUID,
  73                         "name":                runner.Name,
  74                         "container_image":     "lightning-runtime",
  75                         "command":             command,
  76                         "mounts":              mounts,
  77                         "use_existing":        true,
  78                         "output_path":         "/mnt/output",
  79                         "runtime_constraints": rc,
  80                         "priority":            1,
  81                         "state":               arvados.ContainerRequestStateCommitted,
  82                 },
  83         })
  84         log.Print(cr.UUID)
  85         return err
  86 }
  87
  88 func (runner *arvadosContainerRunner) TranslatePaths(paths ...*string) error {
  89         if runner.Mounts == nil {
  90                 runner.Mounts = make(map[string]string)
  91         }
  92         for _, path := range paths {
  93                 if *path == "" {
  94                         continue
  95                 }
  96                 m := collectionInPathRe.FindStringSubmatch(*path)
  97                 if m == nil {
  98                         return fmt.Errorf("cannot find uuid in path: %q", *path)
  99                 }
 100                 uuid := m[2]
 101                 mnt, ok := runner.Mounts[uuid]
 102                 if !ok {
 103                         mnt = "/mnt/" + uuid
 104                         runner.Mounts[uuid] = mnt
 105                 }
 106                 *path = mnt + m[3]
 107         }
 108         return nil
 109 }
 110
 111 func (runner *arvadosContainerRunner) makeCommandCollection() (string, error) {
 112         exe, err := ioutil.ReadFile("/proc/self/exe")
 113         if err != nil {
 114                 return "", err
 115         }
 116         b2 := blake2b.Sum256(exe)
 117         cname := fmt.Sprintf("lightning-%x", b2)
 118         var existing arvados.CollectionList
 119         err = runner.Client.RequestAndDecode(&existing, "GET", "arvados/v1/collections", nil, arvados.ListOptions{
 120                 Limit: 1,
 121                 Count: "none",
 122                 Filters: []arvados.Filter{
 123                         {Attr: "name", Operator: "=", Operand: cname},
 124                         {Attr: "owner_uuid", Operator: "=", Operand: runner.ProjectUUID},
 125                 },
 126         })
 127         if err != nil {
 128                 return "", err
 129         }
 130         if len(existing.Items) > 0 {
 131                 uuid := existing.Items[0].UUID
 132                 log.Printf("using existing collection %q named %q (did not verify whether content matches)", uuid, cname)
 133                 return uuid, nil
 134         }
 135         log.Printf("writing lightning binary to new collection %q", cname)
 136         ac, err := arvadosclient.New(runner.Client)
 137         if err != nil {
 138                 return "", err
 139         }
 140         kc := keepclient.New(ac)
 141         var coll arvados.Collection
 142         fs, err := coll.FileSystem(runner.Client, kc)
 143         if err != nil {
 144                 return "", err
 145         }
 146         f, err := fs.OpenFile("lightning", os.O_CREATE|os.O_WRONLY, 0777)
 147         if err != nil {
 148                 return "", err
 149         }
 150         _, err = f.Write(exe)
 151         if err != nil {
 152                 return "", err
 153         }
 154         err = f.Close()
 155         if err != nil {
 156                 return "", err
 157         }
 158         mtxt, err := fs.MarshalManifest(".")
 159         if err != nil {
 160                 return "", err
 161         }
 162         err = runner.Client.RequestAndDecode(&coll, "POST", "arvados/v1/collections", nil, map[string]interface{}{
 163                 "collection": map[string]interface{}{
 164                         "owner_uuid":    runner.ProjectUUID,
 165                         "manifest_text": mtxt,
 166                         "name":          cname,
 167                 },
 168         })
 169         if err != nil {
 170                 return "", err
 171         }
 172         log.Printf("collection: %#v", coll)
 173         return coll.UUID, nil
 174 }