tlopex opened a new pull request, #19736:
URL: https://github.com/apache/tvm/pull/19736
Problem
-------
1. Disco sessions could no longer be constructed. The tvm_ffi Object
metaclass now gives every subclass `__slots__ = ()`, so the Python-side
attribute assignments the Disco wrappers depend on raise AttributeError:
DPackedFunc/DModule store their owning `session`, and Session caches method
lookups in `self.__dict__`.
2. The disco tests did not gate on the disco runtime actually being present.
A build can report USE_NCCL=ON / USE_NVSHMEM=ON yet not ship the disco
CCL/NVSHMEM runtime (the pip wheel is exactly this case). On such builds the
disco tests errored during collection or skipped on the wrong signal instead of
skipping cleanly. Conversely, several tests silently passed or collected
nothing while exercising no code, so the gaps went unnoticed.
3. TVM could not even be built with USE_NVSHMEM=ON:
src/runtime/extra/contrib/nvshmem/dist_gemm.cu failed to compile (GetDataSize /
LOG / FATAL undefined after the tvm-ffi migration).
Fix
---
Session attribute storage:
- Opt DPackedFunc, DModule, and Session back into instance attributes via
the documented `__slots__ = ("__dict__",)` opt-out.
NVSHMEM runtime build:
- dist_gemm.cu: `GetDataSize` now lives in namespace tvm::ffi, and the `if
(worker == nullptr) LOG(FATAL)` checks referenced macros that were out of
scope. Use ffi::GetDataSize and TVM_FFI_ICHECK(worker != nullptr). LOG(FATAL)
here already throws ffi::Error("InternalError", ...), so this is equivalent and
matches the file's existing tvm-ffi style.
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.
- test_loader.py skips at module level unless `runtime.disco.compiled_ccl`
is present and reports "nccl" (the same source of truth as test_ccl.py), and
requires two GPUs via tvm.testing.requires_multi_gpu.
- test_nvshmem.py gates every test on NVSHMEM support and on the machine
having one CUDA device per worker. Each body now runs in a freshly spawned
process with a timeout, because nvshmem_init may be called only once per
process and disco worker-0 lives in the test process, yet pytest previously ran
every body in that one process. Replaces the hand-rolled __main__ matrix with
tvm.testing.main().
- Move the socket-session helpers duplicated verbatim by test_session.py and
test_nvshmem.py into tvm.testing.disco (following the tvm.testing.tir
precedent). The module stays importable when `tvm.runtime.disco` is None, so
dependent test modules skip instead of erroring at collection.
Bugs surfaced and fixed while validating the above:
- Socket-session deadlock: test_session.py hung indefinitely once a second
socket-session test ran. create_socket_session relied on GC/__del__ to shut
down the previous SocketSessionTester, but a failing test kept the old tester
alive through its traceback, so two socket sessions ran at once and the new
SocketSession server thread blocked forever waiting for the remote node. Add an
idempotent SocketSessionTester.shutdown() and call it explicitly before
constructing the next tester. The shared tester also launches remote workers
with `sys.executable` instead of hard-coded "python3", surfaces SocketSession
constructor failures from the server thread, and initializes the attributes
used by teardown first so it is safe when __init__ raises.
- test_callback.py hung at interpreter teardown: the test created a
ProcessSession with NCCL but never shut it down, so garbage-collecting the
session at exit deadlocked. Shut the session down explicitly in a finally block.
- test_custom_allreduce.py silently skipped all 12 test_allreduce cases:
`compiled_ccl()` returns the backend name string, but the code iterated over it
as if it were a list ('n','c','c','l'), so the `ccl` parametrization was always
empty. Compare the value directly, matching test_loader.py.
- test_session.py::test_num_workers used a bare `return` to skip the
socket-session/num_workers<2 case, which reported as PASSED and hid the lost
coverage. Use pytest.skip with a reason.
- test_nvshmem.py: factor the init/finalize boilerplate, duplicated in four
test bodies, into _init_nvshmem()/_finalize_nvshmem() helpers.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
