Process pool

reditools/tools/analyze/main.py

@@ -2,14 +2,14 @@
 
 import argparse
 import sys
-from multiprocessing import Process, Queue
+from functools import partial
+from multiprocessing.context import TimeoutError
+from multiprocessing.pool import Pool
 
 from reditools.logger import Logger
-from reditools.region import Region
 from reditools.tools.analyze.concat_output import concat_output
-from reditools.tools.analyze.monitor import monitor
 from reditools.tools.analyze.parse_args import parse_args
-from reditools.tools.analyze.redi_thread import redi_thread
+from reditools.tools.analyze.redi_thread import REDIThreadManager
 from reditools.tools.analyze.region_args import region_args
 
 
@@ -51,44 +51,23 @@ def setup_logger(options: argparse.Namespace) -> Logger:
         return Logger(Logger.info_level)
     return Logger(Logger.silent_level)
 
-def fill_queue(options: argparse.Namespace) -> Queue[tuple[int, Region] | None]:
+def pool_error(pool: Pool, debug: bool, exc: Exception) -> None:
     """
-    Fill the input queue with genomic regions to be analyzed.
+    Terminates a multiprocessing Pool.
 
     Parameters
     ----------
-    options : argparse.Namespace
-        The parsed command line options.
-
-    Returns
-    -------
-    Queue[tuple[int, Region] | None]
-        A queue containing indexed Region objects.
-
-    Raises
-    ------
-    SystemExit
-        If a required file is not found.
+    pool : Pool
+        mutliprocessing Pool to terminate.
+    debug : bool
+        If True, raises the exception passed in the third argument.
+    exc : Exception
+        Exception responsible for the pool to terminate.
     """
