[data] iter_batches needs streaming operation

[kekulai-fredchang]

Description

Is there a reason for ray.data.iter_batches to execute lazy transformation of the dataset?

We have an application in which an extremely large ray data is consumed by a generator, so iter_batches is used to yield the data for that use case.

However, iter_batches invokes the total transformation, thus exceeding object storage and getting an OOM error.

• why doesn't the object storage spill to disk to prevent the OOM error?
• why does iter_batches need to execute lazy transformation?
• is there a streaming iter_batches in the pipeline?

Thanks

Use case

another user also reports needing a streaming iter_batches in which we have similar applications in mind
https://discuss.ray.io/t/ray-dataset-from-iterabledataset-no-lazy-implementation/20630

[richardliaw]

Hi @kekulai-fredchang do you have an example script for what you are looking at?

Ray Data has a lazy eval execution pattern, so that's something that is built in by default.

You're probably doing something that is causing a materialization under the hood, and isn't something specific to Ray Data.

Would be great to see an example script for where you're seeing the issue, happy to help!

[kekulai-fredchang]

Hi @richardliaw

According to the iter_batch doc, execution of the lazy transformations will be triggered.

def make_batches(ds, batch_size):
    for batch in ds.iter_batches(batch_size=batch_size, batch_format="numpy", drop_last=False):
        yield stuff_in_batch

ds = ray.data.read_csv(superlargefile)
ds = ds.do_stuffs()
create_table(make_batches(ds))

In Ray dashboard I see the fused operations for ds.do_stuffs(), but the iter_batches step does not appear. Consistent with the documentation of iter_batch, I believe I see the lazy transformations fully executed, but I don't see iter_batches appear because data is not fully materialized. However, the cluster runs out of memory before reaching the step where iter_batches can appear.

Even if this is the case, how come objects don't spill to disk?

And you can see the cluster resource status when it crashed. note that memory is filling at head node.

[richardliaw]

Hmm, well first - iter_batches is a consumption API. so it's not going to appear on the logical plan / UI.

Head node memory filling up but not in the object store suggests that there might be something in your application that is holding memory / leaking memory.

[kekulai-fredchang]

Thanks for your reply

• I now understand your point in regards to iter_batches, that it is indeed should be a streaming operation, but that does not necessarily mean that the dataset is materialized into memory. (so trigger lazy execution means in documentation = streaming not materialize)

• Is there a reason why iter_batches does not seem to be running even though lazy execution was triggered? (i am loading its output to an object storage, so I can see that none of the iter_batches was ever run and loaded into the storage. NB. the iter_batches is used with yield to be a generator for create_table.

e.g. the more complete code flow is as follows:

def make_batches(ds, batch_size):
    for batch in ds.iter_batches(batch_size=batch_size, batch_format="numpy", drop_last=False):
        yield stuff_in_batch

def vec_to_storage(ds):
    schema = ...
    db = connect to a lancedb table
    db.create_table(make_batches(ds),schema=...)  //

ds = ray.data.read_csv(sftp://super_large_file.csv.bz)
ds = ds.do_stuffs()
vec_to_storage(ds)

[kekulai-fredchang]

As an alternative path to load data and handle the OOM issue:

ds = ds.do_stuff()
ds = ds.map_batches(loading_data,...)
ds = ds.select_columns([])  # to not materialize the dataset into memory
ds.materialize() # trigger execution

This does approach does not materialize the dataset into memory and functionally performs the same as iter_batches.

[richardliaw]

Glad that there's a workaround. Unfortunately not clear to me from the provided information why iter_batches isn't working properly. Would be great to have a reproducible script.

[richardliaw]

closing this for now - feel free to reopen with any new questions.