github.com/sentienttechnologies/studio-go-runner@v0.0.0-20201118202441-6d21f2ced8ee/tools/gpu-inventory/main.go (about)

     1  // Copyright 2018-2020 (c) Cognizant Digital Business, Evolutionary AI. All rights reserved. Issued under the Apache 2.0 License.
     2  package main
     3  
     4  import (
     5  	"fmt"
     6  
     7  	"github.com/leaf-ai/studio-go-runner/internal/runner"
     8  	"github.com/leaf-ai/studio-go-runner/pkg/studio"
     9  
    10  	"github.com/davecgh/go-spew/spew"
    11  
    12  	"github.com/go-stack/stack"
    13  	"github.com/jjeffery/kv"
    14  )
    15  
    16  var (
    17  	logger = studio.NewLogger("runner")
    18  )
    19  
    20  func main() {
    21  	if _, free := runner.GPUSlots(); free == 0 {
    22  		if runner.HasCUDA() {
    23  
    24  			msg := fmt.Errorf("no available GPUs could be found using the nvidia management library")
    25  			if runner.CudaInitErr != nil {
    26  				msg = *runner.CudaInitErr
    27  			}
    28  			err := kv.Wrap(msg).With("stack", stack.Trace().TrimRuntime())
    29  			logger.Fatal(fmt.Sprint(err))
    30  		}
    31  		gpuDevices, err := runner.GetCUDAInfo()
    32  		if err != nil {
    33  			logger.Fatal(err.Error())
    34  		}
    35  		fmt.Println(spew.Sdump(gpuDevices))
    36  	} else {
    37  		logger.Fatal("No GPUs present")
    38  	}
    39  }