fix / fix do_write

python/ray/data/dataset.py

@@ -105,9 +105,6 @@
  TFRecordDatasource,
  WriteResult,
 )
-from ray.data.datasource.file_based_datasource import (
- _wrap_arrow_serialization_workaround,
-)
 from ray.data.random_access_dataset import RandomAccessDataset
 from ray.data.row import TableRow
 from ray.types import ObjectRef
@@ -2674,7 +2671,7 @@ def write_datasource(
 
  def transform(blocks: Iterable[Block], task_idx, fn) -> []:
  try:
- datasource.sync_write(blocks, task_idx, **write_args)
+ datasource.do_write(blocks, task_idx, **write_args)
  datasource.on_write_complete([])
  except Exception as e:
  datasource.on_write_failed([], e)
@@ -4414,16 +4411,3 @@ def _sliding_window(iterable: Iterable, n: int):
  for elem in it:
  window.append(elem)
  yield tuple(window)
-
-
-def _do_write(
- ds: Datasource,
- ctx: DatasetContext,
- blocks: List[Block],
- meta: List[BlockMetadata],
- ray_remote_args: Dict[str, Any],
- write_args: Dict[str, Any],
-) -> List[ObjectRef[WriteResult]]:
- write_args = _unwrap_arrow_serialization_workaround(write_args)
- DatasetContext._set_current(ctx)
- return ds.do_write(blocks, meta, ray_remote_args=ray_remote_args, **write_args)

python/ray/data/datasource/datasource.py

@@ -1,5 +1,5 @@
 import builtins
-from typing import Any, Callable, Dict, Generic, Iterable, List, Optional, Tuple, Union
+from typing import Any, Callable, Generic, Iterable, List, Optional, Tuple, Union
 
 import numpy as np
 
@@ -52,22 +52,17 @@ def prepare_read(self, parallelism: int, **read_args) -> List["ReadTask[T]"]:
 
  def do_write(
  self,
- blocks: List[ObjectRef[Block]],
- metadata: List[BlockMetadata],
- ray_remote_args: Dict[str, Any],
+ blocks: Iterable[Block],
  **write_args,
- ) -> List[ObjectRef[WriteResult]]:
- """Launch Ray tasks for writing blocks out to the datasource.
+ ) -> WriteResult:
+ """Write blocks out to the datasource. This is used by a single write task.
 
  Args:
- blocks: List of data block references. It is recommended that one
- write task be generated per block.
- metadata: List of block metadata.
- ray_remote_args: Kwargs passed to ray.remote in the write tasks.
+ blocks: List of data blocks.
  write_args: Additional kwargs to pass to the datasource impl.
 
  Returns:
- A list of the output of the write tasks.
+ The output of the write tasks.
  """
  raise NotImplementedError
 
@@ -346,7 +341,7 @@ def get_num_failed(self):
 
  self.data_sink = DataSink.remote()
 
- def sync_write(
+ def do_write(
  self,
  blocks: Iterable[Block],
  task_idx: int,
@@ -357,18 +352,6 @@ def sync_write(
  tasks.append(self.data_sink.write.remote(b))
  return ray.get(tasks)
 
- def do_write(
- self,
- blocks: List[ObjectRef[Block]],
- metadata: List[BlockMetadata],
- ray_remote_args: Dict[str, Any],
- **write_args,
- ) -> List[ObjectRef[WriteResult]]:
- tasks = []
- for b in blocks:
- tasks.append(self.data_sink.write.remote(b))
- return tasks
-
  def on_write_complete(self, write_results: List[WriteResult]) -> None:
  assert all(w == "ok" for w in write_results), write_results
  self.data_sink.increment_ok.remote()

python/ray/data/datasource/file_based_datasource.py

@@ -16,14 +16,12 @@
 )
 
 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.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
-from ray.data.datasource.datasource import Datasource, Reader, ReadTask, WriteResult
+from ray.data.datasource.datasource import Datasource, Reader, ReadTask
 from ray.data.datasource.file_meta_provider import (
  BaseFileMetadataProvider,
  DefaultFileMetadataProvider,
@@ -258,8 +256,7 @@ def _convert_block_to_tabular_block(
  "then you need to implement `_convert_block_to_tabular_block."
  )
 
- # This doesn't launch Ray tasks to write.
- def sync_write(
+ def do_write(
  self,
  blocks: Iterable[Block],
  task_idx: int,
@@ -323,79 +320,7 @@ def write_block(write_path: str, block: Block):
  block_index=task_idx,
  file_format=file_format,
  )
- write_block(write_path, block)
-
- def do_write(
- self,
- blocks: List[ObjectRef[Block]],
- metadata: List[BlockMetadata],
- path: str,
- dataset_uuid: str,
- filesystem: Optional["pyarrow.fs.FileSystem"] = None,
- try_create_dir: bool = True,
- open_stream_args: Optional[Dict[str, Any]] = None,
- 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."""
- path, filesystem = _resolve_paths_and_filesystem(path, filesystem)
- path = path[0]
- if try_create_dir:
- # Arrow's S3FileSystem doesn't allow creating buckets by default, so we add
- # a query arg enabling bucket creation if an S3 URI is provided.
- tmp = _add_creatable_buckets_param_if_s3_uri(path)
- filesystem.create_dir(tmp, recursive=True)
- filesystem = _wrap_s3_serialization_workaround(filesystem)
-
- _write_block_to_file = self._write_block
-
- 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
- if isinstance(fs, _S3FileSystemWrapper):
- fs = fs.unwrap()
- if _block_udf is not None:
- block = _block_udf(block)
-
- with fs.open_output_stream(write_path, **open_stream_args) as f:
- _write_block_to_file(
- f,
- BlockAccessor.for_block(block),
- writer_args_fn=write_args_fn,
- **write_args,
- )
-
- write_block = cached_remote_fn(write_block).options(**ray_remote_args)
-
- file_format = self._FILE_EXTENSION
- if isinstance(file_format, list):
- file_format = file_format[0]
-
- write_tasks = []
- 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
+ return write_block(write_path, block)
 
  def _write_block(
  self,

python/ray/data/datasource/mongo_datasource.py

@@ -1,13 +1,12 @@
 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.util.annotations import PublicAPI
 from typing import Iterable
@@ -38,7 +37,7 @@ class MongoDatasource(Datasource):
  def create_reader(self, **kwargs) -> Reader:
  return _MongoDatasourceReader(**kwargs)
 
- def sync_write(
+ def do_write(
  self,
  blocks: Iterable[Block],
  task_idx: int,
@@ -65,38 +64,6 @@ def write_block(uri: str, database: str, collection: str, block: Block):
  write_tasks.append(write_task)
  return write_tasks
 
- def do_write(
- self,
- blocks: List[ObjectRef[Block]],
- metadata: List[BlockMetadata],
- ray_remote_args: Optional[Dict[str, Any]],
- uri: str,
- database: str,
- collection: str,
- ) -> List[ObjectRef[WriteResult]]:
- import pymongo
-
- _validate_database_collection_exist(
- pymongo.MongoClient(uri), database, collection
- )
-
- def write_block(uri: str, database: str, collection: str, block: Block):
- from pymongoarrow.api import write
-
- block = BlockAccessor.for_block(block).to_arrow()
- 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 = []
- for block in blocks:
- write_task = write_block.remote(uri, database, collection, block)
- write_tasks.append(write_task)
- return write_tasks
-
 
 class _MongoDatasourceReader(Reader):
  def __init__(

python/ray/data/tests/test_dataset_formats.py

@@ -1,5 +1,5 @@
 import os
-from typing import Any, Dict, List, Union
+from typing import List, Union
 
 import pandas as pd
 import pyarrow as pa
@@ -13,7 +13,7 @@
 
 import ray
 from ray.data._internal.arrow_block import ArrowRow
-from ray.data.block import Block, BlockAccessor, BlockMetadata
+from ray.data.block import Block, BlockAccessor
 from ray.data.datasource import (
  Datasource,
  DummyOutputDatasource,
@@ -24,6 +24,7 @@
 from ray.data.tests.mock_http_server import * # noqa
 from ray.tests.conftest import * # noqa
 from ray.types import ObjectRef
+from typing import Iterable
 
 
 def maybe_pipeline(ds, enabled):
@@ -227,6 +228,8 @@ def __init__(self):
  self.rows_written = 0
  self.enabled = True
  self.node_ids = set()
+ self.num_ok = 0
+ self.num_failed = 0
 
  def write(self, node_id: str, block: Block) -> str:
  block = BlockAccessor.for_block(block)
@@ -245,40 +248,48 @@ def get_node_ids(self):
  def set_enabled(self, enabled):
  self.enabled = enabled
 
+ def increment_ok(self):
+ self.num_ok += 1
+
+ def get_num_ok(self):
+ return self.num_ok
+
+ def increment_failed(self):
+ self.num_failed += 1
+
+ def get_num_failed(self):
+ return self.num_failed
+
  self.data_sink = DataSink.remote()
- self.num_ok = 0
- self.num_failed = 0
 
  def do_write(
  self,
- blocks: List[ObjectRef[Block]],
- metadata: List[BlockMetadata],
- ray_remote_args: Dict[str, Any],
+ blocks: Iterable[Block],
+ task_idx: int,
  **write_args,
- ) -> List[ObjectRef[WriteResult]]:
+ ) -> WriteResult:
  data_sink = self.data_sink
 
- @ray.remote
  def write(b):
  node_id = ray.get_runtime_context().get_node_id()
  return ray.get(data_sink.write.remote(node_id, b))
 
- tasks = []
  for b in blocks:
- tasks.append(write.options(**ray_remote_args).remote(b))
- return tasks
+ result = write(b)
+ return result
 
  def on_write_complete(self, write_results: List[WriteResult]) -> None:
  assert all(w == "ok" for w in write_results), write_results
- self.num_ok += 1
+ self.data_sink.increment_ok.remote()
 
  def on_write_failed(
  self, write_results: List[ObjectRef[WriteResult]], error: Exception
  ) -> None:
- self.num_failed += 1
+ self.data_sink.increment_failed.remote()
 
 
 def test_write_datasource_ray_remote_args(ray_start_cluster):
+ ray.shutdown()
  cluster = ray_start_cluster
  cluster.add_node(
  resources={"foo": 100},
@@ -298,8 +309,8 @@ def get_node_id():
  ds = ray.data.range(100, parallelism=10)
  # Pin write tasks to
  ds.write_datasource(output, ray_remote_args={"resources": {"bar": 1}})
- assert output.num_ok == 1
- assert output.num_failed == 0
+ assert ray.get(output.data_sink.get_num_ok.remote()) == 10
+ assert ray.get(output.data_sink.get_num_failed.remote()) == 0
  assert ray.get(output.data_sink.get_rows_written.remote()) == 100
 
  node_ids = ray.get(output.data_sink.get_node_ids.remote())