Source code for schrodinger.infra.gpgpu_utils

from schrodinger import gpgpu
from schrodinger.infra import mm
from schrodinger.infra import mmjob


def get_local_gpgpus():
    """
    Detect local GPUs via schrodinger.gpgpu

    :return: list of tuple of local gpgpu device number and description
    :rtype: list(tuple(dev_num, desc))
    """
    local_gpulist = []
    for device in gpgpu.get_available_devices():
        num, desc = device.description().split(": ", 1)
        local_gpulist.append((num, desc))
    return local_gpulist


def get_remote_gpgpus(host):
    """
    Return a list of GPGPUs belonging to the specified host.

    :param host: name of host matching a name from the host file
    :type host: str

    :return: list of tuple of gpgpu device number and description
    :rtype: list(tuple(dev_num, desc))
    """
    gpgpu_list = []
    try:
        gpgpu_count = mmjob.mmjob_host_gpgpu_length(host)
    except mm.MmException:
        return []

    for card_idx in range(1, gpgpu_count + 1):
        gpgpu_list.append(mmjob.mmjob_host_gpgpu_item(host, card_idx))
    return gpgpu_list