[AIR/Data] Add collate_fn to iter_torch_batches

[amogkam]

Signed-off-by: amogkam amogkamsetty@yahoo.com

Adds a collate_fn argument to iter_batches and iter_torch_batches. This is useful for any last meter preprocessing that's done directly on the batch to be used for training.

Closes #32224.

Why are these changes needed?

Related issue number

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

[bveeramani]

LGTM

[clarkzinzow]

Hmm I don't know if it makes sense to have this as part of the ds.iter_batches() API, since the primary motivation behind this feature is to (1) make the Torch DataLoader port easier, and (2) provide a hook to override how we convert NumPy ndarrays to Torch Tensors. collate_fn doesn't provide much critical value for the core ds.iter_batches() API: it's functionally exactly equivalent to

for batch in ds.iter_batches():
    batch = collate_fn(batch)

Is the primary motivation for this to be able to include it in the local pipelining operation chain, to improve collation performance? If that's the case, I'm wondering if ds.iter_torch_batches() shouldn't reuse ds.iter_batches(), and we should instead have ds.iter_torch_batches() use the internal batch_block_refs API directly, which can expose collate_fn. Most of the batching logic is captured in batch_block_refs anyway.

[amogkam]

Yeah I was also thinking the same...will make the change.

Easiest thing is just to have a private arg in iter_batches. IMO that would be preferable for future maintenance than both APIs directly using batch_block_refs

[amogkam]

Updated to make _collate_fn private for iter_batches

[clarkzinzow]

Hmm I'd rather not muddy the user-facing .iter_batches() API even with a private _collate_fn arg, if possible, since it's already a pretty wide interface. And since batch_block_refs contains most of the batching logic, I don't think that there's much maintenance or drift risk here.

What are your primary concerns for maintenance here?

[amogkam]

it's mostly about the timing logic, but it's not a big deal. Updated to using batch_block_refs directly.

[clarkzinzow]

If stuff like execution triggering and stats collection becomes involved enough around batch_block_refs, we can always introduce a private Dataset._iter_batches() method that can support all of the other batch-based iterator APIs without making the user-facing Dataset.iter_batches() API any wider.

But for the sake of this PR, duplicating that thin logic around batch_block_refs seems fine to me!

[amogkam]

this actually doesn't work for the pipelined case...for some reason that calls into Dataset.iter_torch_batches

[amogkam]

DatasetPipeline makes an implicit assumption that Dataset.iter_torch_batches calls into self.iter_batches, which no longer holds true if we change this implementation.

[amogkam]

I think private arg is the way to go, because the abstractions here are very messy. We can remove this later once we deprecate DatasetPipeline.

[clarkzinzow]

LGTM, nice work on this one!