13931: un-uriquote keep: URIs when converting to (collection, path) tuple
[arvados.git] / sdk / cwl / arvados_cwl / fsaccess.py
index b249da77dd3e1ac4bda297fcab4fdbbd76d0e231..316a652529b384205661827e2c46d056025d5506 100644 (file)
@@ -1,8 +1,15 @@
+# Copyright (C) The Arvados Authors. All rights reserved.
+#
+# SPDX-License-Identifier: Apache-2.0
+
 import fnmatch
 import os
 import errno
 import urlparse
 import re
+import logging
+import threading
+from collections import OrderedDict
 
 import ruamel.yaml as yaml
 
@@ -13,26 +20,70 @@ import cwltool.resolver
 import arvados.util
 import arvados.collection
 import arvados.arvfile
+import arvados.errors
+
+from googleapiclient.errors import HttpError
 
 from schema_salad.ref_resolver import DefaultFetcher
 
+logger = logging.getLogger('arvados.cwl-runner')
+
+class CollectionCache(object):
+    def __init__(self, api_client, keep_client, num_retries,
+                 cap=256*1024*1024,
+                 min_entries=2):
+        self.api_client = api_client
+        self.keep_client = keep_client
+        self.num_retries = num_retries
+        self.collections = OrderedDict()
+        self.lock = threading.Lock()
+        self.total = 0
+        self.cap = cap
+        self.min_entries = min_entries
+
+    def cap_cache(self):
+        if self.total > self.cap:
+            # ordered list iterates from oldest to newest
+            for pdh, v in self.collections.items():
+                if self.total < self.cap or len(self.collections) < self.min_entries:
+                    break
+                # cut it loose
+                logger.debug("Evicting collection reader %s from cache", pdh)
+                del self.collections[pdh]
+                self.total -= v[1]
+
+    def get(self, pdh):
+        with self.lock:
+            if pdh not in self.collections:
+                logger.debug("Creating collection reader for %s", pdh)
+                cr = arvados.collection.CollectionReader(pdh, api_client=self.api_client,
+                                                         keep_client=self.keep_client,
+                                                         num_retries=self.num_retries)
+                sz = len(cr.manifest_text()) * 128
+                self.collections[pdh] = (cr, sz)
+                self.total += sz
+                self.cap_cache()
+            else:
+                cr, sz = self.collections[pdh]
+                # bump it to the back
+                del self.collections[pdh]
+                self.collections[pdh] = (cr, sz)
+            return cr
+
+
 class CollectionFsAccess(cwltool.stdfsaccess.StdFsAccess):
     """Implement the cwltool FsAccess interface for Arvados Collections."""
 
-    def __init__(self, basedir, api_client=None, keep_client=None):
+    def __init__(self, basedir, collection_cache=None):
         super(CollectionFsAccess, self).__init__(basedir)
-        self.api_client = api_client
-        self.keep_client = keep_client
-        self.collections = {}
+        self.collection_cache = collection_cache
 
     def get_collection(self, path):
-        p = path.split("/")
-        if p[0].startswith("keep:") and arvados.util.keep_locator_pattern.match(p[0][5:]):
-            pdh = p[0][5:]
-            if pdh not in self.collections:
-                self.collections[pdh] = arvados.collection.CollectionReader(pdh, api_client=self.api_client,
-                                                                            keep_client=self.keep_client)
-            return (self.collections[pdh], "/".join(p[1:]))
+        sp = path.split("/", 1)
+        p = sp[0]
+        if p.startswith("keep:") and arvados.util.keep_locator_pattern.match(p[5:]):
+            pdh = p[5:]
+            return (self.collection_cache.get(pdh), urlparse.unquote(sp[1]) if len(sp) == 2 else None)
         else:
             return (None, path)
 
@@ -60,28 +111,48 @@ class CollectionFsAccess(cwltool.stdfsaccess.StdFsAccess):
 
     def glob(self, pattern):
         collection, rest = self.get_collection(pattern)
-        if collection and not rest:
+        if collection is not None and not rest:
             return [pattern]
         patternsegments = rest.split("/")
-        return self._match(collection, patternsegments, "keep:" + collection.manifest_locator())
+        return sorted(self._match(collection, patternsegments, "keep:" + collection.manifest_locator()))
 
     def open(self, fn, mode):
         collection, rest = self.get_collection(fn)
-        if collection:
+        if collection is not None:
             return collection.open(rest, mode)
         else:
             return super(CollectionFsAccess, self).open(self._abs(fn), mode)
 
     def exists(self, fn):
-        collection, rest = self.get_collection(fn)
-        if collection:
-            return collection.exists(rest)
+        try:
+            collection, rest = self.get_collection(fn)
+        except HttpError as err:
+            if err.resp.status == 404:
+                return False
+            else:
+                raise
+        if collection is not None:
+            if rest:
+                return collection.exists(rest)
+            else:
+                return True
         else:
             return super(CollectionFsAccess, self).exists(fn)
 
