4ee603c382
This PR introduces a device hook that retrieves the device mount information for an allocation. It also updates the computed node class computation to take into account devices. TODO Fix the task runner unit test. The environment variable is being lost even though it is being properly set in the prestart hook.
96 lines
2 KiB
Go
96 lines
2 KiB
Go
package taskrunner
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
|
|
log "github.com/hashicorp/go-hclog"
|
|
"github.com/hashicorp/nomad/client/allocrunner/interfaces"
|
|
"github.com/hashicorp/nomad/client/devicemanager"
|
|
"github.com/hashicorp/nomad/plugins/device"
|
|
"github.com/hashicorp/nomad/plugins/drivers"
|
|
)
|
|
|
|
// deviceHook is used to retrieve device mounting information.
|
|
type deviceHook struct {
|
|
logger log.Logger
|
|
dm devicemanager.Manager
|
|
}
|
|
|
|
func newDeviceHook(dm devicemanager.Manager, logger log.Logger) *deviceHook {
|
|
h := &deviceHook{
|
|
dm: dm,
|
|
}
|
|
h.logger = logger.Named(h.Name())
|
|
return h
|
|
}
|
|
|
|
func (*deviceHook) Name() string {
|
|
return "devices"
|
|
}
|
|
|
|
func (h *deviceHook) Prestart(ctx context.Context, req *interfaces.TaskPrestartRequest, resp *interfaces.TaskPrestartResponse) error {
|
|
if len(req.TaskResources.Devices) == 0 {
|
|
resp.Done = true
|
|
return nil
|
|
}
|
|
|
|
// Capture the responses
|
|
var reservations []*device.ContainerReservation
|
|
for _, req := range req.TaskResources.Devices {
|
|
// Ask the device manager for the reservation information
|
|
res, err := h.dm.Reserve(req)
|
|
if err != nil {
|
|
return fmt.Errorf("failed to reserve device %s: %v", req.ID(), err)
|
|
}
|
|
|
|
reservations = append(reservations, res)
|
|
}
|
|
|
|
// Build the response
|
|
for _, res := range reservations {
|
|
for k, v := range res.Envs {
|
|
if resp.Env == nil {
|
|
resp.Env = make(map[string]string)
|
|
}
|
|
|
|
resp.Env[k] = v
|
|
}
|
|
|
|
for _, m := range res.Mounts {
|
|
resp.Mounts = append(resp.Mounts, convertMount(m))
|
|
}
|
|
|
|
for _, d := range res.Devices {
|
|
resp.Devices = append(resp.Devices, convertDevice(d))
|
|
}
|
|
}
|
|
|
|
resp.Done = true
|
|
return nil
|
|
}
|
|
|
|
func convertMount(in *device.Mount) *drivers.MountConfig {
|
|
if in == nil {
|
|
return nil
|
|
}
|
|
|
|
return &drivers.MountConfig{
|
|
TaskPath: in.TaskPath,
|
|
HostPath: in.HostPath,
|
|
Readonly: in.ReadOnly,
|
|
}
|
|
}
|
|
|
|
func convertDevice(in *device.DeviceSpec) *drivers.DeviceConfig {
|
|
if in == nil {
|
|
return nil
|
|
}
|
|
|
|
return &drivers.DeviceConfig{
|
|
TaskPath: in.TaskPath,
|
|
HostPath: in.HostPath,
|
|
Permissions: in.CgroupPerms,
|
|
}
|
|
}
|