Skip to content

Commit

Permalink
Add support for multiprocessing parallelism in jackhmmer
Browse files Browse the repository at this point in the history
  • Loading branch information
althonos committed Oct 14, 2024
1 parent 0f25c81 commit 92babdf
Showing 1 changed file with 19 additions and 5 deletions.
24 changes: 19 additions & 5 deletions src/pyhmmer/hmmer/_jackhmmer.py
Original file line number Diff line number Diff line change
Expand Up @@ -41,7 +41,6 @@ class _JACKHMMERWorker(
DigitalSequenceBlock,
_I,
],
threading.Thread
):
def __init__(
self,
Expand Down Expand Up @@ -127,6 +126,14 @@ def _iterate(
return iteration_checkpoints if checkpoints else iteration


class _JACKHMMERThread(_JACKHMMERWorker, threading.Thread):
pass


class _JACKHMMERProcess(_JACKHMMERWorker, multiprocessing.Process):
pass


# --- Dispatcher ---------------------------------------------------------------

class _JACKHMMERDispatcher(
Expand Down Expand Up @@ -173,9 +180,7 @@ def _new_worker(
query_count: "multiprocessing.Value[int]", # type: ignore
kill_switch: threading.Event,
) -> _JACKHMMERWorker[_I]:
if self.backend != "threading":
raise ValueError(f"Invalid backend for `jackhmmer`: {self.backend!r}")
return _JACKHMMERWorker(
params = [
self.targets,
query_queue,
query_count,
Expand All @@ -186,7 +191,13 @@ def _new_worker(
self.max_iterations,
self.select_hits,
self.checkpoints,
)
]
if self.backend == "threading":
return _JACKHMMERThread(*params)
elif self.backend == "multiprocessing":
return _JACKHMMERProcess(*params)
else:
raise ValueError(f"Invalid backend for `jackhmmer`: {self.backend!r}")


# --- jackhmmer -----------------------------------------------------------------
Expand Down Expand Up @@ -295,6 +306,9 @@ def jackhmmer(
builder (`~pyhmmer.plan7.Builder`, optional): A builder to configure
how the queries are converted to HMMs. Passing `None` will create
a default instance.
backend (`str`): The parallel backend to use for workers to be
executed. Supports ``threading`` to use thread-based parallelism,
or ``multiprocessing`` to use process-based parallelism.
Yields:
`~pyhmmer.plan7.IterationResult`: An *iteration result* instance for
Expand Down

0 comments on commit 92babdf

Please sign in to comment.