[Data] Support async callable classes in `map_batches()` #46129

scottjlee · 2024-06-18T20:09:02Z

Why are these changes needed?

Add support for passing CallableClass with asynchronous generator __call__ method to Dataset.map_batches() API. This is useful for streaming outputs from asynchronous generators as they become available to maximize throughput.

Related issue number

Closes #46235

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 :(

Signed-off-by: Scott Lee <[email protected]>

raulchen

Just realized that we also need to set max_concurrency to allow handling multiple batches a time.
Some additional changes needed for this PR:

the event loop and the thread executor should outlive the transform_fn. they should be global singletons.

raulchen · 2024-06-18T22:49:07Z

python/ray/data/tests/test_map.py

+            tasks = [sleep_and_yield(i) for i in batch["id"]]
+            results = await asyncio.gather(*tasks)
+            for result in results:
+                yield result


Avoid using gather, so we can yield results as soon as they become available.

tasks = [asyncio.create_task(sleep_and_yield(i)) for i in batch["id"]] for task in tasks: yield await task

raulchen · 2024-06-18T23:25:15Z

python/ray/data/_internal/planner/plan_udf_map_op.py

+
+            async def process_all_batches():
+                tasks = [asyncio.create_task(process_batch(x)) for x in input_iterable]
+                for task in asyncio.as_completed(tasks):


as_completed doesn't preserve the order. Shouldn't use it when preserve_order is set.

Signed-off-by: Scott Lee <[email protected]>

raulchen · 2024-06-21T18:21:36Z

python/ray/data/tests/test_map.py

+            pass
+
+        async def __call__(self, batch):
+            tasks = [sleep_and_yield(i) for i in batch["id"]]


Remove this line.

raulchen · 2024-06-21T18:33:19Z

python/ray/data/_internal/planner/plan_udf_map_op.py

+            # Use the existing event loop to create and run
+            # Tasks to process each batch
+            loop = ray.data._cached_loop
+            loop.run_until_complete(process_all_batches())


This is still running in the main thread. Thus cannot run multiple batches at the same time.
We should put loop.run_forever() in a thread, and then call loop.call_soon_threadsafe here.

raulchen · 2024-06-21T18:54:07Z

python/ray/data/_internal/planner/plan_udf_map_op.py

+                    for task in asyncio.as_completed(tasks):
+                        await task
+                # Sentinel to indicate completion.
+                output_batch_queue.put(None)


use a special object, in case the UDF also returns None

raulchen · 2024-06-21T18:55:25Z

python/ray/data/_internal/planner/plan_udf_map_op.py

-                    res = [batch]
+    if inspect.iscoroutinefunction(fn):
+        # UDF is a callable class with async generator `__call__` method.
+        def transform_fn(


This inline function is too long. let's define it as a util function.

raulchen · 2024-06-21T18:57:28Z

python/ray/data/_internal/planner/plan_udf_map_op.py

@@ -104,22 +107,45 @@ def _parse_op_fn(op: AbstractUDFMap):
        fn_constructor_args = op._fn_constructor_args or ()
        fn_constructor_kwargs = op._fn_constructor_kwargs or {}

-        op_fn = make_callable_class_concurrent(op_fn)
+        if inspect.isasyncgenfunction(op._fn.__call__):


Add a TODO that we should 1) support non-generator async functions; 2) make the entire map actor async.

raulchen · 2024-06-21T18:58:18Z

python/ray/data/tests/test_map.py

+
+    n = 5
+    ds = ray.data.range(n, override_num_blocks=1)
+    ds = ds.map_batches(AsyncActor, batch_size=None, concurrency=1)


Add max_concurrency to test concurrently handling many batches.

raulchen · 2024-06-21T22:38:49Z

python/ray/data/_internal/planner/plan_udf_map_op.py

+            async def fn(item: Any) -> Any:
+                assert ray.data._cached_fn is not None
+                assert ray.data._cached_cls == op_fn
+                assert ray.data._cached_loop is not None


to make code cleaner, let's define a class (e.g. MapActorContext) to capture all these variables

Signed-off-by: Scott Lee <[email protected]>

raulchen

LGTM except for some nits

raulchen · 2024-06-24T17:26:12Z

python/ray/data/_internal/planner/plan_udf_map_op.py

+                        loop.run_forever()
+
+                    thread = Thread(target=run_loop)
+                    thread.start()


nit, maybe move the above to MapActorContext to make the code cleaner here.
We can pass in a boolean flag here _MapActorContext(..., is_asyncio=True).
Then the code for sync and async branches can be consolidated.

raulchen · 2024-06-24T17:26:44Z

python/ray/data/_internal/planner/plan_udf_map_op.py

