Merge branch '11017-docker-migration'
[arvados.git] / sdk / cwl / arvados_cwl / pathmapper.py
index 0a11925ca74c4a45838a0e67b3c701ec8dcf1bd6..1f6aa577c180c65a9f8a31b0155c863f17e59b67 100644 (file)
@@ -1,13 +1,23 @@
 import re
+import logging
+import uuid
+import os
 
 import arvados.commands.run
 import arvados.collection
-from cwltool.pathmapper import PathMapper, MapperEnt, abspath
+
+from schema_salad.sourceline import SourceLine
+
+from cwltool.pathmapper import PathMapper, MapperEnt, abspath, adjustFileObjs, adjustDirObjs
+from cwltool.workflow import WorkflowException
+
+logger = logging.getLogger('arvados.cwl-runner')
 
 class ArvPathMapper(PathMapper):
     """Convert container-local paths to and from Keep collection ids."""
 
-    pdh_path = re.compile(r'^keep:[0-9a-f]{32}\+\d+/.+')
+    pdh_path = re.compile(r'^keep:[0-9a-f]{32}\+\d+/.+$')
+    pdh_dirpath = re.compile(r'^keep:[0-9a-f]{32}\+\d+(/.+)?$')
 
     def __init__(self, arvrunner, referenced_files, input_basedir,
                  collection_pattern, file_pattern, name=None, **kwargs):
@@ -18,43 +28,127 @@ class ArvPathMapper(PathMapper):
         self.name = name
         super(ArvPathMapper, self).__init__(referenced_files, input_basedir, None)
 
+    def visit(self, srcobj, uploadfiles):
+        src = srcobj["location"]
+        if srcobj["class"] == "File":
+            if "#" in src:
+                src = src[:src.index("#")]
+            if isinstance(src, basestring) and ArvPathMapper.pdh_path.match(src):
+                self._pathmap[src] = MapperEnt(src, self.collection_pattern % src[5:], "File")
+            if src not in self._pathmap:
+                # Local FS ref, may need to be uploaded or may be on keep
+                # mount.
+                ab = abspath(src, self.input_basedir)
+                st = arvados.commands.run.statfile("", ab, fnPattern="keep:%s/%s")
+                with SourceLine(srcobj, "location", WorkflowException):
+                    if isinstance(st, arvados.commands.run.UploadFile):
+                        uploadfiles.add((src, ab, st))
+                    elif isinstance(st, arvados.commands.run.ArvFile):
+                        self._pathmap[src] = MapperEnt(st.fn, self.collection_pattern % st.fn[5:], "File")
+                    elif src.startswith("_:"):
+                        if "contents" in srcobj:
+                            pass
+                        else:
+                            raise WorkflowException("File literal '%s' is missing contents" % src)
+                    elif src.startswith("arvwf:"):
+                        self._pathmap[src] = MapperEnt(src, src, "File")
+                    else:
+                        raise WorkflowException("Input file path '%s' is invalid" % st)
+            if "secondaryFiles" in srcobj:
+                for l in srcobj["secondaryFiles"]:
+                    self.visit(l, uploadfiles)
+        elif srcobj["class"] == "Directory":
+            if isinstance(src, basestring) and ArvPathMapper.pdh_dirpath.match(src):
+                self._pathmap[src] = MapperEnt(src, self.collection_pattern % src[5:], "Directory")
+            for l in srcobj.get("listing", []):
+                self.visit(l, uploadfiles)
+
+    def addentry(self, obj, c, path, subdirs):
+        if obj["location"] in self._pathmap:
+            src, srcpath = self.arvrunner.fs_access.get_collection(self._pathmap[obj["location"]].resolved)
+            if srcpath == "":
+                srcpath = "."
+            c.copy(srcpath, path + "/" + obj["basename"], source_collection=src, overwrite=True)
+            for l in obj.get("secondaryFiles", []):
+                self.addentry(l, c, path, subdirs)
+        elif obj["class"] == "Directory":
+            for l in obj["listing"]:
+                self.addentry(l, c, path + "/" + obj["basename"], subdirs)
+            subdirs.append((obj["location"], path + "/" + obj["basename"]))
+        elif obj["location"].startswith("_:") and "contents" in obj:
+            with c.open(path + "/" + obj["basename"], "w") as f:
+                f.write(obj["contents"].encode("utf-8"))
+        else:
+            raise SourceLine(obj, "location", WorkflowException).makeError("Don't know what to do with '%s'" % obj["location"])
+
     def setup(self, referenced_files, basedir):
         # type: (List[Any], unicode) -> None
-        self._pathmap = self.arvrunner.get_uploaded()
         uploadfiles = set()
 
+        already_uploaded = self.arvrunner.get_uploaded()
+        for k in referenced_files:
+            loc = k["location"]
+            if loc in already_uploaded:
+                v = already_uploaded[loc]
+                self._pathmap[loc] = MapperEnt(v.resolved, self.collection_pattern % v.resolved[5:], "File")
+
         for srcobj in referenced_files:
