github.com/projecteru2/core@v0.0.0-20240321043226-06bcc1c23f58/resource/cobalt/remap.go (about)

     1  package cobalt
     2  
     3  import (
     4  	"context"
     5  
     6  	"github.com/projecteru2/core/log"
     7  	"github.com/projecteru2/core/resource/plugins"
     8  	plugintypes "github.com/projecteru2/core/resource/plugins/types"
     9  	resourcetypes "github.com/projecteru2/core/resource/types"
    10  	"github.com/projecteru2/core/types"
    11  )
    12  
    13  // Remap remaps resource and returns engine args for workloads. format: {"workload-1": {"cpus": ["1-3"]}}
    14  // remap doesn't change resource args
    15  func (m Manager) Remap(ctx context.Context, nodename string, workloads []*types.Workload) (map[string]resourcetypes.Resources, error) {
    16  	logger := log.WithFunc("resource.cobalt.GetRemapArgs").WithField("node", nodename)
    17  	// call plugins to remap
    18  	resps, err := call(ctx, m.plugins, func(plugin plugins.Plugin) (*plugintypes.CalculateRemapResponse, error) {
    19  		workloadsResourceMap := map[string]plugintypes.WorkloadResource{}
    20  		for _, workload := range workloads {
    21  			workloadsResourceMap[workload.ID] = workload.Resources[plugin.Name()]
    22  		}
    23  		resp, err := plugin.CalculateRemap(ctx, nodename, workloadsResourceMap)
    24  		if err != nil {
    25  			logger.Errorf(ctx, err, "plugin %+v node %+v failed to remap", plugin.Name(), nodename)
    26  		}
    27  		return resp, err
    28  	})
    29  	if err != nil {
    30  		return nil, err
    31  	}
    32  
    33  	enginesParams := map[string]resourcetypes.Resources{}
    34  	// merge engine args
    35  	for plugin, resp := range resps {
    36  		for workloadID, engineParams := range resp.EngineParamsMap {
    37  			if _, ok := enginesParams[workloadID]; !ok {
    38  				enginesParams[workloadID] = resourcetypes.Resources{plugin.Name(): resourcetypes.RawParams{}}
    39  			}
    40  			v, err := m.mergeEngineParams(ctx, enginesParams[workloadID][plugin.Name()], engineParams)
    41  			if err != nil {
    42  				logger.Error(ctx, err, "invalid engine args")
    43  				return nil, err
    44  			}
    45  			enginesParams[workloadID][plugin.Name()] = v
    46  		}
    47  	}
    48  
    49  	return enginesParams, nil
    50  }
    51  
    52  // mergeEngineParams e.g. {"file": ["/bin/sh:/bin/sh"], "cpu": 1.2, "cpu-bind": true} + {"file": ["/bin/ls:/bin/ls"], "mem": "1PB"}
    53  // => {"file": ["/bin/sh:/bin/sh", "/bin/ls:/bin/ls"], "cpu": 1.2, "cpu-bind": true, "mem": "1PB"}
    54  func (m Manager) mergeEngineParams(ctx context.Context, m1 plugintypes.EngineParams, m2 plugintypes.EngineParams) (plugintypes.EngineParams, error) {
    55  	r := plugintypes.EngineParams{}
    56  	for key, value := range m1 {
    57  		r[key] = value
    58  	}
    59  	for key, value := range m2 {
    60  		if _, ok := r[key]; ok {
    61  			// only two string slices can be merged
    62  			_, ok1 := r[key].([]string)
    63  			_, ok2 := value.([]string)
    64  			if !ok1 || !ok2 {
    65  				log.WithFunc("resource.cobalt.mergeEngineParams").Errorf(ctx, types.ErrInvalidEngineArgs, "only two string slices can be merged! error key %+v, m1[key] = %+v, m2[key] = %+v", key, m1[key], m2[key])
    66  				return nil, types.ErrInvalidEngineArgs
    67  			}
    68  			r[key] = append(r[key].([]string), value.([]string)...)
    69  		} else {
    70  			r[key] = value
    71  		}
    72  	}
    73  	return r, nil
    74  }