-
Notifications
You must be signed in to change notification settings - Fork 5.6k
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
[RFC][Dataset] Actor based prefetching #23952
Conversation
f7ce018
to
9ea1089
Compare
Awesome, great benchmarking! It seems like the underlying problem with the current implementation is that I'm assuming that we don't know why |
Also, if we were to bypass We were already planning to do this in order to be able to pipeline some last-last-mile transformations (Torch tensor conversion, local shuffling, device transfer, etc.) with training, so this implementation would stick even if we fixed |
i think ray.wait behavior is not well defined and can only be reasoned by reading the code.
Do you have code snippet on how to do that? Is calling |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Agree that we could probably do this as a separate thread (which would be less likely to be problematic than using a separate actor).
The root cause here seems to be a bug in ray.wait() though, no? Wait is supposed to be fetching the blocks, and somehow that's not going on correctly. cc @stephanie-wang
I'm fine with merging a thread/actor-based hack for now, as long as we document this as a workaround and investigate the root cause.
Agreed, this sounds like a bug in |
@scv119 In case we go with the actor-based prefetcher, I left a review. |
Co-authored-by: Clark Zinzow <clarkzinzow@gmail.com>
Co-authored-by: Clark Zinzow <clarkzinzow@gmail.com>
Co-authored-by: Clark Zinzow <clarkzinzow@gmail.com>
cc @clarkzinzow @jianoaix maybe merge this? |
@scv119 I'll make a pass at a background thread-based solution on Monday, and if that doesn't pan out, we can merge this actor-based solution. |
This PR looks good to me. Actually, from my understanding, using Actors is cleaner since it's a higher level abstraction than using raw threads. If the performance diff is small, Actor should be preferred. Anything I am missing? |
# ray.wait doesn't work as expected so that we use | ||
# the actor based prefetcher as a work around. Read | ||
# https://github.com/ray-project/ray/issues/23983 for details. | ||
if len(block_window) > 1 and context.actor_prefetcher_enabled: |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Are we able to tell whether this is running at Ray Client (maybe from the ray address given to ray.init()?), and if so, we go back to use ray.wait()? I'm looking whether this solution can address the issue of Ray Client as pointed out by @clarkzinzow at #24158.
If so, it seems a reasonable short-term fix before we land a better one (thread based prefetch or anything else). This is changing an internal spot so it should be not too hard to evolve in the future.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
yeah there is an easy way to test if it's ray client:
ray.util.client.ray.is_connected()
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
comments addressed
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
LGTM as a short-term fix and to catch the 1.13 cut today.
and context.actor_prefetcher_enabled | ||
and not ray.util.client.ray.is_connected() | ||
): | ||
prefetcher = get_or_create_prefetcher() |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
This will currently executes a named actor handle fetch on ever prefetch window; could we cache this actor handle so it's fetched once per batch_blocks()
call?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
LGTM for the 1.13 cut as well, although it'd be great if we could do one named actor handle fetch per .iter_batches()
call instead of one per prefetch window, i.e. one call vs. len(blocks) - prefetch_blocks
calls.
I'll look at making this change, along with test coverage of each path.
…per .iter_batches() call; add test coverage.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
LGTM
Datasets tests look good, merging! |
Why are these changes needed?
The prefetch_blocks implementation doesn't work as expected. Due to
ray.wait()
doesn't given us fine grained control, today we block waiting any of the block returns. As I read the code, it may or may not actually fetching all the blocks.A better way to ensure prefetching not blocking is to use ray remote function call, which is not blocking and ensures the blocks are fetched eventually.
The evaluation shows the actor based prefetcher could achieve perfect prefetching if the prefetch speed is faster than consuming speed.
Evaluation
Preliminary test shows the actor based prefetcher improve the throughtput of following script on a 5 cpu node 1 gpu node cluster , which simulate ingesting validation dataset by
0.01 second per batch validation time:
If we simulate per-batch validation time to 0.01 seconds:
validation time: 88.94s -> 68.13s
ray.wait prefetching:
actor based prefetching
0.05 second per batch validation time:
If we simulate per-batch validation time to 0.05 seconds:
validation time: 165.99 -> 104.1s
Actually as we can see from the experiment, we achieved perfect prefetching as the ray.get/ray.wait time is close to 0.
ray.wait prefetching:
actor based prefetching
Related issue number
Checks
scripts/format.sh
to lint the changes in this PR.