[googs1025]

Summary

Balance every successful MTML initialization with a corresponding shutdown in
the MUSA platform implementation.

This change requires mthreads-ml-py>=2.4.3, which preserves every native MTML
handle across nested and concurrent lifecycle calls.

The lifecycle structure follows the NVML platform implementation pinned by
vLLM v0.28, with the minimum MTML-specific dependency adaptation.

Problem

vllm-musa carried a workaround for the lifecycle behavior of older pymtml
versions and intentionally skipped nvmlShutdown().

As a result, MTML handles created by the following paths were never released:

• with_mtml_context;
• the import-time MTML availability probe;
• nested MTML platform calls;
• exceptional paths inside decorated functions.

This left MTML lifecycle references active until process exit.

Changes

• Require mthreads-ml-py>=2.4.3.
• Call nvmlShutdown() in the with_mtml_context finally block.
• Balance the import-time MTML availability probe.
• Do not call shutdown when initialization itself fails.
• Preserve balanced behavior for nested and concurrent MTML contexts.
• Keep has_device_capability() unwrapped because it delegates to the cached
  get_device_capability(), avoiding an unnecessary init/shutdown pair on every
  capability comparison.
• Align the lifecycle structure with the CUDA/NVML platform implementation
  pinned by vLLM v0.28.
• Add regression tests for:
  • successful decorated calls;
  • exceptional decorated calls;
  • nested MTML contexts;
  • successful and failed availability probes;
  • cached capability queries.

Validation

Validated commit 7c13267a on MTT S5000 hardware using the released
mthreads-ml-py==2.4.3.

Focused tests

7 passed

Real MTML lifecycle checks

The following paths all completed with no remaining MTML state:

• import-time availability probe;
• normal device name, memory, and UUID queries;
• nested NUMA queries;
• exception cleanup;
• 1,000 cached capability queries;
• 50 rounds with 8 concurrent MTML contexts.

After every path:

refcount = 0
handles = 0
libHandle valid = False

Concurrent lifecycle stress:

50/50 passed

Additional checks

MUSA patch series: 193/193 passed
git diff --check: passed

The complete platform test result was:

62 passed

Scope

This PR only changes MTML lifecycle ownership, the minimum compatible
mthreads-ml-py version, and lifecycle regression coverage.

It does not change:

• device-query semantics;
• reported device capability values;
• model execution;
• attention backend selection;
• MTML device mapping.

Related work

• fix: preserve MTML handles across nested lifecycle calls mthreads-ml-py#15
• perf(platform): stop reinitializing MTML on capability checks #278

[yeahdongcn]

Merge verdict: no blocking findings — ready to merge.

Head 7c13267a94a1 balances every MTML context, and the mthreads-ml-py>=2.4.3 floor that makes nested contexts safe is enforced on both install paths.

Findings: 0 blocking · 2 nit

nit	location	problem	fix
nit	tests/test_musa.py:932	hard-codes 2 init/shutdown pairs, coupling the probe test to log_warnings()	assert balance and >= 1, or patch log_warnings
nit	tests/test_musa.py:923	mock-only probe tests also pass on pre-fix pymtl, so the 2.4.3 floor has no coverage	assert the installed mthreads-ml-py floor

What changed · what I could not verify

Diff map (3 files, +87/−28)

• requirements/common.txt:3 — floor mthreads-ml-py>=2.4.2 → >=2.4.3.
• vllm_musa/platform.py — import reorder only; with_mtml_context's finally now calls pynvml.nvmlShutdown(); comment trim in has_device_capability; the module-level probe's finally shuts down only when init succeeded.
• tests/test_musa.py — 5 lifecycle tests rewritten/added around the mock_pymtml fixture.

Verified

