Merge branch '287_spawn_pool' into 'dev'

Run process pools with spawn after pythonnet starts

See merge request mass-spectrometry/corems!276

Author: kheal

corems/chroma_peak/calc/subset.py

@@ -1,12 +1,13 @@
 # This file contains functions for subsetting dataframes that contain mass feature data.
 # This is based on the deimos package, found here: https://github.com/pnnl/deimos/blob/master/deimos/subset.py with some modifications.
 
-import multiprocessing as mp
 from functools import partial
 
 import numpy as np
 import pandas as pd
 
+from corems.encapsulation.spawn_pool import spawn_pool
+
 class MultiSamplePartitions:
     '''
     Generator object that will lazily build and return each partition constructed
@@ -157,7 +158,7 @@ def map(self, func, processes=1, **kwargs):
 
         # Parallel
         else:
-            with mp.Pool(processes=processes) as p:
+            with spawn_pool(processes) as p:
                 result = list(p.imap(partial(func, **kwargs), self))
 
         # Add partition index

corems/encapsulation/spawn_pool.py

@@ -0,0 +1,58 @@
+"""Process pools that still return after pythonnet has started."""
+
+from contextlib import contextmanager
+from multiprocessing import get_context
+from multiprocessing.pool import TERMINATE
+
+
+@contextmanager
+def spawn_pool(processes):
+    """Yield a spawn process pool and kill its workers on exit.
+
+    Parameters
+    ----------
+    processes : int
+        Number of worker processes. The caller chooses this count.
+
+    Yields
+    ------
+    multiprocessing.pool.Pool
+        Pool started with the ``spawn`` start method.
+
+    Notes
+    -----
+    Linux ``multiprocessing.Pool`` forks. A forked child deadlocks after
+    pythonnet has started CoreCLR or Mono, so the pool never finishes.
+    Spawn starts a new interpreter. Those workers ignore SIGTERM, and
+    ``Pool.terminate`` waits forever in its ``join``. The caller reads
+    results inside the ``with`` block. On the way out, including when the
+    body raises, the workers are killed and reaped.
+    """
+    pool = get_context("spawn").Pool(processes)
+    try:
+        yield pool
+    finally:
+        _kill_spawn_pool(pool)
+
+
+def _kill_spawn_pool(pool):
+    """Stop a spawn pool without ``Pool.terminate``.
+
+    ``terminate`` signals the workers with SIGTERM and then joins them.
+    A worker that has loaded CoreCLR or Mono ignores that signal, so the
+    join never returns. The results are already in the parent. Kill the
+    workers and cancel the pool finalizer, which would call ``terminate``
+    again at garbage collection.
+    """
+    pool._state = TERMINATE
+    for handler in (pool._worker_handler, pool._task_handler, pool._result_handler):
+        handler._state = TERMINATE
+    pool._change_notifier.put(None)
+    for process in list(pool._pool):
+        if process.is_alive():
+            try:
+                process.kill()
+            except ProcessLookupError:
+                pass
+        process.join()
+    pool._terminate.cancel()

corems/mass_spectra/calc/lc_calc.py

@@ -4,7 +4,7 @@
 
 import numpy as np
 import pandas as pd
-import warnings, scipy, multiprocessing
+import warnings, scipy
 
 if not hasattr(np, 'trapezoid'):  # numpy < 2.0
     np.trapezoid = np.trapz
@@ -18,6 +18,7 @@
 
 from corems.chroma_peak.factory.chroma_peak_classes import LCMSMassFeature
 from corems.encapsulation.plot_utils import _finalize_plot
+from corems.encapsulation.spawn_pool import spawn_pool
 from corems.mass_spectra.calc import SignalProcessing as sp
 from corems.mass_spectra.factory.chromat_data import EIC_Data
 from corems.mass_spectrum.input.numpyArray import ms_from_array_profile
@@ -4478,16 +4479,13 @@ def reload_representative_mass_features(self, add_ms2=False, auto_process_ms2=Tr
                                                   ms2_scan_filter=ms2_scan_filter)
         else:
             # Parallel processing
-            import multiprocessing
             from tqdm import tqdm
 
             if self.parameters.lcms_collection.cores > len(sample_mf_map):
                 ncores = len(sample_mf_map)
             else:
                 ncores = self.parameters.lcms_collection.cores
 
-            pool = multiprocessing.Pool(ncores)
-            
             # Build arguments list for starmap
             args_list = [
                 (sample_id, sample_mf_map[sample_id], add_ms2, auto_process_ms2, 
@@ -4496,9 +4494,8 @@ def reload_representative_mass_features(self, add_ms2=False, auto_process_ms2=Tr
             ]
 
             # Execute in parallel
-            mp_result = pool.starmap(self._reload_sample_mass_features, args_list)
-            pool.close()
-            pool.join()
+            with spawn_pool(ncores) as pool:
+                mp_result = pool.starmap(self._reload_sample_mass_features, args_list)
 
             # Collect results back into samples
             for i, sample_id in enumerate(tqdm(sample_mf_map.keys(), desc="Collecting reloaded mass features", unit="sample")):
@@ -5412,11 +5409,11 @@ def fill_missing_cluster_features(self):
                 ncores = len(self)
             else:
                 ncores = self.parameters.lcms_collection.cores
-            pool = multiprocessing.Pool(ncores)
-            mp_result = pool.starmap(
-                self._search_for_targeted_mass_features_in_sample, 
-                [(x, missingdf, cluster_dict, expand_on_miss, False) for x in range(sample_ct)]
-            )
+            with spawn_pool(ncores) as pool:
+                mp_result = pool.starmap(
+                    self._search_for_targeted_mass_features_in_sample, 
+                    [(x, missingdf, cluster_dict, expand_on_miss, False) for x in range(sample_ct)]
+                )
 
             for i in tqdm(range(sample_ct), desc="Collecting gap-filled features", unit="sample"):
                 self[i].induced_mass_features = mp_result[i]
@@ -5536,15 +5533,11 @@ def process_samples_pipeline(self, operations, description=None, keep_raw_data=F
                         results_by_operation[op_name][sample_id] = result
         else:
             # Parallel processing
-            import multiprocessing
-            
             if self.parameters.lcms_collection.cores > sample_ct:
                 ncores = sample_ct
             else:
                 ncores = self.parameters.lcms_collection.cores
 
-            pool = multiprocessing.Pool(ncores)
-            
             # Build arguments for each sample
             args_list = [
                 (sample_id, operations, runtime_params, False)
@@ -5554,45 +5547,43 @@ def process_samples_pipeline(self, operations, description=None, keep_raw_data=F
             # Execute in parallel with progress tracking
             results_by_operation = {op.name: {} for op in operations}
 
-            if show_progress:
-                from tqdm import tqdm
-                import time
-                
-                # Use starmap_async for parallel execution with progress tracking
-                async_result = pool.starmap_async(self._execute_sample_pipeline, args_list)
-                
-                # Poll for completion and update progress bar
-                print(description)
-                pbar = tqdm(
-                    total=sample_ct, 
-                    desc="",
-                    unit="sample", 
-                    position=0,
-                    leave=True,
-                    dynamic_ncols=True
-                )
-                prev_completed = 0
-                while not async_result.ready():
-                    # Get number of completed tasks by checking remaining
-                    completed = sample_ct - async_result._number_left
-                    if completed > prev_completed:
-                        pbar.update(completed - prev_completed)
-                        prev_completed = completed
-                    time.sleep(0.5)  # Poll every 500ms to avoid spam
-                
-                # Final update to 100%
-                if prev_completed < sample_ct:
-                    pbar.update(sample_ct - prev_completed)
-                pbar.close()
-                
-                # Get all results
-                mp_results = async_result.get()
-            else:
-                # Execute without progress
-                mp_results = pool.starmap(self._execute_sample_pipeline, args_list)
-            
-            pool.close()
-            pool.join()
+            with spawn_pool(ncores) as pool:
+                if show_progress:
+                    from tqdm import tqdm
+                    import time
+                    
+                    # Use starmap_async for parallel execution with progress tracking
+                    async_result = pool.starmap_async(self._execute_sample_pipeline, args_list)
+                    
+                    # Poll for completion and update progress bar
+                    print(description)
+                    pbar = tqdm(
+                        total=sample_ct, 
+                        desc="",
+                        unit="sample", 
+                        position=0,
+                        leave=True,
+                        dynamic_ncols=True
+                    )
+                    prev_completed = 0
+                    while not async_result.ready():
+                        # Get number of completed tasks by checking remaining
+                        completed = sample_ct - async_result._number_left
+                        if completed > prev_completed:
+                            pbar.update(completed - prev_completed)
+                            prev_completed = completed
+                        time.sleep(0.5)  # Poll every 500ms to avoid spam
+                    
+                    # Final update to 100%
+                    if prev_completed < sample_ct:
+                        pbar.update(sample_ct - prev_completed)
+                    pbar.close()
+                    
+                    # Get all results
+                    mp_results = async_result.get()
+                else:
+                    # Execute without progress
+                    mp_results = pool.starmap(self._execute_sample_pipeline, args_list)
 
             # Collect results back into collection
             for result in mp_results:

corems/mass_spectra/input/corems_hdf5.py

@@ -16,6 +16,7 @@
 import warnings
 
 from corems.chroma_peak.factory.chroma_peak_classes import LCMSMassFeature
+from corems.encapsulation.spawn_pool import spawn_pool
 from corems.encapsulation.input.parameter_from_json import (
     load_and_set_json_parameters_lcms,
     load_and_set_toml_parameters_lcms,
@@ -1231,9 +1232,9 @@ def get_lcms_collection(self, load_raw = False, load_light = True, use_original_
             else:
                 ncores = self._cores
             # Create a pool of workers (one for each core or sample, whichever is smaller)
-            pool = multiprocessing.Pool(ncores)
             args = [(sample, load_raw, load_light, use_original_parser) for sample in samples]
-            lcms_objs = pool.starmap(self.get_lcms_obj, args)
+            with spawn_pool(ncores) as pool:
+                lcms_objs = pool.starmap(self.get_lcms_obj, args)
             for sample_name, lcms_obj in zip(samples, lcms_objs):
                 lcms_coll._lcms[sample_name] = lcms_obj
 

corems/molecular_id/factory/MolecularLookupTable.py

@@ -6,7 +6,6 @@
 import io
 import itertools
 import json
-import multiprocessing
 import pstats
 from copy import deepcopy
 from typing import Dict
@@ -17,6 +16,7 @@
 
 from corems import chunks, timeit
 from corems.encapsulation.constant import Atoms
+from corems.encapsulation.spawn_pool import spawn_pool
 from corems.encapsulation.factory.parameters import MSParameters
 from corems.encapsulation.factory.processingSetting import (
     MolecularLookupDictSettings,
@@ -447,13 +447,11 @@ def runworker(self, molecular_search_settings, **kwargs):
                 worker_args = [
                     (chunk, settings.url_database) for chunk in list_insert_chunks
                 ]
-                p = multiprocessing.Pool(settings.db_jobs)
-                for class_list in tqdm(
-                        p.imap_unordered(insert_database_worker, worker_args), disable= not verbose
-                        ):
-                    pass
-                p.close()
-                p.join()
+                with spawn_pool(settings.db_jobs) as p:
+                    for class_list in tqdm(
+                            p.imap_unordered(insert_database_worker, worker_args), disable= not verbose
+                            ):
+                        pass
 
         return classes_list
 

support_code/nmdc/lipidomics/lipidomics_workflow.py

@@ -7,7 +7,6 @@
 import sys
 
 sys.path.append("./")
-from multiprocessing import Pool
 from pathlib import Path
 import datetime
 import tomlkit
@@ -16,6 +15,7 @@
 import pandas as pd
 import time
 
+from corems.encapsulation.spawn_pool import spawn_pool
 from corems.mass_spectra.input.corems_hdf5 import ReadCoreMSHDFMassSpectra
 from corems.mass_spectra.input.mzml import MZMLSpectraParser
 from corems.mass_spectra.input.rawFileReader import ImportMassSpectraThermoMSFileReader
@@ -682,7 +682,6 @@ def run_lipid_workflow(
             )
             mz_dicts.append(mz_dict)
     elif cores > 1:
-        pool = Pool(cores)
         args = [
             (
                 str(file_in),
@@ -694,9 +693,8 @@ def run_lipid_workflow(
             )
             for file_in, file_out in list(zip(files_list, out_paths_list))
         ]
-        mz_dicts = pool.starmap(run_lipid_sp_ms1, args)
-        pool.close()
-        pool.join()
+        with spawn_pool(cores) as pool:
+            mz_dicts = pool.starmap(run_lipid_sp_ms1, args)
 
     # Skip ms2 spectral search for now, will add back once we have an improved MetabRefLCInterface method
     """

support_code/nmdc/metabolomics/gcms_workflow.py

@@ -5,10 +5,10 @@
 sys.path.append("./")
 
 from pathlib import Path
-from multiprocessing import Pool
 import cProfile
 
 
+from corems.encapsulation.spawn_pool import spawn_pool
 from corems.molecular_id.input.nistMSI import ReadNistMSI
 from corems.mass_spectra.input.andiNetCDF import ReadAndiNetCDF
 from corems.molecular_id.search.compoundSearch import LowResMassSpectralMatch
@@ -103,11 +103,9 @@ def calibrate_and_search(out_put_file_name, jobs, dms_file_path):
             file_dialog.close()
 
             # run in multiprocessing mode
-            pool = Pool(jobs)
             args = [(file_path, ref_dict, cal_file_path) for file_path in file_locations[0]]
-            gcmss = pool.map(run, args)
-            pool.close()
-            pool.join()
+            with spawn_pool(jobs) as pool:
+                gcmss = pool.map(run, args)
 
             for file_index, gcms in enumerate(gcmss):
 

support_code/nmdc/metabolomics/metabolomics_collection.py

@@ -2,13 +2,13 @@
 import time
 import pandas as pd
 import numpy as np
-from multiprocessing import Pool
 import shutil
 
 from corems.mass_spectra.input.corems_hdf5 import ReadCoreMSHDFMassSpectraCollection
 from corems.mass_spectra.output.export import LCMSMetabolomicsExport
 from corems.mass_spectra.input.rawFileReader import ImportMassSpectraThermoMSFileReader
 from corems.encapsulation.factory.parameters import LCMSParameters
+from corems.encapsulation.spawn_pool import spawn_pool
 from corems.molecular_id.search.database_interfaces import MSPInterface
 from corems.encapsulation.factory.parameters import hush_output
 from corems.mass_spectra.output.export import LCMSCollectionExport
@@ -309,7 +309,7 @@ def preprocess_raw_samples(raw_data_path, processed_folder, ncores=1, reprocess=
     process_args = [(raw_file, processed_folder, params) for raw_file in raw_files]
 
     # Process samples in parallel
-    with Pool(processes=ncores) as pool:
+    with spawn_pool(ncores) as pool:
         processed_files = pool.map(process_single_sample, process_args)
 
     print(f"Preprocessing complete: {time.time() - start_time:.1f} seconds using {ncores} cores")