X-Git-Url: https://git.arvados.org/arvados.git/blobdiff_plain/5620a47e20e004188e0f1d0245ae9e916575e806..4c5a10ab54877413189c490dc5cd89dcc3d876ae:/sdk/python/arvados.py

diff --git a/sdk/python/arvados.py b/sdk/python/arvados.py
index f017b1e06e..7d99d582a6 100644
--- a/sdk/python/arvados.py
+++ b/sdk/python/arvados.py
@@ -13,6 +13,7 @@ import hashlib
 import string
 import bz2
 import zlib
+import fcntl
 
 from apiclient import errors
 from apiclient.discovery import build
@@ -42,9 +43,17 @@ class CredentialsFromEnv:
 url = ('https://%s/discovery/v1/apis/'
        '{api}/{apiVersion}/rest' % os.environ['ARVADOS_API_HOST'])
 credentials = CredentialsFromEnv()
-http = httplib2.Http()
+
+# Use system's CA certificates (if we find them) instead of httplib2's
+ca_certs = '/etc/ssl/certs/ca-certificates.crt'
+if not os.path.exists(ca_certs):
+    ca_certs = None             # use httplib2 default
+
+http = httplib2.Http(ca_certs=ca_certs)
 http = credentials.authorize(http)
-http.disable_ssl_certificate_validation=True
+if re.match(r'(?i)^(true|1|yes)$',
+            os.environ.get('ARVADOS_API_HOST_INSECURE', '')):
+    http.disable_ssl_certificate_validation=True
 service = build("arvados", "v1", http=http, discoveryServiceUrl=url)
 
 def task_set_output(self,s):
@@ -63,7 +72,7 @@ def current_task():
     t = service.job_tasks().get(uuid=os.environ['TASK_UUID']).execute()
     t = UserDict.UserDict(t)
     t.set_output = types.MethodType(task_set_output, t)
-    t.tmpdir = os.environ['TASK_TMPDIR']
+    t.tmpdir = os.environ['TASK_WORK']
     _current_task = t
     return t
 
@@ -74,7 +83,7 @@ def current_job():
         return _current_job
     t = service.jobs().get(uuid=os.environ['JOB_UUID']).execute()
     t = UserDict.UserDict(t)
-    t.tmpdir = os.environ['CRUNCH_WORK']
+    t.tmpdir = os.environ['JOB_WORK']
     _current_job = t
     return t
 
@@ -82,8 +91,8 @@ def api():
     return service
 
 class JobTask:
-    def __init__(self, parameters=dict(), resource_limits=dict()):
-        print "init jobtask %s %s" % (parameters, resource_limits)
+    def __init__(self, parameters=dict(), runtime_constraints=dict()):
+        print "init jobtask %s %s" % (parameters, runtime_constraints)
 
 class job_setup:
     @staticmethod
@@ -110,13 +119,39 @@ class job_setup:
                                        ).execute()
             exit(0)
 
+    @staticmethod
+    def one_task_per_input_stream(if_sequence=0, and_end_task=True):
+        if if_sequence != current_task()['sequence']:
+            return
+        job_input = current_job()['script_parameters']['input']
+        cr = CollectionReader(job_input)
+        for s in cr.all_streams():
+            task_input = s.tokens()
+            new_task_attrs = {
+                'job_uuid': current_job()['uuid'],
+                'created_by_job_task_uuid': current_task()['uuid'],
+                'sequence': if_sequence + 1,
+                'parameters': {
+                    'input':task_input
+                    }
+                }
+            service.job_tasks().create(job_task=json.dumps(new_task_attrs)).execute()
+        if and_end_task:
+            service.job_tasks().update(uuid=current_task()['uuid'],
+                                       job_task=json.dumps({'success':True})
+                                       ).execute()
+            exit(0)
+
 class util:
     @staticmethod
     def run_command(execargs, **kwargs):
