From 3b7d1057eb55d17878be567cf5b49358e2d393cb Mon Sep 17 00:00:00 2001
From: Wang Siyuan <wsy0227@sjtu.edu.cn>
Date: Tue, 25 Nov 2025 23:17:35 +0800
Subject: [PATCH 01/12] Refine platform detection and add tests

---
 src/keep_gpu/utilities/platform_manager.py | 79 +++++++++++++++-------
 tests/utilities/test_platform_manager.py   | 50 ++++++++++++++
 2 files changed, 106 insertions(+), 23 deletions(-)
 create mode 100644 tests/utilities/test_platform_manager.py

diff --git a/src/keep_gpu/utilities/platform_manager.py b/src/keep_gpu/utilities/platform_manager.py
index 995a60de..4b090fec 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,18 +16,21 @@ 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 (torch or NVML)."""
+    try:
+        if torch.cuda.is_available():
+            return True
+    except Exception as exc:  # pragma: no cover - torch edge cases
+        logger.debug("torch.cuda.is_available() failed: %s", exc)
+
     try:
         import pynvml
 
         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():
@@ -30,33 +38,58 @@ def _check_rocm():
         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..2efde0bd
--- /dev/null
+++ b/tests/utilities/test_platform_manager.py
@@ -0,0 +1,50 @@
+from keep_gpu.utilities import platform_manager as pm
+
+
+def _reset_cache(monkeypatch):
+    monkeypatch.setattr(pm, "_cached_platform", None)
+
+
+def test_env_override_cpu(monkeypatch):
+    _reset_cache(monkeypatch)
+    monkeypatch.setenv("KEEP_GPU_PLATFORM", "cpu")
+    # Ensure no real detection runs
+    monkeypatch.setattr(
+        pm,
+        "_PLATFORM_CHECKS",
+        [(pm.ComputingPlatform.CPU, lambda: True)],
+    )
+    assert pm.get_platform() == pm.ComputingPlatform.CPU
+
+
+def test_invalid_override_falls_back(monkeypatch):
+    _reset_cache(monkeypatch)
+    monkeypatch.setenv("KEEP_GPU_PLATFORM", "invalid")
+    monkeypatch.setattr(
+        pm,
+        "_PLATFORM_CHECKS",
+        [
+            (pm.ComputingPlatform.CUDA, lambda: False),
+            (pm.ComputingPlatform.CPU, lambda: True),
+        ],
+    )
+    assert pm.get_platform() == pm.ComputingPlatform.CPU
+
+
+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

From 94ae4ee556f9c4d4b7330911d6fd91395a48258a Mon Sep 17 00:00:00 2001
From: Wang Siyuan <wsy0227@sjtu.edu.cn>
Date: Tue, 25 Nov 2025 23:21:46 +0800
Subject: [PATCH 02/12] Make platform detection ROCm-aware and quieter

---
 src/keep_gpu/utilities/platform_manager.py | 21 ++++++++++--
 tests/utilities/test_platform_manager.py   | 39 ++++++++++++++++++----
 2 files changed, 50 insertions(+), 10 deletions(-)

diff --git a/src/keep_gpu/utilities/platform_manager.py b/src/keep_gpu/utilities/platform_manager.py
index 4b090fec..34d2bb8d 100644
--- a/src/keep_gpu/utilities/platform_manager.py
+++ b/src/keep_gpu/utilities/platform_manager.py
@@ -1,5 +1,6 @@
 import os
 from enum import Enum
+import warnings
 from typing import Callable, List, Tuple
 
 import torch
@@ -16,15 +17,23 @@ class ComputingPlatform(Enum):
 
 
 def _check_cuda():
-    """Return True if CUDA appears available (torch or NVML)."""
+    """
+    Return True if CUDA appears available.
+
+    - Prefer torch reporting CUDA with no ROCm build.
+    - Fall back to NVML availability.
+    """
     try:
-        if torch.cuda.is_available():
+        # 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
+        with warnings.catch_warnings():
+            warnings.filterwarnings("ignore", category=FutureWarning, module="pynvml")
+            import pynvml
 
         pynvml.nvmlInit()
         return True
@@ -34,6 +43,12 @@ def _check_cuda():
 
 
 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
 
diff --git a/tests/utilities/test_platform_manager.py b/tests/utilities/test_platform_manager.py
index 2efde0bd..08a4ed5f 100644
--- a/tests/utilities/test_platform_manager.py
+++ b/tests/utilities/test_platform_manager.py
@@ -3,32 +3,43 @@
 
 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")
-    # Ensure no real detection runs
-    monkeypatch.setattr(
-        pm,
-        "_PLATFORM_CHECKS",
-        [(pm.ComputingPlatform.CPU, lambda: True)],
-    )
     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, lambda: False),
+            (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):
@@ -48,3 +59,17 @@ def _fake_check():
     assert pm.get_platform() == pm.ComputingPlatform.CPU
     assert pm.get_platform() == pm.ComputingPlatform.CPU
     assert calls["count"] == 1
+
+
+def test_cuda_prefers_torch_over_hip(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

From 12af290514a6b424d04367d4064d7713df2e8dff Mon Sep 17 00:00:00 2001
From: Wang Siyuan <wsy0227@sjtu.edu.cn>
Date: Tue, 25 Nov 2025 23:30:59 +0800
Subject: [PATCH 03/12] Prioritize NVML for CUDA detection

---
 src/keep_gpu/utilities/platform_manager.py | 21 +++++++++++----------
 tests/utilities/test_platform_manager.py   | 17 +++++++++++++++++
 2 files changed, 28 insertions(+), 10 deletions(-)

diff --git a/src/keep_gpu/utilities/platform_manager.py b/src/keep_gpu/utilities/platform_manager.py
index 34d2bb8d..ff6f8387 100644
--- a/src/keep_gpu/utilities/platform_manager.py
+++ b/src/keep_gpu/utilities/platform_manager.py
@@ -20,16 +20,9 @@ def _check_cuda():
     """
     Return True if CUDA appears available.
 
