feat: enhance kubernetes driver

Signed-off-by: Wang Jinglei <morlay.null@gmail.com>

README.md

@@ -411,6 +411,7 @@ Passes additional driver-specific options. Details for each driver:
   - `image=IMAGE` - Sets the container image to be used for running buildkit.
   - `namespace=NS` - Sets the Kubernetes namespace. Defaults to the current namespace.
   - `replicas=N` - Sets the number of `Pod` replicas. Defaults to 1.
+  - `nodeselector="label1=value1,label2=value2"` - Sets the kv of `Pod` nodeSelector. No Defaults. Example `nodeselector=kubernetes.io/arch=arm64`
   - `rootless=(true|false)` - Run the container as a non-root user without `securityContext.privileged`. [Using Ubuntu host kernel is recommended](https://github.com/moby/buildkit/blob/master/docs/rootless.md). Defaults to false.
   - `loadbalance=(sticky|random)` - Load-balancing strategy. If set to "sticky", the pod is chosen using the hash of the context path. Defaults to "sticky"
 

commands/create.go

@@ -142,6 +142,12 @@ func runCreate(dockerCli command.Cli, in createOptions, args []string) error {
 				return err
 			}
 		}
+
+		if in.driver == "kubernetes" {
+			// naming endpoint to make --append works
+			ep = fmt.Sprintf("%s://%s?deployment=%s", in.driver, in.name, in.nodeName)
+		}
+
 		m, err := csvToMap(in.driverOpts)
 		if err != nil {
 			return err

commands/util.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"os"
 	"path/filepath"
+	"strings"
 
 	"github.com/docker/buildx/build"
 	"github.com/docker/buildx/driver"
@@ -192,8 +193,7 @@ func driversForNodeGroup(ctx context.Context, dockerCli command.Cli, ng *store.N
 				if kcc == nil {
 					kcc = driver.KubeClientConfigInCluster{}
 				}
-
-				d, err := driver.GetDriver(ctx, "buildx_buildkit_"+n.Name, f, dockerapi, kcc, n.Flags, n.ConfigFile, n.DriverOpts, contextPathHash)
+				d, err := driver.GetDriver(ctx, "buildx_buildkit_"+n.Name, f, dockerapi, kcc, n.Flags, n.ConfigFile, assignDriverOptsByDriverInfo(n.DriverOpts, di), contextPathHash)
 				if err != nil {
 					di.Err = err
 					return nil
@@ -211,6 +211,20 @@ func driversForNodeGroup(ctx context.Context, dockerCli command.Cli, ng *store.N
 	return dis, nil
 }
 
+// pass platform as driver opts to provide for some drive, like kubernetes
+func assignDriverOptsByDriverInfo(opts map[string]string, driveInfo build.DriverInfo) map[string]string {
+	m := map[string]string{}
+
+	if len(driveInfo.Platform) > 0 {
+		m["platform"] = strings.Join(platformutil.Format(driveInfo.Platform), ",")
+	}
+
+	for key := range opts {
+		m[key] = opts[key]
+	}
+	return m
+}
+
 // clientForEndpoint returns a docker client for an endpoint
 func clientForEndpoint(dockerCli command.Cli, name string) (dockerclient.APIClient, error) {
 	list, err := dockerCli.ContextStore().List()
@@ -355,24 +369,28 @@ func loadNodeGroupData(ctx context.Context, dockerCli command.Cli, ngi *nginfo)
 	if eg.Wait(); err != nil {
 		return err
 	}
-	for _, di := range ngi.drivers {
-		// dynamic nodes are used in Kubernetes driver.
-		// Kubernetes pods are dynamically mapped to BuildKit Nodes.
-		if di.info != nil && len(di.info.DynamicNodes) > 0 {
-			var drivers []dinfo
-			for i := 0; i < len(di.info.DynamicNodes); i++ {
-				// all []dinfo share *build.DriverInfo and *driver.Info
-				diClone := di
-				if pl := di.info.DynamicNodes[i].Platforms; len(pl) > 0 {
-					diClone.platforms = pl
+
+	// skip when multi drivers
+	if len(ngi.drivers) == 1 {
+		for _, di := range ngi.drivers {
+			// dynamic nodes are used in Kubernetes driver.
+			// Kubernetes pods are dynamically mapped to BuildKit Nodes.
+			if di.info != nil && len(di.info.DynamicNodes) > 0 {
+				var drivers []dinfo
+				for i := 0; i < len(di.info.DynamicNodes); i++ {
+					// all []dinfo share *build.DriverInfo and *driver.Info
+					diClone := di
+					if pl := di.info.DynamicNodes[i].Platforms; len(pl) > 0 {
+						diClone.platforms = pl
+					}
+					drivers = append(drivers, di)
 				}
-				drivers = append(drivers, di)
+				// not append (remove the static nodes in the store)
+				ngi.ng.Nodes = di.info.DynamicNodes
+				ngi.ng.Dynamic = true
+				ngi.drivers = drivers
+				return nil
 			}
-			// not append (remove the static nodes in the store)
-			ngi.ng.Nodes = di.info.DynamicNodes
-			ngi.ng.Dynamic = true
-			ngi.drivers = drivers
-			return nil
 		}
 	}
 	return nil

driver/kubernetes/driver.go

@@ -4,12 +4,15 @@ import (
 	"context"
 	"fmt"
 	"net"
+	"strings"
 	"time"
 
 	"github.com/docker/buildx/driver"
 	"github.com/docker/buildx/driver/kubernetes/execconn"
+	"github.com/docker/buildx/driver/kubernetes/manifest"
 	"github.com/docker/buildx/driver/kubernetes/podchooser"
 	"github.com/docker/buildx/store"
+	"github.com/docker/buildx/util/platformutil"
 	"github.com/docker/buildx/util/progress"
 	"github.com/moby/buildkit/client"
 	"github.com/pkg/errors"
@@ -109,6 +112,16 @@ func (d *Driver) Info(ctx context.Context) (*driver.Info, error) {
 			Name: p.Name,
 			// Other fields are unset (TODO: detect real platforms)
 		}
+
+		if p.Annotations != nil {
+			if p, ok := p.Annotations[manifest.AnnotationPlatform]; ok {
+				ps, err := platformutil.Parse(strings.Split(p, ","))
+				if err == nil {
+					node.Platforms = ps
+				}
+			}
+		}
+
 		dynNodes = append(dynNodes, node)
 	}
 	return &driver.Info{

driver/kubernetes/factory.go

@@ -9,6 +9,7 @@ import (
 	"github.com/docker/buildx/driver/bkimage"
 	"github.com/docker/buildx/driver/kubernetes/manifest"
 	"github.com/docker/buildx/driver/kubernetes/podchooser"
+	"github.com/docker/buildx/util/platformutil"
 	dockerclient "github.com/docker/docker/client"
 	"github.com/pkg/errors"
 	"k8s.io/client-go/kubernetes"
@@ -90,6 +91,24 @@ func (f *factory) New(ctx context.Context, cfg driver.InitConfig) (driver.Driver
 				return nil, err
 			}
 			deploymentOpt.Image = bkimage.DefaultRootlessImage
+		case "platform":
+			if v != "" {
+				platforms, err := platformutil.Parse(strings.Split(v, ","))
+				if err != nil {
+					return nil, err
+				}
+				deploymentOpt.Platforms = platforms
+			}
+		case "nodeselector":
+			kvs := strings.Split(strings.Trim(v, `"`), ",")
+			s := map[string]string{}
+			for i := range kvs {
+				kv := strings.Split(kvs[i], "=")
+				if len(kv) == 2 {
+					s[kv[0]] = kv[1]
+				}
+			}
+			deploymentOpt.NodeSelector = s
 		case "loadbalance":
 			switch v {
 			case LoadbalanceSticky:

driver/kubernetes/manifest/manifest.go

@@ -1,6 +1,10 @@
 package manifest
 
 import (
+	"strings"
+
+	"github.com/docker/buildx/util/platformutil"
+	v1 "github.com/opencontainers/image-spec/specs-go/v1"
 	appsv1 "k8s.io/api/apps/v1"
 	corev1 "k8s.io/api/core/v1"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
@@ -13,28 +17,38 @@ type DeploymentOpt struct {
 	Replicas      int
 	BuildkitFlags []string
 	Rootless      bool
+	NodeSelector  map[string]string
+	Platforms     []v1.Platform
 }
 
 const (
-	containerName = "buildkitd"
+	containerName      = "buildkitd"
+	AnnotationPlatform = "buildx.docker.com/platform"
 )
 
 func NewDeployment(opt *DeploymentOpt) (*appsv1.Deployment, error) {
 	labels := map[string]string{
 		"app": opt.Name,
 	}
+	annotations := map[string]string{}
 	replicas := int32(opt.Replicas)
 	privileged := true
 	args := opt.BuildkitFlags
+
+	if len(opt.Platforms) > 0 {
+		annotations[AnnotationPlatform] = strings.Join(platformutil.Format(opt.Platforms), ",")
+	}
+
 	d := &appsv1.Deployment{
 		TypeMeta: metav1.TypeMeta{
 			APIVersion: appsv1.SchemeGroupVersion.String(),
 			Kind:       "Deployment",
 		},
 		ObjectMeta: metav1.ObjectMeta{
-			Namespace: opt.Namespace,
-			Name:      opt.Name,
-			Labels:    labels,
+			Namespace:   opt.Namespace,
+			Name:        opt.Name,
+			Labels:      labels,
+			Annotations: annotations,
 		},
 		Spec: appsv1.DeploymentSpec{
 			Replicas: &replicas,
@@ -43,7 +57,8 @@ func NewDeployment(opt *DeploymentOpt) (*appsv1.Deployment, error) {
 			},
 			Template: corev1.PodTemplateSpec{
 				ObjectMeta: metav1.ObjectMeta{
-					Labels: labels,
+					Labels:      labels,
+					Annotations: annotations,
 				},
 				Spec: corev1.PodSpec{
 					Containers: []corev1.Container{
@@ -72,6 +87,11 @@ func NewDeployment(opt *DeploymentOpt) (*appsv1.Deployment, error) {
 			return nil, err
 		}
 	}
+
+	if len(opt.NodeSelector) > 0 {
+		d.Spec.Template.Spec.NodeSelector = opt.NodeSelector
+	}
+
 	return d, nil
 }