parallel.py

import numpy.typing as npt
import multiprocessing as mp
import logging
import queue
import time
import contextlib
from multiprocessing.managers import BaseProxy
from functools import reduce
from pytom_tm.tmjob import TMJob
from pytom_tm.utils import mute_stdout_stderr

try:  # new processes need to be spawned in order to set cupy to use the correct GPU
    mp.set_start_method('spawn')
except RuntimeError:
    pass


def gpu_runner(
        gpu_id: int, task_queue: BaseProxy, result_queue: BaseProxy, log_level: int, unittest_mute: bool
) -> None:
    """Start a GPU runner, each runner should be initialized to a multiprocessing.Process() and manage running jobs
    on a single GPU. Each runner will grab jobs from the task_queue and assign jobs to the result_queue once they
    finish. When the task_queue is empty the gpu_runner will stop.

    Parameters
    ----------
    gpu_id: int
        a GPU index to assign to the runner
    task_queue: mp.managers.BaseProxy
        shared queue from multiprocessing with jobs to run
    result_queue: mp.manager.BaseProxy
        shared queue from multiprocessing for finished jobs
    log_level: int
        log level for logging
    unittest_mute: Bool
        optional muting of runner to prevent unittests flooding the terminal, only use for development
    """
    if unittest_mute:
        mute_context = mute_stdout_stderr
    else:
        mute_context = contextlib.nullcontext
    with mute_context():
        logging.basicConfig(level=log_level)
        while True:
            try:
                job = task_queue.get_nowait()
                result_queue.put_nowait(job.start_job(gpu_id, return_volumes=False))
            except queue.Empty:
                break


def run_job_parallel(
        main_job: TMJob, volume_splits: tuple[int, int, int], gpu_ids: list[int, ...], unittest_mute: bool = False,
) -> tuple[npt.NDArray[float], npt.NDArray[float]]:
    """Run a job in parallel over a single or multiple GPUs. If no volume_splits are given the search is parallelized
    by splitting the angular search. If volume_splits are provided the job will first be split by volume,
    if there are still more GPUs available, the subvolume jobs are still further split by angular search.

    Parameters
    ----------
    main_job: pytom_tm.tmjob.TMJob
        a TMJob object from pytom_tm that contains all data for a search
    volume_splits: tuple[int, int, int]
        tuple of len 3 with splits in x, y, and z
    gpu_ids: list[int, ...]
        list of gpu indices to spread jobs over
    unittest_mute: bool, default False
        boolean to mute spawned process terminal output, only set to True for unittesting

    Returns
    -------
    result: tuple[npt.NDArray[float], npt.NDArray[float]]
        the volumes with the LCCmax and angle ids
    """

    n_pieces = reduce(lambda x, y: x * y, volume_splits)
    jobs = []

    # =================== Splitting into subjobs ===============
    if n_pieces == 1:
        if len(gpu_ids) > 1:  # split rotation search

            jobs = main_job.split_rotation_search(len(gpu_ids))

        else:  # we just run the whole tomo on a single gpu

            jobs.append(main_job)

    elif n_pieces > 1:

        rotation_split_factor = len(gpu_ids) // n_pieces

        if rotation_split_factor >= 2:  # we can split the rotation search for the subvolumes

            for j in main_job.split_volume_search(volume_splits):

                jobs += j.split_rotation_search(rotation_split_factor)

        else:  # only split the subvolume search

            jobs = main_job.split_volume_search(volume_splits)

    else:
        raise ValueError('Invalid number of pieces in split volume')

    # ================== Execution of jobs =========================
    if len(jobs) == 1:

        return main_job.start_job(gpu_ids[0], return_volumes=True)

    elif len(jobs) >= len(gpu_ids):

        results = []

        with mp.Manager() as manager:
            task_queue = manager.Queue()  # the list of tasks where processes can get there next task from
            result_queue = manager.Queue()  # this will accumulate results from the processes

            [task_queue.put_nowait(j) for j in jobs]  # put all tasks

            # set the processes and start them!
            procs = [mp.Process(target=gpu_runner, args=(
                g, task_queue, result_queue, main_job.log_level, unittest_mute)) for g in gpu_ids]
            [p.start() for p in procs]

            while True:
                while not result_queue.empty():
                    results.append(result_queue.get_nowait())

                if len(results) == len(jobs):  # its done if all the results from the spawn were send back
                    logging.debug('Got all results from the child processes')
                    break

                for p in procs:  # if one of the processes is no longer alive and has a failed exit we should error
                    if not p.is_alive() and p.exitcode == 1:  # to prevent a deadlock
                        [x.terminate() for x in procs]  # kill all spawned processes if something broke
                        raise RuntimeError('One or more of the processes stopped unexpectedly.')

                time.sleep(1)

            [p.join() for p in procs]
            logging.debug('Terminated the processes')

        # merge split jobs; pass along the list of stats to annotate them in main_job
        return main_job.merge_sub_jobs(stats=results)

    else:
        ValueError('For some reason there are more gpu_ids than split job, this should never happen.')