github.com/netdata/go.d.plugin@v0.58.1/modules/nvidia_smi/nvidia_smi.go (about)

     1  // SPDX-License-Identifier: GPL-3.0-or-later
     2  
     3  package nvidia_smi
     4  
     5  import (
     6  	_ "embed"
     7  	"time"
     8  
     9  	"github.com/netdata/go.d.plugin/agent/module"
    10  	"github.com/netdata/go.d.plugin/pkg/web"
    11  )
    12  
    13  //go:embed "config_schema.json"
    14  var configSchema string
    15  
    16  func init() {
    17  	module.Register("nvidia_smi", module.Creator{
    18  		JobConfigSchema: configSchema,
    19  		Defaults: module.Defaults{
    20  			Disabled:    true,
    21  			UpdateEvery: 10,
    22  		},
    23  		Create: func() module.Module { return New() },
    24  	})
    25  }
    26  
    27  func New() *NvidiaSMI {
    28  	return &NvidiaSMI{
    29  		Config: Config{
    30  			Timeout:      web.Duration{Duration: time.Second * 10},
    31  			UseCSVFormat: true,
    32  		},
    33  		binName: "nvidia-smi",
    34  		charts:  &module.Charts{},
    35  		gpus:    make(map[string]bool),
    36  		migs:    make(map[string]bool),
    37  	}
    38  
    39  }
    40  
    41  type Config struct {
    42  	Timeout      web.Duration
    43  	BinaryPath   string `yaml:"binary_path"`
    44  	UseCSVFormat bool   `yaml:"use_csv_format"`
    45  }
    46  
    47  type (
    48  	NvidiaSMI struct {
    49  		module.Base
    50  		Config `yaml:",inline"`
    51  
    52  		charts *module.Charts
    53  
    54  		binName string
    55  		exec    nvidiaSMI
    56  
    57  		gpuQueryProperties []string
    58  
    59  		gpus map[string]bool
    60  		migs map[string]bool
    61  	}
    62  	nvidiaSMI interface {
    63  		queryGPUInfoXML() ([]byte, error)
    64  		queryGPUInfoCSV(properties []string) ([]byte, error)
    65  		queryHelpQueryGPU() ([]byte, error)
    66  	}
    67  )
    68  
    69  func (nv *NvidiaSMI) Init() bool {
    70  	if nv.exec == nil {
    71  		smi, err := nv.initNvidiaSMIExec()
    72  		if err != nil {
    73  			nv.Error(err)
    74  			return false
    75  		}
    76  		nv.exec = smi
    77  	}
    78  
    79  	return true
    80  }
    81  
    82  func (nv *NvidiaSMI) Check() bool {
    83  	return len(nv.Collect()) > 0
    84  }
    85  
    86  func (nv *NvidiaSMI) Charts() *module.Charts {
    87  	return nv.charts
    88  }
    89  
    90  func (nv *NvidiaSMI) Collect() map[string]int64 {
    91  	mx, err := nv.collect()
    92  	if err != nil {
    93  		nv.Error(err)
    94  	}
    95  
    96  	if len(mx) == 0 {
    97  		return nil
    98  	}
    99  	return mx
   100  }
   101  
   102  func (nv *NvidiaSMI) Cleanup() {}