+        if 'stdin' not in kwargs:
+            kwargs['stdin'] = subprocess.PIPE
+        if 'stdout' not in kwargs:
+            kwargs['stdout'] = subprocess.PIPE
+        if 'stderr' not in kwargs:
+            kwargs['stderr'] = subprocess.PIPE
         p = subprocess.Popen(execargs, close_fds=True, shell=False,
-                             stdin=subprocess.PIPE,
-                             stdout=subprocess.PIPE,
-                             stderr=subprocess.PIPE,
                              **kwargs)
         stdoutdata, stderrdata = p.communicate(None)
         if p.returncode != 0:
@@ -135,6 +170,262 @@ class util:
                          cwd=path)
         return path
 
+    @staticmethod
+    def tar_extractor(path, decompress_flag):
+        return subprocess.Popen(["tar",
+                                 "-C", path,
+                                 ("-x%sf" % decompress_flag),
+                                 "-"],
+                                stdout=None,
+                                stdin=subprocess.PIPE, stderr=sys.stderr,
+                                shell=False, close_fds=True)
+
+    @staticmethod
+    def tarball_extract(tarball, path):
+        """Retrieve a tarball from Keep and extract it to a local
+        directory.  Return the absolute path where the tarball was
+        extracted. If the top level of the tarball contained just one
+        file or directory, return the absolute path of that single
+        item.
+
+        tarball -- collection locator
+        path -- where to extract the tarball: absolute, or relative to job tmp
+        """
+        if not re.search('^/', path):
+            path = os.path.join(current_job().tmpdir, path)
+        lockfile = open(path + '.lock', 'w')
+        fcntl.flock(lockfile, fcntl.LOCK_EX)
+        try:
+            os.stat(path)
+        except OSError:
+            os.mkdir(path)
+        already_have_it = False
+        try:
+            if os.readlink(os.path.join(path, '.locator')) == tarball:
+                already_have_it = True
+        except OSError:
+            pass
+        if not already_have_it:
+
+            # emulate "rm -f" (i.e., if the file does not exist, we win)
+            try:
+                os.unlink(os.path.join(path, '.locator'))
+            except OSError:
+                if os.path.exists(os.path.join(path, '.locator')):
+                    os.unlink(os.path.join(path, '.locator'))
+
+            for f in CollectionReader(tarball).all_files():
+                if re.search('\.(tbz|tar.bz2)$', f.name()):
+                    p = util.tar_extractor(path, 'j')
+                elif re.search('\.(tgz|tar.gz)$', f.name()):
+                    p = util.tar_extractor(path, 'z')
+                elif re.search('\.tar$', f.name()):
+                    p = util.tar_extractor(path, '')
+                else:
+                    raise Exception("tarball_extract cannot handle filename %s"
+                                    % f.name())
+                while True:
+                    buf = f.read(2**20)
+                    if len(buf) == 0:
+                        break
+                    p.stdin.write(buf)
+                p.stdin.close()
+                p.wait()
+                if p.returncode != 0:
+                    lockfile.close()
+                    raise Exception("tar exited %d" % p.returncode)
+            os.symlink(tarball, os.path.join(path, '.locator'))
+        tld_extracts = filter(lambda f: f != '.locator', os.listdir(path))
+        lockfile.close()
+        if len(tld_extracts) == 1:
+            return os.path.join(path, tld_extracts[0])
+        return path
+
+    @staticmethod
+    def zipball_extract(zipball, path):
+        """Retrieve a zip archive from Keep and extract it to a local
+        directory.  Return the absolute path where the archive was
+        extracted. If the top level of the archive contained just one
+        file or directory, return the absolute path of that single
+        item.
+
+        zipball -- collection locator
+        path -- where to extract the archive: absolute, or relative to job tmp
+        """
+        if not re.search('^/', path):
+            path = os.path.join(current_job().tmpdir, path)
+        lockfile = open(path + '.lock', 'w')
+        fcntl.flock(lockfile, fcntl.LOCK_EX)
+        try:
+            os.stat(path)
+        except OSError:
+            os.mkdir(path)
+        already_have_it = False
+        try:
+            if os.readlink(os.path.join(path, '.locator')) == zipball:
+                already_have_it = True
+        except OSError:
+            pass
+        if not already_have_it:
+
+            # emulate "rm -f" (i.e., if the file does not exist, we win)
+            try:
+                os.unlink(os.path.join(path, '.locator'))
+            except OSError:
+                if os.path.exists(os.path.join(path, '.locator')):
+                    os.unlink(os.path.join(path, '.locator'))
+
+            for f in CollectionReader(zipball).all_files():
+                if not re.search('\.zip$', f.name()):
+                    raise Exception("zipball_extract cannot handle filename %s"
+                                    % f.name())
+                zip_filename = os.path.join(path, os.path.basename(f.name()))
+                zip_file = open(zip_filename, 'wb')
+                while True:
+                    buf = f.read(2**20)
+                    if len(buf) == 0:
+                        break
+                    zip_file.write(buf)
+                zip_file.close()
+                
+                p = subprocess.Popen(["unzip",
+                                      "-q", "-o",
+                                      "-d", path,
+                                      zip_filename],
+                                     stdout=None,
+                                     stdin=None, stderr=sys.stderr,
+                                     shell=False, close_fds=True)
+                p.wait()
+                if p.returncode != 0:
+                    lockfile.close()
+                    raise Exception("unzip exited %d" % p.returncode)
+                os.unlink(zip_filename)
+            os.symlink(zipball, os.path.join(path, '.locator'))
+        tld_extracts = filter(lambda f: f != '.locator', os.listdir(path))
+        lockfile.close()
+        if len(tld_extracts) == 1:
+            return os.path.join(path, tld_extracts[0])
+        return path
+
+    @staticmethod
+    def collection_extract(collection, path, files=[], decompress=True):
+        """Retrieve a collection from Keep and extract it to a local
+        directory.  Return the absolute path where the collection was
+        extracted.
+
+        collection -- collection locator
+        path -- where to extract: absolute, or relative to job tmp
+        """
+        if not re.search('^/', path):
+            path = os.path.join(current_job().tmpdir, path)
+        lockfile = open(path + '.lock', 'w')
+        fcntl.flock(lockfile, fcntl.LOCK_EX)
+        try:
+            os.stat(path)
+        except OSError:
+            os.mkdir(path)
+        already_have_it = False
+        try:
+            if os.readlink(os.path.join(path, '.locator')) == collection:
+                already_have_it = True
+        except OSError:
+            pass
+
+        # emulate "rm -f" (i.e., if the file does not exist, we win)
+        try:
+            os.unlink(os.path.join(path, '.locator'))
+        except OSError:
+            if os.path.exists(os.path.join(path, '.locator')):
+                os.unlink(os.path.join(path, '.locator'))
+
+        files_got = []
+        for s in CollectionReader(collection).all_streams():
+            stream_name = s.name()
+            for f in s.all_files():
+                if (files == [] or
+                    ((f.name() not in files_got) and
+                     (f.name() in files or
+                      (decompress and f.decompressed_name() in files)))):
+                    outname = f.decompressed_name() if decompress else f.name()
+                    files_got += [outname]
+                    if os.path.exists(os.path.join(path, stream_name, outname)):
+                        continue
+                    util.mkdir_dash_p(os.path.dirname(os.path.join(path, stream_name, outname)))
+                    outfile = open(os.path.join(path, stream_name, outname), 'wb')
+                    for buf in (f.readall_decompressed() if decompress
+                                else f.readall()):
+                        outfile.write(buf)
+                    outfile.close()
+        if len(files_got) < len(files):
+            raise Exception("Wanted files %s but only got %s from %s" % (files, files_got, map(lambda z: z.name(), list(CollectionReader(collection).all_files()))))
+        os.symlink(collection, os.path.join(path, '.locator'))
+
+        lockfile.close()
+        return path
+
+    @staticmethod
+    def mkdir_dash_p(path):
+        if not os.path.exists(path):
+            util.mkdir_dash_p(os.path.dirname(path))
+            try:
+                os.mkdir(path)
+            except OSError:
+                if not os.path.exists(path):
+                    os.mkdir(path)
+
+    @staticmethod
+    def stream_extract(stream, path, files=[], decompress=True):
+        """Retrieve a stream from Keep and extract it to a local
+        directory.  Return the absolute path where the stream was
+        extracted.
+
+        stream -- StreamReader object
+        path -- where to extract: absolute, or relative to job tmp
+        """
+        if not re.search('^/', path):
+            path = os.path.join(current_job().tmpdir, path)
+        lockfile = open(path + '.lock', 'w')
+        fcntl.flock(lockfile, fcntl.LOCK_EX)
+        try:
+            os.stat(path)
+        except OSError:
+            os.mkdir(path)
+
+        files_got = []
+        for f in stream.all_files():
+            if (files == [] or
+                ((f.name() not in files_got) and
+                 (f.name() in files or
+                  (decompress and f.decompressed_name() in files)))):
+                outname = f.decompressed_name() if decompress else f.name()
+                files_got += [outname]
+                if os.path.exists(os.path.join(path, outname)):
+                    os.unlink(os.path.join(path, outname))
+                util.mkdir_dash_p(os.path.dirname(os.path.join(path, outname)))
+                outfile = open(os.path.join(path, outname), 'wb')
+                for buf in (f.readall_decompressed() if decompress
+                            else f.readall()):
+                    outfile.write(buf)
+                outfile.close()
+        if len(files_got) < len(files):
+            raise Exception("Wanted files %s but only got %s from %s" %
+                            (files, files_got, map(lambda z: z.name(),
+                                                   list(stream.all_files()))))
+        lockfile.close()
+        return path
+
+    @staticmethod
+    def listdir_recursive(dirname, base=None):
+        allfiles = []
+        for ent in sorted(os.listdir(dirname)):
+            ent_path = os.path.join(dirname, ent)
+            ent_base = os.path.join(base, ent) if base else ent
+            if os.path.isdir(ent_path):
+                allfiles += util.listdir_recursive(ent_path, ent_base)
+            else:
+                allfiles += [ent_base]
+        return allfiles
+
 class DataReader:
     def __init__(self, data_locator):
         self.data_locator = data_locator