-    - Prefer torch reporting CUDA with no ROCm build.
-    - Fall back to NVML availability.
+    - Prefer NVML availability (matches vLLM approach).
+    - Fall back to torch reporting CUDA with no ROCm build.
     """
-    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:
         with warnings.catch_warnings():
             warnings.filterwarnings("ignore", category=FutureWarning, module="pynvml")
@@ -39,7 +32,15 @@ def _check_cuda():
         return True
     except Exception as exc:
         logger.debug("NVML unavailable: %s", exc)
-        return False
+
+    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)
+
+    return False
 
 
 def _check_rocm():
diff --git a/tests/utilities/test_platform_manager.py b/tests/utilities/test_platform_manager.py
index 08a4ed5f..e5915e39 100644
--- a/tests/utilities/test_platform_manager.py
+++ b/tests/utilities/test_platform_manager.py
@@ -1,3 +1,5 @@
+import sys
+
 from keep_gpu.utilities import platform_manager as pm
 
 
@@ -73,3 +75,18 @@ def test_rocm_detects_hip(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_prefers_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

From 9c0fe387fbdfbb73c063df1683ecb7b232f16f11 Mon Sep 17 00:00:00 2001
From: Wang Siyuan <wsy0227@sjtu.edu.cn>
Date: Tue, 25 Nov 2025 23:33:34 +0800
Subject: [PATCH 04/12] Use nvidia-ml-py NVML for CUDA detection

---
 src/keep_gpu/utilities/platform_manager.py | 20 +++++++++-----------
 1 file changed, 9 insertions(+), 11 deletions(-)

diff --git a/src/keep_gpu/utilities/platform_manager.py b/src/keep_gpu/utilities/platform_manager.py
index ff6f8387..cfdc14b9 100644
--- a/src/keep_gpu/utilities/platform_manager.py
+++ b/src/keep_gpu/utilities/platform_manager.py
@@ -1,6 +1,5 @@
 import os
 from enum import Enum
-import warnings
 from typing import Callable, List, Tuple
 
 import torch
@@ -23,16 +22,6 @@ def _check_cuda():
     - Prefer NVML availability (matches vLLM approach).
     - Fall back to torch reporting CUDA with no ROCm build.
     """
-    try:
-        with warnings.catch_warnings():
-            warnings.filterwarnings("ignore", category=FutureWarning, module="pynvml")
-            import pynvml
-
-        pynvml.nvmlInit()
-        return True
-    except Exception as exc:
-        logger.debug("NVML unavailable: %s", exc)
-
     try:
         # ROCm builds set torch.version.hip; treat those as non-CUDA.
         if torch.cuda.is_available() and torch.version.hip is None:
@@ -40,6 +29,15 @@ def _check_cuda():
     except Exception as exc:  # pragma: no cover - torch edge cases
         logger.debug("torch.cuda.is_available() failed: %s", exc)
 
