added nodeSelector support to opaque server (#1203)

namespacelabs · Jul 13, 2023 · 8bd035c · 8bd035c
2 parents 7a7eb69 + 901fbbb
commit 8bd035c
Show file tree

Hide file tree

Showing 6 changed files with 320 additions and 206 deletions.
diff --git a/internal/frontend/cuefrontendopaque/server.go b/internal/frontend/cuefrontendopaque/server.go
@@ -25,7 +25,7 @@ var (
 	extensionFields = []string{
 		"args", "env", "services", "ports", "unstable_permissions", "permissions", "probe", "probes", "security",
 		"sidecars", "mounts", "resources", "requires", "tolerations", "annotations",
-		"resourceLimits", "resourceRequests", "extensions",
+		"resourceLimits", "resourceRequests", "extensions", "nodeSelector",
 		// This is needed for the "spec" in server templates. This can't be a private field, otherwise it can't be overridden.
 		"spec"}
 
@@ -60,6 +60,7 @@ type cueServerExtension struct {
 	Security       *cueServerSecurity              `json:"security,omitempty"`
 	Tolerations    []*schema.Server_Toleration     `json:"tolerations,omitempty"`
 	Annotations    map[string]args.ResolvableValue `json:"annotations,omitempty"`
+	NodeSelector   map[string]string               `json:"nodeSelector,omitempty"`
 
 	Extensions []string `json:"extensions,omitempty"`
 }
@@ -334,6 +335,19 @@ func parseServerExtension(ctx context.Context, env *schema.Environment, pl parsi
 		}
 	}
 
+	if len(bits.NodeSelector) > 0 {
+		if err := parsing.RequireFeature(loc.Module, "experimental/container/nodeSelector"); err != nil {
+			return nil, fnerrors.AttachLocation(loc, err)
+		}
+
+		for k, v := range bits.NodeSelector {
+			out.NodeSelector = append(out.NodeSelector, &schema.NodeSelectorItem{
+				Key:   k,
+				Value: v,
+			})
+		}
+	}
+
 	if len(bits.Tolerations) > 0 {
 		if err := parsing.RequireFeature(loc.Module, "experimental/container/tolerations"); err != nil {
 			return nil, fnerrors.AttachLocation(loc, err)

diff --git a/internal/planning/deploy/deploy.go b/internal/planning/deploy/deploy.go
@@ -872,6 +872,7 @@ func PrepareRunOpts(ctx context.Context, stack *planning.Stack, srv planning.Pla
 	out.Probes = frag.Probe
 	out.Tolerations = frag.Toleration
 	out.Annotations = frag.Annotation
+	out.NodeSelector = frag.NodeSelector
 
 	if err := integrations.IntegrationFor(srv.Framework()).PrepareRun(ctx, srv, &out.MainContainer); err != nil {
 		return err

diff --git a/internal/runtime/interface.go b/internal/runtime/interface.go
@@ -237,15 +237,16 @@ type DeployableSpec struct {
 	Focused    bool // Set to true if the user explicitly asked for this object to be deployed.
 	Attachable AttachableKind
 
-	Description string
-	Class       schema.DeployableClass
-	Id          string // Must not be empty.
-	Name        string // Can be empty.
-	Volumes     []*schema.Volume
-	Permissions *schema.ServerPermissions
-	Replicas    int32
-	Tolerations []*schema.Server_Toleration
-	Annotations []*schema.NamedResolvable // Annotations that apply to individual pods.
+	Description  string
+	Class        schema.DeployableClass
+	Id           string // Must not be empty.
+	Name         string // Can be empty.
+	Volumes      []*schema.Volume
+	Permissions  *schema.ServerPermissions
+	Replicas     int32
+	Tolerations  []*schema.Server_Toleration
+	Annotations  []*schema.NamedResolvable // Annotations that apply to individual pods.
+	NodeSelector []*schema.NodeSelectorItem
 
 	MainContainer ContainerRunOpts
 	Sidecars      []SidecarRunOpts

diff --git a/internal/runtime/kubernetes/deployment.go b/internal/runtime/kubernetes/deployment.go
@@ -178,6 +178,13 @@ func prepareDeployment(ctx context.Context, target BoundNamespace, deployable ru
 		spec = spec.WithTolerations(t)
 	}
 
+	// nodeSelector := spec.NodeSelector
+	nodeSelector := map[string]string{}
+	for _, nodeSelectorItem := range deployable.NodeSelector {
+		nodeSelector[nodeSelectorItem.Key] = nodeSelectorItem.Value
+	}
+	spec = spec.WithNodeSelector(nodeSelector)
+
 	// Explicitly allow all pods on all available platforms.
 	// On GKE, workloads are not allowed on ARM nodes by default, even if all nodes are ARM.
 	// https://cloud.google.com/kubernetes-engine/docs/how-to/prepare-arm-workloads-for-deployment#overview