Improve memory allocation (#1581)

Author: tsalo

xcp_d/data/tests/config.toml

@@ -26,17 +26,15 @@ run_uuid = "20240205-123456"
 templateflow_home = "~/.cache/templateflow"
 work_dir = "work/"
 write_graph = false
+atlases = ["Gordon"]
 
 [workflow]
-cifti = false
+file_format = "nifti"
 dummy_scans = 0
 input_type = "fmriprep"
 despike = false
 smoothing = 6
 combine_runs = false
-motion_filter_type = false
-band_stop_min = false
-band_stop_max = false
 motion_filter_order = 4
 head_radius = 50
 fd_thresh = 0.3
@@ -45,7 +43,6 @@ bandpass_filter = true
 high_pass = 0.01
 low_pass = 0.1
 bpf_order = 2
-atlases = []
 min_coverage = 0.5
 correlation_lengths = []
 process_surfaces = false

xcp_d/workflows/anatomical/plotting.py

@@ -32,6 +32,7 @@ def init_brainsprite_figures_wf(
     t1w_available,
     t2w_available,
     apply_transform,
+    mem_gb=None,
     name='brainsprite_figures_wf',
 ):
     """Create mosaic and PNG files for executive summary brainsprite.
@@ -61,6 +62,8 @@ def init_brainsprite_figures_wf(
         True if a T2w image is available.
     apply_transform : bool
         Whether to apply the transform to the surfaces.
+    mem_gb : :obj:`dict` or None
+        Memory size in GB. If None, a default of 1 GB per volume is used.
     %(name)s
         Default is "init_brainsprite_figures_wf".
 
@@ -76,6 +79,9 @@ def init_brainsprite_figures_wf(
     rh_pial_surf
     template_to_anat_xfm
     """
+    if mem_gb is None:
+        mem_gb = {'volume': 1}
+
     workflow = Workflow(name=name)
 
     inputnode = pe.Node(
@@ -140,7 +146,7 @@ def init_brainsprite_figures_wf(
         plot_slices = pe.Node(
             PlotSlicesForBrainSprite(n_procs=config.nipype.omp_nthreads),
             name=f'plot_slices_{image_type}',
-            mem_gb=config.DEFAULT_MEMORY_MIN_GB,
+            mem_gb=mem_gb['volume'],
             n_procs=config.nipype.omp_nthreads,
         )
         workflow.connect([

xcp_d/workflows/base.py

@@ -40,7 +40,7 @@
 )
 from xcp_d.utils.doc import fill_doc
 from xcp_d.utils.modified_data import calculate_exact_scans, flag_bad_run
-from xcp_d.utils.utils import estimate_brain_radius, is_number
+from xcp_d.utils.utils import _create_mem_gb, estimate_brain_radius, is_number
 from xcp_d.workflows.anatomical.parcellation import init_parcellate_surfaces_wf
 from xcp_d.workflows.anatomical.surface import init_postprocess_surfaces_wf
 from xcp_d.workflows.anatomical.volume import init_postprocess_anat_wf
@@ -116,7 +116,7 @@ def init_single_subject_wf(subject_id: str, anat_session: str, func_sessions: li
             from xcp_d.workflows.base import init_single_subject_wf
 
             with mock_config():
-                wf = init_single_subject_wf("01", "01", ["01"])
+                wf = init_single_subject_wf("01", "", [""])
 
     Parameters
     ----------
@@ -517,6 +517,8 @@ def init_single_subject_wf(subject_id: str, anat_session: str, func_sessions: li
             }
 
         n_processed_task_runs = 0
+        concat_bold_gb = 0
+        concat_volume_gb = 0
         for j_run, bold_file in enumerate(task_files):
             run_data = collect_run_data(
                 layout=config.execution.layout,
@@ -568,6 +570,14 @@ def init_single_subject_wf(subject_id: str, anat_session: str, func_sessions: li
                     bold_file=bold_file,
                 )
 
+            # Compute memory estimates from the BOLD file header
+            mem_gbx = _create_mem_gb(bold_file)
+
+            # Accumulate memory estimates for the concatenation workflow.
+            # The concatenated file will be the sum of all runs.
+            concat_bold_gb += mem_gbx['bold']
+            concat_volume_gb = max(concat_volume_gb, mem_gbx['volume'])
+
             postprocess_bold_wf = init_postprocess_bold_wf(
                 bold_file=bold_file,
                 head_radius=head_radius,
@@ -577,6 +587,7 @@ def init_single_subject_wf(subject_id: str, anat_session: str, func_sessions: li
                 n_runs=n_runs,
                 has_multiple_runs=multiscans,
                 exact_scans=exact_scans,
+                mem_gb=mem_gbx,
                 name=f'postprocess_{run_counter}_wf',
             )
             run_counter += 1
@@ -621,9 +632,14 @@ def init_single_subject_wf(subject_id: str, anat_session: str, func_sessions: li
                     ])  # fmt:skip
 
         if config.workflow.combine_runs and (n_processed_task_runs > 0) and multiscans:
+            concat_mem_gb = {
+                'bold': concat_bold_gb,
+                'volume': concat_volume_gb,
+            }
             concatenate_data_wf = init_concatenate_data_wf(
                 TR=TR,
                 head_radius=head_radius,
+                mem_gb=concat_mem_gb,
                 name=f'concatenate_entity_set_{ent_set}_wf',
             )
 

xcp_d/workflows/bold/cifti.py

@@ -11,7 +11,6 @@
 from xcp_d import config
 from xcp_d.interfaces.utils import ConvertTo32
 from xcp_d.utils.doc import fill_doc
-from xcp_d.utils.utils import _create_mem_gb
 from xcp_d.workflows.bold.connectivity import init_functional_connectivity_cifti_wf
 from xcp_d.workflows.bold.metrics import init_alff_wf, init_reho_cifti_wf
 from xcp_d.workflows.bold.outputs import init_postproc_derivatives_wf
@@ -38,6 +37,7 @@ def init_postprocess_cifti_wf(
     n_runs,
     has_multiple_runs,
     exact_scans,
+    mem_gb,
     name='cifti_postprocess_wf',
 ):
     """Organize the cifti processing workflow.
@@ -63,10 +63,13 @@ def init_postprocess_cifti_wf(
 
                 run_data = collect_run_data(
                     layout=layout,
-                    input_type="fmriprep",
                     bold_file=bold_file,
-                    cifti=True,
+                    file_format="cifti",
+                    target_space="MNI152NLin2009cAsym",
                 )
+                run_data['confounds'] = None
+
+                from xcp_d.utils.utils import _create_mem_gb
 
                 wf = init_postprocess_cifti_wf(
                     bold_file=bold_file,
@@ -77,6 +80,7 @@ def init_postprocess_cifti_wf(
                     n_runs=1,
                     has_multiple_runs=False,
                     exact_scans=[],
+                    mem_gb=_create_mem_gb(bold_file),
                     name="cifti_postprocess_wf",
                 )
 
@@ -95,6 +99,8 @@ def init_postprocess_cifti_wf(
         Whether there are multiple runs for this task or not.
         Interacts with the output_run_wise_correlations parameter.
     %(exact_scans)s
+    mem_gb : :obj:`dict`
+        Dictionary of memory allocations with keys ``'bold'`` and ``'volume'``.
     %(name)s
         Default is "cifti_postprocess_wf".
 
@@ -207,7 +213,7 @@ def init_postprocess_cifti_wf(
         name='outputnode',
     )
 
-    mem_gbx = _create_mem_gb(bold_file)
+    mem_gbx = mem_gb
 
     downcast_data = pe.Node(
         ConvertTo32(),
@@ -227,6 +233,7 @@ def init_postprocess_cifti_wf(
         TR=TR,
         exact_scans=exact_scans,
         head_radius=head_radius,
+        mem_gb=mem_gbx,
     )
 
     workflow.connect([
@@ -257,7 +264,7 @@ def init_postprocess_cifti_wf(
     ])  # fmt:skip
 
     if despike:
-        despike_wf = init_despike_wf(TR=TR)
+        despike_wf = init_despike_wf(TR=TR, mem_gb=mem_gbx)
 
         workflow.connect([
             (prepare_confounds_wf, despike_wf, [

xcp_d/workflows/bold/concatenation.py

@@ -20,7 +20,7 @@
 
 
 @fill_doc
-def init_concatenate_data_wf(TR, head_radius, name='concatenate_data_wf'):
+def init_concatenate_data_wf(TR, head_radius, mem_gb, name='concatenate_data_wf'):
     """Concatenate postprocessed data across runs and directions.
 
     Workflow Graph
@@ -36,13 +36,17 @@ def init_concatenate_data_wf(TR, head_radius, name='concatenate_data_wf'):
                 wf = init_concatenate_data_wf(
                     TR=2,
                     head_radius=50,
+                    mem_gb={"bold": 2.0, "volume": 0.01},
                     name="concatenate_data_wf",
                 )
 
     Parameters
     ----------
     %(TR)s
     %(head_radius)s
+    mem_gb : :obj:`dict`
+        Dictionary of memory allocations with keys ``'bold'`` and ``'volume'``.
+        This should represent the *total* memory across all runs being concatenated.
     %(name)s
         Default is "concatenate_data_wf".
 
@@ -83,9 +87,6 @@ def init_concatenate_data_wf(TR, head_radius, name='concatenate_data_wf'):
     fd_thresh = config.workflow.fd_thresh
     atlases = config.execution.atlases
 
-    # Guess memory needs since they can't be estimated from the inputs
-    mem_gb = {'bold': 6.0, 'volume': 1.0}
-
     workflow.__desc__ = """
 Postprocessing derivatives from multi-run tasks were then concatenated across runs and directions.
 """

xcp_d/workflows/bold/connectivity.py

@@ -122,7 +122,7 @@ def init_functional_connectivity_nifti_wf(
         NiftiParcellate(min_coverage=min_coverage),
         name='parcellate_data',
         iterfield=['atlas', 'atlas_labels'],
-        mem_gb=mem_gb['bold'],
+        mem_gb=2 * mem_gb['bold'],
     )
     workflow.connect([
         (inputnode, parcellate_data, [

xcp_d/workflows/bold/metrics.py

@@ -147,7 +147,7 @@ def init_alff_wf(
             high_pass=high_pass,
             n_threads=config.nipype.omp_nthreads,
         ),
-        mem_gb=mem_gb['bold'],
+        mem_gb=2 * mem_gb['bold'],
         name='alff_compt',
         n_procs=config.nipype.omp_nthreads,
     )
@@ -274,7 +274,7 @@ def init_reho_cifti_wf(
             with mock_config():
                 wf = init_reho_cifti_wf(
                     name_source="/path/to/bold.dtseries.nii",
-                    mem_gb={"volume": 0.1},
+                    mem_gb={"bold": 0.1},
                     name="cifti_reho_wf",
                 )
 
@@ -359,12 +359,12 @@ def init_reho_cifti_wf(
     lh_reho = pe.Node(
         SurfaceReHo(surf_hemi='L'),
         name='reho_lh',
-        mem_gb=mem_gb['bold'],
+        mem_gb=2 * mem_gb['bold'],
     )
     rh_reho = pe.Node(
         SurfaceReHo(surf_hemi='R'),
         name='reho_rh',
-        mem_gb=mem_gb['bold'],
+        mem_gb=2 * mem_gb['bold'],
     )
     subcortical_reho = pe.Node(
         ReHoNamePatch(neighborhood='vertices'),
@@ -442,7 +442,7 @@ def init_reho_nifti_wf(name_source, mem_gb, name='reho_nifti_wf'):
             with mock_config():
                 wf = init_reho_nifti_wf(
                     name_source="/path/to/bold.nii.gz",
-                    mem_gb={"volume": 0.1},
+                    mem_gb={"bold": 0.1},
                     name="nifti_reho_wf",
                 )
 
@@ -484,7 +484,7 @@ def init_reho_nifti_wf(name_source, mem_gb, name='reho_nifti_wf'):
     compute_reho = pe.Node(
         ReHoNamePatch(neighborhood='vertices'),
         name='reho_3d',
-        mem_gb=mem_gb['bold'],
+        mem_gb=2 * mem_gb['bold'],
         n_procs=1,
     )
     # Get the svg

xcp_d/workflows/bold/nifti.py

@@ -11,7 +11,6 @@
 from xcp_d import config
 from xcp_d.interfaces.utils import ConvertTo32
 from xcp_d.utils.doc import fill_doc
-from xcp_d.utils.utils import _create_mem_gb
 from xcp_d.workflows.bold.connectivity import init_functional_connectivity_nifti_wf
 from xcp_d.workflows.bold.metrics import init_alff_wf, init_reho_nifti_wf
 from xcp_d.workflows.bold.outputs import init_postproc_derivatives_wf
@@ -38,6 +37,7 @@ def init_postprocess_nifti_wf(
     n_runs,
     has_multiple_runs,
     exact_scans,
+    mem_gb,
     name='bold_postprocess_wf',
 ):
     """Organize the bold processing workflow.
@@ -66,10 +66,13 @@ def init_postprocess_nifti_wf(
 
                 run_data = collect_run_data(
                     layout=layout,
-                    input_type="fmriprep",
                     bold_file=bold_file,
-                    cifti=False,
+                    file_format="nifti",
+                    target_space="MNI152NLin2009cAsym",
                 )
+                run_data['confounds'] = None
+
+                from xcp_d.utils.utils import _create_mem_gb
 
                 wf = init_postprocess_nifti_wf(
                     bold_file=bold_file,
@@ -80,6 +83,7 @@ def init_postprocess_nifti_wf(
                     n_runs=1,
                     has_multiple_runs=False,
                     exact_scans=[],
+                    mem_gb=_create_mem_gb(bold_file),
                     name="nifti_postprocess_wf",
                 )
 
@@ -97,6 +101,8 @@ def init_postprocess_nifti_wf(
         Whether there are multiple runs for this task or not.
         Interacts with the output_run_wise_correlations parameter.
     %(exact_scans)s
+    mem_gb : :obj:`dict`
+        Dictionary of memory allocations with keys ``'bold'`` and ``'volume'``.
     %(name)s
         Default is "nifti_postprocess_wf".
 
@@ -215,7 +221,7 @@ def init_postprocess_nifti_wf(
         name='outputnode',
     )
 
-    mem_gbx = _create_mem_gb(bold_file)
+    mem_gbx = mem_gb
 
     downcast_data = pe.Node(
         ConvertTo32(),
@@ -240,6 +246,7 @@ def init_postprocess_nifti_wf(
         TR=TR,
         exact_scans=exact_scans,
         head_radius=head_radius,
+        mem_gb=mem_gbx,
     )
 
     workflow.connect([
@@ -271,7 +278,7 @@ def init_postprocess_nifti_wf(
     ])  # fmt:skip
 
     if despike:
-        despike_wf = init_despike_wf(TR=TR)
+        despike_wf = init_despike_wf(TR=TR, mem_gb=mem_gbx)
 
         workflow.connect([
             (prepare_confounds_wf, despike_wf, [