[googs1025]

Summary

Fix nested and concurrent MTML lifecycle handling in the Python binding.

This change:

• calls native mtmlLibraryInit() for every successful Python init;
• retains every opaque MTML library handle returned by native init;
• calls native mtmlLibraryShutDown(handle) once for each Python shutdown;
• keeps the first live handle stable for existing implicit query APIs;
• releases additional retained handles before the stable active handle;
• serializes lifecycle changes with a dedicated lock;
• updates Python state only after the native operation succeeds;
• preserves the handle list and active handle when native init or shutdown fails.
• runs the hardware-independent lifecycle suite in CI.

Problem

MTML and NVML expose different native lifecycle models.

NVML exposes process-global initialization APIs:

nvmlInit()
nvmlShutdown()

MTML returns an opaque library handle for each initialization:

MtmlReturn mtmlLibraryInit(MtmlLibrary **lib);
MtmlReturn mtmlLibraryShutDown(MtmlLibrary *lib);

The previous Python binding called native mtmlLibraryInit() repeatedly but
stored every result in one global libHandle:

init #1 -> handle1
init #2 -> handle2, overwriting handle1
shutdown #1 -> releases handle2 and clears libHandle

At that point a Python lifecycle reference still existed, but the binding had
lost handle1 and no longer had a valid active handle.

MTML native behavior

The MTML implementation already supports multiple handles:

• every mtmlLibraryInit() returns a new MtmlLibrary handle;
• those handles share the underlying library ownership;
• every returned handle must be matched by a native shutdown;
• shutdown order does not need to match initialization order.

The issue is therefore in Python handle management, not in the MTML C API.

Alignment with pynvml

The public lifecycle control flow follows pynvml:

1. every public init calls the native init API;
2. the native result is checked before Python bookkeeping is updated;
3. every public shutdown calls the native shutdown API;
4. Python bookkeeping is updated only after native shutdown succeeds.

MTML additionally returns one opaque handle per init. The binding must retain
those handles while providing one stable libHandle for existing query APIs
that do not accept an explicit library handle.

Implementation

The binding now maintains all successfully initialized native handles:

_mtmlLib_handles = []

On init:

native init into a local handle
-> validate the native result
-> append the returned handle
-> publish it as libHandle only when it is the first live handle

On shutdown, when additional handles exist:

select the newest additional handle
-> native shutdown(handle)
-> validate the native result
-> remove the handle
-> keep the first live handle unchanged as libHandle

When only the active handle remains:

native shutdown(active handle)
-> validate the native result
-> clear the handle list and libHandle

Lifecycle state transitions are protected by libLifecycleLock. A failed
native init or shutdown leaves _mtmlLib_handles, libHandle, and
_mtmlLib_refcount unchanged.

Keeping the first handle stable also prevents a concurrent implicit query from
using an additional handle while another client releases that handle.

Tests

Hardware-independent

22 passed

Coverage includes:

• two init calls create two distinct native handles;
• nested shutdown pairs and releases every native handle;
• the first live handle remains stable for implicit query APIs;
• reinitialization after final shutdown;
• shutdown without initialization reports MTML_ERROR_UNINITIALIZED;
• failed initial and nested initialization rollback;
• failed shutdown rollback;
• concurrent lifecycle clients;
• bounded barrier waits so concurrency regressions fail instead of hanging.

Concurrent lifecycle stress:

50/50 passed
8 clients per iteration

Additional checks:

Python syntax check: passed
git diff --check: passed
black/isort/flake8 for the lifecycle test: passed

Additional MUSA compatibility validation

Lifecycle implementation tested at:
9aeb2f608578e29d09411166fef9514b4e8d7552

The current head only adds timeouts to the test barriers; it does not change
the lifecycle implementation.

Hardware: MTT S5000

MUSA SDK/runtime	Native C++ lifecycle	Python lifecycle	Compatibility scripts
4.3.x	Passed: sequential, nested/non-LIFO, and concurrent lifecycle calls	Passed: nested handles, shutdown/reinit, and concurrent clients	All passed
5.1.x	Passed: sequential, nested/non-LIFO, and concurrent lifecycle calls	Passed: nested handles, shutdown/reinit, and concurrent clients	All passed
5.2.x	Passed: sequential, nested/non-LIFO, and concurrent lifecycle calls	Passed: nested handles, shutdown/reinit, and concurrent clients	All passed
5.3.x (current master)	Passed: sequential, nested/non-LIFO, and concurrent lifecycle calls	Passed: nested handles, shutdown/reinit, fork, and concurrent clients	All passed

Compatibility scripts:

• python3 test_pymtml.py
• python3 test_pynvml.py
• python3 test_sglang_compat.py

The native C++ tests called libmtml directly without going through
pymtml. The Python tests verified that every successful initialization
returns and retains a distinct native handle, every handle is paired with one
native shutdown, and the stable implicit-query handle remains valid until the
final shutdown.

Full pytest on the S5000 environment:

24 passed, 8 errors

The remaining errors are pre-existing pytest collection issues: script-style
helper functions in test_pymtml.py, test_pynvml.py, and
test_sglang_compat.py are collected as fixture-based tests. There were no
functional test failures.

No new GPU fault, XID, timeout, kernel Oops, or call trace was observed during
the validation window.

Scope

This PR changes:

• pymtml.py
• test_library_lifecycle.py
• .github/workflows/main.yml, to run the hardware-independent tests in CI

It does not modify the MTML C API or the existing hardware compatibility
scripts. The CI step uses an explicit unit-test list because the script-style
hardware helpers are not pytest fixture tests and cannot safely be collected by
the broad test_*.py pattern.

