[data] torch datasource streaming read

[Zandew]

Why are these changes needed?

Provides a parallel and streaming implementation for the from_torch read api.

In this PR, we set parallelism=1, so we only utilize the streaming portion of this api.

Not included in this PR:

• Parallel read across multiple nodes
• Blocks may be reordered because of parallel read tasks. Add option to preserve order. Not an issue here because we use parallelism=1.
• Block metadata, there is no good way of getting the size/schema of a torch dataset. We cannot autodetect parallelism well. Also not a problem for now because of parallelism=1

Benchmarks (For parallel read implementation, not streaming)

S3 Images

Image data taks3://air-example-data-2/20G-image-data-synthetic-raw/ was read with parallelism=-1 and iterated over with iter_torch_batches on 32vCPU, 128GiB nodes.

6 nodes
time: ~180s
cpu usage: peak ~56 cores

1 node
time: ~240s
cpu usage: peak ~27 cores

Speed up between cases: ~1.33x

The bottleneck here is while deserializing the images afterwards on the driver. Parallelism does not help with this.

Fake Data

We can generate fake data roughly the same size as the S3 images using torch.randint(...) and spinning for some time for i in range(1e7) to emulate the scenario above but without the need to deserialize the images.

6 nodes
time: ~100s
cpu usage: peak ~170 cores

1 node
time: ~550s
cpu usage: peak ~27 cores

Speed up between cases: ~5.5x

Related issue number

Closes torch read #39287

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 :(

[raulchen]

This will serialize the subset as part of the read task function. I think there may be 2 issues:

1. Some user-custom datasets may not be serializable.
2. If the data is already in memory, serializing and copying it to a remote node would lead to even bad perf compared with just using from_items.

I'm thinking of 2 solutions:

1. fall back to from_items in these cases.
2. Pass a factory method as the parameter, instead of passing the concrete dataset object.