Skip to content

Commit 9a9ec5f

Browse files
authored
refactor PodSetInfo flow to enable immutability (#45)
Store the PodSetInfos provided by Kueue's RunWithPodSetInfo method in a separate array in the AppWrapperComponent instead of immediately injecting them into the raw templates. This simplifies both the Workload controller (which we hope to upstream to Kueue) and the enforcement of RBAC by allowing us to require that Component.Template is immutable.
1 parent efd67f8 commit 9a9ec5f

11 files changed

+224
-208
lines changed

api/v1beta2/appwrapper_types.go

+11
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@ limitations under the License.
1717
package v1beta2
1818

1919
import (
20+
corev1 "k8s.io/api/core/v1"
2021
metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
2122
runtime "k8s.io/apimachinery/pkg/runtime"
2223
)
@@ -35,6 +36,9 @@ type AppWrapperComponent struct {
3536
// PodSets contained in the component
3637
PodSets []AppWrapperPodSet `json:"podSets"`
3738

39+
// PodSetInfos assigned to the Component by Kueue
40+
PodSetInfos []AppWrapperPodSetInfo `json:"podSetInfos,omitempty"`
41+
3842
// +kubebuilder:pruning:PreserveUnknownFields
3943
// +kubebuilder:validation:EmbeddedResource
4044
// Template for the component
@@ -50,6 +54,13 @@ type AppWrapperPodSet struct {
5054
Path string `json:"path"`
5155
}
5256

57+
type AppWrapperPodSetInfo struct {
58+
Annotations map[string]string `json:"annotations,omitempty"`
59+
Labels map[string]string `json:"labels,omitempty"`
60+
NodeSelector map[string]string `json:"nodeSelector,omitempty"`
61+
Tolerations []corev1.Toleration `json:"tolerations,omitempty"`
62+
}
63+
5364
// AppWrapperStatus defines the observed state of the appwrapper
5465
type AppWrapperStatus struct {
5566
// Phase of the AppWrapper object

api/v1beta2/zz_generated.deepcopy.go

+53-2
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

cmd/main.go

-1
Original file line numberDiff line numberDiff line change
@@ -73,7 +73,6 @@ func main() {
7373
flag.BoolVar(&enableHTTP2, "enable-http2", false,
7474
"If set, HTTP/2 will be enabled for the metrics and webhook servers")
7575
flag.BoolVar(&config.ManageJobsWithoutQueueName, "manage-no-queue", true, "Manage AppWrappers without queue names")
76-
flag.StringVar(&config.ServiceAccountName, "service-account", "", "Service account name for controller")
7776
opts := zap.Options{
7877
Development: true,
7978
}

config/crd/bases/workload.codeflare.dev_appwrappers.yaml

+56
Original file line numberDiff line numberDiff line change
@@ -48,6 +48,62 @@ spec:
4848
items:
4949
description: AppWrapperComponent describes a wrapped resource
5050
properties:
51+
podSetInfos:
52+
description: PodSetInfos assigned to the Component by Kueue
53+
items:
54+
properties:
55+
annotations:
56+
additionalProperties:
57+
type: string
58+
type: object
59+
labels:
60+
additionalProperties:
61+
type: string
62+
type: object
63+
nodeSelector:
64+
additionalProperties:
65+
type: string
66+
type: object
67+
tolerations:
68+
items:
69+
description: |-
70+
The pod this Toleration is attached to tolerates any taint that matches
71+
the triple <key,value,effect> using the matching operator <operator>.
72+
properties:
73+
effect:
74+
description: |-
75+
Effect indicates the taint effect to match. Empty means match all taint effects.
76+
When specified, allowed values are NoSchedule, PreferNoSchedule and NoExecute.
77+
type: string
78+
key:
79+
description: |-
80+
Key is the taint key that the toleration applies to. Empty means match all taint keys.
81+
If the key is empty, operator must be Exists; this combination means to match all values and all keys.
82+
type: string
83+
operator:
84+
description: |-
85+
Operator represents a key's relationship to the value.
86+
Valid operators are Exists and Equal. Defaults to Equal.
87+
Exists is equivalent to wildcard for value, so that a pod can
88+
tolerate all taints of a particular category.
89+
type: string
90+
tolerationSeconds:
91+
description: |-
92+
TolerationSeconds represents the period of time the toleration (which must be
93+
of effect NoExecute, otherwise this field is ignored) tolerates the taint. By default,
94+
it is not set, which means tolerate the taint forever (do not evict). Zero and
95+
negative values will be treated as 0 (evict immediately) by the system.
96+
format: int64
97+
type: integer
98+
value:
99+
description: |-
100+
Value is the taint value the toleration matches to.
101+
If the operator is Exists, the value should be empty, otherwise just a regular string.
102+
type: string
103+
type: object
104+
type: array
105+
type: object
106+
type: array
51107
podSets:
52108
description: PodSets contained in the component
53109
items:

config/default/manager_auth_proxy_patch.yaml

-1
Original file line numberDiff line numberDiff line change
@@ -34,7 +34,6 @@ spec:
3434
memory: 64Mi
3535
- name: manager
3636
args:
37-
- "--service-account=system:serviceaccount:$(MY_NAMESPACE):$(MY_SERVICE_ACCOUNT_NAME)"
3837
- "--health-probe-bind-address=:8081"
3938
- "--metrics-bind-address=127.0.0.1:8080"
4039
- "--leader-elect"

config/manager/manager.yaml

-10
Original file line numberDiff line numberDiff line change
@@ -66,7 +66,6 @@ spec:
6666
- /manager
6767
args:
6868
- --leader-elect
69-
- "--service-account=system:serviceaccount:$(MY_NAMESPACE):$(MY_SERVICE_ACCOUNT_NAME)"
7069
image: controller:latest
7170
name: manager
7271
securityContext:
@@ -93,14 +92,5 @@ spec:
9392
requests:
9493
cpu: 10m
9594
memory: 64Mi
96-
env:
97-
- name: MY_SERVICE_ACCOUNT_NAME
98-
valueFrom:
99-
fieldRef:
100-
fieldPath: spec.serviceAccountName
101-
- name: MY_NAMESPACE
102-
valueFrom:
103-
fieldRef:
104-
fieldPath: metadata.namespace
10595
serviceAccountName: controller-manager
10696
terminationGracePeriodSeconds: 10

internal/controller/appwrapper_config.go

+1-2
Original file line numberDiff line numberDiff line change
@@ -26,8 +26,7 @@ import (
2626
)
2727

2828
type AppWrapperConfig struct {
29-
ManageJobsWithoutQueueName bool `json:"manageJobsWithoutQueueName,omitempty"`
30-
ServiceAccountName string `json:"serviceAccountName,omitempty"`
29+
ManageJobsWithoutQueueName bool `json:"manageJobsWithoutQueueName,omitempty"`
3130
}
3231

3332
// SetupWithManager creates and configures all components of the AppWrapper controller

internal/controller/appwrapper_controller.go

+85-25
Original file line numberDiff line numberDiff line change
@@ -34,8 +34,8 @@ import (
3434
"sigs.k8s.io/controller-runtime/pkg/handler"
3535
"sigs.k8s.io/controller-runtime/pkg/log"
3636
"sigs.k8s.io/controller-runtime/pkg/reconcile"
37-
"sigs.k8s.io/kueue/pkg/controller/constants"
38-
"sigs.k8s.io/kueue/pkg/controller/jobframework"
37+
"sigs.k8s.io/kueue/pkg/podset"
38+
utilmaps "sigs.k8s.io/kueue/pkg/util/maps"
3939

4040
workloadv1beta2 "github.com/project-codeflare/appwrapper/api/v1beta2"
4141
)
@@ -299,38 +299,98 @@ func parseComponent(aw *workloadv1beta2.AppWrapper, raw []byte) (*unstructured.U
299299
return obj, nil
300300
}
301301

302-
func parseComponents(aw *workloadv1beta2.AppWrapper) ([]client.Object, error) {
303-
components := aw.Spec.Components
304-
objects := make([]client.Object, len(components))
305-
for i, component := range components {
306-
obj, err := parseComponent(aw, component.Template.Raw)
307-
if err != nil {
308-
return nil, err
302+
func materializeObject(aw *workloadv1beta2.AppWrapper, component *workloadv1beta2.AppWrapperComponent) (client.Object, error) {
303+
toMap := func(x interface{}) map[string]string {
304+
if x == nil {
305+
return nil
306+
} else {
307+
if sm, ok := x.(map[string]string); ok {
308+
return sm
309+
} else if im, ok := x.(map[string]interface{}); ok {
310+
sm := make(map[string]string, len(im))
311+
for k, v := range im {
312+
str, ok := v.(string)
313+
if ok {
314+
sm[k] = str
315+
} else {
316+
sm[k] = fmt.Sprint(v)
317+
}
318+
}
319+
return sm
320+
} else {
321+
return nil
322+
}
309323
}
310-
objects[i] = obj
311324
}
312-
return objects, nil
313-
}
325+
awLabels := map[string]string{AppWrapperLabel: aw.Name}
314326

315-
func (r *AppWrapperReconciler) createComponents(ctx context.Context, aw *workloadv1beta2.AppWrapper) (error, bool) {
316-
objects, err := parseComponents(aw)
327+
obj, err := parseComponent(aw, component.Template.Raw)
317328
if err != nil {
318-
return err, true // fatal
329+
return nil, err
319330
}
320331

321-
ref := &metav1.OwnerReference{APIVersion: GVK.GroupVersion().String(), Kind: GVK.Kind, Name: aw.Name, UID: aw.UID}
322-
myWorkloadName, err := jobframework.GetWorkloadNameForOwnerRef(ref)
323-
if err != nil {
324-
return err, true
332+
for podSetsIdx, podSet := range component.PodSets {
333+
toInject := component.PodSetInfos[podSetsIdx]
334+
335+
p, err := getRawTemplate(obj.UnstructuredContent(), podSet.Path)
336+
if err != nil {
337+
return nil, err // Should not happen, path validity is enforced by validateAppWrapperInvariants
338+
}
339+
if md, ok := p["metadata"]; !ok || md == nil {
340+
p["metadata"] = make(map[string]interface{})
341+
}
342+
metadata := p["metadata"].(map[string]interface{})
343+
spec := p["spec"].(map[string]interface{}) // Must exist, enforced by validateAppWrapperInvariants
344+
345+
// Annotations
346+
if len(toInject.Annotations) > 0 {
347+
existing := toMap(metadata["annotations"])
348+
if err := utilmaps.HaveConflict(existing, toInject.Annotations); err != nil {
349+
return nil, podset.BadPodSetsUpdateError("annotations", err)
350+
}
351+
metadata["annotations"] = utilmaps.MergeKeepFirst(existing, toInject.Annotations)
352+
}
353+
354+
// Labels
355+
mergedLabels := utilmaps.MergeKeepFirst(toInject.Labels, awLabels)
356+
existing := toMap(metadata["labels"])
357+
if err := utilmaps.HaveConflict(existing, mergedLabels); err != nil {
358+
return nil, podset.BadPodSetsUpdateError("labels", err)
359+
}
360+
metadata["labels"] = utilmaps.MergeKeepFirst(existing, mergedLabels)
361+
362+
// NodeSelectors
363+
if len(toInject.NodeSelector) > 0 {
364+
existing := toMap(metadata["nodeSelector"])
365+
if err := utilmaps.HaveConflict(existing, toInject.NodeSelector); err != nil {
366+
return nil, podset.BadPodSetsUpdateError("nodeSelector", err)
367+
}
368+
metadata["nodeSelector"] = utilmaps.MergeKeepFirst(existing, toInject.NodeSelector)
369+
}
370+
371+
// Tolerations
372+
if len(toInject.Tolerations) > 0 {
373+
if _, ok := spec["tolerations"]; !ok {
374+
spec["tolerations"] = []interface{}{}
375+
}
376+
tolerations := spec["tolerations"].([]interface{})
377+
for _, addition := range toInject.Tolerations {
378+
tolerations = append(tolerations, addition)
379+
}
380+
spec["tolerations"] = tolerations
381+
}
325382
}
326383

327-
for _, obj := range objects {
328-
annotations := obj.GetAnnotations()
329-
if annotations == nil {
330-
annotations = make(map[string]string)
384+
return obj, nil
385+
}
386+
387+
func (r *AppWrapperReconciler) createComponents(ctx context.Context, aw *workloadv1beta2.AppWrapper) (error, bool) {
388+
for _, component := range aw.Spec.Components {
389+
obj, err := materializeObject(aw, &component)
390+
if err != nil {
391+
return err, true
331392
}
332-
annotations[constants.ParentWorkloadAnnotation] = myWorkloadName
333-
obj.SetAnnotations(annotations)
393+
334394
if err := controllerutil.SetControllerReference(aw, obj, r.Scheme); err != nil {
335395
return err, true
336396
}

internal/controller/appwrapper_webhook.go

+3-11
Original file line numberDiff line numberDiff line change
@@ -84,7 +84,7 @@ func (w *AppWrapperWebhook) ValidateUpdate(ctx context.Context, oldObj, newObj r
8484
oldAW := oldObj.(*workloadv1beta2.AppWrapper)
8585
newAW := newObj.(*workloadv1beta2.AppWrapper)
8686

87-
allErrors := w.validateAppWrapperUpdate(ctx, oldAW, newAW)
87+
allErrors := w.validateAppWrapperUpdate(oldAW, newAW)
8888
if w.Config.ManageJobsWithoutQueueName || jobframework.QueueName((*AppWrapper)(newAW)) != "" {
8989
allErrors = append(allErrors, jobframework.ValidateUpdateForQueueName((*AppWrapper)(oldAW), (*AppWrapper)(newAW))...)
9090
allErrors = append(allErrors, jobframework.ValidateUpdateForWorkloadPriorityClassName((*AppWrapper)(oldAW), (*AppWrapper)(newAW))...)
@@ -201,16 +201,8 @@ func (w *AppWrapperWebhook) validateAppWrapperCreate(ctx context.Context, aw *wo
201201
return allErrors
202202
}
203203

204-
// validateAppWrapperUpdate enforces that AppWrapper.Spec.Components is deeply immutable
205-
func (w *AppWrapperWebhook) validateAppWrapperUpdate(ctx context.Context, old *workloadv1beta2.AppWrapper, new *workloadv1beta2.AppWrapper) field.ErrorList {
206-
// The AppWrapper controller must be allowed to mutate Spec.Components
207-
// to enable it to implement RunWithPodSetsInfo and RestorePodSetsInfo
208-
if request, err := admission.RequestFromContext(ctx); err == nil {
209-
if w.Config.ServiceAccountName != "" && request.UserInfo.Username == w.Config.ServiceAccountName {
210-
return field.ErrorList{}
211-
}
212-
}
213-
204+
// validateAppWrapperUpdate enforces deep immutablity of all fields that were validated by validateAppWrapperCreate
205+
func (w *AppWrapperWebhook) validateAppWrapperUpdate(old *workloadv1beta2.AppWrapper, new *workloadv1beta2.AppWrapper) field.ErrorList {
214206
allErrors := field.ErrorList{}
215207
msg := "attempt to change immutable field"
216208
componentsPath := field.NewPath("spec").Child("components")

internal/controller/suite_test.go

+1-1
Original file line numberDiff line numberDiff line change
@@ -161,7 +161,7 @@ var _ = BeforeSuite(func() {
161161
})
162162
Expect(err).NotTo(HaveOccurred())
163163

164-
awConfig := AppWrapperConfig{ManageJobsWithoutQueueName: true, ServiceAccountName: ctrlUserName}
164+
awConfig := AppWrapperConfig{ManageJobsWithoutQueueName: true}
165165
err = (&AppWrapperWebhook{Config: &awConfig}).SetupWebhookWithManager(mgr)
166166
Expect(err).NotTo(HaveOccurred())
167167

0 commit comments

Comments
 (0)