1010
1111from one .api import ONE
1212
13- from ibllib .io .extractors .base import get_session_extractor_type , get_pipeline
14- from ibllib .pipes import ephys_preprocessing , training_preprocessing , tasks
13+ from ibllib .io .extractors .base import get_pipeline , get_task_protocol , get_session_extractor_type
14+ from ibllib .pipes import tasks , training_preprocessing , ephys_preprocessing
1515from ibllib .time import date2isostr
1616import ibllib .oneibl .registration as registration
1717
1818_logger = logging .getLogger ('ibllib' )
19+ LARGE_TASKS = ['EphysVideoCompress' , 'TrainingVideoCompress' , 'EphysDLC' , 'TrainingDLC' , 'SpikeSorting' ]
20+
21+
22+ def _get_pipeline_class (session_path , one ):
23+ pipeline = get_pipeline (session_path )
24+ if pipeline == 'training' :
25+ PipelineClass = training_preprocessing .TrainingExtractionPipeline
26+ elif pipeline == 'ephys' :
27+ PipelineClass = ephys_preprocessing .EphysExtractionPipeline
28+ else :
29+ # try and look if there is a custom extractor in the personal projects extraction class
30+ import projects .base
31+ task_type = get_session_extractor_type (session_path )
32+ PipelineClass = projects .base .get_pipeline (task_type )
33+ _logger .info (f"Using { PipelineClass } pipeline for { session_path } " )
34+ return PipelineClass (session_path = session_path , one = one )
1935
2036
2137def _get_lab (one ):
@@ -101,16 +117,10 @@ def job_creator(root_path, one=None, dry=False, rerun=False, max_md5_size=None):
101117 session_path , one = one , max_md5_size = max_md5_size )
102118 if dsets is not None :
103119 all_datasets .extend (dsets )
104- pipeline = get_pipeline (session_path )
105- if pipeline == 'training' :
106- pipe = training_preprocessing .TrainingExtractionPipeline (session_path , one = one )
107- # only start extracting ephys on a raw_session.flag
108- elif pipeline == 'ephys' and flag_file .name == 'raw_session.flag' :
109- pipe = ephys_preprocessing .EphysExtractionPipeline (session_path , one = one )
110- else :
111- _logger .info (f'Session type { get_session_extractor_type (session_path )} '
112- f'as no matching pipeline pattern { session_path } ' )
113- continue
120+ pipe = _get_pipeline_class (session_path , one )
121+ if pipe is None :
122+ task_protocol = get_task_protocol (session_path )
123+ _logger .info (f'Session task protocol { task_protocol } has no matching pipeline pattern { session_path } ' )
114124 if rerun :
115125 rerun__status__in = '__all__'
116126 else :
@@ -125,11 +135,12 @@ def job_creator(root_path, one=None, dry=False, rerun=False, max_md5_size=None):
125135 return all_datasets
126136
127137
128- def job_runner (subjects_path , lab = None , dry = False , one = None , count = 5 ):
138+ def job_runner (subjects_path , mode = 'all' , lab = None , dry = False , one = None , count = 5 ):
129139 """
130140 Function to be used as a process to run the jobs as they are created on the database
131141 This will query waiting jobs from the specified Lab
132142 :param subjects_path: on servers: /mnt/s0/Data/Subjects. Contains sessions
143+ :param mode: Whether to run all jobs, or only small or large (video compression, DLC, spike sorting) jobs
133144 :param lab: lab name as per Alyx
134145 :param dry:
135146 :param count:
@@ -141,8 +152,18 @@ def job_runner(subjects_path, lab=None, dry=False, one=None, count=5):
141152 lab = _get_lab (one )
142153 if lab is None :
143154 return # if the lab is none, this will return empty tasks each time
144- tasks = one .alyx .rest ('tasks' , 'list' , status = 'Waiting' ,
145- django = f'session__lab__name__in,{ lab } ' )
155+ # Filter for tasks
156+ if mode == 'all' :
157+ tasks = one .alyx .rest ('tasks' , 'list' , status = 'Waiting' ,
158+ django = f'session__lab__name__in,{ lab } ' , no_cache = True )
159+ elif mode == 'small' :
160+ tasks_all = one .alyx .rest ('tasks' , 'list' , status = 'Waiting' ,
161+ django = f'session__lab__name__in,{ lab } ' , no_cache = True )
162+ tasks = [t for t in tasks_all if t ['name' ] not in LARGE_TASKS ]
163+ elif mode == 'large' :
164+ tasks = one .alyx .rest ('tasks' , 'list' , status = 'Waiting' ,
165+ django = f'session__lab__name__in,{ lab } ,name__in,{ LARGE_TASKS } ' , no_cache = True )
166+
146167 tasks_runner (subjects_path , tasks , one = one , count = count , time_out = 3600 , dry = dry )
147168
148169
0 commit comments