[Data] Remove run_by_consumer parameter of ExecutionPlan (ray-project#46052)

This PR removes dead functions and parameters after ray-project#46054.

Signed-off-by: Balaji Veeramani <balaji@anyscale.com>

Author: bveeramani

python/ray/data/_internal/execution/legacy_compat.py

@@ -15,7 +15,7 @@
 from ray.data._internal.logical.util import record_operators_usage
 from ray.data._internal.plan import ExecutionPlan
 from ray.data._internal.stats import DatasetStats
-from ray.data.block import Block, BlockMetadata, List
+from ray.data.block import Block, BlockMetadata
 from ray.types import ObjectRef
 
 # Warn about tasks larger than this.
@@ -25,13 +25,9 @@
 def execute_to_legacy_block_iterator(
     executor: Executor,
     plan: ExecutionPlan,
-    allow_clear_input_blocks: bool,
-    dataset_uuid: str,
 ) -> Iterator[Tuple[ObjectRef[Block], BlockMetadata]]:
     """Same as execute_to_legacy_bundle_iterator but returning blocks and metadata."""
-    bundle_iter = execute_to_legacy_bundle_iterator(
-        executor, plan, allow_clear_input_blocks, dataset_uuid
-    )
+    bundle_iter = execute_to_legacy_bundle_iterator(executor, plan)
     for bundle in bundle_iter:
         for block, metadata in bundle.blocks:
             yield block, metadata
@@ -40,17 +36,13 @@ def execute_to_legacy_block_iterator(
 def execute_to_legacy_bundle_iterator(
     executor: Executor,
     plan: ExecutionPlan,
-    allow_clear_input_blocks: bool,
-    dataset_uuid: str,
     dag_rewrite=None,
 ) -> Iterator[RefBundle]:
     """Execute a plan with the new executor and return a bundle iterator.
 
     Args:
         executor: The executor to use.
         plan: The legacy plan to execute.
-        allow_clear_input_blocks: Whether the executor may consider clearing blocks.
-        dataset_uuid: UUID of the dataset for this execution.
         dag_rewrite: Callback that can be used to mutate the DAG prior to execution.
             This is currently used as a legacy hack to inject the OutputSplit operator
             for `Dataset.streaming_split()`.
@@ -73,7 +65,6 @@ def execute_to_legacy_bundle_iterator(
 def execute_to_legacy_block_list(
     executor: Executor,
     plan: ExecutionPlan,
-    allow_clear_input_blocks: bool,
     dataset_uuid: str,
     preserve_order: bool,
 ) -> BlockList:
@@ -82,7 +73,6 @@ def execute_to_legacy_block_list(
     Args:
         executor: The executor to use.
         plan: The legacy plan to execute.
-        allow_clear_input_blocks: Whether the executor may consider clearing blocks.
         dataset_uuid: UUID of the dataset for this execution.
         preserve_order: Whether to preserve order in execution.
 
@@ -149,23 +139,6 @@ def _bundles_to_block_list(bundles: Iterator[RefBundle]) -> BlockList:
     return BlockList(blocks, metadata, owned_by_consumer=owns_blocks)
 
 
-def _block_list_to_bundles(blocks: BlockList, owns_blocks: bool) -> List[RefBundle]:
-    output = []
-    for block, meta in blocks.iter_blocks_with_metadata():
-        output.append(
-            RefBundle(
-                [
-                    (
-                        block,
-                        meta,
-                    )
-                ],
-                owns_blocks=owns_blocks,
-            )
-        )
-    return output
-
-
 def _set_stats_uuid_recursive(stats: DatasetStats, dataset_uuid: str) -> None:
     if not stats.dataset_uuid:
         stats.dataset_uuid = dataset_uuid

python/ray/data/_internal/iterator/stream_split_iterator.py

@@ -175,8 +175,6 @@ def add_split_op(dag):
                 output_iterator = execute_to_legacy_bundle_iterator(
                     executor,
                     dataset._plan,
-                    True,
-                    dataset._plan._dataset_uuid,
                     dag_rewrite=add_split_op,
                 )
                 yield output_iterator

python/ray/data/_internal/plan.py

@@ -58,16 +58,13 @@ def __init__(
         self,
         stats: DatasetStats,
         *,
-        run_by_consumer: bool,
         data_context: Optional[DataContext] = None,
     ):
         """Create a plan with no transformation operators.
 
         Args:
             stats: Stats for the base blocks.
             dataset_uuid: Dataset's UUID.
-            run_by_consumer: Whether this plan is invoked to run by the consumption
-            APIs (e.g. .iter_batches()).
         """
         self._in_stats = stats
         # A computed snapshot of some prefix of operators and their corresponding
@@ -81,7 +78,6 @@ def __init__(
         # Set when a Dataset is constructed with this plan
         self._dataset_uuid = None
 
-        self._run_by_consumer = run_by_consumer
         self._dataset_name = None
 
         self._has_started_execution = False
@@ -97,7 +93,6 @@ def __repr__(self) -> str:
         return (
             f"ExecutionPlan("
             f"dataset_uuid={self._dataset_uuid}, "
-            f"run_by_consumer={self._run_by_consumer}, "
             f"snapshot_operator={self._snapshot_operator}"
         )
 
@@ -163,9 +158,7 @@ def generate_logical_plan_string(
                     count = None
                 else:
                     assert len(sources) == 1
-                    plan = ExecutionPlan(
-                        DatasetStats(metadata={}, parent=None), run_by_consumer=False
-                    )
+                    plan = ExecutionPlan(DatasetStats(metadata={}, parent=None))
                     plan.link_logical_plan(LogicalPlan(sources[0]))
                     schema = plan.schema()
                     count = plan.meta_count()
@@ -292,7 +285,6 @@ def copy(self) -> "ExecutionPlan":
         """
         plan_copy = ExecutionPlan(
             self._in_stats,
-            run_by_consumer=self._run_by_consumer,
             data_context=self._context,
         )
         if self._snapshot_bundle is not None:
@@ -311,10 +303,7 @@ def deep_copy(self) -> "ExecutionPlan":
         Returns:
             A deep copy of this execution plan.
         """
-        plan_copy = ExecutionPlan(
-            copy.copy(self._in_stats),
-            run_by_consumer=self._run_by_consumer,
-        )
+        plan_copy = ExecutionPlan(copy.copy(self._in_stats))
         if self._snapshot_bundle:
             # Copy over the existing snapshot.
             plan_copy._snapshot_bundle = copy.copy(self._snapshot_bundle)
@@ -397,7 +386,6 @@ def meta_count(self) -> Optional[int]:
     @omit_traceback_stdout
     def execute_to_iterator(
         self,
-        allow_clear_input_blocks: bool = True,
     ) -> Tuple[
         Iterator[Tuple[ObjectRef[Block], BlockMetadata]],
         DatasetStats,
@@ -407,10 +395,6 @@ def execute_to_iterator(
 
         This will use streaming execution to generate outputs.
 
-        Args:
-            allow_clear_input_blocks: Whether we should try to clear the input blocks
-                for each operator.
-
         Returns:
             Tuple of iterator over output blocks and the executor.
         """
@@ -420,7 +404,7 @@ def execute_to_iterator(
         ctx = self._context
 
         if self.has_computed_output():
-            bundle = self.execute(allow_clear_input_blocks)
+            bundle = self.execute()
             return iter(bundle.blocks), self._snapshot_stats, None
 
         from ray.data._internal.execution.legacy_compat import (
@@ -433,8 +417,6 @@ def execute_to_iterator(
         block_iter = execute_to_legacy_block_iterator(
             executor,
             self,
-            allow_clear_input_blocks=allow_clear_input_blocks,
-            dataset_uuid=self._dataset_uuid,
         )
         # Since the generator doesn't run any code until we try to fetch the first
         # value, force execution of one bundle before we call get_stats().
@@ -449,14 +431,11 @@ def execute_to_iterator(
     @omit_traceback_stdout
     def execute(
         self,
-        allow_clear_input_blocks: bool = True,
         preserve_order: bool = False,
     ) -> RefBundle:
         """Execute this plan.
 
         Args:
-            allow_clear_input_blocks: Whether we should try to clear the input blocks
-                for each operator.
             preserve_order: Whether to preserve order in execution.
 
         Returns:
@@ -515,7 +494,6 @@ def execute(
                 blocks = execute_to_legacy_block_list(
                     executor,
                     self,
-                    allow_clear_input_blocks=allow_clear_input_blocks,
                     dataset_uuid=self._dataset_uuid,
                     preserve_order=preserve_order,
                 )

python/ray/data/dataset.py

@@ -1403,10 +1403,7 @@ def train(self, data_iterator):
                 logical_plan = LogicalPlan(InputData(input_data=ref_bundles))
                 split_datasets.append(
                     MaterializedDataset(
-                        ExecutionPlan(
-                            stats,
-                            run_by_consumer=owned_by_consumer,
-                        ),
+                        ExecutionPlan(stats),
                         logical_plan,
                     )
                 )
@@ -1524,10 +1521,7 @@ def build_node_id_by_actor(actors: List[Any]) -> Dict[Any, str]:
             logical_plan = LogicalPlan(InputData(input_data=[bundle]))
             split_datasets.append(
                 MaterializedDataset(
-                    ExecutionPlan(
-                        stats,
-                        run_by_consumer=owned_by_consumer,
-                    ),
+                    ExecutionPlan(stats),
                     logical_plan,
                 )
             )
@@ -1601,10 +1595,7 @@ def split_at_indices(self, indices: List[int]) -> List["MaterializedDataset"]:
 
             splits.append(
                 MaterializedDataset(
-                    ExecutionPlan(
-                        stats,
-                        run_by_consumer=bundle.owns_blocks,
-                    ),
+                    ExecutionPlan(stats),
                     logical_plan,
                 )
             )
@@ -1793,7 +1784,7 @@ def union(self, *other: List["Dataset"]) -> "Dataset":
         )
         stats.time_total_s = time.perf_counter() - start_time
         return Dataset(
-            ExecutionPlan(stats, run_by_consumer=False),
+            ExecutionPlan(stats),
             logical_plan,
         )
 
@@ -4526,10 +4517,7 @@ def materialize(self) -> "MaterializedDataset":
         ]
         logical_plan = LogicalPlan(InputData(input_data=ref_bundles))
         output = MaterializedDataset(
-            ExecutionPlan(
-                copy._plan.stats(),
-                run_by_consumer=False,
-            ),
+            ExecutionPlan(copy._plan.stats()),
             logical_plan,
         )
         # Metrics are tagged with `copy`s uuid, update the output uuid with

python/ray/data/iterator.py

@@ -875,10 +875,7 @@ def materialize(self) -> "MaterializedDataset":
         ]
         logical_plan = LogicalPlan(InputData(input_data=ref_bundles))
         return MaterializedDataset(
-            ExecutionPlan(
-                stats,
-                run_by_consumer=owned_by_consumer,
-            ),
+            ExecutionPlan(stats),
             logical_plan,
         )
 

python/ray/data/read_api.py

@@ -124,10 +124,7 @@ def from_blocks(blocks: List[Block]):
     from_blocks_op = FromBlocks(block_refs, metadata)
     logical_plan = LogicalPlan(from_blocks_op)
     return MaterializedDataset(
-        ExecutionPlan(
-            DatasetStats(metadata={"FromBlocks": metadata}, parent=None),
-            run_by_consumer=False,
-        ),
+        ExecutionPlan(DatasetStats(metadata={"FromBlocks": metadata}, parent=None)),
         logical_plan,
     )
 
@@ -207,10 +204,7 @@ def from_items(
     from_items_op = FromItems(blocks, metadata)
     logical_plan = LogicalPlan(from_items_op)
     return MaterializedDataset(
-        ExecutionPlan(
-            DatasetStats(metadata={"FromItems": metadata}, parent=None),
-            run_by_consumer=False,
-        ),
+        ExecutionPlan(DatasetStats(metadata={"FromItems": metadata}, parent=None)),
         logical_plan,
     )
 
@@ -441,7 +435,7 @@ def read_datasource(
     )
     logical_plan = LogicalPlan(read_op)
     return Dataset(
-        plan=ExecutionPlan(stats, run_by_consumer=False),
+        plan=ExecutionPlan(stats),
         logical_plan=logical_plan,
     )
 
@@ -2481,10 +2475,7 @@ def from_pandas_refs(
         metadata = ray.get([get_metadata.remote(df) for df in dfs])
         logical_plan = LogicalPlan(FromPandas(dfs, metadata))
         return MaterializedDataset(
-            ExecutionPlan(
-                DatasetStats(metadata={"FromPandas": metadata}, parent=None),
-                run_by_consumer=False,
-            ),
+            ExecutionPlan(DatasetStats(metadata={"FromPandas": metadata}, parent=None)),
             logical_plan,
         )
 
@@ -2495,10 +2486,7 @@ def from_pandas_refs(
     metadata = ray.get(metadata)
     logical_plan = LogicalPlan(FromPandas(blocks, metadata))
     return MaterializedDataset(
-        ExecutionPlan(
-            DatasetStats(metadata={"FromPandas": metadata}, parent=None),
-            run_by_consumer=False,
-        ),
+        ExecutionPlan(DatasetStats(metadata={"FromPandas": metadata}, parent=None)),
         logical_plan,
     )
 
@@ -2581,10 +2569,7 @@ def from_numpy_refs(
     logical_plan = LogicalPlan(FromNumpy(blocks, metadata))
 
     return MaterializedDataset(
-        ExecutionPlan(
-            DatasetStats(metadata={"FromNumpy": metadata}, parent=None),
-            run_by_consumer=False,
-        ),
+        ExecutionPlan(DatasetStats(metadata={"FromNumpy": metadata}, parent=None)),
         logical_plan,
     )
 
@@ -2660,10 +2645,7 @@ def from_arrow_refs(
     logical_plan = LogicalPlan(FromArrow(tables, metadata))
 
     return MaterializedDataset(
-        ExecutionPlan(
-            DatasetStats(metadata={"FromArrow": metadata}, parent=None),
-            run_by_consumer=False,
-        ),
+        ExecutionPlan(DatasetStats(metadata={"FromArrow": metadata}, parent=None)),
         logical_plan,
     )
 

python/ray/data/tests/test_split.py

@@ -101,7 +101,7 @@ def _test_equal_split_balanced(block_sizes, num_splits):
     logical_plan = LogicalPlan(InputData(input_data=ref_bundles))
     stats = DatasetStats(metadata={"TODO": []}, parent=None)
     ds = Dataset(
-        ExecutionPlan(stats, run_by_consumer=True),
+        ExecutionPlan(stats),
         logical_plan,
     )