[Data] Make memory an explicit map parameter

doc/source/data/transforming-data.rst

@@ -312,6 +312,28 @@ To transform data with a Python class, complete these steps:
 
             ds.materialize()
 
+Avoiding out-of-memory errors
+=============================
+
+If your user defined function uses lots of memory, you might encounter out-of-memory 
+errors. To avoid these errors, configure the ``memory`` parameter. It tells Ray how much 
+memory your function uses, and prevents Ray from scheduling too many tasks on a node.
+
+.. testcode::
+    :hide:
+
+    import ray
+
+    ds = ray.data.range(1)
+
+.. testcode::
+
+    def uses_lots_of_memory(batch: Dict[str, np.ndarray]) -> Dict[str, np.ndarray]:
+        ...
+
+    # Tell Ray that the function uses 1 GiB of memory
+    ds.map_batches(uses_lots_of_memory, memory=1 * 1024 * 1024)
+
 .. _transforming_groupby:
 
 Groupby and transforming groups

python/ray/data/dataset.py

@@ -264,6 +264,7 @@ def map(
         fn_constructor_kwargs: Optional[Dict[str, Any]] = None,
         num_cpus: Optional[float] = None,
         num_gpus: Optional[float] = None,
+        memory: Optional[float] = None,
         concurrency: Optional[Union[int, Tuple[int, int]]] = None,
         ray_remote_args_fn: Optional[Callable[[], Dict[str, Any]]] = None,
         **ray_remote_args,
@@ -335,6 +336,7 @@ def parse_filename(row: Dict[str, Any]) -> Dict[str, Any]:
             num_gpus: The number of GPUs to reserve for each parallel map worker. For
                 example, specify `num_gpus=1` to request 1 GPU for each parallel map
                 worker.
+            memory: The heap memory in bytes to reserve for each parallel map worker.
             concurrency: The number of Ray workers to use concurrently. For a fixed-sized
                 worker pool of size ``n``, specify ``concurrency=n``. For an autoscaling
                 worker pool from ``m`` to ``n`` workers, specify ``concurrency=(m, n)``.
@@ -370,6 +372,9 @@ def parse_filename(row: Dict[str, Any]) -> Dict[str, Any]:
         if num_gpus is not None:
             ray_remote_args["num_gpus"] = num_gpus
 
+        if memory is not None:
+            ray_remote_args["memory"] = memory
+
         plan = self._plan.copy()
         map_op = MapRows(
             self._logical_plan.dag,
@@ -412,6 +417,7 @@ def map_batches(
         fn_constructor_kwargs: Optional[Dict[str, Any]] = None,
         num_cpus: Optional[float] = None,
         num_gpus: Optional[float] = None,
+        memory: Optional[float] = None,
         concurrency: Optional[Union[int, Tuple[int, int]]] = None,
         ray_remote_args_fn: Optional[Callable[[], Dict[str, Any]]] = None,
         **ray_remote_args,
@@ -560,6 +566,7 @@ def __call__(self, batch: Dict[str, np.ndarray]) -> Dict[str, np.ndarray]:
             num_cpus: The number of CPUs to reserve for each parallel map worker.
             num_gpus: The number of GPUs to reserve for each parallel map worker. For
                 example, specify `num_gpus=1` to request 1 GPU for each parallel map worker.
+            memory: The heap memory in bytes to reserve for each parallel map worker.
             concurrency: The number of Ray workers to use concurrently. For a fixed-sized
                 worker pool of size ``n``, specify ``concurrency=n``. For an autoscaling
                 worker pool from ``m`` to ``n`` workers, specify ``concurrency=(m, n)``.
@@ -629,6 +636,7 @@ def __call__(self, batch: Dict[str, np.ndarray]) -> Dict[str, np.ndarray]:
             fn_constructor_kwargs=fn_constructor_kwargs,
             num_cpus=num_cpus,
             num_gpus=num_gpus,
+            memory=memory,
             concurrency=concurrency,
             ray_remote_args_fn=ray_remote_args_fn,
             **ray_remote_args,
@@ -648,6 +656,7 @@ def _map_batches_without_batch_size_validation(
         fn_constructor_kwargs: Optional[Dict[str, Any]],
         num_cpus: Optional[float],
         num_gpus: Optional[float],
+        memory: Optional[float],
         concurrency: Optional[Union[int, Tuple[int, int]]],
         ray_remote_args_fn: Optional[Callable[[], Dict[str, Any]]],
         **ray_remote_args,
@@ -671,6 +680,9 @@ def _map_batches_without_batch_size_validation(
         if num_gpus is not None:
             ray_remote_args["num_gpus"] = num_gpus
 
+        if memory is not None:
+            ray_remote_args["memory"] = memory
+
         batch_format = _apply_batch_format(batch_format)
 
         min_rows_per_bundled_input = None
@@ -1103,6 +1115,7 @@ def flat_map(
         fn_constructor_kwargs: Optional[Dict[str, Any]] = None,
         num_cpus: Optional[float] = None,
         num_gpus: Optional[float] = None,
+        memory: Optional[float] = None,
         concurrency: Optional[Union[int, Tuple[int, int]]] = None,
         ray_remote_args_fn: Optional[Callable[[], Dict[str, Any]]] = None,
         **ray_remote_args,
@@ -1168,6 +1181,7 @@ def duplicate_row(row: Dict[str, Any]) -> List[Dict[str, Any]]:
             num_gpus: The number of GPUs to reserve for each parallel map worker. For
                 example, specify `num_gpus=1` to request 1 GPU for each parallel map
                 worker.
+            memory: The heap memory in bytes to reserve for each parallel map worker.
             concurrency: The number of Ray workers to use concurrently. For a
                 fixed-sized worker pool of size ``n``, specify ``concurrency=n``.
                 For an autoscaling worker pool from ``m`` to ``n`` workers, specify
@@ -1202,6 +1216,9 @@ def duplicate_row(row: Dict[str, Any]) -> List[Dict[str, Any]]:
         if num_gpus is not None:
             ray_remote_args["num_gpus"] = num_gpus
 
+        if memory is not None:
+            ray_remote_args["memory"] = memory
+
         plan = self._plan.copy()
         op = FlatMap(
             input_op=self._logical_plan.dag,

python/ray/data/grouped_data.py

@@ -105,6 +105,7 @@ def map_groups(
         fn_constructor_kwargs: Optional[Dict[str, Any]] = None,
         num_cpus: Optional[float] = None,
         num_gpus: Optional[float] = None,
+        memory: Optional[float] = None,
         concurrency: Optional[Union[int, Tuple[int, int]]] = None,
         **ray_remote_args,
     ) -> "Dataset":
@@ -175,6 +176,7 @@ def map_groups(
             num_gpus: The number of GPUs to reserve for each parallel map worker. For
                 example, specify `num_gpus=1` to request 1 GPU for each parallel map
                 worker.
+            memory: The heap memory in bytes to reserve for each parallel map worker.
             ray_remote_args: Additional resource requirements to request from
                 Ray (e.g., num_gpus=1 to request GPUs for the map tasks). See
                 :func:`ray.remote` for details.
@@ -257,6 +259,7 @@ def wrapped_fn(batch, *args, **kwargs):
             fn_constructor_kwargs=fn_constructor_kwargs,
             num_cpus=num_cpus,
             num_gpus=num_gpus,
+            memory=memory,
             concurrency=concurrency,
             ray_remote_args_fn=None,
             **ray_remote_args,

python/ray/data/tests/test_map.py

@@ -15,6 +15,7 @@
 
 import ray
 from ray._private.test_utils import wait_for_condition
+from ray.data import Dataset
 from ray.data._internal.execution.interfaces.ref_bundle import (
     _ref_bundles_iterator_to_block_refs_list,
 )
@@ -1137,7 +1138,8 @@ def __call__(self, x):
     assert values == [11, 15, 19]
 
 
-def test_map_with_memory_resources(shutdown_only):
+@pytest.mark.parametrize("method", [Dataset.map, Dataset.map_batches, Dataset.flat_map])
+def test_map_with_memory_resources(method, shutdown_only):
     """Test that we can use memory resource to limit the concurrency."""
     num_blocks = 50
     memory_per_task = 100 * 1024**2
@@ -1146,19 +1148,35 @@ def test_map_with_memory_resources(shutdown_only):
 
     concurrency_counter = ConcurrencyCounter.remote()
 
-    def map_batches(batch):
+    def map_fn(row_or_batch):
         ray.get(concurrency_counter.inc.remote())
         time.sleep(0.5)
         ray.get(concurrency_counter.decr.remote())
-        return batch
+        if method is Dataset.flat_map:
+            return [row_or_batch]
+        else:
+            return row_or_batch
 
     ds = ray.data.range(num_blocks, override_num_blocks=num_blocks)
-    ds = ds.map_batches(
-        map_batches,
-        batch_size=None,
-        num_cpus=1,
-        memory=memory_per_task,
-    )
+    if method is Dataset.map:
+        ds = ds.map(
+            map_fn,
+            num_cpus=1,
+            memory=memory_per_task,
+        )
+    elif method is Dataset.map_batches:
+        ds = ds.map_batches(
+            map_fn,
+            batch_size=None,
+            num_cpus=1,
+            memory=memory_per_task,
+        )
+    elif method is Dataset.flat_map:
+        ds = ds.flat_map(
+            map_fn,
+            num_cpus=1,
+            memory=memory_per_task,
+        )
     assert len(ds.take(num_blocks)) == num_blocks
 
     actual_max_concurrency = ray.get(concurrency_counter.get_max_concurrency.remote())