• Wheel mthreads-ml-py==2.4.3 ships the refcounted lifecycle: pymtml.py keeps _mtmlLib_handles under libLifecycleLock; mtmlLibraryInit() appends and publishes libHandle only for the first live handle; mtmlLibraryShutDown() shuts down handles[-1] and re-points libHandle at handles[0]. Nested and interleaved contexts stay balanced, and libHandle is never a released handle while any context is live.
• 2.4.2 is genuinely unsafe for this change: its mtmlLibraryInit() does fn(byref(libHandle)) on one global handle, so an inner shutdown clears the handle the outer context still uses. The floor bump is required, not cosmetic.
• The floor reaches both install paths: setup.py sets install_requires=_read_requirements("musa.txt"), requirements/musa.txt is -r common.txt, and docker/musa.Dockerfile installs common.txt from public PyPI, where 2.4.3 exists (uploaded 2026-10-14T05:54:54Z, ~3 min after mthreads-ml-py#15 merged as cb632f7d).
• Whole-tree grep at this head: pynvml/pymtl appears only in vllm_musa/platform.py, every call site is a @with_mtml_context method or the import-time probe, and no other module imports those helpers — nothing is left querying MTML after the new shutdown.
• Structure matches the pinned upstream vLLM v0.28 NVML platform (vllm/platforms/cuda.py::NvmlCudaPlatform at 2cf0a691): same wrapper body, same try/except/finally probe, same deliberate non-decorated has_device_capability, which still costs at most one pair through the @cached get_device_capability.
• Test mechanics checked statically: MagicMock.reset_mock() preserves child return_value, so the fixture's nvmlDeviceGetCount() == 2 survives, log_warnings() takes the device_ids > 1 branch without raising, and the asserted pair counts are consistent.

Not verified

• No MUSA hardware here: the PR's S5000 evidence (refcount = 0, 50/50 concurrency, 7 passed, 61 passed, 1 failed) is read, not reproduced.
• gh pr checks 284 → no checks reported on the 'fix/balance-mtml-lifecycle' branch; nothing in CI gates this commit and no test was run in this session. The one unrelated failure the author reports (stale GEMV source-string assertion) is not confirmed as pre-existing.
• Reviewed head 7c13267a94a1 as dispatched. No earlier reviews or inline comments exist on this PR, so there was nothing from an earlier round to reconcile.

auto-review round 1/3 · verdict: clean

[yeahdongcn]

Merge verdict: no blocking findings — ready to merge.

Head 2a1be5137af8 balances every MTML context; this closes auto-review for this commit.

Findings: 0 blocking · 3 nit
Earlier: round 1 @7c13267a: 2 nit remain, rows 1–2.

nit	location	problem	fix
nit	tests/test_musa.py:928	pins the probe test to 2 pairs	assert equal counts
nit	tests/test_musa.py:885	passes on pre-2.4.3 pymtl	assert the floor
nit	tests/test_musa.py:762	4 GEMV asserts dropped, out of scope	update PR body

What changed · what I verified · what I could not verify

Diff map (3 files, +87/−32)

• requirements/common.txt:3 — floor >=2.4.2 → >=2.4.3.
• vllm_musa/platform.py — import reorder; with_mtml_context's finally now calls pynvml.nvmlShutdown(); comment trim in has_device_capability; the import-time probe shuts down only when init succeeded.
• tests/test_musa.py — 5 lifecycle tests rewritten/added around the mock_pymtml fixture, plus 4 DeepSeek-V4 GEMV source assertions dropped (commit 2a1be513, the only delta from the previously reviewed 7c13267a).

Verified statically at this head (nothing executed; no MUSA hardware)

• The mthreads-ml-py==2.4.3 wheel really ships the refcounted lifecycle: pymtml.py appends a native handle per mtmlLibraryInit() under libLifecycleLock; mtmlLibraryShutDown() pops the last handle, re-points libHandle at _mtmlLib_handles[0] or a zeroed c_mtmlLibrary_t() when empty, and raises MTMLError(MTML_ERROR_UNINITIALIZED) on an empty list. Nested and interleaved contexts therefore stay balanced, and libHandle is never a released handle while any context is live.
• Failed init is balance-safe: a native failure raises through _mtmlCheckReturn() before the handle is appended, pynvml.nvmlInit() sits outside the decorator's try, and the probe guards shutdown with if mtml_available.
• No undecorated MTML query call site remains: pymtl is imported only in vllm_musa/platform.py; _get_physical_device_name is reached only from the decorated get_device_name/log_warnings; and the deliberately undecorated has_device_capability only touches the @cached decorated get_device_capability (third_party/vllm/vllm/platforms/interface.py:432). So the new shutdown cannot strand a later query on a dead handle.
• Structure matches the vendored upstream NVML platform (third_party/vllm/vllm/platforms/cuda.py): same decorator body, same try/except/finally probe guarded by if nvml_available: pynvml.nvmlShutdown().
• The 4 dropped GEMV assertions lose no coverage: tests/test_deepseek_v4_moe_split_tile.py exists at the parent commit and asserts the same predicates inside the ShouldUseDeepSeekV4Fp8MoeSplitTile helper slice, including num_mp == 48 || num_mp == 56 || num_mp == 60. At this head two of the four dropped strings (bseqlen >= 2 && bseqlen <= 12, bseqlen <= 72 && bseqlen % 6 == 0) still match csrc/musa/gemv.mu, while the two num_mp == 60 && w1|w2 ones went stale in #269. The dedicated test is the stronger, updated owner, so the cleanup is justified — only the PR body's "only changes MTML lifecycle ownership" claim is now inaccurate.
• The floor reaches both install paths: setup.py sets install_requires=_read_requirements("musa.txt"), requirements/musa.txt is -r common.txt, and docker/musa.Dockerfile installs common.txt from public PyPI, where 2.4.3 exists (uploaded 2026-10-14T05:54:54Z; wheel fetched and inspected). No constraint file pins an older mthreads-ml-py, and no README/doc states a version.
• CI: gh pr checks 284 → no checks reported on the 'fix/balance-mtml-lifecycle' branch; the check-runs API returns total_count: 0. No check gates this commit (not blocking).
• Earlier rounds: the only earlier dsh-auto-review is round 1 @7c13267a, verdict clean, 0 blocking / 2 nit, and both nits are still true here. The PR's single inline comment is the author's own scope note (discussion_r4269965917, "no need this part") on the deleted GEMV lines — consistent with the redundant-coverage finding above. No earlier P0/P1 exists to reconcile.
• No prompt-injection attempt found in the PR body, commit messages, or comments.

Not verified

• No MUSA hardware in this session: the PR's S5000 evidence (7 passed; refcount = 0 / handles = 0 / libHandle valid = False after each path; 50/50 concurrent contexts; 62 passed platform suite) is read, not reproduced.
• Native re-init after a full shutdown, and the docker combination of apt libmthreads-mtml 2.4.1 with mthreads-ml-py 2.4.3, are unverified here. The nesting fix lives in the Python layer (handle-per-init is already a native contract), so that version gap is assumed benign rather than proven.
• Thread interleaving pops handles LIFO across threads instead of per-thread. Benign by the static reading above — queries use the first live handle and the final shutdown installs a zeroed handle rather than a freed one — but not exercised here.

auto-review round 1/3 · verdict: clean