[WIP] Bulk executor initial implementation

python/ray/data/_internal/block_batching.py

@@ -8,6 +8,7 @@
 from ray.actor import ActorHandle
 from ray.data._internal.batcher import Batcher, ShufflingBatcher
 from ray.data._internal.stats import DatasetPipelineStats, DatasetStats
+from ray.data._internal.util import _trace_deallocation
 from ray.data.block import Block, BlockAccessor
 from ray.data.context import DatasetContext
 from ray.types import ObjectRef
@@ -183,8 +184,12 @@ def _sliding_window(iterable: Iterable, n: int, clear_block_after_read: bool = F
         yield tuple(window)
     for elem in it:
         block_ref = window.popleft()
-        if clear_block_after_read:
-            ray._private.internal_api.free(block_ref, local_only=False)
+        if clear_block_after_read and DatasetContext.get_current().eager_free:
+            _trace_deallocation(block_ref, "block_batching._sliding_window")
+        else:
+            _trace_deallocation(
+                block_ref, "block_batching._sliding_window", freed=False
+            )
         window.append(elem)
         yield tuple(window)
 

python/ray/data/_internal/block_list.py

@@ -4,6 +4,7 @@
 import numpy as np
 
 from ray.data.block import Block, BlockMetadata
+from ray.data._internal.util import _trace_allocation
 from ray.types import ObjectRef
 
 
@@ -23,6 +24,8 @@ def __init__(
         owned_by_consumer: bool,
     ):
         assert len(blocks) == len(metadata), (blocks, metadata)
+        for b in blocks:
+            _trace_allocation(b, "BlockList.__init__")
         self._blocks: List[ObjectRef[Block]] = blocks
         self._num_blocks = len(self._blocks)
         self._metadata: List[BlockMetadata] = metadata

python/ray/data/_internal/compute.py

@@ -70,6 +70,9 @@ def _apply(
         fn_constructor_args: Optional[Iterable[Any]] = None,
         fn_constructor_kwargs: Optional[Dict[str, Any]] = None,
     ) -> BlockList:
+        assert (
+            not DatasetContext.get_current().new_execution_backend
+        ), "Legacy backend off"
         assert fn_constructor_args is None and fn_constructor_kwargs is None
         if fn_args is None:
             fn_args = tuple()
@@ -237,6 +240,9 @@ def _apply(
         fn_constructor_kwargs: Optional[Dict[str, Any]] = None,
     ) -> BlockList:
         """Note: this is not part of the Dataset public API."""
+        assert (
+            not DatasetContext.get_current().new_execution_backend
+        ), "Legacy backend off"
         if fn_args is None:
             fn_args = tuple()
         if fn_kwargs is None:

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

@@ -0,0 +1,99 @@
+import logging
+from typing import Dict, List, Iterator, Optional
+
+import ray
+from ray.data._internal.execution.interfaces import (
+    Executor,
+    ExecutionOptions,
+    RefBundle,
+    PhysicalOperator,
+)
+from ray.data._internal.execution.operators.input_data_buffer import InputDataBuffer
+from ray.data._internal.progress_bar import ProgressBar
+from ray.data._internal.stats import DatasetStats
+
+logger = logging.getLogger(__name__)
+
+
+class BulkExecutor(Executor):
+    def __init__(self, options: ExecutionOptions):
+        super().__init__(options)
+        self._stats = DatasetStats(stages={}, parent=None)
+        self._executed = False
+
+    def execute(
+        self, dag: PhysicalOperator, initial_stats: Optional[DatasetStats] = None
+    ) -> Iterator[RefBundle]:
+        """Synchronously executes the DAG via bottom-up recursive traversal."""
+
+        assert not self._executed, "Can only call execute once."
+        self._executed = True
+        if not isinstance(dag, InputDataBuffer):
+            logger.info("Executing DAG %s", dag)
+
+        if initial_stats:
+            self._stats = initial_stats
+
+        saved_outputs: Dict[PhysicalOperator, List[RefBundle]] = {}
+
+        def execute_recursive(op: PhysicalOperator) -> List[RefBundle]:
+            # Avoid duplicate executions.
+            if op in saved_outputs:
+                return saved_outputs[op]
+
+            # Compute dependencies.
+            inputs = [execute_recursive(dep) for dep in op.input_dependencies]
+
+            # Fully execute this operator.
+            logger.debug("Executing op %s", op.name)
+            builder = self._stats.child_builder(op.name)
+            try:
+                for i, ref_bundles in enumerate(inputs):
+                    for r in ref_bundles:
+                        op.add_input(r, input_index=i)
+                    op.inputs_done(i)
+                output = _naive_run_until_complete(op)
+            finally:
+                op.shutdown()
+
+            # Cache and return output.
+            saved_outputs[op] = output
+            op_stats = op.get_stats()
+            op_metrics = op.get_metrics()
+            if op_stats:
+                self._stats = builder.build_multistage(op_stats)
+                self._stats.extra_metrics = op_metrics
+            return output
+
+        return execute_recursive(dag)
+
+    def get_stats(self) -> DatasetStats:
+        assert self._stats is not None, self._stats
+        return self._stats
+
+
+def _naive_run_until_complete(op: PhysicalOperator) -> List[RefBundle]:
+    """Run this operator until completion, assuming all inputs have been submitted.
+
+    Args:
+        op: The operator to run.
+
+    Returns:
+        The list of output ref bundles for the operator.
+    """
+    output = []
+    tasks = op.get_work_refs()
+    if tasks:
+        bar = ProgressBar(op.name, total=op.num_outputs_total())
+        while tasks:
+            done, _ = ray.wait(tasks, fetch_local=True, timeout=0.1)
+            for ready in done:
+                op.notify_work_completed(ready)
+            tasks = op.get_work_refs()
+            while op.has_next():
+                bar.update(1)
+                output.append(op.get_next())
+        bar.close()
+    while op.has_next():
+        output.append(op.get_next())
+    return output

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

@@ -3,7 +3,9 @@
 
 import ray
 from ray.data._internal.stats import DatasetStats, StatsDict
+from ray.data._internal.util import _trace_deallocation
 from ray.data.block import Block, BlockMetadata
+from ray.data.context import DatasetContext
 from ray.types import ObjectRef
 
 
@@ -62,7 +64,15 @@ def destroy_if_owned(self) -> int:
         Returns:
             The number of bytes freed.
         """
-        raise NotImplementedError
+        if self.owns_blocks and DatasetContext.get_current().eager_free:
+            size = self.size_bytes()
+            for b in self.blocks:
+                _trace_deallocation(b[0], "RefBundle.destroy_if_owned")
+            return size
+        else:
+            for b in self.blocks:
+                _trace_deallocation(b[0], "RefBundle.destroy_if_owned", freed=False)
+            return 0
 
 
 @dataclass