[data] [streaming] Add streaming_split() API

[ericl]

Why are these changes needed?

TODO:

• Implement backpressure
• Implement split locality
• Implement output locality
• Performance testing
• Unit tests

import numpy as np
import ray
import time

ray.data.context.DatasetContext.get_current().use_streaming_executor = True
ray.data.context.DatasetContext.get_current().execution_options.resource_limits.object_store_memory = 40e9

NUM_CONSUMERS = 10
DATA_GB = 1000

@ray.remote(scheduling_strategy="SPREAD", num_cpus=1)
class Consumer:
    def consume(self, it):
        total = 0
        for block in it.iter_batches(batch_format="pyarrow"):
            total += len(block)
        print("consume done", total)

    def get_location(self):
        return ray.get_runtime_context().get_node_id()

consumers = [Consumer.remote() for c in range(NUM_CONSUMERS)]
hints = [ray.get(c.get_location.remote()) for c in consumers]

def preprocess(x):
    return x

shards = ray.data.range_tensor(DATA_GB * 7000, shape=(80, 80, 3), parallelism=2 * DATA_GB) \
    .map_batches(preprocess) \
    .streaming_split(NUM_CONSUMERS, equal=True, locality_hints=hints)

start = time.time()
ray.get([consumers[i].consume.remote(shards[i]) for i in range(NUM_CONSUMERS)])

time.sleep(1)
print("Throughput", DATA_GB / (time.time() - start), "GiB/s")

Output on 5-node cluster:

Throughput 13.839650118615584 GiB/s

Output on 50 nodes, 50 consumers:

Throughput 58.442277295083365 GiB/s

[amogkam]

Ah I missed that future is ObjectRef[ObjectRef[Block]].

[ericl]

Added explicit types here.

[jianoaix]

Look exciting to get this!

[ericl]

Addressed!