RF/ENH: Rework workflow generation (#219)

mgxd · web-flow · commit e8e9054a850f · 2022-05-25T09:35:11.000-04:00
* WIP: Clean up excessively complicated build workflow code

* RF: Clean up workflow build

- Builds workflow / boilerplate within the same process
- Less warnings by default
- Set OMP_NUM_THREADS to 1 to reduce VMS load

* FIX: Import

* FIX: Remove forgotten variables

* ENH: Add analysis level attribute to workflow

* FIX: Track cwd on entrypoint

* FIX: Ensure config is properly initialized

* FIX: Suppress warnings on releases
diff --git a/nibabies/cli/parser.py b/nibabies/cli/parser.py
@@ -796,8 +796,13 @@ def parse_args(args=None, namespace=None):
     if missing_subjects:
         parser.error(
             "One or more participant labels were not found in the BIDS directory: "
-            "%s." % ", ".join(missing_subjects)
+            f"{', '.join(missing_subjects)}."
         )
 
     config.execution.participant_label = sorted(participant_label)
     config.workflow.skull_strip_template = config.workflow.skull_strip_template[0]
+
+    # finally, write config to file
+    config_file = config.execution.work_dir / config.execution.run_uuid / "config.toml"
+    config_file.parent.mkdir(exist_ok=True, parents=True)
+    config.to_filename(config_file)
diff --git a/nibabies/cli/run.py b/nibabies/cli/run.py
@@ -6,164 +6,136 @@
 
 def main():
     """Entry point."""
+    import atexit
     import gc
+    import os
     import sys
-    from multiprocessing import Manager, Process
-    from os import EX_SOFTWARE
     from pathlib import Path
 
     from ..utils.bids import write_bidsignore, write_derivative_description
     from .parser import parse_args
+    from .workflow import build_boilerplate, build_workflow
+
+    _cwd = os.getcwd()
+    # Revert OMP_NUM_THREADS + other runtime set environment variables
+    atexit.register(config.restore_env)
 
     parse_args()
 
-    # sentry_sdk = None
-    # if not config.execution.notrack:
-    #     import sentry_sdk
-    #     from ..utils.sentry import sentry_setup
-
-    #     sentry_setup()
-
-    # CRITICAL Save the config to a file. This is necessary because the execution graph
-    # is built as a separate process to keep the memory footprint low. The most
-    # straightforward way to communicate with the child process is via the filesystem.
-    config_file = config.execution.work_dir / config.execution.run_uuid / "config.toml"
-    config_file.parent.mkdir(exist_ok=True, parents=True)
-    config.to_filename(config_file)
-
-    # CRITICAL Call build_workflow(config_file, retval) in a subprocess.
-    # Because Python on Linux does not ever free virtual memory (VM), running the
-    # workflow construction jailed within a process preempts excessive VM buildup.
-    with Manager() as mgr:
-        from .workflow import build_workflow
-
-        retval = mgr.dict()
-        p = Process(target=build_workflow, args=(str(config_file), retval))
-        p.start()
-        p.join()
-
-        retcode = p.exitcode or retval.get("return_code", 0)
-        nibabies_wf = retval.get("workflow", None)
-
-    # CRITICAL Load the config from the file. This is necessary because the ``build_workflow``
-    # function executed constrained in a process may change the config (and thus the global
-    # state of NiBabies).
-    config.load(config_file)
-
-    if config.execution.reports_only:
-        sys.exit(int(retcode > 0))
-
-    if nibabies_wf and config.execution.write_graph:
-        nibabies_wf.write_graph(graph2use="colored", format="svg", simple_form=True)
-
-    retcode = retcode or (nibabies_wf is None) * EX_SOFTWARE
-    if retcode != 0:
-        sys.exit(retcode)
-
-    # Generate boilerplate
-    with Manager() as mgr:
-        from .workflow import build_boilerplate
-
-        p = Process(target=build_boilerplate, args=(str(config_file), nibabies_wf))
-        p.start()
-        p.join()
-
-    if config.execution.boilerplate_only:
-        sys.exit(int(retcode > 0))
-
-    # Clean up master process before running workflow, which may create forks
-    gc.collect()
-
-    # Sentry tracking
-    # if sentry_sdk is not None:
-    #     with sentry_sdk.configure_scope() as scope:
-    #         scope.set_tag("run_uuid", config.execution.run_uuid)
-    #         scope.set_tag("npart", len(config.execution.participant_label))
-    #     sentry_sdk.add_breadcrumb(message="nibabies started", level="info")
-    #     sentry_sdk.capture_message("nibabies started", level="info")
-
-    config.loggers.workflow.log(
-        15,
-        "\n".join(["nibabies config:"] + ["\t\t%s" % s for s in config.dumps().splitlines()]),
-    )
-    config.loggers.workflow.log(25, "nibabies started!")
-    # errno = 1  # Default is error exit unless otherwise set
-    try:
-        nibabies_wf.run(**config.nipype.get_plugin())
-    except Exception as e:
-        # if not config.execution.notrack:
-        #     from ..utils.sentry import process_crashfile
-
-        #     crashfolders = [
-        #         config.execution.nibabies_dir,
-        #         / "sub-{}".format(s)
-        #         / "log"
-        #         / config.execution.run_uuid
-        #         for s in config.execution.participant_label
-        #     ]
-        #     for crashfolder in crashfolders:
-        #         for crashfile in crashfolder.glob("crash*.*"):
-        #             process_crashfile(crashfile)
-
-        #     if "Workflow did not execute cleanly" not in str(e):
-        #         sentry_sdk.capture_exception(e)
-        config.loggers.workflow.critical("nibabies failed: %s", e)
-        raise
-    else:
-        config.loggers.workflow.log(25, "nibabies finished successfully!")
-        # if not config.execution.notrack:
-        #     success_message = "nibabies finished without errors"
-        #     sentry_sdk.add_breadcrumb(message=success_message, level="info")
-        #     sentry_sdk.capture_message(success_message, level="info")
-
-        # Bother users with the boilerplate only iff the workflow went okay.
-        boiler_file = config.execution.nibabies_dir / "logs" / "CITATION.md"
-        if boiler_file.exists():
-            if config.environment.exec_env in (
-                "singularity",
-                "docker",
-                "nibabies-docker",
-            ):
-                boiler_file = Path("<OUTPUT_PATH>") / boiler_file.relative_to(
-                    config.execution.output_dir
-                )
-            config.loggers.workflow.log(
-                25,
-                "Works derived from this nibabies execution should include the "
-                f"boilerplate text found in {boiler_file}.",
+    if "participant" in config.workflow.analysis_level:
+        _pool = None
+        if config.nipype.plugin == "MultiProc":
+            import multiprocessing as mp
+            from concurrent.futures import ProcessPoolExecutor
+            from contextlib import suppress
+
+            # should drastically reduce VMS
+            # see https://github.com/nipreps/mriqc/pull/984 for more details
+            os.environ["OMP_NUM_THREADS"] = "1"
+
+            with suppress(RuntimeError):
+                mp.set_start_method("fork")
+            gc.collect()
+
+            _pool = ProcessPoolExecutor(
+                max_workers=config.nipype.nprocs,
+                initializer=config._process_initializer,
+                initargs=(_cwd, config.nipype.omp_nthreads),
             )
 
-        if config.workflow.run_reconall:
-            from niworkflows.utils.misc import _copy_any
-            from templateflow import api
+        config_file = config.execution.work_dir / config.execution.run_uuid / "config.toml"
+        config_file.parent.mkdir(exist_ok=True, parents=True)
+        config.to_filename(config_file)
 
-            dseg_tsv = str(api.get("fsaverage", suffix="dseg", extension=[".tsv"]))
-            _copy_any(dseg_tsv, str(config.execution.nibabies_dir / "desc-aseg_dseg.tsv"))
-            _copy_any(dseg_tsv, str(config.execution.nibabies_dir / "desc-aparcaseg_dseg.tsv"))
-        # errno = 0
-    finally:
-        from pkg_resources import resource_filename as pkgrf
-
-        from ..reports.core import generate_reports
-
-        # Generate reports phase
-        generate_reports(
-            config.execution.participant_label,
-            config.execution.session_id,
-            config.execution.nibabies_dir,
-            config.execution.run_uuid,
-            config=pkgrf("nibabies", "data/reports-spec.yml"),
-            packagename="nibabies",
+        # build the workflow within the same process
+        # it still needs to be saved / loaded to be properly initialized
+        retval = build_workflow(config_file)
+        retcode = retval['return_code']
+        nibabies_wf = retval['workflow']
+
+        if nibabies_wf is None:
+            if config.execution.reports_only:
+                sys.exit(int(retcode > 0))
+            sys.exit(os.EX_SOFTWARE)
+
+        if config.execution.write_graph:
+            nibabies_wf.write_graph(graph2use="colored", format="svg", simple_form=True)
+
+        if retcode != 0:
+            sys.exit(retcode)
+
+        # generate boilerplate
+        build_boilerplate(nibabies_wf)
+        if config.execution.boilerplate_only:
+            sys.exit(0)
+
+        gc.collect()
+
+        config.loggers.workflow.log(
+            15,
+            "\n".join(["nibabies config:"] + ["\t\t%s" % s for s in config.dumps().splitlines()]),
         )
-        write_derivative_description(config.execution.bids_dir, config.execution.nibabies_dir)
-        write_bidsignore(config.execution.nibabies_dir)
-
-        # if failed_reports and not config.execution.notrack:
-        #     sentry_sdk.capture_message(
-        #         "Report generation failed for %d subjects" % failed_reports,
-        #         level="error",
-        #     )
-        # sys.exit(int((errno + failed_reports) > 0))
+        config.loggers.workflow.log(25, "nibabies started!")
+
+        # Hack MultiProc's pool to reduce VMS
+        _plugin = config.nipype.get_plugin()
+        if _pool:
+            from nipype.pipeline.plugins.multiproc import MultiProcPlugin
+
+            multiproc = MultiProcPlugin(plugin_args=config.nipype.plugin_args)
+            multiproc.pool = _pool
+            _plugin = {"plugin": multiproc}
+
+        gc.collect()
+        try:
+            nibabies_wf.run(**_plugin)
+        except Exception as e:
+            config.loggers.workflow.critical("nibabies failed: %s", e)
+            raise
+        else:
+            config.loggers.workflow.log(25, "nibabies finished successfully!")
+
+            # Bother users with the boilerplate only iff the workflow went okay.
+            boiler_file = config.execution.nibabies_dir / "logs" / "CITATION.md"
+            if boiler_file.exists():
+                if config.environment.exec_env in (
+                    "singularity",
+                    "docker",
+                    "nibabies-docker",
+                ):
+                    boiler_file = Path("<OUTPUT_PATH>") / boiler_file.relative_to(
+                        config.execution.output_dir
+                    )
+                config.loggers.workflow.log(
+                    25,
+                    "Works derived from this nibabies execution should include the "
+                    f"boilerplate text found in {boiler_file}.",
+                )
+
+            if config.workflow.run_reconall:
+                from niworkflows.utils.misc import _copy_any
+                from templateflow import api
+
+                dseg_tsv = str(api.get("fsaverage", suffix="dseg", extension=[".tsv"]))
+                _copy_any(dseg_tsv, str(config.execution.nibabies_dir / "desc-aseg_dseg.tsv"))
+                _copy_any(dseg_tsv, str(config.execution.nibabies_dir / "desc-aparcaseg_dseg.tsv"))
+        # errno = 0
+        finally:
+            from pkg_resources import resource_filename as pkgrf
+
+            from ..reports.core import generate_reports
+
+            # Generate reports phase
+            generate_reports(
+                config.execution.participant_label,
+                config.execution.session_id,
+                config.execution.nibabies_dir,
+                config.execution.run_uuid,
+                config=pkgrf("nibabies", "data/reports-spec.yml"),
+                packagename="nibabies",
+            )
+            write_derivative_description(config.execution.bids_dir, config.execution.nibabies_dir)
+            write_bidsignore(config.execution.nibabies_dir)
 
 
 if __name__ == "__main__":
diff --git a/nibabies/cli/workflow.py b/nibabies/cli/workflow.py
@@ -10,7 +10,7 @@
 """
 
 
-def build_workflow(config_file, retval):
+def build_workflow(config_file):
     """Create the Nipype Workflow that supports the whole execution graph."""
     from niworkflows.utils.bids import check_pipeline_version, collect_participants
     from niworkflows.utils.misc import check_valid_fs_license
@@ -20,19 +20,19 @@ def build_workflow(config_file, retval):
     from ..utils.misc import check_deps
     from ..workflows.base import init_nibabies_wf
 
+    # initalize config
     config.load(config_file)
-    build_log = config.loggers.workflow
+    build_logger = config.loggers.workflow
 
     nibabies_dir = config.execution.nibabies_dir
     version = config.environment.version
 
-    retval["return_code"] = 1
-    retval["workflow"] = None
+    retval = {"return_code": 1, "workflow": None}
 
     # warn if older results exist: check for dataset_description.json in output folder
     msg = check_pipeline_version(version, nibabies_dir / "dataset_description.json")
     if msg is not None:
-        build_log.warning(msg)
+        build_logger.warning(msg)
 
     # Please note this is the input folder's dataset_description.json
     dset_desc_path = config.execution.bids_dir / "dataset_description.json"
@@ -57,7 +57,7 @@ def build_workflow(config_file, retval):
     if config.execution.reports_only:
         from pkg_resources import resource_filename as pkgrf
 
-        build_log.log(25, "Running --reports-only on participants %s", ", ".join(subject_list))
+        build_logger.log(25, "Running --reports-only on participants %s", ", ".join(subject_list))
         retval["return_code"] = generate_reports(
             subject_list,
             nibabies_dir,
@@ -82,13 +82,13 @@ def build_workflow(config_file, retval):
     if config.execution.fs_subjects_dir:
         init_msg += f"""
       * Pre-run FreeSurfer's SUBJECTS_DIR: {config.execution.fs_subjects_dir}."""
-    build_log.log(25, init_msg)
+    build_logger.log(25, init_msg)
 
     retval["workflow"] = init_nibabies_wf(subjects_sessions)
 
     # Check for FS license after building the workflow
     if not check_valid_fs_license():
-        build_log.critical(
+        build_logger.critical(
             """\
 ERROR: a valid license file is required for FreeSurfer to run. nibabies looked for an existing \
 license file at several paths, in this order: 1) command line argument ``--fs-license-file``; \
@@ -101,32 +101,29 @@ def build_workflow(config_file, retval):
     # Check workflow for missing commands
     missing = check_deps(retval["workflow"])
     if missing:
-        build_log.critical(
+        build_logger.critical(
             "Cannot run nibabies. Missing dependencies:%s",
             "\n\t* ".join([""] + [f"{cmd} (Interface: {iface})" for iface, cmd in missing]),
         )
         retval["return_code"] = 127  # 127 == command not found.
         return retval
 
-    config.to_filename(config_file)
-    build_log.info(
+    # config.to_filename(config_file)
+    build_logger.info(
         "NiBabies workflow graph with %d nodes built successfully.",
         len(retval["workflow"]._get_all_nodes()),
     )
     retval["return_code"] = 0
     return retval
 
 
-def build_boilerplate(config_file, workflow):
+def build_boilerplate(workflow):
     """Write boilerplate in an isolated process."""
     from .. import config
 
-    config.load(config_file)
     logs_path = config.execution.nibabies_dir / "logs"
     boilerplate = workflow.visit_desc()
-    citation_files = {
-        ext: logs_path / ("CITATION.%s" % ext) for ext in ("bib", "tex", "md", "html")
-    }
+    citation_files = {ext: logs_path / f"CITATION.{ext}" for ext in ("bib", "tex", "md", "html")}
 
     if boilerplate:
         # To please git-annex users and also to guarantee consistency
diff --git a/nibabies/config.py b/nibabies/config.py
