Merge pull request #214 from chrisfilo/enh/threads

[RTM] multithreading tweaks

fmriprep/run_workflow.py

@@ -59,7 +59,7 @@ def main():
 
     # ANTs options
     g_ants = parser.add_argument_group('specific settings for ANTs registrations')
-    g_ants.add_argument('--ants-nthreads', action='store', type=int,
+    g_ants.add_argument('--ants-nthreads', action='store', type=int, default=0,
                         help='number of threads that will be set in ANTs processes')
     g_ants.add_argument('--skull-strip-ants', dest="skull_strip_ants",
                         action='store_true',
@@ -85,6 +85,7 @@ def create_workflow(opts):
         'write_graph': opts.write_graph,
         'nthreads': opts.nthreads,
         'debug': opts.debug,
+        'ants_nthreads': opts.ants_nthreads,
         'skull_strip_ants': opts.skull_strip_ants,
         'output_dir': op.abspath(opts.output_dir),
         'work_dir': op.abspath(opts.work_dir)
@@ -97,9 +98,6 @@ def create_workflow(opts):
         settings['ants_t1-mni_settings'] = 't1-mni_registration_test'
         logger.setLevel(logging.DEBUG)
 
-    if opts.ants_nthreads is not None:
-        settings['ants_threads'] = opts.ants_nthreads
-
     log_dir = op.join(settings['output_dir'], 'log')
     derivatives = op.join(settings['output_dir'], 'derivatives')
 
@@ -134,6 +132,9 @@ def create_workflow(opts):
             plugin_settings['plugin'] = 'MultiProc'
             plugin_settings['plugin_args'] = {'n_procs': settings['nthreads']}
 
+    if settings['ants_nthreads'] == 0:
+        settings['ants_nthreads'] = cpu_count()
+
     # Determine subjects to be processed
     subject_list = opts.participant_label
 

fmriprep/workflows/anatomical.py

@@ -68,12 +68,15 @@ def t1w_preprocessing(name='t1w_preprocessing', settings=None):
     t1_2_mni = pe.Node(
         RobustMNINormalizationRPT(
             generate_report=True,
-            num_threads=settings.get('ants_threads', 6),
+            num_threads=settings['ants_nthreads'],
             testing=settings.get('debug', False),
             template='mni_icbm152_nlin_asym_09c'
         ),
         name='T1_2_MNI_Registration'
     )
+    # should not be necesssary byt does not hurt - make sure the multiproc
+    # scheduler knows the resource limits
+    t1_2_mni.interface.num_threads = settings['ants_nthreads']
 
     # Resample the brain mask and the tissue probability maps into mni space
     bmask_mni = pe.Node(
@@ -307,8 +310,14 @@ def skullstrip_ants(name='ANTsBrainExtraction', settings=None):
 
     t1_skull_strip = pe.Node(BrainExtractionRPT(
         dimension=3, use_floatingpoint_precision=1,
-        debug=settings['debug'], generate_report=True),
+        debug=settings['debug'], generate_report=True,
+        num_threads=settings['ants_nthreads']),
         name='Ants_T1_Brain_Extraction')
+
+    # should not be necesssary byt does not hurt - make sure the multiproc
+    # scheduler knows the resource limits
+    t1_skull_strip.interface.num_threads = settings['ants_nthreads']
+
     t1_skull_strip.inputs.brain_template = op.join(
         get_ants_oasis_template_ras(),
         'T_template0.nii.gz'

test/workflows/test_base.py

@@ -10,7 +10,8 @@ class TestBase(TestWorkflow):
     def test_wf_ds054_type(self, _):
         # set up
         mock_subject_data = {'t1w': ['um'], 'sbref': ['um'], 'func': 'um'}
-        mock_settings = {'output_dir': '.', 'work_dir': '.'}
+        mock_settings = {'output_dir': '.', 'work_dir': '.',
+                         'ants_nthreads': 1}
 
         # run
         wf054 = wf_ds054_type(mock_subject_data, mock_settings)
@@ -37,7 +38,7 @@ def test_wf_ds054_type(self, _):
     def test_wf_ds005_type(self, _):
         # set up
         mock_subject_data = {'func': ''}
-        mock_settings = {'output_dir': '.'}
+        mock_settings = {'output_dir': '.', 'ants_nthreads': 1}
 
         # run
         wf005 = wf_ds005_type(mock_subject_data, mock_settings)