Merge pull request #1058 from int-brain-lab/datahandler

allow custom data handler to be passed to a task

Author: oliche

brainbox/tests/test_behavior.py

@@ -259,7 +259,7 @@ def test_get_movement_onset(self):
         expected = [np.nan, 79.66293334, 100.73593334, 129.26693334, np.nan]
         np.testing.assert_array_almost_equal(times, expected)
         with self.assertRaises(ValueError):
-            wheel.get_movement_onset(intervals, np.random.permutation(self.trials['feedback_times']))
+            wheel.get_movement_onset(intervals, np.flipud(self.trials['feedback_times']))
 
 
 class TestTraining(unittest.TestCase):

ibllib/oneibl/patcher.py

@@ -243,8 +243,10 @@ def __init__(self, client_name='default', one=None, label='ibllib patch'):
         # transfers/delete from the current computer to the flatiron: mandatory and executed first
         local_id = self.endpoints['local']['id']
         self.globus_transfer = globus_sdk.TransferData(
-            self.client, local_id, flatiron_id, verify_checksum=True, sync_level='checksum', label=label)
-        self.globus_delete = globus_sdk.DeleteData(self.client, flatiron_id, label=label)
+            source_endpoint=local_id, destination_endpoint=flatiron_id,
+            verify_checksum=True, sync_level='checksum', label=label
+        )
+        self.globus_delete = globus_sdk.DeleteData(endpoint=flatiron_id, label=label)
         # transfers/delete from flatiron to optional third parties to synchronize / delete
         self.globus_transfers_locals = {}
         self.globus_deletes_locals = {}
@@ -303,7 +305,7 @@ def patch_datasets(self, file_list, **kwargs):
                 # if there is no transfer already created, initialize it
                 if repo_gid not in self.globus_transfers_locals:
                     self.globus_transfers_locals[repo_gid] = globus_sdk.TransferData(
-                        self.client, flatiron_id, repo_gid, verify_checksum=True,
+                        source_endpoint=flatiron_id, destination_endpoint=repo_gid, verify_checksum=True,
                         sync_level='checksum', label=f"{self.label} on {fr['data_repository']}")
                 # get the local server path and create the transfer item
                 local_server_path = self.to_address(fr['relative_path'], fr['data_repository'])
@@ -343,17 +345,15 @@ def _wait_for_task(resp):
             _wait_for_task(gtc.submit_transfer(self.globus_transfer))
             # re-initialize the globus_transfer property
             self.globus_transfer = globus_sdk.TransferData(
-                gtc,
-                self.globus_transfer['source_endpoint'],
-                self.globus_transfer['destination_endpoint'],
+                source_endpoint=self.globus_transfer['source_endpoint'],
+                destination_endpoint=self.globus_transfer['destination_endpoint'],
                 label=self.globus_transfer['label'],
                 verify_checksum=True, sync_level='checksum')
 
         # do the same for deletes
         if len(self.globus_delete['DATA']) > 0:
             _wait_for_task(gtc.submit_delete(self.globus_delete))
             self.globus_delete = globus_sdk.DeleteData(
-                gtc,
                 endpoint=self.globus_delete['endpoint'],
                 label=self.globus_delete['label'])
 

ibllib/pipes/tasks.py

@@ -116,7 +116,8 @@ class Task(abc.ABC):
     on_error = 'continue'  # whether to raise an exception on error ('raise') or report the error and continue ('continue')
 
     def __init__(self, session_path, parents=None, taskid=None, one=None,
-                 machine=None, clobber=True, location='server', scratch_folder=None, on_error='continue', force=False, **kwargs):
+                 machine=None, clobber=True, location='server', scratch_folder=None, on_error='continue',
+                 force=False, data_handler_class=None, **kwargs):
         """
         Base task class
         :param session_path: session path
@@ -127,7 +128,10 @@ def __init__(self, session_path, parents=None, taskid=None, one=None,
         :param clobber: whether or not to overwrite log on rerun
         :param location: location where task is run. Options are 'server' (lab local servers'), 'remote' (remote compute node,
         data required for task downloaded via one), 'AWS' (remote compute node, data required for task downloaded via AWS),
-        or 'SDSC' (SDSC flatiron compute node)
+        or 'SDSC' (SDSC flatiron compute node). The data_handler_class parameter will override the location in
+         determining the handler class.
+        :param data_handler_class: custom class to handle data. If not provided, the location will
+         be used to infer the handler class.
         :param scratch_folder: optional: Path where to write intermediate temporary data
         :param force: whether to re-download missing input files on local server if not present
         :param args: running arguments
@@ -149,6 +153,7 @@ def __init__(self, session_path, parents=None, taskid=None, one=None,
         self.plot_tasks = []  # Plotting task/ tasks to create plot outputs during the task
         self.scratch_folder = scratch_folder
         self.kwargs = kwargs
+        self.data_handler_class = data_handler_class
 
     @property
     def signature(self) -> Dict[str, List]:
@@ -551,27 +556,28 @@ def get_data_handler(self, location=None):
         Gets the relevant data handler based on location argument
         :return:
         """
