[Datasets] Support out-of-band serialization.

[clarkzinzow]

This PR adds support for out-of-band serialization of datasets, which is required for tuning a training dataset hyperparameter with cross-cluster stopping and resuming of experiments.

In the process of adding this feature, a refactor of the execution plan and LazyBlockList seemed prudent to meet the following set of requirements:

1. [P0] The full stage lineage must be maintained - we must always be able to reconstruct a Datasets from a lazy datasource.
2. [P0] In order to be able to eagerly reclaim memory, we can’t hold on to object references from previous stages.
3. [P0] We need to preserve existing eager mode semantics, both for plain datasets and for each stage in a dataset pipeline.
4. [P0] To keep from unnecessarily duplicating reads, we should reuse computed block partitions when possible.

while adhering to the following constraints:

1. Refactoring the pipelining implementation is out-of-scope.
2. Deviating from the core execution model paradigm of BlockList --> BlockList stages is out-of-scope.
3. Shoe-horning raw (no ray.put()) read tasks into a BlockList is untenable.

Solution

In addition to adding out-of-band serialization support, this PR:

• Adds tracking of the full lineage for all Datasets, storing any intermediate or end-results as a snapshot of some stage in the lineage.
• Makes ReadTasks a first-class concept in LazyBlockList.
• Modifies read stage fusion to utilize already-computed blocks, eliminating redundant reads.
• Optimization passes no longer mutate the input blocks or stages, allowing us to leverage read stage fusion while maintaining the requisite lineage to reconstruct the dataset from the lazy datasource.
• Adds more careful logic around LazyBlockList ramp-up, including around progressive schema/metadata fetching.
• Eagerly unlinks input block references to allow that memory to be reclaimed after the first stage finishes execution.

TODO

• Fix transient bug in stats.
• Add more test coverage for ExecutionPlan and LazyBlockList.
• Try reverting back to always shoe-horning read tasks/block partitions into a plain block list at execution time to consolidate fusion logic into the stages.
• Validate/clean up implementation.
• Split into multiple easier-to-review PRs, if needed.

Closes #22778

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]

A few things to simplify the logic here:

1. Can we avoid the "shoe-horning" of read tasks into a block list? The plan object can have an explicit list of read tasks we extract from the LazyBlockList. This avoids needing to change the block list classes.
2. Can we avoid having an "index" of completed vs not? It would be clearer to instead split the stages into "prev_stages" and "stages".

[scv119]

ping @clarkzinzow any update on this PR?

[clarkzinzow]

@scv119 Actively working on it and the integration with Xiaowei, ran into some complications with the suggested refactor and I'm working on a solution that doesn't increase the scope of the PR.

We should merge this by EOD Monday to make sure AIR is unblocked.

[ericl]

Can you split this up into smaller PRs?

[clarkzinzow]

Superseded by stacked PRs, supported added in #23932. Closing!