package main import ( "encoding/json" "errors" "fmt" "io/ioutil" "net/url" "os" "regexp" "strings" "sync" "time" "git.arvados.org/arvados.git/sdk/go/arvados" "git.arvados.org/arvados.git/sdk/go/arvadosclient" "git.arvados.org/arvados.git/sdk/go/keepclient" log "github.com/sirupsen/logrus" "golang.org/x/crypto/blake2b" "golang.org/x/net/websocket" ) type eventMessage struct { Status int ObjectUUID string `json:"object_uuid"` EventType string `json:"event_type"` Properties struct { Text string } } type arvadosClient struct { *arvados.Client notifying map[string]map[chan<- eventMessage]int wantClose chan struct{} wsconn *websocket.Conn mtx sync.Mutex } // Listen for events concerning the given uuids. When an event occurs // (and after connecting/reconnecting to the event stream), send each // uuid to ch. If a {ch, uuid} pair is subscribed twice, the uuid will // be sent only once for each update, but two Unsubscribe calls will // be needed to stop sending them. func (client *arvadosClient) Subscribe(ch chan<- eventMessage, uuid string) { client.mtx.Lock() defer client.mtx.Unlock() if client.notifying == nil { client.notifying = map[string]map[chan<- eventMessage]int{} client.wantClose = make(chan struct{}) go client.runNotifier() } chmap := client.notifying[uuid] if chmap == nil { chmap = map[chan<- eventMessage]int{} client.notifying[uuid] = chmap } needSub := true for _, nch := range chmap { if nch > 0 { needSub = false break } } chmap[ch]++ if needSub && client.wsconn != nil { go json.NewEncoder(client.wsconn).Encode(map[string]interface{}{ "method": "subscribe", "filters": [][]interface{}{ {"object_uuid", "=", uuid}, {"event_type", "in", []string{"stderr", "crunch-run", "update"}}, }, }) } } func (client *arvadosClient) Unsubscribe(ch chan<- eventMessage, uuid string) { client.mtx.Lock() defer client.mtx.Unlock() chmap := client.notifying[uuid] if n := chmap[ch] - 1; n == 0 { delete(chmap, ch) if len(chmap) == 0 { delete(client.notifying, uuid) } if client.wsconn != nil { go json.NewEncoder(client.wsconn).Encode(map[string]interface{}{ "method": "unsubscribe", "filters": [][]interface{}{ {"object_uuid", "=", uuid}, {"event_type", "in", []string{"stderr", "crunch-run", "update"}}, }, }) } } else if n > 0 { chmap[ch] = n } } func (client *arvadosClient) Close() { client.mtx.Lock() defer client.mtx.Unlock() if client.notifying != nil { client.notifying = nil close(client.wantClose) } } func (client *arvadosClient) runNotifier() { reconnect: for { var cluster arvados.Cluster err := client.RequestAndDecode(&cluster, "GET", arvados.EndpointConfigGet.Path, nil, nil) if err != nil { log.Warnf("error getting cluster config: %s", err) time.Sleep(5 * time.Second) continue reconnect } wsURL := cluster.Services.Websocket.ExternalURL wsURL.Scheme = strings.Replace(wsURL.Scheme, "http", "ws", 1) wsURL.Path = "/websocket" wsURLNoToken := wsURL.String() wsURL.RawQuery = url.Values{"api_token": []string{client.AuthToken}}.Encode() conn, err := websocket.Dial(wsURL.String(), "", cluster.Services.Controller.ExternalURL.String()) if err != nil { log.Warnf("websocket connection error: %s", err) time.Sleep(5 * time.Second) continue reconnect } log.Printf("connected to websocket at %s", wsURLNoToken) client.mtx.Lock() client.wsconn = conn resubscribe := make([]string, 0, len(client.notifying)) for uuid := range client.notifying { resubscribe = append(resubscribe, uuid) } client.mtx.Unlock() go func() { w := json.NewEncoder(conn) for _, uuid := range resubscribe { w.Encode(map[string]interface{}{ "method": "subscribe", "filters": [][]interface{}{ {"object_uuid", "=", uuid}, {"event_type", "in", []string{"stderr", "crunch-run", "crunchstat", "update"}}, }, }) } }() r := json.NewDecoder(conn) for { var msg eventMessage err := r.Decode(&msg) select { case <-client.wantClose: return default: if err != nil { log.Printf("error decoding websocket message: %s", err) client.mtx.Lock() client.wsconn = nil client.mtx.Unlock() go conn.Close() continue reconnect } client.mtx.Lock() for ch := range client.notifying[msg.ObjectUUID] { ch <- msg } client.mtx.Unlock() } } } } type arvadosContainerRunner struct { Client *arvados.Client Name string ProjectUUID string VCPUs int RAM int64 Prog string // if empty, run /proc/self/exe Args []string Mounts map[string]map[string]interface{} Priority int } func (runner *arvadosContainerRunner) Run() (string, error) { if runner.ProjectUUID == "" { return "", errors.New("cannot run arvados container: ProjectUUID not provided") } mounts := map[string]map[string]interface{}{ "/mnt/output": { "kind": "collection", "writable": true, }, } for path, mnt := range runner.Mounts { mounts[path] = mnt } prog := runner.Prog if prog == "" { prog = "/mnt/cmd/lightning" cmdUUID, err := runner.makeCommandCollection() if err != nil { return "", err } mounts["/mnt/cmd"] = map[string]interface{}{ "kind": "collection", "uuid": cmdUUID, } } command := append([]string{prog}, runner.Args...) priority := runner.Priority if priority < 1 { priority = 500 } rc := arvados.RuntimeConstraints{ VCPUs: runner.VCPUs, RAM: runner.RAM, KeepCacheRAM: (1 << 26) * 2 * int64(runner.VCPUs), } var cr arvados.ContainerRequest err := runner.Client.RequestAndDecode(&cr, "POST", "arvados/v1/container_requests", nil, map[string]interface{}{ "container_request": map[string]interface{}{ "owner_uuid": runner.ProjectUUID, "name": runner.Name, "container_image": "lightning-runtime", "command": command, "mounts": mounts, "use_existing": true, "output_path": "/mnt/output", "runtime_constraints": rc, "priority": runner.Priority, "state": arvados.ContainerRequestStateCommitted, }, }) if err != nil { return "", err } log.Printf("container request UUID: %s", cr.UUID) log.Printf("container UUID: %s", cr.ContainerUUID) logch := make(chan eventMessage) client := arvadosClient{Client: runner.Client} defer client.Close() subscribedUUID := "" defer func() { if subscribedUUID != "" { client.Unsubscribe(logch, subscribedUUID) } }() ticker := time.NewTicker(5 * time.Second) defer ticker.Stop() neednewline := "" lastState := cr.State refreshCR := func() { err = runner.Client.RequestAndDecode(&cr, "GET", "arvados/v1/container_requests/"+cr.UUID, nil, nil) if err != nil { fmt.Fprint(os.Stderr, neednewline) log.Printf("error getting container request: %s", err) return } if lastState != cr.State { fmt.Fprint(os.Stderr, neednewline) log.Printf("container request state: %s", cr.State) lastState = cr.State } if subscribedUUID != cr.ContainerUUID { fmt.Fprint(os.Stderr, neednewline) neednewline = "" if subscribedUUID != "" { client.Unsubscribe(logch, subscribedUUID) } client.Subscribe(logch, cr.ContainerUUID) subscribedUUID = cr.ContainerUUID } } var reCrunchstat = regexp.MustCompile(`mem .* rss`) for cr.State != arvados.ContainerRequestStateFinal { select { case <-ticker.C: refreshCR() case msg := <-logch: switch msg.EventType { case "update": refreshCR() case "stderr": for _, line := range strings.Split(msg.Properties.Text, "\n") { if line != "" { fmt.Fprint(os.Stderr, neednewline) neednewline = "" log.Print(line) } } case "crunchstat": for _, line := range strings.Split(msg.Properties.Text, "\n") { mem := reCrunchstat.FindString(line) if mem != "" { fmt.Fprintf(os.Stderr, "%s \r", mem) neednewline = "\n" } } } } } fmt.Fprint(os.Stderr, neednewline) var c arvados.Container err = runner.Client.RequestAndDecode(&c, "GET", "arvados/v1/containers/"+cr.ContainerUUID, nil, nil) if err != nil { return "", err } else if c.State != arvados.ContainerStateComplete { return "", fmt.Errorf("container did not complete: %s", c.State) } else if c.ExitCode != 0 { return "", fmt.Errorf("container exited %d", c.ExitCode) } return cr.OutputUUID, err } var collectionInPathRe = regexp.MustCompile(`^(.*/)?([0-9a-f]{32}\+[0-9]+|[0-9a-z]{5}-[0-9a-z]{5}-[0-9a-z]{15})(/.*)?$`) func (runner *arvadosContainerRunner) TranslatePaths(paths ...*string) error { if runner.Mounts == nil { runner.Mounts = make(map[string]map[string]interface{}) } for _, path := range paths { if *path == "" || *path == "-" { continue } m := collectionInPathRe.FindStringSubmatch(*path) if m == nil { return fmt.Errorf("cannot find uuid in path: %q", *path) } uuid := m[2] mnt, ok := runner.Mounts["/mnt/"+uuid] if !ok { mnt = map[string]interface{}{ "kind": "collection", "uuid": uuid, } runner.Mounts["/mnt/"+uuid] = mnt } *path = "/mnt/" + uuid + m[3] } return nil } func (runner *arvadosContainerRunner) makeCommandCollection() (string, error) { exe, err := ioutil.ReadFile("/proc/self/exe") if err != nil { return "", err } b2 := blake2b.Sum256(exe) cname := fmt.Sprintf("lightning-%x", b2) var existing arvados.CollectionList err = runner.Client.RequestAndDecode(&existing, "GET", "arvados/v1/collections", nil, arvados.ListOptions{ Limit: 1, Count: "none", Filters: []arvados.Filter{ {Attr: "name", Operator: "=", Operand: cname}, {Attr: "owner_uuid", Operator: "=", Operand: runner.ProjectUUID}, }, }) if err != nil { return "", err } if len(existing.Items) > 0 { uuid := existing.Items[0].UUID log.Printf("using lightning binary in existing collection %s (name is %q; did not verify whether content matches)", uuid, cname) return uuid, nil } log.Printf("writing lightning binary to new collection %q", cname) ac, err := arvadosclient.New(runner.Client) if err != nil { return "", err } kc := keepclient.New(ac) var coll arvados.Collection fs, err := coll.FileSystem(runner.Client, kc) if err != nil { return "", err } f, err := fs.OpenFile("lightning", os.O_CREATE|os.O_WRONLY, 0777) if err != nil { return "", err } _, err = f.Write(exe) if err != nil { return "", err } err = f.Close() if err != nil { return "", err } mtxt, err := fs.MarshalManifest(".") if err != nil { return "", err } err = runner.Client.RequestAndDecode(&coll, "POST", "arvados/v1/collections", nil, map[string]interface{}{ "collection": map[string]interface{}{ "owner_uuid": runner.ProjectUUID, "manifest_text": mtxt, "name": cname, }, }) if err != nil { return "", err } log.Printf("stored lightning binary in new collection %s", coll.UUID) return coll.UUID, nil }