diff --git a/README.md b/README.md index ea5b79f2..99aaaa59 100644 --- a/README.md +++ b/README.md @@ -18,6 +18,7 @@ On many clusters, idle GPUs are reaped or silently shared after a short grace pe - **Portable** – Typer/Rich CLI for humans; Python API for orchestrators and notebooks. - **Observable** – Structured logging and optional file logs for auditing what kept the GPU alive. - **Power-aware** – Uses intervalled elementwise ops instead of heavy matmul floods to present “busy” utilization while keeping power and thermals lower (see `CudaGPUController._run_mat_batch` for the loop). +- **NVML-backed** – GPU telemetry comes from `nvidia-ml-py` (the `pynvml` module), with optional `rocm-smi` support when you install the `rocm` extra. ## Quick start (CLI) @@ -28,6 +29,24 @@ pip install keep-gpu keep-gpu --gpu-ids 0 --vram 1GiB --busy-threshold 25 --interval 60 ``` +### Platform installs at a glance + +- **CUDA (example: cu121)** + ```bash + pip install --index-url https://download.pytorch.org/whl/cu121 torch + pip install keep-gpu + ``` +- **ROCm (example: rocm6.1)** + ```bash + pip install --index-url https://download.pytorch.org/whl/rocm6.1 torch + pip install keep-gpu[rocm] + ``` +- **CPU-only** + ```bash + pip install torch + pip install keep-gpu + ``` + Flags that matter: - `--vram` (`1GiB`, `750MB`, or bytes): how much memory to pin. @@ -58,7 +77,7 @@ with GlobalGPUController(gpu_ids=[0, 1], vram_to_keep="750MB", interval=90, busy ## What you get - Battle-tested keep-alive loop built on PyTorch. -- NVML-based utilization monitoring (by way of `nvidia-ml-py`) to avoid hogging busy GPUs. +- NVML-based utilization monitoring (by way of `nvidia-ml-py`) to avoid hogging busy GPUs; optional ROCm SMI support by way of `pip install keep-gpu[rocm]`. - CLI + API parity: same controllers power both code paths. - Continuous docs + CI: mkdocs + mkdocstrings build in CI to keep guidance up to date. diff --git a/docs/concepts/architecture.md b/docs/concepts/architecture.md index 562f77a5..3576ea79 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 `nvidia-ml-py` (the `pynvml` module) so controllers - can read utilization data without shelling out to `nvidia-smi`. +4. **GPU monitor (NVML/ROCm)** – Wraps `nvidia-ml-py` (the `pynvml` module) for CUDA + telemetry and optionally `rocm-smi` when installed by way of the `rocm` extra. 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/getting-started.md b/docs/getting-started.md index 9594c0d8..aefd689c 100644 --- a/docs/getting-started.md +++ b/docs/getting-started.md @@ -7,7 +7,7 @@ understand the minimum knobs you need to keep a GPU occupied. - NVIDIA drivers + CUDA runtime visible to PyTorch. - Python 3.9+ (matching the version in your environment/cluster image). -- Optional but recommended: `nvidia-smi` in `PATH` for utilization monitoring. +- Optional but recommended: `nvidia-smi` in `PATH` for utilization monitoring (CUDA) or `rocm-smi` if you install the `rocm` extra. !!! warning "ROCm & multi-tenant clusters" The current release focuses on CUDA devices. ROCm/AMD support is experimental; @@ -20,6 +20,25 @@ understand the minimum knobs you need to keep a GPU occupied. pip install keep-gpu ``` +=== "CUDA (example: cu121)" + ```bash + pip install --index-url https://download.pytorch.org/whl/cu121 torch + pip install keep-gpu + ``` + +=== "ROCm (example: rocm6.1)" + ```bash + pip install --index-url https://download.pytorch.org/whl/rocm6.1 torch + pip install keep-gpu[rocm] + ``` + Install the ROCm-compatible PyTorch build that matches your runtime. + +=== "CPU-only" + ```bash + pip install torch + pip install keep-gpu + ``` + === "Editable dev install" ```bash git clone https://github.com/Wangmerlyn/KeepGPU.git diff --git a/pyproject.toml b/pyproject.toml index 0b2ece44..de91dbfd 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -45,6 +45,9 @@ dev = [ "pre-commit", # Git hooks "bump-my-version", # Version management ] +rocm = [ + "rocm-smi", +] [project.urls] bugs = "https://github.com/Wangmerlyn/KeepGPU/issues" diff --git a/src/keep_gpu/utilities/platform_manager.py b/src/keep_gpu/utilities/platform_manager.py index 995a60de..0063b793 100644 --- a/src/keep_gpu/utilities/platform_manager.py +++ b/src/keep_gpu/utilities/platform_manager.py @@ -1,4 +1,9 @@ +import os from enum import Enum +from typing import Callable, List, Tuple + +import torch + from keep_gpu.utilities.logger import setup_logger logger = setup_logger(__name__) @@ -11,52 +16,92 @@ class ComputingPlatform(Enum): def _check_cuda(): - # 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. + """ + Return True if CUDA appears available. + + - Prefer torch reporting CUDA with no ROCm build. + - Fall back to NVML availability. + """ + try: + # ROCm builds set torch.version.hip; treat those as non-CUDA. + if torch.cuda.is_available() and torch.version.hip is None: + return True + except Exception as exc: # pragma: no cover - torch edge cases + logger.debug("torch.cuda.is_available() failed: %s", exc) + try: - import pynvml + import pynvml # provided by nvidia-ml-py pynvml.nvmlInit() - return ComputingPlatform.CUDA - except Exception as e: - logger.debug(f"CUDA not available: {e}") - return None + return True + except Exception as exc: + logger.debug("NVML unavailable: %s", exc) + return False def _check_rocm(): + try: + if torch.cuda.is_available() and torch.version.hip: + return True + except Exception as exc: # pragma: no cover - torch edge cases + logger.debug("torch ROCm detection failed: %s", exc) + try: import rocm_smi rocm_smi.rocm_smi_init() - return ComputingPlatform.ROCM - except Exception as e: - logger.debug(f"ROCM not available: {e}") - return None + return True + except Exception as exc: + logger.debug("ROCm SMI unavailable: %s", exc) + return False def _check_cpu(): - return ComputingPlatform.CPU + return True -platform_check_funcs = { - ComputingPlatform.CUDA: _check_cuda, - ComputingPlatform.ROCM: _check_rocm, - ComputingPlatform.CPU: _check_cpu, -} +_PLATFORM_CHECKS: List[Tuple[ComputingPlatform, Callable[[], bool]]] = [ + (ComputingPlatform.CUDA, _check_cuda), + (ComputingPlatform.ROCM, _check_rocm), + (ComputingPlatform.CPU, _check_cpu), +] + +_cached_platform: ComputingPlatform | None = None def get_platform(): """ Return the current computing platform. """ - for platform, check_func in platform_check_funcs.items(): - if check_func() is not None: - logger.info(f"Detected computing platform: {platform.value}") + global _cached_platform + + if _cached_platform is not None: + return _cached_platform + + override = os.getenv("KEEP_GPU_PLATFORM") + if override: + try: + platform = ComputingPlatform(override.lower()) + logger.info("Using KEEP_GPU_PLATFORM=%s override", platform.value) + _cached_platform = platform return platform - logger.info("No specific computing platform detected, defaulting to CPU.") - return ComputingPlatform.CPU # Default to CPU if no other platform is available + except ValueError: + logger.warning( + "Invalid KEEP_GPU_PLATFORM=%s; falling back to auto-detect", override + ) + + for platform, check_func in _PLATFORM_CHECKS: + try: + if check_func(): + logger.info("Detected computing platform: %s", platform.value) + _cached_platform = platform + return platform + except Exception as exc: # pragma: no cover - defensive + logger.debug("Platform check %s failed: %s", platform.value, exc) + + logger.info("No specific platform detected, defaulting to CPU.") + _cached_platform = ComputingPlatform.CPU + return _cached_platform # Default to CPU if no other platform is available if __name__ == "__main__": diff --git a/tests/utilities/test_platform_manager.py b/tests/utilities/test_platform_manager.py new file mode 100644 index 00000000..8184d860 --- /dev/null +++ b/tests/utilities/test_platform_manager.py @@ -0,0 +1,92 @@ +import sys + +from keep_gpu.utilities import platform_manager as pm + + +def _reset_cache(monkeypatch): + monkeypatch.setattr(pm, "_cached_platform", None) + # isolate checks for each test + monkeypatch.setattr( + pm, + "_PLATFORM_CHECKS", + [ + (pm.ComputingPlatform.CUDA, lambda: False), + (pm.ComputingPlatform.ROCM, lambda: False), + (pm.ComputingPlatform.CPU, lambda: True), + ], + ) + + +def test_env_override_cpu(monkeypatch): + _reset_cache(monkeypatch) + monkeypatch.setenv("KEEP_GPU_PLATFORM", "cpu") + assert pm.get_platform() == pm.ComputingPlatform.CPU + + +def test_invalid_override_falls_back(monkeypatch): + _reset_cache(monkeypatch) + monkeypatch.setenv("KEEP_GPU_PLATFORM", "invalid") + calls = {"count": 0} + + def fake_cuda(): + calls["count"] += 1 + return False + + monkeypatch.setattr( + pm, + "_PLATFORM_CHECKS", + [ + (pm.ComputingPlatform.CUDA, fake_cuda), + (pm.ComputingPlatform.CPU, lambda: True), + ], + ) + assert pm.get_platform() == pm.ComputingPlatform.CPU + assert calls["count"] == 1 + + +def test_cached_result_reused(monkeypatch): + _reset_cache(monkeypatch) + calls = {"count": 0} + + def _fake_check(): + calls["count"] += 1 + return True + + monkeypatch.setattr( + pm, + "_PLATFORM_CHECKS", + [(pm.ComputingPlatform.CPU, _fake_check)], + ) + + assert pm.get_platform() == pm.ComputingPlatform.CPU + assert pm.get_platform() == pm.ComputingPlatform.CPU + assert calls["count"] == 1 + + +def test_cuda_detected_via_torch_non_hip_build(monkeypatch): + _reset_cache(monkeypatch) + monkeypatch.setattr(pm.torch.cuda, "is_available", lambda: True) + monkeypatch.setattr(pm.torch, "version", type("v", (), {"hip": None})) + assert pm._check_cuda() is True + + +def test_rocm_detects_hip(monkeypatch): + _reset_cache(monkeypatch) + monkeypatch.setattr(pm.torch.cuda, "is_available", lambda: True) + monkeypatch.setattr(pm.torch, "version", type("v", (), {"hip": "6.0"})) + assert pm._check_rocm() is True + + +def test_cuda_detection_falls_back_to_nvml(monkeypatch): + _reset_cache(monkeypatch) + # Force torch to look like ROCm build to ensure NVML takes precedence + monkeypatch.setattr(pm.torch.cuda, "is_available", lambda: True) + monkeypatch.setattr(pm.torch, "version", type("v", (), {"hip": "6.0"})) + + class DummyNVML: + @staticmethod + def nvmlInit(): + return None + + monkeypatch.setitem(sys.modules, "pynvml", DummyNVML) + assert pm._check_cuda() is True