[Datasets] Add logical operator for randomize_block_order() #31977

c21 · 2023-01-26T22:51:02Z

Signed-off-by: Cheng Su scnju13@gmail.com

Why are these changes needed?

This PR adds logical operator for randomize_block_order(). The change includes:

Introduce AbstractAllToAll for all logical operators converted to AllToAllOperator
RandomizeBlocks logical operator for randomize_block_order().
_internal/planner to move logic for Planner here and have generated function for randomize_blocks. This can be used later to create MapOperator/AllToAllOperator.

Related issue number

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

c21 · 2023-01-26T22:53:41Z

python/ray/data/_internal/logical/operators/map_operator.py

@@ -62,6 +45,24 @@ def __init__(
        fn_constructor_kwargs: Optional[Dict[str, Any]] = None,
        ray_remote_args: Optional[Dict[str, Any]] = None,
    ):
+        """


this is just moving around, no actual change for words.

c21 · 2023-01-26T23:02:08Z

python/ray/data/_internal/logical/operators/all_to_all_operator.py

+    op: AbstractAllToAll,
+    input_physical_dag: PhysicalOperator
+) -> AllToAllOperator:
+    """Get the corresponding physical operators DAG for AbstractAllToAll operators."""


This method is mostly copied from legacy_compat.py:_stage_to_operator()

ericl · 2023-01-27T01:12:12Z

python/ray/data/_internal/logical/operators/all_to_all_operator.py

+        fn: Callable[[BlockList, bool, Callable, dict], Tuple[BlockList, dict]],
+        num_blocks: Optional[int] = None,
+        supports_block_udf: bool = False,
+        block_udf: Optional[BlockTransform] = None,


Should we remove supports_block_udf / block_udf from the logical plan (afaik it's only used to represent fused stages, and we can instead represent it as Map->AllToAll)?

I think the planner can generate the code implementing the block udf fusion.

@ericl - yeah removed.

ericl · 2023-01-27T01:15:20Z

python/ray/data/_internal/logical/operators/all_to_all_operator.py

+        self,
+        name: str,
+        input_op: LogicalOperator,
+        fn: Callable[[BlockList, bool, Callable, dict], Tuple[BlockList, dict]],


We shouldn't reference the legacy BlockList abstractions in the new packages.

@ericl - removed.

ericl · 2023-01-27T01:38:51Z

python/ray/data/_internal/logical/operators/all_to_all_operator.py

+from ray.data._internal.stats import StatsDict
+
+
+class AbstractAllToAll(LogicalOperator):


Should logical operators not specify a fn at all? That should be generated during planning time.

@ericl - yeah agreed. Created a separate planner directory to put the generated functions for each operator.

jianoaix · 2023-01-27T01:59:49Z

python/ray/data/_internal/logical/operators/all_to_all_operator.py

+
+
+def plan_all_to_all_op(
+    op: AbstractAllToAll, input_physical_dag: PhysicalOperator


Can we make it a bit more clear (by naming and/or comments) that this is translating a single operator node (instead of entire dag)?

@jianoaix - sure, updated.

jianoaix · 2023-01-27T18:54:17Z

python/ray/data/_internal/logical/operators/map_operator.py

+                inspecting the logical plan of a Dataset.
+            input_op: The operator preceding this operator in the plan DAG. The outputs
+                of `input_op` will be the inputs to this operator.
+            block_fn: The transform function to apply to each input block to produce


Shall we name this as transform_fn? It's consistent (and IMO more clear) with the naming in new backend.

yeah, actually after second thought, this argument should not be part of logical operator. Logical operator should only care about user-defined function. This is internal function we generate for execution. Plan to remove it and put into planner directory in a separate PR.

c21 · 2023-01-27T22:06:15Z

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

+from ray.data.block import Block, CallableClass
+
+
+def _plan_map_op(op: AbstractMap, input_physical_dag: PhysicalOperator) -> MapOperator:


This is moved from map_operator.py.

c21 · 2023-01-27T22:06:30Z

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

+from ray.data.datasource.datasource import ReadTask
+
+
+def _plan_read_op(op: Read) -> PhysicalOperator:


This is moved from read_operator.py

c21 · 2023-01-27T22:08:07Z

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

+) -> Callable[[List[RefBundle]], Tuple[List[RefBundle], StatsDict]]:
+    """Generate function to randomize order of blocks."""
+
+    def fn(refs: List[RefBundle]) -> Tuple[List[RefBundle], StatsDict]:


This function is refactored from
https://github.com/ray-project/ray/blob/master/python/ray/data/_internal/stage_impl.py#L80 and
https://github.com/ray-project/ray/blob/master/python/ray/data/_internal/block_list.py#L228

ericl · 2023-01-27T22:12:01Z

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

+    See Planner.plan() for more details.
+    """
+    compute = get_compute(op._compute)
+    block_fn = op._block_fn


