[PR #223] [MERGED] fix(api/libnvml): drain in-flight NVML queries before atexit shutdown #222

New issue

Closed

opened 2026-10-02 23:15:06 -06:00 by gitea-mirror · 0 comments

gitea-mirror commented

2026-10-02 23:15:06 -06:00

Owner

📋 Pull Request Information

Original PR: https://github.com/XuehaiPan/nvitop/pull/223
Author: @XuehaiPan
Created: 7/9/2026
Status: ✅ Merged
Merged: 7/10/2026
Merged by: @XuehaiPan

Base: main ← Head: fix/nvml-shutdown-drain

📝 Commits (8)

77b78ea chore(pre-commit): update pre-commit hooks
6888d7e fix(api/libnvml): drain in-flight NVML queries before atexit shutdown
d611b1f fix(api/libnvml): avoid re-initializing NVML during atexit shutdown
051d0ca fix(api/libnvml): guard atexit nvmlShutdown() and register the hook once
6189dab fix(api/libnvml): cover lazy initialization in the shutdown drain
06b5acc fix(api/libnvml): reset the shutdown-drain state after os.fork()
b0821c6 docs(api/libnvml): correct nvmlQuery Raises and the __shutting_down comment
5ade547 docs(changelog): record the atexit NVML shutdown-drain fix

📊 Changes

3 files changed (+162 additions, -44 deletions)

View changed files

📝 .pre-commit-config.yaml (+2 -2)
📝 CHANGELOG.md (+1 -1)
📝 nvitop/api/libnvml.py (+159 -41)

📄 Description

Issue Type

Bug fix

Runtime Environment

Operating system and version: Ubuntu 24.04.3 LTS (x86_64) — affected environment from #222; also verified on another Linux x86_64 host
Terminal emulator and version: N/A (the race is independent of the terminal emulator)
Python version: 3.14.6 (reporter) / 3.12.3 (verification host)
NVML version (driver version): 580.126.20
nvitop version or commit: fix/nvml-shutdown-drain@b0821c6
python-ml-py version: 13.595.45
Locale: N/A

Description

Pressing q to quit the TUI intermittently crashed with a SIGSEGV inside libnvidia-ml. The atexit nvmlShutdown() could free NVML internal state while a background thread was still inside a libnvidia-ml call — a use-after-free. This affects any concurrent NVML user of nvitop.api, including the TUI snapshot daemons and the reproducer in #222 (which spawns query threads with no TUI involved).

The fix makes the implicit atexit shutdown safe rather than dropping it, by draining in-flight queries before the shutdown. All changes are confined to nvitop/api/libnvml.py and to nvitop's own additions on top of pynvml — the re-exported pynvml API is left untouched.

Drain in-flight queries at atexit:

nvmlQuery bails on an unlocked __shutting_down fast path, then reserves an in-flight slot — incrementing __active_queries under a threading.Condition — before any NVML-touching work (including lazy nvmlInit), and re-checks __shutting_down after initialization. The NVML call itself runs outside the lock, so concurrent NVML access is never serialized; the slot is released and the condition notified in a finally.
_atexit_shutdown sets the one-way __shutting_down latch, then waits on the condition for __active_queries to reach zero before calling nvmlShutdown(). Once the latch is set no new query starts, so the counter only decreases and the wait wakes on the last decrement (event-driven, not polling). The wait is bounded (120 s) and skips nvmlShutdown() if queries do not drain in time — the OS reclaims NVML resources at process exit anyway.

Reset the drain state across os.fork():

os.fork() copies the module globals but leaves only the calling thread alive in the child, so the child can inherit a non-zero __active_queries (a query owned by a thread that no longer exists) or an already-held internal lock — making the inherited atexit hook block for its full timeout, or deadlock. _reset_after_fork, registered via os.register_at_fork(after_in_child=...), rebuilds the lock/condition and zeroes the drain state in the child (guarded with hasattr(os, 'register_at_fork') since it is unavailable on Windows).

The re-exported pynvml API is unchanged: nvmlShutdown() keeps its original signature, and an explicit nvmlShutdown() — or a raw libnvml.nvml* call — racing in-flight queries remains the caller's responsibility, since the drain only tracks calls made through nvmlQuery.

The chore(pre-commit) commit bumps the ruff and mypy pre-commit hooks.

Motivation and Context

Fixes #222.

Root cause: the atexit nvmlShutdown() races with in-flight NVML calls from background threads still running at interpreter exit; the shutdown frees NVML internal state under a concurrent call. Draining in-flight nvmlQuery calls before the shutdown closes the race for any concurrent caller — including the no-TUI reproducer in the issue, which a TUI-level fix could not help.

Testing

Verified on a Linux host with 4× NVIDIA H100 80 GB GPUs (Python 3.12.3) using the nvml_race.py reproducer from #222, run 25 times per build:

Build	`nvml_race.py` × 25
`main` (no fix)	15 / 25 `SIGSEGV` (`rc=139`)
`fix/nvml-shutdown-drain`	0 / 25

$ git switch main
$ python3 nvml_race.py
[1]    920641 segmentation fault (core dumped)  python3 nvml_race.py

$ git switch fix/nvml-shutdown-drain
$ python3 nvml_race.py
$   # exits cleanly, no coredump

On main the reproducer segfaults on ~60% of runs; on this branch it never crashes across 25 runs.
The crash is Linux-specific: on Windows, main does not coredump.
The os.fork() reset was verified GPU-free with a real fork(): a child that inherits a non-zero in-flight counter is reset to zero by the registered hook, so its atexit drain returns immediately instead of blocking for the timeout.
The coordination logic was also checked GPU-free during development: nvmlQuery refuses new NVML calls once __shutting_down is set, and _atexit_shutdown waits for an in-flight query before calling nvmlShutdown().
pre-commit (ruff, ruff-format, codespell, mypy, pylint) passes. No automated tests are included (the repository has no test harness yet).

_{🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.}

## 📋 Pull Request Information **Original PR:** https://github.com/XuehaiPan/nvitop/pull/223 **Author:** [@XuehaiPan](https://github.com/XuehaiPan) **Created:** 7/9/2026 **Status:** ✅ Merged **Merged:** 7/10/2026 **Merged by:** [@XuehaiPan](https://github.com/XuehaiPan) **Base:** `main` ← **Head:** `fix/nvml-shutdown-drain` --- ### 📝 Commits (8) - [`77b78ea`](https://github.com/XuehaiPan/nvitop/commit/77b78ea80e5a6046da55fe46be048ff426d08fa1) chore(pre-commit): update pre-commit hooks - [`6888d7e`](https://github.com/XuehaiPan/nvitop/commit/6888d7e51241a7ce748a5d146db7874966f1a3b9) fix(api/libnvml): drain in-flight NVML queries before atexit shutdown - [`d611b1f`](https://github.com/XuehaiPan/nvitop/commit/d611b1fcdad52fc292be48471f7d7cfdd241c7e9) fix(api/libnvml): avoid re-initializing NVML during atexit shutdown - [`051d0ca`](https://github.com/XuehaiPan/nvitop/commit/051d0ca74b93e260150a3a46590df9ca9f813fc8) fix(api/libnvml): guard atexit nvmlShutdown() and register the hook once - [`6189dab`](https://github.com/XuehaiPan/nvitop/commit/6189dab15e3a82a90c5a5c5d235a507adbae91e1) fix(api/libnvml): cover lazy initialization in the shutdown drain - [`06b5acc`](https://github.com/XuehaiPan/nvitop/commit/06b5acc29f4d1dd459e7547624bb94893ab79d76) fix(api/libnvml): reset the shutdown-drain state after os.fork() - [`b0821c6`](https://github.com/XuehaiPan/nvitop/commit/b0821c619fad4e9af68494c292ade3f25010cb51) docs(api/libnvml): correct nvmlQuery Raises and the __shutting_down comment - [`5ade547`](https://github.com/XuehaiPan/nvitop/commit/5ade54791d3bd1feb165b92d43b7103d152779e1) docs(changelog): record the atexit NVML shutdown-drain fix ### 📊 Changes **3 files changed** (+162 additions, -44 deletions) <details> <summary>View changed files</summary> 📝 `.pre-commit-config.yaml` (+2 -2) 📝 `CHANGELOG.md` (+1 -1) 📝 `nvitop/api/libnvml.py` (+159 -41) </details> ### 📄 Description #### Issue Type - Bug fix #### Runtime Environment  - Operating system and version: Ubuntu 24.04.3 LTS (x86_64) — affected environment from #222; also verified on another Linux x86_64 host - Terminal emulator and version: N/A (the race is independent of the terminal emulator) - Python version: `3.14.6` (reporter) / `3.12.3` (verification host) - NVML version (driver version): `580.126.20` - `nvitop` version or commit: `fix/nvml-shutdown-drain@b0821c6` - `python-ml-py` version: `13.595.45` - Locale: N/A #### Description Pressing <kbd>q</kbd> to quit the TUI intermittently crashed with a `SIGSEGV` inside `libnvidia-ml`. The `atexit` `nvmlShutdown()` could free NVML internal state while a background thread was still inside a `libnvidia-ml` call — a use-after-free. This affects any concurrent NVML user of `nvitop.api`, including the TUI snapshot daemons and the reproducer in #222 (which spawns query threads with no TUI involved). The fix makes the implicit `atexit` shutdown safe rather than dropping it, by draining in-flight queries before the shutdown. All changes are confined to `nvitop/api/libnvml.py` and to nvitop's own additions on top of `pynvml` — the re-exported `pynvml` API is left untouched. **Drain in-flight queries at `atexit`:** - `nvmlQuery` bails on an unlocked `__shutting_down` fast path, then reserves an in-flight slot — incrementing `__active_queries` under a `threading.Condition` — **before** any NVML-touching work (including lazy `nvmlInit`), and re-checks `__shutting_down` after initialization. The NVML call itself runs outside the lock, so concurrent NVML access is never serialized; the slot is released and the condition notified in a `finally`. - `_atexit_shutdown` sets the one-way `__shutting_down` latch, then waits on the condition for `__active_queries` to reach zero before calling `nvmlShutdown()`. Once the latch is set no new query starts, so the counter only decreases and the wait wakes on the last decrement (event-driven, not polling). The wait is bounded (120 s) and skips `nvmlShutdown()` if queries do not drain in time — the OS reclaims NVML resources at process exit anyway. **Reset the drain state across `os.fork()`:** - `os.fork()` copies the module globals but leaves only the calling thread alive in the child, so the child can inherit a non-zero `__active_queries` (a query owned by a thread that no longer exists) or an already-held internal lock — making the inherited `atexit` hook block for its full timeout, or deadlock. `_reset_after_fork`, registered via `os.register_at_fork(after_in_child=...)`, rebuilds the lock/condition and zeroes the drain state in the child (guarded with `hasattr(os, 'register_at_fork')` since it is unavailable on Windows). The re-exported `pynvml` API is unchanged: `nvmlShutdown()` keeps its original signature, and an explicit `nvmlShutdown()` — or a raw `libnvml.nvml*` call — racing in-flight queries remains the caller's responsibility, since the drain only tracks calls made through `nvmlQuery`. The `chore(pre-commit)` commit bumps the `ruff` and `mypy` pre-commit hooks. #### Motivation and Context Fixes #222. Root cause: the `atexit` `nvmlShutdown()` races with in-flight NVML calls from background threads still running at interpreter exit; the shutdown frees NVML internal state under a concurrent call. Draining in-flight `nvmlQuery` calls before the shutdown closes the race for **any** concurrent caller — including the no-TUI reproducer in the issue, which a TUI-level fix could not help. #### Testing Verified on a Linux host with 4× NVIDIA H100 80 GB GPUs (Python 3.12.3) using the `nvml_race.py` reproducer from #222, run 25 times per build: | Build | `nvml_race.py` × 25 | | --- | --- | | `main` (no fix) | **15 / 25 `SIGSEGV`** (`rc=139`) | | `fix/nvml-shutdown-drain` | **0 / 25** | ```console $ git switch main $ python3 nvml_race.py [1] 920641 segmentation fault (core dumped) python3 nvml_race.py $ git switch fix/nvml-shutdown-drain $ python3 nvml_race.py $ # exits cleanly, no coredump ``` - On `main` the reproducer segfaults on ~60% of runs; on this branch it never crashes across 25 runs. - The crash is Linux-specific: on Windows, `main` does not coredump. - The `os.fork()` reset was verified GPU-free with a real `fork()`: a child that inherits a non-zero in-flight counter is reset to zero by the registered hook, so its `atexit` drain returns immediately instead of blocking for the timeout. - The coordination logic was also checked GPU-free during development: `nvmlQuery` refuses new NVML calls once `__shutting_down` is set, and `_atexit_shutdown` waits for an in-flight query before calling `nvmlShutdown()`. - `pre-commit` (ruff, ruff-format, codespell, mypy, pylint) passes. No automated tests are included (the repository has no test harness yet). --- <sub>🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.</sub>

gitea-mirror

2026-10-02 23:15:06 -06:00