[core] add ray.util.concurrent.futures.RayExecutor

[GuillaumeDesforges]

RayExecutor is a drop-in replacement for ProcessPoolExecutor and ThreadPoolExecutor from concurrent.futures but distributes and executes the specified tasks over a pool of dedicated actors belonging to a Ray cluster instead of multiple processes or threads, respectively.

Why are these changes needed?

This change provides an easy way to use Ray for futures, allowing to scale an existing codebase that relies on them.

Checks

• I've signed off every commit(by using the -s flag, i.e., git commit -s) in this PR.
• I've run scripts/format.sh to lint the changes in this PR.
• I've included any doc changes needed for https://docs.ray.io/en/master/.
  • I've added any new APIs to the API Reference. For example, if I added a
    method in Tune, I've added it in doc/source/tune/api/ under the
    corresponding .rst file.
• I've made sure the tests are passing. Note that there might be a few flaky tests, see the recent failures at https://flakey-tests.ray.io/
• Testing Strategy
  • Unit tests
  • Release tests
  • This PR is not tested :(

[aslonnie]

asking core team to review first.

[judahrand]

This would be phenomenally useful! Is this likely to get reviewed soon? Anyway I can do to help?

[judahrand]

@GuillaumeDesforges I made a few fixes in #46249. Feel free to pull them into this PR if you like and then I will close the other one.

[GuillaumeDesforges]

The changes you propose make sense to me, I'll pull them into this PR asap and try to fix tests 👍

[GuillaumeDesforges]

Sorry it took so much time. I've rebased on #46249 to include your changes @judahrand 🙏

I'll try to fix failed tests if I can.

[anyscalesam]

@GuillaumeDesforges any progress on this?

[anyscalesam]

@GuillaumeDesforges any update here?

[GuillaumeDesforges]

I have been short on time, but I do plan to make these tests pass. 👍

EDIT 2024-10-04: still on my TODO 😄

[GuillaumeDesforges]

I have looked into the doc error. I have no clue how to solve them.

For the Python tests, we seem to have relied on an API that is not available in Python 3.9.

[elijahbenizzy]

Hey folks! Any chance this PR is likely to get merged at some point? Would be very useful, and probably better than my (extremely hacky) prototype:

class RayExecutor(concurrent.futures.ThreadPoolExecutor):
    def __init__(self, threadpool_kwargs: Dict[str, Any], ray_kwargs: Dict[str, Any]):
        super(RayExecutor, self).__init__(**threadpool_kwargs)
        self.ray_kwargs = ray_kwargs

    def init(self):
        # TODO - synchronize number of workers
        ray.init(**self.ray_kwargs)

    def submit(self, __fn, *args, **kwargs):
        return ray.remote(__fn).remote(*args, **kwargs).future()
        ```

[GuillaumeDesforges]

@elijahbenizzy happy to see interest in our work 😄

I won't have time anytime soon to try to satisfy the CI unfortunately. Anyone is free to pick this up. 🙂

EDIT: just saw someone was assigned to it. Nice!

[davidwagnerkc]

Not sure about the docs / linter, but for the actual core tests I think could use something like this to make the _result_or_cancel available to Python 3.9.

diff --git a/python/ray/util/concurrent/futures/ray_executor.py b/python/ray/util/concurrent/futures/ray_executor.py
index 237a9f17cd..62e0a4cfee 100644
--- a/python/ray/util/concurrent/futures/ray_executor.py
+++ b/python/ray/util/concurrent/futures/ray_executor.py
@@ -2,7 +2,21 @@ from abc import ABC, abstractmethod
 import time
 from functools import partial
 from concurrent.futures import Executor, Future
-from concurrent.futures._base import _result_or_cancel  # type: ignore
+
+try:
+    from concurrent.futures._base import _result_or_cancel  # type: ignore
+except ImportError:
+    # Backport private Python 3.12 function to Python 3.9
+    # https://github.com/python/cpython/tree/main/Lib/concurrent/futures/_base.py#L306
+    def _result_or_cancel(fut, timeout=None):
+        try:
+            try:
+                return fut.result(timeout)
+            finally:
+                fut.cancel()
+        finally:
+            del fut
+
 from typing import (
     Any,
     ParamSpec,

[davidwagnerkc]

@jjyao @GuillaumeDesforges Not sure if we want to get this merged or not, but #51933 makes CI all green for this PR. Changes in the description.

[davidwagnerkc]

@GuillaumeDesforges

So this pattern of accessing results:

from concurrent.futures import ProcessPoolExecutor
with ProcessPoolExecutor() as ex:
    results = ex.map(str, list(range(25)))
print(list(results))

Will hang for RayExecutor:

from ray.util.concurrent.futures.ray_executor import RayExecutor
with RayExecutor() as ex:
    results = ex.map(str, list(range(25)))
print(list(results))

In fact you quite nicely have the test_results_are_not_accessible_after_shutdown test to test for this fact. But what prevents us from keeping this behavior the same? You should not need anything Ray, any refs, and futures, etc. after the exit of the context manager, so if wait=True on shutdown as it is by default we should not only future.result() to wait on the computations, but also return those results correct?

I can poke at this if we want to change it, but wanted to ask first since I imagine you thought through this given you wrote the test for it.

Cool PR by the way. I have done serious damage with the ProcessPoolExecutor in many different contexts so think it is a great utility API to provide.

[davidwagnerkc]

I made a change to address the above comment - moving discussion to #51933

[codope]

Closing this PR in favor of #51933 as it fixes the CI.