[UIL] Fix triton LSA test library loading and allocator robustness#484
Open
MC952-arch wants to merge 1 commit into
Open
[UIL] Fix triton LSA test library loading and allocator robustness#484MC952-arch wants to merge 1 commit into
MC952-arch wants to merge 1 commit into
Conversation
MC952-arch
requested review from
Copilot
and removed request for
aoyulong and
mikethegoblin
There was a problem hiding this comment.
Pull request overview
This PR improves reliability of the Triton-based LSA (Local Shared Access) test by making FlagCX shared library discovery and the CUDA pluggable allocator setup more robust, reducing failures due to missing/incorrect library paths and partially-initialized wrapper objects.
Changes:
- Add
FLAGCX_LIB_PATHsupport and cache the compiled CUDA pluggable allocator +torch.cuda.MemPool, with clearer failure signaling in the LSA test. - Add default
libflagcx.sodiscovery logic inFLAGCXLibrary(via$FLAGCX_PATHor repo-localbuild/lib) and harden destructor cleanup to avoid crashes on partial initialization. - Adjust shared library caching in
FLAGCXLibraryto reuse loadedctypes.CDLLinstances per path.
Reviewed changes
Copilot reviewed 2 out of 2 changed files in this pull request and generated 2 comments.
| File | Description |
|---|---|
| plugin/interservice/test_triton_lsa.py | Adds configurable linker path for libflagcx.so, caches allocator/mempool compilation results, and fails fast if the pool can’t be initialized. |
| plugin/interservice/flagcx_wrapper.py | Adds default library search logic and improves robustness around library loading and cleanup. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Comment on lines
+139
to
+141
| import atexit | ||
| atexit.register(_cleanup_flagcx_mem_pool) | ||
| atexit.register(_cleanup_flagcx_allocator_wrapper) |
Comment on lines
468
to
+471
| def __del__(self): | ||
| # free flagcx handler | ||
| self.FLAGCX_CHECK(self._funcs["flagcxHandleFree"](self.handler)) | ||
| if hasattr(self, '_funcs') and hasattr(self, 'handler'): | ||
| self.FLAGCX_CHECK(self._funcs["flagcxHandleFree"](self.handler)) |
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
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.
PR Category
UIL
PR Types
Bug Fixes
PR Description
This PR improves reliability of the Triton-based LSA (Local Shared Access) test by making FlagCX shared library discovery and the CUDA pluggable allocator setup more robust, reducing failures due to missing/incorrect library paths and partially-initialized wrapper objects.