[Docs, Dependencies] Docs update and add pynvml dependencies

.pre-commit-config.yaml

@@ -1,4 +1,8 @@
 repos:
+  - repo: https://github.com/google/pre-commit-tool-hooks
+    rev: v1.2.5
+    hooks:
+      - id: check-google-doc-style
   - repo: https://github.com/charliermarsh/ruff-pre-commit
     rev: v0.4.4
     hooks:

README.md

@@ -10,19 +10,19 @@
 
 ---
 
-## Features
+Contributions Welcome!
 
-- Simple command-line interface
-- Uses PyTorch and `nvidia-smi` to monitor and load GPUs
-- Easy to extend for your own keep-alive logic
+If you have ideas for new features or improvements, feel free to open an issue or submit a pull request.
+
+This project does not yet fully support ROCm GPUs, so any contributions, suggestions, or testing help in that area are especially welcome!
 
 ---
 
-## TODO ✅
+## Features
 
-- [ ] Add more CLI args (e.g. `--gpu-id`, `--gpu-ids`, `--gpu-keep-threshold`, `--gpu-keep-time`, `--gpu-keep-vram-usage`)
-- [ ] Add documentation
-- [ ] Add importable Python functions
+- Simple command-line interface
+- Uses PyTorch and `nvidia-smi` to monitor and load GPUs
+- Easy to extend for your own keep-alive logic
 
 ---
 
@@ -52,3 +52,9 @@ keep-gpu --gpu-ids 0,1,2
 ## Credits
 
 This package was created with [Cookiecutter](https://github.com/audreyr/cookiecutter) and the [audreyr/cookiecutter-pypackage](https://github.com/audreyr/cookiecutter-pypackage) project template.
+
+## Contributors
+
+<a href="https://github.com/Wangmerlyn/KeepGPU/graphs/contributors">
+  <img src="https://contrib.rocks/image?repository=Wangmerlyn/KeepGPU" />
+</a>

docs/api.md

@@ -4,6 +4,14 @@
     options:
       show_source: true
 
+::: keep_gpu.single_gpu_controller.cuda_gpu_controller
+    options:
+      show_source: true
+
+::: keep_gpu.global_gpu_controller.global_gpu_controller
+    options:
+      show_source: true
+
 ::: keep_gpu.utilities
     options:
       show_source: true

docs/usage.md

@@ -6,7 +6,7 @@ This page provides instructions on how to install, configure, and use KeepGPU.
 
 ## 🔧 Installation
 
-You can install KeepGPU via pip:
+You can install KeepGPU by way of pip:
 
 ```bash
 pip install keep-gpu

pyproject.toml

@@ -26,6 +26,7 @@ classifiers = [
 ]
 license = {text = "MIT license"}
 dependencies = [
+  "pynvml",
   "typer",
   "torch",
   "colorlog",

src/keep_gpu/single_gpu_controller/cuda_gpu_controller.py

@@ -29,23 +29,27 @@
 
 
 class CudaGPUController(BaseGPUController):
-    """
+    """CudaGPUController
     Keep a single CUDA GPU busy by repeatedly running lightweight
     matrix-multiplication workloads in a background thread.
 
-    Typical usage pattern
-    ---------------------
-    >>> ctrl = CudaGPUController(rank=0, interval=0.5)
-    >>> ctrl.start()          # occupy GPU while you do CPU-only work
-    >>> dataset.process()
-    >>> ctrl.release()        # give GPU memory back
-    >>> model.train_start()   # now run real GPU training
+    Typical usage:
+
+    ```python
+    ctrl = CudaGPUController(rank=0, interval=0.5)
+    ctrl.start()          # occupy GPU while you do CPU-only work
+    dataset.process()
+    ctrl.release()        # give GPU memory back
+    model.train_start()   # now run real GPU training
+    ```
 
-    You can also use the controller as a context manager:
+    Or as a context manager:
 
-    >>> with CudaGPUController(rank=0, interval=0.5):
-    ...     dataset.process()  # GPU occupied inside this block
-    >>> model.train_start()    # GPU free after exiting block
+    ```python
+    with CudaGPUController(rank=0, interval=0.5):
+        dataset.process()  # GPU occupied inside this block
+    model.train_start()    # GPU free after exiting block
+    ```
     """
 
     def __init__(
@@ -58,20 +62,19 @@ def __init__(
         busy_threshold: int = 10,
     ):
         """
-        Parameters
-        ----------
-        rank : int
-            Local CUDA device index to occupy.
-        interval : float, optional
-            Sleep time (seconds) between workload batches.
-        matmul_iterations : int, optional
-            Number of matmul ops per batch.
-        vram_to_keep : str | int, optional
-            Amount of VRAM to keep busy, e.g. "1000 MB", "20 GB" or 1000 * 1000.
-            This is the total size of the matrix allocated to keep the GPU busy.
-        busy_threshold : int, optional
-            If current utilisation (%) exceeds this value, the worker will
-            insert extra sleeps to avoid hogging the GPU.
+        Args:
+            rank (int): Local CUDA device index to occupy.
+            interval (float, optional): Sleep time (seconds) between workload
+                batches. Defaults to 0.5.
+            matmul_iterations (int, optional): Number of matmul ops per batch.
+            vram_to_keep (int or str, optional): Amount of VRAM to keep busy,
+                e.g. `"1000 MB"`, `"20 GB"`, or an integer like `1000 * 1000`.
+                This represents the total size of the matrix allocated to
+                occupy the GPU.
+            busy_threshold (int, optional): If current utilisation (%) exceeds
+                this threshold, the worker will insert extra sleeps to avoid
+                hogging the GPU.
+
         """
         if isinstance(vram_to_keep, str):
             vram_to_keep = self.parse_size(vram_to_keep)