+        cached_cls: UserDefinedFunction,
+        cached_fn: Callable[[Any], Any],
+        cached_loop: Optional[asyncio.AbstractEventLoop] = None,
+        cached_asyncio_thread: Optional[Thread] = None,


remove the cached_ prefixes?

raulchen · 2024-06-24T22:37:03Z

python/ray/data/_internal/planner/plan_udf_map_op.py

+                for task in asyncio.as_completed(tasks):
+                    await task
+            # Sentinel to indicate completion.
+            output_batch_queue.put(OutputQueueSentinel())


OutputQueueSentinel isn't needed now because of while not future.done()

raulchen · 2024-06-24T22:38:31Z

python/ray/data/tests/test_map.py

+            pass
+
+        async def __call__(self, batch):
+            tasks = [sleep_and_yield(i) for i in batch["id"]]


this line is redundant.

thanks for the catch, i had thought i removed it but somehow ended up back again...

raulchen · 2024-06-25T21:46:19Z

python/ray/data/_internal/planner/plan_udf_map_op.py

+
+        else:
+
+            def fn(item: Any) -> Any:


these 2 fns are almost identical, except the assertion? better avoid duplicating the code.

ah that's what i initially thought, but one of them is async, and the other is not. is there a way to combine the inner implementation but create an async version?

oh, nvm, I didn't notice the async prefix

tests

599be0b

Signed-off-by: Scott Lee <[email protected]>

scottjlee marked this pull request as ready for review June 18, 2024 20:51

scottjlee requested review from ericl, scv119, c21, amogkam, bveeramani, raulchen, stephanie-wang and omatthew98 as code owners June 18, 2024 20:51

scottjlee assigned raulchen and c21 Jun 18, 2024

raulchen reviewed Jun 18, 2024

View reviewed changes

scottjlee added 5 commits June 20, 2024 12:58

Merge branch 'master' into 0610-async-map

bdd812e

Signed-off-by: Scott Lee <[email protected]>

initialize loop in init_fn

dec32a1

Signed-off-by: Scott Lee <[email protected]>

Merge branch 'master' into 0610-async-map

13b980f

Signed-off-by: Scott Lee <[email protected]>

import

4292b93

Signed-off-by: Scott Lee <[email protected]>

avoid using asyncio.gather in test

de61a0d

Signed-off-by: Scott Lee <[email protected]>

raulchen reviewed Jun 21, 2024

View reviewed changes

scottjlee added 2 commits June 21, 2024 21:07

address comments

d217c1c

Signed-off-by: Scott Lee <[email protected]>

add _map_actor_context

9175edd

Signed-off-by: Scott Lee <[email protected]>

scottjlee assigned scottjlee and unassigned raulchen and c21 Jun 24, 2024

scottjlee added the data Ray Data-related issues label Jun 24, 2024

scottjlee assigned raulchen and unassigned scottjlee Jun 24, 2024

raulchen approved these changes Jun 24, 2024

View reviewed changes

comments

c5af4ae

scottjlee requested a review from raulchen June 25, 2024 21:35

raulchen approved these changes Jun 25, 2024

View reviewed changes

raulchen enabled auto-merge (squash) June 25, 2024 22:00

github-actions bot added the go add ONLY when ready to merge, run all tests label Jun 25, 2024

raulchen merged commit f75ad5d into ray-project:master Jun 25, 2024
7 of 8 checks passed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[Data] Support async callable classes in `map_batches()` #46129

[Data] Support async callable classes in `map_batches()` #46129

scottjlee commented Jun 18, 2024 •

edited

Loading

raulchen left a comment

raulchen Jun 18, 2024

raulchen Jun 18, 2024

raulchen Jun 21, 2024

raulchen Jun 21, 2024

raulchen Jun 21, 2024

raulchen Jun 21, 2024

raulchen Jun 21, 2024

raulchen Jun 21, 2024

raulchen Jun 21, 2024

raulchen left a comment

raulchen Jun 24, 2024

raulchen Jun 24, 2024

raulchen Jun 24, 2024

raulchen Jun 24, 2024

scottjlee Jun 25, 2024

raulchen Jun 25, 2024

scottjlee Jun 25, 2024

raulchen Jun 25, 2024

[Data] Support async callable classes in map_batches() #46129

[Data] Support async callable classes in map_batches() #46129

Conversation

scottjlee commented Jun 18, 2024 • edited Loading

Why are these changes needed?

Related issue number

Checks

raulchen left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

raulchen left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

[Data] Support async callable classes in `map_batches()` #46129

[Data] Support async callable classes in `map_batches()` #46129

scottjlee commented Jun 18, 2024 •

edited

Loading