diff --git a/docs/concepts/architecture.md b/docs/concepts/architecture.md index d81e3cd6..562f77a5 100644 --- a/docs/concepts/architecture.md +++ b/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 diff --git a/docs/guides/cli.md b/docs/guides/cli.md index 0b3f7bb7..966a5faf 100644 --- a/docs/guides/cli.md +++ b/docs/guides/cli.md @@ -6,7 +6,7 @@ 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 | @@ -14,12 +14,18 @@ keep-gpu --interval 120 --gpu-ids 0,1 --vram 2GiB --threshold 25 | `--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. diff --git a/docs/reference/cli.md b/docs/reference/cli.md index 4df45e2e..a7464d11 100644 --- a/docs/reference/cli.md +++ b/docs/reference/cli.md @@ -20,7 +20,7 @@ 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" @@ -28,6 +28,11 @@ controller, clears the CUDA cache, and exits with status `0`. 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 | diff --git a/pyproject.toml b/pyproject.toml index f62fcbf3..493c77ad 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -26,7 +26,7 @@ classifiers = [ ] license = {text = "MIT license"} dependencies = [ - "pynvml", + "nvidia-ml-py", "typer", "torch", "colorlog", diff --git a/src/keep_gpu/cli.py b/src/keep_gpu/cli.py index be56c324..39666621 100644 --- a/src/keep_gpu/cli.py +++ b/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: diff --git a/src/keep_gpu/single_gpu_controller/cuda_gpu_controller.py b/src/keep_gpu/single_gpu_controller/cuda_gpu_controller.py index f927a0b9..7ff794a0 100644 --- a/src/keep_gpu/single_gpu_controller/cuda_gpu_controller.py +++ b/src/keep_gpu/single_gpu_controller/cuda_gpu_controller.py @@ -1,5 +1,3 @@ -import re -import subprocess import threading import time from typing import Optional @@ -7,6 +5,7 @@ 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 diff --git a/src/keep_gpu/utilities/gpu_monitor.py b/src/keep_gpu/utilities/gpu_monitor.py new file mode 100644 index 00000000..1c944728 --- /dev/null +++ b/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"] diff --git a/src/keep_gpu/utilities/platform_manager.py b/src/keep_gpu/utilities/platform_manager.py index 0881f7b6..995a60de 100644 --- a/src/keep_gpu/utilities/platform_manager.py +++ b/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: diff --git a/tests/test_cli_thresholds.py b/tests/test_cli_thresholds.py new file mode 100644 index 00000000..07a72056 --- /dev/null +++ b/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" diff --git a/tests/utilities/test_gpu_monitor.py b/tests/utilities/test_gpu_monitor.py new file mode 100644 index 00000000..0d6a777b --- /dev/null +++ b/tests/utilities/test_gpu_monitor.py @@ -0,0 +1,51 @@ +import types + +from keep_gpu.utilities.gpu_monitor import NVMLMonitor + + +class DummyNVML: + """Minimal stand-in for the NVML module used in tests.""" + + class NVMLError(Exception): + pass + + def __init__(self, should_fail: bool = False, gpu_util: int = 50) -> None: + self.should_fail = should_fail + self.gpu_util = gpu_util + self.init_calls = 0 + + def nvmlInit(self): + self.init_calls += 1 + if self.should_fail: + raise self.NVMLError("init failure") + + def nvmlShutdown(self): + pass + + def nvmlDeviceGetHandleByIndex(self, index: int): + if self.should_fail: + raise self.NVMLError("handle failure") + return types.SimpleNamespace(index=index) + + def nvmlDeviceGetUtilizationRates(self, handle): + return types.SimpleNamespace(gpu=self.gpu_util) + + +def test_monitor_returns_none_when_nvml_missing(): + monitor = NVMLMonitor(None) + assert monitor.get_gpu_utilization(0) is None + + +def test_monitor_reads_gpu_utilization(): + dummy = DummyNVML(gpu_util=73) + monitor = NVMLMonitor(dummy) + assert monitor.get_gpu_utilization(1) == 73 + # second call reuses initialization + assert monitor.get_gpu_utilization(2) == 73 + assert dummy.init_calls == 1 + + +def test_monitor_handles_nvml_errors(): + dummy = DummyNVML(should_fail=True) + monitor = NVMLMonitor(dummy) + assert monitor.get_gpu_utilization(0) is None