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.

def get_gpus(num_gpu=1, worker_index=-1):
  """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.
  """
  # get list of gpus (index, uuid)
  list_gpus = subprocess.check_output(["nvidia-smi", "--list-gpus"]).decode()
  logging.debug("all GPUs:\n{0}".format(list_gpus))

  # parse index and guid
  gpus = [x for x in list_gpus.split('\n') if len(x) > 0]

  def parse_gpu(gpu_str):
    cols = gpu_str.split(' ')
    return cols[5].split(')')[0], cols[1].split(':')[0]
  gpu_list = [parse_gpu(gpu) for gpu in gpus]

  free_gpus = []
  retries = 0
  while len(free_gpus) < num_gpu and retries < MAX_RETRIES:
    smi_output = subprocess.check_output(["nvidia-smi", "--format=csv,noheader,nounits", "--query-compute-apps=gpu_uuid"]).decode()
    logging.debug("busy GPUs:\n{0}".format(smi_output))
    busy_uuids = [x for x in smi_output.split('\n') if len(x) > 0]
    for uuid, index in gpu_list:
      if uuid not in busy_uuids:
        free_gpus.append(index)

    if len(free_gpus) < num_gpu:
      logging.warn("Unable to find available GPUs: requested={0}, available={1}".format(num_gpu, len(free_gpus)))
      retries += 1
      time.sleep(30 * retries)
      free_gpus = []

  logging.info("Available GPUs: {}".format(free_gpus))

  # if still can't find available GPUs, raise exception
  if len(free_gpus) < num_gpu:
    smi_output = subprocess.check_output(["nvidia-smi", "--format=csv", "--query-compute-apps=gpu_uuid,pid,process_name,used_gpu_memory"]).decode()
    logging.info(": {0}".format(smi_output))
    raise Exception("Unable to find {} free GPU(s)\n{}".format(num_gpu, smi_output))

  # Get logical placement
  num_available = len(free_gpus)
  if worker_index == -1:
    # use original random placement
    random.shuffle(free_gpus)
    proposed_gpus = free_gpus[:num_gpu]
  else:
    # ordered by worker index
    if worker_index * num_gpu + num_gpu > num_available:
      worker_index = worker_index * num_gpu % num_available
    proposed_gpus = free_gpus[worker_index * num_gpu:(worker_index * num_gpu + num_gpu)]
  logging.info("Proposed GPUs: {}".format(proposed_gpus))

  return ','.join(str(x) for x in proposed_gpus)