Object GC for block splitting inside the dataset splitting

[jianoaix]

Why are these changes needed?

The pipeline will spill objects when splitting the dataset into multiple equal parts.

Related issue number

#25249

Checks

• I've run scripts/format.sh to lint the changes in this PR.
• I've included any doc changes needed for https://docs.ray.io/en/master/.
• I've made sure the tests are passing. Note that there might be a few flaky tests, see the recent failures at https://flakey-tests.ray.io/
• Testing Strategy
  • Unit tests
  • Release tests
  • This PR is not tested :(

[ericl]

Hmm isn't it copied in both cases?

[jianoaix]

At the level of this method, the right half may not copied and conditioned on this boolean: https://sourcegraph.com/github.com/ray-project/ray@master/-/blob/python/ray/data/dataset.py?L3689

[ericl]

But the left half is always copied. And the right half is either copied or None. It seems good to clear in either case.

[ericl]

Should we more generally auto-free on any transform in lazy mode? It feels a little ad-hoc right now, like we should improve our block handling abstractions. Though we can defer this to future work.

[jianoaix]

Yes, we should do it more cleanly. We may need to have an overhaul of execution semantics, e.g. 1) we execute the plan when it's split(), regardless of whether it's lazy or not; 2) we need to have the clear semantics of eager v.s. cached/pinned dataset (via running ds.fully_executed()), the later will make the dataset like an eager one, but running fully_executed() isn't going to make a lazy dataset eager.

[clarkzinzow]

Since this is force-clearing on the base blocks of the dataset, we should be careful that this won't break certain fan-out patterns. E.g.:

ds = ray.data.from_items(list(range(1000))).experimental_lazy()
ds.limit(10).show()
ds.show()

ds2 = ray.data.from_items(list(range(1000))).experimental_lazy()
dses = ds.split_at_indices([250, 750])
for ds in dses:
    ds.show()
ds2.show()

The first should work since we're only clearing if return_right_half=True (although agreed with @ericl that that shouldn't matter), but I think the second might fail.

[jianoaix]

I looked more into this regarding safety. It can happen actually quite easily (not even need to fanout or use from_items()). For example, this will crash:

ds1 = ray.data.range(100, parallelism=10).map_batches(lambda x: x).experimental_lazy() # map_batches() can be dropped
ds1_splits = split(3, equal=True)
for ds in ds1_splits:
    for batch in ds.iter_batches():
        pass
for batch in ds1.iter_batches():
    pass

It's generally more tricky to do eager object gc for dataset than pipeline, as the later is protected by the API semantics that it can be consumed once at most.

We might still be able to do eager gc for lazy dataset, if we move to a semantics like "each time to execute a lazy dataset, it will be executed from the beginning" ("beginning" would mean, for example, the source input blocks, or some intermediate dataset which is known to be pinned and cannot be cleared before unpinned, the output of eager dataset). This will need a revamp work (filed #26264) first (at least part of it).

One hack we may try is to flag the Dataset (with an extra parameter) whether it's used in DatasetPipeline processing, and if so we clear the input block after split.

[jianoaix]

@ericl @clarkzinzow This should be ready to review.

[clarkzinzow]

LGTM!