nbqol.devices#

Functions

`count_cuda_devices`()	Count the number of available CUDA devices using nvidia-smi.
`cuda_device_report`([device_ids, simple, ...])
`cuda_visibles`(*args[, report])	Set visible CUDA devices and optionally report their status.
`format_key`(key[, pretty])	Format key names for consistent display.
`get_nvidia_smi`([device_ids, accessible, ...])	Get NVIDIA GPU information using nvidia-smi.
`memory_to_gb`(value_kb)	Convert memory value in KB to GB.
`normalize_device_ids`(device_ids)	Normalize device IDs to a consistent list format.
`safe_subprocess_call`(command[, shell])	Execute subprocess command with proper error handling.
`set_cuda_visibles`(*args[, report])	Set visible CUDA devices and optionally report their status.
`torch_device_report`(device[, readable, ...])

normalize_device_ids(device_ids: int | List[int] | None) → List[int] | None[source]#

Normalize device IDs to a consistent list format.

Parameters:: device_ids – Single device ID or list of device IDs
Returns:: List of device IDs or None if the input was None

memory_to_gb(value_kb: int | str) → float[source]#

Convert memory value in KB to GB.

Parameters:: value_kb – Memory value in KB, can be int or string
Returns:: Memory value in GB, rounded to 2 decimal places

format_key(key: str, pretty: bool = True) → str[source]#

Format key names for consistent display.

Parameters:

key – The key to format
pretty – If True, use human-readable format; otherwise, use code-friendly format

Returns:

Formatted key

safe_subprocess_call(command: List[str] | str, shell: bool = False) → str | None[source]#

Execute subprocess command with proper error handling.

Parameters:

command – Command to execute (list or string)
shell – Whether to use shell execution

Returns:

Command output as string or None if command failed

count_cuda_devices() → int[source]#

Count the number of available CUDA devices using nvidia-smi.

Returns:: Number of available CUDA devices, or 0 if none found
Return type:: int

Examples

>>> from nbqol import devices
>>> num_devices = devices.count_cuda_devices()
>>> print(f"Found {num_devices} CUDA devices")

get_nvidia_smi(device_ids: int | List[int] | None = None, accessible: bool = False, to_pandas: bool = False, **kwargs) → Dict[int, Dict[str, Any]] | DataFrame | Series[source]#

Get NVIDIA GPU information using nvidia-smi.

This function queries the NVIDIA System Management Interface to get detailed information about available GPUs, including memory usage, device names, and running processes.

Parameters:

device_ids – Specific device ID(s) to query. If None, query all devices.
accessible – If True, use more code-friendly key names for the output.
to_pandas – If True, return the results as a pandas DataFrame or Series.
**kwargs – check_visible (bool): If True, only return devices in CUDA_VISIBLE_DEVICES. device_prefix (str): Prefix for device names (default: ‘cuda:’).

Returns:

Dictionary, DataFrame, or Series containing GPU information.

Examples

>>> from nbqol import devices
>>> # Get info for all GPUs
>>> gpu_info = devices.get_nvidia_smi()
>>> # Get info for specific GPU
>>> gpu_info = devices.get_nvidia_smi(0)
>>> # Get as pandas DataFrame
>>> gpu_df = devices.get_nvidia_smi(to_pandas=True)

cuda_device_report(device_ids=None, simple=True, pretty=True, to_pandas=False, to_dict=False, **kwargs)[source]#

torch_device_report(device, readable=True, pretty=True, to_pandas=True, **extra_kwargs)[source]#

set_cuda_visibles(*args, report='torch', **kwargs) → DataFrame | None[source]#

Set visible CUDA devices and optionally report their status.

This function sets the CUDA_VISIBLE_DEVICES environment variable to limit which GPUs are accessible to CUDA applications. It can also generate a report of the selected devices.

Parameters:

*args – Device IDs to make visible. Can be individual numbers or lists.
report – Report type (‘torch’, ‘nvidia’, or None for no report).
**kwargs – verbose (bool): Print device selection message. to_pandas (bool): Return report as pandas DataFrame (torch report only).

Returns:

DataFrame of device info if to_pandas=True, None otherwise.

Return type:

Optional[pd.DataFrame]

Examples

>>> from nbqol import devices
>>> # Make only first GPU visible
>>> devices.set_cuda_visibles(0)
>>> # Make GPUs 0 and 2 visible
>>> devices.set_cuda_visibles(0, 2)
>>> # Disable all GPUs (use CPU only)
>>> devices.set_cuda_visibles()

cuda_visibles(*args, report='torch', **kwargs) → DataFrame | None#

Set visible CUDA devices and optionally report their status.

This function sets the CUDA_VISIBLE_DEVICES environment variable to limit which GPUs are accessible to CUDA applications. It can also generate a report of the selected devices.

Parameters:

*args – Device IDs to make visible. Can be individual numbers or lists.
report – Report type (‘torch’, ‘nvidia’, or None for no report).
**kwargs – verbose (bool): Print device selection message. to_pandas (bool): Return report as pandas DataFrame (torch report only).

Returns:

DataFrame of device info if to_pandas=True, None otherwise.

Return type:

Optional[pd.DataFrame]

Examples

>>> from nbqol import devices
>>> # Make only first GPU visible
>>> devices.set_cuda_visibles(0)
>>> # Make GPUs 0 and 2 visible
>>> devices.set_cuda_visibles(0, 2)
>>> # Disable all GPUs (use CPU only)
>>> devices.set_cuda_visibles()