+    def size(self, fn):  # type: (unicode) -> bool
+        collection, rest = self.get_collection(fn)
+        if collection is not None:
+            if rest:
+                arvfile = collection.find(rest)
+                if isinstance(arvfile, arvados.arvfile.ArvadosFile):
+                    return arvfile.size()
+            raise IOError(errno.EINVAL, "Not a path to a file %s" % (fn))
+        else:
+            return super(CollectionFsAccess, self).size(fn)
+
     def isfile(self, fn):  # type: (unicode) -> bool
         collection, rest = self.get_collection(fn)
-        if collection:
+        if collection is not None:
             if rest:
                 return isinstance(collection.find(rest), arvados.arvfile.ArvadosFile)
             else:
@@ -91,7 +162,7 @@ class CollectionFsAccess(cwltool.stdfsaccess.StdFsAccess):
 
     def isdir(self, fn):  # type: (unicode) -> bool
         collection, rest = self.get_collection(fn)
-        if collection:
+        if collection is not None:
             if rest:
                 return isinstance(collection.find(rest), arvados.collection.RichCollectionBase)
             else:
@@ -101,7 +172,7 @@ class CollectionFsAccess(cwltool.stdfsaccess.StdFsAccess):
 
     def listdir(self, fn):  # type: (unicode) -> List[unicode]
         collection, rest = self.get_collection(fn)
-        if collection:
+        if collection is not None:
             if rest:
                 dir = collection.find(rest)
             else:
@@ -123,33 +194,42 @@ class CollectionFsAccess(cwltool.stdfsaccess.StdFsAccess):
         if path.startswith("$(task.tmpdir)") or path.startswith("$(task.outdir)"):
             return path
         collection, rest = self.get_collection(path)
-        if collection:
+        if collection is not None:
             return path
         else:
             return os.path.realpath(path)
 
 class CollectionFetcher(DefaultFetcher):
-    def __init__(self, cache, session, api_client=None, keep_client=None):
+    def __init__(self, cache, session, api_client=None, fs_access=None, num_retries=4):
         super(CollectionFetcher, self).__init__(cache, session)
         self.api_client = api_client
-        self.fsaccess = CollectionFsAccess("", api_client=api_client, keep_client=keep_client)
+        self.fsaccess = fs_access
+        self.num_retries = num_retries
 
     def fetch_text(self, url):
         if url.startswith("keep:"):
             with self.fsaccess.open(url, "r") as f:
                 return f.read()
         if url.startswith("arvwf:"):
-            record = self.api_client.workflows().get(uuid=url[6:]).execute()
+            record = self.api_client.workflows().get(uuid=url[6:]).execute(num_retries=self.num_retries)
             definition = record["definition"] + ('\nlabel: "%s"\n' % record["name"].replace('"', '\\"'))
             return definition
         return super(CollectionFetcher, self).fetch_text(url)
 
     def check_exists(self, url):
-        if url.startswith("keep:"):
-            return self.fsaccess.exists(url)
-        if url.startswith("arvwf:"):
-            if self.fetch_text(url):
+        try:
+            if url.startswith("http://arvados.org/cwl"):
                 return True
+            if url.startswith("keep:"):
+                return self.fsaccess.exists(url)
+            if url.startswith("arvwf:"):
+                if self.fetch_text(url):
+                    return True
+        except arvados.errors.NotFoundError:
+            return False
+        except:
+            logger.exception("Got unexpected exception checking if file exists:")
+            return False
         return super(CollectionFetcher, self).check_exists(url)
 
     def urljoin(self, base_url, url):
@@ -188,12 +268,15 @@ class CollectionFetcher(DefaultFetcher):
 workflow_uuid_pattern = re.compile(r'[a-z0-9]{5}-7fd4e-[a-z0-9]{15}')
 pipeline_template_uuid_pattern = re.compile(r'[a-z0-9]{5}-p5p6p-[a-z0-9]{15}')
 
-def collectionResolver(api_client, document_loader, uri):
+def collectionResolver(api_client, document_loader, uri, num_retries=4):
+    if uri.startswith("keep:") or uri.startswith("arvwf:"):
+        return uri
+
     if workflow_uuid_pattern.match(uri):
         return "arvwf:%s#main" % (uri)
 
     if pipeline_template_uuid_pattern.match(uri):
-        pt = api_client.pipeline_templates().get(uuid=uri).execute()
+        pt = api_client.pipeline_templates().get(uuid=uri).execute(num_retries=num_retries)
         return "keep:" + pt["components"].values()[0]["script_parameters"]["cwl:tool"]
 
     p = uri.split("/")