[pjrt] ensure client destruction on process exit #1999
Merged
Conversation
This file contains hidden or 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
pilkicTT
requested review from
acolicTT,
ajakovljevicTT,
hshahTT,
jameszianxuTT,
jnie-TT,
mrakitaTT,
nvukobratTT,
sdjukicTT,
sgligorijevicTT and
sshonTT
as code owners
Contributor
|
cc @kmabeeTT |
Contributor
|
Have we filed an issue with the torch-xla folks about them not calling |
pilkicTT
force-pushed
the
pilkic/workaround-destroy
branch
from
3257cff to
fe091fd
Compare
Contributor
|
@hshahTT torch-xla folks are not a fan of implementing proper destructor logic / calling Several parties have encountered this issue, not just us. |
Contributor
|
||||||||||||||
sshonTT
reviewed
Nov 5, 2025
pilkicTT
force-pushed
the
pilkic/workaround-destroy
branch
from
d7c948d to
19cd3a7
Compare
sshonTT
approved these changes
Nov 5, 2025
jameszianxuTT
approved these changes
Nov 5, 2025
Contributor
|
Thank you! Works good for me, am able to run back to back pytest cmds on n300-llmbox now without hangs (and needing to call tt-smi -r between tests). I still get DRAM leak between tests in same pytest cmd, but that's a separate issue to figure out. |
sgligorijevicTT
approved these changes
Nov 5, 2025
alcolic-to
approved these changes
Nov 6, 2025
`torch_xla` doesn't call `PJRT_ClientDestroy` properly. This means that we are not closing the devices properly. Recently, this started causing hangs on `n300` boards on subsequent execution of tests. This PR introduces a global singleton object which will ensure that we are properly destroying the client instance on process shutdown. The singleton serves as a fallback mechanism if the framework doesn't call `PJRT_ClientDestroy` - like in the case of `torch_xla`. Additionally, optimizer sub-meshes are now closed after each compilation; this previously was needed to avoid hangs, but now it causes them. Leaving the mechanism of persisting optimizer submesh in the code base, so that we can play with it if needed. Obviously, we need to dig deep into these issues to fix them properly. NOTE: this does not solve the case when the process terminates abruptly, e.g. in case of `SIGSEGV` (segmentation fault). For this, ideally we would want a fix on `tt-metal` side. Closes #1824
pilkicTT
force-pushed
the
pilkic/workaround-destroy
branch
from
19cd3a7 to
aa67e43
Compare
This was referenced Nov 7, 2025
1 task
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.
torch_xladoesn't callPJRT_ClientDestroyproperly. This means that we are not closing the devices properly.Recently, this started causing hangs on
n300boards on subsequent execution of tests.This PR introduces a global singleton object which will ensure that we are properly destroying the client instance on process shutdown. The singleton serves as a fallback mechanism if the framework doesn't call
PJRT_ClientDestroy- like in the case oftorch_xla.Additionally, optimizer sub-meshes are now closed after each compilation; this previously was needed to avoid hangs, but now it causes them. Leaving the mechanism of persisting optimizer submesh in the code base, so that we can play with it if needed. Obviously, we need to dig deep into these issues to fix them properly.
NOTE: this does not solve the case when the process terminates abruptly, e.g. in case of
SIGSEGV(segmentation fault). For this, ideally we would want a fix ontt-metalside.Closes #1824