From 338b4cc1d18b4f4b74a34741c87a6bf5780de70f Mon Sep 17 00:00:00 2001 From: Julien Mancuso Date: Mon, 31 Aug 2026 13:14:15 -0600 Subject: [PATCH] feat(api): propagate SnapshotJob PodSnapshot metadata Signed-off-by: Julien Mancuso --- api/v1alpha1/crds/embed_test.go | 88 +++++++++++++------ .../crds/nvidia.com_snapshotjobs.yaml | 36 ++++++++ api/v1alpha1/snapshotjob_types.go | 32 ++++++- api/v1alpha1/snapshotjob_types_test.go | 16 +++- api/v1alpha1/zz_generated.deepcopy.go | 34 +++++++ .../crds/nvidia.com_snapshotjobs.yaml | 36 ++++++++ .../internal/controller/snapshotjob_job.go | 3 + .../controller/snapshotjob_job_test.go | 23 +++++ .../controller/snapshotjob_podsnapshot.go | 33 +++++-- .../snapshotjob_podsnapshot_test.go | 75 ++++++++++++++++ .../controller/snapshotjob_validation.go | 31 +++++++ 11 files changed, 369 insertions(+), 38 deletions(-) diff --git a/api/v1alpha1/crds/embed_test.go b/api/v1alpha1/crds/embed_test.go index 92466c54..bf30b7b8 100644 --- a/api/v1alpha1/crds/embed_test.go +++ b/api/v1alpha1/crds/embed_test.go @@ -105,24 +105,8 @@ func TestAllReturnsACopy(t *testing.T) { } func TestSnapshotJobConditionsUseMapListSchema(t *testing.T) { - manifestJSON, err := utilyaml.ToJSON([]byte(SnapshotJobCRD())) - if err != nil { - t.Fatalf("convert SnapshotJob CRD to JSON: %v", err) - } - var crd map[string]any - if err := json.Unmarshal(manifestJSON, &crd); err != nil { - t.Fatalf("decode SnapshotJob CRD: %v", err) - } - - versions := nestedSlice(t, crd, "spec", "versions") - if len(versions) == 0 { - t.Fatal("SnapshotJob CRD has no versions") - } - version, ok := versions[0].(map[string]any) - if !ok { - t.Fatalf("SnapshotJob CRD version has type %T, want object", versions[0]) - } - conditions := nestedMap(t, version, "schema", "openAPIV3Schema", "properties", "status", "properties", "conditions") + schema := snapshotJobOpenAPISchema(t) + conditions := nestedMap(t, schema, "properties", "status", "properties", "conditions") if got := conditions["x-kubernetes-list-type"]; got != "map" { t.Errorf("conditions x-kubernetes-list-type = %v, want map", got) @@ -134,6 +118,57 @@ func TestSnapshotJobConditionsUseMapListSchema(t *testing.T) { } func TestSnapshotJobRejectsMissingSpecAndLongNamesAtAdmission(t *testing.T) { + schema := snapshotJobOpenAPISchema(t) + + required := nestedSlice(t, schema, "required") + if !slices.Contains(required, any("spec")) { + t.Errorf("SnapshotJob required fields = %v, want spec", required) + } + + const nameRule = "size(self.metadata.name) <= 63" + validations := nestedSlice(t, schema, "x-kubernetes-validations") + if !containsValidationRule(validations, nameRule) { + t.Errorf("SnapshotJob root validations = %v, want rule %q", validations, nameRule) + } +} + +func TestSnapshotJobIncludesTemplateMetadataAdmissionRules(t *testing.T) { + schema := snapshotJobOpenAPISchema(t) + template := nestedMap(t, schema, + "properties", "spec", "properties", "podSnapshotTemplate") + metadata := nestedMap(t, template, "properties", "metadata") + for _, field := range []string{"labels", "annotations"} { + property := nestedMap(t, metadata, "properties", field) + if got := property["maxProperties"]; got != float64(4096) { + t.Errorf("PodSnapshot template metadata %s maxProperties = %v, want 4096", field, got) + } + } + validations := nestedSlice(t, metadata, "x-kubernetes-validations") + wantRules := []string{ + "!has(self.labels) || self.labels.all(k, !format.qualifiedName().validate(k).hasValue())", + "!has(self.labels) || self.labels.all(k, size(self.labels[k]) <= 63)", + "!has(self.annotations) || self.annotations.all(k, !format.qualifiedName().validate(k).hasValue())", + } + for _, rule := range wantRules { + if !containsValidationRule(validations, rule) { + t.Errorf("PodSnapshot template metadata validations = %v, want rule %q", validations, rule) + } + } + + templateValidations := nestedSlice(t, template, "x-kubernetes-validations") + wantTemplateRules := []string{ + "!has(self.metadata) || !has(self.metadata.labels) || !('nvidia.com/snapshot-job' in self.metadata.labels)", + "!has(self.metadata) || !has(self.metadata.labels) || !('nvidia.com/snapshot-job-uid' in self.metadata.labels)", + } + for _, rule := range wantTemplateRules { + if !containsValidationRule(templateValidations, rule) { + t.Errorf("PodSnapshot template validations = %v, want rule %q", templateValidations, rule) + } + } +} + +func snapshotJobOpenAPISchema(t *testing.T) map[string]any { + t.Helper() manifestJSON, err := utilyaml.ToJSON([]byte(SnapshotJobCRD())) if err != nil { t.Fatalf("convert SnapshotJob CRD to JSON: %v", err) @@ -151,22 +186,17 @@ func TestSnapshotJobRejectsMissingSpecAndLongNamesAtAdmission(t *testing.T) { if !ok { t.Fatalf("SnapshotJob CRD version has type %T, want object", versions[0]) } - schema := nestedMap(t, version, "schema", "openAPIV3Schema") - - required := nestedSlice(t, schema, "required") - if !slices.Contains(required, any("spec")) { - t.Errorf("SnapshotJob required fields = %v, want spec", required) - } + return nestedMap(t, version, "schema", "openAPIV3Schema") +} - const nameRule = "size(self.metadata.name) <= 63" - validations := nestedSlice(t, schema, "x-kubernetes-validations") +func containsValidationRule(validations []any, rule string) bool { for _, validation := range validations { entry, ok := validation.(map[string]any) - if ok && entry["rule"] == nameRule { - return + if ok && entry["rule"] == rule { + return true } } - t.Errorf("SnapshotJob root validations = %v, want rule %q", validations, nameRule) + return false } func nestedMap(t *testing.T, object map[string]any, fields ...string) map[string]any { diff --git a/api/v1alpha1/crds/nvidia.com_snapshotjobs.yaml b/api/v1alpha1/crds/nvidia.com_snapshotjobs.yaml index 5bec2335..46b08405 100644 --- a/api/v1alpha1/crds/nvidia.com_snapshotjobs.yaml +++ b/api/v1alpha1/crds/nvidia.com_snapshotjobs.yaml @@ -77,6 +77,35 @@ spec: PodSnapshotTemplate defines the properties of the PodSnapshot produced by this job. The controller fills in spec.source from the pod it creates. properties: + metadata: + description: |- + Metadata contains labels and annotations copied to the generated + PodSnapshot. SnapshotJob ownership labels are reserved for the controller + and cannot be supplied here. + properties: + annotations: + additionalProperties: + type: string + description: Annotations are copied to the generated PodSnapshot. + maxProperties: 4096 + type: object + labels: + additionalProperties: + type: string + description: Labels are copied to the generated PodSnapshot. + maxProperties: 4096 + type: object + type: object + x-kubernetes-validations: + - message: metadata.labels keys must be valid Kubernetes qualified + names + rule: '!has(self.labels) || self.labels.all(k, !format.qualifiedName().validate(k).hasValue())' + - message: metadata.labels values must be no more than 63 characters + rule: '!has(self.labels) || self.labels.all(k, size(self.labels[k]) + <= 63)' + - message: metadata.annotations keys must be valid Kubernetes + qualified names + rule: '!has(self.annotations) || self.annotations.all(k, !format.qualifiedName().validate(k).hasValue())' targetContainers: default: - main @@ -100,6 +129,13 @@ spec: minItems: 1 type: array type: object + x-kubernetes-validations: + - message: metadata.labels must not set controller-owned label nvidia.com/snapshot-job + rule: '!has(self.metadata) || !has(self.metadata.labels) || !(''nvidia.com/snapshot-job'' + in self.metadata.labels)' + - message: metadata.labels must not set controller-owned label nvidia.com/snapshot-job-uid + rule: '!has(self.metadata) || !has(self.metadata.labels) || !(''nvidia.com/snapshot-job-uid'' + in self.metadata.labels)' podTemplate: description: |- PodTemplate defines the workload to run and capture. The controller injects diff --git a/api/v1alpha1/snapshotjob_types.go b/api/v1alpha1/snapshotjob_types.go index 9b21b157..b60c6a9e 100644 --- a/api/v1alpha1/snapshotjob_types.go +++ b/api/v1alpha1/snapshotjob_types.go @@ -134,9 +134,17 @@ type SnapshotJobSpec struct { PodSnapshotTemplate PodSnapshotTemplate `json:"podSnapshotTemplate"` } -// PodSnapshotTemplate mirrors the PodSnapshot spec fields the user controls. The -// controller fills in spec.source (the pod reference) automatically. +// PodSnapshotTemplate mirrors the PodSnapshot fields the user controls. The +// controller fills in identity, ownership, and spec.source automatically. +// +kubebuilder:validation:XValidation:rule="!has(self.metadata) || !has(self.metadata.labels) || !('nvidia.com/snapshot-job' in self.metadata.labels)",message="metadata.labels must not set controller-owned label nvidia.com/snapshot-job" +// +kubebuilder:validation:XValidation:rule="!has(self.metadata) || !has(self.metadata.labels) || !('nvidia.com/snapshot-job-uid' in self.metadata.labels)",message="metadata.labels must not set controller-owned label nvidia.com/snapshot-job-uid" type PodSnapshotTemplate struct { + // Metadata contains labels and annotations copied to the generated + // PodSnapshot. SnapshotJob ownership labels are reserved for the controller + // and cannot be supplied here. + // +optional + Metadata *PodSnapshotTemplateMetadata `json:"metadata,omitempty"` + // TargetContainers names the container(s) to checkpoint with CRIU. The pod // may contain any number of additional containers (helpers, sidecars, etc.) // — this field controls only the CRIU dump target. @@ -157,6 +165,26 @@ type PodSnapshotTemplate struct { TargetContainers []string `json:"targetContainers,omitempty"` } +// PodSnapshotTemplateMetadata is the caller-owned metadata propagated to a +// SnapshotJob's generated PodSnapshot. It intentionally exposes only labels and +// annotations; identity and ownership remain controller-owned. The SnapshotJob +// spec is immutable, so metadata rejected by the controller as InvalidSpec +// requires deleting and recreating the SnapshotJob. +// +kubebuilder:validation:XValidation:rule="!has(self.labels) || self.labels.all(k, !format.qualifiedName().validate(k).hasValue())",message="metadata.labels keys must be valid Kubernetes qualified names" +// +kubebuilder:validation:XValidation:rule="!has(self.labels) || self.labels.all(k, size(self.labels[k]) <= 63)",message="metadata.labels values must be no more than 63 characters" +// +kubebuilder:validation:XValidation:rule="!has(self.annotations) || self.annotations.all(k, !format.qualifiedName().validate(k).hasValue())",message="metadata.annotations keys must be valid Kubernetes qualified names" +type PodSnapshotTemplateMetadata struct { + // Labels are copied to the generated PodSnapshot. + // +optional + // +kubebuilder:validation:MaxProperties=4096 + Labels map[string]string `json:"labels,omitempty"` + + // Annotations are copied to the generated PodSnapshot. + // +optional + // +kubebuilder:validation:MaxProperties=4096 + Annotations map[string]string `json:"annotations,omitempty"` +} + // SnapshotJobStatus defines the observed state of SnapshotJob. type SnapshotJobStatus struct { // SourceJobUID identifies the one source batch/v1 Job incarnation accepted diff --git a/api/v1alpha1/snapshotjob_types_test.go b/api/v1alpha1/snapshotjob_types_test.go index 09ed2d1f..4c775edb 100644 --- a/api/v1alpha1/snapshotjob_types_test.go +++ b/api/v1alpha1/snapshotjob_types_test.go @@ -39,7 +39,13 @@ func TestSnapshotJobDeepCopyIsIndependent(t *testing.T) { Containers: []corev1.Container{{Name: "worker"}}, }, }, - PodSnapshotTemplate: PodSnapshotTemplate{TargetContainers: []string{"worker"}}, + PodSnapshotTemplate: PodSnapshotTemplate{ + Metadata: &PodSnapshotTemplateMetadata{ + Labels: map[string]string{"dynamo.nvidia.com/worker-generation": "abc123"}, + Annotations: map[string]string{"dynamo.nvidia.com/gms-mode": "enabled"}, + }, + TargetContainers: []string{"worker"}, + }, }, Status: SnapshotJobStatus{ PodSnapshotName: "warm-worker-snapshot", @@ -56,6 +62,8 @@ func TestSnapshotJobDeepCopyIsIndependent(t *testing.T) { clone.Spec.PodTemplate.Spec.Containers[0].Name = "mutated" clone.Spec.PodSnapshotTemplate.TargetContainers[0] = "mutated" + clone.Spec.PodSnapshotTemplate.Metadata.Labels["dynamo.nvidia.com/worker-generation"] = "changed" + clone.Spec.PodSnapshotTemplate.Metadata.Annotations["dynamo.nvidia.com/gms-mode"] = "disabled" clone.Status.Conditions[0].Reason = "Changed" if original.Spec.PodTemplate.Spec.Containers[0].Name != "worker" { t.Errorf("mutating clone PodTemplate changed original: got %q", original.Spec.PodTemplate.Spec.Containers[0].Name) @@ -63,6 +71,12 @@ func TestSnapshotJobDeepCopyIsIndependent(t *testing.T) { if original.Spec.PodSnapshotTemplate.TargetContainers[0] != "worker" { t.Errorf("mutating clone TargetContainers changed original: got %q", original.Spec.PodSnapshotTemplate.TargetContainers[0]) } + if original.Spec.PodSnapshotTemplate.Metadata.Labels["dynamo.nvidia.com/worker-generation"] != "abc123" { + t.Errorf("mutating clone labels changed original: got %q", original.Spec.PodSnapshotTemplate.Metadata.Labels["dynamo.nvidia.com/worker-generation"]) + } + if original.Spec.PodSnapshotTemplate.Metadata.Annotations["dynamo.nvidia.com/gms-mode"] != "enabled" { + t.Errorf("mutating clone annotations changed original: got %q", original.Spec.PodSnapshotTemplate.Metadata.Annotations["dynamo.nvidia.com/gms-mode"]) + } if original.Status.Conditions[0].Reason != ReasonPodReady { t.Errorf("mutating clone condition changed original: got %q", original.Status.Conditions[0].Reason) } diff --git a/api/v1alpha1/zz_generated.deepcopy.go b/api/v1alpha1/zz_generated.deepcopy.go index 7f89aca2..3e5592ab 100644 --- a/api/v1alpha1/zz_generated.deepcopy.go +++ b/api/v1alpha1/zz_generated.deepcopy.go @@ -284,6 +284,11 @@ func (in *PodSnapshotStatus) DeepCopy() *PodSnapshotStatus { // DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil. func (in *PodSnapshotTemplate) DeepCopyInto(out *PodSnapshotTemplate) { *out = *in + if in.Metadata != nil { + in, out := &in.Metadata, &out.Metadata + *out = new(PodSnapshotTemplateMetadata) + (*in).DeepCopyInto(*out) + } if in.TargetContainers != nil { in, out := &in.TargetContainers, &out.TargetContainers *out = make([]string, len(*in)) @@ -301,6 +306,35 @@ func (in *PodSnapshotTemplate) DeepCopy() *PodSnapshotTemplate { return out } +// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil. +func (in *PodSnapshotTemplateMetadata) DeepCopyInto(out *PodSnapshotTemplateMetadata) { + *out = *in + if in.Labels != nil { + in, out := &in.Labels, &out.Labels + *out = make(map[string]string, len(*in)) + for key, val := range *in { + (*out)[key] = val + } + } + if in.Annotations != nil { + in, out := &in.Annotations, &out.Annotations + *out = make(map[string]string, len(*in)) + for key, val := range *in { + (*out)[key] = val + } + } +} + +// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new PodSnapshotTemplateMetadata. +func (in *PodSnapshotTemplateMetadata) DeepCopy() *PodSnapshotTemplateMetadata { + if in == nil { + return nil + } + out := new(PodSnapshotTemplateMetadata) + in.DeepCopyInto(out) + return out +} + // DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil. func (in *SnapshotJob) DeepCopyInto(out *SnapshotJob) { *out = *in diff --git a/charts/snapshot/crds/nvidia.com_snapshotjobs.yaml b/charts/snapshot/crds/nvidia.com_snapshotjobs.yaml index 5bec2335..46b08405 100644 --- a/charts/snapshot/crds/nvidia.com_snapshotjobs.yaml +++ b/charts/snapshot/crds/nvidia.com_snapshotjobs.yaml @@ -77,6 +77,35 @@ spec: PodSnapshotTemplate defines the properties of the PodSnapshot produced by this job. The controller fills in spec.source from the pod it creates. properties: + metadata: + description: |- + Metadata contains labels and annotations copied to the generated + PodSnapshot. SnapshotJob ownership labels are reserved for the controller + and cannot be supplied here. + properties: + annotations: + additionalProperties: + type: string + description: Annotations are copied to the generated PodSnapshot. + maxProperties: 4096 + type: object + labels: + additionalProperties: + type: string + description: Labels are copied to the generated PodSnapshot. + maxProperties: 4096 + type: object + type: object + x-kubernetes-validations: + - message: metadata.labels keys must be valid Kubernetes qualified + names + rule: '!has(self.labels) || self.labels.all(k, !format.qualifiedName().validate(k).hasValue())' + - message: metadata.labels values must be no more than 63 characters + rule: '!has(self.labels) || self.labels.all(k, size(self.labels[k]) + <= 63)' + - message: metadata.annotations keys must be valid Kubernetes + qualified names + rule: '!has(self.annotations) || self.annotations.all(k, !format.qualifiedName().validate(k).hasValue())' targetContainers: default: - main @@ -100,6 +129,13 @@ spec: minItems: 1 type: array type: object + x-kubernetes-validations: + - message: metadata.labels must not set controller-owned label nvidia.com/snapshot-job + rule: '!has(self.metadata) || !has(self.metadata.labels) || !(''nvidia.com/snapshot-job'' + in self.metadata.labels)' + - message: metadata.labels must not set controller-owned label nvidia.com/snapshot-job-uid + rule: '!has(self.metadata) || !has(self.metadata.labels) || !(''nvidia.com/snapshot-job-uid'' + in self.metadata.labels)' podTemplate: description: |- PodTemplate defines the workload to run and capture. The controller injects diff --git a/operator/internal/controller/snapshotjob_job.go b/operator/internal/controller/snapshotjob_job.go index df2e58c7..b42d3c6d 100644 --- a/operator/internal/controller/snapshotjob_job.go +++ b/operator/internal/controller/snapshotjob_job.go @@ -34,6 +34,9 @@ func buildSourceJob(sj *snapshotv1alpha1.SnapshotJob) (*batchv1.Job, error) { if errs := contentvalidation.IsLabelValue(sj.Name); len(errs) > 0 { return nil, fmt.Errorf("metadata.name %q is not a valid label value: %s", sj.Name, strings.Join(errs, "; ")) } + if err := validatePodSnapshotTemplateMetadata(sj); err != nil { + return nil, err + } targetContainer, err := snapshotJobTargetContainer(sj) if err != nil { diff --git a/operator/internal/controller/snapshotjob_job_test.go b/operator/internal/controller/snapshotjob_job_test.go index a06e867b..aefe1be9 100644 --- a/operator/internal/controller/snapshotjob_job_test.go +++ b/operator/internal/controller/snapshotjob_job_test.go @@ -162,6 +162,29 @@ func TestBuildSourceJob(t *testing.T) { require.Error(t, err) }) + t.Run("invalid PodSnapshot metadata is a terminal spec error", func(t *testing.T) { + tests := map[string]*snapshotv1alpha1.PodSnapshotTemplateMetadata{ + "invalid label": { + Labels: map[string]string{"example.com/team": strings.Repeat("x", 64)}, + }, + "invalid annotation": { + Annotations: map[string]string{"not a qualified annotation key": "value"}, + }, + "reserved owner label": { + Labels: map[string]string{snapshotv1alpha1.SnapshotJobOwnerLabel: "caller"}, + }, + } + for name, metadata := range tests { + t.Run(name, func(t *testing.T) { + sj := minimalSnapshotJob() + sj.Spec.PodSnapshotTemplate.Metadata = metadata + + _, err := buildSourceJob(sj) + require.Error(t, err) + }) + } + }) + t.Run("more than one targetContainers entry is a terminal spec error", func(t *testing.T) { // The CRD caps this at MaxItems=1, but this is defense in depth for an // object that bypassed CEL validation — v1alpha1 supports exactly one diff --git a/operator/internal/controller/snapshotjob_podsnapshot.go b/operator/internal/controller/snapshotjob_podsnapshot.go index 0128c4b1..64b43eb3 100644 --- a/operator/internal/controller/snapshotjob_podsnapshot.go +++ b/operator/internal/controller/snapshotjob_podsnapshot.go @@ -7,6 +7,7 @@ import ( "context" "errors" "fmt" + "maps" "slices" "time" @@ -190,18 +191,20 @@ func buildPodSnapshot(sj *snapshotv1alpha1.SnapshotJob, pod *corev1.Pod) (*snaps if _, err := snapshotJobTargetContainer(sj); err != nil { return nil, err } + labels, annotations, err := podSnapshotTemplateMetadata(sj) + if err != nil { + return nil, err + } return &snapshotv1alpha1.PodSnapshot{ TypeMeta: metav1.TypeMeta{ APIVersion: snapshotv1alpha1.GroupVersion.String(), Kind: "PodSnapshot", }, ObjectMeta: metav1.ObjectMeta{ - Name: sj.Name, - Namespace: sj.Namespace, - Labels: map[string]string{ - snapshotv1alpha1.SnapshotJobOwnerLabel: sj.Name, - snapshotv1alpha1.SnapshotJobOwnerUIDLabel: string(sj.UID), - }, + Name: sj.Name, + Namespace: sj.Namespace, + Labels: labels, + Annotations: annotations, }, Spec: snapshotv1alpha1.PodSnapshotSpec{ Source: snapshotv1alpha1.PodSnapshotSource{ @@ -215,6 +218,24 @@ func buildPodSnapshot(sj *snapshotv1alpha1.SnapshotJob, pod *corev1.Pod) (*snaps }, nil } +func podSnapshotTemplateMetadata(sj *snapshotv1alpha1.SnapshotJob) (map[string]string, map[string]string, error) { + if err := validatePodSnapshotTemplateMetadata(sj); err != nil { + return nil, nil, err + } + metadata := sj.Spec.PodSnapshotTemplate.Metadata + var labels, annotations map[string]string + if metadata != nil { + labels = maps.Clone(metadata.Labels) + annotations = maps.Clone(metadata.Annotations) + } + if labels == nil { + labels = make(map[string]string, 2) + } + labels[snapshotv1alpha1.SnapshotJobOwnerLabel] = sj.Name + labels[snapshotv1alpha1.SnapshotJobOwnerUIDLabel] = string(sj.UID) + return labels, annotations, nil +} + // createPodSnapshot creates this SnapshotJob's PodSnapshot. The caller has // confirmed via findOwnedPodSnapshot that none exists, so this is a pure create. // On AlreadyExists the object at the deterministic name is classified: cache lag diff --git a/operator/internal/controller/snapshotjob_podsnapshot_test.go b/operator/internal/controller/snapshotjob_podsnapshot_test.go index 8c3469d0..d800e605 100644 --- a/operator/internal/controller/snapshotjob_podsnapshot_test.go +++ b/operator/internal/controller/snapshotjob_podsnapshot_test.go @@ -30,6 +30,10 @@ import ( func TestBuildPodSnapshot(t *testing.T) { sj := minimalSnapshotJob() sj.UID = types.UID("sj-uid") + sj.Spec.PodSnapshotTemplate.Metadata = &snapshotv1alpha1.PodSnapshotTemplateMetadata{ + Labels: map[string]string{"dynamo.nvidia.com/worker-generation": "abc123"}, + Annotations: map[string]string{"dynamo.nvidia.com/gms-mode": "enabled"}, + } pod := &corev1.Pod{ ObjectMeta: metav1.ObjectMeta{Name: "warm-worker-abcde", Namespace: "inference", UID: types.UID("pod-uid")}, } @@ -47,6 +51,11 @@ func TestBuildPodSnapshot(t *testing.T) { assert.Equal(t, string(sj.UID), snap.Labels[snapshotv1alpha1.SnapshotJobOwnerUIDLabel]) }) + t.Run("propagates caller metadata", func(t *testing.T) { + assert.Equal(t, "abc123", snap.Labels["dynamo.nvidia.com/worker-generation"]) + assert.Equal(t, "enabled", snap.Annotations["dynamo.nvidia.com/gms-mode"]) + }) + t.Run("pins the source pod name and UID", func(t *testing.T) { assert.Equal(t, pod.Name, snap.Spec.Source.PodRef.Name) assert.Equal(t, pod.UID, snap.Spec.Source.PodRef.UID) @@ -90,6 +99,30 @@ func TestBuildPodSnapshot(t *testing.T) { assert.Equal(t, original, src.Spec.PodSnapshotTemplate.TargetContainers, "mutating the PodSnapshot's copy must not affect the SnapshotJob's own spec slice") }) + + t.Run("does not share metadata maps with the SnapshotJob", func(t *testing.T) { + got, err := buildPodSnapshot(sj, pod) + require.NoError(t, err) + got.Labels["dynamo.nvidia.com/worker-generation"] = "changed" + got.Annotations["dynamo.nvidia.com/gms-mode"] = "disabled" + + assert.Equal(t, "abc123", sj.Spec.PodSnapshotTemplate.Metadata.Labels["dynamo.nvidia.com/worker-generation"]) + assert.Equal(t, "enabled", sj.Spec.PodSnapshotTemplate.Metadata.Annotations["dynamo.nvidia.com/gms-mode"]) + }) + + for _, reserved := range []string{ + snapshotv1alpha1.SnapshotJobOwnerLabel, + snapshotv1alpha1.SnapshotJobOwnerUIDLabel, + } { + t.Run("rejects reserved label "+reserved, func(t *testing.T) { + bad := minimalSnapshotJob() + bad.Spec.PodSnapshotTemplate.Metadata = &snapshotv1alpha1.PodSnapshotTemplateMetadata{ + Labels: map[string]string{reserved: "caller-value"}, + } + _, err := buildPodSnapshot(bad, pod) + require.ErrorContains(t, err, "controller-owned") + }) + } } // ---- reconciler-level PodSnapshot creation ---- @@ -373,6 +406,25 @@ func TestClassifyExistingPodSnapshot(t *testing.T) { assert.ErrorIs(t, err, errPodSnapshotNameConflict) }) + t.Run("caller metadata drift: adopted", func(t *testing.T) { + withMetadata := sj.DeepCopy() + withMetadata.Spec.PodSnapshotTemplate.Metadata = &snapshotv1alpha1.PodSnapshotTemplateMetadata{ + Labels: map[string]string{"dynamo.nvidia.com/worker-generation": "abc123"}, + Annotations: map[string]string{"dynamo.nvidia.com/gms-mode": "enabled", "empty-value": ""}, + } + desired, err := buildPodSnapshot(withMetadata, pod) + require.NoError(t, err) + existing := desired.DeepCopy() + existing.Labels["dynamo.nvidia.com/worker-generation"] = "mutated" + delete(existing.Annotations, "dynamo.nvidia.com/gms-mode") + delete(existing.Annotations, "empty-value") + r := makeSnapshotJobReconciler(s, existing) + + got, err := r.classifyExistingPodSnapshot(context.Background(), withMetadata, desired, errors.New("AlreadyExists")) + require.NoError(t, err) + assert.Equal(t, existing.Name, got.Name) + }) + t.Run("lookup rejects a stale SnapshotJob UID", func(t *testing.T) { stale, err := buildPodSnapshot(sj, pod) require.NoError(t, err) @@ -614,6 +666,29 @@ func TestSnapshotJobReconcileRejectsPodSnapshotSpecDriftBeforeUIDBinding(t *test } } +func TestValidatePodSnapshotForAdoptionIgnoresMutableMetadata(t *testing.T) { + s := snapshotJobReconcilerScheme() + sj := minimalSnapshotJob() + sj.UID = types.UID("sj-uid") + sj.Spec.PodSnapshotTemplate.Metadata = &snapshotv1alpha1.PodSnapshotTemplateMetadata{ + Labels: map[string]string{"dynamo.nvidia.com/worker-generation": "abc123"}, + Annotations: map[string]string{"dynamo.nvidia.com/gms-mode": "enabled"}, + } + job, err := buildSourceJob(sj) + require.NoError(t, err) + job.UID = types.UID("job-uid") + pod := sourcePodForJob(job) + snap, err := buildPodSnapshot(sj, pod) + require.NoError(t, err) + snap.Labels["dynamo.nvidia.com/worker-generation"] = "mutated" + delete(snap.Annotations, "dynamo.nvidia.com/gms-mode") + r := makeSnapshotJobReconciler(s, pod) + + failure, err := r.validatePodSnapshotForAdoption(context.Background(), sj, job, snap) + require.NoError(t, err) + assert.Nil(t, failure) +} + // ---- findSourcePod ---- func TestFindSourcePod(t *testing.T) { diff --git a/operator/internal/controller/snapshotjob_validation.go b/operator/internal/controller/snapshotjob_validation.go index 7010b906..af55eaeb 100644 --- a/operator/internal/controller/snapshotjob_validation.go +++ b/operator/internal/controller/snapshotjob_validation.go @@ -6,6 +6,10 @@ package controller import ( "fmt" + apivalidation "k8s.io/apimachinery/pkg/api/validation" + metav1validation "k8s.io/apimachinery/pkg/apis/meta/v1/validation" + "k8s.io/apimachinery/pkg/util/validation/field" + snapshotv1alpha1 "github.com/ai-dynamo/snapshot/api/v1alpha1" ) @@ -19,3 +23,30 @@ func snapshotJobTargetContainer(sj *snapshotv1alpha1.SnapshotJob) (string, error } return targets[0], nil } + +// validatePodSnapshotTemplateMetadata rejects invalid Kubernetes metadata and +// labels reserved for the SnapshotJob controller before the source workload is +// started. The generated PodSnapshot remains the final API-server validation +// boundary; this check turns malformed immutable input into InvalidSpec early. +func validatePodSnapshotTemplateMetadata(sj *snapshotv1alpha1.SnapshotJob) error { + metadata := sj.Spec.PodSnapshotTemplate.Metadata + if metadata == nil { + return nil + } + labelsPath := field.NewPath("spec", "podSnapshotTemplate", "metadata", "labels") + errs := metav1validation.ValidateLabels(metadata.Labels, labelsPath) + annotationErrs := apivalidation.ValidateAnnotations( + metadata.Annotations, + field.NewPath("spec", "podSnapshotTemplate", "metadata", "annotations"), + ) + errs = append(errs, annotationErrs...) + for _, reserved := range []string{ + snapshotv1alpha1.SnapshotJobOwnerLabel, + snapshotv1alpha1.SnapshotJobOwnerUIDLabel, + } { + if _, found := metadata.Labels[reserved]; found { + errs = append(errs, field.Forbidden(labelsPath.Key(reserved), "label is controller-owned")) + } + } + return errs.ToAggregate() +}