github.com/filecoin-project/bacalhau@v0.3.23-0.20230228154132-45c989550ace/pkg/compute/capacity/system/provider.go (about)

     1  package system
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"os/exec"
     7  	"runtime"
     8  	"strings"
     9  
    10  	"github.com/filecoin-project/bacalhau/pkg/compute/capacity"
    11  	"github.com/filecoin-project/bacalhau/pkg/config"
    12  	"github.com/filecoin-project/bacalhau/pkg/model"
    13  	"github.com/pbnjay/memory"
    14  	"github.com/ricochet2200/go-disk-usage/du"
    15  )
    16  
    17  // NvidiaCLI is the path to the Nvidia helper binary
    18  const NvidiaCLI = "nvidia-container-cli"
    19  
    20  type PhysicalCapacityProvider struct {
    21  }
    22  
    23  func NewPhysicalCapacityProvider() *PhysicalCapacityProvider {
    24  	return &PhysicalCapacityProvider{}
    25  }
    26  
    27  func (p *PhysicalCapacityProvider) GetAvailableCapacity(ctx context.Context) (model.ResourceUsageData, error) {
    28  	diskSpace, err := getFreeDiskSpace(config.GetStoragePath())
    29  	if err != nil {
    30  		return model.ResourceUsageData{}, err
    31  	}
    32  	gpus, err := numSystemGPUs()
    33  	if err != nil {
    34  		return model.ResourceUsageData{}, err
    35  	}
    36  
    37  	// the actual resources we have
    38  	return model.ResourceUsageData{
    39  		CPU:    float64(runtime.NumCPU()) * 0.8,
    40  		Memory: memory.TotalMemory() * 80 / 100,
    41  		Disk:   diskSpace * 80 / 100,
    42  		GPU:    gpus,
    43  	}, nil
    44  }
    45  
    46  // get free disk space for storage path
    47  // returns bytes
    48  func getFreeDiskSpace(path string) (uint64, error) {
    49  	usage := du.NewDiskUsage(path)
    50  	if usage == nil {
    51  		return 0, fmt.Errorf("getFreeDiskSpace: unable to get disk space for path %s", path)
    52  	}
    53  	return usage.Free(), nil
    54  }
    55  
    56  // numSystemGPUs wraps nvidia-container-cli to get the number of GPUs
    57  func numSystemGPUs() (uint64, error) {
    58  	nvidiaPath, err := exec.LookPath(NvidiaCLI)
    59  	if err != nil {
    60  		// If the NVIDIA CLI is not installed, we can't know the number of GPUs, assume zero
    61  		if (err.(*exec.Error)).Unwrap() == exec.ErrNotFound {
    62  			return 0, nil
    63  		}
    64  		return 0, err
    65  	}
    66  	args := []string{
    67  		"info",
    68  		"--csv",
    69  	}
    70  	cmd := exec.Command(nvidiaPath, args...)
    71  	resp, err := cmd.Output()
    72  	if err != nil {
    73  		return 0, err
    74  	}
    75  
    76  	// Parse output of nvidia-container-cli command
    77  	lines := strings.Split(string(resp), "\n")
    78  	deviceInfoFlag := false
    79  	numDevices := uint64(0)
    80  	for _, line := range lines {
    81  		if strings.TrimSpace(line) == "" {
    82  			continue
    83  		}
    84  		if strings.HasPrefix(line, "Device Index") {
    85  			deviceInfoFlag = true
    86  			continue
    87  		}
    88  		if deviceInfoFlag {
    89  			numDevices += 1
    90  		}
    91  	}
    92  
    93  	fmt.Println(numDevices)
    94  	return numDevices, nil
    95  }
    96  
    97  // compile-time check that the provider implements the interface
    98  var _ capacity.Provider = (*PhysicalCapacityProvider)(nil)