From 2b0f3169cb394335c092955598756d9a10a93951 Mon Sep 17 00:00:00 2001 From: CYJiang Date: Mon, 12 Oct 2026 23:43:54 +0800 Subject: [PATCH 1/4] fix: make MTML library lifecycle reference counted Signed-off-by: CYJiang --- pymtml.py | 53 ++++++------- test_library_lifecycle.py | 154 ++++++++++++++++++++++++++++++++++++++ 2 files changed, 177 insertions(+), 30 deletions(-) create mode 100644 test_library_lifecycle.py diff --git a/pymtml.py b/pymtml.py index 75c6bc3..78cc298 100644 --- a/pymtml.py +++ b/pymtml.py @@ -757,6 +757,9 @@ def timestamp(self): ## Lib loading ## mtmlLib = None libLoadLock = threading.Lock() +# The Python API exposes one process-wide handle. Keep its native lifetime +# balanced across nested and concurrent clients independently of library loading. +libLifecycleLock = threading.Lock() libHandle = c_mtmlLibrary_t() _mtmlLib_refcount = 0 # Incremented on each mtmlInit and decremented on mtmlShutdown @@ -983,44 +986,34 @@ def _LoadMtmlLibrary(): def mtmlLibraryInit(): _LoadMtmlLibrary() - # - # Initialize the library - # - global libHandle fn = _mtmlGetFunctionPointer("mtmlLibraryInit") - ret = fn(byref(libHandle)) - _mtmlCheckReturn(ret) - - # Atomically update refcount + global libHandle global _mtmlLib_refcount - libLoadLock.acquire() - _mtmlLib_refcount += 1 - libLoadLock.release() + with libLifecycleLock: + if _mtmlLib_refcount == 0: + handle = c_mtmlLibrary_t() + ret = fn(byref(handle)) + _mtmlCheckReturn(ret) + libHandle = handle + _mtmlLib_refcount += 1 return None def mtmlLibraryShutDown(): - # - # Leave the library loaded, but shutdown the interface - # global libHandle - if libHandle is None: - return None - - fn = _mtmlGetFunctionPointer("mtmlLibraryShutDown") - ret = fn(libHandle) - _mtmlCheckReturn(ret) - - # Reset libHandle to a fresh instance to allow reinitialization - # and prevent dangling references during garbage collection - libHandle = c_mtmlLibrary_t() - - # Atomically update refcount global _mtmlLib_refcount - libLoadLock.acquire() - if 0 < _mtmlLib_refcount: - _mtmlLib_refcount -= 1 - libLoadLock.release() + with libLifecycleLock: + if _mtmlLib_refcount == 0: + return None + if _mtmlLib_refcount > 1: + _mtmlLib_refcount -= 1 + return None + + fn = _mtmlGetFunctionPointer("mtmlLibraryShutDown") + ret = fn(libHandle) + _mtmlCheckReturn(ret) + libHandle = c_mtmlLibrary_t() + _mtmlLib_refcount = 0 return None diff --git a/test_library_lifecycle.py b/test_library_lifecycle.py new file mode 100644 index 0000000..d428e8e --- /dev/null +++ b/test_library_lifecycle.py @@ -0,0 +1,154 @@ +"""Unit tests for process-wide MTML library lifecycle ownership.""" + +from concurrent.futures import ThreadPoolExecutor +from ctypes import POINTER, c_uint, cast, pointer +from threading import Barrier + +import pytest + +import pymtml + + +class FakeMtmlLibrary: + def __init__(self): + self.init_calls = [] + self.shutdown_calls = [] + self._handles = [] + self.init_result = pymtml.MTML_SUCCESS + self.shutdown_result = pymtml.MTML_SUCCESS + + def init(self, handle_out): + handle_storage = pymtml.struct_c_mtmlLibrary_t() + handle = pointer(handle_storage) + cast(handle_out, POINTER(pymtml.c_mtmlLibrary_t))[0] = handle + self._handles.append((handle_storage, handle)) + self.init_calls.append(handle) + return self.init_result + + def shutdown(self, handle): + self.shutdown_calls.append(handle) + return self.shutdown_result + + def count_devices(self, _handle, count_out): + cast(count_out, POINTER(c_uint))[0] = 2 + return pymtml.MTML_SUCCESS + + +@pytest.fixture +def fake_mtml(monkeypatch): + native = FakeMtmlLibrary() + functions = { + "mtmlLibraryInit": native.init, + "mtmlLibraryShutDown": native.shutdown, + "mtmlLibraryCountDevice": native.count_devices, + } + + monkeypatch.setattr(pymtml, "_LoadMtmlLibrary", lambda: None) + monkeypatch.setattr(pymtml, "_mtmlGetFunctionPointer", functions.__getitem__) + monkeypatch.setattr(pymtml, "libHandle", pymtml.c_mtmlLibrary_t()) + monkeypatch.setattr(pymtml, "_mtmlLib_refcount", 0) + return native + + +def test_nested_init_keeps_native_handle_until_last_shutdown(fake_mtml): + pymtml.mtmlLibraryInit() + pymtml.mtmlLibraryInit() + + assert len(fake_mtml.init_calls) == 1 + assert pymtml._mtmlLib_refcount == 2 + assert bool(pymtml.libHandle) + + pymtml.mtmlLibraryShutDown() + + assert fake_mtml.shutdown_calls == [] + assert pymtml._mtmlLib_refcount == 1 + assert bool(pymtml.libHandle) + assert pymtml.mtmlLibraryCountDevice() == 2 + + pymtml.mtmlLibraryShutDown() + + assert len(fake_mtml.shutdown_calls) == 1 + assert pymtml._mtmlLib_refcount == 0 + assert not bool(pymtml.libHandle) + + +def test_library_can_reinitialize_after_final_shutdown(fake_mtml): + pymtml.mtmlLibraryInit() + pymtml.mtmlLibraryShutDown() + pymtml.mtmlLibraryInit() + + assert len(fake_mtml.init_calls) == 2 + assert len(fake_mtml.shutdown_calls) == 1 + assert pymtml._mtmlLib_refcount == 1 + assert bool(pymtml.libHandle) + assert pymtml.mtmlLibraryCountDevice() == 2 + + pymtml.mtmlLibraryShutDown() + + assert len(fake_mtml.shutdown_calls) == 2 + assert pymtml._mtmlLib_refcount == 0 + assert not bool(pymtml.libHandle) + + +def test_shutdown_without_reference_is_idempotent(fake_mtml): + pymtml.mtmlLibraryShutDown() + + assert fake_mtml.shutdown_calls == [] + assert pymtml._mtmlLib_refcount == 0 + assert not bool(pymtml.libHandle) + + +def test_failed_native_init_does_not_publish_handle_or_reference(fake_mtml): + fake_mtml.init_result = pymtml.MTML_ERROR_UNKNOWN + + with pytest.raises(pymtml.MTMLError): + pymtml.mtmlLibraryInit() + + assert len(fake_mtml.init_calls) == 1 + assert pymtml._mtmlLib_refcount == 0 + assert not bool(pymtml.libHandle) + + +def test_failed_final_shutdown_preserves_handle_and_reference(fake_mtml): + pymtml.mtmlLibraryInit() + active_handle = pymtml.libHandle + fake_mtml.shutdown_result = pymtml.MTML_ERROR_UNKNOWN + + with pytest.raises(pymtml.MTMLError): + pymtml.mtmlLibraryShutDown() + + assert len(fake_mtml.shutdown_calls) == 1 + assert pymtml._mtmlLib_refcount == 1 + assert pymtml.libHandle == active_handle + + fake_mtml.shutdown_result = pymtml.MTML_SUCCESS + pymtml.mtmlLibraryShutDown() + + assert len(fake_mtml.shutdown_calls) == 2 + assert pymtml._mtmlLib_refcount == 0 + assert not bool(pymtml.libHandle) + + +def test_concurrent_clients_share_one_native_lifecycle(fake_mtml): + worker_count = 8 + start = Barrier(worker_count) + initialized = Barrier(worker_count) + queried = Barrier(worker_count) + + def use_library(): + start.wait() + pymtml.mtmlLibraryInit() + initialized.wait() + assert pymtml.mtmlLibraryCountDevice() == 2 + queried.wait() + pymtml.mtmlLibraryShutDown() + + with ThreadPoolExecutor(max_workers=worker_count) as executor: + futures = [executor.submit(use_library) for _ in range(worker_count)] + for future in futures: + future.result() + + assert len(fake_mtml.init_calls) == 1 + assert len(fake_mtml.shutdown_calls) == 1 + assert pymtml._mtmlLib_refcount == 0 + assert not bool(pymtml.libHandle) From 9aeb2f608578e29d09411166fef9514b4e8d7552 Mon Sep 17 00:00:00 2001 From: CYJiang Date: Tue, 13 Oct 2026 18:15:10 +0800 Subject: [PATCH 2/4] fix: preserve MTML handles across nested lifecycle calls Signed-off-by: CYJiang --- pymtml.py | 36 ++++++++------- test_library_lifecycle.py | 95 +++++++++++++++++++++++++++------------ 2 files changed, 88 insertions(+), 43 deletions(-) diff --git a/pymtml.py b/pymtml.py index 78cc298..1c89b17 100644 --- a/pymtml.py +++ b/pymtml.py @@ -757,10 +757,12 @@ def timestamp(self): ## Lib loading ## mtmlLib = None libLoadLock = threading.Lock() -# The Python API exposes one process-wide handle. Keep its native lifetime -# balanced across nested and concurrent clients independently of library loading. +# Keep init/shutdown flow aligned with pynvml. MTML additionally returns one +# opaque handle per init, so retain every native handle while keeping the first +# live handle stable for implicit query APIs. libLifecycleLock = threading.Lock() libHandle = c_mtmlLibrary_t() +_mtmlLib_handles = [] _mtmlLib_refcount = 0 # Incremented on each mtmlInit and decremented on mtmlShutdown @@ -988,32 +990,36 @@ def mtmlLibraryInit(): fn = _mtmlGetFunctionPointer("mtmlLibraryInit") global libHandle + global _mtmlLib_handles global _mtmlLib_refcount with libLifecycleLock: - if _mtmlLib_refcount == 0: - handle = c_mtmlLibrary_t() - ret = fn(byref(handle)) - _mtmlCheckReturn(ret) + handle = c_mtmlLibrary_t() + ret = fn(byref(handle)) + _mtmlCheckReturn(ret) + _mtmlLib_handles.append(handle) + if len(_mtmlLib_handles) == 1: libHandle = handle - _mtmlLib_refcount += 1 + _mtmlLib_refcount = len(_mtmlLib_handles) return None def mtmlLibraryShutDown(): global libHandle + global _mtmlLib_handles global _mtmlLib_refcount with libLifecycleLock: - if _mtmlLib_refcount == 0: - return None - if _mtmlLib_refcount > 1: - _mtmlLib_refcount -= 1 - return None + if not _mtmlLib_handles: + raise MTMLError(MTML_ERROR_UNINITIALIZED) + handle = _mtmlLib_handles[-1] fn = _mtmlGetFunctionPointer("mtmlLibraryShutDown") - ret = fn(libHandle) + ret = fn(handle) _mtmlCheckReturn(ret) - libHandle = c_mtmlLibrary_t() - _mtmlLib_refcount = 0 + _mtmlLib_handles.pop() + _mtmlLib_refcount = len(_mtmlLib_handles) + libHandle = ( + _mtmlLib_handles[0] if _mtmlLib_handles else c_mtmlLibrary_t() + ) return None diff --git a/test_library_lifecycle.py b/test_library_lifecycle.py index d428e8e..f2bff3e 100644 --- a/test_library_lifecycle.py +++ b/test_library_lifecycle.py @@ -1,7 +1,7 @@ """Unit tests for process-wide MTML library lifecycle ownership.""" from concurrent.futures import ThreadPoolExecutor -from ctypes import POINTER, c_uint, cast, pointer +from ctypes import POINTER, c_uint, c_void_p, cast, pointer from threading import Barrier import pytest @@ -9,10 +9,15 @@ import pymtml +def handle_address(handle): + return cast(handle, c_void_p).value + + class FakeMtmlLibrary: def __init__(self): self.init_calls = [] self.shutdown_calls = [] + self.query_calls = [] self._handles = [] self.init_result = pymtml.MTML_SUCCESS self.shutdown_result = pymtml.MTML_SUCCESS @@ -22,14 +27,15 @@ def init(self, handle_out): handle = pointer(handle_storage) cast(handle_out, POINTER(pymtml.c_mtmlLibrary_t))[0] = handle self._handles.append((handle_storage, handle)) - self.init_calls.append(handle) + self.init_calls.append(handle_address(handle)) return self.init_result def shutdown(self, handle): - self.shutdown_calls.append(handle) + self.shutdown_calls.append(handle_address(handle)) return self.shutdown_result - def count_devices(self, _handle, count_out): + def count_devices(self, handle, count_out): + self.query_calls.append(handle_address(handle)) cast(count_out, POINTER(c_uint))[0] = 2 return pymtml.MTML_SUCCESS @@ -46,90 +52,120 @@ def fake_mtml(monkeypatch): monkeypatch.setattr(pymtml, "_LoadMtmlLibrary", lambda: None) monkeypatch.setattr(pymtml, "_mtmlGetFunctionPointer", functions.__getitem__) monkeypatch.setattr(pymtml, "libHandle", pymtml.c_mtmlLibrary_t()) + monkeypatch.setattr(pymtml, "_mtmlLib_handles", [], raising=False) monkeypatch.setattr(pymtml, "_mtmlLib_refcount", 0) return native -def test_nested_init_keeps_native_handle_until_last_shutdown(fake_mtml): +def test_nested_init_pairs_handles_and_keeps_first_handle_active(fake_mtml): pymtml.mtmlLibraryInit() pymtml.mtmlLibraryInit() - assert len(fake_mtml.init_calls) == 1 + assert len(fake_mtml.init_calls) == 2 + assert len(set(fake_mtml.init_calls)) == 2 + assert [handle_address(handle) for handle in pymtml._mtmlLib_handles] == ( + fake_mtml.init_calls + ) + assert handle_address(pymtml.libHandle) == fake_mtml.init_calls[0] assert pymtml._mtmlLib_refcount == 2 - assert bool(pymtml.libHandle) pymtml.mtmlLibraryShutDown() - assert fake_mtml.shutdown_calls == [] + assert fake_mtml.shutdown_calls == [fake_mtml.init_calls[1]] + assert [handle_address(handle) for handle in pymtml._mtmlLib_handles] == [ + fake_mtml.init_calls[0] + ] + assert handle_address(pymtml.libHandle) == fake_mtml.init_calls[0] assert pymtml._mtmlLib_refcount == 1 - assert bool(pymtml.libHandle) assert pymtml.mtmlLibraryCountDevice() == 2 + assert fake_mtml.query_calls == [fake_mtml.init_calls[0]] pymtml.mtmlLibraryShutDown() - assert len(fake_mtml.shutdown_calls) == 1 + assert fake_mtml.shutdown_calls == list(reversed(fake_mtml.init_calls)) + assert pymtml._mtmlLib_handles == [] assert pymtml._mtmlLib_refcount == 0 assert not bool(pymtml.libHandle) def test_library_can_reinitialize_after_final_shutdown(fake_mtml): pymtml.mtmlLibraryInit() + first_handle = handle_address(pymtml.libHandle) pymtml.mtmlLibraryShutDown() pymtml.mtmlLibraryInit() + second_handle = handle_address(pymtml.libHandle) + assert first_handle != second_handle assert len(fake_mtml.init_calls) == 2 assert len(fake_mtml.shutdown_calls) == 1 assert pymtml._mtmlLib_refcount == 1 - assert bool(pymtml.libHandle) assert pymtml.mtmlLibraryCountDevice() == 2 pymtml.mtmlLibraryShutDown() assert len(fake_mtml.shutdown_calls) == 2 + assert pymtml._mtmlLib_handles == [] assert pymtml._mtmlLib_refcount == 0 assert not bool(pymtml.libHandle) -def test_shutdown_without_reference_is_idempotent(fake_mtml): - pymtml.mtmlLibraryShutDown() +def test_shutdown_without_init_reports_uninitialized(fake_mtml): + with pytest.raises(pymtml.MTMLError) as exc_info: + pymtml.mtmlLibraryShutDown() + assert exc_info.value.value == pymtml.MTML_ERROR_UNINITIALIZED assert fake_mtml.shutdown_calls == [] + assert pymtml._mtmlLib_handles == [] assert pymtml._mtmlLib_refcount == 0 assert not bool(pymtml.libHandle) -def test_failed_native_init_does_not_publish_handle_or_reference(fake_mtml): +def test_failed_initial_init_does_not_publish_handle(fake_mtml): fake_mtml.init_result = pymtml.MTML_ERROR_UNKNOWN with pytest.raises(pymtml.MTMLError): pymtml.mtmlLibraryInit() assert len(fake_mtml.init_calls) == 1 + assert pymtml._mtmlLib_handles == [] assert pymtml._mtmlLib_refcount == 0 assert not bool(pymtml.libHandle) -def test_failed_final_shutdown_preserves_handle_and_reference(fake_mtml): +def test_failed_nested_init_preserves_existing_handle(fake_mtml): pymtml.mtmlLibraryInit() - active_handle = pymtml.libHandle - fake_mtml.shutdown_result = pymtml.MTML_ERROR_UNKNOWN + active_handle = handle_address(pymtml.libHandle) + fake_mtml.init_result = pymtml.MTML_ERROR_UNKNOWN with pytest.raises(pymtml.MTMLError): - pymtml.mtmlLibraryShutDown() + pymtml.mtmlLibraryInit() - assert len(fake_mtml.shutdown_calls) == 1 + assert len(fake_mtml.init_calls) == 2 + assert [handle_address(handle) for handle in pymtml._mtmlLib_handles] == [ + active_handle + ] + assert handle_address(pymtml.libHandle) == active_handle assert pymtml._mtmlLib_refcount == 1 - assert pymtml.libHandle == active_handle - fake_mtml.shutdown_result = pymtml.MTML_SUCCESS - pymtml.mtmlLibraryShutDown() - assert len(fake_mtml.shutdown_calls) == 2 - assert pymtml._mtmlLib_refcount == 0 - assert not bool(pymtml.libHandle) +def test_failed_shutdown_preserves_handle_stack(fake_mtml): + pymtml.mtmlLibraryInit() + pymtml.mtmlLibraryInit() + handles_before = [handle_address(handle) for handle in pymtml._mtmlLib_handles] + fake_mtml.shutdown_result = pymtml.MTML_ERROR_UNKNOWN + with pytest.raises(pymtml.MTMLError): + pymtml.mtmlLibraryShutDown() + + assert fake_mtml.shutdown_calls == [handles_before[-1]] + assert [handle_address(handle) for handle in pymtml._mtmlLib_handles] == ( + handles_before + ) + assert handle_address(pymtml.libHandle) == handles_before[0] + assert pymtml._mtmlLib_refcount == 2 -def test_concurrent_clients_share_one_native_lifecycle(fake_mtml): + +def test_concurrent_clients_pair_all_native_handles(fake_mtml): worker_count = 8 start = Barrier(worker_count) initialized = Barrier(worker_count) @@ -148,7 +184,10 @@ def use_library(): for future in futures: future.result() - assert len(fake_mtml.init_calls) == 1 - assert len(fake_mtml.shutdown_calls) == 1 + assert len(fake_mtml.init_calls) == worker_count + assert len(set(fake_mtml.init_calls)) == worker_count + assert fake_mtml.shutdown_calls == list(reversed(fake_mtml.init_calls)) + assert set(fake_mtml.query_calls) == {fake_mtml.init_calls[0]} + assert pymtml._mtmlLib_handles == [] assert pymtml._mtmlLib_refcount == 0 assert not bool(pymtml.libHandle) From d5242fec15b7edb46484a2ab846d2ea1f990d5bf Mon Sep 17 00:00:00 2001 From: CYJiang Date: Wed, 14 Oct 2026 11:23:35 +0800 Subject: [PATCH 3/4] test: bound lifecycle barrier waits Signed-off-by: CYJiang --- test_library_lifecycle.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/test_library_lifecycle.py b/test_library_lifecycle.py index f2bff3e..145113c 100644 --- a/test_library_lifecycle.py +++ b/test_library_lifecycle.py @@ -172,11 +172,11 @@ def test_concurrent_clients_pair_all_native_handles(fake_mtml): queried = Barrier(worker_count) def use_library(): - start.wait() + start.wait(timeout=10) pymtml.mtmlLibraryInit() - initialized.wait() + initialized.wait(timeout=10) assert pymtml.mtmlLibraryCountDevice() == 2 - queried.wait() + queried.wait(timeout=10) pymtml.mtmlLibraryShutDown() with ThreadPoolExecutor(max_workers=worker_count) as executor: From 52995256db44b72795288b970352f61a7d5cacf7 Mon Sep 17 00:00:00 2001 From: CYJiang Date: Wed, 14 Oct 2026 13:41:00 +0800 Subject: [PATCH 4/4] ci: run hardware-independent pytest suite Signed-off-by: CYJiang --- .github/workflows/main.yml | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml index 91527eb..4944194 100644 --- a/.github/workflows/main.yml +++ b/.github/workflows/main.yml @@ -10,6 +10,14 @@ jobs: run: mthreads-gmi - name: Checkout code uses: actions/checkout@v4 + - name: Test hardware-independent units + run: | + python -m pytest \ + test_abi_layout.py \ + test_library_lifecycle.py \ + test_library_loading.py \ + test_mtlink_visibility.py \ + -q - name: Test mtml sglang compatibility run: python test_sglang_compat.py - name: Test mtml @@ -18,4 +26,4 @@ jobs: python test_pymtml.py - name: Test pynvml run: | - python test_pynvml.py \ No newline at end of file + python test_pynvml.py