Added fakes3 integration with datasink and started adding a local_copy flag to the output generation logic

Author: pintohutch

nipype/interfaces/io.py

@@ -205,6 +205,11 @@ class DataSinkInputSpec(DynamicTraitedSpec, BaseInterfaceInputSpec):
                               'access')
     encrypt_bucket_keys = traits.Bool(desc='Flag indicating whether to use S3 '\
                                         'server-side AES-256 encryption')
+    # Set this if user wishes to override the bucket with their own
+    bucket = traits.Generic(mandatory=False,
+                            desc='Boto3 S3 bucket for manual override of bucket')
+    # Set this if user wishes to have local copy of files as well
+    local_dir = traits.Str(desc='Copy files locally as well as to S3 bucket')
 
     # Set call-able inputs attributes
     def __setattr__(self, key, value):
@@ -385,7 +390,6 @@ def _check_s3_base_dir(self):
 
         # Init variables
         s3_str = 's3://'
-        sep = os.path.sep
         base_directory = self.inputs.base_directory
 
         # Explicitly lower-case the "s3"
@@ -396,11 +400,16 @@ def _check_s3_base_dir(self):
 
         # Check if 's3://' in base dir
         if base_directory.startswith(s3_str):
+            # Attempt to access bucket
             try:
                 # Expects bucket name to be 's3://bucket_name/base_dir/..'
-                bucket_name = base_directory.split(s3_str)[1].split(sep)[0]
+                bucket_name = base_directory.split(s3_str)[1].split('/')[0]
                 # Get the actual bucket object
-                self.bucket = self._fetch_bucket(bucket_name)
+                if self.inputs.bucket:
+                    self.bucket = self.inputs.bucket
+                else:
+                    self.bucket = self._fetch_bucket(bucket_name)
+            # Report error in case of exception
             except Exception as exc:
                 err_msg = 'Unable to access S3 bucket. Error:\n%s. Exiting...'\
                           % exc
@@ -566,7 +575,7 @@ def _upload_to_s3(self, src, dst):
         bucket = self.bucket
         iflogger = logging.getLogger('interface')
         s3_str = 's3://'
-        s3_prefix = os.path.join(s3_str, bucket.name)
+        s3_prefix = s3_str + bucket.name
 
         # Explicitly lower-case the "s3"
         if dst.lower().startswith(s3_str):
@@ -629,41 +638,53 @@ def _list_outputs(self):
         iflogger = logging.getLogger('interface')
         outputs = self.output_spec().get()
         out_files = []
-        outdir = self.inputs.base_directory
+        # Use hardlink
         use_hardlink = str2bool(config.get('execution', 'try_hard_link_datasink'))
 
-        # If base directory isn't given, assume current directory
-        if not isdefined(outdir):
-            outdir = '.'
+        # Set local output directory if specified
+        if isdefined(self.inputs.local_copy):
+            outdir = self.inputs.local_copy
+        else:
+            outdir = self.inputs.base_directory
+            # If base directory isn't given, assume current directory
+            if not isdefined(outdir):
+                outdir = '.'
 
-        # Check if base directory reflects S3-bucket upload
+        # Check if base directory reflects S3 bucket upload
         try:
             s3_flag = self._check_s3_base_dir()
+            s3dir = self.inputs.base_directory
+            if isdefined(self.inputs.container):
+                s3dir = os.path.join(s3dir, self.inputs.container)
         # If encountering an exception during bucket access, set output
         # base directory to a local folder
         except Exception as exc:
-            local_out_exception = os.path.join(os.path.expanduser('~'),
-                                               'data_output')
+            if not isdefined(self.inputs.local_copy):
+                local_out_exception = os.path.join(os.path.expanduser('~'),
+                                                   's3_datasink_' + self.bucket.name)
+                outdir = local_out_exception
+            else:
+                outdir = self.inputs.local_copy
+            # Log local copying directory
             iflogger.info('Access to S3 failed! Storing outputs locally at: '\
-                          '%s\nError: %s' %(local_out_exception, exc))
-            self.inputs.base_directory = local_out_exception
-
-        # If not accessing S3, just set outdir to local absolute path 
-        if not s3_flag:
-            outdir = os.path.abspath(outdir)
+                          '%s\nError: %s' %(outdir, exc))
 
         # If container input is given, append that to outdir
         if isdefined(self.inputs.container):
             outdir = os.path.join(outdir, self.inputs.container)
-        # Create the directory if it doesn't exist
-        if not os.path.exists(outdir):
-            try:
-                os.makedirs(outdir)
-            except OSError, inst:
-                if 'File exists' in inst:
-                    pass
-                else:
-                    raise(inst)
+
+        # If doing a localy output
+        if not outdir.lower().startswith('s3://'):
+            outdir = os.path.abspath(outdir)
+            # Create the directory if it doesn't exist
+            if not os.path.exists(outdir):
+                try:
+                    os.makedirs(outdir)
+                except OSError, inst:
+                    if 'File exists' in inst:
+                        pass
+                    else:
+                        raise(inst)
 
         # Iterate through outputs attributes {key : path(s)}
         for key, files in self.inputs._outputs.items():
@@ -672,10 +693,14 @@ def _list_outputs(self):
             iflogger.debug("key: %s files: %s" % (key, str(files)))
             files = filename_to_list(files)
             tempoutdir = outdir
+            if s3_flag:
+                s3tempoutdir = s3dir
             for d in key.split('.'):
                 if d[0] == '@':
                     continue
                 tempoutdir = os.path.join(tempoutdir, d)