Related work

• perf(platform): stop reinitializing MTML on capability checks vllm-musa#278

[googs1025]

cc @yeahdongcn /PTAL

[yeahdongcn]

Merge verdict: no blocking findings — ready to merge.
Head 9aeb2f6 verified against the native contract in mtml_2.2.0.h: one handle per init, one shutdown per handle, released in any order.

Findings: 0 blocking · 2 nit

nit	location	problem	fix
nit	test_library_lifecycle.py:174	barrier workers wait forever when one assertion fails, so a regression hangs make test instead of failing it	give every Barrier.wait() a timeout=
nit	.github/workflows/main.yml:13	CI runs three scripts rather than pytest, so this new test file never runs there	add a python -m pytest test_*.py -q step

What changed · what I could not verify

The contract this now follows. mtml_2.2.0.h:865-866 — several mtmlLibraryInit() pointers keep the library alive "until the corresponding number of mtmlLibraryShutDown() calls have been made with those pointers", and the release order "does not necessarily have to match" the init order. mtml_2.2.0.h:872 — "Providing this function with a lib that has already been initialized causes memory leak". The old code passed the same libHandle into every init, which is exactly that warning; allocating one handle per init and pinning libHandle to the first live one fixes the leak and keeps the implicit query APIs (mtmlLibraryCountDevice and friends) on a handle that cannot have been released, because the stack only ever pops its last entry.

Verified.

• python -m pytest test_library_lifecycle.py -v at this head: 7 passed in 0.03s. The fixture replaces _LoadMtmlLibrary and _mtmlGetFunctionPointer, so the tests are hardware-independent as claimed, and monkeypatch restores the module globals afterwards.
• No existing caller regresses. test_pymtml.py:376-383 runs three balanced init/shutdown cycles and both teardowns swallow exceptions (:405-412); example.py:29/57 pairs; test_pynvml.py:338 shuts down before re-initializing, which main()'s leading nvmlInit() keeps balanced, and that worker catches Exception. An unbalanced shutdown used to reach the native call with a NULL handle — MTML_ERROR_INVALID_ARGUMENT per the header — so it already raised; MTML_ERROR_UNINITIALIZED is a more accurate code, not a new failure mode.
• Locking: the native init/shutdown call itself is now serialized by libLifecycleLock (pymtml.py:995-999), while _LoadMtmlLibrary and _mtmlGetFunctionPointer stay outside it, so the libLifecycleLock → libLoadLock order has no inverse path.

Not verified. No MUSA hardware run, no driver smoke; the change is exercised only against the fake library. Native behaviour after a failed shutdown and after a release that follows a failed init is assumed, not observed. gh pr checks 15 --repo MooreThreads/mthreads-ml-py reports no checks reported on the 'fix/mtml-lifecycle' branch: the two workflow runs for this branch (2026-10-13T04:41Z, 10:24Z) sit at action_required on the self-hosted runner, and that workflow would not run pytest even when approved — see the second nit.

auto-review round 1/3 · verdict: clean

[yeahdongcn]

@googs1025 Please add more tests (MUSA SDK 4.3.x/5.2.x/5.3.x) to PR description. Thanks!

[googs1025]

Updated the PR description with additional S5000 validation across MUSA 4.3.x, 5.1.x, 5.2.x, and 5.3.x.

The validation includes direct native C++ lifecycle tests against libmtml, Python nested/concurrent/fork lifecycle tests, and the existing pymtml, pynvml-compatibility, and sglang-compatibility scripts.

[yeahdongcn]

Merge verdict: no blocking findings — ready to merge.
Head d5242fe fixes round 1's barrier hang; pymtml.py is unchanged from the approved commit.

Findings: 0 blocking · 1 nit
Earlier: round 1's hanging barrier test is fixed @d5242fec; the CI nit below still stands

nit	location	problem	fix
nit	.github/workflows/main.yml:13	CI runs three scripts rather than pytest, so this test file never runs there	add a python -m pytest test_*.py -q step

What changed · what I could not verify

What changed since the approved commit. One commit, d5242fec ("test: bound lifecycle barrier waits"), and it is exactly the first nit of round 1: all three Barrier.wait() calls in test_concurrent_clients_pair_all_native_handles now pass timeout=10 (test_library_lifecycle.py:175,177,179). A threading.Barrier timeout breaks the barrier, so every waiting worker raises BrokenBarrierError instead of blocking the executor — a regression now fails the test rather than stalling make test. The timeout is far above the 0.04 s this suite needs. pymtml.py is byte-identical to 9aeb2f608578, so round 1's contract reading still holds: one handle per mtmlLibraryInit() (mtml_2.2.0.h:865-866), released once each in any order, and libHandle pinned to a live entry so the implicit query APIs never address a released object.

Verified. python -m pytest test_library_lifecycle.py -v at this head: 7 passed in 0.04s. The changed lines are the only diff between the two heads; no callers of mtmlLibraryShutDown were touched. gh pr checks 15 --repo MooreThreads/mthreads-ml-py reports no checks reported on the 'fix/mtml-lifecycle' branch.

Not verified. No MUSA hardware run, no driver smoke; native behaviour after a failed shutdown remains assumed, not observed.

auto-review round 1/3 · verdict: clean

[googs1025]

Addressed the remaining CI nit in 5299525. The workflow now runs the hardware-independent pytest suite, including test_library_lifecycle.py.

I used an explicit unit-test list instead of test_*.py because the repository also contains script-style hardware helpers whose test_* functions require positional arguments and are mis-collected as pytest fixtures by the broad glob. The exact CI command passes locally with 22 tests.