Should we also remove block_fn from the op?

@ericl - yes, plan to do in a separate PR - #31977 (comment) .

Signed-off-by: Cheng Su <scnju13@gmail.com>

…ect#31977) This PR adds logical operator for randomize_block_order(). The change includes: Introduce AbstractAllToAll for all logical operators converted to AllToAllOperator RandomizeBlocks logical operator for randomize_block_order(). _internal/planner to move logic for Planner here and have generated function for randomize_blocks. This can be used later to create MapOperator/AllToAllOperator. Signed-off-by: Edward Oakes <ed.nmi.oakes@gmail.com>

c21 requested review from ericl, scv119, clarkzinzow, jjyao and jianoaix as code owners January 26, 2023 22:51

c21 assigned ericl, clarkzinzow and jianoaix Jan 26, 2023

c21 commented Jan 26, 2023

View reviewed changes

ericl reviewed Jan 27, 2023

View reviewed changes

ericl added the @author-action-required The PR author is responsible for the next step. Remove tag to send back to the reviewer. label Jan 27, 2023

ericl reviewed Jan 27, 2023

View reviewed changes

jianoaix reviewed Jan 27, 2023

View reviewed changes

c21 changed the title ~~[Datasets] Add logical operator for randomize_block_order() and random_shuffle()~~ [Datasets] Add logical operator for randomize_block_order() Jan 27, 2023

jianoaix reviewed Jan 27, 2023

View reviewed changes

c21 force-pushed the randomize-block branch 2 times, most recently from c669ff6 to 4ca2252 Compare January 27, 2023 19:56

c21 commented Jan 27, 2023

View reviewed changes

ericl approved these changes Jan 27, 2023

View reviewed changes

c21 removed the @author-action-required The PR author is responsible for the next step. Remove tag to send back to the reviewer. label Jan 27, 2023

jianoaix approved these changes Jan 28, 2023

View reviewed changes

c21 added 6 commits January 27, 2023 16:48

Add logical operator for randomize_block_order and random_shuffle

8c56630

Signed-off-by: Cheng Su <scnju13@gmail.com>

Remove unused import

90e9a43

Signed-off-by: Cheng Su <scnju13@gmail.com>

Fix lint

bc92695

Signed-off-by: Cheng Su <scnju13@gmail.com>

Address comments, and revert change for random_shuffle

82c264e

Signed-off-by: Cheng Su <scnju13@gmail.com>

Fix lint

811723f

Signed-off-by: Cheng Su <scnju13@gmail.com>

Create separate planner package

fcc8557

Signed-off-by: Cheng Su <scnju13@gmail.com>

Fix unit tests

f98f2f7

Signed-off-by: Cheng Su <scnju13@gmail.com>

c21 force-pushed the randomize-block branch from 2296b70 to f98f2f7 Compare January 28, 2023 00:50

ericl merged commit e44a7d0 into ray-project:master Jan 28, 2023

c21 deleted the randomize-block branch January 28, 2023 01:54

c21 mentioned this pull request Jan 28, 2023

[Datasets] Refactor map-like functions into planner package #32021

Merged

7 tasks

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[Datasets] Add logical operator for randomize_block_order() #31977

[Datasets] Add logical operator for randomize_block_order() #31977

c21 commented Jan 26, 2023 •

edited

Loading

c21 Jan 26, 2023

c21 Jan 26, 2023

ericl Jan 27, 2023

c21 Jan 27, 2023

ericl Jan 27, 2023

c21 Jan 27, 2023

ericl Jan 27, 2023

c21 Jan 27, 2023 •

edited

Loading

jianoaix Jan 27, 2023

c21 Jan 27, 2023

jianoaix Jan 27, 2023

c21 Jan 27, 2023

c21 Jan 27, 2023

c21 Jan 27, 2023

c21 Jan 27, 2023

ericl Jan 27, 2023

c21 Jan 27, 2023

		from ray.data._internal.stats import StatsDict


		class AbstractAllToAll(LogicalOperator):



		def plan_all_to_all_op(
		op: AbstractAllToAll, input_physical_dag: PhysicalOperator

		from ray.data.block import Block, CallableClass


		def _plan_map_op(op: AbstractMap, input_physical_dag: PhysicalOperator) -> MapOperator:

		from ray.data.datasource.datasource import ReadTask


		def _plan_read_op(op: Read) -> PhysicalOperator:

[Datasets] Add logical operator for randomize_block_order() #31977

[Datasets] Add logical operator for randomize_block_order() #31977

Conversation

c21 commented Jan 26, 2023 • edited Loading

Why are these changes needed?

Related issue number

Checks

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

c21 Jan 27, 2023 • edited Loading

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

c21 commented Jan 26, 2023 •

edited

Loading

c21 Jan 27, 2023 •

edited

Loading