Change shuffle metadata messages to use UCX Active Messages #2409
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
|
If anyone would like me to walk them throw the code, I am happy to do so. |
Signed-off-by: Alessandro Bellina <abellina@nvidia.com>
abellina
force-pushed
the
shuffle/active_messages_refactor
branch
from
f5ad142 to
dcc173b
Compare
abellina
changed the title
jlowe
requested changes
May 13, 2021
sql-plugin/src/main/scala/com/nvidia/spark/rapids/RapidsConf.scala
Outdated
Show resolved
Hide resolved
sql-plugin/src/main/scala/com/nvidia/spark/rapids/RapidsConf.scala
Outdated
Show resolved
Hide resolved
shuffle-plugin/src/main/scala/com/nvidia/spark/rapids/shuffle/ucx/UCX.scala
Outdated
Show resolved
Hide resolved
shuffle-plugin/src/main/scala/com/nvidia/spark/rapids/shuffle/ucx/UCX.scala
Outdated
Show resolved
Hide resolved
jlowe
reviewed
May 13, 2021
shuffle-plugin/src/main/scala/com/nvidia/spark/rapids/shuffle/ucx/UCX.scala
Outdated
Show resolved
Hide resolved
shuffle-plugin/src/main/scala/com/nvidia/spark/rapids/shuffle/ucx/UCX.scala
Outdated
Show resolved
Hide resolved
shuffle-plugin/src/main/scala/com/nvidia/spark/rapids/shuffle/ucx/UCX.scala
Show resolved
Hide resolved
shuffle-plugin/src/main/scala/com/nvidia/spark/rapids/shuffle/ucx/UCX.scala
Show resolved
Hide resolved
shuffle-plugin/src/main/scala/com/nvidia/spark/rapids/shuffle/ucx/UCX.scala
Outdated
Show resolved
Hide resolved
shuffle-plugin/src/main/scala/com/nvidia/spark/rapids/shuffle/ucx/UCX.scala
Outdated
Show resolved
Hide resolved
shuffle-plugin/src/main/scala/com/nvidia/spark/rapids/shuffle/ucx/UCX.scala
Outdated
Show resolved
Hide resolved
shuffle-plugin/src/main/scala/com/nvidia/spark/rapids/shuffle/ucx/UCX.scala
Show resolved
Hide resolved
shuffle-plugin/src/main/scala/com/nvidia/spark/rapids/shuffle/ucx/UCX.scala
Outdated
Show resolved
Hide resolved
shuffle-plugin/src/main/scala/com/nvidia/spark/rapids/shuffle/ucx/UCX.scala
Outdated
Show resolved
Hide resolved
jlowe
reviewed
May 13, 2021
shuffle-plugin/src/main/scala/com/nvidia/spark/rapids/shuffle/ucx/UCXConnection.scala
Outdated
Show resolved
Hide resolved
shuffle-plugin/src/main/scala/com/nvidia/spark/rapids/shuffle/ucx/UCXConnection.scala
Outdated
Show resolved
Hide resolved
shuffle-plugin/src/main/scala/com/nvidia/spark/rapids/shuffle/ucx/UCXConnection.scala
Outdated
Show resolved
Hide resolved
shuffle-plugin/src/main/scala/com/nvidia/spark/rapids/shuffle/ucx/UCXConnection.scala
Outdated
Show resolved
Hide resolved
shuffle-plugin/src/main/scala/com/nvidia/spark/rapids/shuffle/ucx/UCXConnection.scala
Outdated
Show resolved
Hide resolved
sql-plugin/src/main/scala/com/nvidia/spark/rapids/shuffle/BufferSendState.scala
Outdated
Show resolved
Hide resolved
sql-plugin/src/main/scala/com/nvidia/spark/rapids/shuffle/RapidsShuffleServer.scala
Outdated
Show resolved
Hide resolved
tests/src/test/scala/com/nvidia/spark/rapids/shuffle/RapidsShuffleClientSuite.scala
Outdated
Show resolved
Hide resolved
tests/src/test/scala/com/nvidia/spark/rapids/shuffle/RapidsShuffleServerSuite.scala
Outdated
Show resolved
Hide resolved
jlowe
reviewed
May 13, 2021
shuffle-plugin/src/main/scala/com/nvidia/spark/rapids/shuffle/ucx/UCX.scala
Outdated
Show resolved
Hide resolved
shuffle-plugin/src/main/scala/com/nvidia/spark/rapids/shuffle/ucx/UCX.scala
Outdated
Show resolved
Hide resolved
shuffle-plugin/src/main/scala/com/nvidia/spark/rapids/shuffle/ucx/UCX.scala
Outdated
Show resolved
Hide resolved
| tx.completeWithSuccess(requestType, Option(am.header), Option(buff)) | ||
| } | ||
|
|
||
| override def onHostMessageReceived(size: Long): RefCountedDirectByteBuffer = { |
There was a problem hiding this comment.
I would consider initially handle both metadata and data buffers with AM. Maybe split to some public API, that would handle both cases. Even left onDeviceMessageReceived unimplemented, but to see the whole picture. Trying to falling my or with AM: #2416 - want it to work with both GPU and host memory
There was a problem hiding this comment.
Should RNDV for GPU buffers work in 1.10 for Active Messages? Just checking, I have done mostly host buffers.
I would really like to add GPU support as an different PR, need to discuss with others as well on the timing.
|
@jlowe I think I have addressed your comments. Could you take another look? |
jlowe
reviewed
May 20, 2021
shuffle-plugin/src/main/scala/com/nvidia/spark/rapids/shuffle/ucx/UCX.scala
Outdated
Show resolved
Hide resolved
jlowe
approved these changes
May 20, 2021
|
build |
abellina
added a commit
to abellina/spark-rapids
that referenced
this pull request
May 24, 2021
) * Change metadata messages to use Active Messages Signed-off-by: Alessandro Bellina <abellina@nvidia.com> * Code review comments * Comments: copyright, visibility, tests, conf check * Comments: spacing * Comments: private final, no side effects in BSS * Comments: simplify UCXActiveMessage * Comments: UCX.scala better comments, interface changes * Comments: UCX.scala putIfAbsent * Comments: small cleanup in UCX.scala * Move some static tag-handling functions out of the UCXConnection instances. To help testing * Fix a couple of bugs I introduced during refactorings in this PR * Comments: Refactor Active Message registrations a bit * Update comment * UCXConnectionSuite * Comments: shift-then-mask while extracting executorId * Add a few test cases touching the active message header generation * Comments: Make request am handler a constructor argument * Fix RequestActiveMessageRegistration
nartal1
pushed a commit
to nartal1/spark-rapids
that referenced
this pull request
Jun 9, 2021
) * Change metadata messages to use Active Messages Signed-off-by: Alessandro Bellina <abellina@nvidia.com> * Code review comments * Comments: copyright, visibility, tests, conf check * Comments: spacing * Comments: private final, no side effects in BSS * Comments: simplify UCXActiveMessage * Comments: UCX.scala better comments, interface changes * Comments: UCX.scala putIfAbsent * Comments: small cleanup in UCX.scala * Move some static tag-handling functions out of the UCXConnection instances. To help testing * Fix a couple of bugs I introduced during refactorings in this PR * Comments: Refactor Active Message registrations a bit * Update comment * UCXConnectionSuite * Comments: shift-then-mask while extracting executorId * Add a few test cases touching the active message header generation * Comments: Make request am handler a constructor argument * Fix RequestActiveMessageRegistration
nartal1
pushed a commit
to nartal1/spark-rapids
that referenced
this pull request
Jun 9, 2021
) * Change metadata messages to use Active Messages Signed-off-by: Alessandro Bellina <abellina@nvidia.com> * Code review comments * Comments: copyright, visibility, tests, conf check * Comments: spacing * Comments: private final, no side effects in BSS * Comments: simplify UCXActiveMessage * Comments: UCX.scala better comments, interface changes * Comments: UCX.scala putIfAbsent * Comments: small cleanup in UCX.scala * Move some static tag-handling functions out of the UCXConnection instances. To help testing * Fix a couple of bugs I introduced during refactorings in this PR * Comments: Refactor Active Message registrations a bit * Update comment * UCXConnectionSuite * Comments: shift-then-mask while extracting executorId * Add a few test cases touching the active message header generation * Comments: Make request am handler a constructor argument * Fix RequestActiveMessageRegistration
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.
Closes #1588
Closes #980
Please set whitespace off when reviewing the bulk of this, it makes it easier to understand. Also there are changes here to the flatbuffer java files, and those are generated. I did do some minor
withResourcescleanup, can split that up into another PR.This PR changes the metadata/transfer request messaging from UCX tags to UCX Active Messages. The advantage of Active Messages is that error handling in UCX is going this way, and when things go wrong Active Messages will clean up resources that UCX is using (like fragments from their pools for pipelined protocols). The other big one is that it removes the need to worry about the max metadata size. I haven't refactored the metadata pool yet to make it remove this setting alltogether, can do as a follow up.
It was tested with UCX 1.10.x, and 1.11. With UCX 1.10.x the RNDV mode is the only mode supported (Active Messages could also be EAGER, and they have an automatic mode to switch based on size), so I set "rndv" set as default. The config is internal, but I should say also there doesn't seem to be a big need for EAGER mode right now, not from the 3TB tests I've run.
I had to tweak client/server/iterator tests here, and they all pass. I made some changes and need to retest in a cluster setup, but from my testing at 3TB I see no performance regressions. There is some bit manipulation I'll write tests for in this PR.