[PJRT] Release the GIL during TransferFromServer
#4504
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
will-cromar
added
DO_NOT_MERGE_YET
JackCaoG
approved these changes
Jan 25, 2023
|
This looks like a real error from the CPU test: The linter is just getting confused. I'll expand those |
|
I'm not sure where this test actually comes from: I get this error: Also nothing if I search this test name: https://github.com/pytorch/pytorch/search?q=TestViewOpsXLA |
|
I can reproduce the error locally like this: |
will-cromar
force-pushed
the
wcromar/release-gil
branch
from
202c765
to
1b00a10
Compare
|
The next step is to rebase onto #4503 and confirm performance on v4 and to double-check this on XRT. I don't expect either to be impacted since both cases have replicas in separate processes. |
|
Tested this with PJRT and XRT on both v3 and v4. I don't see any noticeable impact on performance in any case.* * For some reason, the performance is lower with XRT than it is in our automated test, but that's true with or without this change. I'll keep an eye on the benchmarks after this is merged. |
|
The CI test failure looks real. I can reproduce it with |
|
I have the CPP tests passing locally again. 🤞 |
alanwaketan
reviewed
Jan 25, 2023
|
I agree, and this was my first thought as well. But I traced the offending call back all the way, and it didn't pass through our python binding layer ( Attaching the full gdb output for reference. Full `gdb` trace |
|
Yea, after eyeballing the trace. It looks like the GIL is held by the PyTorch layer. I guess it's okay then. Did you consider that maybe we can move the NoGilSection section out from pybind layer to a shared header, and reuse that code here? |
|
I thought about bringing in NoGilSection, but in this case we have to explicitly check if Python is initialized and that the GIL is not already released before we save the thread state. What do you think of modifying |
|
It looks like the upstream is using |
|
Good find! Skimming the code, I'm not sure how it will handle our case where this function is called from a non-python context (C++ tests mainly) or the GIL has already been released. I found that |
|
Haha, not even pybind people has thought about our use cases... Let's leave it as it is now. But probably add some notes suggesting that this code is more like an exception. I just try to prevent people overusing them in the cpp layer. |
will-cromar
force-pushed
the
wcromar/release-gil
branch
from
1ea78b7
to
798dca5
Compare
|
Confirmed performance after rebasing and added a warning in the comments. We have other risky changes going into the next nightly build (especially the TF pin update), so I'm going to hold off on merging this until tests run again. |
will-cromar
force-pushed
the
wcromar/release-gil
branch
from
798dca5
to
fbbc6c6
Compare
|
Implemented offline request from @yeounoh to add an easy switch to disable this change |
will-cromar
removed
the
DO_NOT_MERGE_YET
ManfeiBai
pushed a commit
that referenced
this pull request
Jan 30, 2023
* Release the GIL during `TransferFromServer` * Check for GIL before releasing it * Add comment * formatting * Check for Python intitialization before GIL status * Add explanation and warning * Add temporary switch to hold onto GIL * Formatting
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.
Because execution in PJRT is asynchronous, if we call
TransferFromServeron a buffer produced by a computation that is in progress,TransferToServerwill block until the underlying buffer is ready. However, in this case, the thread callingTransferFromServeris also holding the GIL. That means that in a multithreaded context, we can end up with a deadlock where one thread is holding the GIL indefinitely while it waits for another replica to participate in a computation, and the thread that must run that computation is blocked because the first thread is holding the GIL.This started happening consistently with PJRT on TPU v3 some time in mid-December, triggered by an unrelated commit. I confirmed the GIL deadlock by inspecting hanging threads with
gdb:`gdb` output
This problem was irrelevant to XRT for two reasons:
TransferFromServerwill never block and wait for a result.We have to do one of the following to prevent this deadlock:
TransferFromServerto allow other threads to move forward.This PR implements (2) to preserve performance.