-        location = str.lower(location or self.location)
-        if location == 'local':
-            return data_handlers.LocalDataHandler(self.session_path, self.signature, one=self.one)
-        self.one = self.one or ONE()
-        if location == 'server':
-            dhandler = data_handlers.ServerDataHandler(self.session_path, self.signature, one=self.one)
-        elif location == 'serverglobus':
-            dhandler = data_handlers.ServerGlobusDataHandler(self.session_path, self.signature, one=self.one)
-        elif location == 'remote':
-            dhandler = data_handlers.RemoteHttpDataHandler(self.session_path, self.signature, one=self.one)
-        elif location == 'aws':
-            dhandler = data_handlers.RemoteAwsDataHandler(self.session_path, self.signature, one=self.one)
-        elif location == 'sdsc':
-            dhandler = data_handlers.SDSCDataHandler(self.session_path, self.signature, one=self.one)
-        elif location == 'popeye':
-            dhandler = data_handlers.PopeyeDataHandler(self.session_path, self.signature, one=self.one)
-        elif location == 'ec2':
-            dhandler = data_handlers.RemoteEC2DataHandler(self.session_path, self.signature, one=self.one)
-        else:
-            raise ValueError(f'Unknown location "{location}"')
-        return dhandler
+        if self.data_handler_class is None:
+            location = str.lower(location or self.location)
+            if location == 'local':
+                return data_handlers.LocalDataHandler(self.session_path, self.signature, one=self.one)
+            self.one = self.one or ONE()
+            if location == 'server':
+                self.data_handler_class = data_handlers.ServerDataHandler
+            elif location == 'serverglobus':
+                self.data_handler_class = data_handlers.ServerGlobusDataHandler
+            elif location == 'remote':
+                self.data_handler_class = data_handlers.RemoteHttpDataHandler
+            elif location == 'aws':
+                self.data_handler_class = data_handlers.RemoteAwsDataHandler
+            elif location == 'sdsc':
+                self.data_handler_class = data_handlers.SDSCDataHandler
+            elif location == 'popeye':
+                self.data_handler_class = data_handlers.PopeyeDataHandler
+            elif location == 'ec2':
+                self.data_handler_class = data_handlers.RemoteEC2DataHandler
+            else:
+                raise ValueError(f'Unknown location "{location}"')
+        return self.data_handler_class(self.session_path, self.signature, one=self.one)
 
     @staticmethod
     def make_lock_file(taskname='', time_out_secs=7200):

ibllib/tests/test_tasks.py

@@ -10,6 +10,7 @@
 import numpy as np
 
 import ibllib.pipes.tasks
+import ibllib.oneibl.data_handlers
 from ibllib.pipes.base_tasks import ExperimentDescriptionRegisterRaw
 from ibllib.pipes.video_tasks import VideoConvert
 from ibllib.io import session_params
@@ -401,6 +402,16 @@ def test_input_files_to_register(self):
             for f in ('alf/foo.bar.*', 'alf/bar.baz.npy', 'alf/baz.foo.npy'):
                 self.assertRegex(cm.output[-1], re.escape(f))
 
+    def test_location_data_handler(self):
+        task = Task00(self.session_path)
+        self.assertIsInstance(task.get_data_handler(), ibllib.oneibl.data_handlers.ServerDataHandler)
+
+        class TotoDataHandler(ibllib.oneibl.data_handlers.DataHandler):
+            pass
+
+        task = Task00(self.session_path, data_handler_class=TotoDataHandler)
+        self.assertIsInstance(task.get_data_handler(), TotoDataHandler)
+
 
 class TestMisc(unittest.TestCase):
     """Tests for misc functions in ibllib.pipes.tasks module."""

requirements.txt

@@ -2,7 +2,7 @@ boto3
 click>=7.0.0
 colorlog>=4.0.2
 flake8>=3.7.8
-globus-sdk
+globus-sdk>=4.0.0
 graphviz
 matplotlib>=3.0.3
 numba>=0.56
@@ -28,7 +28,7 @@ iblutil>=1.13.0
 iblqt>=0.8.2
 mtscomp>=1.0.1
 ONE-api>=3.2.0
-phylib>=2.6.0
+phylib>=2.6.2
 psychofit
 slidingRP>=1.1.1  # steinmetz lab refractory period metrics
 pyqt5