[cuda, cli] feat: add NVML monitor and optimize cli

docs/concepts/architecture.md

@@ -11,7 +11,9 @@ schedulers that the GPU is still busy, without burning a full training workload.
    instantiates one single-GPU controller per selected device.
 3. **`CudaGPUController`** – Owns the background thread, VRAM allocation, and small
    matmul loops that tick every `interval` seconds.
-4. **Utilities** – `parse_size` turns strings like `1GiB` into bytes, while
+4. **GPU monitor (NVML)** – Wraps `nvidia-ml-py` (the `pynvml` module) so controllers
+   can read utilization data without shelling out to `nvidia-smi`.
+5. **Utilities** – `parse_size` turns strings like `1GiB` into bytes, while
    `setup_logger` wires both console and file logging with optional colors.
 
 ```text
@@ -26,7 +28,7 @@ CLI args ──▶ GlobalGPUController ──▶ [CudaGPUController rank=0]
 2. During `keep()` / `__enter__`, each Cuda worker:
    - Allocates a tensor sized by way of `vram_to_keep`.
    - Starts a daemon thread that performs `matmul_iterations` fused activations.
-   - Calls `_monitor_utilization` (by way of `nvidia-smi`) to detect real activity.
+   - Calls `_monitor_utilization` (by way of NVML) to detect real activity.
 3. If utilization exceeds `busy_threshold`, the worker just sleeps for one more
    `interval`. Otherwise it runs a new batch of ops.
 4. When you call `release()` (or exit the context), every worker sets a stop

docs/guides/cli.md

@@ -6,20 +6,26 @@ Jupyter environments.
 ## Command anatomy
 
 ```bash
-keep-gpu --interval 120 --gpu-ids 0,1 --vram 2GiB --threshold 25
+keep-gpu --interval 120 --gpu-ids 0,1 --vram 2GiB --busy-threshold 25
 ```
 
 | Flag | Meaning | Default |
 | --- | --- | --- |
 | `--interval` | Sleep between keep-alive cycles (seconds). Lower = tighter lock. | `300` |
 | `--gpu-ids` | Comma-separated visible IDs. Leave unset to keep every detected GPU busy. | all |
 | `--vram` | Amount of memory each controller allocates. Accepts `800MB`, `1GiB`, `1073741824`, etc. | `1GiB` |
-| `--threshold` | Skip work when utilization is already above this percentage. | `-1` (never skip) |
+| `--busy-threshold` | Skip work when utilization is already above this percentage (`--threshold` still works for legacy scripts). | `-1` (never skip) |
+
+!!! note "Still using `--threshold`?"
+    Values passed to `--threshold` are auto-detected: numbers override
+    `--busy-threshold`, while strings such as `1GiB` override `--vram`. Prefer the explicit
+    flags going forward, but old commands continue to run.
 
 !!! info "What happens under the hood?"
     Each GPU gets a `CudaGPUController` that allocates one tensor sized by
-    `--vram` and runs a lightweight matmul loop. Controllers watch `nvidia-smi`
-    to avoid hogging a device that is already busy (see `--threshold`).
+    `--vram` and runs a lightweight matmul loop. Controllers use NVML
+    (`nvidia-ml-py` / `pynvml` module) to read utilization so they back off when a device is already
+    busy (see `--busy-threshold`).
 
 ## Scenarios
 
@@ -45,7 +51,7 @@ keep-gpu --interval 180 --vram 512MB
 ### 3. Share the node without starving teammates
 
 ```bash
-keep-gpu --gpu-ids 0,1 --interval 90 --threshold 35
+keep-gpu --gpu-ids 0,1 --interval 90 --busy-threshold 35
 ```
 
 - Controllers pause their work whenever utilization exceeds 35%.
@@ -54,7 +60,7 @@ keep-gpu --gpu-ids 0,1 --interval 90 --threshold 35
 ### 4. Run from Jupyter or VS Code terminal
 
 ```bash
