Repository: wearepal/EthicML

Issue:
Replace ray with joblib.Parallel
`ray` is complete overkill for just running things in parallel. [joblib.Parallel](https://joblib.readthedocs.io/en/latest/parallel.html) will do just fine. (The only thing that joblib can't do is manage access to GPUs, but I don't think we need that.)


Discussion:

Resolving patch:
diff --git a/ethicml/evaluators/evaluate_models.py b/ethicml/evaluators/evaluate_models.py
--- a/ethicml/evaluators/evaluate_models.py
+++ b/ethicml/evaluators/evaluate_models.py
@@ -153,7 +153,7 @@ def evaluate_models(
     splitter: Optional[DataSplitter] = None,
     topic: Optional[str] = None,
     fair_pipeline: bool = True,
-    num_cpus: Optional[int] = None,
+    num_jobs: Optional[int] = None,
     scaler: Optional[ScalerType] = None,
 ) -> Results:
     """Evaluate all the given models for all the given datasets and compute all the given metrics.
@@ -171,7 +171,7 @@ def evaluate_models(
         splitter: (optional) custom train-test splitter
         topic: (optional) a string that identifies the run; the string is prepended to the filename
         fair_pipeline: if True, run fair inprocess algorithms on the output of preprocessing
-        num_cpus: number of CPUs to use; if None, use all
+        num_jobs: number of parallel jobs; if None, the number of CPUs is used
     """
     from .parallelism import run_in_parallel
 
@@ -222,7 +222,7 @@ def evaluate_models(
         all_results.append_from_csv(csv_file)
 
     # ============================= inprocess models on untransformed =============================
-    all_predictions = run_in_parallel(inprocess_models, data_splits, num_cpus)
+    all_predictions = run_in_parallel(inprocess_models, data_splits, num_jobs)
     inprocess_untransformed = _gather_metrics(
         all_predictions, test_data, inprocess_models, metrics, per_sens_metrics, outdir, topic
     )
@@ -230,7 +230,7 @@ def evaluate_models(
 
     # ===================================== preprocess models =====================================
     # run all preprocess models
-    all_transformed = run_in_parallel(preprocess_models, data_splits, num_cpus)
+    all_transformed = run_in_parallel(preprocess_models, data_splits, num_jobs)
 
     # append the transformed data to `transformed_data`
     transformed_data: List[TrainTestPair] = []
@@ -255,7 +255,7 @@ def evaluate_models(
         # if not fair pipeline, run only the non-fair models on the transformed data
         run_on_transformed = [model for model in inprocess_models if not model.is_fairness_algo]
 
-    transf_preds = run_in_parallel(run_on_transformed, transformed_data, num_cpus)
+    transf_preds = run_in_parallel(run_on_transformed, transformed_data, num_jobs)
     transf_results = _gather_metrics(
         transf_preds, transformed_test, run_on_transformed, metrics, per_sens_metrics, outdir, topic
     )
diff --git a/ethicml/evaluators/parallelism.py b/ethicml/evaluators/parallelism.py
--- a/ethicml/evaluators/parallelism.py
+++ b/ethicml/evaluators/parallelism.py
@@ -2,15 +2,12 @@
 from typing import List, Optional, Sequence, Tuple, TypeVar, Union, cast, overload
 from typing_extensions import Protocol
 
-import ray
-from ray.actor import ActorHandle
+from joblib import Parallel, delayed
 
 from ethicml.algorithms.inprocess.in_algorithm import InAlgorithm
 from ethicml.algorithms.preprocess.pre_algorithm import PreAlgorithm
 from ethicml.utility import DataTuple, Prediction, TestTuple, TrainTestPair
 
-from .progressbar import ProgressBar
-
 __all__ = ["arrange_in_parallel", "run_in_parallel"]
 
 
@@ -22,25 +19,24 @@
 
 
 @overload
-def run_in_parallel(algos: InSeq, data: DataSeq, num_cpus: Optional[int] = None) -> InResult:
+def run_in_parallel(algos: InSeq, data: DataSeq, num_jobs: Optional[int] = None) -> InResult:
     ...
 
 
 @overload
-def run_in_parallel(algos: PreSeq, data: DataSeq, num_cpus: Optional[int] = None) -> PreResult:
+def run_in_parallel(algos: PreSeq, data: DataSeq, num_jobs: Optional[int] = None) -> PreResult:
     ...
 
 
 def run_in_parallel(
-    algos: Union[InSeq, PreSeq], data: DataSeq, num_cpus: Optional[int] = None
+    algos: Union[InSeq, PreSeq], data: DataSeq, num_jobs: Optional[int] = None
 ) -> Union[InResult, PreResult]:
     """Run the given algorithms (embarrassingly) parallel.
 
     Args:
         algos: list of algorithms
         data: list of pairs of data tuples (train and test)
-        num_cpus: how many processes can run in parallel at most. if None, then
-            there is no maximum
+        num_jobs: how many jobs can run in parallel at most. if None, use number of CPUs
 
     Returns:
         list of the results
@@ -51,10 +47,10 @@ def run_in_parallel(
     # but that's completely fine because this check is only here for mypy anyway.
     if isinstance(algos[0], InAlgorithm):
         in_algos = cast(Sequence[InAlgorithm], algos)
-        return arrange_in_parallel(algos=in_algos, data=data, num_cpus=num_cpus)
+        return arrange_in_parallel(algos=in_algos, data=data, num_jobs=num_jobs)
     else:
         pre_algos = cast(Sequence[PreAlgorithm], algos)
-        return arrange_in_parallel(algos=pre_algos, data=data, num_cpus=num_cpus)
+        return arrange_in_parallel(algos=pre_algos, data=data, num_jobs=num_jobs)
 
 
 _RT = TypeVar("_RT", Prediction, Tuple[DataTuple, TestTuple], covariant=True)  # the return type
@@ -68,45 +64,34 @@ def run(self, train: DataTuple, test: TestTuple) -> _RT:
 
 
 def arrange_in_parallel(
-    algos: Sequence[Algorithm[_RT]], data: DataSeq, num_cpus: Optional[int] = None
+    algos: Sequence[Algorithm[_RT]], data: DataSeq, num_jobs: Optional[int] = None
 ) -> List[List[_RT]]:
     """Arrange the given algorithms to run (embarrassingly) parallel.
 
     Args:
         algos: list of tuples consisting of a `run_async` function of an algorithm and a name
         data: list of pairs of data tuples (train and test)
-        num_cpus: number of CPUs to use. `None` means all.
+        num_jobs: number of parallel jobs. `None` means as many as available CPUs.
 
     Returns:
         list of the results
     """
-    ray.shutdown()  # (not sure why this is needed but it seems to be needed...)
-    ray.init(num_cpus=num_cpus)
+    runner = Parallel(n_jobs=num_jobs, verbose=10, backend="loky")
     assert len(algos) >= 1
     assert len(data) >= 1
     assert isinstance(data[0][0], DataTuple)
     assert isinstance(data[0][1], TestTuple)
     # ================================== create queue of tasks ====================================
-    num_tasks = len(algos) * len(data)
-    pbar = ProgressBar(total=num_tasks)
-    futures: List[_RT] = []
     # for each algorithm, first loop over all available datasets and then go on to the next algo
-    for algo in algos:
-        futures.extend(_run.remote(algo, data_item, pbar.actor) for data_item in data)
-    pbar.print_until_done()
-    # actually run everything
-    results = ray.get(futures)
-    ray.shutdown()
+    results = runner(_run(algo, data_item) for algo in algos for data_item in data)
     # return [[result_dict[(i, j)] for j in range(len(data))] for i in range(len(algos))]
     # we have to reconstruct the nested list from the flattened list
     return [[results[i * len(data) + j] for j in range(len(data))] for i in range(len(algos))]
 
 
-@ray.remote
-def _run(algo: Algorithm[_RT], train_test_pair: TrainTestPair, pbar: ActorHandle) -> _RT:
+@delayed
+def _run(algo: Algorithm[_RT], train_test_pair: TrainTestPair) -> _RT:
     train, test = train_test_pair
     # do the work
     result: _RT = algo.run(train, test)
-    # put result into results dictionary
-    pbar.update.remote(1)
     return result
diff --git a/ethicml/evaluators/progressbar.py b/ethicml/evaluators/progressbar.py
deleted file mode 100644
--- a/ethicml/evaluators/progressbar.py
+++ /dev/null
@@ -1,119 +0,0 @@
-"""Progress Bar for Ray Actors.
-
-Tracking progress of distributed tasks can be tricky.
-
-This script will demonstrate how to implement a simple
-progress bar for a Ray actor to track progress across various
-different distributed components.
-
-Original source: `Link <https://github.com/votingworks/arlo-e2e>`_
-
-Setup: Dependencies
--------------------
-
-First, import some dependencies.
-"""
-
-# Inspiration: https://github.com/honnibal/spacy-ray/pull/
-# 1/files#diff-7ede881ddc3e8456b320afb958362b2aR12-R45
-from asyncio import Event
-from typing import Tuple
-
-import ray
-from ray.actor import ActorHandle
-from tqdm import tqdm
-
-__all__ = ["ProgressBar"]
-
-############################################################
-# This is the Ray "actor" that can be called from anywhere to update
-# our progress. You'll be using the `update` method. Don't
-# instantiate this class yourself. Instead,
-# it's something that you'll get from a `ProgressBar`.
-
-
-@ray.remote
-class ProgressBarActor:
-    counter: int
-    delta: int
-    event: Event
-
-    def __init__(self) -> None:
-        self.counter = 0
-        self.delta = 0
-        self.event = Event()
-
-    def update(self, num_items_completed: int) -> None:
-        """Updates the ProgressBar with the incremental number of items that were just completed."""
-        self.counter += num_items_completed
-        self.delta += num_items_completed
-        self.event.set()
-
-    async def wait_for_update(self) -> Tuple[int, int]:
-        """Blocking call.
-
-        Waits until somebody calls `update`, then returns a tuple of
-        the number of updates since the last call to
-        `wait_for_update`, and the total number of completed items.
-        """
-        await self.event.wait()
-        self.event.clear()
-        saved_delta = self.delta
-        self.delta = 0
-        return saved_delta, self.counter
-
-    def get_counter(self) -> int:
-        """Returns the total number of complete items."""
-        return self.counter
-
-
-######################################################################
-# This is where the progress bar starts. You create one of these
-# on the head node, passing in the expected total number of items,
-# and an optional string description.
-# Pass along the `actor` reference to any remote task,
-# and if they complete ten
-# tasks, they'll call `actor.update.remote(10)`.
-
-# Back on the local node, once you launch your remote Ray tasks, call
-# `print_until_done`, which will feed everything back into a `tqdm` counter.
-
-
-class ProgressBar:
-    """Progress bar for ray."""
-
-    progress_actor: ActorHandle
-    total: int
-    description: str
-    pbar: tqdm
-
-    def __init__(self, total: int, description: str = ""):
-        # Ray actors don't seem to play nice with mypy, generating
-        # a spurious warning for the following line,
-        # which we need to suppress. The code is fine.
-        self.progress_actor = ProgressBarActor.remote()  # type: ignore
-        self.total = total
-        self.description = description
-
-    @property
-    def actor(self) -> ActorHandle:
-        """Returns a reference to the remote `ProgressBarActor`.
-
-        When you complete tasks, call `update` on the actor.
-        """
-        return self.progress_actor
-
-    def print_until_done(self) -> None:
-        """Blocking call.
-
-        Do this after starting a series of remote Ray tasks, to which you've
-        passed the actor handle. Each of them calls `update` on the actor.
-        When the progress meter reaches 100%, this method returns.
-        """
-        pbar = tqdm(desc=self.description, total=self.total)
-        while True:
-            delta, counter = ray.get(self.actor.wait_for_update.remote())
-            pbar.update(delta)
-            if counter >= self.total:
-                pbar.close()
-                return


Test patch:
diff --git a/tests/run_algorithm_test.py b/tests/run_algorithm_test.py
--- a/tests/run_algorithm_test.py
+++ b/tests/run_algorithm_test.py
@@ -28,7 +28,7 @@ def test_run_parallel(toy_train_test: em.TrainTestPair):
     result = parallelism.run_in_parallel(
         [em.LR(), em.SVM(), em.Majority()],
         [em.TrainTestPair(*data0), em.TrainTestPair(*data1)],
-        num_cpus=2,
+        num_jobs=2,
     )
     # LR
     assert np.count_nonzero(result[0][0].hard.values == 1) == 44