Fix and improve map reduce

CHANGELOG.md

@@ -14,6 +14,14 @@
 - Fixes bug in Influence calculation with multi-dimensional input and adds
   new example notebook
   [PR #195](https://github.com/appliedAI-Initiative/pyDVL/pull/195)
+- **Breaking change**: Passes the input to `MapReduceJob` at initialization,
+  removes `chunkify_inputs` argument from `MapReduceJob`,
+  removes `n_runs` argument from `MapReduceJob`,
+  calls the parallel backend's `put()` method for each generated chunk in `_chunkify()`,
+  renames ParallelConfig's `num_workers` attribute to `n_local_workers`,
+  fixes a bug in `MapReduceJob`'s chunkification when `n_runs` >= `n_jobs`,
+  and defines a sequential parallel backend to run all jobs in the current thread
+  [PR #232](https://github.com/appliedAI-Initiative/pyDVL/pull/232)
 
 ## 0.3.0 - 💥 Breaking changes
 

notebooks/shapley_basic_spotify.ipynb

@@ -95,12 +95,12 @@
    "source": [
     "from pydvl.reporting.plots import plot_shapley\n",
     "from pydvl.utils import (\n",
-    "    available_cpus,\n",
     "    load_spotify_dataset,\n",
     "    Dataset,\n",
     "    GroupedDataset,\n",
     "    Utility,\n",
     ")\n",
+    "from pydvl.utils.parallel.backend import available_cpus\n",
     "from pydvl.value.shapley import compute_shapley_values"
    ]
   },

notebooks/shapley_utility_learning.ipynb

@@ -149,8 +149,8 @@
     "    Utility,\n",
     "    DataUtilityLearning,\n",
     "    top_k_value_accuracy,\n",
-    "    available_cpus,\n",
     ")\n",
+    "from pydvl.utils.parallel.backend import available_cpus\n",
     "from pydvl.reporting.plots import shaded_mean_std\n",
     "from pydvl.value.shapley import compute_shapley_values"
    ]

src/pydvl/utils/config.py

@@ -1,5 +1,5 @@
 from dataclasses import dataclass, field
-from typing import Iterable, Optional, Tuple, Union
+from typing import Iterable, Literal, Optional, Tuple, Union
 
 from pymemcache.serde import PickleSerde
 
@@ -16,12 +16,12 @@ class ParallelConfig:
 
     :param backend: Type of backend to use. For now only 'ray' is supported.
     :param address: Address of existing remote or local cluster to use.
