From a2f12d4f8cf42b6b86f298f924c65d84c6db41a2 Mon Sep 17 00:00:00 2001 From: Wang Siyuan Date: Wed, 12 Nov 2025 14:33:58 +0800 Subject: [PATCH 1/8] Use NVML for GPU utilization --- .../cuda_gpu_controller.py | 23 ++---- src/keep_gpu/utilities/gpu_monitor.py | 82 +++++++++++++++++++ 2 files changed, 87 insertions(+), 18 deletions(-) create mode 100644 src/keep_gpu/utilities/gpu_monitor.py 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..a8325515 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 @@ -197,20 +196,8 @@ def _run_mat_batch(self, matrix: torch.Tensor) -> None: @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..ded34033 --- /dev/null +++ b/src/keep_gpu/utilities/gpu_monitor.py @@ -0,0 +1,82 @@ +"""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 + 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 Exception 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"] From ede2b33d18ca82719b9bbe13df9a25401a76455f Mon Sep 17 00:00:00 2001 From: Wang Siyuan Date: Wed, 12 Nov 2025 14:34:10 +0800 Subject: [PATCH 2/8] Add tests for NVML monitor --- tests/utilities/test_gpu_monitor.py | 51 +++++++++++++++++++++++++++++ 1 file changed, 51 insertions(+) create mode 100644 tests/utilities/test_gpu_monitor.py diff --git a/tests/utilities/test_gpu_monitor.py b/tests/utilities/test_gpu_monitor.py new file mode 100644 index 00000000..51ca6f9f --- /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 pynvml 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 From 0a0cc0c2270cb823129bd938e070d542e5b77df1 Mon Sep 17 00:00:00 2001 From: Wang Siyuan Date: Wed, 12 Nov 2025 14:34:22 +0800 Subject: [PATCH 3/8] Document NVML-based monitoring --- docs/concepts/architecture.md | 6 ++++-- docs/guides/cli.md | 5 +++-- 2 files changed, 7 insertions(+), 4 deletions(-) diff --git a/docs/concepts/architecture.md b/docs/concepts/architecture.md index d81e3cd6..7b1b417a 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 `pynvml` 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` (via 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..977f7b08 100644 --- a/docs/guides/cli.md +++ b/docs/guides/cli.md @@ -18,8 +18,9 @@ keep-gpu --interval 120 --gpu-ids 0,1 --vram 2GiB --threshold 25 !!! 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 + (`pynvml`) to read utilization so they back off when a device is already + busy (see `--threshold`). ## Scenarios From 9d0a3bb4fc5b91ca1f31661d8d86b6a75f4b8dc1 Mon Sep 17 00:00:00 2001 From: Wang Siyuan Date: Wed, 12 Nov 2025 14:35:56 +0800 Subject: [PATCH 4/8] Tweak doc wording per lint --- docs/concepts/architecture.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/concepts/architecture.md b/docs/concepts/architecture.md index 7b1b417a..685a7b5a 100644 --- a/docs/concepts/architecture.md +++ b/docs/concepts/architecture.md @@ -28,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` (via NVML) 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 From 25dd07eb0e815beb692bb53a9b1380236275adc2 Mon Sep 17 00:00:00 2001 From: Wang Siyuan Date: Wed, 12 Nov 2025 14:36:55 +0800 Subject: [PATCH 5/8] Clarify utilization monitor section --- src/keep_gpu/single_gpu_controller/cuda_gpu_controller.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 a8325515..7ff794a0 100644 --- a/src/keep_gpu/single_gpu_controller/cuda_gpu_controller.py +++ b/src/keep_gpu/single_gpu_controller/cuda_gpu_controller.py @@ -191,7 +191,7 @@ 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: From 6ed5b3814193956661b132a018058c86111b1db9 Mon Sep 17 00:00:00 2001 From: Wang Siyuan Date: Wed, 12 Nov 2025 14:41:54 +0800 Subject: [PATCH 6/8] Depend on nvidia-ml-py for NVML --- docs/concepts/architecture.md | 4 ++-- docs/guides/cli.md | 2 +- pyproject.toml | 2 +- src/keep_gpu/utilities/gpu_monitor.py | 1 + src/keep_gpu/utilities/platform_manager.py | 3 ++- tests/utilities/test_gpu_monitor.py | 2 +- 6 files changed, 8 insertions(+), 6 deletions(-) diff --git a/docs/concepts/architecture.md b/docs/concepts/architecture.md index 685a7b5a..562f77a5 100644 --- a/docs/concepts/architecture.md +++ b/docs/concepts/architecture.md @@ -11,8 +11,8 @@ 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. **GPU monitor (NVML)** – Wraps `pynvml` so controllers can read utilization - data without shelling out to `nvidia-smi`. +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. diff --git a/docs/guides/cli.md b/docs/guides/cli.md index 977f7b08..c2392337 100644 --- a/docs/guides/cli.md +++ b/docs/guides/cli.md @@ -19,7 +19,7 @@ keep-gpu --interval 120 --gpu-ids 0,1 --vram 2GiB --threshold 25 !!! 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 use NVML - (`pynvml`) to read utilization so they back off when a device is already + (`nvidia-ml-py` / `pynvml` module) to read utilization so they back off when a device is already busy (see `--threshold`). ## Scenarios 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/utilities/gpu_monitor.py b/src/keep_gpu/utilities/gpu_monitor.py index ded34033..a5573f0f 100644 --- a/src/keep_gpu/utilities/gpu_monitor.py +++ b/src/keep_gpu/utilities/gpu_monitor.py @@ -11,6 +11,7 @@ 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 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/utilities/test_gpu_monitor.py b/tests/utilities/test_gpu_monitor.py index 51ca6f9f..0d6a777b 100644 --- a/tests/utilities/test_gpu_monitor.py +++ b/tests/utilities/test_gpu_monitor.py @@ -4,7 +4,7 @@ class DummyNVML: - """Minimal stand-in for pynvml module used in tests.""" + """Minimal stand-in for the NVML module used in tests.""" class NVMLError(Exception): pass From d412b29363116fbf14fd835aaf4fb95fb62c4220 Mon Sep 17 00:00:00 2001 From: Wang Siyuan Date: Wed, 12 Nov 2025 14:55:49 +0800 Subject: [PATCH 7/8] Improve CLI help and threshold handling --- docs/guides/cli.md | 15 ++++++---- docs/reference/cli.md | 7 ++++- src/keep_gpu/cli.py | 56 ++++++++++++++++++++++++++++++++---- tests/test_cli_thresholds.py | 22 ++++++++++++++ 4 files changed, 88 insertions(+), 12 deletions(-) create mode 100644 tests/test_cli_thresholds.py diff --git a/docs/guides/cli.md b/docs/guides/cli.md index c2392337..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,13 +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 use NVML (`nvidia-ml-py` / `pynvml` module) to read utilization so they back off when a device is already - busy (see `--threshold`). + busy (see `--busy-threshold`). ## Scenarios @@ -46,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%. @@ -55,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/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/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" From 2c3ab15e80390d30b65845b5574458aa90dac2eb Mon Sep 17 00:00:00 2001 From: Wang Siyuan Date: Wed, 12 Nov 2025 15:33:09 +0800 Subject: [PATCH 8/8] Update src/keep_gpu/utilities/gpu_monitor.py Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> --- src/keep_gpu/utilities/gpu_monitor.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/keep_gpu/utilities/gpu_monitor.py b/src/keep_gpu/utilities/gpu_monitor.py index a5573f0f..1c944728 100644 --- a/src/keep_gpu/utilities/gpu_monitor.py +++ b/src/keep_gpu/utilities/gpu_monitor.py @@ -67,7 +67,7 @@ def get_gpu_utilization(self, index: int) -> Optional[int]: handle = self._nvml.nvmlDeviceGetHandleByIndex(index) rates = self._nvml.nvmlDeviceGetUtilizationRates(handle) return int(rates.gpu) - except Exception as exc: + except self._nvml.NVMLError as exc: logger.debug("NVML query failed for GPU %s: %s", index, exc) return None