Fix TMCS starts too many processes and dies

CHANGELOG.md

@@ -2,6 +2,10 @@
 
 ## Unreleased
 
+- Create new `RayExecutor` class based on the concurrent.futures API,
+  use the new class to fix an issue with Truncated Monte Carlo Shapley
+  (TMCS) starting too many processes and dying, plus other small changes
+  [PR #329](https://github.com/appliedAI-Initiative/pyDVL/pull/329)
 - Fix creation of GroupedDataset objects using the `from_arrays`
   and `from_sklearn` class methods 
   [PR #324](https://github.com/appliedAI-Initiative/pyDVL/pull/334)

src/pydvl/utils/config.py

@@ -13,17 +13,23 @@
 class ParallelConfig:
     """Configuration for parallel computation backend.
 
-    :param backend: Type of backend to use. For now only 'ray' is supported.
+    :param backend: Type of backend to use.
+        Defaults to 'ray'
     :param address: Address of existing remote or local cluster to use.
-    :param n_local_workers: Number of workers (CPUs) to use when using a local ray cluster.
+    :param n_cpus_local: Number of CPUs to use when creating a local ray cluster.
+        This has no effect when using an existing ray cluster.
     :param logging_level: Logging level for the parallel backend's worker.
     """
 
     backend: Literal["sequential", "ray"] = "ray"
     address: Optional[Union[str, Tuple[str, int]]] = None
-    n_local_workers: Optional[int] = None
+    n_cpus_local: Optional[int] = None
     logging_level: int = logging.WARNING
 
+    def __post_init__(self) -> None:
+        if self.address is not None and self.n_cpus_local is not None:
+            raise ValueError("When `address` is set, `n_cpus_local` should be None.")
+
 
 @dataclass
 class MemcachedClientConfig:

src/pydvl/utils/parallel/__init__.py

@@ -1,3 +1,3 @@
-from .actor import *
 from .backend import *
+from .futures import *
 from .map_reduce import *

src/pydvl/utils/parallel/backend.py

@@ -16,7 +16,6 @@
 
 import ray
 from ray import ObjectRef
-from ray.remote_function import RemoteFunction
 
 from ..config import ParallelConfig
 
@@ -93,18 +92,16 @@ def __repr__(self) -> str:
 
 
 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`.
+    """Class used to run jobs sequentially and locally.
+
+    It shouldn't be initialized directly. You should instead call
+    :func:`~pydvl.utils.parallel.backend.init_parallel_backend`.
 
     :param config: instance of :class:`~pydvl.utils.config.ParallelConfig` with number of cpus
     """
 
     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
+        self.config = {}
 
     def get(self, v: Any, *args, **kwargs):
         return v
@@ -126,8 +123,10 @@ def _effective_n_jobs(self, n_jobs: int) -> int:
 
 
 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`.
+    """Class used to wrap ray to make it transparent to algorithms.
+
+    It shouldn't be initialized directly. You should instead call
+    :func:`~pydvl.utils.parallel.backend.init_parallel_backend`.
 
     :param config: instance of :class:`~pydvl.utils.config.ParallelConfig` with
         cluster address, number of cpus, etc.
@@ -136,7 +135,9 @@ class RayParallelBackend(BaseParallelBackend, backend_name="ray"):
     def __init__(self, config: ParallelConfig):
         config_dict = asdict(config)
         config_dict.pop("backend")
-        config_dict["num_cpus"] = config_dict.pop("n_local_workers")
+        n_cpus_local = config_dict.pop("n_cpus_local")
+        if config_dict.get("address", None) is None:
+            config_dict["num_cpus"] = n_cpus_local
         self.config = config_dict
         if not ray.is_initialized():
             ray.init(**self.config)
@@ -169,7 +170,7 @@ def wrap(self, fun: Callable, **kwargs) -> Callable:
 
         :return: The `.remote` method of the ray `RemoteFunction`.
         """
-        if len(kwargs) > 1:
+        if len(kwargs) > 0:
             return ray.remote(**kwargs)(fun).remote  # type: ignore
         return ray.remote(fun).remote  # type: ignore
 
@@ -201,7 +202,8 @@ def init_parallel_backend(
 ) -> 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.
+    :param config: instance of :class:`~pydvl.utils.config.ParallelConfig`
+        with cluster address, number of cpus, etc.
 
     :Example:
 

src/pydvl/utils/parallel/futures/__init__.py

@@ -0,0 +1,54 @@
+from concurrent.futures import Executor, ThreadPoolExecutor
+from contextlib import contextmanager
+from typing import Generator, Optional
+
+from pydvl.utils.config import ParallelConfig
+from pydvl.utils.parallel.futures.ray import RayExecutor
+
+__all__ = ["init_executor"]
+
+
+@contextmanager
+def init_executor(
+    max_workers: Optional[int] = None,
+    config: ParallelConfig = ParallelConfig(),
+    **kwargs,
+) -> Generator[Executor, None, None]:
+    """Initializes a futures executor based on the passed parallel configuration object.
+
+    :param max_workers: Maximum number of concurrent tasks.
+    :param config: instance of :class:`~pydvl.utils.config.ParallelConfig` with cluster address, number of cpus, etc.
+    :param kwargs: Other optional parameter that will be passed to the executor.
+
+    :Example:
+
+    >>> from pydvl.utils.parallel.futures import init_executor
+    >>> from pydvl.utils.config import ParallelConfig
+    >>> config = ParallelConfig(backend="ray")
+    >>> with init_executor(max_workers=3, config=config) as executor:
+    ...     pass
+
+    >>> from pydvl.utils.parallel.futures import init_executor
+    >>> with init_executor() as executor:
+    ...     future = executor.submit(lambda x: x + 1, 1)
+    ...     result = future.result()
+    ...
+    >>> print(result)
+    2
+
+    >>> from pydvl.utils.parallel.futures import init_executor
+    >>> with init_executor() as executor:
+    ...     results = list(executor.map(lambda x: x + 1, range(5)))
+    ...
+    >>> print(results)
+    [1, 2, 3, 4, 5]
+
+    """
+    if config.backend == "ray":
+        with RayExecutor(max_workers, config=config, **kwargs) as executor:
+            yield executor
+    elif config.backend == "sequential":
+        with ThreadPoolExecutor(1) as executor:
+            yield executor
+    else:
+        raise NotImplementedError(f"Unexpected parallel type {config.backend}")