Make write an operator as part of the execution plan (#32015)

python/ray/data/dataset.py

@@ -2640,8 +2640,6 @@ def write_datasource(
  ray_remote_args: Kwargs passed to ray.remote in the write tasks.
  write_args: Additional write args to pass to the datasource.
  """
-
- ctx = DatasetContext.get_current()
  if ray_remote_args is None:
  ray_remote_args = {}
  path = write_args.get("path", None)
@@ -2655,37 +2653,71 @@ def write_datasource(
  soft=False,
  )
 
- blocks, metadata = zip(*self._plan.execute().get_blocks_with_metadata())
-
- # TODO(ekl) remove this feature flag.
- if "RAY_DATASET_FORCE_LOCAL_METADATA" in os.environ:
- write_results: List[ObjectRef[WriteResult]] = datasource.do_write(
- blocks, metadata, ray_remote_args=ray_remote_args, **write_args
- )
- else:
- # Prepare write in a remote task so that in Ray client mode, we
- # don't do metadata resolution from the client machine.
- do_write = cached_remote_fn(_do_write, retry_exceptions=False, num_cpus=0)
- write_results: List[ObjectRef[WriteResult]] = ray.get(
- do_write.remote(
- datasource,
- ctx,
- blocks,
- metadata,
+ if hasattr(datasource, "write"):
+ # If the write operator succeeds, the resulting Dataset is a list of
+ # WriteResult (one element per write task). Otherwise, an error will
+ # be raised. The Datasource can handle execution outcomes with the
+ # on_write_complete() and on_write_failed().
+ def transform(blocks: Iterable[Block], ctx, fn) -> Iterable[Block]:
+ return [[datasource.write(blocks, ctx, **write_args)]]
+
+ plan = self._plan.with_stage(
+ OneToOneStage(
+ "write",
+ transform,
+ "tasks",
  ray_remote_args,
- _wrap_arrow_serialization_workaround(write_args),
+ fn=lambda x: x,
  )
  )
+ try:
+ self._write_ds = Dataset(plan, self._epoch, self._lazy).fully_executed()
+ datasource.on_write_complete(
+ ray.get(self._write_ds._plan.execute().get_blocks())
+ )
+ except Exception as e:
+ datasource.on_write_failed([], e)
+ raise
+ else:
+ ctx = DatasetContext.get_current()
+ blocks, metadata = zip(*self._plan.execute().get_blocks_with_metadata())
 
- progress = ProgressBar("Write Progress", len(write_results))
- try:
- progress.block_until_complete(write_results)
- datasource.on_write_complete(ray.get(write_results))
- except Exception as e:
- datasource.on_write_failed(write_results, e)
- raise
- finally:
- progress.close()
+ # TODO(ekl) remove this feature flag.
+ if "RAY_DATASET_FORCE_LOCAL_METADATA" in os.environ:
+ write_results: List[ObjectRef[WriteResult]] = datasource.do_write(
+ blocks, metadata, ray_remote_args=ray_remote_args, **write_args
+ )
+ else:
+ logger.warning(
+ "The Datasource.do_write() is deprecated in "
+ "Ray 2.4 and will be removed in future release. Use "
+ "Datasource.write() instead."
+ )
+ # Prepare write in a remote task so that in Ray client mode, we
+ # don't do metadata resolution from the client machine.
+ do_write = cached_remote_fn(
+ _do_write, retry_exceptions=False, num_cpus=0
+ )
+ write_results: List[ObjectRef[WriteResult]] = ray.get(
+ do_write.remote(
+ datasource,
+ ctx,
+ blocks,
+ metadata,
+ ray_remote_args,
+ _wrap_arrow_serialization_workaround(write_args),
+ )
+ )
+
+ progress = ProgressBar("Write Progress", len(write_results))
+ try:
+ progress.block_until_complete(write_results)
+ datasource.on_write_complete(ray.get(write_results))
+ except Exception as e:
+ datasource.on_write_failed(write_results, e)
+ raise
+ finally:
+ progress.close()
 
  def iterator(self) -> DatasetIterator:
  """Return a :class:`~ray.data.DatasetIterator` that

python/ray/data/datasource/datasource.py

@@ -6,6 +6,7 @@
 import ray
 from ray.data._internal.arrow_block import ArrowRow
 from ray.data._internal.delegating_block_builder import DelegatingBlockBuilder
+from ray.data._internal.execution.interfaces import TaskContext
 from ray.data._internal.util import _check_pyarrow_version
 from ray.data.block import (
  Block,
@@ -31,7 +32,7 @@ class Datasource(Generic[T]):
  of how to implement readable and writable datasources.
 
  Datasource instances must be serializable, since ``create_reader()`` and
- ``do_write()`` are called in remote tasks.
+ ``write()`` are called in remote tasks.
  """
 
  def create_reader(self, **read_args) -> "Reader[T]":
@@ -50,6 +51,25 @@ def prepare_read(self, parallelism: int, **read_args) -> List["ReadTask[T]"]:
  """Deprecated: Please implement create_reader() instead."""
  raise NotImplementedError
 
+ def write(
+ self,
+ blocks: Iterable[Block],
+ **write_args,
+ ) -> WriteResult:
+ """Write blocks out to the datasource. This is used by a single write task.
+
+ Args:
+ blocks: List of data blocks.
+ write_args: Additional kwargs to pass to the datasource impl.
+
+ Returns:
+ The output of the write task.
+ """
+ raise NotImplementedError
+
+ @Deprecated(
+ message="do_write() is deprecated in Ray 2.4. Use write() instead", warning=True
+ )
  def do_write(
  self,
  blocks: List[ObjectRef[Block]],
@@ -319,35 +339,33 @@ def __init__(self):
 
  def write(self, block: Block) -> str:
  block = BlockAccessor.for_block(block)
- if not self.enabled:
- raise ValueError("disabled")
  self.rows_written += block.num_rows()
  return "ok"
 
  def get_rows_written(self):
  return self.rows_written
 
- def set_enabled(self, enabled):
- self.enabled = enabled
-
  self.data_sink = DataSink.remote()
  self.num_ok = 0
  self.num_failed = 0
+ self.enabled = True
 
- def do_write(
+ def write(
  self,
- blocks: List[ObjectRef[Block]],
- metadata: List[BlockMetadata],
- ray_remote_args: Dict[str, Any],
+ blocks: Iterable[Block],
+ ctx: TaskContext,
  **write_args,
- ) -> List[ObjectRef[WriteResult]]:
+ ) -> WriteResult:
  tasks = []
+ if not self.enabled:
+ raise ValueError("disabled")
  for b in blocks:
  tasks.append(self.data_sink.write.remote(b))
- return tasks
+ ray.get(tasks)
+ return "ok"
 
  def on_write_complete(self, write_results: List[WriteResult]) -> None:
- assert all(w == "ok" for w in write_results), write_results
+ assert all(w == ["ok"] for w in write_results), write_results
  self.num_ok += 1
 
  def on_write_failed(

python/ray/data/datasource/file_based_datasource.py

@@ -16,9 +16,9 @@
 )
 
 from ray.data._internal.arrow_block import ArrowRow
-from ray.data._internal.block_list import BlockMetadata
+from ray.data._internal.delegating_block_builder import DelegatingBlockBuilder
+from ray.data._internal.execution.interfaces import TaskContext
 from ray.data._internal.output_buffer import BlockOutputBuffer
-from ray.data._internal.remote_fn import cached_remote_fn
 from ray.data._internal.util import _check_pyarrow_version, _resolve_custom_scheme
 from ray.data.block import Block, BlockAccessor
 from ray.data.context import DatasetContext
@@ -60,7 +60,7 @@ def _get_write_path_for_block(
  *,
  filesystem: Optional["pyarrow.fs.FileSystem"] = None,
  dataset_uuid: Optional[str] = None,
- block: Optional[ObjectRef[Block]] = None,
+ block: Optional[Block] = None,
  block_index: Optional[int] = None,
  file_format: Optional[str] = None,
  ) -> str:
@@ -77,7 +77,7 @@ def _get_write_path_for_block(
  write a file out to the write path returned.
  dataset_uuid: Unique identifier for the dataset that this block
  belongs to.
- block: Object reference to the block to write.
+ block: The block to write.
  block_index: Ordered index of the block to write within its parent
  dataset.
  file_format: File format string for the block that can be used as
@@ -94,7 +94,7 @@ def __call__(
  *,
  filesystem: Optional["pyarrow.fs.FileSystem"] = None,
  dataset_uuid: Optional[str] = None,
- block: Optional[ObjectRef[Block]] = None,
+ block: Optional[Block] = None,
  block_index: Optional[int] = None,
  file_format: Optional[str] = None,
  ) -> str:
@@ -257,10 +257,10 @@ def _convert_block_to_tabular_block(
  "then you need to implement `_convert_block_to_tabular_block."
  )
 
- def do_write(
+ def write(
  self,
- blocks: List[ObjectRef[Block]],
- metadata: List[BlockMetadata],
+ blocks: Iterable[Block],
+ ctx: TaskContext,
  path: str,
  dataset_uuid: str,
  filesystem: Optional["pyarrow.fs.FileSystem"] = None,
@@ -269,10 +269,9 @@ def do_write(
  block_path_provider: BlockWritePathProvider = DefaultBlockWritePathProvider(),
  write_args_fn: Callable[[], Dict[str, Any]] = lambda: {},
  _block_udf: Optional[Callable[[Block], Block]] = None,
- ray_remote_args: Dict[str, Any] = None,
  **write_args,
- ) -> List[ObjectRef[WriteResult]]:
- """Creates and returns write tasks for a file-based datasource."""
+ ) -> WriteResult:
+ """Write blocks for a file-based datasource."""
  path, filesystem = _resolve_paths_and_filesystem(path, filesystem)
  path = path[0]
  if try_create_dir:
@@ -287,9 +286,6 @@ def do_write(
  if open_stream_args is None:
  open_stream_args = {}
 
- if ray_remote_args is None:
- ray_remote_args = {}
-
  def write_block(write_path: str, block: Block):
  logger.debug(f"Writing {write_path} file.")
  fs = filesystem
@@ -305,29 +301,30 @@ def write_block(write_path: str, block: Block):
  writer_args_fn=write_args_fn,
  **write_args,
  )
-
- write_block = cached_remote_fn(write_block).options(**ray_remote_args)
+ # TODO: decide if we want to return richer object when the task
+ # succeeds.
+ return "ok"
 
  file_format = self._FILE_EXTENSION
  if isinstance(file_format, list):
  file_format = file_format[0]
 
- write_tasks = []
+ builder = DelegatingBlockBuilder()
+ for block in blocks:
+ builder.add_block(block)
+ block = builder.build()
+
  if not block_path_provider:
  block_path_provider = DefaultBlockWritePathProvider()
- for block_idx, block in enumerate(blocks):
- write_path = block_path_provider(
- path,
- filesystem=filesystem,
- dataset_uuid=dataset_uuid,
- block=block,
- block_index=block_idx,
- file_format=file_format,
- )
- write_task = write_block.remote(write_path, block)
- write_tasks.append(write_task)
-
- return write_tasks
+ write_path = block_path_provider(
+ path,
+ filesystem=filesystem,
+ dataset_uuid=dataset_uuid,
+ block=block,
+ block_index=ctx.task_idx,
+ file_format=file_format,
+ )
+ return write_block(write_path, block)
 
  def _write_block(
  self,

python/ray/data/datasource/mongo_datasource.py

@@ -1,15 +1,16 @@
 import logging
-from typing import Any, Dict, List, Optional, TYPE_CHECKING
+from typing import Dict, List, Optional, TYPE_CHECKING
 
 from ray.data.datasource.datasource import Datasource, Reader, ReadTask, WriteResult
 from ray.data.block import (
  Block,
  BlockAccessor,
  BlockMetadata,
 )
-from ray.data._internal.remote_fn import cached_remote_fn
-from ray.types import ObjectRef
+from ray.data._internal.delegating_block_builder import DelegatingBlockBuilder
+from ray.data._internal.execution.interfaces import TaskContext
 from ray.util.annotations import PublicAPI
+from typing import Iterable
 
 if TYPE_CHECKING:
  import pymongoarrow.api
@@ -37,15 +38,14 @@ class MongoDatasource(Datasource):
  def create_reader(self, **kwargs) -> Reader:
  return _MongoDatasourceReader(**kwargs)
 
- def do_write(
+ def write(
  self,
- blocks: List[ObjectRef[Block]],
- metadata: List[BlockMetadata],
- ray_remote_args: Optional[Dict[str, Any]],
+ blocks: Iterable[Block],
+ ctx: TaskContext,
  uri: str,
  database: str,
  collection: str,
- ) -> List[ObjectRef[WriteResult]]:
+ ) -> WriteResult:
  import pymongo
 
  _validate_database_collection_exist(
@@ -59,15 +59,16 @@ def write_block(uri: str, database: str, collection: str, block: Block):
  client = pymongo.MongoClient(uri)
  write(client[database][collection], block)
 
- if ray_remote_args is None:
- ray_remote_args = {}
-
- write_block = cached_remote_fn(write_block).options(**ray_remote_args)
- write_tasks = []
+ builder = DelegatingBlockBuilder()
  for block in blocks:
- write_task = write_block.remote(uri, database, collection, block)
- write_tasks.append(write_task)
- return write_tasks
+ builder.add_block(block)
+ block = builder.build()
+
+ write_block(uri, database, collection, block)
+
+ # TODO: decide if we want to return richer object when the task
+ # succeeds.
+ return "ok"
 
 
 class _MongoDatasourceReader(Reader):

python/ray/data/tests/conftest.py

@@ -162,7 +162,7 @@ def _get_write_path_for_block(
  block_index=None,
  file_format=None,
  ):
- num_rows = BlockAccessor.for_block(ray.get(block)).num_rows()
+ num_rows = BlockAccessor.for_block(block).num_rows()
  suffix = (
  f"{block_index:06}_{num_rows:02}_{dataset_uuid}" f".test.{file_format}"
  )