genome_entropy.encode3di.gpu_utils

GPU discovery and management utilities for multi-GPU encoding.

Functions

discover_available_gpus()	Discover available GPU devices from environment variables and CUDA.
select_device_for_gpu(gpu_id)	Get the device string for a specific GPU.
validate_gpu_availability(gpu_ids)	Validate that specified GPUs are actually available.

genome_entropy.encode3di.gpu_utils.discover_available_gpus()

Discover available GPU devices from environment variables and CUDA.

Checks multiple sources in order of priority: 1. SLURM_JOB_GPUS - SLURM allocated GPU IDs 2. SLURM_GPUS - Alternative SLURM GPU specification 3. CUDA_VISIBLE_DEVICES - User-specified visible devices 4. torch.cuda - Query CUDA directly if available

Returns:

List of GPU device IDs available for use. Empty list if no GPUs found.

Return type:

List[int]

Examples

>>> # With SLURM_JOB_GPUS="0,1,2"
>>> discover_available_gpus()
[0, 1, 2]

>>> # With CUDA_VISIBLE_DEVICES="2,3"
>>> discover_available_gpus()
[0, 1]  # Remapped to local indices

genome_entropy.encode3di.gpu_utils.select_device_for_gpu(gpu_id)

Get the device string for a specific GPU.

Parameters:

gpu_id (int) – GPU device ID

Returns:

0”, “cuda:1”)

Return type:

Device string (e.g., “cuda

genome_entropy.encode3di.gpu_utils.validate_gpu_availability(gpu_ids)

Validate that specified GPUs are actually available.

Parameters:

gpu_ids (List[int]) – List of GPU IDs to validate

Returns:

List of valid GPU IDs (subset of input)

Return type:

List[int]