[Data] Progress bar estimates are initially incorrect it read tasks yields multiple outputs

[bveeramani]

What happened + What you expected to happen

Ray Data initially assumes that each read task produces exactly one block.

Running: 1/1 CPU, 0/0 GPU, 384.0MB/1.0GB object_store_memory:  80%|███████████████████████▏     | 8/10 [00:10<00:02,  1.13s/it]
- ReadRange->MapBatches(sleep): 1 active, 9 queued, [cpu: 1.0, objects: 256.0MB]:  90%|████████ | 9/10 [00:10<00:01,  1.10s/it

When a read tasks produces multiple output, you might see both the numerator and the denominator increasing by one.

Running: 1/1 CPU, 0/0 GPU, 384.0MB/1.0GB object_store_memory: 100%|████████████████████████████| 61/61 [01:08<00:00,  1.09s/it]

Until Ray Data finally corrects it's estimate when a task completes.

Running: 1/1 CPU, 0/0 GPU, 384.0MB/1.0GB object_store_memory:  27%|██████▊                  | 274/1000 [00:23<01:01, 11.87it/s]

This behavior is janky and confusing if you don't know what's going on under the hood.

Versions / Dependencies

d14c95c

Reproduction script

import time

import numpy as np

import ray

ray.init(num_cpus=1)


target_block_size = ray.data.DataContext.get_current().target_max_block_size


def sleep(batch):
    for _ in range(100):
        time.sleep(0.1)
        yield {"batch": np.zeros((target_block_size,), dtype=np.uint8)}


ray.data.range(10, override_num_blocks=10).map_batches(
    sleep, batch_size=None
).materialize()

Issue Severity

None