-!keep-gpu --interval 45 --vram 768MB --threshold 50
+!keep-gpu --interval 45 --vram 768MB --busy-threshold 50
 ```
 
 - Prefix with `!` (Jupyter) or use the integrated terminal.

docs/reference/cli.md

@@ -20,14 +20,19 @@ controller, clears the CUDA cache, and exits with status `0`.
 | `--interval INTEGER` | seconds | Sleep duration between utilization checks and keep-alive batches. Lower values keep the GPU hotter; higher values save power. Default: `300`. |
 | `--gpu-ids TEXT` | comma-separated ints | Subset of GPUs to guard (for example, `0,2`). If omitted, KeepGPU enumerates `torch.cuda.device_count()` and protects every visible device. |
 | `--vram TEXT` | human size or bytes | Amount of memory each GPU controller allocates. Accept formats like `512MB`, `1GiB`, or `1073741824`. Default: `1GiB`. |
-| `--threshold INTEGER` | percent | Upper bound on observed utilization. When `nvidia-smi` reports a higher number, the controller adds extra sleeps so it will not interfere with legitimate workloads. Default: `-1` (never throttle). |
+| `--busy-threshold INTEGER` / `--util-threshold INTEGER` | percent | Upper bound on observed utilization. When NVML reports a higher number, the controller adds extra sleeps so it will not interfere with legitimate workloads. Default: `-1` (never throttle). |
 | `--help` | flag | Show Typer-generated help and exit. |
 
 !!! tip "Choosing a VRAM value"
     Allocating 20‑30 % of a GPU’s memory is usually enough for schedulers that
     only watch the “memory in use” column. If your cluster monitors utilization,
     pair a higher `--vram` with a shorter `--interval`.
 
+!!! note "`--threshold` legacy flag"
+    Older scripts may still pass `--threshold`. Numeric values map to
+    `--busy-threshold`; strings such as `1GiB` override `--vram`. Prefer the
+    explicit flags going forward.
+
 ## Environment variables
 
 | Variable | Effect |

pyproject.toml

@@ -26,7 +26,7 @@ classifiers = [
 ]
 license = {text = "MIT license"}
 dependencies = [
-  "pynvml",
+  "nvidia-ml-py",
   "typer",
   "torch",
   "colorlog",

src/keep_gpu/cli.py

@@ -2,7 +2,7 @@
 
 import os
 import time
-from typing import Optional
+from typing import Optional, Tuple
 
 import torch
 import typer
@@ -11,11 +11,31 @@
 from keep_gpu.global_gpu_controller.global_gpu_controller import GlobalGPUController
 from keep_gpu.utilities.logger import setup_logger
 
-app = typer.Typer()
+app = typer.Typer(context_settings={"help_option_names": ["-h", "--help"]})
 console = Console()
 logger = setup_logger(__name__)
 
 
+def _apply_legacy_threshold(
+    vram_value: str, legacy_threshold: Optional[str], busy_threshold: int
+) -> Tuple[str, int, Optional[str]]:
+    """
+    Interpret the deprecated --threshold flag:
+    - If it parses as int, treat it as busy-threshold override.
+    - Otherwise treat it as a VRAM override.
+    Returns (vram, busy_threshold, mode) where mode is 'busy', 'vram', or None.
+    """
+    if legacy_threshold is None:
+        return vram_value, busy_threshold, None
+
+    try:
+        parsed_threshold = int(legacy_threshold)
+    except ValueError:
+        return legacy_threshold, busy_threshold, "vram"
+    else:
+        return vram_value, parsed_threshold, "busy"
+
+
 @app.command()
 def main(
     interval: int = typer.Option(
@@ -27,16 +47,40 @@ def main(
     ),
     vram: str = typer.Option(
         "1GiB",
-        help="Amount of VRAM to keep occupied (e.g., '500MB', '1GiB', or integer in bytes)",
+        "--vram",
+        help=(
+            "Amount of VRAM to keep occupied (e.g., '500MB', '1GiB', or integer in bytes). "
+            "Legacy flag '--threshold' remains supported as an alias."
+        ),
+    ),
+    legacy_threshold: Optional[str] = typer.Option(
+        None,
+        "--threshold",
+        hidden=True,
+        help="Deprecated alias. If numeric, overrides --busy-threshold; otherwise overrides --vram.",
     ),
-    threshold: int = typer.Option(
+    busy_threshold: int = typer.Option(
         -1,
+        "--busy-threshold",
+        "--util-threshold",
         help="Max GPU utilization threshold to trigger keeping GPU awake",
     ),
 ):
     """
     Keep specified GPUs awake by allocating VRAM and monitoring usage.
     """
+    vram, busy_threshold, legacy_mode = _apply_legacy_threshold(
+        vram, legacy_threshold, busy_threshold
+    )
+    if legacy_mode == "vram":
+        console.print(
+            "[yellow]`--threshold` for VRAM is deprecated; please use `--vram` going forward.[/yellow]"
+        )
+    elif legacy_mode == "busy":
+        console.print(
+            "[yellow]`--threshold` for utilization is deprecated; please use `--busy-threshold`.[/yellow]"
+        )
+
     # Process GPU IDs
     if gpu_ids:
         try:
@@ -58,14 +102,14 @@ def main(
     logger.info(f"GPU count: {gpu_count}")
     logger.info(f"VRAM to keep occupied: {vram}")
     logger.info(f"Check interval: {interval} seconds")
-    logger.info(f"Busy threshold: {threshold}%")
+    logger.info(f"Busy threshold: {busy_threshold}%")
 
     # Create and start Global GPU Controller
     global_controller = GlobalGPUController(
         gpu_ids=gpu_id_list,
         interval=interval,
         vram_to_keep=vram,
-        busy_threshold=threshold,
+        busy_threshold=busy_threshold,
     )
 
     with global_controller:

src/keep_gpu/single_gpu_controller/cuda_gpu_controller.py

@@ -1,12 +1,11 @@
-import re
-import subprocess
 import threading
 import time
 from typing import Optional
 
 import torch
 
 from keep_gpu.single_gpu_controller.base_gpu_controller import BaseGPUController
+from keep_gpu.utilities.gpu_monitor import get_gpu_utilization
 from keep_gpu.utilities.humanized_input import parse_size
 from keep_gpu.utilities.logger import setup_logger
 from keep_gpu.utilities.platform_manager import ComputingPlatform
@@ -192,25 +191,13 @@ def _run_mat_batch(self, matrix: torch.Tensor) -> None:
         )
 
     # ------------------------------------------------------------------
-    # Optional: simple nvidia-smi monitor (not used in thread version)
+    # Utilization monitor
     # ------------------------------------------------------------------
     @staticmethod
     def _monitor_utilization(rank: int) -> int:
         """
-        Return current GPU utilization (%) for `rank`
-        by parsing `nvidia-smi` output. Can be plugged into
-        `_keep_loop` if you want adaptive sleeping.
+        Return current GPU utilization (%) for `rank`.
+        Falls back to 0 when NVML is unavailable.
         """
-        proc = subprocess.Popen(
-            ["nvidia-smi", "-i", str(rank)],
-            stdout=subprocess.PIPE,
-            stderr=subprocess.PIPE,
-        )
-        stdout, _ = proc.communicate()
-        for line in stdout.decode().split("\n")[::-1]:
-            if "Default" in line:
-                try:
-                    return int(re.findall(r"\d+", line)[-1])
-                except (IndexError, ValueError):
-                    break
-        return 0
+        utilization = get_gpu_utilization(rank)
+        return utilization if utilization is not None else 0

src/keep_gpu/utilities/gpu_monitor.py

@@ -0,0 +1,83 @@
+"""Utilities for querying GPU utilization in a Pythonic way."""
+
+from __future__ import annotations
+
+import atexit
+import threading
+from typing import Optional
+
+from keep_gpu.utilities.logger import setup_logger
+
+logger = setup_logger(__name__)
+
+try:  # pragma: no cover - import guard
+    # Provided by the maintained `nvidia-ml-py` package.
+    import pynvml  # type: ignore
+except Exception:  # pragma: no cover - env without NVML
+    pynvml = None
+
+
+class NVMLMonitor:
+    """Lightweight wrapper around NVML to read GPU utilization."""
+
+    def __init__(self, nvml_module) -> None:
+        self._nvml = nvml_module
+        self._lock = threading.Lock()
+        self._initialized = False
+        self._shutdown_registered = False
+
+    def _ensure_initialized(self) -> bool:
+        if self._nvml is None:
+            return False
+        if self._initialized:
+            return True
+
+        with self._lock:
+            if self._initialized:
+                return True
+            try:
+                self._nvml.nvmlInit()
+            except Exception as exc:  # pragma: no cover - passthrough
+                logger.debug("NVML init failed: %s", exc)
+                return False
+
+            if not self._shutdown_registered:
+                atexit.register(self._safe_shutdown)
+                self._shutdown_registered = True
+
+            self._initialized = True
+            return True
+
+    def _safe_shutdown(self) -> None:
+        if not self._nvml or not self._initialized:
+            return
+        try:
+            self._nvml.nvmlShutdown()
+        except Exception as exc:  # pragma: no cover - best effort
+            logger.debug("NVML shutdown failed: %s", exc)
+        finally:
+            self._initialized = False
+
+    def get_gpu_utilization(self, index: int) -> Optional[int]:
+        """Return utilization percentage for `index`, or None when unavailable."""
+        if not self._ensure_initialized():
+            return None
+
+        try:
+            handle = self._nvml.nvmlDeviceGetHandleByIndex(index)
+            rates = self._nvml.nvmlDeviceGetUtilizationRates(handle)
+            return int(rates.gpu)
+        except self._nvml.NVMLError as exc:
+            logger.debug("NVML query failed for GPU %s: %s", index, exc)
+            return None
+
+
+_nvml_monitor = NVMLMonitor(pynvml)
+
+
+def get_gpu_utilization(index: int) -> Optional[int]:
+    """Return utilization percentage for `index`, or None when unavailable."""
+    return _nvml_monitor.get_gpu_utilization(index)
+
+
+__all__ = ["get_gpu_utilization", "NVMLMonitor"]

src/keep_gpu/utilities/platform_manager.py

@@ -11,7 +11,8 @@ class ComputingPlatform(Enum):
 
 
 def _check_cuda():
-    # NOTE: This function checks for CUDA availability by trying to import pynvml.
+    # NOTE: This function checks for CUDA availability by trying to import pynvml
+    # from nvidia-ml-py (the maintained NVML bindings).
     # See https://github.com/vllm-project/vllm/blob/536fd330036b0406786c847f68e4f67cba06f421/vllm/platforms/__init__.py#L58
     # for related discussion.
     try:

tests/test_cli_thresholds.py

@@ -0,0 +1,22 @@
+from keep_gpu import cli
+
+
+def test_apply_legacy_threshold_none():
+    vram, threshold, mode = cli._apply_legacy_threshold("1GiB", None, -1)
+    assert vram == "1GiB"
+    assert threshold == -1
+    assert mode is None
+
+
+def test_apply_legacy_threshold_numeric():
+    vram, threshold, mode = cli._apply_legacy_threshold("1GiB", "25", -1)
+    assert vram == "1GiB"
+    assert threshold == 25
+    assert mode == "busy"
+
+
+def test_apply_legacy_threshold_memory_string():
+    vram, threshold, mode = cli._apply_legacy_threshold("1GiB", "2GiB", -1)
+    assert vram == "2GiB"
+    assert threshold == -1
+    assert mode == "vram"