@@ -178,7 +469,7 @@ class StreamFileReader:
         data = self._stream.read(min(size, self._size - self._filepos))
         self._filepos += len(data)
         return data
-    def readall(self, size, **kwargs):
+    def readall(self, size=2**20, **kwargs):
         while True:
             data = self.read(size, **kwargs)
             if data == '':
@@ -196,14 +487,20 @@ class StreamFileReader:
             data = decompressor.decompress(decompressor.unconsumed_tail + chunk)
             if data and data != '':
                 yield data
-    def readlines(self, decompress=True):
+    def readall_decompressed(self, size=2**20):
         self._stream.seek(self._pos + self._filepos)
-        if decompress and re.search('\.bz2$', self._name):
-            datasource = self.bunzip2(2**10)
-        elif decompress and re.search('\.gz$', self._name):
-            datasource = self.gunzip(2**10)
+        if re.search('\.bz2$', self._name):
+            return self.bunzip2(size)
+        elif re.search('\.gz$', self._name):
+            return self.gunzip(size)
+        else:
+            return self.readall(size)
+    def readlines(self, decompress=True):
+        if decompress:
+            datasource = self.readall_decompressed()
         else:
-            datasource = self.readall(2**10)
+            self._stream.seek(self._pos + self._filepos)
+            datasource = self.readall()
         data = ''
         for newdata in datasource:
             data += newdata
