[Data] Async iter_batches
#33510
[Data] Async iter_batches
#33510
Conversation
Signed-off-by: amogkam <[email protected]>
Signed-off-by: amogkam <[email protected]>
Signed-off-by: amogkam <[email protected]>
Signed-off-by: amogkam <[email protected]>
Signed-off-by: amogkam <[email protected]>
Signed-off-by: amogkam <[email protected]>
Signed-off-by: amogkam <[email protected]>
Signed-off-by: amogkam <[email protected]>
Signed-off-by: amogkam <[email protected]>
amogkam
requested review from
richardliaw,
gjoliver,
krfricke,
xwjiang2010,
matthewdeng,
Yard1,
maxpumperla,
a team,
ericl,
scv119,
clarkzinzow,
jjyao,
jianoaix and
c21
as code owners
Signed-off-by: amogkam <[email protected]>
Signed-off-by: amogkam <[email protected]>
Signed-off-by: amogkam <[email protected]>
…atches Signed-off-by: amogkam <[email protected]>
Signed-off-by: amogkam <[email protected]>
Signed-off-by: amogkam <[email protected]>
Signed-off-by: amogkam <[email protected]>
|
|
||
| logger = logging.getLogger(__name__) | ||
|
|
||
| PREFETCHER_ACTOR_NAMESPACE = "ray.dataset" |
There was a problem hiding this comment.
We should really try to get rid of the actor prefetcher in 2.5...
| yield enter_result | ||
|
|
||
|
|
||
| def iter_batches( |
There was a problem hiding this comment.
This function sounds very innocuous when it's doing a lot under the hood (creating threads, etc.).
How about making it a class, such as ThreadPoolBatcher, to make it more clear when reading code that we're setting up this state and executing batching in parallel?
| # Step 5: Make sure to preserve order from threadpool results. | ||
| yield from _preserve_order(batch_iter) | ||
| else: | ||
| # If no batch prefetching is specified, then don't use a threadpool. |
There was a problem hiding this comment.
It may be preferable to always use a threadpool, but use size 1 if there isn't any prefetching. This avoids two different code paths being in play.
| also be None, meaning the entirety of the last block is included in this | ||
| batch. If this value is None, this allows us to eagerly clear the last | ||
| block in this batch after reading, since the last block is not included in | ||
| any other batches. |
There was a problem hiding this comment.
Nice, was wondering how the GC would work.
|
|
||
| def _async_iter_batches(block_refs): | ||
| # Step 1: Construct logical batches based on the metadata. | ||
| batch_iter = _bundle_block_refs_to_logical_batches( |
There was a problem hiding this comment.
I think you need to add type annotations for each of these iterators.
| yield from async_batch_iter | ||
|
|
||
|
|
||
| def legacy_iter_batches( |
There was a problem hiding this comment.
Should move legacy code to its own file.
|
|
||
| def threadpool_computations(logical_batch_iter: Iterator[LogicalBatch]): | ||
| # Step 4.1: Resolve the blocks. | ||
| resolved_batch_iter = _resolve_blocks( |
There was a problem hiding this comment.
It seems you do block deletion here, but isn't there a risk of a race condition if a block gets deleted that another thread still needs? I would expect you would have to do deletion at the end in serial order / from a single thread.
| fetch_queue.put(e, block=True) | ||
|
|
||
| threads = [ | ||
| threading.Thread(target=execute_computation, args=(i,)) |
There was a problem hiding this comment.
Do you want to mark these as daemon threads?
| batch_iter, fn=threadpool_computations, num_workers=prefetch_batches | ||
| ) | ||
| # Step 5: Make sure to preserve order from threadpool results. | ||
| yield from _preserve_order(batch_iter) |
There was a problem hiding this comment.
| yield from _preserve_order(batch_iter) | |
| yield from _restore_original_order(batch_iter) |
ericl
added
the
@author-action-required
Signed-off-by: amogkam <[email protected]>
Signed-off-by: amogkam <[email protected]>
TODO:
Why are these changes needed?
Related issue number
Closes #33508
Checks
git commit -s) in this PR.scripts/format.shto lint the changes in this PR.method in Tune, I've added it in
doc/source/tune/api/under thecorresponding
.rstfile.