cdi: support optional devices

Signed-off-by: CrazyMax <[email protected]>

Author: crazy-max

client/client_test.go

@@ -11027,8 +11027,9 @@ devices:
 		st = busybox.Run(append(ro, llb.Shlex(cmd), llb.Dir("/wd"))...).AddMount("/wd", st)
 	}
 
-	run(`sh -c 'env|sort | tee foo.env'`, llb.AddCDIDevice("vendor1.com/device=foo"))
-	run(`sh -c 'env|sort | tee bar.env'`, llb.AddCDIDevice("vendor2.com/device=bar"))
+	run(`sh -c 'env|sort | tee foo.env'`, llb.AddCDIDevice(llb.CDIDeviceName("vendor1.com/device=foo")))
+	run(`sh -c 'env|sort | tee bar.env'`, llb.AddCDIDevice(llb.CDIDeviceName("vendor2.com/device=bar")))
+	run(`ls`, llb.AddCDIDevice(llb.CDIDeviceName("vendor3.com/device=baz"), llb.CDIDeviceOptional))
 
 	def, err := st.Marshal(sb.Context())
 	require.NoError(t, err)

client/llb/exec.go

@@ -60,6 +60,7 @@ type ExecOp struct {
 	isValidated bool
 	secrets     []SecretInfo
 	ssh         []SSHInfo
+	cdiDevices  []CDIDeviceInfo
 }
 
 func (e *ExecOp) AddMount(target string, source Output, opt ...MountOption) Output {
@@ -267,21 +268,6 @@ func (e *ExecOp) Marshal(ctx context.Context, c *Constraints) (digest.Digest, []
 		Security: security,
 	}
 
-	cdiDevices, err := getCDIDevice(e.base)(ctx, c)
-	if err != nil {
-		return "", nil, nil, nil, err
-	}
-	if len(cdiDevices) > 0 {
-		addCap(&e.constraints, pb.CapExecMetaCDI)
-		cd := make([]*pb.CDIDevice, len(cdiDevices))
-		for i, d := range cdiDevices {
-			cd[i] = &pb.CDIDevice{
-				Name: d.Name,
-			}
-		}
-		peo.CdiDevices = cd
-	}
-
 	if network != NetModeSandbox {
 		addCap(&e.constraints, pb.CapExecMetaNetwork)
 	}
@@ -337,6 +323,18 @@ func (e *ExecOp) Marshal(ctx context.Context, c *Constraints) (digest.Digest, []
 		addCap(&e.constraints, pb.CapExecMountSSH)
 	}
 
+	if len(e.cdiDevices) > 0 {
+		addCap(&e.constraints, pb.CapExecMetaCDI)
+		cd := make([]*pb.CDIDevice, len(e.cdiDevices))
+		for i, d := range e.cdiDevices {
+			cd[i] = &pb.CDIDevice{
+				Name:     d.Name,
+				Optional: d.Optional,
+			}
+		}
+		peo.CdiDevices = cd
+	}
+
 	if e.constraints.Platform == nil {
 		p, err := getPlatform(e.base)(ctx, c)
 		if err != nil {
@@ -640,12 +638,41 @@ func AddUlimit(name UlimitName, soft int64, hard int64) RunOption {
 	})
 }
 
-func AddCDIDevice(name string) RunOption {
+func AddCDIDevice(opts ...CDIDeviceOption) RunOption {
 	return runOptionFunc(func(ei *ExecInfo) {
-		ei.State = ei.State.AddCDIDevice(name)
+		c := &CDIDeviceInfo{}
+		for _, opt := range opts {
+			opt.SetCDIDeviceOption(c)
+		}
+		ei.CDIDevices = append(ei.CDIDevices, *c)
+	})
+}
+
+type CDIDeviceOption interface {
+	SetCDIDeviceOption(*CDIDeviceInfo)
+}
+
+type cdiDeviceOptionFunc func(*CDIDeviceInfo)
+
+func (fn cdiDeviceOptionFunc) SetCDIDeviceOption(ci *CDIDeviceInfo) {
+	fn(ci)
+}
+
+func CDIDeviceName(name string) CDIDeviceOption {
+	return cdiDeviceOptionFunc(func(ci *CDIDeviceInfo) {
+		ci.Name = name
 	})
 }
 
+var CDIDeviceOptional = cdiDeviceOptionFunc(func(ci *CDIDeviceInfo) {
+	ci.Optional = true
+})
+
+type CDIDeviceInfo struct {
+	Name     string
+	Optional bool
+}
+
 func ValidExitCodes(codes ...int) RunOption {
 	return runOptionFunc(func(ei *ExecInfo) {
 		ei.State = validExitCodes(codes...)(ei.State)
@@ -837,6 +864,7 @@ type ExecInfo struct {
 	ProxyEnv       *ProxyEnv
 	Secrets        []SecretInfo
 	SSH            []SSHInfo
+	CDIDevices     []CDIDeviceInfo
 }
 
 type MountInfo struct {

client/llb/meta.go

@@ -24,7 +24,6 @@ var (
 	keyExtraHost      = contextKeyT("llb.exec.extrahost")
 	keyHostname       = contextKeyT("llb.exec.hostname")
 	keyUlimit         = contextKeyT("llb.exec.ulimit")
-	keyDevice         = contextKeyT("llb.exec.device")
 	keyCgroupParent   = contextKeyT("llb.exec.cgroup.parent")
 	keyUser           = contextKeyT("llb.exec.user")
 	keyValidExitCodes = contextKeyT("llb.exec.validexitcodes")
@@ -306,33 +305,6 @@ func getUlimit(s State) func(context.Context, *Constraints) ([]*pb.Ulimit, error
 	}
 }
 
-func cdiDevice(name string) StateOption {
-	return func(s State) State {
-		return s.withValue(keyDevice, func(ctx context.Context, c *Constraints) (interface{}, error) {
-			v, err := getCDIDevice(s)(ctx, c)
-			if err != nil {
-				return nil, err
-			}
-			return append(v, &pb.CDIDevice{
-				Name: name,
-			}), nil
-		})
-	}
-}
-
-func getCDIDevice(s State) func(context.Context, *Constraints) ([]*pb.CDIDevice, error) {
-	return func(ctx context.Context, c *Constraints) ([]*pb.CDIDevice, error) {
-		v, err := s.getValue(keyDevice)(ctx, c)
-		if err != nil {
-			return nil, err
-		}
-		if v != nil {
-			return v.([]*pb.CDIDevice), nil
-		}
-		return nil, nil
-	}
-}
-
 func cgroupParent(cp string) StateOption {
 	return func(s State) State {
 		return s.WithValue(keyCgroupParent, cp)

client/llb/state.go

@@ -295,6 +295,7 @@ func (s State) Run(ro ...RunOption) ExecState {
 	}
 	exec.secrets = ei.Secrets
 	exec.ssh = ei.SSH
+	exec.cdiDevices = ei.CDIDevices
 
 	return ExecState{
 		State: s.WithOutput(exec.Output()),
@@ -476,12 +477,6 @@ func (s State) AddUlimit(name UlimitName, soft int64, hard int64) State {
 	return ulimit(name, soft, hard)(s)
 }
 
-// AddCDIDevice sets the fully qualified CDI device name.
-// https://github.com/cncf-tags/container-device-interface/blob/main/SPEC.md
-func (s State) AddCDIDevice(name string) State {
-	return cdiDevice(name)(s)
-}
-
 // WithCgroupParent sets the parent cgroup for any containers created from this state.
 // This is useful when you want to apply resource constraints to a group of containers.
 // Cgroups are Linux specific and only applies to containers created from this state such as via `[State.Run]`

executor/oci/spec_linux.go

@@ -157,29 +157,43 @@ func generateCDIOpts(manager *cdi.Cache, devices []*pb.CDIDevice) ([]oci.SpecOpt
 	if len(devices) == 0 {
 		return nil, nil
 	}
-	var dd []string
-	for _, d := range devices {
-		if d == nil {
-			continue
-		}
-		if _, _, _, err := parser.ParseQualifiedName(d.Name); err != nil {
-			return nil, errors.Wrapf(err, "invalid CDI device name %s", d.Name)
-		}
-		dd = append(dd, d.Name)
-	}
 
-	withCDIDevices := func(devices ...string) oci.SpecOpts {
+	withCDIDevices := func(devices []*pb.CDIDevice) oci.SpecOpts {
 		return func(ctx context.Context, _ oci.Client, c *containers.Container, s *specs.Spec) error {
-			if len(devices) == 0 {
-				return nil
-			}
 			if err := manager.Refresh(); err != nil {
 				bklog.G(ctx).Warnf("CDI registry refresh failed: %v", err)
 			}
-			bklog.G(ctx).Debugf("Injecting CDI devices %v", devices)
-			if _, err := manager.InjectDevices(s, devices...); err != nil {
+
+			registeredDevices := manager.ListDevices()
+			isDeviceRegistered := func(device *pb.CDIDevice) bool {
+				for _, name := range registeredDevices {
+					if device.Name == name {
+						return true
+					}
+				}
+				return false
+			}
+
+			var dd []string
+			for _, d := range devices {
+				if d == nil {
+					continue
+				}
+				if _, _, _, err := parser.ParseQualifiedName(d.Name); err != nil {
+					return errors.Wrapf(err, "invalid CDI device name %s", d.Name)
+				}
+				if !isDeviceRegistered(d) && d.Optional {
+					bklog.G(ctx).Warnf("Optional CDI device %q is not registered", d.Name)
+					continue
+				}
+				dd = append(dd, d.Name)
+			}
+
+			bklog.G(ctx).Debugf("Injecting CDI devices %v", dd)
+			if _, err := manager.InjectDevices(s, dd...); err != nil {
 				return errors.Wrapf(err, "CDI device injection failed")
 			}
+
 			// One crucial thing to keep in mind is that CDI device injection
 			// might add OCI Spec environment variables, hooks, and mounts as
 			// well. Therefore, it is important that none of the corresponding
@@ -189,7 +203,7 @@ func generateCDIOpts(manager *cdi.Cache, devices []*pb.CDIDevice) ([]oci.SpecOpt
 	}
 
 	return []oci.SpecOpts{
-		withCDIDevices(dd...),
+		withCDIDevices(devices),
 	}, nil
 }
 

frontend/dockerfile/dockerfile2llb/convert.go

@@ -1305,10 +1305,15 @@ func dispatchRun(d *dispatchState, c *instructions.RunCommand, proxy *llb.ProxyE
 		}
 	}
 
-	// TODO: use entitlements
 	if dopt.llbCaps != nil && dopt.llbCaps.Supports(pb.CapExecMetaCDI) == nil {
-		for _, u := range dopt.devices {
-			opt = append(opt, llb.AddCDIDevice(u.Name))
+		for _, device := range dopt.devices {
+			deviceOpts := []llb.CDIDeviceOption{
+				llb.CDIDeviceName(device.Name),
+			}
+			if device.Optional {
+				deviceOpts = append(deviceOpts, llb.CDIDeviceOptional)
+			}
+			opt = append(opt, llb.AddCDIDevice(deviceOpts...))
 		}
 		runDevices, err := dispatchRunDevices(c)
 		if err != nil {

frontend/dockerfile/dockerfile2llb/convert_rundevice.go

@@ -11,7 +11,7 @@ func dispatchRunDevices(c *instructions.RunCommand) ([]llb.RunOption, error) {
 	var out []llb.RunOption
 	devices := instructions.GetDevices(c)
 	for _, device := range devices {
-		out = append(out, llb.AddCDIDevice(device))
+		out = append(out, llb.AddCDIDevice(llb.CDIDeviceName(device), llb.CDIDeviceOptional))
 	}
 	return out, nil
 }