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 diff --git a/pymtml.py b/pymtml.py index 75c6bc3..1c89b17 100644 --- a/pymtml.py +++ b/pymtml.py @@ -757,7 +757,12 @@ def timestamp(self): ## Lib loading ## mtmlLib = None libLoadLock = threading.Lock() +# 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 @@ -983,44 +988,38 @@ 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_handles global _mtmlLib_refcount - libLoadLock.acquire() - _mtmlLib_refcount += 1 - libLoadLock.release() + with libLifecycleLock: + handle = c_mtmlLibrary_t() + ret = fn(byref(handle)) + _mtmlCheckReturn(ret) + _mtmlLib_handles.append(handle) + if len(_mtmlLib_handles) == 1: + libHandle = handle + _mtmlLib_refcount = len(_mtmlLib_handles) 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_handles global _mtmlLib_refcount - libLoadLock.acquire() - if 0 < _mtmlLib_refcount: - _mtmlLib_refcount -= 1 - libLoadLock.release() + with libLifecycleLock: + if not _mtmlLib_handles: + raise MTMLError(MTML_ERROR_UNINITIALIZED) + + handle = _mtmlLib_handles[-1] + fn = _mtmlGetFunctionPointer("mtmlLibraryShutDown") + ret = fn(handle) + _mtmlCheckReturn(ret) + _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 new file mode 100644 index 0000000..145113c --- /dev/null +++ b/test_library_lifecycle.py @@ -0,0 +1,193 @@ +"""Unit tests for process-wide MTML library lifecycle ownership.""" + +from concurrent.futures import ThreadPoolExecutor +from ctypes import POINTER, c_uint, c_void_p, cast, pointer +from threading import Barrier + +import pytest + +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 + + 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_address(handle)) + return self.init_result + + def shutdown(self, handle): + self.shutdown_calls.append(handle_address(handle)) + return self.shutdown_result + + 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 + + +@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_handles", [], raising=False) + monkeypatch.setattr(pymtml, "_mtmlLib_refcount", 0) + return native + + +def test_nested_init_pairs_handles_and_keeps_first_handle_active(fake_mtml): + pymtml.mtmlLibraryInit() + pymtml.mtmlLibraryInit() + + 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 + + pymtml.mtmlLibraryShutDown() + + 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 pymtml.mtmlLibraryCountDevice() == 2 + assert fake_mtml.query_calls == [fake_mtml.init_calls[0]] + + pymtml.mtmlLibraryShutDown() + + 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 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_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_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_nested_init_preserves_existing_handle(fake_mtml): + pymtml.mtmlLibraryInit() + active_handle = handle_address(pymtml.libHandle) + fake_mtml.init_result = pymtml.MTML_ERROR_UNKNOWN + + with pytest.raises(pymtml.MTMLError): + pymtml.mtmlLibraryInit() + + 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 + + +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_pair_all_native_handles(fake_mtml): + worker_count = 8 + start = Barrier(worker_count) + initialized = Barrier(worker_count) + queried = Barrier(worker_count) + + def use_library(): + start.wait(timeout=10) + pymtml.mtmlLibraryInit() + initialized.wait(timeout=10) + assert pymtml.mtmlLibraryCountDevice() == 2 + queried.wait(timeout=10) + 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) == 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)