[Datasets] Enable lazy execution by default

[c21]

Signed-off-by: Cheng Su scnju13@gmail.com

Why are these changes needed?

This PR is to enable lazy execution by default. See ray-project/enhancements#19 for motivation. The change includes:

• Change Dataset constructor: Dataset.__init__(lazy: bool = True). Also remove defer_execution field, as it's no longer needed.
• read_api.py:read_datasource() returns a lazy Dataset with computing the first input block.
• Add ds.fully_executed() calls to required unit tests, to make sure they are passing.

TODO:

• Fix all unit tests
• [Datasets] Fix the bug of eagerly clearing up input blocks #31459
• [Datasets] Make data nightly tests still work with lazy execution #31460
• Remove the behavior to eagerly compute first block for read
• [data] Improve str/repr of lazy Datasets #31417
• Update documentation

Related issue number

Checks

• I've signed off every commit(by using the -s flag, i.e., git commit -s) in this PR.
• 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]

Nice. Shall we also update the documentation in the same change?

[c21]

Shall we also update the documentation in the same change?

@ericl - I am thinking about in a separate PR for easier doc review (assuming there're more code change for fixing unit tests). But I can also do in same PR if people prefer.

[ericl]

Just tried this out, seems to work well! A few thoughts:

1. Currently, any execution will "cache" a snapshot of the final stage of blocks. Should we change this behavior to only "cache" on a call to fully_executed() or add an explicit cache() action?
2. The str-form of Datasets will include "num_rows=?" and "schema=Unknown schema" a lot now. Can we change this to num_rows=<Pending execution> and schema=<Pending execution> for clarity? Better yet, we could improve the str-form to show the pending stages that will be executed.

These changes could go into separate PRs.

Edit: filed #31417 for (2)

[clarkzinzow]

@c21 We should also audit all benchmarks that involve Datasets, to ensure that setup operations that were previously executed eagerly are still executed eagerly, so we're not accidentally including e.g. reading or setup transformations when we're trying to benchmark a single downstream operation. The .map_batches() benchmarks come to mind.

[c21]

We should also audit all benchmarks that involve Datasets, to ensure that setup operations that were previously executed eagerly are still executed eagerly, so we're not accidentally including e.g. reading or setup transformations when we're trying to benchmark a single downstream operation. The .map_batches() benchmarks come to mind.

@clarkzinzow - thanks, will go over all nightly tests.

[c21]

Currently, any execution will "cache" a snapshot of the final stage of blocks. Should we change this behavior to only "cache" on a call to fully_executed() or add an explicit cache() action?

Discussed offline with @ericl, we shall postpone it later given impact is low.

The str-form of Datasets will include "num_rows=?" and "schema=Unknown schema" a lot now. Can we change this to num_rows= and schema= for clarity? Better yet, we could improve the str-form to show the pending stages that will be executed.

Will do it in a separate PR.

[c21]

All CI tests are passed. The failed nightly tests will be addressed in #31460 .

[c21]

This is the major behavior change needed to call out - batch_predictor.predict() becomes lazy now, and not force the execution. Users need to call batch_predictor.predict().fully_executed() to force the prediction to actually run.

The pros is we can chain multiple predictors now for free - batch_predictor2.predict(batch_predictor1.predict()).

The cons is this will probably be a surprising behavior change for current users.

If the impact of pros is not big, we can change to force execution inside batch_predictor.predict().

cc @ericl, @amogkam.

[amogkam]

I'm strongly in favor of having the default behavior being calling fully_executed inside predict, at least for now.

Laziness for the chained case is something we can handle separately.

[ericl]

Let's force execution inside batch predictor? Chaining is not really a use case here, and you can always fall back to using the Data API directly.

[c21]

Cool, let me make the change to force execution inside predictor.

[c21]

Updated to force execution inside batch predictor.

[clarkzinzow]

Mostly LGTM, the big thing that we need resolve IMO is the semi-lazy read behavior. Semi-lazy reading (eagerly reading the first block/file) has been confusing for users, results in redundant reading when there's stage fusion or conversion to a pipeline, and I think we should try to get rid of it when switching to lazy execution by default, if possible. I think that it would be much better if we had fully lazy mode and eager mode, where semi-lazy execution is an execution optimization that is enabled when we determine we only need to compute a subset of blocks; these semantics should translate pretty well to our new execution planner and streaming execution model (e.g. limit pushdown, metadata peeking, streaming consumption of narrow op chains).

Instead of always computing the first block, I think that we should make the read fully lazy, and if the user immediately calls ds.schema() right after reading, we only then trigger reading of the first block (and only if the schema isn't already available from e.g. file metadata).

I think it should be pretty straightforward to move this lazy_block_list.compute_first_block()/progressive computation logic to the ExecutionPlan, since we can have ExecutionPlan.schema() (and other ops that only require the first block) trigger semi-lazy execution with e.g. a certain flag set, and if that flag is set AND the plan only consists of a read stage AND the input blocks are a lazy block list, we trigger minimal computation of that read stage.

[clarkzinzow]

Should we change the fetch_if_missing default value to True in ds.schema()? I.e. should we transparently trigger execution when fetching metadata on a lazy dataset?

[c21]

should we transparently trigger execution when fetching metadata on a lazy dataset?

Yes agreed. Do we want to do in a separate PR ? #31286 (comment) .

[clarkzinzow]

Yep doing that in a separate PR sounds good!

[clarkzinzow]

IMO we should get rid of this default behavior of always computing the first block and ensuring the metadata for the first block, and instead make the dataset fully lazy by default. We can still progressively launch read tasks to e.g. fetch the schema, see the first few rows, streaming iteration directly on the read, etc. We'd just move that logic into ExecutionPlan.schema().

[ericl]

Hmm, that seems fine as long as we still resolve the metadata for parquet and so on. This would only apply to JSON/CSV presumably. Maybe we can also add an optimized schema resolver for these file types that peeks at the header of the file only.

[c21]

Yes agreed, we shall improve the schema resolver in the long term. I am 100% agree to get rid of always computing the first block. Do we want to do it in a separate PR? #31286 (comment)

[ericl]

FWIW, getting rid of the first block reading would also help with integrating fully streaming execution (right now actually that breaks streaming actually). How about we do that as a separate PR followup from this though? I think this PR is already a pretty extensive change, and we should generally avoid mixing complex changes.

[c21]

FWIW, getting rid of the first block reading would also help with integrating fully streaming execution (right now actually that breaks streaming actually). How about we do that as a separate PR followup from this though? I think this PR is already a pretty extensive change, and we should generally avoid mixing complex changes.

SGTM. @clarkzinzow - WDYT?

btw I will make the batch_predictor.predict change in this PR.

[clarkzinzow]

How about we do that as a separate PR followup from this though?

@ericl @c21 As long as it's done as a P0 follow-up PR that we're sure will get in before the next release, that sounds good to me! I would normally say that we shouldn't enable lazy execution by default in master until we have the fully lazy semantics, but since we're actively iterating on the execution model and we have a good bit of time before the release, we can be pragmatic here.

In hindsight, it probably would have been better to do the following sequence of PRs:

1. Move semi-lazy reading from an eager computing of the first block in read_datasource to a metadata peeking and streaming read optimization in the ExecutionPlan, which would keep the existing "eagerly compute first block" semantics for eager mode and would make lazy mode fully lazy.
2. Enable lazy execution by default.
3. Port metadata peeking optimization to plan optimization + streaming executor, since the policy of "only compute as many blocks as needed for operation" is really just a pushdown optimization rule + streaming execution.

[clarkzinzow]

LGTM! Are we going to update documentation and/or improve the repr in this PR?

[c21]

As long as it's done as a P0 follow-up PR that we're sure will get in before the next release, that sounds good to me! I would normally say that we shouldn't enable lazy execution by default in master until we have the fully lazy semantics, but since we're actively iterating on the execution model and we have a good bit of time before the release, we can be pragmatic here.

@clarkzinzow - yeah agree here. The TODOs for this PR (1.Remove the behavior to eagerly compute first block for read, 2.improve the str/repr of dataset, 3.update documentation) are definitely P0 which I will work on next week immediately.