choosegpu 0.0.3
Choose GPU (supports NVIDIA CUDA and Apple Silicon MPS)
License: MIT license
Requires Python: >=3.8

# Choose GPU

[![](https://img.shields.io/pypi/v/choosegpu.svg)](https://pypi.python.org/pypi/choosegpu)
[![CI](https://github.com/maximz/choosegpu/actions/workflows/ci.yaml/badge.svg?branch=master)](https://github.com/maximz/choosegpu/actions/workflows/ci.yaml)
[![](https://img.shields.io/badge/docs-here-blue.svg)](https://choosegpu.maximz.com)
[![](https://img.shields.io/github/stars/maximz/choosegpu?style=social)](https://github.com/maximz/choosegpu)

Automatically configure GPU usage for PyTorch and TensorFlow. Supports both NVIDIA CUDA GPUs and Apple Silicon (M series) MPS.

## Overview

`choosegpu` helps you manage GPU device selection and configuration before importing deep learning frameworks. It:

- **On NVIDIA systems**: Automatically detects available CUDA GPUs and sets `CUDA_VISIBLE_DEVICES` to use free GPUs
- **On Apple Silicon (Mac M series)**: Configures PyTorch to use the Metal Performance Shaders (MPS) backend
- Respects your preferences for which GPUs to use
- Can enable memory pooling on NVIDIA GPUs (requires RMM and CuPy)

**Important behavioral difference between platforms:**

- **NVIDIA CUDA**: When you call `configure_gpu(enable=False)`, it sets `CUDA_VISIBLE_DEVICES="-1"`, which makes `torch.cuda.is_available()` return `False`. The GPU is truly "hidden" from PyTorch/TensorFlow.

- **Apple Silicon MPS**: When you call `configure_gpu(enable=False)`, there is no effect because the hardware is always available. **The GPU is not actually disabled at the PyTorch level** - your code must check `choosegpu.get_gpu_config()` to determine whether to use it.

## Installation

```bash
pip install choosegpu
```

This installs `choosegpu` with all dependencies, including `nvitop` for NVIDIA GPU detection. On Mac Silicon, `nvitop` will be installed but not used (the library automatically detects the platform and uses MPS instead).

### Optional: Memory Pooling (NVIDIA only)

If you want to use `memory_pool=True` on NVIDIA GPUs, also install RMM and CuPy:

```bash
pip install rmm-cu11 cupy-cuda11x  # Adjust CUDA version as needed
```

## Usage

### Basic Usage

```python
import choosegpu

# Enable GPU (uses MPS on Mac Silicon, CUDA on NVIDIA)
choosegpu.configure_gpu(enable=True)

# Now import your deep learning framework
import torch

# The appropriate GPU backend will be configured
```

### Disable GPU (only effective on NVIDIA, not on Mac Silicon)

```python
import choosegpu

# Disable GPU - use CPU only
choosegpu.configure_gpu(enable=False)

import torch

# On Mac Silicon: check choosegpu.is_gpu_enabled() to see if GPU should be used
```

### Check GPU Configuration

```python
import choosegpu

choosegpu.configure_gpu(enable=True)

# Check if GPU is enabled (recommended)
gpu_enabled = choosegpu.is_gpu_enabled()
# Returns: True if GPU is enabled
# Returns: False if GPU is disabled (CPU mode)
# Returns: None if not yet configured

if gpu_enabled:
    print("GPU is enabled")
elif gpu_enabled is False:
    print("GPU is disabled (CPU mode)")
else:
    print("GPU not yet configured")

# Advanced: Get the raw GPU configuration
gpu_config = choosegpu.get_gpu_config()
# Returns: ["mps"] on Mac Silicon, or ["GPU-UUID"] on NVIDIA
# Returns: ["-1"] when disabled
# Returns: None if not configured
```

### Check Hardware Availability (PyTorch)

```python
import choosegpu

choosegpu.configure_gpu(enable=False)  # Disable GPU

# Check configuration state
print(f"GPU enabled: {choosegpu.is_gpu_enabled()}")  # False

# Check hardware availability (requires PyTorch)
print(f"GPU hardware available: {choosegpu.check_if_gpu_libraries_see_gpu()}")
# On Mac Silicon: True (hardware cannot be hidden)
# On NVIDIA: False (CUDA_VISIBLE_DEVICES="-1" hides hardware)
```

### Advanced: Prefer Specific GPUs (NVIDIA only)

```python
import choosegpu

# Prefer specific GPU IDs if they're available
choosegpu.configure_gpu(enable=True, gpu_device_ids=[2, 3])

# Or set global preference
choosegpu.preferred_gpu_ids = [2, 3]
choosegpu.configure_gpu(enable=True)
```

### Memory Pooling (NVIDIA only)

```python
import choosegpu

# Enable memory pooling with RMM (requires rmm and cupy)
choosegpu.configure_gpu(enable=True, memory_pool=True)
```

## Platform Support

| Platform | GPU Backend | Hardware Detection | GPU Disable Behavior |
|----------|-------------|-------------------|---------------------|
| Apple Silicon (M series) | MPS (Metal Performance Shaders) | Always available if hardware supports it | Sets configuration flag, but `torch.backends.mps.is_available()` remains `True` |
| NVIDIA | CUDA | Uses `nvitop` to detect available GPUs | Sets `CUDA_VISIBLE_DEVICES="-1"`, making `torch.cuda.is_available()` return `False` |

## API Reference

### `configure_gpu(enable=True, desired_number_of_gpus=1, memory_pool=False, gpu_device_ids=None, overwrite_existing_configuration=True)`

Configure GPU usage. Must be called before importing PyTorch/TensorFlow.

**Parameters:**
- `enable` (bool): Enable or disable GPU
- `desired_number_of_gpus` (int): Number of GPUs to use (NVIDIA only, ignored on Mac Silicon)
- `memory_pool` (bool): Enable memory pooling with RMM (NVIDIA only)
- `gpu_device_ids` (list): Preferred GPU IDs to use if available (NVIDIA only)
- `overwrite_existing_configuration` (bool): Whether to overwrite existing GPU configuration

**Returns:** List of GPU device identifiers (e.g., `["mps"]` or `["GPU-UUID-123"]`)

### `is_gpu_enabled()`

Check if GPU is enabled. This is the recommended way to check GPU configuration status.

**Returns:**
- `True`: GPU is enabled (configured to use GPU)
- `False`: GPU is disabled (configured to use CPU only)
- `None`: GPU configuration has not been set yet

### `check_if_gpu_libraries_see_gpu()`

Check if PyTorch can see GPU hardware (CUDA or MPS). This requires PyTorch to be installed.

**Important platform differences:**
- **Mac Silicon**: Returns `True` if PyTorch MPS is available, REGARDLESS of `configure_gpu()` settings. Mac hardware cannot be "hidden" like CUDA can.
- **NVIDIA**: Returns `True` if CUDA is available AND not disabled by `configure_gpu()`. When `configure_gpu(enable=False)` is called, `CUDA_VISIBLE_DEVICES="-1"` makes `torch.cuda.is_available()` return `False`.

**Returns:** `True` if GPU hardware is available to PyTorch, `False` otherwise (including when PyTorch is not installed)

### `get_gpu_config()`

Get current GPU configuration. For most use cases, prefer `is_gpu_enabled()` instead.

**Returns:**
- `["mps"]` if MPS is enabled on Mac Silicon
- `["-1"]` if GPU is disabled
- `["GPU-UUID", ...]` for NVIDIA GPUs
- `None` if not configured

### `are_gpu_settings_configured()`

Check if GPU settings have been configured by this library.

**Returns:** `True` if `configure_gpu()` has been called, `False` otherwise

## Development

Submit PRs against `develop` branch, then make a release pull request to `master`.

```bash
# Install requirements
pip install --upgrade pip wheel
pip install -r requirements_dev.txt

# Install local package
pip install -e .

# Install pre-commit
pre-commit install

# Run tests
make test

# Run lint
make lint

# bump version before submitting a PR against master (all master commits are deployed)
bump2version patch # possible: major / minor / patch

# also ensure CHANGELOG.md updated
```


# Changelog

## 0.0.1

* First release on PyPI.