-    in_queue: Queue[tuple[int, Region] | None] = Queue()
-    try:
-        for _ in enumerate(region_args(options)):  # noqa: WPS468
-            in_queue.put(_)
-    except FileNotFoundError as exc:
-        sys.stderr.write(f'[ERROR] {exc}\n')
-        sys.exit(1)
-
-    # Check thread count
-    if in_queue.qsize() < options.threads:
-        sys.stderr.write(
-            "[WARNING] You have assigned more threads "
-            f"({options.threads}) than there are genomic ranges "
-            f"({in_queue.qsize()})\n",
-        )
-        options.threads = in_queue.qsize()
-    for _ in range(options.threads):
-        in_queue.put(None)
-    return in_queue
+    pool.terminate()
+    if debug:
+        raise exc.__cause__  # type: ignore[misc]
+    sys.stderr.write(f'[ERROR] ({type(exc)}) {exc}\n')
 
 def main() -> None:
     """
@@ -107,19 +86,36 @@ def main() -> None:
 
     options.encoding = 'utf-8'
 
-    in_queue = fill_queue(options)
+    regions = region_args(options)
 
-    # Start parallel jobs
-    out_queue: Queue[tuple[int, str]] = Queue()
-    processes = []
-    for _ in range(options.threads):
-        processes.append(Process(
-            target=redi_thread,
-            args=(options, in_queue, out_queue),
-        ))
+    if options.threads > len(regions):
+        sys.stderr.write(
+            f"[WARNING] You have assigned {options.threads} threads, "
+            f"But there are only {len(regions)} genomic range(s). "
+            "Consider change the value of --window\n"
+        )
+        options.threads = len(regions)
+    try:
+        with Pool(
+            options.threads,
+            REDIThreadManager.init_thread,
+            (options,),
+        ) as pool:
+            imap_iter = [
+                pool.apply_async(
+                    REDIThreadManager.analyze,
+                    args=(region,),
+                    error_callback=partial(pool_error, pool, options.debug),
+                ) for region in regions
+            ]
+            pool.close()
+            pool.join()
+            temp_files = [_.get(1) for _ in imap_iter]
+    except (TimeoutError, IndexError):
+        sys.exit(1)
 
     concat_output(
-        monitor(processes, out_queue, in_queue.qsize()),
+        temp_files,
         options.output_file,
         'a' if options.append_file else 'w',
         options.encoding,

reditools/tools/analyze/redi_thread.py

@@ -1,10 +1,5 @@
 import argparse
-import sys
-import traceback
-from multiprocessing import Queue
 
-from reditools.alignment_manager import AlignmentManager
-from reditools.reditools import REDItools
 from reditools.region import Region
 from reditools.tools.analyze.rtchecks import RTChecks
 from reditools.tools.analyze.setup_alignment_manager import \
@@ -13,82 +8,81 @@
 from reditools.tools.analyze.write_results import write_results
 
 
-def analyze(
-        options: argparse.Namespace,
-        rtools: REDItools,
-        sam_manager: AlignmentManager,
-        region: Region,
-        rtqc: RTChecks,
-) -> str:
-    """Analyze a specific genomic region.
-
-    Parameters
-    ----------
-    options : argparse.Namespace
-        The command-line options.
-    rtools : REDItools
-        The REDItools analysis engine.
-    sam_manager : AlignmentManager
-        The alignment file manager.
-    region : Region
-        The genomic region to analyze.
-    rtqc : RTChecks
-        The quality control checks to apply.
-
-    Returns
-    -------
-    str
-        The path to the temporary file containing the results.
-    """
-    rtresults = rtools.analyze(sam_manager, region)
-    return write_results(
-        rtresults,
-        options.temp_dir,
-        rtqc,
-        rtools.log,
-    )
-
-def redi_thread(
-        options: argparse.Namespace,
-        in_queue: Queue,
-        out_queue: Queue,
-) -> bool:
-    """Worker thread function for parallel REDItools analysis.
-
-    Parameters
-    ----------
-    options : argparse.Namespace
-        The command-line options.
-    in_queue : Queue
-        The queue containing genomic regions to analyze.
-    out_queue : Queue
-        The queue to put analysis results into.
-
-    Returns
-    -------
-    bool
-        True when the worker has finished processing all regions.
-    """
-    rtools = setup_rtools(options)
-    sam_manager = setup_alignment_manager(
-        options.file,
-        options.min_read_quality,
-        options.min_read_length,
-        options.exclude_reads,
-    )
-    rtqc = RTChecks(options)
-    while True:
-        args = in_queue.get()
-        if args is None:
-            return True
-        idx, region = args
-        try:  # noqa: WPS229
-            out_queue.put((
-                idx,
-                analyze(options, rtools, sam_manager, region, rtqc),
-            ))
-        except Exception as exc:
-            if options.debug:
-                traceback.print_exception(*sys.exc_info())
-            sys.stderr.write(f'[ERROR] ({type(exc)}) {exc}\n')
-            sys.exit(1)
+class REDIThread:
+    def __init__(self, options: argparse.Namespace) -> None:
+        """Worker thread function for parallel REDItools analysis.
+
+        Parameters
+        ----------
+        options : argparse.Namespace
+            The command-line options.
+        """
+        self.rtools = setup_rtools(options)
+        self.sam_manager = setup_alignment_manager(
+            options.file,
+            options.min_read_quality,
+            options.min_read_length,
+            options.exclude_reads,
+        )
+        self.rtqc = RTChecks(options)
+        self.temp_dir = options.temp_dir
+
+    def analyze(
+            self,
+            region: Region,
+    ) -> str:
+        """Analyze a specific genomic region.
+
+        Parameters
+        ----------
+        region : Region
+            The genomic region to analyze.
+
+        Returns
+        -------
+        str
+            The path to the temporary file containing the results.
+        """
+        rtresults = self.rtools.analyze(self.sam_manager, region)
+        return write_results(
+            rtresults,
+            self.temp_dir,
+            self.rtqc,
+            self.rtools.log,
+        )
+
+class REDIThreadManager:
+    """Manages a worker thread function for parallel REDItools analysis."""
+
+    thread: None | REDIThread = None
+
+    @classmethod
+    def init_thread(cls, options: argparse.Namespace) -> None:
+        """Initialize a REDIThread.
+
+        Parameters
+        ----------
+        options : argparse.Namespace
+            The command-line options.
+        """
+
+        cls.thread = REDIThread(options)
+
+    @classmethod
+    def analyze(cls, region: Region) -> str:
+        """Instruct thread to analyze a specific genomic region.
+
+        Parameters
+        ----------
+        region : Region
+            The genomic region to analyze.
+
+        Returns
+        -------
+        str
+            The path to the temporary file containing the results.
+        """
+
+        if cls.thread is None:
+            raise AttributeError('REDIThreadManager not initialized.')
+        return cls.thread.analyze(region)