-    :param num_workers: Number of workers (CPUs) to use.
+    :param n_local_workers: Number of workers (CPUs) to use when using a local ray cluster
     """
 
-    backend: str = "ray"
+    backend: Literal["sequential", "ray"] = "ray"
     address: Optional[Union[str, Tuple[str, int]]] = None
-    num_workers: Optional[int] = None
+    n_local_workers: Optional[int] = None
 
 
 @unpackable

src/pydvl/utils/numeric.py

@@ -17,6 +17,7 @@
 )
 
 import numpy as np
+from scipy.special import expit
 
 from pydvl.utils.types import compose_score
 
@@ -252,7 +253,8 @@ def top_k_value_accuracy(y_true: "NDArray", y_pred: "NDArray", k: int = 3) -> fl
 
 
 def sigmoid(x: float) -> float:
-    return float(1 / (1 + np.exp(-x)))
+    result: float = expit(x).item()
+    return result
 
 
 squashed_r2 = compose_score("r2", sigmoid, "squashed r2")

src/pydvl/utils/parallel/actor.py

@@ -1,16 +1,12 @@
 import abc
 import inspect
 import logging
-from typing import TYPE_CHECKING, Any, Dict, Optional, Union
+from typing import Any, Dict, Optional, Union
 
 from ray import ObjectRef
 
 from .backend import RayParallelBackend
 
-if TYPE_CHECKING:
-    from numpy.typing import NDArray
-
-
 __all__ = ["RayActorWrapper", "Coordinator", "Worker"]
 
 
@@ -25,7 +21,7 @@ class RayActorWrapper:
 
     :Example:
 
-    >>> from pydvl.utils.parallel import init_parallel_backend
+    >>> from pydvl.utils.parallel.backend import RayParallelBackend, init_parallel_backend
     >>> from pydvl.utils.config import ParallelConfig
     >>> from pydvl.utils.parallel.actor import RayActorWrapper
     >>> class Actor:
@@ -35,8 +31,9 @@ class RayActorWrapper:
     ...     def get(self):
     ...         return self.x
     ...
-    >>> config = ParallelConfig()
+    >>> config = ParallelConfig(backend="ray")
     >>> parallel_backend = init_parallel_backend(config)
+    >>> assert isinstance(parallel_backend, RayParallelBackend)
     >>> actor_handle = parallel_backend.wrap(Actor).remote(5)
     >>> parallel_backend.get(actor_handle.get.remote())
     5

src/pydvl/utils/parallel/backend.py

@@ -1,6 +1,8 @@
+import functools
 import os
+from abc import ABCMeta, abstractmethod
 from dataclasses import asdict
-from typing import Any, Iterable, List, Optional, Tuple, TypeVar, Union
+from typing import Any, Dict, Iterable, List, Optional, Tuple, Type, TypeVar, Union
 
 import ray
 from ray import ObjectRef
@@ -10,86 +12,182 @@
 
 __all__ = [
     "init_parallel_backend",
-    "available_cpus",
 ]
 
 T = TypeVar("T")
 
-_PARALLEL_BACKEND: Optional["RayParallelBackend"] = None
+_PARALLEL_BACKENDS: Dict[str, "Type[BaseParallelBackend]"] = {}
 
 
-class RayParallelBackend:
-    """Class used to wrap ray to make it transparent to algorithms. It shouldn't
+class NoPublicConstructor(ABCMeta):
+    """Metaclass that ensures a private constructor
+
+    If a class uses this metaclass like this:
+
+        class SomeClass(metaclass=NoPublicConstructor):
+            pass
+
+    If you try to instantiate your class (`SomeClass()`),
+    a `TypeError` will be thrown.
+
+    Taken almost verbatim from:
+    https://stackoverflow.com/a/64682734
+    """
+
+    def __call__(cls, *args, **kwargs):
+        raise TypeError(
+            f"{cls.__module__}.{cls.__qualname__} cannot be initialized directly. "
+            "Use init_parallel_backend() instead."
+        )
+
+    def _create(cls, *args: Any, **kwargs: Any):
+        return super().__call__(*args, **kwargs)
+
+
+class BaseParallelBackend(metaclass=NoPublicConstructor):
+    """Abstract base class for all parallel backends"""
+
+    config: Dict[str, Any] = {}
+
+    def __init_subclass__(cls, *, backend_name: str, **kwargs):
+        global _PARALLEL_BACKENDS
+        _PARALLEL_BACKENDS[backend_name] = cls
+        super().__init_subclass__(**kwargs)
+
+    @abstractmethod
+    def get(self, v: Any, *args, **kwargs):
+        ...
+
+    @abstractmethod
+    def put(self, v: Any, *args, **kwargs) -> Any:
+        ...
+
+    @abstractmethod
+    def wrap(self, *args, **kwargs) -> Any:
+        ...
+
+    @abstractmethod
+    def wait(self, v: Any, *args, **kwargs) -> Any:
+        ...
+
+    @abstractmethod
+    def _effective_n_jobs(self, n_jobs: int) -> int:
+        ...
+
+    def effective_n_jobs(self, n_jobs: int = -1) -> int:
+        if n_jobs == 0:
+            raise ValueError("n_jobs == 0 in Parallel has no meaning")
+        n_jobs = self._effective_n_jobs(n_jobs)
+        return n_jobs
+
+    def __repr__(self) -> str:
+        return f"<{self.__class__.__name__}: {self.config}>"
+
+
+class SequentialParallelBackend(BaseParallelBackend, backend_name="sequential"):
+    """Class used to run jobs sequentially and locally. It shouldn't
     be initialized directly. You should instead call `init_parallel_backend`.
 
-    :param config: instance of :class:`~pydvl.utils.config.ParallelConfig` with
-        cluster address, number of cpus, etc.
+    :param config: instance of :class:`~pydvl.utils.config.ParallelConfig` with number of cpus
+    """
 
-    :Example:
+    def __init__(self, config: ParallelConfig):
+        config_dict = asdict(config)
+        config_dict.pop("backend")
+        config_dict.pop("address")
+        config_dict["num_cpus"] = config_dict.pop("n_local_workers")
+        self.config = config_dict
 
-    >>> from pydvl.utils.parallel.backend import RayParallelBackend
-    >>> from pydvl.utils.config import ParallelConfig
-    >>> config = ParallelConfig(backend="ray")
-    >>> parallel_backend = RayParallelBackend(config)
-    >>> parallel_backend
-    <RayParallelBackend: {'address': None, 'num_cpus': None}>
+    def get(self, v: Any, *args, **kwargs):
+        return v
+
+    def put(self, v: Any, *args, **kwargs) -> Any:
+        return v
+
+    def wrap(self, *args, **kwargs) -> Any:
+        assert len(args) == 1
+        return functools.partial(args[0], **kwargs)
 
