tensorflowonspark.gpu_info module

MAX_RETRIES = 3

Maximum retries to allocate GPUs

get_gpus(num_gpu=1, worker_index=- 1, format='string')

Get list of free GPUs according to nvidia-smi.

This will retry for MAX_RETRIES times until the requested number of GPUs are available.

Args:

num_gpu

number of GPUs desired.

worker_index

index “hint” for allocation of available GPUs.

Returns:

Comma-delimited string of GPU ids, or raises an Exception if the requested number of GPUs could not be found.

is_gpu_available()

Determine if GPUs are available on the host