[Datasets] BlockPrefetcher not colocated with iterator

[matthewdeng]

What happened + What you expected to happen

When calling Dataset.iter_batches(prefetch_blocks=1), I expect the _BlockPrefetcher actor to be colocated with the Actor that's calling iter_batches, so that the data is actually prefetched on the same node.

From the repro below, the _BlockPrefetcher is being placed on a different node.

Versions / Dependencies

Ray nightly

Reproduction script

import time

import ray
from ray.data import Dataset


@ray.remote
class InfiniteIterator:

    def iterate(self, ds: Dataset, batch_size=32):
        epoch = 0
        # while True:
        for i in range(2):
            batch_id = 0
            for batch in ds.iter_batches(batch_size=batch_size, prefetch_blocks=2):
                print(f"Epoch {epoch} batch {batch_id}.")
                batch_id += 1
                time.sleep(1)
            epoch += 1


def main():
    ds = ray.data.range(1000)
    iterator = InfiniteIterator.remote()
    ray.get(iterator.iterate.remote(ds))


if __name__ == "__main__":
    main()

Issue Severity

High: It blocks me from completing my task.

[zhe-thoughts]

Thanks for reporting this @matthewdeng .

• Is this affecting any existing production users?
• Is this a regression from 1.12 (I suppose not, but confirming here)

[scv119]

Seems a scheduling bug in Ray core, see details in #25329