♻️ Attach guardrails directly to Nodes

Author: shnizzedy

CPAC/pipeline/nipype_pipeline_engine/engine.py

@@ -469,8 +469,7 @@ class GuardrailedNode:
         ``node = wf.guardrailed_node(node, reference, registered, pipe_num)``
         to automatically build guardrails
         """
-        def __init__(self, wf, node, reference, registered, pipe_num,
-                     retry=True):
+        def __init__(self, wf, node, reference, registered, pipe_num, retry):
             '''A Node with guardrails
 
             Parameters
@@ -493,39 +492,34 @@ def __init__(self, wf, node, reference, registered, pipe_num,
             retry : bool
                 retry if run is so configured
             '''
-            self.guardrails = [registration_guardrail_node(
-                f'{node.name}_guardrail_{pipe_num}')]
             self.node = node
+            self.node.guardrail = registration_guardrail_node(
+                f'{node.name}_guardrail_{pipe_num}')
             self.reference = reference
             self.registered = registered
-            self.retries = []
+            self.tries = [node]
             self.wf = wf
-            self.wf.connect(self.node, registered,
-                            self.guardrails[0], 'registered')
             if retry and self.wf.num_tries > 1:
                 if registration_guardrails.retry_on_first_failure:
-                    self.guardrails.append(registration_guardrail_node(
-                        f'{node.name}_guardrail'))
-                    self.retries.append(retry_clone(self.node))
-                    self.retries[0].interface.inputs.add_trait(
+                    self.tries.append(retry_clone(self.node))
+                    self.tries[1].interface.inputs.add_trait(
                         'previous_failure', traits.Bool())
-                    self.guardrails.append(registration_guardrail_node(
-                        f'{self.retries[0].name}_guardrail',
-                        raise_on_failure=True))
-                    self.wf.connect(self.guardrails[0], 'failed_qc',
-                                    self.retries[0], 'previous_failure')
+                    self.tries[1].guardrail = registration_guardrail_node(
+                        f'{self.tries[1].name}_guardrail',
+                        raise_on_failure=True)
+                    self.wf.connect(self.tries[0].guardrail, 'failed_qc',
+                                    self.tries[1], 'previous_failure')
                 else:
                     num_retries = self.wf.num_tries - 1
                     for i in range(num_retries):
-                        self.retries.append(retry_clone(self.node, i + 2))
-                        self.guardrails.append(registration_guardrail_node(
-                            f'{self.retries[-1].name}_guardrail',
-                            raise_on_failure=(i + 1 == num_retries)))
-                for i, _retry in enumerate(self.retries):
-                    self.wf.connect(_retry, registered,
-                                    self.guardrails[i + 1], 'registered')
-            for guardrail in self.guardrails:
-                guardrail.inputs.reference = self.reference
+                        self.tries.append(retry_clone(self.node, i + 2))
+                        self.tries[i + 1].guardrail = (
+                            registration_guardrail_node(
+                                f'{self.tries[i + 1].name}_guardrail',
+                                raise_on_failure=(i + 1 == num_retries)))
+            for i, _try in enumerate(self.tries):
+                self.wf.connect(_try, registered, _try.guardrail, 'registered')
+                _try.guardrail.inputs.reference = self.reference
 
         def __getattr__(self, __name):
             """Get attributes from the node that is guardrailed if that
@@ -684,7 +678,8 @@ def guardrail(self):
         """
         return any(registration_guardrails.thresholds.values())
 
-    def guardrailed_node(self, node, reference, registered, pipe_num):
+    def guardrailed_node(self, node, reference, registered, pipe_num,
+                         retry=True):
         """Method to return a GuardrailedNode in the given Workflow.
 
         .. seealso:: Workflow.GuardrailedNode
@@ -702,9 +697,12 @@ def guardrailed_node(self, node, reference, registered, pipe_num):
 
         pipe_num : int
             int
+
+        retry : bool
+            retry if run is so configured
         """
         return self.GuardrailedNode(self, node, reference, registered,
-                                    pipe_num)
+                                    pipe_num, retry)
 
     def guardrail_selection(self, node: 'Workflow.GuardrailedNode',
                             output_key: str) -> Node:

CPAC/registration/registration.py

@@ -22,7 +22,6 @@
 from nipype.interfaces.afni import utils as afni_utils
 from CPAC.anat_preproc.lesion_preproc import create_lesion_preproc
 from CPAC.func_preproc.utils import chunk_ts, split_ts_chunks
-from CPAC.qc.globals import registration_guardrails
 from CPAC.registration.utils import seperate_warps_list, \
                                     check_transforms, \
                                     generate_inverse_transform_flags, \
@@ -922,7 +921,6 @@ def create_bbregister_func_to_anat(phase_diff_distcor, name, bbreg_status,
             Functional data in anatomical space
     """
     suffix = f'{bbreg_status.title()}_{pipe_num}'