+    try:
+        import pynvml  # provided by nvidia-ml-py
+
+        pynvml.nvmlInit()
+        return True
+    except Exception as exc:
+        logger.debug("NVML unavailable: %s", exc)
+        return False
+
     return False
 
 

From 3e0b27b15bb13addb05b75bc2b2d90874d43a09b Mon Sep 17 00:00:00 2001
From: Wang Siyuan <wsy0227@sjtu.edu.cn>
Date: Tue, 25 Nov 2025 23:55:34 +0800
Subject: [PATCH 05/12] Add ROCm extra and clarify NVML source

---
 README.md      | 1 +
 pyproject.toml | 3 +++
 2 files changed, 4 insertions(+)

diff --git a/README.md b/README.md
index ea5b79f2..93b5ebad 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)
 
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"

From 86a572336c6015512111b3f6fe710e2093639ee5 Mon Sep 17 00:00:00 2001
From: Wang Siyuan <wsy0227@sjtu.edu.cn>
Date: Wed, 26 Nov 2025 00:00:05 +0800
Subject: [PATCH 06/12] Document platform support and ROCm extra

---
 README.md                     | 2 +-
 docs/concepts/architecture.md | 4 ++--
 docs/getting-started.md       | 8 +++++++-
 3 files changed, 10 insertions(+), 4 deletions(-)

diff --git a/README.md b/README.md
index 93b5ebad..ad8d4238 100644
--- a/README.md
+++ b/README.md
@@ -59,7 +59,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 via `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..692c78b9 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 via 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..333aae0c 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,12 @@ understand the minimum knobs you need to keep a GPU occupied.
     pip install keep-gpu
     ```
 
+=== "ROCm extra (utilities for ROCm telemetry)"
+    ```bash
+    pip install keep-gpu[rocm]
+    ```
+    Install your ROCm-compatible PyTorch build separately (per upstream instructions).
+
 === "Editable dev install"
     ```bash
     git clone https://github.com/Wangmerlyn/KeepGPU.git

From b76ec3ce61c63830c2e87f0f06edc8ea941b93e7 Mon Sep 17 00:00:00 2001
From: Wang Siyuan <wsy0227@sjtu.edu.cn>
Date: Wed, 26 Nov 2025 00:05:23 +0800
Subject: [PATCH 07/12] Add platform-specific install quick steps

---
 README.md               | 18 ++++++++++++++++++
 docs/getting-started.md | 17 +++++++++++++++--
 2 files changed, 33 insertions(+), 2 deletions(-)

diff --git a/README.md b/README.md
index ad8d4238..05f4fb76 100644
--- a/README.md
+++ b/README.md
@@ -29,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.
diff --git a/docs/getting-started.md b/docs/getting-started.md
index 333aae0c..aefd689c 100644
--- a/docs/getting-started.md
+++ b/docs/getting-started.md
@@ -20,11 +20,24 @@ understand the minimum knobs you need to keep a GPU occupied.
     pip install keep-gpu
     ```
 
