Generate numpy matrices from slices.

[lightning.git] / arvados.go

diff --git a/arvados.go b/arvados.go
index 811d512e65339bbfcade4f93121256d809b07be5..c952a9870c77cab73c35237440555fb2b551ce23 100644 (file)
--- a/arvados.go
+++ b/arvados.go
@@ -1,3 +1,7 @@
+// Copyright (C) The Lightning Authors. All rights reserved.
+//
+// SPDX-License-Identifier: AGPL-3.0
+
 package lightning
 
 import (
@@ -11,7 +15,6 @@ import (
        "net/url"
        "os"
        "regexp"
-       "runtime"
        "strings"
        "sync"
        "time"
@@ -268,6 +271,9 @@ func (runner *arvadosContainerRunner) RunContext(ctx context.Context) (string, e
                                Preemptible: true,
                                Partitions:  []string{},
                        },
+                       "environment": map[string]string{
+                               "GOMAXPROCS": fmt.Sprintf("%d", rc.VCPUs),
+                       },
                },
        })
        if err != nil {
@@ -282,6 +288,7 @@ func (runner *arvadosContainerRunner) RunContext(ctx context.Context) (string, e
        subscribedUUID := ""
        defer func() {
                if subscribedUUID != "" {
+                       log.Printf("unsubscribe container UUID: %s", subscribedUUID)
                        client.Unsubscribe(logch, subscribedUUID)
                }
        }()
@@ -305,8 +312,10 @@ func (runner *arvadosContainerRunner) RunContext(ctx context.Context) (string, e
                        fmt.Fprint(os.Stderr, neednewline)
                        neednewline = ""
                        if subscribedUUID != "" {
+                               log.Printf("unsubscribe container UUID: %s", subscribedUUID)
                                client.Unsubscribe(logch, subscribedUUID)
                        }
+                       log.Printf("subscribe container UUID: %s", cr.ContainerUUID)
                        client.Subscribe(logch, cr.ContainerUUID)
                        subscribedUUID = cr.ContainerUUID
                }
@@ -504,11 +513,17 @@ func (gr gzipr) Close() error {
 
 var (
        arvadosClientFromEnv = arvados.NewClientFromEnv()
+       keepClient           *keepclient.KeepClient
        siteFS               arvados.CustomFileSystem
        siteFSMtx            sync.Mutex
 )
 
-func open(fnm string) (io.ReadCloser, error) {
+type file interface {
+       io.ReadCloser
+       Readdir(n int) ([]os.FileInfo, error)
+}
+
+func open(fnm string) (file, error) {
        if os.Getenv("ARVADOS_API_HOST") == "" {
                return os.Open(fnm)
        }
@@ -517,8 +532,8 @@ func open(fnm string) (io.ReadCloser, error) {
                return os.Open(fnm)
        }
        uuid := m[2]
-       mnt := "/mnt/" + uuid + "/"
-       if !strings.HasPrefix(fnm, mnt) {
+       mnt := "/mnt/" + uuid
+       if fnm != mnt && !strings.HasPrefix(fnm, mnt+"/") {
                return os.Open(fnm)
        }
 
@@ -531,15 +546,15 @@ func open(fnm string) (io.ReadCloser, error) {
                        return nil, err
                }
                ac.Client = arvados.DefaultSecureClient
-               kc := keepclient.New(ac)
+               keepClient = keepclient.New(ac)
                // Don't use keepclient's default short timeouts.
-               kc.HTTPClient = arvados.DefaultSecureClient
-               // Guess max concurrent readers, hope to avoid cache
-               // thrashing.
-               kc.BlockCache = &keepclient.BlockCache{MaxBlocks: runtime.NumCPU() * 3}
-               siteFS = arvadosClientFromEnv.SiteFileSystem(kc)
+               keepClient.HTTPClient = arvados.DefaultSecureClient
+               keepClient.BlockCache = &keepclient.BlockCache{MaxBlocks: 4}
+               siteFS = arvadosClientFromEnv.SiteFileSystem(keepClient)
+       } else {
+               keepClient.BlockCache.MaxBlocks++
        }
 
        log.Infof("reading %q from %s using Arvados client", fnm[len(mnt):], uuid)
-       return siteFS.Open("by_id/" + uuid + "/" + fnm[len(mnt):])
+       return siteFS.Open("by_id/" + uuid + fnm[len(mnt):])
 }