-    retry = bbreg_status == 'On'
     register_bbregister_func_to_anat = pe.Workflow(name=f'{name}_{suffix}')
     inputspec = pe.Node(util.IdentityInterface(fields=['func',
                                                        'anat',
@@ -961,7 +959,8 @@ def bbreg_args(bbreg_target):
                                  name=f'bbreg_func_to_anat_{suffix}')
     bbreg_func_to_anat.inputs.dof = 6
     bbreg_func_to_anat = register_bbregister_func_to_anat.guardrailed_node(
-        bbreg_func_to_anat, 'reference', 'out_file', pipe_num)
+        bbreg_func_to_anat, 'reference', 'out_file', pipe_num,
+        retry=bbreg_status == 'On')
     register_bbregister_func_to_anat.connect([
         (inputspec, bbreg_func_to_anat, [
             ('bbr_schedule', 'schedule'),
@@ -979,21 +978,14 @@ def bbreg_args(bbreg_target):
                 ('fieldmapmask', 'fieldmapmask')]),
             (inputNode_echospacing, bbreg_func_to_anat, [
                 ('echospacing', 'echospacing')])])
-    if retry and registration_guardrails.retry_on_first_failure:
-        outfile = register_bbregister_func_to_anat.guardrail_selection(
-            bbreg_func_to_anat, 'out_file')
-        matrix = register_bbregister_func_to_anat.guardrail_selection(
-            bbreg_func_to_anat, 'out_matrix_file')
-        register_bbregister_func_to_anat.connect(
-            matrix, 'out', outputspec, 'func_to_anat_linear_xfm')
-        register_bbregister_func_to_anat.connect(outfile, 'out',
-                                                 outputspec, 'anat_func')
-    else:
-        register_bbregister_func_to_anat.connect(
-            bbreg_func_to_anat, 'out_matrix_file',
-            outputspec, 'func_to_anat_linear_xfm')
-        register_bbregister_func_to_anat.connect(
-            bbreg_func_to_anat, 'out_file', outputspec, 'anat_func')
+    outfile = register_bbregister_func_to_anat.guardrail_selection(
+        bbreg_func_to_anat, 'out_file')
+    matrix = register_bbregister_func_to_anat.guardrail_selection(
+        bbreg_func_to_anat, 'out_matrix_file')
+    register_bbregister_func_to_anat.connect(
+        matrix, 'out', outputspec, 'func_to_anat_linear_xfm')
+    register_bbregister_func_to_anat.connect(outfile, 'out',
+                                                outputspec, 'anat_func')
     return register_bbregister_func_to_anat
 
 
@@ -1827,7 +1819,7 @@ def bold_to_T1template_xfm_connector(wf_name, cfg, reg_tool, symmetric=False,
             name='change_transform_type')
 
         wf.connect(fsl_reg_2_itk, 'itk_transform',
-                         change_transform, 'input_affine_file')
+                   change_transform, 'input_affine_file')
 
         # combine ALL xfm's into one - makes it easier downstream
         write_composite_xfm = pe.Node(
@@ -2862,10 +2854,6 @@ def coregistration(wf, cfg, strat_pool, pipe_num, opt=None):
             cfg.registration_workflows['functional_registration'][
                 'coregistration']['boundary_based_registration'][
                 'bbr_wm_mask_args']
-        if fallback:
-            bbreg_guardrail = pe.registration_guardrail_node(
-                f'bbreg{bbreg_status}_guardrail_{pipe_num}',
-                raise_on_failure=False)
 
         node, out = strat_pool.get_data('sbref')
         wf.connect(node, out, func_to_anat_bbreg, 'inputspec.func')
@@ -2875,16 +2863,12 @@ def coregistration(wf, cfg, strat_pool, pipe_num, opt=None):
                 'reference'] == 'whole-head':
             node, out = strat_pool.get_data('desc-head_T1w')
             wf.connect(node, out, func_to_anat_bbreg, 'inputspec.anat')
-            if fallback:
-                wf.connect(node, out, bbreg_guardrail, 'reference')
 
         elif cfg.registration_workflows['functional_registration'][
                 'coregistration']['boundary_based_registration'][
                 'reference'] == 'brain':
             node, out = strat_pool.get_data('desc-preproc_T1w')
             wf.connect(node, out, func_to_anat_bbreg, 'inputspec.anat')
-            if fallback:
-                wf.connect(node, out, bbreg_guardrail, 'reference')
 
         wf.connect(func_to_anat, 'outputspec.func_to_anat_linear_xfm_nobbreg',
                    func_to_anat_bbreg, 'inputspec.linear_reg_matrix')