Add objects GC in dataset iterator

python/ray/data/_internal/dataset_iterator/dataset_iterator_impl.py

@@ -30,7 +30,7 @@ def _to_block_iterator(
         ds = self._base_dataset
         block_iterator, stats, executor = ds._plan.execute_to_iterator()
         ds._current_executor = executor
-        return block_iterator, stats
+        return block_iterator, stats, False
 
     def stats(self) -> str:
         return self._base_dataset.stats()

python/ray/data/_internal/dataset_iterator/pipelined_dataset_iterator.py

@@ -35,11 +35,20 @@ def _to_block_iterator(
     ]:
         epoch_pipeline = self._get_next_dataset()
 
+        if epoch_pipeline._first_dataset is not None:
+            blocks_owned_by_consumer = (
+                epoch_pipeline._first_dataset._plan.execute()._owned_by_consumer
+            )
+        else:
+            blocks_owned_by_consumer = (
+                epoch_pipeline._peek()._plan.execute()._owned_by_consumer
+            )
+
         def block_iter():
             for ds in epoch_pipeline.iter_datasets():
                 yield from ds._plan.execute().iter_blocks_with_metadata()
 
-        return block_iter(), None
+        return block_iter(), None, blocks_owned_by_consumer
 
     def iter_batches(
         self,

python/ray/data/_internal/dataset_iterator/stream_split_dataset_iterator.py

@@ -100,7 +100,7 @@ def gen_blocks() -> Iterator[Tuple[ObjectRef[Block], BlockMetadata]]:
                     )
                     yield block_ref
 
-        return gen_blocks(), None
+        return gen_blocks(), None, False
 
     def stats(self) -> str:
         """Implements DatasetIterator."""

python/ray/data/dataset_iterator.py

@@ -153,7 +153,7 @@ def iter_batches(
 
         time_start = time.perf_counter()
 
-        block_iterator, stats = self._to_block_iterator()
+        block_iterator, stats, blocks_owned_by_consumer = self._to_block_iterator()
         if use_legacy:
             # Legacy iter_batches does not use metadata.
             def drop_metadata(block_iterator):
@@ -164,6 +164,7 @@ def drop_metadata(block_iterator):
                 drop_metadata(block_iterator),
                 stats=stats,
                 prefetch_blocks=prefetch_blocks,
+                clear_block_after_read=blocks_owned_by_consumer,
                 batch_size=batch_size,
                 batch_format=batch_format,
                 drop_last=drop_last,
@@ -175,6 +176,7 @@ def drop_metadata(block_iterator):
             yield from iter_batches(
                 block_iterator,
                 stats=stats,
+                clear_block_after_read=blocks_owned_by_consumer,
                 batch_size=batch_size,
                 batch_format=batch_format,
                 drop_last=drop_last,

python/ray/data/tests/test_dataset_pipeline.py

@@ -393,6 +393,12 @@ def test_iter_batches_basic(ray_start_regular_shared):
     assert all(len(e) == 1 for e in batches)
 
 
+def test_to_torch(ray_start_regular_shared):
+    pipe = ray.data.range(10, parallelism=10).window(blocks_per_window=2)
+    batches = list(pipe.to_torch(batch_size=None))
+    assert len(batches) == 10
+
+
 def test_iter_batches_batch_across_windows(ray_start_regular_shared):
     # 3 windows, each containing 3 blocks, each containing 3 rows.
     pipe = ray.data.range(27, parallelism=9).window(blocks_per_window=3)

python/ray/data/tests/test_object_gc.py

@@ -16,6 +16,17 @@ def check_no_spill(ctx, pipe):
     assert "Spilled" not in meminfo, meminfo
 
 
+def check_to_torch_no_spill(ctx, pipe):
+    # Run up to 10 epochs of the pipeline to stress test that
+    # no spilling will happen.
+    max_epoch = 10
+    for p in pipe.iter_epochs(max_epoch):
+        for _ in p.to_torch(batch_size=None):
+            pass
+    meminfo = memory_summary(ctx.address_info["address"], stats_only=True)
+    assert "Spilled" not in meminfo, meminfo
+
+
 def test_iter_batches_no_spilling_upon_no_transformation(shutdown_only):
     # The object store is about 300MB.
     ctx = ray.init(num_cpus=1, object_store_memory=300e6)
@@ -24,6 +35,8 @@ def test_iter_batches_no_spilling_upon_no_transformation(shutdown_only):
 
     check_no_spill(ctx, ds.repeat())
     check_no_spill(ctx, ds.window(blocks_per_window=20))
+    check_to_torch_no_spill(ctx, ds.repeat())
+    check_to_torch_no_spill(ctx, ds.window(blocks_per_window=20))
 
 
 def test_iter_batches_no_spilling_upon_rewindow(shutdown_only):
@@ -35,6 +48,9 @@ def test_iter_batches_no_spilling_upon_rewindow(shutdown_only):
     check_no_spill(
         ctx, ds.window(blocks_per_window=20).repeat().rewindow(blocks_per_window=10)
     )
+    check_to_torch_no_spill(
+        ctx, ds.window(blocks_per_window=20).repeat().rewindow(blocks_per_window=10)
+    )
 
 
 def test_iter_batches_no_spilling_upon_prior_transformation(shutdown_only):