fix: preserve MTML handles across nested lifecycle calls

.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
+          python test_pynvml.py

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
 
 

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)