+                if s3_flag:
+                    s3tempoutdir = os.path.join(s3tempoutdir, d)
 
             # flattening list
             if isinstance(files, list):
@@ -690,25 +715,26 @@ def _list_outputs(self):
                     src = os.path.join(src, '')
                 dst = self._get_dst(src)
                 dst = os.path.join(tempoutdir, dst)
+                s3dst = os.path.join(s3tempoutdir, dst)
                 dst = self._substitute(dst)
                 path, _ = os.path.split(dst)
 
-                # Create output directory if it doesnt exist
-                if not os.path.exists(path):
-                    try:
-                        os.makedirs(path)
-                    except OSError, inst:
-                        if 'File exists' in inst:
-                            pass
-                        else:
-                            raise(inst)
-
                 # If we're uploading to S3
                 if s3_flag:
+                    dst = dst.replace(outdir, self.inputs.base_directory)
                     self._upload_to_s3(src, dst)
                     out_files.append(dst)
                 # Otherwise, copy locally src -> dst
                 else:
+                    # Create output directory if it doesnt exist
+                    if not os.path.exists(path):
+                        try:
+                            os.makedirs(path)
+                        except OSError, inst:
+                            if 'File exists' in inst:
+                                pass
+                            else:
+                                raise(inst)
                     # If src is a file, copy it to dst
                     if os.path.isfile(src):
                         iflogger.debug('copyfile: %s %s' % (src, dst))

nipype/interfaces/tests/test_io.py

@@ -13,13 +13,21 @@
 import nipype.interfaces.io as nio
 from nipype.interfaces.base import Undefined
 
+# Check for boto
 noboto = False
 try:
     import boto
     from boto.s3.connection import S3Connection, OrdinaryCallingFormat
 except:
     noboto = True
 
+# Check for boto3
+noboto3 = False
+try:
+    import boto3
+    from botocore.utils import fix_s3_host
+except:
+    noboto3 = True
 
 def test_datagrabber():
     dg = nio.DataGrabber()
@@ -155,6 +163,102 @@ def test_datasink():
     ds = nio.DataSink(infields=['test'])
     yield assert_true, 'test' in ds.inputs.copyable_trait_names()
 
+# Function to check for fakes3
+def _check_for_fakes3():
+    '''
+    Function used internally to check for fakes3 installation
+    '''
+
+    # Import packages
+    import subprocess
+
+    # Init variables
+    fakes3_found = False
+
+    # Check for fakes3
+    try:
+        ret_code = subprocess.check_call(['which', 'fakes3'])
+        if ret_code == 0:
+            fakes3_found = True
+    except subprocess.CalledProcessError as exc:
+        print 'fakes3 not found, install via \'gem install fakes3\', skipping test...'
+    except:
+        print 'Unable to check for fakes3 installation, skipping test...'
+
+    # Return if found
+    return fakes3_found
+
+@skipif(noboto3)
+# Test datasink writes to s3 properly
+def test_datasink_to_s3():
+    '''
+    This function tests to see if the S3 functionality of a DataSink
+    works properly
+    '''
+
+    # Import packages
+    import hashlib
+    import tempfile
+
+    # Init variables
+    ds = nio.DataSink()
+    bucket_name = 'test'
+    container = 'outputs'
+    attr_folder = 'text_file'
+    output_dir = 's3://' + bucket_name
+    # Local temporary filepaths for testing
+    fakes3_dir = tempfile.mkdtemp()
+    input_dir = tempfile.mkdtemp()
+    input_path = os.path.join(input_dir, 'datasink_test_s3.txt')
+
+    # Check for fakes3
+    fakes3_found = _check_for_fakes3()
+    if not fakes3_found:
+        return
+
+    # Start up fake-S3 server
+    proc = Popen(['fakes3', '-r', fakes3_dir, '-p', '4567'], stdout=open(os.devnull, 'wb'))
+
+    # Init boto3 s3 resource to talk with fakes3
+    resource = boto3.resource(aws_access_key_id='mykey',
+                              aws_secret_access_key='mysecret',
+                              service_name='s3',
+                              endpoint_url='http://localhost:4567',
+                              use_ssl=False)
+    resource.meta.client.meta.events.unregister('before-sign.s3', fix_s3_host)
+
+    # Create bucket
+    bucket = resource.create_bucket(Bucket=bucket_name)
+
+    # Create input file
+    with open(input_path, 'wb') as f:
+        f.write('ABCD1234')
+
+    # Prep datasink
+    ds.inputs.base_directory = output_dir
+    ds.inputs.container = container
+    ds.inputs.bucket = bucket
+    setattr(ds.inputs, attr_folder, input_path)
+
+    # Run datasink
+    ds.run()
+
+    # Get MD5sums and compare
+    key = '/'.join([container, attr_folder, os.path.basename(input_path)])
+    obj = bucket.Object(key=key)
+    dst_md5 = obj.e_tag.replace('"', '')
+    src_md5 = hashlib.md5(open(input_path, 'rb').read()).hexdigest()
+
+    # Make sure md5sums match
+    yield assert_equal, src_md5, dst_md5
+
+    # Kill fakes3
+    proc.kill()
+
+    # Delete fakes3 folder and input file
+    shutil.rmtree(fakes3_dir)
+    shutil.rmtree(input_dir)
+
 @skipif(noboto)
 def test_s3datasink():
     ds = nio.S3DataSink()