-            if srcobj["class"] == "File":
-                src = srcobj["location"]
-                if "#" in src:
-                    src = src[:src.index("#")]
-                if isinstance(src, basestring) and ArvPathMapper.pdh_path.match(src):
-                    self._pathmap[src] = MapperEnt(src, self.collection_pattern % src[5:], "File")
-                if src not in self._pathmap:
-                    # Local FS ref, may need to be uploaded or may be on keep
-                    # mount.
-                    ab = abspath(src, self.input_basedir)
-                    st = arvados.commands.run.statfile("", ab, fnPattern=self.file_pattern)
-                    if isinstance(st, arvados.commands.run.UploadFile):
-                        uploadfiles.add((src, ab, st))
-                    elif isinstance(st, arvados.commands.run.ArvFile):
-                        self._pathmap[src] = MapperEnt(ab, st.fn, "File")
-                    else:
-                        raise cwltool.workflow.WorkflowException("Input file path '%s' is invalid" % st)
+            self.visit(srcobj, uploadfiles)
 
         if uploadfiles:
             arvados.commands.run.uploadfiles([u[2] for u in uploadfiles],
                                              self.arvrunner.api,
                                              dry_run=False,
-                                             num_retries=3,
-                                             fnPattern=self.file_pattern,
+                                             num_retries=self.arvrunner.num_retries,
+                                             fnPattern="keep:%s/%s",
                                              name=self.name,
                                              project=self.arvrunner.project_uuid)
 
         for src, ab, st in uploadfiles:
-            self._pathmap[src] = MapperEnt(ab, st.fn, "File")
+            self._pathmap[src] = MapperEnt(st.fn, self.collection_pattern % st.fn[5:], "File")
             self.arvrunner.add_uploaded(src, self._pathmap[src])
 
+        for srcobj in referenced_files:
+            if srcobj["class"] == "Directory":
+                if srcobj["location"] not in self._pathmap:
+                    c = arvados.collection.Collection(api_client=self.arvrunner.api,
+                                                      keep_client=self.arvrunner.keep_client,
+                                                      num_retries=self.arvrunner.num_retries)
+                    subdirs = []
+                    for l in srcobj["listing"]:
+                        self.addentry(l, c, ".", subdirs)
+
+                    check = self.arvrunner.api.collections().list(filters=[["portable_data_hash", "=", c.portable_data_hash()]], limit=1).execute(num_retries=self.arvrunner.num_retries)
+                    if not check["items"]:
+                        c.save_new(owner_uuid=self.arvrunner.project_uuid)
+
+                    ab = self.collection_pattern % c.portable_data_hash()
+                    self._pathmap[srcobj["location"]] = MapperEnt(ab, ab, "Directory")
+                    for loc, sub in subdirs:
+                        ab = self.file_pattern % (c.portable_data_hash(), sub[2:])
+                        self._pathmap[loc] = MapperEnt(ab, ab, "Directory")
+            elif srcobj["class"] == "File" and (srcobj.get("secondaryFiles") or
+                (srcobj["location"].startswith("_:") and "contents" in srcobj)):
+
+                c = arvados.collection.Collection(api_client=self.arvrunner.api,
+                                                  keep_client=self.arvrunner.keep_client,
+                                                  num_retries=self.arvrunner.num_retries                                                  )
+                subdirs = []
+                self.addentry(srcobj, c, ".", subdirs)
+
+                check = self.arvrunner.api.collections().list(filters=[["portable_data_hash", "=", c.portable_data_hash()]], limit=1).execute(num_retries=self.arvrunner.num_retries)
+                if not check["items"]:
+                    c.save_new(owner_uuid=self.arvrunner.project_uuid)
+
+                ab = self.file_pattern % (c.portable_data_hash(), srcobj["basename"])
+                self._pathmap[srcobj["location"]] = MapperEnt(ab, ab, "File")
+                if srcobj.get("secondaryFiles"):
+                    ab = self.collection_pattern % c.portable_data_hash()
+                    self._pathmap["_:" + unicode(uuid.uuid4())] = MapperEnt(ab, ab, "Directory")
+                for loc, sub in subdirs:
+                    ab = self.file_pattern % (c.portable_data_hash(), sub[2:])
+                    self._pathmap[loc] = MapperEnt(ab, ab, "Directory")
+
         self.keepdir = None
 
     def reversemap(self, target):
@@ -64,3 +158,46 @@ class ArvPathMapper(PathMapper):
             return (target, "keep:" + target[len(self.keepdir)+1:])
         else:
             return super(ArvPathMapper, self).reversemap(target)
+
+class StagingPathMapper(PathMapper):
+    _follow_dirs = True
+
+    def visit(self, obj, stagedir, basedir, copy=False):
+        # type: (Dict[unicode, Any], unicode, unicode, bool) -> None
+        loc = obj["location"]
+        tgt = os.path.join(stagedir, obj["basename"])
+        if obj["class"] == "Directory":
+            self._pathmap[loc] = MapperEnt(loc, tgt, "Directory")
+            if loc.startswith("_:") or self._follow_dirs:
+                self.visitlisting(obj.get("listing", []), tgt, basedir)
+        elif obj["class"] == "File":
+            if loc in self._pathmap:
+                return
+            if "contents" in obj and loc.startswith("_:"):
+                self._pathmap[loc] = MapperEnt(obj["contents"], tgt, "CreateFile")
+            else:
+                if copy:
+                    self._pathmap[loc] = MapperEnt(loc, tgt, "WritableFile")
+                else:
+                    self._pathmap[loc] = MapperEnt(loc, tgt, "File")
+                self.visitlisting(obj.get("secondaryFiles", []), stagedir, basedir)
+
+
+class VwdPathMapper(StagingPathMapper):
+    def setup(self, referenced_files, basedir):
+        # type: (List[Any], unicode) -> None
+
+        # Go through each file and set the target to its own directory along
+        # with any secondary files.
+        self.visitlisting(referenced_files, self.stagedir, basedir)
+
+        for path, (ab, tgt, type) in self._pathmap.items():
+            if type in ("File", "Directory") and ab.startswith("keep:"):
+                self._pathmap[path] = MapperEnt("$(task.keep)/%s" % ab[5:], tgt, type)
+
+
+class NoFollowPathMapper(StagingPathMapper):
+    _follow_dirs = False
+    def setup(self, referenced_files, basedir):
+        # type: (List[Any], unicode) -> None
+        self.visitlisting(referenced_files, self.stagedir, basedir)