+    def wait(self, v: Any, *args, **kwargs) -> Tuple[list, list]:
+        return v, []
+
+    def _effective_n_jobs(self, n_jobs: int) -> int:
+        if n_jobs < 0:
+            if self.config["num_cpus"]:
+                eff_n_jobs: int = self.config["num_cpus"]
+            else:
+                eff_n_jobs = available_cpus()
+        else:
+            eff_n_jobs = n_jobs
+        return eff_n_jobs
+
+
+class RayParallelBackend(BaseParallelBackend, backend_name="ray"):
+    """Class used to wrap ray to make it transparent to algorithms. It shouldn't
+    be initialized directly. You should instead call `init_parallel_backend`.
+
+    :param config: instance of :class:`~pydvl.utils.config.ParallelConfig` with
+        cluster address, number of cpus, etc.
     """
 
     def __init__(self, config: ParallelConfig):
         config_dict = asdict(config)
         config_dict.pop("backend")
-        config_dict["num_cpus"] = config_dict.pop("num_workers")
+        config_dict["num_cpus"] = config_dict.pop("n_local_workers")
         self.config = config_dict
+        if self.config["address"] is None:
+            self.config["ignore_reinit_error"] = True
         ray.init(**self.config)
 
     def get(
         self,
         v: Union[ObjectRef, Iterable[ObjectRef], T],
-        *,
-        timeout: Optional[float] = None,
+        *args,
+        **kwargs,
     ) -> Union[T, Any]:
+        timeout: Optional[float] = kwargs.get("timeout", None)
         if isinstance(v, ObjectRef):
             return ray.get(v, timeout=timeout)
         elif isinstance(v, Iterable):
             return [self.get(x, timeout=timeout) for x in v]
         else:
             return v
 
-    def put(self, x: Any, **kwargs) -> ObjectRef:
-        return ray.put(x, **kwargs)  # type: ignore
+    def put(self, v: T, *args, **kwargs) -> Union["ObjectRef[T]", T]:
+        try:
+            return ray.put(v, **kwargs)  # type: ignore
+        except TypeError:
+            return v  # type: ignore
 
     def wrap(self, *args, **kwargs) -> RemoteFunction:
         return ray.remote(*args, **kwargs)  # type: ignore
 
     def wait(
         self,
-        object_refs: List["ray.ObjectRef"],
-        *,
-        num_returns: int = 1,
-        timeout: Optional[float] = None,
+        v: List["ObjectRef"],
+        *args,
+        **kwargs,
     ) -> Tuple[List[ObjectRef], List[ObjectRef]]:
+        num_returns: int = kwargs.get("num_returns", 1)
+        timeout: Optional[float] = kwargs.get("timeout", None)
         return ray.wait(  # type: ignore
-            object_refs,
+            v,
             num_returns=num_returns,
             timeout=timeout,
         )
 
-    def effective_n_jobs(self, n_jobs: Optional[int]) -> int:
-        if n_jobs == 0:
-            raise ValueError("n_jobs == 0 in Parallel has no meaning")
-        elif n_jobs is None or n_jobs < 0:
+    def _effective_n_jobs(self, n_jobs: int) -> int:
+        if n_jobs < 0:
             ray_cpus = int(ray._private.state.cluster_resources()["CPU"])  # type: ignore
             eff_n_jobs = ray_cpus
         else:
             eff_n_jobs = n_jobs
         return eff_n_jobs
 
-    def __repr__(self) -> str:
-        return f"<RayParallelBackend: {self.config}>"
-
 
-def init_parallel_backend(config: ParallelConfig) -> "RayParallelBackend":
+def init_parallel_backend(
+    config: ParallelConfig,
+) -> BaseParallelBackend:
     """Initializes the parallel backend and returns an instance of it.
 
     :param config: instance of :class:`~pydvl.utils.config.ParallelConfig` with cluster address, number of cpus, etc.
@@ -101,16 +199,15 @@ def init_parallel_backend(config: ParallelConfig) -> "RayParallelBackend":
     >>> config = ParallelConfig(backend="ray")
     >>> parallel_backend = init_parallel_backend(config)
     >>> parallel_backend
-    <RayParallelBackend: {'address': None, 'num_cpus': None}>
+    <RayParallelBackend: {'address': None, 'num_cpus': None, 'ignore_reinit_error': True}>
 
     """
-    global _PARALLEL_BACKEND
-    if _PARALLEL_BACKEND is None:
-        if config.backend == "ray":
-            _PARALLEL_BACKEND = RayParallelBackend(config)
-        else:
-            raise NotImplementedError(f"Unexpected parallel type {config.backend}")
-    return _PARALLEL_BACKEND
+    try:
+        parallel_backend_cls = _PARALLEL_BACKENDS[config.backend]
+    except KeyError:
+        raise NotImplementedError(f"Unexpected parallel backend {config.backend}")
+    parallel_backend = parallel_backend_cls._create(config)
+    return parallel_backend  # type: ignore
 
 
 def available_cpus() -> int: