labdata 0.0.3__py3-none-any.whl

labdata/rules/utils.py

@@ -0,0 +1,290 @@
+from ..utils import *
+from ..s3 import copy_to_s3
+
+# has utilities needed by other rules
+def _checksum_files(filepath, local_path):
+    '''
+    Checksum files that don't need to be copied
+    '''
+    # construct the path:
+    src = Path(local_path)/filepath
+    hash = compute_md5_hash(src)  # computes the hash
+    srcstat = src.stat()
+    file_size = srcstat.st_size
+    return dict(src_path = filepath,
+                src_md5 = hash,
+                src_size = file_size,
+                src_datetime = datetime.fromtimestamp(srcstat.st_ctime))
+
+class UploadRule():
+    def __init__(self,job_id,prefs = None):
+        '''
+
+Rule to apply on upload. 
+
+        1) Checksum on the files; compare with provided (reserve job if use_db)
+        2) Apply function
+        3) Checksum on the output - the files that changed
+        4) Submit upload
+        5) Update tables 
+
+Can submit job on slurm, some of these can be long or take resources.
+
+        '''
+        self.rule_name = 'default'
+        
+        self.job_id = job_id
+        self.src_paths = None
+        self.processed_paths = None
+        self.dst_paths = None
+        self.inserted_files = []
+        if prefs is None:
+            from ..utils import prefs
+        self.prefs = prefs
+        self.local_path = self.prefs['local_paths'][0]
+        self.dataset_key = None # will get written on upload, use in _post_upload
+        self.max_concurrent = -1  # maximum number of concurrent UploadJobs that can run [-1 is infinite].
+
+    def apply(self):
+        # parse inputs
+        from ..schema import UploadJob, File, dj
+
+        if not self.job_id is None:
+            with dj.conn().transaction:
+                self.jobquery = (UploadJob() & dict(job_id = self.job_id))
+                job_status = self.jobquery.fetch(as_dict = True)
+                if len(job_status):
+                    if job_status[0]['job_waiting']:
+                        # check here if there are other jobs running for the same rule on this host
+                        if self.max_concurrent > 0:
+                            kk = dict(job_host = self.prefs['hostname'],
+                                      job_waiting = 0,
+                                      job_status = 'WORKING',
+                                      job_rule = None)
+                            if not self.rule_name == 'default':
+                                kk['job_rule'] = self.rule_name
+                            number_of_running_jobs = len(UploadJob() & kk)
+                            if number_of_running_jobs >= self.max_concurrent:
+                                print(f"Job {self.job_id} can not run because there are {number_of_running_jobs} for the same rule.")
+                                self.set_job_status(job_status = 'WAITING', job_waiting = 1, job_log = f'Waiting for {number_of_running_jobs}.')
+                                return
+                        self.set_job_status(job_status = 'WORKING', job_waiting = 0) # take the job
+                    else:
+                        print(f"Job {self.job_id} is already taken.")
+                        #print(job_status, flush = True)
+                        return # exit.
+                else:
+                    raise ValueError(f'job_id {self.job_id} does not exist.')
+        # get the paths
+        self.src_paths = pd.DataFrame((UploadJob.AssignedFiles() & dict(job_id = self.job_id)).fetch())
+        if not len(self.src_paths):
+            self.set_job_status(job_status = 'FAILED',
+                                job_log = f'Could not find files for {self.job_id} in Upload.AssignedFiles.')
+            raise ValueError(f'Could not find files for {self.job_id} in Upload.AssignedFiles.')
+        self.upload_storage = self.jobquery.fetch('upload_storage')[0]
+        if self.upload_storage is None:
+            if 'upload_storage' in self.prefs.keys():
+                self.upload_storage = self.prefs['upload_storage']
+        
+        # this should not fail because we have to keep track of errors, should update the table
+        src = [Path(self.local_path) / p for p in self.src_paths.src_path.values]
+        try:
+            comparison = compare_md5s(src,self.src_paths.src_md5.values)
+        except FileNotFoundError:
+            print('File not found for {0}'.format(Path(self.src_paths.src_path.iloc[0]).parent))
+            self.set_job_status(job_status = 'FAILED', job_log = 'FILE NOT FOUND; check file transfer.')
+            return
+        
+        if not comparison:
+            print('CHECKSUM FAILED for {0}'.format(Path(self.src_paths.src_path.iloc[0]).parent))
+            self.set_job_status(job_status = 'FAILED',job_log = 'MD5 CHECKSUM failed; check file transfer.')
+            return # exit.
+        import traceback
+        try:
+            paths = self._apply_rule() # can use the src_paths
+            self._upload()             # compare the hashes after
+        except Exception as err:
+            # log the error
+            print('There was an error processing or uploading this dataset.')
+            print(err)
+            self.set_job_status(job_status = 'FAILED',job_log = f'{traceback.print_exc()}')
+            return
+        try:
+            self._post_upload()        # so the rules can insert tables and all.
+        except Exception as err:
+            # log the error
+            print('There was an error with the post-upload this dataset.')
+            print(err)
+            self.set_job_status(job_status = 'FAILED',job_log = f'POST-UPLOAD: {traceback.print_exc()}')
+            return 
+        return paths # so apply can output paths.
+    
+    def set_job_status(self, job_status = 'FAILED',job_log = '',job_waiting = 0):
+        from ..schema import UploadJob
+        if not self.job_id is None:
+            kk = dict(job_id = self.job_id,
+                      job_waiting = job_waiting,
+                      job_status = job_status,
+                      job_log = job_log)
+            if job_status == 'FAILED':
+                print(f'Check job_id {self.job_id} : {job_status}')
+                kk['job_host'] = self.prefs['hostname'], # write the hostname so we know where it failed.
+            UploadJob.update1(kk)
+            
+
+    def _handle_processed_and_src_paths(self, processed_files,new_files):
+        '''
+        Put the files in the proper place and compute checksums for new files.
+        Call this from the apply method.
+        '''
+        n_jobs = DEFAULT_N_JOBS
+        self.processed_paths = []
+        for f in processed_files:
+            i = np.where(self.src_paths.src_path == f)[0][0]
+            self.processed_paths.append(self.src_paths.iloc[i])
+            self.src_paths.drop(self.src_paths.iloc[i].name,axis = 0,inplace = True)
+            self.src_paths.reset_index(drop=True,inplace = True)
+        self.processed_paths = pd.DataFrame(self.processed_paths).reset_index(drop=True)        
+
+        res = Parallel(n_jobs = n_jobs)(delayed(_checksum_files)(
+            path,
+            local_path = self.local_path) for path in new_files)
+        for r in res:
+            r['job_id'] = self.job_id
+        self.src_paths = pd.concat([self.src_paths,pd.DataFrame(res)], ignore_index=True)
+        # drop duplicate paths in case there are any
+        self.src_paths = self.src_paths.drop_duplicates(subset=['src_path'], keep='last')
+        
+    def _post_upload(self):
+        return
+    
+    def _upload(self):
+        # this reads the attributes and uploads
+        # It also puts the files in the Tables
+        
+        # destination in the bucket is actually the path
+        dst = [k for k in self.src_paths.src_path.values]
+        # source is the place where data are
+        src = [Path(self.local_path) / p for p in self.src_paths.src_path.values] # same as md5
+        # s3 copy in parallel hashes were compared before so no need to do it now.
+        self.set_job_status(job_status = 'WORKING',
+                            job_log = datetime.now().strftime(f'Uploading {len(src)} files %Y %m %d %H:%M:%S'),
+                            job_waiting = 0)
+        copy_to_s3(src, dst, md5_checksum=None,storage_name=self.upload_storage)
+        from ..schema import UploadJob, File, dj, ProcessedFile, Dataset
+        self.set_job_status(job_status = 'WORKING',
+                            job_log = datetime.now().strftime('%Y %m %d %H:%M:%S'),
+                            job_waiting = 0)
+        with dj.conn().transaction:  # make it all update at the same time
+            # insert to Files so we know where to get the data
+            insfiles = []
+            for i,f in self.src_paths.iterrows():
+                insfiles.append(dict(file_path = f.src_path,
+                                     storage = self.upload_storage,
+                                     file_datetime = f.src_datetime,
+                                     file_size = f.src_size,
+                                     file_md5 = f.src_md5))
+            if len(insfiles):
+                File.insert(insfiles)
+            # Add to dataset?
+            job = self.jobquery.fetch(as_dict=True)[0]
+            # check if it has a dataset
+            if all([not job[a] is None for a in ['subject_name','session_name','dataset_name']]):
+                for i,p in enumerate(insfiles):
+                    insfiles[i] = dict(subject_name = job['subject_name'],
+                                       session_name = job['session_name'],
+                                       dataset_name = job['dataset_name'],
+                                       file_path = p['file_path'],
+                                       storage = self.upload_storage)
+                if len(insfiles):
+                    Dataset.DataFiles.insert(insfiles)
+                self.dataset_key = dict(subject_name = job['subject_name'],
+                                        session_name = job['session_name'],
+                                        dataset_name = job['dataset_name'])
+            self.inserted_files += insfiles
+            # Insert the processed files so the deletions are safe
+            if not self.processed_paths is None:
+                ins = []
+                for i,f in self.processed_paths.iterrows():
+                    ins.append(dict(file_path = f.src_path,
+                                    file_datetime = f.src_datetime,
+                                    file_size = f.src_size,
+                                    file_md5 = f.src_md5))
+                if len(ins):
+                    ProcessedFile.insert(ins)
+            if len(self.inserted_files): # keep the job in the queue 
+                # completed
+                self.set_job_status(job_status = 'COMPLETED',
+                                    job_log = datetime.now().strftime('UPLOADED %Y %m %d %H:%M:%S'),
+                                    job_waiting = 0)
+            else:
+                self.set_job_status(job_status = 'FAILED',
+                                    job_log = datetime.now().strftime('No files inserted %Y %m %d %H:%M:%S'),
+                                    job_waiting = 0)
+            #(UploadJob & dict(job_id = self.job_id)).delete(safemode = False) # only delete if on completed
+        
+    def _apply_rule(self):
+        # this rule does nothing, so the src_paths are going to be empty,
+        # and the "paths" are going to be the src_paths
+        self.processed_paths = None # processed paths are just the same, no file changed, so no need to do anything.
+        # needs to compute the checksum on all the new files
+        return self.src_paths.src_path.values
+    
+class ReplaceRule(UploadRule):
+    '''
+    Check if files exist in storage
+    Upload new versions of the files
+    Update MD5 checksums in database
+    '''
+
+    def __init__(self, job_id, prefs = None):
+        super(ReplaceRule,self).__init__(job_id = job_id, prefs = prefs)
+        self.rule_name = 'replace'
+
+    def _apply_rule(self):
+        # can only replace files that were uploaded, the files that are not there will be added to the dataset.
+        from ..s3 import s3_delete_file, copyfile_to_s3
+        from ..schema import File
+        self.processed_paths = None # processed paths are just the same, no file changed, so no need to do anything.
+
+        same_files = []
+        replace_files = []
+        new_files = []
+        for i,f in self.src_paths.iterrows():
+            existing = (File() & dict(file_path = f.src_path,
+                          storage = self.upload_storage)).fetch(as_dict = True)
+            if len(existing):
+                existing = existing[0] # there is only one file
+                if existing['file_md5'] == f.src_md5:
+                    same_files.append(dict(f))
+                else:
+                    replace_files.append(dict(f))
+            else:
+                new_files.append(dict(f))
+        txt = f'There are {len(new_files)} new files, {len(replace_files)} files to be replaced and {len(same_files)} unmodified files.'
+        print(txt,flush = True)
+        self.set_job_status(job_status = 'WORKING',
+                            job_log = txt,
+                            job_waiting = 0)
+        ##################### replace files happens here #########################
+         # destination in the bucket is actually the path
+        dst = [k['src_path'] for k in replace_files]
+        # source is the place where data are (this only works for a UNIX server at the moment)
+        src = [Path(self.local_path) / p['src_path'] for p in replace_files] # same as md5
+        # s3 copy in parallel hashes were compared before so no need to do it now.
+        copy_to_s3(src, dst, md5_checksum=None, storage_name=self.upload_storage)
+        for f in replace_files:
+            File.update1(dict(file_path = f['src_path'],
+                              storage = self.upload_storage,
+                              file_datetime = f['src_datetime'],
+                              file_size = f['src_size'],
+                              file_md5 = f['src_md5']))
+            print(f'Updated {f["src_path"]}')
+        
+        # replace the src_paths with the new files otherwise delete all rows
+        if len(new_files):
+            self.src_paths = pd.DataFrame(new_files)
+        else:
+            self.src_paths.drop(self.src_paths.index, inplace = True)
+        
+        return new_files+replace_files

labdata/s3.py

@@ -0,0 +1,317 @@
+# Functions for interacting with S3 storage.
+# This file provides utilities for uploading, downloading and managing files on S3 storage.
+# Key functionality includes:
+#   - Validating S3 storage configurations 
+#   - Copying files to/from S3 buckets
+#   - Deleting files from S3
+#   - Parallel file transfer operations
+#
+# labdata 2024
+
+from .utils import *
+
+__all__ = ['validate_storage',
+           'copyfile_to_s3',
+           'copyfile_from_s3',
+           'copy_to_s3',
+           'copy_from_s3',
+           's3_delete_file']
+
+def validate_storage(storage):
+    '''
+    storage = validate_storage(storage)
+
+    Checks that there is an s3 access_key and secret_key in the storage dictionary.
+    Updates the labdata preference file.
+
+    Joao Couto - labdata 2024
+
+    '''
+    # lets find the storage in the preferences
+    if 'storage' in prefs.keys():
+        for key in prefs['storage'].keys():
+            if prefs['storage'][key] == storage:
+                save_prefs = False
+                storage = prefs['storage'][key]
+    if not 'protocol' in storage.keys():
+        # assume S3
+        storage['protocol'] = 's3'
+        if 'save_prefs' in dir():
+            save_prefs = True
+    if storage['protocol'] == 's3':
+        # then try to find an access key
+        for k in ['access_key','secret_key']:
+            if not k in storage.keys():
+                storage[k] = None
+            if storage[k] is None:
+                import getpass
+                # get the password and write to file
+                storage[k] = getpass.getpass(prompt=f'S3 {k}:')
+                if 'save_prefs' in dir():
+                    save_prefs = True
+    if 'save_prefs' in dir():
+        if save_prefs:
+            save_labdata_preferences(prefs, LABDATA_FILE)                
+    return storage
+
+
+def copyfile_to_s3(source_file,
+                   destination_file,
+                   storage,
+                   md5_checksum = None):
+    '''
+    Copy a single file to S3 and do a checksum comparisson.
+
+    Joao Couto - 2024
+    '''
+    from minio import Minio
+    client = Minio(endpoint = storage['endpoint'],
+                   access_key = storage['access_key'],
+                   secret_key = storage['secret_key'])
+
+    if 'folder' in storage.keys():
+        if len(storage['folder']):
+            destination_file = storage['folder'] + '/' + destination_file
+
+    if not md5_checksum is None:
+        if not md5_checksum == compute_md5_hash(source_file):
+            raise OSError(f'Checksum {md5_checksum} does not match {source_file}.')
+    res = client.fput_object(
+        storage['bucket'], destination_file, source_file)
+    return res
+
+def copy_to_s3(source_files, destination_files,
+               storage = None,
+               storage_name = None,
+               md5_checksum = None,
+               n_jobs = DEFAULT_N_JOBS):
+    '''
+    Copy S3 and do a checksum comparisson.
+    Copy occurs in parallel for multiple files.
+
+    Joao Couto - 2024
+    '''
+    if storage is None:
+        if storage_name is None:
+            raise ValueError("Specify a storage to copy to - either pass the storage dictionary or specify a name from the prefs.")
+        storage = prefs['storage'][storage_name] # link to preferences storage from storage_name
+    storage = validate_storage(storage) # validate and update keys
+
+    if not type(source_files) is list: # check type of source
+        raise ValueError('source_files has to be a list of paths')
+    
+    if not type(destination_files) is list:  # check type of destination
+        raise ValueError('destination_files has to be a list of paths')
+    # Check if the source and the destination are the correct sizes
+    assert len(source_files) == len(destination_files),ValueError('source and destination are the wrong size')
+    
+    if md5_checksum is None:
+        md5_checksum = [None]*len(source_files)
+    from tqdm import tqdm
+    n_jobs = validate_num_jobs_joblib(n_jobs) # check if running inside a joblib worker.
+    res = Parallel(n_jobs = n_jobs)(delayed(copyfile_to_s3)(src,
+                                                            dst,
+                                                            storage = storage,
+                                                            md5_checksum = md5)
+                                    for src,dst,md5 in tqdm(zip(source_files,destination_files,md5_checksum),
+                                                            desc = f'Pushing to S3 [{storage["bucket"]}]'))
+    return res
+
+def boto3_copy_from_s3(src,dst,store):
+    import boto3
+    Path(dst).parent.mkdir(exist_ok=True,parents = True)
+    s3 = boto3.resource('s3',aws_access_key_id = store['access_key'],
+                    aws_secret_access_key = store['secret_key'])
+    obj = s3.Object(bucket_name = store['bucket'],
+                    key = str(src))
+    obj.download_file(str(dst))
+    return True
+
+def copyfile_from_s3(source_file,
+                     destination_file,
+                     storage,
+                     md5_checksum = None,
+                     engine = 'boto3'):
+    '''
+    Copy a single file to S3 and do a checksum comparisson.
+    
+    This function copies a single file from an S3 bucket to local storage. 
+    Supports both boto3 and minio engines for S3 operations and handles creating destination directories.
+    The function can optionally verify file integrity using an MD5 checksum [NOT IMPLEMENTED]. 
+    
+    Note: If 'use_awscli' is set to TRUE in the preferences, it will use AWS CLI to copy from S3
+
+    Parameters
+    ----------
+    source_file : str
+        Path to file in S3 bucket to copy
+    destination_file : str
+        Local path to copy file to
+    storage : dict
+        Dictionary containing S3 storage configuration with:
+        - endpoint: S3 endpoint URL 
+        - access_key: AWS access key
+        - secret_key: AWS secret key
+        - bucket: S3 bucket name
+    md5_checksum : str, optional
+        MD5 hash to verify file integrity
+    engine : str, default 'boto3'
+        Engine to use for S3 operations ('boto3' or 'minio')
+
+    Returns
+    -------
+    bool or MinioResponse
+        True if copy successful with boto3
+        MinioResponse object if using minio engine
+
+    Joao Couto - 2024
+    '''
+    if 'use_awscli' in prefs.keys():  # to copy from inside ec2 
+        if prefs['use_awscli']:
+            cmd = f'AWS_ACCESS_KEY_ID={storage["access_key"]} AWS_SECRET_ACCESS_KEY={storage["secret_key"]} aws s3 cp s3://{storage["bucket"]}/{source_file} {destination_file}'
+            os.system(cmd)
+            return True
+    if engine == 'boto3':
+        res = boto3_copy_from_s3(source_file,destination_file,storage)
+    else:
+        from minio import Minio
+        client = Minio(endpoint = storage['endpoint'],
+                       access_key = storage['access_key'],
+                       secret_key = storage['secret_key'])
+
+        if 'folder' in storage.keys():
+            if len(storage['folder']):
+                destination_file = storage['folder'] + '/' + destination_file
+
+        res = client.fget_object(
+            storage['bucket'], source_file, destination_file)
+    return res
+
+def copy_from_s3(source_files, destination_files,
+                 storage = None,
+                 storage_name = None,
+                 n_jobs = DEFAULT_N_JOBS,
+                 engine = 'boto3'):
+    '''
+    copy_from_s3(source_files, destination_files,
+                 storage = None,
+                 storage_name = None,
+                 n_jobs = DEFAULT_N_JOBS,
+                 engine = 'boto3')
+
+    This function handles copying multiple files in parallel from an S3 bucket to local storage.
+    It supports both boto3 and minio engines for S3 operations. Files can be copied using
+    credentials directly provided in a storage dict or from named storage configurations in
+    preferences.
+
+    Parameters
+    ----------
+    source_files : list
+        List of file paths to copy from S3
+    destination_files : list 
+        List of local file paths to copy to
+    storage : dict, optional
+        Dictionary containing S3 storage configuration. Must include:
+        - endpoint: S3 endpoint URL
+        - access_key: AWS access key
+        - secret_key: AWS secret key
+        - bucket: S3 bucket name
+    storage_name : str, optional
+        Name of storage configuration to use from preferences. Required if storage not provided.
+    n_jobs : int, default DEFAULT_N_JOBS
+        Number of parallel jobs for copying files
+    engine : str, default 'boto3'
+        Engine to use for S3 operations ('boto3' or 'minio')
+
+    Returns
+    -------
+    list
+        List of results from copy operations
+
+    Raises
+    ------
+    ValueError
+        If storage or storage_name not provided
+        If source_files or destination_files not lists
+        If source_files and destination_files different lengths
+    Joao Couto - 2024
+    '''
+    if storage is None:
+        if storage_name is None:
+            raise ValueError("Specify a storage to copy to - either pass the storage dictionary or specify a name from the prefs.")
+        storage = prefs['storage'][storage_name] # link to preferences storage from storage_name
+    storage = validate_storage(storage) # validate and update keys, this will store the credentials
+
+    if not type(source_files) is list: # check type of source
+        raise ValueError('source_files has to be a list of paths from s3')
+    
+    if not type(destination_files) is list:  # check type of destination
+        raise ValueError(f'destination_files has to be a list of paths {destination_files}')
+    # Check if the source and the destination are the correct sizes
+    assert len(source_files) == len(destination_files),ValueError(f'source {source_files} and destination {destination_files} are the wrong size')
+    if len(source_files) == 1: # don't run in parallel if only copying one file.
+        print(f'[Downloading] {source_files[0]}')
+        res = [copyfile_from_s3(source_files[0],destination_files[0],storage = storage)]
+    else:
+        from tqdm import tqdm
+        n_jobs = validate_num_jobs_joblib(n_jobs) # avoid nested parallelism.
+        res = Parallel(n_jobs = n_jobs)(delayed(copyfile_from_s3)(src,
+                                                                  dst,
+                                                                  storage = storage,
+                                                                  engine = engine)
+                                        for src,dst in tqdm(zip(source_files,destination_files),
+                                                            desc = f'Pulling from S3 [{storage["bucket"]}]'))
+    return res
+
+
+def s3_delete_file(filepath,storage, remove_versions = False):
+    '''
+    s3_delete_file(filepath,storage, remove_versions = False)
+
+    Deletes files from s3.
+    TODO: make this parallel or a parallel version of the function.
+    !!Warning!! this does not wait for confirmation.
+    remove_versions = True will delete all versions from the bucket.
+
+    Parameters
+    ----------
+    filepath : str
+        Path to file to delete on S3
+    storage : dict
+        Dictionary containing S3 storage configuration:
+        - endpoint: S3 endpoint URL
+        - access_key: S3 access key
+        - secret_key: S3 secret key
+        - bucket: S3 bucket name
+    remove_versions : bool, default False
+        Whether to remove all versions of the file
+        If True, deletes all versions
+        If False, only deletes latest version
+
+    Returns
+    -------
+    object
+        Response from S3 delete operation
+
+    Raises
+    ------
+    MinioException
+        If delete operation fails
+        If file does not exist
+        If credentials are invalid
+
+    Joao Couto - 2024
+    '''
+    from minio import Minio
+    client = Minio(endpoint = storage['endpoint'],
+                   access_key = storage['access_key'],
+                   secret_key = storage['secret_key'])
+
+    if remove_versions:
+        objects = client.list_objects(storage['bucket'], prefix=filepath,include_version=True)
+        for obj in objects:
+            res = client.remove_object(storage['bucket'], obj.object_name,version_id = obj.version_id)
+        return 
+    else:
+        res = client.remove_object(storage['bucket'], filepath)
+    return res

labdata/schema/__init__.py

@@ -0,0 +1,24 @@
+'''Schema for lab data management.
+
+This package provides DataJoint schemas (tables) for accessing and managing laboratory data.
+The schemas are organized into modules by data type:
+
+- `general`     - Core tables for files, subjects, sessions, datasets
+- `procedures`  - Tables for experimental procedures and protocols
+- `ephys`       - Tables for electrophysiology recordings and analysis
+- `twophoton`   - Tables for two-photon microscopy data
+- `onephoton`   - Tables for one-photon imaging (widefield and miniscope)
+- `tasks`       - Tables for behavioral task data
+- `video`       - Tables for video recordings
+- `histology`   - Tables for histology and anatomy data
+'''
+
+from .general import *
+from .procedures import *
+from .ephys import *
+from .twophoton import * 
+from .onephoton import *  # includes widefield and miniscope
+from .tasks import *
+from .video import *
+from .histology import *
+