diff --git a/reditools/tools/analyze/main.py b/reditools/tools/analyze/main.py index 2a489ba..fd7a2bc 100644 --- a/reditools/tools/analyze/main.py +++ b/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, diff --git a/reditools/tools/analyze/monitor.py b/reditools/tools/analyze/monitor.py deleted file mode 100644 index f1a8434..0000000 --- a/reditools/tools/analyze/monitor.py +++ /dev/null @@ -1,59 +0,0 @@ -from __future__ import annotations - -import sys -from multiprocessing import Process, Queue -from queue import Empty as EmptyQueueException - - -def check_dead(processes: list[Process]) -> None: - """Check if any of the processes have failed. - - If a process has exited with code 1, all other processes are killed - and the program exits. - - Parameters - ---------- - processes : list[Process] - The list of processes to monitor. - """ - for proc in processes: - if proc.exitcode == 1: - for to_kill in processes: - to_kill.kill() - sys.stderr.write('[ERROR] Killing job\n') - sys.exit(1) - -def monitor( - processes: list[Process], - out_queue: Queue[tuple[int, str]], - chunks: int, -) -> list[str]: - """Monitor progress of parallel analysis processes. - - Parameters - ---------- - processes : list[Process] - The list of worker processes. - out_queue : Queue[tuple[int, str]] - The queue containing analysis result filenames and their indices. - chunks : int - The total number of work chunks. - - Returns - ------- - list[str] - A list of filenames containing analysis results, ordered by chunk index. - """ - tfs = ['' for _ in range(chunks - len(processes))] - - for prc in processes: - prc.start() - - while '' in tfs: - try: - idx, fname = out_queue.get(block=False, timeout=1) - except EmptyQueueException: - check_dead(processes) - else: - tfs[idx] = fname - return tfs diff --git a/reditools/tools/analyze/redi_thread.py b/reditools/tools/analyze/redi_thread.py index 63e7b79..fd53197 100644 --- a/reditools/tools/analyze/redi_thread.py +++ b/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)