-=== "ROCm extra (utilities for ROCm telemetry)"
+=== "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 your ROCm-compatible PyTorch build separately (per upstream instructions).
+    Install the ROCm-compatible PyTorch build that matches your runtime.
+
+=== "CPU-only"
+    ```bash
+    pip install torch
+    pip install keep-gpu
+    ```
 
 === "Editable dev install"
     ```bash

From d594119171ae166c4a25c7f359cd285fa3a6a509 Mon Sep 17 00:00:00 2001
From: Wang Siyuan <sywang0227@gmail.com>
Date: Wed, 26 Nov 2025 00:24:07 +0800
Subject: [PATCH 08/12] Update src/keep_gpu/utilities/platform_manager.py

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
---
 src/keep_gpu/utilities/platform_manager.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/keep_gpu/utilities/platform_manager.py b/src/keep_gpu/utilities/platform_manager.py
index cfdc14b9..680ad947 100644
--- a/src/keep_gpu/utilities/platform_manager.py
+++ b/src/keep_gpu/utilities/platform_manager.py
@@ -19,8 +19,8 @@ def _check_cuda():
     """
     Return True if CUDA appears available.
 
-    - Prefer NVML availability (matches vLLM approach).
-    - Fall back to torch reporting CUDA with no ROCm build.
+    - 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.

From c0c0bc8f51d5edf7ecbf87138a3dd2da634e100d Mon Sep 17 00:00:00 2001
From: Wang Siyuan <sywang0227@gmail.com>
Date: Wed, 26 Nov 2025 00:24:13 +0800
Subject: [PATCH 09/12] Update tests/utilities/test_platform_manager.py

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
---
 tests/utilities/test_platform_manager.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/utilities/test_platform_manager.py b/tests/utilities/test_platform_manager.py
index e5915e39..c0330b05 100644
--- a/tests/utilities/test_platform_manager.py
+++ b/tests/utilities/test_platform_manager.py
@@ -77,7 +77,7 @@ def test_rocm_detects_hip(monkeypatch):
     assert pm._check_rocm() is True
 
 
-def test_cuda_prefers_nvml(monkeypatch):
+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)

From 6bf17b441be898e0dbf30bb478f94da7e8332f23 Mon Sep 17 00:00:00 2001
From: Wang Siyuan <sywang0227@gmail.com>
Date: Wed, 26 Nov 2025 00:24:18 +0800
Subject: [PATCH 10/12] Update tests/utilities/test_platform_manager.py

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
---
 tests/utilities/test_platform_manager.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/utilities/test_platform_manager.py b/tests/utilities/test_platform_manager.py
index c0330b05..8184d860 100644
--- a/tests/utilities/test_platform_manager.py
+++ b/tests/utilities/test_platform_manager.py
@@ -63,7 +63,7 @@ def _fake_check():
     assert calls["count"] == 1
 
 
-def test_cuda_prefers_torch_over_hip(monkeypatch):
+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}))

From 55c7d2e891c81beee714c46be4e6905a862f349e Mon Sep 17 00:00:00 2001
From: Wang Siyuan <wsy0227@sjtu.edu.cn>
Date: Wed, 26 Nov 2025 00:27:19 +0800
Subject: [PATCH 11/12] Remove unreachable return in CUDA check

---
 src/keep_gpu/utilities/platform_manager.py | 2 --
 1 file changed, 2 deletions(-)

diff --git a/src/keep_gpu/utilities/platform_manager.py b/src/keep_gpu/utilities/platform_manager.py
index 680ad947..0063b793 100644
--- a/src/keep_gpu/utilities/platform_manager.py
+++ b/src/keep_gpu/utilities/platform_manager.py
@@ -38,8 +38,6 @@ def _check_cuda():
         logger.debug("NVML unavailable: %s", exc)
         return False
 
-    return False
-
 
 def _check_rocm():
     try:

From 677983ec9631d3f75ba556091f32b3bbee1da242 Mon Sep 17 00:00:00 2001
From: Wang Siyuan <wsy0227@sjtu.edu.cn>
Date: Wed, 26 Nov 2025 00:29:26 +0800
Subject: [PATCH 12/12] Fix formatting after lint

---
 README.md                     | 8 ++++----
 docs/concepts/architecture.md | 2 +-
 2 files changed, 5 insertions(+), 5 deletions(-)

diff --git a/README.md b/README.md
index 05f4fb76..99aaaa59 100644
--- a/README.md
+++ b/README.md
@@ -31,17 +31,17 @@ keep-gpu --gpu-ids 0 --vram 1GiB --busy-threshold 25 --interval 60
 
 ### Platform installs at a glance
 
-- **CUDA (example: cu121)**  
+- **CUDA (example: cu121)**
   ```bash
   pip install --index-url https://download.pytorch.org/whl/cu121 torch
   pip install keep-gpu
   ```
-- **ROCm (example: rocm6.1)**  
+- **ROCm (example: rocm6.1)**
   ```bash
   pip install --index-url https://download.pytorch.org/whl/rocm6.1 torch
   pip install keep-gpu[rocm]
   ```
-- **CPU-only**  
+- **CPU-only**
   ```bash
   pip install torch
   pip install keep-gpu
@@ -77,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; optional ROCm SMI support via `pip install keep-gpu[rocm]`.
+- 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 692c78b9..3576ea79 100644
--- a/docs/concepts/architecture.md
+++ b/docs/concepts/architecture.md
@@ -12,7 +12,7 @@ schedulers that the GPU is still busy, without burning a full training workload.
 3. **`CudaGPUController`** – Owns the background thread, VRAM allocation, and small
    matmul loops that tick every `interval` seconds.
 4. **GPU monitor (NVML/ROCm)** – Wraps `nvidia-ml-py` (the `pynvml` module) for CUDA
-   telemetry and optionally `rocm-smi` when installed via the `rocm` extra.
+   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.