def get_num_gpus()

in src/vw-serving/src/vw_serving/sagemaker/gpu.py [0:0]

• 26 lines of code
• 8 McCabe index (conditional complexity)

def get_num_gpus(num_gpus=AUTODETECT_GPU_COUNT, **kwargs):
    """
    Returns the number of available GPUs based on configuration parameters and available hardware GPU devices.

    :param num_gpus: (int or "auto")
        If set to "auto", the function queries and returns the number of available GPUs.
        If set to an integer value, the function returns the value of min(num_gpus, auto_detected_gpu_count)
        Otherwise raises ValueError.
    :param kwargs: additional configuration parameters, not used
    :return: (int) number of GPUs
    """

    # Shortcut execution if what we want is 0 gpu, i.e. only cpu
    if num_gpus == 0:
        return 0

    try:
        num_available_gpus = _query_num_gpus()
    except TimeoutError:
        if num_gpus == AUTODETECT_GPU_COUNT:
            return 0
        else:
            return num_gpus

    if num_gpus == AUTODETECT_GPU_COUNT:
        return num_available_gpus
    else:
        try:
            num_requested_gpus = int(num_gpus)
        except ValueError:
            raise CustomerValueError(
                "Invalid value '{}' provided for hyperparameter '_num_gpus'. '_num_gpus' must be an integer or 'auto'. "
                "Please set the value of '_num_gpus' hyperparameter to 'auto' or an integer value and try again."
                .format(num_gpus))

        if num_requested_gpus > num_available_gpus:
            logging.warning("Request number of gpus: %d, Number of GPUs found: %d",
                            num_requested_gpus, num_available_gpus)
            return num_available_gpus
        else:
            return num_requested_gpus