github.com/juju/juju@v0.0.0-20240430160146-1752b71fcf00/state/devices.go (about)

     1  // Copyright 2018 Canonical Ltd.
     2  // Licensed under the AGPLv3, see LICENCE file for details.
     3  
     4  package state
     5  
     6  import (
     7  	"github.com/juju/charm/v12"
     8  	"github.com/juju/errors"
     9  	"github.com/juju/mgo/v3"
    10  	"github.com/juju/mgo/v3/txn"
    11  
    12  	"github.com/juju/juju/environs/config"
    13  )
    14  
    15  type DeviceType string
    16  
    17  // DeviceConstraints describes a set of device constraints.
    18  type DeviceConstraints struct {
    19  
    20  	// Type is the device type or device-class.
    21  	// currently supported types are
    22  	// - gpu
    23  	// - nvidia.com/gpu
    24  	// - amd.com/gpu
    25  	Type DeviceType `bson:"type"`
    26  
    27  	// Count is the number of devices that the user has asked for - count min and max are the
    28  	// number of devices the charm requires.
    29  	Count int64 `bson:"count"`
    30  
    31  	// Attributes is a collection of key value pairs device related (node affinity labels/tags etc.).
    32  	Attributes map[string]string `bson:"attributes"`
    33  }
    34  
    35  // NewDeviceBackend creates a backend for managing device.
    36  func NewDeviceBackend(st *State) (*deviceBackend, error) {
    37  	m, err := st.Model()
    38  	if err != nil {
    39  		return nil, errors.Trace(err)
    40  	}
    41  	return &deviceBackend{
    42  		mb:          st,
    43  		settings:    NewStateSettings(st),
    44  		modelType:   m.Type(),
    45  		config:      m.ModelConfig,
    46  		application: st.Application,
    47  		unit:        st.Unit,
    48  		machine:     st.Machine,
    49  	}, nil
    50  }
    51  
    52  type deviceBackend struct {
    53  	mb          modelBackend
    54  	config      func() (*config.Config, error)
    55  	application func(string) (*Application, error)
    56  	unit        func(string) (*Unit, error)
    57  	machine     func(string) (*Machine, error)
    58  
    59  	modelType ModelType
    60  	settings  *StateSettings
    61  }
    62  
    63  // deviceConstraintsDoc contains device constraints for an entity.
    64  type deviceConstraintsDoc struct {
    65  	DocID       string                       `bson:"_id"`
    66  	Constraints map[string]DeviceConstraints `bson:"constraints"`
    67  }
    68  
    69  func createDeviceConstraintsOp(id string, cons map[string]DeviceConstraints) txn.Op {
    70  	return txn.Op{
    71  		C:      deviceConstraintsC,
    72  		Id:     id,
    73  		Assert: txn.DocMissing,
    74  		Insert: &deviceConstraintsDoc{
    75  			Constraints: cons,
    76  		},
    77  	}
    78  }
    79  
    80  func removeDeviceConstraintsOp(id string) txn.Op {
    81  	return txn.Op{
    82  		C:      deviceConstraintsC,
    83  		Id:     id,
    84  		Remove: true,
    85  	}
    86  }
    87  
    88  // DeviceConstraints returns the device constraints for the specified application.
    89  func (db *deviceBackend) DeviceConstraints(id string) (map[string]DeviceConstraints, error) {
    90  	devices, err := readDeviceConstraints(db.mb, id)
    91  	if err == nil {
    92  		return devices, nil
    93  	} else if errors.IsNotFound(err) {
    94  		return map[string]DeviceConstraints{}, nil
    95  	}
    96  	return nil, err
    97  }
    98  
    99  func readDeviceConstraints(mb modelBackend, id string) (map[string]DeviceConstraints, error) {
   100  	coll, closer := mb.db().GetCollection(deviceConstraintsC)
   101  	defer closer()
   102  
   103  	var doc deviceConstraintsDoc
   104  	err := coll.FindId(id).One(&doc)
   105  	if err == mgo.ErrNotFound {
   106  		return nil, errors.NotFoundf("device constraints for %q", id)
   107  	}
   108  	if err != nil {
   109  		return nil, errors.Annotatef(err, "cannot get device constraints for %q", id)
   110  	}
   111  	return doc.Constraints, nil
   112  }
   113  
   114  func validateDeviceConstraints(db *deviceBackend, allCons map[string]DeviceConstraints, charmMeta *charm.Meta) error {
   115  	err := validateDeviceConstraintsAgainstCharm(db, allCons, charmMeta)
   116  	if err != nil {
   117  		return errors.Trace(err)
   118  	}
   119  	// Ensure all devices have constraints specified. Defaults should have
   120  	// been set by this point, if the user didn't specify constraints.
   121  	for name, charmDevice := range charmMeta.Devices {
   122  		if _, ok := allCons[name]; !ok && charmDevice.CountMin > 0 {
   123  			return errors.Errorf("no constraints specified for device %q", name)
   124  		}
   125  	}
   126  	return nil
   127  }
   128  
   129  func validateDeviceConstraintsAgainstCharm(
   130  	db *deviceBackend,
   131  	allCons map[string]DeviceConstraints,
   132  	charmMeta *charm.Meta,
   133  ) error {
   134  	for name, cons := range allCons {
   135  		charmDevice, ok := charmMeta.Devices[name]
   136  		if !ok {
   137  			return errors.Errorf("charm %q has no device called %q", charmMeta.Name, name)
   138  		}
   139  		if err := validateCharmDeviceCount(charmDevice, cons.Count); err != nil {
   140  			return errors.Annotatef(err, "charm %q device %q", charmMeta.Name, name)
   141  		}
   142  
   143  	}
   144  	return nil
   145  }
   146  
   147  func validateCharmDeviceCount(charmDevice charm.Device, count int64) error {
   148  	if charmDevice.CountMin > 0 && count < charmDevice.CountMin {
   149  		return errors.Errorf("minimum device size is %d, %d specified", charmDevice.CountMin, count)
   150  	}
   151  	return nil
   152  }