diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index ca0322be..64fe8964 100644 --- a/.pre-commit-config.yaml +++ b/.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: diff --git a/README.md b/README.md index 9052ee2e..20e1c3e0 100644 --- a/README.md +++ b/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 + + + + diff --git a/docs/api.md b/docs/api.md index 104f1fcf..028f0506 100644 --- a/docs/api.md +++ b/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 diff --git a/docs/usage.md b/docs/usage.md index d49d3baa..dfd0b932 100644 --- a/docs/usage.md +++ b/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 diff --git a/pyproject.toml b/pyproject.toml index fa4a17b3..a750462b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -26,6 +26,7 @@ classifiers = [ ] license = {text = "MIT license"} dependencies = [ + "pynvml", "typer", "torch", "colorlog", diff --git a/src/keep_gpu/single_gpu_controller/__init__.py b/src/keep_gpu/single_gpu_controller/__init__.py new file mode 100644 index 00000000..e69de29b 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 fda50142..2fc69b8c 100644 --- a/src/keep_gpu/single_gpu_controller/cuda_gpu_controller.py +++ b/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)