@@ -246,6 +543,9 @@ class StreamReader:
                 self.files += [[int(pos), int(size), name]]
             else:
                 raise Exception("Invalid manifest format")
+
+    def tokens(self):
+        return self._tokens
     def tokens_for_range(self, range_start, range_size):
         resp = [self._stream_name]
         return_all_tokens = False
@@ -353,8 +653,9 @@ class CollectionReader:
             self._manifest_text = Keep.get(self._manifest_locator)
         self._streams = []
         for stream_line in self._manifest_text.split("\n"):
-            stream_tokens = stream_line.split()
-            self._streams += [stream_tokens]
+            if stream_line != '':
+                stream_tokens = stream_line.split()
+                self._streams += [stream_tokens]
     def all_streams(self):
         self._populate()
         resp = []
@@ -365,6 +666,9 @@ class CollectionReader:
         for s in self.all_streams():
             for f in s.all_files():
                 yield f
+    def manifest_text(self):
+        self._populate()
+        return self._manifest_text
 
 class CollectionWriter:
     KEEP_BLOCK_SIZE = 2**26
@@ -382,6 +686,31 @@ class CollectionWriter:
         pass
     def __exit__(self):
         self.finish()
+    def write_directory_tree(self,
+                             path, stream_name='.', max_manifest_depth=-1):
+        self.start_new_stream(stream_name)
+        todo = []
+        if max_manifest_depth == 0:
+            dirents = util.listdir_recursive(path)
+        else:
+            dirents = sorted(os.listdir(path))
+        for dirent in dirents:
+            target = os.path.join(path, dirent)
+            if os.path.isdir(target):
+                todo += [[target,
+                          os.path.join(stream_name, dirent),
+                          max_manifest_depth-1]]
+            else:
+                self.start_new_file(dirent)
+                with open(target, 'rb') as f:
+                    while True:
+                        buf = f.read(2**26)
+                        if len(buf) == 0:
+                            break
+                        self.write(buf)
+        self.finish_current_stream()
+        map(lambda x: self.write_directory_tree(*x), todo)
+
     def write(self, newdata):
         self._data_buffer += [newdata]
         self._data_buffer_len += len(newdata)
