github.com/iqoqo/nomad@v0.11.3-0.20200911112621-d7021c74d101/client/allocrunner/taskrunner/device_hook.go (about)

     1  package taskrunner
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  
     7  	log "github.com/hashicorp/go-hclog"
     8  	"github.com/hashicorp/nomad/client/allocrunner/interfaces"
     9  	"github.com/hashicorp/nomad/client/devicemanager"
    10  	"github.com/hashicorp/nomad/plugins/device"
    11  	"github.com/hashicorp/nomad/plugins/drivers"
    12  )
    13  
    14  const (
    15  	// HookNameDevices is the name of the devices hook
    16  	HookNameDevices = "devices"
    17  )
    18  
    19  // deviceHook is used to retrieve device mounting information.
    20  type deviceHook struct {
    21  	logger log.Logger
    22  	dm     devicemanager.Manager
    23  }
    24  
    25  func newDeviceHook(dm devicemanager.Manager, logger log.Logger) *deviceHook {
    26  	h := &deviceHook{
    27  		dm: dm,
    28  	}
    29  	h.logger = logger.Named(h.Name())
    30  	return h
    31  }
    32  
    33  func (*deviceHook) Name() string {
    34  	return HookNameDevices
    35  }
    36  
    37  func (h *deviceHook) Prestart(ctx context.Context, req *interfaces.TaskPrestartRequest, resp *interfaces.TaskPrestartResponse) error {
    38  	//TODO Can the nil check be removed once the TODO in NewTaskRunner
    39  	//     where this is set is addressed?
    40  	if req.TaskResources == nil || len(req.TaskResources.Devices) == 0 {
    41  		resp.Done = true
    42  		return nil
    43  	}
    44  
    45  	// Capture the responses
    46  	var reservations []*device.ContainerReservation
    47  	for _, req := range req.TaskResources.Devices {
    48  		// Ask the device manager for the reservation information
    49  		res, err := h.dm.Reserve(req)
    50  		if err != nil {
    51  			return fmt.Errorf("failed to reserve device %s: %v", req.ID(), err)
    52  		}
    53  
    54  		reservations = append(reservations, res)
    55  	}
    56  
    57  	// Build the response
    58  	for _, res := range reservations {
    59  		for k, v := range res.Envs {
    60  			if resp.Env == nil {
    61  				resp.Env = make(map[string]string)
    62  			}
    63  
    64  			resp.Env[k] = v
    65  		}
    66  
    67  		for _, m := range res.Mounts {
    68  			resp.Mounts = append(resp.Mounts, convertMount(m))
    69  		}
    70  
    71  		for _, d := range res.Devices {
    72  			resp.Devices = append(resp.Devices, convertDevice(d))
    73  		}
    74  	}
    75  
    76  	resp.Done = true
    77  	return nil
    78  }
    79  
    80  func convertMount(in *device.Mount) *drivers.MountConfig {
    81  	return &drivers.MountConfig{
    82  		TaskPath: in.TaskPath,
    83  		HostPath: in.HostPath,
    84  		Readonly: in.ReadOnly,
    85  	}
    86  }
    87  
    88  func convertDevice(in *device.DeviceSpec) *drivers.DeviceConfig {
    89  	return &drivers.DeviceConfig{
    90  		TaskPath:    in.TaskPath,
    91  		HostPath:    in.HostPath,
    92  		Permissions: in.CgroupPerms,
    93  	}
    94  }