[Data] Optimize block prefetching

[raulchen]

Why are these changes needed?

WaitBlockPrefetcher will blockingly wait for the first block. When prefetch size is small, this can cause latency on the critical path. This PR moves the wait to a background thread.

Related issue number

closes #35521

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 added any new APIs to the API Reference. For example, if I added a
    method in Tune, I've added it in doc/source/tune/api/ under the
    corresponding .rst file.
• 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 :(

[amogkam]

kicking off the release test here: https://buildkite.com/ray-project/release-tests-pr/builds/39284

[amogkam]

why is this change required instead of triggering the fetch of the entire sliding window?

[raulchen]

It's a bug previously. There is no need to redundantly prefetching the same blocks (both for wait and actor-based prefetchers).

[raulchen]

kicking off the release test here: https://buildkite.com/ray-project/release-tests-pr/builds/39284

@amogkam How to check the benchmark results from this?

[amogkam]

iter-torch-batches-bs-32-prefetch-0-shuffleNone = {'time': 92.82726407200005}
--
  | iter-torch-batches-bs-32-prefetch-0-shuffle64 = {'time': 94.301909722}
  | iter-torch-batches-bs-32-prefetch-1-shuffleNone = {'time': 92.12948181000002}
  | iter-torch-batches-bs-32-prefetch-1-shuffle64 = {'time': 91.936336773}
  | iter-torch-batches-bs-32-prefetch-4-shuffleNone = {'time': 92.30885751599999}
  | iter-torch-batches-bs-32-prefetch-4-shuffle64 = {'time': 92.08876677199999}

looks like we still see the regression. I would expect that with prefetching we should be down to 70 seconds. Running it again to confirm.

[ericl]

Catch and log exceptions from this loop?

[raulchen]

Fixed a few bugs. Now prefetch_batches=1 is more performant. But weirdly, prefetch_batches=4 is the same as no prefetching.

iter-torch-batches-bs-32-prefetch-0-shuffleNone = {'time': 94.05068665299996}
--
iter-torch-batches-bs-32-prefetch-0-shuffle64 = {'time': 88.09606919100008}
iter-torch-batches-bs-32-prefetch-1-shuffleNone = {'time': 65.197580313}
iter-torch-batches-bs-32-prefetch-1-shuffle64 = {'time': 64.72387075799998}
iter-torch-batches-bs-32-prefetch-4-shuffleNone = {'time': 89.3722670200001}
iter-torch-batches-bs-32-prefetch-4-shuffle64 = {'time': 86.70470300799991}

[ericl]

Hmm, previously we waited for all blocks queued. That might turn out to be important. I think we should follow the previous call pattern of passing all blocks to be prefetched to the wait call, and using num returns 1 still.

[raulchen]

After reverting to prefetching the entire sliding window every time, prefetch_batches=4 is effective as well now. I didn't figure out why though.

iter-torch-batches-bs-32-prefetch-0-shuffleNone = {'time': 86.87967626900002}
--
iter-torch-batches-bs-32-prefetch-0-shuffle64 = {'time': 86.52172062299996}
iter-torch-batches-bs-32-prefetch-1-shuffleNone = {'time': 63.789988204}
iter-torch-batches-bs-32-prefetch-1-shuffle64 = {'time': 63.34117246599999}
iter-torch-batches-bs-32-prefetch-4-shuffleNone = {'time': 60.843809012000065}
iter-torch-batches-bs-32-prefetch-4-shuffle64 = {'time': 59.433647494999946}

@ericl another question. currently each prefetcher instance will create a new thread (I added a "Prefetcher.stop" method to make sure the thread will stop asap). Do you think this will create too many threads in practice? If so, we may want to use a static thread pool. I think it should be fine. because users are not likely to use multiple "iter_batches" simultaneously.

[ericl]

Interesting. It could be we auto cancel previous wait requests. Unfortunately the code is very old here.

Also agree a thread per active iterator is totally fine.

[amogkam]

awesome, thanks @raulchen!

[raulchen]

Per offline discussions, some possible explanations are:

1. the blocks that have been prefetched get evicted again. the second "ray.wait" either prevents them from being evicted or fetches them again.
2. A second ray.wait somehow cancels the previous ray.wait.

Either way, we can merge this PR first to fix the regression.