@@ -393,10 +722,12 @@ class CollectionWriter:
         if data_buffer != '':
             self._current_stream_locators += [Keep.put(data_buffer[0:self.KEEP_BLOCK_SIZE])]
             self._data_buffer = [data_buffer[self.KEEP_BLOCK_SIZE:]]
+            self._data_buffer_len = len(self._data_buffer[0])
     def start_new_file(self, newfilename=None):
         self.finish_current_file()
         self.set_current_file_name(newfilename)
     def set_current_file_name(self, newfilename):
+        newfilename = re.sub(r' ', '\\\\040', newfilename)
         if re.search(r'[ \t\n]', newfilename):
             raise AssertionError("Manifest filenames cannot contain whitespace")
         self._current_file_name = newfilename
@@ -411,7 +742,7 @@ class CollectionWriter:
                                        self._current_stream_length - self._current_file_pos,
                                        self._current_file_name]]
         self._current_file_pos = self._current_stream_length
-    def start_new_stream(self, newstreamname=None):
+    def start_new_stream(self, newstreamname='.'):
         self.finish_current_stream()
         self.set_current_stream_name(newstreamname)
     def set_current_stream_name(self, newstreamname):
@@ -443,6 +774,8 @@ class CollectionWriter:
         self.finish_current_stream()
         manifest = ''
         for stream in self._finished_streams:
+            if not re.search(r'^\.(/.*)?$', stream[0]):
+                manifest += './'
             manifest += stream[0]
             if len(stream[1]) == 0:
                 manifest += " d41d8cd98f00b204e9800998ecf8427e+0"