[platform] refactor: refine detection and platform-specific installs

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.
 

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.
 

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

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"

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__":

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