[Data] Remove prefetch_blocks parameter of batch-based iteration APIs

python/ray/air/util/check_ingest.py

@@ -44,15 +44,13 @@ def __init__(
         num_epochs: int = 1,
         prefetch_batches: int = 1,
         batch_size: Optional[int] = 4096,
-        # Deprecated.
-        prefetch_blocks: int = 0,
         **kwargs,
     ):
         if not scaling_config:
             scaling_config = ScalingConfig(num_workers=1)
         super().__init__(
             train_loop_per_worker=DummyTrainer.make_train_loop(
-                num_epochs, prefetch_batches, prefetch_blocks, batch_size
+                num_epochs, prefetch_batches, batch_size
             ),
             *args,
             scaling_config=scaling_config,
@@ -63,7 +61,6 @@ def __init__(
     def make_train_loop(
         num_epochs: int,
         prefetch_batches: int,
-        prefetch_blocks: int,
         batch_size: Optional[int],
     ):
         """Make a debug train loop that runs for the given amount of epochs."""
@@ -83,7 +80,6 @@ def train_loop_per_worker():
                 batch_start = time.perf_counter()
                 for batch in data_shard.iter_batches(
                     prefetch_batches=prefetch_batches,
-                    prefetch_blocks=prefetch_blocks,
                     batch_size=batch_size,
                 ):
                     batch_delay = time.perf_counter() - batch_start

python/ray/data/dataset.py

@@ -3656,8 +3656,6 @@ def iter_batches(
         local_shuffle_buffer_size: Optional[int] = None,
         local_shuffle_seed: Optional[int] = None,
         _collate_fn: Optional[Callable[[DataBatch], CollatedData]] = None,
-        # Deprecated.
-        prefetch_blocks: int = 0,
     ) -> Iterable[DataBatch]:
         """Return an iterable over batches of data.
 
@@ -3709,7 +3707,6 @@ def iter_batches(
         batch_format = _apply_batch_format(batch_format)
         return self.iterator().iter_batches(
             prefetch_batches=prefetch_batches,
-            prefetch_blocks=prefetch_blocks,
             batch_size=batch_size,
             batch_format=batch_format,
             drop_last=drop_last,
@@ -3730,8 +3727,6 @@ def iter_torch_batches(
         drop_last: bool = False,
         local_shuffle_buffer_size: Optional[int] = None,
         local_shuffle_seed: Optional[int] = None,
-        # Deprecated
-        prefetch_blocks: int = 0,
     ) -> Iterable[TorchBatchType]:
         """Return an iterable over batches of data represented as Torch tensors.
 
@@ -3817,7 +3812,6 @@ def iter_torch_batches(
         """  # noqa: E501
         return self.iterator().iter_torch_batches(
             prefetch_batches=prefetch_batches,
-            prefetch_blocks=prefetch_blocks,
             batch_size=batch_size,
             dtypes=dtypes,
             device=device,
@@ -3837,8 +3831,6 @@ def iter_tf_batches(
         drop_last: bool = False,
         local_shuffle_buffer_size: Optional[int] = None,
         local_shuffle_seed: Optional[int] = None,
-        # Deprecated
-        prefetch_blocks: int = 0,
     ) -> Iterable[TensorFlowTensorBatchType]:
         """Return an iterable over batches of data represented as TensorFlow tensors.
 
@@ -3904,7 +3896,6 @@ def iter_tf_batches(
         """  # noqa: E501
         return self.iterator().iter_tf_batches(
             prefetch_batches=prefetch_batches,
-            prefetch_blocks=prefetch_blocks,
             batch_size=batch_size,
             dtypes=dtypes,
             drop_last=drop_last,
@@ -3931,8 +3922,6 @@ def to_torch(
         local_shuffle_seed: Optional[int] = None,
         unsqueeze_label_tensor: bool = True,
         unsqueeze_feature_tensors: bool = True,
-        # Deprecated
-        prefetch_blocks: int = 0,
     ) -> "torch.utils.data.IterableDataset":
         """Return a
         `Torch IterableDataset <https://pytorch.org/docs/stable/data.html#torch.utils.data.IterableDataset>`_
@@ -4032,7 +4021,6 @@ def to_torch(
             label_column_dtype=label_column_dtype,
             feature_column_dtypes=feature_column_dtypes,
             batch_size=batch_size,
-            prefetch_blocks=prefetch_blocks,
             prefetch_batches=prefetch_batches,
             drop_last=drop_last,
             local_shuffle_buffer_size=local_shuffle_buffer_size,
@@ -4054,8 +4042,6 @@ def to_tf(
         local_shuffle_seed: Optional[int] = None,
         feature_type_spec: Union["tf.TypeSpec", Dict[str, "tf.TypeSpec"]] = None,
         label_type_spec: Union["tf.TypeSpec", Dict[str, "tf.TypeSpec"]] = None,
-        # Deprecated
-        prefetch_blocks: int = 0,
     ) -> "tf.data.Dataset":
         """Return a `TensorFlow Dataset <https://www.tensorflow.org/api_docs/python/tf/data/Dataset/>`_
         over this :class:`~ray.data.Dataset`.
@@ -4160,7 +4146,6 @@ def to_tf(
             feature_columns=feature_columns,
             label_columns=label_columns,
             prefetch_batches=prefetch_batches,
-            prefetch_blocks=prefetch_blocks,
             drop_last=drop_last,
             batch_size=batch_size,
             local_shuffle_buffer_size=local_shuffle_buffer_size,

python/ray/data/iterator.py

@@ -113,8 +113,6 @@ def iter_batches(
         local_shuffle_seed: Optional[int] = None,
         _collate_fn: Optional[Callable[[DataBatch], "CollatedData"]] = None,
         _finalize_fn: Optional[Callable[[Any], Any]] = None,
-        # Deprecated.
-        prefetch_blocks: int = 0,
     ) -> Iterable[DataBatch]:
         """Return a batched iterable over the dataset.
 
@@ -152,14 +150,6 @@ def iter_batches(
         Returns:
             An iterable over record batches.
         """
-
-        if prefetch_blocks > 0:
-            raise DeprecationWarning(
-                "`prefetch_blocks` arg is deprecated in Ray 2.4. Use "
-                "the `prefetch_batches` arg instead to specify the amount of "
-                "prefetching in terms of batches instead of blocks."
-            )
-
         batch_format = _apply_batch_format(batch_format)
 
         def _create_iterator() -> Iterator[DataBatch]:
@@ -261,8 +251,6 @@ def iter_torch_batches(
         drop_last: bool = False,
         local_shuffle_buffer_size: Optional[int] = None,
         local_shuffle_seed: Optional[int] = None,
-        # Deprecated.
-        prefetch_blocks: int = 0,
     ) -> Iterable["TorchBatchType"]:
         """Return a batched iterable of Torch Tensors over the dataset.
 
@@ -389,7 +377,6 @@ def finalize_fn(batch: Union["torch.Tensor", Dict[str, "torch.Tensor"]]):
 
         return self.iter_batches(
             prefetch_batches=prefetch_batches,
-            prefetch_blocks=prefetch_blocks,
             batch_size=batch_size,
             drop_last=drop_last,
             local_shuffle_buffer_size=local_shuffle_buffer_size,
@@ -407,8 +394,6 @@ def iter_tf_batches(
         drop_last: bool = False,
         local_shuffle_buffer_size: Optional[int] = None,
         local_shuffle_seed: Optional[int] = None,
-        # Deprecated.
-        prefetch_blocks: int = 0,
     ) -> Iterable["TensorFlowTensorBatchType"]:
         """Return a batched iterable of TensorFlow Tensors over the dataset.
 
@@ -464,7 +449,6 @@ def iter_tf_batches(
 
         batch_iterable = self.iter_batches(
             prefetch_batches=prefetch_batches,
-            prefetch_blocks=prefetch_blocks,
             batch_size=batch_size,
             drop_last=drop_last,
             local_shuffle_buffer_size=local_shuffle_buffer_size,
@@ -497,8 +481,6 @@ def to_torch(
         local_shuffle_seed: Optional[int] = None,
         unsqueeze_label_tensor: bool = True,
         unsqueeze_feature_tensors: bool = True,
-        # Deprecated.
-        prefetch_blocks: int = 0,
     ) -> "torch.utils.data.IterableDataset":
         """Return a Torch IterableDataset over this dataset.
 
@@ -632,7 +614,6 @@ def make_generator():
             for batch in self.iter_batches(
                 batch_size=batch_size,
                 batch_format="pandas",
-                prefetch_blocks=prefetch_blocks,
                 prefetch_batches=prefetch_batches,
                 drop_last=drop_last,
                 local_shuffle_buffer_size=local_shuffle_buffer_size,
@@ -687,8 +668,6 @@ def to_tf(
         local_shuffle_seed: Optional[int] = None,
         feature_type_spec: Union["tf.TypeSpec", Dict[str, "tf.TypeSpec"]] = None,
         label_type_spec: Union["tf.TypeSpec", Dict[str, "tf.TypeSpec"]] = None,
-        # Deprecated.
-        prefetch_blocks: int = 0,
     ) -> "tf.data.Dataset":
         """Return a TF Dataset over this dataset.
 
@@ -828,7 +807,6 @@ def convert_batch_to_tensors(
         def generator():
             for batch in self.iter_batches(
                 prefetch_batches=prefetch_batches,
-                prefetch_blocks=prefetch_blocks,
                 batch_size=batch_size,
                 drop_last=drop_last,
                 local_shuffle_buffer_size=local_shuffle_buffer_size,