hanzo-ml/trainer
1
0
Fork 0
mirror of https://github.com/hanzo-ml/trainer.git synced 2026-09-24 00:31:34 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat(api): Set RuntimePatch.Time field automatically during admission (#3319)

Browse Source 
* feat(api): Set RuntimePatch.Time field automatically during admission

Signed-off-by: Antonin Stefanutti <[email protected]>

* Address @tenzen-y comments

Signed-off-by: Antonin Stefanutti <[email protected]>

* fix: Replace RuntimePatches time on UPDATE call

Signed-off-by: Andrey Velichkevich <[email protected]>

---------

Signed-off-by: Antonin Stefanutti <[email protected]>
Signed-off-by: Andrey Velichkevich <[email protected]>
Co-authored-by: Andrey Velichkevich <[email protected]>
This commit is contained in:
Antonin Stefanutti
2026-03-18 11:53:24 +00:00
committed by GitHub
co-authored by Andrey Velichkevich
parent 64a01ebbba
commit dc67f4e356
16 changed files with 645 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files
charts/kubeflow-trainer/templates/rbac/clusterrole.yaml
+1
View File
@@ -54,6 +54,7 @@ rules:
- apiGroups:
- admissionregistration.k8s.io
resources:
- mutatingwebhookconfigurations
- validatingwebhookconfigurations
verbs:
- get
charts/kubeflow-trainer/templates/webhook/_helpers.tpl
+7
View File
@@ -44,6 +44,13 @@ Create the name of the webhook service.
{{ include "trainer.manager.service.name" . }}
{{- end -}}
{{/*
Create the name of the mutating webhook configuration.
*/}}
{{- define "trainer.webhook.mutatingWebhookConfiguration.name" -}}
defaulter.trainer.kubeflow.org
{{- end -}}
{{/*
Create the name of the validating webhook configuration.
*/}}
charts/kubeflow-trainer/templates/webhook/mutating_webhook_configuration.yaml
+45
View File
@@ -0,0 +1,45 @@
{{- /*
Copyright 2025 The Kubeflow authors.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
https://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/ -}}
apiVersion: admissionregistration.k8s.io/v1
kind: MutatingWebhookConfiguration
metadata:
name: {{ include "trainer.webhook.mutatingWebhookConfiguration.name" . }}
labels:
{{- include "trainer.webhook.labels" . | nindent 4 }}
webhooks:
- name: defaulter.trainjob.trainer.kubeflow.org
admissionReviewVersions:
- v1
clientConfig:
service:
name: {{ include "trainer.webhook.service.name" . }}
namespace: {{ .Release.Namespace }}
path: /mutate-trainer-kubeflow-org-v1alpha1-trainjob
sideEffects: None
{{- with .Values.webhook.failurePolicy }}
failurePolicy: {{ . }}
{{- end }}
rules:
- apiGroups:
- trainer.kubeflow.org
apiVersions:
- v1alpha1
resources:
- trainjobs
operations:
- CREATE
- UPDATE
charts/kubeflow-trainer/tests/webhook/mutating_webhook_configuration_test.yaml
+41
View File
@@ -0,0 +1,41 @@
#
# Copyright 2025 The Kubeflow authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
suite: Test MutatingWebhookConfiguration
templates:
- webhook/mutating_webhook_configuration.yaml
release:
name: kubeflow-trainer
namespace: kubeflow-system
tests:
- it: Should create MutatingWebhookConfiguration
asserts:
- containsDocument:
apiVersion: admissionregistration.k8s.io/v1
kind: MutatingWebhookConfiguration
name: defaulter.trainer.kubeflow.org
- it: Should use the specified failurePolicy if `webhook.failurePolicy` is set
set:
webhook:
failurePolicy: Ignore
asserts:
- equal:
path: webhooks[*].failurePolicy
value: Ignore
cmd/trainer-controller-manager/main.go
+6 -4
View File
@@ -49,7 +49,8 @@ import (
)
const (
webhookConfigurationName = "validator.trainer.kubeflow.org"
validatingWebhookConfigurationName = "validator.trainer.kubeflow.org"
mutatingWebhookConfigurationName = "defaulter.trainer.kubeflow.org"
)
var (
@@ -128,9 +129,10 @@ func main() {
if config.IsCertManagementEnabled(&cfg) {
setupLog.Info("Setting up certificate management")
if err = cert.ManageCerts(mgr, cert.Config{
WebhookSecretName: cfg.CertManagement.WebhookSecretName,
WebhookServiceName: cfg.CertManagement.WebhookServiceName,
WebhookConfigurationName: webhookConfigurationName,
WebhookSecretName: cfg.CertManagement.WebhookSecretName,
WebhookServiceName: cfg.CertManagement.WebhookServiceName,
ValidatingWebhookConfigurationName: validatingWebhookConfigurationName,
MutatingWebhookConfigurationName: mutatingWebhookConfigurationName,
}, certsReady); err != nil {
setupLog.Error(err, "unable to set up cert rotation")
os.Exit(1)
manifests/base/rbac/role.yaml
+1
View File
@@ -37,6 +37,7 @@ rules:
- apiGroups:
- admissionregistration.k8s.io
resources:
- mutatingwebhookconfigurations
- validatingwebhookconfigurations
verbs:
- get
manifests/base/webhook/kustomization.yaml
+6 -1
View File
@@ -3,7 +3,12 @@ kind: Kustomization
resources:
- manifests.yaml
patches:
- path: patch.yaml
- path: patch_mutating.yaml
target:
group: admissionregistration.k8s.io
version: v1
kind: MutatingWebhookConfiguration
- path: patch_validating.yaml
target:
group: admissionregistration.k8s.io
version: v1
manifests/base/webhook/kustomizeconfig.yaml
+5
View File
@@ -1,6 +1,11 @@
# the following config is for teaching kustomize where to look at when substituting vars.
# It requires kustomize v2.1.0 or newer to work properly.
namespace:
- kind: MutatingWebhookConfiguration
group: admissionregistration.k8s.io
path: webhooks/clientConfig/service/namespace
create: true
- kind: ValidatingWebhookConfiguration
group: admissionregistration.k8s.io
path: webhooks/clientConfig/service/namespace
manifests/base/webhook/manifests.yaml
+26
View File
@@ -1,5 +1,31 @@
---
apiVersion: admissionregistration.k8s.io/v1
kind: MutatingWebhookConfiguration
metadata:
name: mutating-webhook-configuration
webhooks:
- admissionReviewVersions:
- v1
clientConfig:
service:
name: webhook-service
namespace: system
path: /mutate-trainer-kubeflow-org-v1alpha1-trainjob
failurePolicy: Fail
name: defaulter.trainjob.trainer.kubeflow.org
rules:
- apiGroups:
- trainer.kubeflow.org
apiVersions:
- v1alpha1
operations:
- CREATE
- UPDATE
resources:
- trainjobs
sideEffects: None
---
apiVersion: admissionregistration.k8s.io/v1
kind: ValidatingWebhookConfiguration
metadata:
name: validating-webhook-configuration
manifests/base/webhook/patch_mutating.yaml
+6
View File
@@ -0,0 +1,6 @@
- op: replace
path: /webhooks/0/clientConfig/service/name
value: kubeflow-trainer-controller-manager
- op: replace
path: /metadata/name
value: defaulter.trainer.kubeflow.org
manifests/base/webhook/patch.yaml → manifests/base/webhook/patch_validating.yaml
View File
pkg/util/cert/cert.go
+15 -7
View File
@@ -45,13 +45,15 @@ func GetOperatorNamespace() string {
}
type Config struct {
WebhookServiceName string
WebhookSecretName string
WebhookConfigurationName string
WebhookServiceName string
WebhookSecretName string
ValidatingWebhookConfigurationName string
MutatingWebhookConfigurationName string
}
//+kubebuilder:rbac:groups="",resources=secrets,verbs=get;list;watch;update
//+kubebuilder:rbac:groups="admissionregistration.k8s.io",resources=validatingwebhookconfigurations,verbs=get;list;watch;update
//+kubebuilder:rbac:groups="admissionregistration.k8s.io",resources=mutatingwebhookconfigurations,verbs=get;list;watch;update
// ManageCerts creates all certs for webhooks.
func ManageCerts(mgr ctrl.Manager, cfg Config, setupFinished chan struct{}) error {
@@ -70,10 +72,16 @@ func ManageCerts(mgr ctrl.Manager, cfg Config, setupFinished chan struct{}) erro
CAOrganization: caOrganization,
DNSName: dnsName,
IsReady: setupFinished,
Webhooks: []cert.WebhookInfo{{
Type: cert.Validating,
Name: cfg.WebhookConfigurationName,
}},
Webhooks: []cert.WebhookInfo{
{
Type: cert.Validating,
Name: cfg.ValidatingWebhookConfigurationName,
},
{
Type: cert.Mutating,
Name: cfg.MutatingWebhookConfigurationName,
},
},
// When Kubeflow Trainer is running in the leader election mode,
// we expect webhook server will run in primary and secondary instance
RequireLeaderElection: false,
pkg/webhooks/trainjob_webhook.go
+64
View File
@@ -18,9 +18,14 @@ package webhooks
import (
"context"
"encoding/json"
"fmt"
admissionv1 "k8s.io/api/admission/v1"
"k8s.io/apimachinery/pkg/api/equality"
metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
"k8s.io/klog/v2"
"k8s.io/utils/clock"
ctrl "sigs.k8s.io/controller-runtime"
"sigs.k8s.io/controller-runtime/pkg/webhook/admission"
@@ -28,6 +33,64 @@ import (
"github.com/kubeflow/trainer/v2/pkg/runtime"
)
// +kubebuilder:webhook:path=/mutate-trainer-kubeflow-org-v1alpha1-trainjob,mutating=true,failurePolicy=fail,sideEffects=None,groups=trainer.kubeflow.org,resources=trainjobs,verbs=create;update,versions=v1alpha1,name=defaulter.trainjob.trainer.kubeflow.org,admissionReviewVersions=v1
// TrainJobDefaulter defaults TrainJobs.
type TrainJobDefaulter struct {
clock clock.PassiveClock
}
var _ admission.Defaulter[*trainer.TrainJob] = (*TrainJobDefaulter)(nil)
func (d *TrainJobDefaulter) Default(ctx context.Context, trainJob *trainer.TrainJob) error {
log := ctrl.LoggerFrom(ctx).WithName("trainJob-webhook")
log.V(5).Info("Defaulting", "TrainJob", klog.KObj(trainJob))
now := metav1.NewTime(d.clock.Now())
req, err := admission.RequestFromContext(ctx)
if err != nil {
return err
}
var oldObj *trainer.TrainJob
if req.Operation == admissionv1.Update {
oldObj = &trainer.TrainJob{}
if err := json.Unmarshal(req.OldObject.Raw, oldObj); err != nil {
return err
}
}
if oldObj == nil {
for i := range trainJob.Spec.RuntimePatches {
if trainJob.Spec.RuntimePatches[i].Time == nil {
trainJob.Spec.RuntimePatches[i].Time = &now
}
}
return nil
}
oldByManager := make(map[string]trainer.RuntimePatch, len(oldObj.Spec.RuntimePatches))
for _, p := range oldObj.Spec.RuntimePatches {
oldByManager[p.Manager] = p
}
for i := range trainJob.Spec.RuntimePatches {
patch := &trainJob.Spec.RuntimePatches[i]
if old, ok := oldByManager[patch.Manager]; ok {
oldCmp, newCmp := old, *patch
oldCmp.Time, newCmp.Time = nil, nil
if equality.Semantic.DeepEqual(oldCmp, newCmp) {
patch.Time = old.Time
} else {
patch.Time = &now
}
} else if patch.Time == nil {
patch.Time = &now
}
}
return nil
}
// +kubebuilder:webhook:path=/validate-trainer-kubeflow-org-v1alpha1-trainjob,mutating=false,failurePolicy=fail,sideEffects=None,groups=trainer.kubeflow.org,resources=trainjobs,verbs=create;update,versions=v1alpha1,name=validator.trainjob.trainer.kubeflow.org,admissionReviewVersions=v1
// TrainJobValidator validates TrainJobs
@@ -39,6 +102,7 @@ var _ admission.Validator[*trainer.TrainJob] = (*TrainJobValidator)(nil)
func setupWebhookForTrainJob(mgr ctrl.Manager, run map[string]runtime.Runtime) error {
return ctrl.NewWebhookManagedBy(mgr, &trainer.TrainJob{}).
WithDefaulter(&TrainJobDefaulter{clock: clock.RealClock{}}).
WithValidator(&TrainJobValidator{runtimes: run}).
Complete()
}
pkg/webhooks/trainjob_webhook_test.go
+278
View File
@@ -18,12 +18,18 @@ package webhooks
import (
"context"
"encoding/json"
"testing"
"time"
"github.com/google/go-cmp/cmp"
"github.com/google/go-cmp/cmp/cmpopts"
admissionv1 "k8s.io/api/admission/v1"
metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
apiruntime "k8s.io/apimachinery/pkg/runtime"
"k8s.io/apimachinery/pkg/util/validation/field"
"k8s.io/klog/v2/ktesting"
clocktesting "k8s.io/utils/clock/testing"
"k8s.io/utils/ptr"
"sigs.k8s.io/controller-runtime/pkg/webhook/admission"
@@ -33,6 +39,278 @@ import (
testingutil "github.com/kubeflow/trainer/v2/pkg/util/testing"
)
func TestDefault(t *testing.T) {
fakeNow := time.Date(2025, 6, 15, 12, 0, 0, 0, time.UTC)
fakeClock := clocktesting.NewFakeClock(fakeNow)
expectedTime := metav1.NewTime(fakeNow)
oldTime := metav1.NewTime(time.Date(2025, 1, 1, 0, 0, 0, 0, time.UTC))
cases := map[string]struct {
oldObj *trainer.TrainJob
newObj *trainer.TrainJob
wantTime func(patches []trainer.RuntimePatch)
}{
"CREATE: all patches get timestamped": {
newObj: testingutil.MakeTrainJobWrapper("default", "test").
RuntimeRef(trainer.SchemeGroupVersion.WithKind(trainer.ClusterTrainingRuntimeKind), "test-runtime").
RuntimePatches([]trainer.RuntimePatch{
{Manager: "acme.io/one"},
{Manager: "acme.io/two"},
}).Obj(),
wantTime: func(patches []trainer.RuntimePatch) {
for _, p := range patches {
if !p.Time.Equal(&expectedTime) {
t.Errorf("patch %q: expected Time %v, got %v", p.Manager, expectedTime, p.Time)
}
}
},
},
"UPDATE, patch unchanged: Time preserved": {
oldObj: func() *trainer.TrainJob {
obj := testingutil.MakeTrainJobWrapper("default", "test").
RuntimeRef(trainer.SchemeGroupVersion.WithKind(trainer.ClusterTrainingRuntimeKind), "test-runtime").
RuntimePatches([]trainer.RuntimePatch{
{
Manager: "acme.io/one",
Time: &oldTime,
TrainingRuntimeSpec: &trainer.TrainingRuntimeSpecPatch{
Template: &trainer.JobSetTemplatePatch{
Spec: &trainer.JobSetSpecPatch{
ReplicatedJobs: []trainer.ReplicatedJobPatch{{
Name: "node",
Template: &trainer.JobTemplatePatch{
Spec: &trainer.JobSpecPatch{
Template: &trainer.PodTemplatePatch{
Spec: &trainer.PodSpecPatch{
ServiceAccountName: ptr.To("sa"),
},
},
},
},
}},
},
},
},
},
}).Obj()
return obj
}(),
newObj: testingutil.MakeTrainJobWrapper("default", "test").
RuntimeRef(trainer.SchemeGroupVersion.WithKind(trainer.ClusterTrainingRuntimeKind), "test-runtime").
RuntimePatches([]trainer.RuntimePatch{
{
Manager: "acme.io/one",
TrainingRuntimeSpec: &trainer.TrainingRuntimeSpecPatch{
Template: &trainer.JobSetTemplatePatch{
Spec: &trainer.JobSetSpecPatch{
ReplicatedJobs: []trainer.ReplicatedJobPatch{{
Name: "node",
Template: &trainer.JobTemplatePatch{
Spec: &trainer.JobSpecPatch{
Template: &trainer.PodTemplatePatch{
Spec: &trainer.PodSpecPatch{
ServiceAccountName: ptr.To("sa"),
},
},
},
},
}},
},
},
},
},
}).Obj(),
wantTime: func(patches []trainer.RuntimePatch) {
if !patches[0].Time.Equal(&oldTime) {
t.Errorf("expected Time to be preserved as %v, got %v", oldTime, patches[0].Time)
}
},
},
"UPDATE, patch changed: Time updated": {
oldObj: func() *trainer.TrainJob {
obj := testingutil.MakeTrainJobWrapper("default", "test").
RuntimeRef(trainer.SchemeGroupVersion.WithKind(trainer.ClusterTrainingRuntimeKind), "test-runtime").
RuntimePatches([]trainer.RuntimePatch{
{
Manager: "acme.io/one",
Time: &oldTime,
TrainingRuntimeSpec: &trainer.TrainingRuntimeSpecPatch{
Template: &trainer.JobSetTemplatePatch{
Spec: &trainer.JobSetSpecPatch{
ReplicatedJobs: []trainer.ReplicatedJobPatch{{
Name: "node",
Template: &trainer.JobTemplatePatch{
Spec: &trainer.JobSpecPatch{
Template: &trainer.PodTemplatePatch{
Spec: &trainer.PodSpecPatch{
ServiceAccountName: ptr.To("old-sa"),
},
},
},
},
}},
},
},
},
},
}).Obj()
return obj
}(),
newObj: testingutil.MakeTrainJobWrapper("default", "test").
RuntimeRef(trainer.SchemeGroupVersion.WithKind(trainer.ClusterTrainingRuntimeKind), "test-runtime").
RuntimePatches([]trainer.RuntimePatch{
{
Manager: "acme.io/one",
TrainingRuntimeSpec: &trainer.TrainingRuntimeSpecPatch{
Template: &trainer.JobSetTemplatePatch{
Spec: &trainer.JobSetSpecPatch{
ReplicatedJobs: []trainer.ReplicatedJobPatch{{
Name: "node",
Template: &trainer.JobTemplatePatch{
Spec: &trainer.JobSpecPatch{
Template: &trainer.PodTemplatePatch{
Spec: &trainer.PodSpecPatch{
ServiceAccountName: ptr.To("new-sa"),
},
},
},
},
}},
},
},
},
},
}).Obj(),
wantTime: func(patches []trainer.RuntimePatch) {
if !patches[0].Time.Equal(&expectedTime) {
t.Errorf("expected Time to be updated to %v, got %v", expectedTime, patches[0].Time)
}
},
},
"CREATE: externally-set Time is preserved": {
newObj: testingutil.MakeTrainJobWrapper("default", "test").
RuntimeRef(trainer.SchemeGroupVersion.WithKind(trainer.ClusterTrainingRuntimeKind), "test-runtime").
RuntimePatches([]trainer.RuntimePatch{
{Manager: "acme.io/one", Time: &oldTime},
{Manager: "acme.io/two"},
}).Obj(),
wantTime: func(patches []trainer.RuntimePatch) {
if !patches[0].Time.Equal(&oldTime) {
t.Errorf("patch %q: expected externally-set Time %v preserved, got %v", patches[0].Manager, oldTime, patches[0].Time)
}
if !patches[1].Time.Equal(&expectedTime) {
t.Errorf("patch %q: expected Time %v, got %v", patches[1].Manager, expectedTime, patches[1].Time)
}
},
},
"UPDATE, patch changed with pre-existing Time: Time updated": {
oldObj: func() *trainer.TrainJob {
obj := testingutil.MakeTrainJobWrapper("default", "test").
RuntimeRef(trainer.SchemeGroupVersion.WithKind(trainer.ClusterTrainingRuntimeKind), "test-runtime").
RuntimePatches([]trainer.RuntimePatch{
{Manager: "acme.io/one", Time: &oldTime},
}).Obj()
return obj
}(),
newObj: testingutil.MakeTrainJobWrapper("default", "test").
RuntimeRef(trainer.SchemeGroupVersion.WithKind(trainer.ClusterTrainingRuntimeKind), "test-runtime").
RuntimePatches([]trainer.RuntimePatch{
{
Manager: "acme.io/one",
Time: &oldTime,
TrainingRuntimeSpec: &trainer.TrainingRuntimeSpecPatch{
Template: &trainer.JobSetTemplatePatch{
Spec: &trainer.JobSetSpecPatch{
ReplicatedJobs: []trainer.ReplicatedJobPatch{{
Name: "node",
}},
},
},
},
},
}).Obj(),
wantTime: func(patches []trainer.RuntimePatch) {
if !patches[0].Time.Equal(&expectedTime) {
t.Errorf("expected Time to be updated to %v, got %v", expectedTime, patches[0].Time)
}
},
},
"UPDATE, new patch added alongside unchanged patch": {
oldObj: func() *trainer.TrainJob {
obj := testingutil.MakeTrainJobWrapper("default", "test").
RuntimeRef(trainer.SchemeGroupVersion.WithKind(trainer.ClusterTrainingRuntimeKind), "test-runtime").
RuntimePatches([]trainer.RuntimePatch{
{Manager: "acme.io/existing", Time: &oldTime},
}).Obj()
return obj
}(),
newObj: testingutil.MakeTrainJobWrapper("default", "test").
RuntimeRef(trainer.SchemeGroupVersion.WithKind(trainer.ClusterTrainingRuntimeKind), "test-runtime").
RuntimePatches([]trainer.RuntimePatch{
{Manager: "acme.io/existing"},
{Manager: "acme.io/new"},
}).Obj(),
wantTime: func(patches []trainer.RuntimePatch) {
if !patches[0].Time.Equal(&oldTime) {
t.Errorf("existing patch: expected Time preserved as %v, got %v", oldTime, patches[0].Time)
}
if !patches[1].Time.Equal(&expectedTime) {
t.Errorf("new patch: expected Time %v, got %v", expectedTime, patches[1].Time)
}
},
},
"UPDATE, new patch added with externally-set Time: Time preserved": {
oldObj: func() *trainer.TrainJob {
obj := testingutil.MakeTrainJobWrapper("default", "test").
RuntimeRef(trainer.SchemeGroupVersion.WithKind(trainer.ClusterTrainingRuntimeKind), "test-runtime").
RuntimePatches([]trainer.RuntimePatch{
{Manager: "acme.io/existing", Time: &oldTime},
}).Obj()
return obj
}(),
newObj: testingutil.MakeTrainJobWrapper("default", "test").
RuntimeRef(trainer.SchemeGroupVersion.WithKind(trainer.ClusterTrainingRuntimeKind), "test-runtime").
RuntimePatches([]trainer.RuntimePatch{
{Manager: "acme.io/existing"},
{Manager: "acme.io/new", Time: &oldTime},
}).Obj(),
wantTime: func(patches []trainer.RuntimePatch) {
if !patches[0].Time.Equal(&oldTime) {
t.Errorf("existing patch: expected Time preserved as %v, got %v", oldTime, patches[0].Time)
}
if !patches[1].Time.Equal(&oldTime) {
t.Errorf("new patch: expected externally-set Time %v preserved, got %v", oldTime, patches[1].Time)
}
},
},
}
for name, tc := range cases {
t.Run(name, func(t *testing.T) {
_, ctx := ktesting.NewTestContext(t)
operation := admissionv1.Create
req := admissionv1.AdmissionRequest{}
if tc.oldObj != nil {
operation = admissionv1.Update
raw, err := json.Marshal(tc.oldObj)
if err != nil {
t.Fatal(err)
}
req.OldObject = apiruntime.RawExtension{Raw: raw}
}
req.Operation = operation
ctx = admission.NewContextWithRequest(ctx, admission.Request{AdmissionRequest: req})
defaulter := &TrainJobDefaulter{clock: fakeClock}
if err := defaulter.Default(ctx, tc.newObj); err != nil {
t.Fatalf("Default returned unexpected error: %v", err)
}
tc.wantTime(tc.newObj.Spec.RuntimePatches)
})
}
}
func TestValidateCreate(t *testing.T) {
cases := map[string]struct {
obj *trainer.TrainJob
test/e2e/e2e_test.go
+124 -1
View File
@@ -1,6 +1,8 @@
package e2e
import (
"time"
"github.com/onsi/ginkgo/v2"
"github.com/onsi/gomega"
corev1 "k8s.io/api/core/v1"
@@ -304,11 +306,14 @@ var _ = ginkgo.Describe("TrainJob e2e", func() {
ginkgo.When("Creating a TrainJob with RuntimePatches", func() {
ginkgo.It("should preserve user-provided manager fields", func() {
userTime := metav1.NewTime(time.Now().Add(-time.Hour).Truncate(time.Second))
trainJob := testingutil.MakeTrainJobWrapper(ns.Name, "e2e-test").
RuntimeRef(trainer.SchemeGroupVersion.WithKind(trainer.ClusterTrainingRuntimeKind), torchRuntime).
RuntimePatches([]trainer.RuntimePatch{
{
Manager: "test.io/manager-one",
Time: &userTime,
TrainingRuntimeSpec: &trainer.TrainingRuntimeSpecPatch{
Template: &trainer.JobSetTemplatePatch{
Spec: &trainer.JobSetSpecPatch{
@@ -356,13 +361,131 @@ var _ = ginkgo.Describe("TrainJob e2e", func() {
gomega.Expect(k8sClient.Create(ctx, trainJob)).Should(gomega.Succeed())
})
ginkgo.By("Verify manager fields are preserved", func() {
ginkgo.By("Verify user-provided Time is preserved for test.io/manager-one and Time is set for kueue.k8s.io/manager", func() {
gomega.Eventually(func(g gomega.Gomega) {
gotTrainJob := &trainer.TrainJob{}
g.Expect(k8sClient.Get(ctx, client.ObjectKeyFromObject(trainJob), gotTrainJob)).Should(gomega.Succeed())
g.Expect(gotTrainJob.Spec.RuntimePatches).Should(gomega.HaveLen(2))
g.Expect(gotTrainJob.Spec.RuntimePatches[0].Manager).To(gomega.Equal("test.io/manager-one"))
g.Expect(gotTrainJob.Spec.RuntimePatches[0].Time).ShouldNot(gomega.BeNil())
g.Expect(gotTrainJob.Spec.RuntimePatches[0].Time.Equal(&userTime)).To(gomega.BeTrue())
g.Expect(gotTrainJob.Spec.RuntimePatches[1].Manager).To(gomega.Equal("kueue.k8s.io/manager"))
g.Expect(gotTrainJob.Spec.RuntimePatches[1].Time).ShouldNot(gomega.BeNil())
}, util.Timeout, util.Interval).Should(gomega.Succeed())
})
})
ginkgo.It("should update Time only for the changed patch on update", func() {
userTime := metav1.NewTime(time.Now().Add(-time.Hour).Truncate(time.Second))
newPatchTime := metav1.NewTime(time.Now().Add(-30 * time.Minute).Truncate(time.Second))
trainJob := testingutil.MakeTrainJobWrapper(ns.Name, "e2e-update-time").
RuntimeRef(trainer.SchemeGroupVersion.WithKind(trainer.ClusterTrainingRuntimeKind), torchRuntime).
Suspend(true).
RuntimePatches([]trainer.RuntimePatch{
{
Manager: "test.io/unchanged",
TrainingRuntimeSpec: &trainer.TrainingRuntimeSpecPatch{
Template: &trainer.JobSetTemplatePatch{
Spec: &trainer.JobSetSpecPatch{
ReplicatedJobs: []trainer.ReplicatedJobPatch{{
Name: constants.Node,
Template: &trainer.JobTemplatePatch{
Spec: &trainer.JobSpecPatch{
Template: &trainer.PodTemplatePatch{
Spec: &trainer.PodSpecPatch{
NodeSelector: map[string]string{"zone": "keep"},
},
},
},
},
}},
},
},
},
},
{
Manager: "test.io/will-change",
Time: &userTime,
TrainingRuntimeSpec: &trainer.TrainingRuntimeSpecPatch{
Template: &trainer.JobSetTemplatePatch{
Spec: &trainer.JobSetSpecPatch{
ReplicatedJobs: []trainer.ReplicatedJobPatch{{
Name: constants.Node,
Template: &trainer.JobTemplatePatch{
Spec: &trainer.JobSpecPatch{
Template: &trainer.PodTemplatePatch{
Spec: &trainer.PodSpecPatch{
NodeSelector: map[string]string{"zone": "old"},
},
},
},
},
}},
},
},
},
},
}).
Obj()
ginkgo.By("Create a suspended TrainJob with two RuntimePatches", func() {
gomega.Expect(k8sClient.Create(ctx, trainJob)).Should(gomega.Succeed())
})
var unchangedTime *metav1.Time
ginkgo.By("Record the original Time values and verify user-provided Time is preserved", func() {
gomega.Eventually(func(g gomega.Gomega) {
gotTrainJob := &trainer.TrainJob{}
g.Expect(k8sClient.Get(ctx, client.ObjectKeyFromObject(trainJob), gotTrainJob)).Should(gomega.Succeed())
g.Expect(gotTrainJob.Spec.RuntimePatches).Should(gomega.HaveLen(2))
g.Expect(gotTrainJob.Spec.RuntimePatches[0].Time).ShouldNot(gomega.BeNil())
// User-provided Time for test.io/will-change must be preserved on create.
g.Expect(gotTrainJob.Spec.RuntimePatches[1].Time).ShouldNot(gomega.BeNil())
g.Expect(gotTrainJob.Spec.RuntimePatches[1].Time.Equal(&userTime)).To(gomega.BeTrue())
unchangedTime = gotTrainJob.Spec.RuntimePatches[0].Time.DeepCopy()
}, util.Timeout, util.Interval).Should(gomega.Succeed())
})
// metav1.Time serializes to second precision (RFC3339), so we must
// wait at least one second to guarantee distinguishable timestamps.
time.Sleep(time.Second)
ginkgo.By("Update the second RuntimePatch and add two new patches", func() {
gomega.Eventually(func(g gomega.Gomega) {
gotTrainJob := &trainer.TrainJob{}
g.Expect(k8sClient.Get(ctx, client.ObjectKeyFromObject(trainJob), gotTrainJob)).Should(gomega.Succeed())
gotTrainJob.Spec.RuntimePatches[1].TrainingRuntimeSpec.Template.Spec.ReplicatedJobs[0].
Template.Spec.Template.Spec.NodeSelector = map[string]string{"zone": "new"}
// New patch without Time: webhook should stamp it.
gotTrainJob.Spec.RuntimePatches = append(gotTrainJob.Spec.RuntimePatches, trainer.RuntimePatch{
Manager: "test.io/new-no-time",
})
// New patch with user-provided Time: should be preserved.
gotTrainJob.Spec.RuntimePatches = append(gotTrainJob.Spec.RuntimePatches, trainer.RuntimePatch{
Manager: "test.io/new-with-time",
Time: &newPatchTime,
})
g.Expect(k8sClient.Update(ctx, gotTrainJob)).Should(gomega.Succeed())
}, util.Timeout, util.Interval).Should(gomega.Succeed())
})
ginkgo.By("Verify Time behaviour for all patches after update", func() {
gomega.Eventually(func(g gomega.Gomega) {
gotTrainJob := &trainer.TrainJob{}
g.Expect(k8sClient.Get(ctx, client.ObjectKeyFromObject(trainJob), gotTrainJob)).Should(gomega.Succeed())
g.Expect(gotTrainJob.Spec.RuntimePatches).Should(gomega.HaveLen(4))
// Unchanged pre-existing patch: Time must be preserved.
g.Expect(gotTrainJob.Spec.RuntimePatches[0].Time).ShouldNot(gomega.BeNil())
g.Expect(gotTrainJob.Spec.RuntimePatches[0].Time.Equal(unchangedTime)).To(gomega.BeTrue())
// Changed pre-existing patch: user-provided Time must be replaced.
g.Expect(gotTrainJob.Spec.RuntimePatches[1].Time).ShouldNot(gomega.BeNil())
g.Expect(gotTrainJob.Spec.RuntimePatches[1].Time.Equal(&userTime)).To(gomega.BeFalse())
// New patch without Time: webhook must stamp it.
g.Expect(gotTrainJob.Spec.RuntimePatches[2].Time).ShouldNot(gomega.BeNil())
// New patch with user-provided Time: must be preserved.
g.Expect(gotTrainJob.Spec.RuntimePatches[3].Time).ShouldNot(gomega.BeNil())
g.Expect(gotTrainJob.Spec.RuntimePatches[3].Time.Equal(&newPatchTime)).To(gomega.BeTrue())
}, util.Timeout, util.Interval).Should(gomega.Succeed())
})
})
test/integration/webhooks/trainjob_test.go
+20
View File
@@ -468,6 +468,26 @@ var _ = ginkgo.Describe("TrainJob marker validations and defaulting", ginkgo.Ord
)
})
ginkgo.When("Defaulting RuntimePatch.Time on creation", func() {
ginkgo.It("Should set Time on each RuntimePatch when creating a TrainJob", func() {
trainJob := testingutil.MakeTrainJobWrapper(ns.Name, "time-defaulting").
RuntimeRef(trainer.SchemeGroupVersion.WithKind(trainer.TrainingRuntimeKind), "testing").
RuntimePatches([]trainer.RuntimePatch{
{Manager: "acme.io/manager-one"},
{Manager: "acme.io/manager-two"},
}).
Obj()
gomega.Expect(k8sClient.Create(ctx, trainJob)).Should(gomega.Succeed())
gomega.Eventually(func(g gomega.Gomega) {
got := &trainer.TrainJob{}
g.Expect(k8sClient.Get(ctx, client.ObjectKeyFromObject(trainJob), got)).Should(gomega.Succeed())
g.Expect(got.Spec.RuntimePatches).Should(gomega.HaveLen(2))
g.Expect(got.Spec.RuntimePatches[0].Time).ShouldNot(gomega.BeNil())
g.Expect(got.Spec.RuntimePatches[1].Time).ShouldNot(gomega.BeNil())
}, util.Timeout, util.Interval).Should(gomega.Succeed())
})
})
ginkgo.When("Updating TrainJob", func() {
ginkgo.DescribeTable("Validate TrainJob on update", func(old func() *trainer.TrainJob, new func(*trainer.TrainJob) *trainer.TrainJob, errorMatcher gomega.OmegaMatcher) {
oldTrainJob := old()