Repository navigation
[Runtime][Disco] Fix session attribute storage, NVSHMEM build, and test gating - #19736
Merged
Merged
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Code Review
This pull request refactors and improves the robustness of the Disco distributed runtime tests. Key changes include adding __slots__ = ("__dict__",) to several Disco classes to allow dynamic Python-side attributes, introducing a centralized SocketSessionTester utility to manage socket sessions and remote nodes, and refactoring NVSHMEM tests to run in fresh processes to prevent undefined behavior from repeated initializations. Feedback on these changes highlights potential hang risks in the new SocketSessionTester implementation, specifically recommending configuring the server thread as a daemon and adding a timeout to thread.join() to prevent indefinite blocking in CI environments.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
tqchen
requested changes
Jun 11, 2026
tlopex
force-pushed
the
fix-disco-wheel-tests
branch
4 times, most recently
from
June 12, 2026 04:50
147bea9 to
ea9a9f6
Compare
tlopex
changed the title
tlopex
force-pushed
the
fix-disco-wheel-tests
branch
from
June 12, 2026 05:01
ea9a9f6 to
f54f538
Compare
tlopex
changed the title
…st gating
The tvm_ffi Object metaclass now gives every subclass `__slots__ = ()`, so
the Disco Python wrappers can no longer store instance attributes and every
session construction fails with AttributeError. Declare the attributes each
wrapper actually stores as named slots, fix the NVSHMEM dist_gemm.cu so TVM
builds with USE_NVSHMEM = ON, and gate the disco tests on the disco runtime
being present so they skip cleanly on builds (e.g. the pip wheel) that report
USE_NCCL / USE_NVSHMEM = ON without shipping it.
Session attribute storage
- DPackedFunc / DModule: `__slots__ = ("session",)`.
- Session: `__slots__ = ("_cache", "_import_python_module")`, and look up
the cache slot with `hasattr` instead of reading `self.__dict__`.
NVSHMEM build
- src/runtime/extra/contrib/nvshmem/dist_gemm.cu failed to compile under
USE_NVSHMEM = ON after the tvm-ffi migration: `GetDataSize` now lives in
namespace tvm::ffi, and the `LOG(FATAL)` macros are no longer in scope.
Qualify the call as `ffi::GetDataSize` and replace the
`if (worker == nullptr) LOG(FATAL) << ...` checks with
`TVM_FFI_ICHECK(worker != nullptr) << ...` (the idiom already used in this
file; it throws ffi::Error, so the behavior is unchanged). No logic change.
Test gating
- Add a `tvm.testing.requires_nvshmem` feature that probes the
`runtime.disco.nvshmem.init_nvshmem_uid` runtime function rather than the
cmake flag (a flag can be ON in a build that does not ship the runtime).
- test_session.py / test_nvshmem.py: skip at module level when the disco
runtime is absent (`tvm.runtime.disco is None`), so the socket-session
helpers they define are never reached and the modules skip cleanly instead
of erroring at collection.
- test_loader.py: skip at module level unless `runtime.disco.compiled_ccl`
reports "nccl" (the same source of truth as test_ccl.py), and require two
GPUs via tvm.testing.requires_multi_gpu.
- test_nvshmem.py: gate every test on NVSHMEM support and on the machine
having one CUDA device per worker; run each body in a freshly spawned
process with a timeout, because nvshmem_init may run only once per process
and disco worker-0 lives in the test process (pytest previously ran all
bodies in one process); replace the hand-rolled __main__ matrix with
tvm.testing.main().
tlopex
force-pushed
the
fix-disco-wheel-tests
branch
from
June 12, 2026 05:05
f54f538 to
7f5c3bd
Compare
tqchen
approved these changes
Jun 12, 2026
MasterJH5574
pushed a commit
to MasterJH5574/tvm
that referenced
this pull request
Jun 15, 2026
…st gating (apache#19736) The tvm_ffi Object metaclass now gives every subclass `__slots__ = ()`, so the Disco Python wrappers can no longer store instance attributes and every session construction fails with AttributeError. Declare the attributes each wrapper actually stores as named slots, fix the NVSHMEM `dist_gemm.cu` so TVM builds with `USE_NVSHMEM = ON`, and gate the disco tests on the disco runtime being present so they skip cleanly on builds (e.g. the pip wheel) that report `USE_NCCL` / `USE_NVSHMEM = ON` without shipping it. ### Session attribute storage - `DPackedFunc` / `DModule`: `__slots__ = ("session",)`. - `Session`: `__slots__ = ("_cache", "_import_python_module")` (cherry picked from commit b57d0b3)
MasterJH5574
pushed a commit
to MasterJH5574/tvm
that referenced
this pull request
Jun 15, 2026
…st gating (apache#19736) The tvm_ffi Object metaclass now gives every subclass `__slots__ = ()`, so the Disco Python wrappers can no longer store instance attributes and every session construction fails with AttributeError. Declare the attributes each wrapper actually stores as named slots, fix the NVSHMEM `dist_gemm.cu` so TVM builds with `USE_NVSHMEM = ON`, and gate the disco tests on the disco runtime being present so they skip cleanly on builds (e.g. the pip wheel) that report `USE_NCCL` / `USE_NVSHMEM = ON` without shipping it. ### Session attribute storage - `DPackedFunc` / `DModule`: `__slots__ = ("session",)`. - `Session`: `__slots__ = ("_cache", "_import_python_module")` (cherry picked from commit b57d0b3)
tqchen
pushed a commit
that referenced
this pull request
Oct 16, 2026
…st gating (#19736) The tvm_ffi Object metaclass now gives every subclass `__slots__ = ()`, so the Disco Python wrappers can no longer store instance attributes and every session construction fails with AttributeError. Declare the attributes each wrapper actually stores as named slots, fix the NVSHMEM `dist_gemm.cu` so TVM builds with `USE_NVSHMEM = ON`, and gate the disco tests on the disco runtime being present so they skip cleanly on builds (e.g. the pip wheel) that report `USE_NCCL` / `USE_NVSHMEM = ON` without shipping it. ### Session attribute storage - `DPackedFunc` / `DModule`: `__slots__ = ("session",)`. - `Session`: `__slots__ = ("_cache", "_import_python_module")`
tqchen
pushed a commit
that referenced
this pull request
Oct 18, 2026
…st gating (#19736) The tvm_ffi Object metaclass now gives every subclass `__slots__ = ()`, so the Disco Python wrappers can no longer store instance attributes and every session construction fails with AttributeError. Declare the attributes each wrapper actually stores as named slots, fix the NVSHMEM `dist_gemm.cu` so TVM builds with `USE_NVSHMEM = ON`, and gate the disco tests on the disco runtime being present so they skip cleanly on builds (e.g. the pip wheel) that report `USE_NCCL` / `USE_NVSHMEM = ON` without shipping it. ### Session attribute storage - `DPackedFunc` / `DModule`: `__slots__ = ("session",)`. - `Session`: `__slots__ = ("_cache", "_import_python_module")`
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.
The tvm_ffi Object metaclass now gives every subclass
__slots__ = (), so the Disco Python wrappers can no longer store instance attributes and every session construction fails with AttributeError. Declare the attributes eachwrapper actually stores as named slots, fix the NVSHMEM
dist_gemm.cuso TVM builds withUSE_NVSHMEM = ON, and gate the disco tests on the disco runtime being present so they skip cleanly on builds (e.g. the pip wheel) that reportUSE_NCCL/USE_NVSHMEM = ONwithout shipping it.Session attribute storage
DPackedFunc/DModule:__slots__ = ("session",).Session:__slots__ = ("_cache", "_import_python_module")