15 "git.arvados.org/arvados.git/sdk/go/arvados"
16 "git.arvados.org/arvados.git/sdk/go/arvadosclient"
17 "git.arvados.org/arvados.git/sdk/go/keepclient"
18 log "github.com/sirupsen/logrus"
19 "golang.org/x/crypto/blake2b"
20 "golang.org/x/net/websocket"
23 type eventMessage struct {
25 ObjectUUID string `json:"object_uuid"`
26 EventType string `json:"event_type"`
32 type arvadosClient struct {
34 notifying map[string]map[chan<- eventMessage]int
35 wantClose chan struct{}
36 wsconn *websocket.Conn
40 // Listen for events concerning the given uuids. When an event occurs
41 // (and after connecting/reconnecting to the event stream), send each
42 // uuid to ch. If a {ch, uuid} pair is subscribed twice, the uuid will
43 // be sent only once for each update, but two Unsubscribe calls will
44 // be needed to stop sending them.
45 func (client *arvadosClient) Subscribe(ch chan<- eventMessage, uuid string) {
47 defer client.mtx.Unlock()
48 if client.notifying == nil {
49 client.notifying = map[string]map[chan<- eventMessage]int{}
50 client.wantClose = make(chan struct{})
51 go client.runNotifier()
53 chmap := client.notifying[uuid]
55 chmap = map[chan<- eventMessage]int{}
56 client.notifying[uuid] = chmap
59 for _, nch := range chmap {
66 if needSub && client.wsconn != nil {
67 go json.NewEncoder(client.wsconn).Encode(map[string]interface{}{
68 "method": "subscribe",
69 "filters": [][]interface{}{
70 {"object_uuid", "=", uuid},
71 {"event_type", "in", []string{"stderr", "crunch-run", "update"}},
77 func (client *arvadosClient) Unsubscribe(ch chan<- eventMessage, uuid string) {
79 defer client.mtx.Unlock()
80 chmap := client.notifying[uuid]
81 if n := chmap[ch] - 1; n == 0 {
84 delete(client.notifying, uuid)
86 if client.wsconn != nil {
87 go json.NewEncoder(client.wsconn).Encode(map[string]interface{}{
88 "method": "unsubscribe",
89 "filters": [][]interface{}{
90 {"object_uuid", "=", uuid},
91 {"event_type", "in", []string{"stderr", "crunch-run", "update"}},
100 func (client *arvadosClient) Close() {
102 defer client.mtx.Unlock()
103 if client.notifying != nil {
104 client.notifying = nil
105 close(client.wantClose)
109 func (client *arvadosClient) runNotifier() {
112 var cluster arvados.Cluster
113 err := client.RequestAndDecode(&cluster, "GET", arvados.EndpointConfigGet.Path, nil, nil)
115 log.Warnf("error getting cluster config: %s", err)
116 time.Sleep(5 * time.Second)
119 wsURL := cluster.Services.Websocket.ExternalURL
120 wsURL.Scheme = strings.Replace(wsURL.Scheme, "http", "ws", 1)
121 wsURL.Path = "/websocket"
122 wsURL.RawQuery = url.Values{"api_token": []string{client.AuthToken}}.Encode()
123 conn, err := websocket.Dial(wsURL.String(), "", cluster.Services.Controller.ExternalURL.String())
125 log.Warnf("websocket connection error: %s", err)
126 time.Sleep(5 * time.Second)
129 log.Printf("connected to websocket at %s", wsURL)
133 resubscribe := make([]string, 0, len(client.notifying))
134 for uuid := range client.notifying {
135 resubscribe = append(resubscribe, uuid)
140 w := json.NewEncoder(conn)
141 for _, uuid := range resubscribe {
142 w.Encode(map[string]interface{}{
143 "method": "subscribe",
144 "filters": [][]interface{}{
145 {"object_uuid", "=", uuid},
146 {"event_type", "in", []string{"stderr", "crunch-run", "update"}},
152 r := json.NewDecoder(conn)
155 err := r.Decode(&msg)
157 case <-client.wantClose:
161 log.Printf("error decoding websocket message: %s", err)
169 for ch := range client.notifying[msg.ObjectUUID] {
178 type arvadosContainerRunner struct {
179 Client *arvados.Client
184 Prog string // if empty, run /proc/self/exe
186 Mounts map[string]map[string]interface{}
190 func (runner *arvadosContainerRunner) Run() (string, error) {
191 if runner.ProjectUUID == "" {
192 return "", errors.New("cannot run arvados container: ProjectUUID not provided")
195 mounts := map[string]map[string]interface{}{
197 "kind": "collection",
201 for path, mnt := range runner.Mounts {
207 prog = "/mnt/cmd/lightning"
208 cmdUUID, err := runner.makeCommandCollection()
212 mounts["/mnt/cmd"] = map[string]interface{}{
213 "kind": "collection",
217 command := append([]string{prog}, runner.Args...)
219 priority := runner.Priority
223 rc := arvados.RuntimeConstraints{
226 KeepCacheRAM: (1 << 26) * 2 * int64(runner.VCPUs),
228 var cr arvados.ContainerRequest
229 err := runner.Client.RequestAndDecode(&cr, "POST", "arvados/v1/container_requests", nil, map[string]interface{}{
230 "container_request": map[string]interface{}{
231 "owner_uuid": runner.ProjectUUID,
233 "container_image": "lightning-runtime",
236 "use_existing": true,
237 "output_path": "/mnt/output",
238 "runtime_constraints": rc,
239 "priority": runner.Priority,
240 "state": arvados.ContainerRequestStateCommitted,
246 log.Printf("container request UUID: %s", cr.UUID)
247 log.Printf("container UUID: %s", cr.ContainerUUID)
249 logch := make(chan eventMessage)
250 client := arvadosClient{Client: runner.Client}
254 if subscribedUUID != "" {
255 client.Unsubscribe(logch, subscribedUUID)
259 ticker := time.NewTicker(5 * time.Second)
262 lastState := cr.State
263 refreshCR := func() {
264 err = runner.Client.RequestAndDecode(&cr, "GET", "arvados/v1/container_requests/"+cr.UUID, nil, nil)
266 log.Printf("error getting container request: %s", err)
269 if lastState != cr.State {
270 log.Printf("container request state: %s", cr.State)
273 if subscribedUUID != cr.ContainerUUID {
274 if subscribedUUID != "" {
275 client.Unsubscribe(logch, subscribedUUID)
277 client.Subscribe(logch, cr.ContainerUUID)
278 subscribedUUID = cr.ContainerUUID
282 for cr.State != arvados.ContainerRequestStateFinal {
287 switch msg.EventType {
291 for _, line := range strings.Split(msg.Properties.Text, "\n") {
300 var c arvados.Container
301 err = runner.Client.RequestAndDecode(&c, "GET", "arvados/v1/containers/"+cr.ContainerUUID, nil, nil)
304 } else if c.State != arvados.ContainerStateComplete {
305 return "", fmt.Errorf("container did not complete: %s", c.State)
306 } else if c.ExitCode != 0 {
307 return "", fmt.Errorf("container exited %d", c.ExitCode)
309 return cr.OutputUUID, err
312 var collectionInPathRe = regexp.MustCompile(`^(.*/)?([0-9a-f]{32}\+[0-9]+|[0-9a-z]{5}-[0-9a-z]{5}-[0-9a-z]{15})(/.*)?$`)
314 func (runner *arvadosContainerRunner) TranslatePaths(paths ...*string) error {
315 if runner.Mounts == nil {
316 runner.Mounts = make(map[string]map[string]interface{})
318 for _, path := range paths {
319 if *path == "" || *path == "-" {
322 m := collectionInPathRe.FindStringSubmatch(*path)
324 return fmt.Errorf("cannot find uuid in path: %q", *path)
327 mnt, ok := runner.Mounts["/mnt/"+uuid]
329 mnt = map[string]interface{}{
330 "kind": "collection",
333 runner.Mounts["/mnt/"+uuid] = mnt
335 *path = "/mnt/" + uuid + m[3]
340 func (runner *arvadosContainerRunner) makeCommandCollection() (string, error) {
341 exe, err := ioutil.ReadFile("/proc/self/exe")
345 b2 := blake2b.Sum256(exe)
346 cname := fmt.Sprintf("lightning-%x", b2)
347 var existing arvados.CollectionList
348 err = runner.Client.RequestAndDecode(&existing, "GET", "arvados/v1/collections", nil, arvados.ListOptions{
351 Filters: []arvados.Filter{
352 {Attr: "name", Operator: "=", Operand: cname},
353 {Attr: "owner_uuid", Operator: "=", Operand: runner.ProjectUUID},
359 if len(existing.Items) > 0 {
360 uuid := existing.Items[0].UUID
361 log.Printf("using lightning binary in existing collection %s (name is %q; did not verify whether content matches)", uuid, cname)
364 log.Printf("writing lightning binary to new collection %q", cname)
365 ac, err := arvadosclient.New(runner.Client)
369 kc := keepclient.New(ac)
370 var coll arvados.Collection
371 fs, err := coll.FileSystem(runner.Client, kc)
375 f, err := fs.OpenFile("lightning", os.O_CREATE|os.O_WRONLY, 0777)
379 _, err = f.Write(exe)
387 mtxt, err := fs.MarshalManifest(".")
391 err = runner.Client.RequestAndDecode(&coll, "POST", "arvados/v1/collections", nil, map[string]interface{}{
392 "collection": map[string]interface{}{
393 "owner_uuid": runner.ProjectUUID,
394 "manifest_text": mtxt,
401 log.Printf("stored lightning binary in new collection %s", coll.UUID)
402 return coll.UUID, nil