-
-
Notifications
You must be signed in to change notification settings - Fork 67
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
Implement task queues for configurable task launch #398
Merged
Conversation
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
jpsamaroo
force-pushed
the
jps/task-queues
branch
from
July 15, 2023 23:16
a74efdf
to
19c1646
Compare
To allow Dagger to fulfill more use cases, it would be helpful if we could perform some basic transformations on tasks before they get submitted to the scheduler. For example, we would like to enable in-order execution semantics for regions of code that execute GPU kernels, as this matches GPU synchronization semantics, and thus makes it easier to upgrade GPU code to use Dagger. Separately, to enable reliable DAG optimizations, we need to be able to guarantee that a region of user code can be seen as a whole within the scheduler. The one-at-a-time task submission semantics that we currently have are insufficient to achieve this, as some tasks in the DAG region of interest may already have been launched before the optimization can see enough of the DAG to be useful. To support these and other use cases, this commit adds a flexible pre-submit task queueing system, as well as making it possible to add additional tasks as synchronization dependencies (instead of the default set from task arguments). The task queueing system allows a custom task queue to be set in TLS, which will be used by `@spawn`/`spawn` when submitting tasks. The task queue is provided one or more task specifications and `EagerThunk` handles, and is free to delay and/or batch task submission, as well as to modify the task specification arbitrarily to match the desired semantics. Task queues are nestable, and tasks submitted within sets of nested task queues should inherit the semantics of the queue they are contained within most directly (with further transformations occuring as tasks move upwards within the nest of queues). The most upstream task queue submits tasks to the worker 1 eager scheduler, but this is also expected to be flexible to allow unique task submission semantics. To support the goal of predictable optimizations, a `LazyTaskQueue` is added (available via `spawn_bulk`) which batches up multiple task submissions into just one, and locks the scheduler until all tasks have been submitted, allowing the scheduler to see the entire DAG structure all at once. Nesting of `spawn_bulk` queues allows multiple DAG regions to be combined into a single total region which is submitted all at once. The ability to specify additional task synchronization dependencies is also a key piece that is orthogonal to task queues. This feature enables the goal of in-order execution semantics by enabling the creation of an `InOrderTaskQueue` (available via `spawn_sequential`), which tracks the last-submitted task or set of tasks, and adds those tasks as additional synchronization dependencies to the next submitted task or set of tasks, effectively causing serializing behavior. Nesting of `spawn_sequential` queues allows separate sequential chains of tasks to be specified, with deeper-nested chains sequencing after previously-submitted tasks or chains in shallower-nested queues. Interestingly, the nesting of `spawn_bulk` within `spawn_sequential` allows entire DAG regions to explicitly synchronize against each other (such that one region executes before another), while allowing tasks within each region to still expose parallelism. The inverse nesting of `spawn_sequential` within `spawn_bulk` allows a chain of sequential tasks to be submitted all at once, adding interesting optimization opportunities. Alongside these enhancements, the eager submission pipeline is optimized by removing the `eager_thunk` submission pathway (which submitted all tasks into a `Channel`), and allows tasks to be directly submitted into the scheduler without redirection. It is expected that this will improve task submission performance and reduce memory usage.
jpsamaroo
force-pushed
the
jps/task-queues
branch
from
July 16, 2023 19:12
261ae2d
to
1ec8563
Compare
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
To allow Dagger to fulfill more use cases, it would be helpful if we
could perform some basic transformations on tasks before they get
submitted to the scheduler. For example, we would like to enable
in-order execution semantics for regions of code that execute GPU
kernels, as this matches GPU synchronization semantics, and thus makes
it easier to upgrade GPU code to use Dagger.
Separately, to enable reliable DAG optimizations, we need to be able to
guarantee that a region of user code can be seen as a whole within the
scheduler. The one-at-a-time task submission semantics that we currently
have are insufficient to achieve this, as some tasks in the DAG region
of interest may already have been launched before the optimization can
see enough of the DAG to be useful.
To support these and other use cases, this commit adds a flexible
pre-submit task queueing system, as well as making it possible to add
additional tasks as synchronization dependencies (instead of the default
set from task arguments).
The task queueing system allows a custom task queue to be set in TLS,
which will be used by
@spawn
/spawn
when submitting tasks. The taskqueue is provided one or more task specifications and
EagerThunk
handles, and is free to delay and/or batch task submission, as well as
to modify the task specification arbitrarily to match the desired
semantics. Task queues are nestable, and tasks submitted within sets of
nested task queues should inherit the semantics of the queue they are
contained within most directly (with further transformations occuring as
tasks move upwards within the nest of queues). The most upstream task
queue submits tasks to the worker 1 eager scheduler, but this is also
expected to be flexible to allow unique task submission semantics.
To support the goal of predictable optimizations, a
LazyTaskQueue
isadded (available via
spawn_bulk
) which batches up multiple tasksubmissions into just one, and locks the scheduler until all tasks have
been submitted, allowing the scheduler to see the entire DAG structure
all at once. Nesting of
spawn_bulk
queues allows multiple DAG regionsto be combined into a single total region which is submitted all at
once.
The ability to specify additional task synchronization dependencies is also a
key piece that is orthogonal to task queues. This feature enables the
goal of in-order execution semantics by enabling the creation of an
InOrderTaskQueue
(available viaspawn_sequential
), which tracks thelast-submitted task or set of tasks, and adds those tasks as additional
synchronization dependencies to the next submitted task or set of tasks,
effectively causing serializing behavior. Nesting of
spawn_sequential
queues allows separate sequential chains of tasks to be specified, with
deeper-nested chains sequencing after previously-submitted tasks or
chains in shallower-nested queues.
Interestingly, the nesting of
spawn_bulk
withinspawn_sequential
allows entire DAG regions to explicitly synchronize against each other
(such that one region executes before another), while allowing tasks
within each region to still expose parallelism. The inverse nesting of
spawn_sequential
withinspawn_bulk
allows a chain of sequentialtasks to be submitted all at once, adding interesting optimization
opportunities.
Alongside these enhancements, the eager submission pipeline is optimized
by removing the
eager_thunk
submission pathway (which submitted alltasks into a
Channel
), and allows tasks to be directly submitted intothe scheduler without redirection. It is expected that this will improve
task submission performance and reduce memory usage.