mirror of
https://github.com/hanzo-ml/trainer.git
synced 2026-09-24 00:31:34 +00:00
Peter PanandGitHub
15b0a9fc1d
* feat: inject PET envs into init containers via envInjection config - Add EnvInjection and EnvInjectionTarget types to TorchMLPolicySource - Add FindContainerByPodSetName helper in runtime package - Update torch plugin to inject PET envs into additional containers - Update jobset plugin to sync init containers - Add 5 comprehensive unit tests for envInjection scenarios - Add Python API models for envInjection types Signed-off-by: Peter Pan <[email protected]> * fix: add missing +optional markers for kubeapilinter validation Signed-off-by: Peter Pan <[email protected]> * fix: change JobName to optional field for kubeapilinter Change JobName from required to optional with omitempty tag to satisfy kubeapilinter validation. Add default marker to CRD schema for listMapKey requirement. Signed-off-by: Peter Pan <[email protected]> * fix test Signed-off-by: Peter Pan <[email protected]> * fix: address envInjection review feedback Reject missing envInjection targets instead of silently skipping them, align the API schema with required fields, and harden init container propagation so misconfigured runtimes fail predictably. Signed-off-by: Peter Pan <[email protected]> * fix: sync envInjection generated files Add the generated boilerplate and Python model updates expected by make generate so CI does not report stale generated assets. Signed-off-by: Peter Pan <[email protected]> * fix: satisfy kube API lint for envInjection fields Keep the envInjection fields marked as required while preserving omitempty tags required by kubeapilinter. Signed-off-by: Peter Pan <[email protected]> * fix: update envInjection OpenAPI defaults Sync generated OpenAPI and Python model output after restoring omitempty on required envInjection fields. Signed-off-by: Peter Pan <[email protected]> * fix: address latest review comments Signed-off-by: Peter Pan <[email protected]> * chore: remove test plan file per review feedback The KEP README already has a Test Plan section; a separate file with local environment details does not belong in the upstream PR. Signed-off-by: Peter Pan <[email protected]> --------- Signed-off-by: Peter Pan <[email protected]>
2194 lines
68 KiB
Go
2194 lines
68 KiB
Go
/*
|
|
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
|
|
|
|
http://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.
|
|
*/
|
|
|
|
package jobset
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"testing"
|
|
|
|
"github.com/google/go-cmp/cmp"
|
|
"github.com/google/go-cmp/cmp/cmpopts"
|
|
batchv1 "k8s.io/api/batch/v1"
|
|
corev1 "k8s.io/api/core/v1"
|
|
apierrors "k8s.io/apimachinery/pkg/api/errors"
|
|
metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
|
|
apiruntime "k8s.io/apimachinery/pkg/runtime"
|
|
"k8s.io/apimachinery/pkg/util/validation/field"
|
|
batchv1ac "k8s.io/client-go/applyconfigurations/batch/v1"
|
|
corev1ac "k8s.io/client-go/applyconfigurations/core/v1"
|
|
"k8s.io/klog/v2/ktesting"
|
|
"k8s.io/utils/ptr"
|
|
"sigs.k8s.io/controller-runtime/pkg/client"
|
|
"sigs.k8s.io/controller-runtime/pkg/client/interceptor"
|
|
"sigs.k8s.io/controller-runtime/pkg/webhook/admission"
|
|
jobsetv1alpha2 "sigs.k8s.io/jobset/api/jobset/v1alpha2"
|
|
jobsetv1alpha2ac "sigs.k8s.io/jobset/client-go/applyconfiguration/jobset/v1alpha2"
|
|
|
|
trainer "github.com/kubeflow/trainer/v2/pkg/apis/trainer/v1alpha1"
|
|
"github.com/kubeflow/trainer/v2/pkg/constants"
|
|
"github.com/kubeflow/trainer/v2/pkg/runtime"
|
|
"github.com/kubeflow/trainer/v2/pkg/runtime/framework"
|
|
utiltesting "github.com/kubeflow/trainer/v2/pkg/util/testing"
|
|
)
|
|
|
|
// TODO: Add tests for all Interfaces.
|
|
// REF: https://github.com/kubeflow/trainer/issues/2468
|
|
|
|
func TestJobSet(t *testing.T) {
|
|
cases := map[string]struct {
|
|
trainJob *trainer.TrainJob
|
|
info *runtime.Info
|
|
wantInfo *runtime.Info
|
|
wantError error
|
|
}{
|
|
"no action when info is nil": {
|
|
trainJob: utiltesting.MakeTrainJobWrapper(metav1.NamespaceDefault, "trainJob").
|
|
Obj(),
|
|
},
|
|
"no action when trainJob is not nil": {
|
|
info: &runtime.Info{
|
|
Labels: map[string]string{"key": "value"},
|
|
},
|
|
wantInfo: &runtime.Info{
|
|
Labels: map[string]string{"key": "value"},
|
|
},
|
|
},
|
|
"no action when template.spec is not JobSet": {
|
|
info: &runtime.Info{
|
|
Labels: map[string]string{"key": "value"},
|
|
TemplateSpec: runtime.TemplateSpec{
|
|
ObjApply: batchv1ac.JobSpec(),
|
|
},
|
|
},
|
|
trainJob: utiltesting.MakeTrainJobWrapper(metav1.NamespaceDefault, "trainJob").
|
|
Obj(),
|
|
wantInfo: &runtime.Info{
|
|
Labels: map[string]string{"key": "value"},
|
|
TemplateSpec: runtime.TemplateSpec{
|
|
ObjApply: batchv1ac.JobSpec(),
|
|
},
|
|
},
|
|
},
|
|
"trainer numNodes is respected rather than parallelism when replicatedJob name is node": {
|
|
trainJob: utiltesting.MakeTrainJobWrapper(metav1.NamespaceDefault, "trainJob").
|
|
Obj(),
|
|
info: &runtime.Info{
|
|
RuntimePolicy: runtime.RuntimePolicy{
|
|
MLPolicySource: utiltesting.MakeMLPolicySourceWrapper().
|
|
MPIPolicy(nil, trainer.MPIImplementationOpenMPI, nil, nil).
|
|
Obj(),
|
|
},
|
|
TemplateSpec: runtime.TemplateSpec{
|
|
PodSets: []runtime.PodSet{
|
|
{
|
|
Name: constants.Launcher,
|
|
Containers: make([]runtime.Container, 1),
|
|
},
|
|
{
|
|
Name: constants.Node,
|
|
Count: ptr.To[int32](2),
|
|
Containers: make([]runtime.Container, 1),
|
|
},
|
|
},
|
|
ObjApply: jobsetv1alpha2ac.JobSetSpec().
|
|
WithReplicatedJobs(
|
|
jobsetv1alpha2ac.ReplicatedJob().
|
|
WithName(constants.Launcher).
|
|
WithTemplate(batchv1ac.JobTemplateSpec().
|
|
WithSpec(batchv1ac.JobSpec().
|
|
WithParallelism(1).
|
|
WithTemplate(corev1ac.PodTemplateSpec().
|
|
WithSpec(corev1ac.PodSpec().
|
|
WithContainers(
|
|
corev1ac.Container().WithName("sidecar"),
|
|
corev1ac.Container().WithName(constants.Node),
|
|
),
|
|
),
|
|
),
|
|
),
|
|
),
|
|
jobsetv1alpha2ac.ReplicatedJob().
|
|
WithName(constants.Node).
|
|
WithTemplate(batchv1ac.JobTemplateSpec().
|
|
WithSpec(batchv1ac.JobSpec().
|
|
WithParallelism(2).
|
|
WithTemplate(corev1ac.PodTemplateSpec().
|
|
WithSpec(corev1ac.PodSpec().
|
|
WithContainers(
|
|
corev1ac.Container().WithName(constants.Node),
|
|
),
|
|
),
|
|
),
|
|
),
|
|
),
|
|
),
|
|
},
|
|
},
|
|
wantInfo: &runtime.Info{
|
|
RuntimePolicy: runtime.RuntimePolicy{
|
|
MLPolicySource: utiltesting.MakeMLPolicySourceWrapper().
|
|
MPIPolicy(nil, trainer.MPIImplementationOpenMPI, nil, nil).
|
|
Obj(),
|
|
},
|
|
TemplateSpec: runtime.TemplateSpec{
|
|
PodSets: []runtime.PodSet{
|
|
{
|
|
Name: constants.Launcher,
|
|
Containers: make([]runtime.Container, 1),
|
|
Endpoints: func(yield func(string) bool) {
|
|
yield("trainJob-launcher-0-0.trainJob")
|
|
},
|
|
},
|
|
{
|
|
Name: constants.Node,
|
|
Count: ptr.To[int32](2),
|
|
Containers: make([]runtime.Container, 1),
|
|
Endpoints: func(yield func(string) bool) {
|
|
yield("trainJob-node-0-0.trainJob")
|
|
yield("trainJob-node-0-1.trainJob")
|
|
},
|
|
},
|
|
},
|
|
ObjApply: jobsetv1alpha2ac.JobSetSpec().
|
|
WithReplicatedJobs(
|
|
jobsetv1alpha2ac.ReplicatedJob().
|
|
WithName(constants.Launcher).
|
|
WithTemplate(batchv1ac.JobTemplateSpec().
|
|
WithSpec(batchv1ac.JobSpec().
|
|
WithParallelism(1).
|
|
WithTemplate(corev1ac.PodTemplateSpec().
|
|
WithSpec(corev1ac.PodSpec().
|
|
WithContainers(
|
|
corev1ac.Container().WithName("sidecar"),
|
|
corev1ac.Container().WithName(constants.Node),
|
|
),
|
|
),
|
|
),
|
|
),
|
|
),
|
|
jobsetv1alpha2ac.ReplicatedJob().
|
|
WithName(constants.Node).
|
|
WithTemplate(batchv1ac.JobTemplateSpec().
|
|
WithSpec(batchv1ac.JobSpec().
|
|
WithParallelism(2).
|
|
WithTemplate(corev1ac.PodTemplateSpec().
|
|
WithSpec(corev1ac.PodSpec().
|
|
WithContainers(
|
|
corev1ac.Container().WithName(constants.Node),
|
|
),
|
|
),
|
|
),
|
|
),
|
|
),
|
|
),
|
|
},
|
|
},
|
|
},
|
|
"subDomain in jobSetSpec is used to endpoint": {
|
|
trainJob: utiltesting.MakeTrainJobWrapper(metav1.NamespaceDefault, "trainJob").
|
|
Obj(),
|
|
info: &runtime.Info{
|
|
RuntimePolicy: runtime.RuntimePolicy{
|
|
MLPolicySource: utiltesting.MakeMLPolicySourceWrapper().Obj(),
|
|
},
|
|
TemplateSpec: runtime.TemplateSpec{
|
|
PodSets: []runtime.PodSet{
|
|
{
|
|
Name: constants.Launcher,
|
|
Containers: make([]runtime.Container, 1),
|
|
},
|
|
{
|
|
Name: constants.Node,
|
|
Containers: make([]runtime.Container, 1),
|
|
},
|
|
},
|
|
ObjApply: jobsetv1alpha2ac.JobSetSpec().
|
|
WithNetwork(jobsetv1alpha2ac.Network().
|
|
WithSubdomain("kubeflow.org")).
|
|
WithReplicatedJobs(
|
|
jobsetv1alpha2ac.ReplicatedJob().
|
|
WithName(constants.Launcher).
|
|
WithTemplate(batchv1ac.JobTemplateSpec().
|
|
WithSpec(batchv1ac.JobSpec().
|
|
WithParallelism(1).
|
|
WithTemplate(corev1ac.PodTemplateSpec().
|
|
WithSpec(corev1ac.PodSpec().
|
|
WithContainers(
|
|
corev1ac.Container().WithName(constants.Node),
|
|
),
|
|
),
|
|
),
|
|
),
|
|
),
|
|
jobsetv1alpha2ac.ReplicatedJob().
|
|
WithName(constants.Node).
|
|
WithTemplate(batchv1ac.JobTemplateSpec().
|
|
WithSpec(batchv1ac.JobSpec().
|
|
WithParallelism(1).
|
|
WithTemplate(corev1ac.PodTemplateSpec().
|
|
WithSpec(corev1ac.PodSpec().
|
|
WithContainers(
|
|
corev1ac.Container().WithName(constants.Node),
|
|
),
|
|
),
|
|
),
|
|
),
|
|
),
|
|
),
|
|
},
|
|
},
|
|
wantInfo: &runtime.Info{
|
|
RuntimePolicy: runtime.RuntimePolicy{
|
|
MLPolicySource: utiltesting.MakeMLPolicySourceWrapper().Obj(),
|
|
},
|
|
TemplateSpec: runtime.TemplateSpec{
|
|
PodSets: []runtime.PodSet{
|
|
{
|
|
Name: constants.Launcher,
|
|
Containers: make([]runtime.Container, 1),
|
|
Endpoints: func(yield func(string) bool) {
|
|
yield("trainJob-launcher-0-0.kubeflow.org")
|
|
},
|
|
},
|
|
{
|
|
Name: constants.Node,
|
|
Containers: make([]runtime.Container, 1),
|
|
Endpoints: func(yield func(string) bool) {
|
|
yield("trainJob-node-0-0.kubeflow.org")
|
|
},
|
|
},
|
|
},
|
|
ObjApply: jobsetv1alpha2ac.JobSetSpec().
|
|
WithNetwork(jobsetv1alpha2ac.Network().
|
|
WithSubdomain("kubeflow.org")).
|
|
WithReplicatedJobs(
|
|
jobsetv1alpha2ac.ReplicatedJob().
|
|
WithName(constants.Launcher).
|
|
WithTemplate(batchv1ac.JobTemplateSpec().
|
|
WithSpec(batchv1ac.JobSpec().
|
|
WithParallelism(1).
|
|
WithTemplate(corev1ac.PodTemplateSpec().
|
|
WithSpec(corev1ac.PodSpec().
|
|
WithContainers(
|
|
corev1ac.Container().WithName(constants.Node),
|
|
),
|
|
),
|
|
),
|
|
),
|
|
),
|
|
jobsetv1alpha2ac.ReplicatedJob().
|
|
WithName(constants.Node).
|
|
WithTemplate(batchv1ac.JobTemplateSpec().
|
|
WithSpec(batchv1ac.JobSpec().
|
|
WithParallelism(1).
|
|
WithTemplate(corev1ac.PodTemplateSpec().
|
|
WithSpec(corev1ac.PodSpec().
|
|
WithContainers(
|
|
corev1ac.Container().WithName(constants.Node),
|
|
),
|
|
),
|
|
),
|
|
),
|
|
),
|
|
),
|
|
},
|
|
},
|
|
},
|
|
}
|
|
for name, tc := range cases {
|
|
t.Run(name, func(t *testing.T) {
|
|
_, ctx := ktesting.NewTestContext(t)
|
|
var cancel func()
|
|
ctx, cancel = context.WithCancel(ctx)
|
|
t.Cleanup(cancel)
|
|
cli := utiltesting.NewClientBuilder().Build()
|
|
p, err := New(ctx, cli, nil, nil)
|
|
if err != nil {
|
|
t.Fatalf("Failed to initialize JobSet plugin: %v", err)
|
|
}
|
|
err = p.(framework.PodNetworkPlugin).IdentifyPodNetwork(tc.info, tc.trainJob)
|
|
if diff := cmp.Diff(tc.wantError, err, cmpopts.EquateErrors()); len(diff) != 0 {
|
|
t.Errorf("Unexpected error (-want,+got):\n%s", diff)
|
|
}
|
|
if diff := cmp.Diff(tc.wantInfo, tc.info,
|
|
cmpopts.SortSlices(func(a, b string) bool { return a < b }),
|
|
cmpopts.SortMaps(func(a, b string) bool { return a < b }),
|
|
utiltesting.PodSetEndpointsCmpOpts,
|
|
); len(diff) != 0 {
|
|
t.Errorf("Unexpected Info from IdentifyPodNetwork (-want,+got):\n%s", diff)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestValidate(t *testing.T) {
|
|
cases := map[string]struct {
|
|
info *runtime.Info
|
|
oldObj *trainer.TrainJob
|
|
newObj *trainer.TrainJob
|
|
jobSet *jobsetv1alpha2.JobSet
|
|
clientErr error
|
|
wantError field.ErrorList
|
|
wantWarnings admission.Warnings
|
|
}{
|
|
"no initializer job": {
|
|
info: &runtime.Info{TemplateSpec: runtime.TemplateSpec{
|
|
ObjApply: &jobsetv1alpha2ac.JobSetSpecApplyConfiguration{},
|
|
}},
|
|
newObj: utiltesting.MakeTrainJobWrapper(metav1.NamespaceDefault, "test").Initializer(nil).
|
|
Obj(),
|
|
},
|
|
"no dataset initializer job": {
|
|
info: &runtime.Info{TemplateSpec: runtime.TemplateSpec{
|
|
ObjApply: &jobsetv1alpha2ac.JobSetSpecApplyConfiguration{},
|
|
}},
|
|
newObj: utiltesting.MakeTrainJobWrapper(metav1.NamespaceDefault, "test").
|
|
Initializer(&trainer.Initializer{Dataset: nil}).
|
|
Obj(),
|
|
},
|
|
"must have dataset initializer job when trainJob is configured with input datasetConfig": {
|
|
info: &runtime.Info{
|
|
TemplateSpec: runtime.TemplateSpec{
|
|
ObjApply: &jobsetv1alpha2ac.JobSetSpecApplyConfiguration{
|
|
ReplicatedJobs: []jobsetv1alpha2ac.ReplicatedJobApplyConfiguration{
|
|
{
|
|
Name: ptr.To("random"),
|
|
Template: &batchv1ac.JobTemplateSpecApplyConfiguration{
|
|
Spec: &batchv1ac.JobSpecApplyConfiguration{
|
|
Template: &corev1ac.PodTemplateSpecApplyConfiguration{
|
|
Spec: &corev1ac.PodSpecApplyConfiguration{
|
|
Containers: []corev1ac.ContainerApplyConfiguration{
|
|
{
|
|
Name: ptr.To("random"),
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
newObj: utiltesting.MakeTrainJobWrapper("default", "test").
|
|
Initializer(&trainer.Initializer{
|
|
Dataset: &trainer.DatasetInitializer{},
|
|
}).Obj(),
|
|
wantError: field.ErrorList{
|
|
field.Invalid(runtimeRefPath,
|
|
utiltesting.MakeTrainJobWrapper("default", "test").Obj().Spec.RuntimeRef,
|
|
fmt.Sprintf("must have %s job when trainJob is configured with input datasetConfig", constants.DatasetInitializer)),
|
|
},
|
|
},
|
|
"must have container with name - dataset initializer in the dataset initializer job": {
|
|
info: &runtime.Info{
|
|
TemplateSpec: runtime.TemplateSpec{
|
|
ObjApply: &jobsetv1alpha2ac.JobSetSpecApplyConfiguration{
|
|
ReplicatedJobs: []jobsetv1alpha2ac.ReplicatedJobApplyConfiguration{
|
|
{
|
|
Name: ptr.To(constants.DatasetInitializer),
|
|
Template: &batchv1ac.JobTemplateSpecApplyConfiguration{
|
|
Spec: &batchv1ac.JobSpecApplyConfiguration{
|
|
Template: &corev1ac.PodTemplateSpecApplyConfiguration{
|
|
Spec: &corev1ac.PodSpecApplyConfiguration{
|
|
Containers: []corev1ac.ContainerApplyConfiguration{},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
newObj: utiltesting.MakeTrainJobWrapper("default", "test").
|
|
Initializer(&trainer.Initializer{
|
|
Dataset: &trainer.DatasetInitializer{},
|
|
}).Obj(),
|
|
wantError: field.ErrorList{
|
|
field.Invalid(runtimeRefPath,
|
|
utiltesting.MakeTrainJobWrapper("default", "test").Obj().Spec.RuntimeRef,
|
|
fmt.Sprintf("must have container with name - %s in the %s job", constants.DatasetInitializer, constants.DatasetInitializer)),
|
|
},
|
|
},
|
|
"no model initializer job": {
|
|
info: &runtime.Info{
|
|
TemplateSpec: runtime.TemplateSpec{
|
|
ObjApply: &jobsetv1alpha2ac.JobSetSpecApplyConfiguration{
|
|
ReplicatedJobs: []jobsetv1alpha2ac.ReplicatedJobApplyConfiguration{
|
|
{
|
|
Name: ptr.To(constants.DatasetInitializer),
|
|
Template: &batchv1ac.JobTemplateSpecApplyConfiguration{
|
|
Spec: &batchv1ac.JobSpecApplyConfiguration{
|
|
Template: &corev1ac.PodTemplateSpecApplyConfiguration{
|
|
Spec: &corev1ac.PodSpecApplyConfiguration{
|
|
Containers: []corev1ac.ContainerApplyConfiguration{
|
|
{
|
|
Name: ptr.To(constants.DatasetInitializer),
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
newObj: utiltesting.MakeTrainJobWrapper(metav1.NamespaceDefault, "test").
|
|
Initializer(&trainer.Initializer{Dataset: nil}).
|
|
Obj(),
|
|
},
|
|
"must have model initializer job when trainJob is configured with input modelConfig": {
|
|
info: &runtime.Info{
|
|
TemplateSpec: runtime.TemplateSpec{
|
|
ObjApply: &jobsetv1alpha2ac.JobSetSpecApplyConfiguration{
|
|
ReplicatedJobs: []jobsetv1alpha2ac.ReplicatedJobApplyConfiguration{
|
|
{
|
|
Name: ptr.To("random"),
|
|
Template: &batchv1ac.JobTemplateSpecApplyConfiguration{
|
|
Spec: &batchv1ac.JobSpecApplyConfiguration{
|
|
Template: &corev1ac.PodTemplateSpecApplyConfiguration{
|
|
Spec: &corev1ac.PodSpecApplyConfiguration{
|
|
Containers: []corev1ac.ContainerApplyConfiguration{
|
|
{
|
|
Name: ptr.To("random"),
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
newObj: utiltesting.MakeTrainJobWrapper("default", "test").
|
|
Initializer(&trainer.Initializer{
|
|
Model: &trainer.ModelInitializer{},
|
|
}).Obj(),
|
|
wantError: field.ErrorList{
|
|
field.Invalid(runtimeRefPath,
|
|
utiltesting.MakeTrainJobWrapper("default", "test").Obj().Spec.RuntimeRef,
|
|
fmt.Sprintf("must have %s job when trainJob is configured with input modelConfig", constants.ModelInitializer)),
|
|
},
|
|
},
|
|
"must have container with name - model initializer in the model initializer job": {
|
|
info: &runtime.Info{
|
|
TemplateSpec: runtime.TemplateSpec{
|
|
ObjApply: &jobsetv1alpha2ac.JobSetSpecApplyConfiguration{
|
|
ReplicatedJobs: []jobsetv1alpha2ac.ReplicatedJobApplyConfiguration{
|
|
{
|
|
Name: ptr.To(constants.ModelInitializer),
|
|
Template: &batchv1ac.JobTemplateSpecApplyConfiguration{
|
|
Spec: &batchv1ac.JobSpecApplyConfiguration{
|
|
Template: &corev1ac.PodTemplateSpecApplyConfiguration{
|
|
Spec: &corev1ac.PodSpecApplyConfiguration{
|
|
Containers: []corev1ac.ContainerApplyConfiguration{},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
newObj: utiltesting.MakeTrainJobWrapper("default", "test").
|
|
Initializer(&trainer.Initializer{
|
|
Model: &trainer.ModelInitializer{},
|
|
}).Obj(),
|
|
wantError: field.ErrorList{
|
|
field.Invalid(runtimeRefPath,
|
|
utiltesting.MakeTrainJobWrapper("default", "test").Obj().Spec.RuntimeRef,
|
|
fmt.Sprintf("must have container with name - %s in the %s job", constants.ModelInitializer, constants.ModelInitializer)),
|
|
},
|
|
},
|
|
"runtimePatches contain invalid replicated job": {
|
|
info: &runtime.Info{
|
|
TemplateSpec: runtime.TemplateSpec{
|
|
ObjApply: &jobsetv1alpha2ac.JobSetSpecApplyConfiguration{
|
|
ReplicatedJobs: []jobsetv1alpha2ac.ReplicatedJobApplyConfiguration{
|
|
{
|
|
Name: ptr.To(constants.Node),
|
|
Template: &batchv1ac.JobTemplateSpecApplyConfiguration{
|
|
Spec: &batchv1ac.JobSpecApplyConfiguration{
|
|
Template: &corev1ac.PodTemplateSpecApplyConfiguration{
|
|
Spec: &corev1ac.PodSpecApplyConfiguration{
|
|
Containers: []corev1ac.ContainerApplyConfiguration{},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
newObj: utiltesting.MakeTrainJobWrapper("default", "test").
|
|
RuntimePatches([]trainer.RuntimePatch{
|
|
{
|
|
Manager: "test.io/manager",
|
|
TrainingRuntimeSpec: &trainer.TrainingRuntimeSpecPatch{
|
|
Template: &trainer.JobSetTemplatePatch{
|
|
Spec: &trainer.JobSetSpecPatch{
|
|
ReplicatedJobs: []trainer.ReplicatedJobPatch{{Name: "invalid"}},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
}).Obj(),
|
|
wantError: field.ErrorList{
|
|
field.Invalid(runtimePatchesPath,
|
|
[]trainer.RuntimePatch{
|
|
{
|
|
Manager: "test.io/manager",
|
|
TrainingRuntimeSpec: &trainer.TrainingRuntimeSpecPatch{
|
|
Template: &trainer.JobSetTemplatePatch{
|
|
Spec: &trainer.JobSetSpecPatch{
|
|
ReplicatedJobs: []trainer.ReplicatedJobPatch{{Name: "invalid"}},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
"must not have replicated job that doesn't exist in the runtime job template"),
|
|
},
|
|
},
|
|
"runtimePatches contain invalid initContainer": {
|
|
info: &runtime.Info{
|
|
TemplateSpec: runtime.TemplateSpec{
|
|
ObjApply: &jobsetv1alpha2ac.JobSetSpecApplyConfiguration{
|
|
ReplicatedJobs: []jobsetv1alpha2ac.ReplicatedJobApplyConfiguration{
|
|
{
|
|
Name: ptr.To(constants.Node),
|
|
Template: &batchv1ac.JobTemplateSpecApplyConfiguration{
|
|
Spec: &batchv1ac.JobSpecApplyConfiguration{
|
|
Template: &corev1ac.PodTemplateSpecApplyConfiguration{
|
|
Spec: &corev1ac.PodSpecApplyConfiguration{
|
|
InitContainers: []corev1ac.ContainerApplyConfiguration{
|
|
{
|
|
Name: ptr.To("custom-init"),
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
newObj: utiltesting.MakeTrainJobWrapper("default", "test").
|
|
RuntimePatches([]trainer.RuntimePatch{
|
|
{
|
|
Manager: "test.io/manager",
|
|
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{
|
|
InitContainers: []trainer.ContainerPatch{
|
|
{Name: "invalid"},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
}},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
}).Obj(),
|
|
wantError: field.ErrorList{
|
|
field.Invalid(runtimePatchesPath,
|
|
[]trainer.RuntimePatch{
|
|
{
|
|
Manager: "test.io/manager",
|
|
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{
|
|
InitContainers: []trainer.ContainerPatch{
|
|
{Name: "invalid"},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
}},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
fmt.Sprintf("must not have initContainer that doesn't exist in the runtime job %s", constants.Node)),
|
|
},
|
|
},
|
|
"runtimePatches contain invalid container": {
|
|
info: &runtime.Info{
|
|
TemplateSpec: runtime.TemplateSpec{
|
|
ObjApply: &jobsetv1alpha2ac.JobSetSpecApplyConfiguration{
|
|
ReplicatedJobs: []jobsetv1alpha2ac.ReplicatedJobApplyConfiguration{
|
|
{
|
|
Name: ptr.To(constants.Node),
|
|
Template: &batchv1ac.JobTemplateSpecApplyConfiguration{
|
|
Spec: &batchv1ac.JobSpecApplyConfiguration{
|
|
Template: &corev1ac.PodTemplateSpecApplyConfiguration{
|
|
Spec: &corev1ac.PodSpecApplyConfiguration{
|
|
Containers: []corev1ac.ContainerApplyConfiguration{
|
|
{
|
|
Name: ptr.To(constants.Node),
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
newObj: utiltesting.MakeTrainJobWrapper("default", "test").
|
|
RuntimePatches([]trainer.RuntimePatch{
|
|
{
|
|
Manager: "test.io/manager",
|
|
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{
|
|
Containers: []trainer.ContainerPatch{
|
|
{Name: "invalid"},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
}},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
}).Obj(),
|
|
wantError: field.ErrorList{
|
|
field.Invalid(runtimePatchesPath,
|
|
[]trainer.RuntimePatch{
|
|
{
|
|
Manager: "test.io/manager",
|
|
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{
|
|
Containers: []trainer.ContainerPatch{
|
|
{Name: "invalid"},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
}},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
fmt.Sprintf("must not have container that doesn't exist in the runtime job %s", constants.Node)),
|
|
},
|
|
},
|
|
"runtimePatches contain envs for reserved container": {
|
|
info: &runtime.Info{
|
|
TemplateSpec: runtime.TemplateSpec{
|
|
ObjApply: &jobsetv1alpha2ac.JobSetSpecApplyConfiguration{
|
|
ReplicatedJobs: []jobsetv1alpha2ac.ReplicatedJobApplyConfiguration{
|
|
{
|
|
Name: ptr.To(constants.Node),
|
|
Template: &batchv1ac.JobTemplateSpecApplyConfiguration{
|
|
Spec: &batchv1ac.JobSpecApplyConfiguration{
|
|
Template: &corev1ac.PodTemplateSpecApplyConfiguration{
|
|
Spec: &corev1ac.PodSpecApplyConfiguration{
|
|
Containers: []corev1ac.ContainerApplyConfiguration{
|
|
{
|
|
Name: ptr.To(constants.Node),
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
newObj: utiltesting.MakeTrainJobWrapper("default", "test").
|
|
RuntimePatches([]trainer.RuntimePatch{
|
|
{
|
|
Manager: "test.io/manager",
|
|
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{
|
|
Containers: []trainer.ContainerPatch{
|
|
{
|
|
Name: constants.Node,
|
|
Env: []corev1.EnvVar{
|
|
{
|
|
Name: "ENV_NAME",
|
|
Value: "OVERRIDE",
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
}},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
}).Obj(),
|
|
wantError: field.ErrorList{
|
|
field.Invalid(runtimePatchesPath,
|
|
[]trainer.RuntimePatch{
|
|
{
|
|
Manager: "test.io/manager",
|
|
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{
|
|
Containers: []trainer.ContainerPatch{
|
|
{
|
|
Name: constants.Node,
|
|
Env: []corev1.EnvVar{
|
|
{
|
|
Name: "ENV_NAME",
|
|
Value: "OVERRIDE",
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
}},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
fmt.Sprintf("must not have envs for the %s, %s, %s containers", constants.DatasetInitializer, constants.ModelInitializer, constants.Node)),
|
|
},
|
|
},
|
|
"allow runtimePatches when creating a new trainJob": {
|
|
info: &runtime.Info{
|
|
TemplateSpec: runtime.TemplateSpec{
|
|
ObjApply: &jobsetv1alpha2ac.JobSetSpecApplyConfiguration{
|
|
ReplicatedJobs: []jobsetv1alpha2ac.ReplicatedJobApplyConfiguration{
|
|
{
|
|
Name: ptr.To(constants.Node),
|
|
Template: &batchv1ac.JobTemplateSpecApplyConfiguration{
|
|
Spec: &batchv1ac.JobSpecApplyConfiguration{
|
|
Template: &corev1ac.PodTemplateSpecApplyConfiguration{
|
|
Spec: &corev1ac.PodSpecApplyConfiguration{
|
|
Containers: []corev1ac.ContainerApplyConfiguration{
|
|
{
|
|
Name: ptr.To(constants.Node),
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
oldObj: nil,
|
|
newObj: utiltesting.MakeTrainJobWrapper(metav1.NamespaceDefault, "test").
|
|
RuntimePatches([]trainer.RuntimePatch{
|
|
{
|
|
Manager: "test.io/manager",
|
|
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{
|
|
ServiceAccountName: ptr.To("service-account"),
|
|
},
|
|
},
|
|
},
|
|
},
|
|
}},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
}).
|
|
Obj(),
|
|
wantError: nil,
|
|
},
|
|
"allow updates to trainJob with no changes to runtimePatches": {
|
|
info: &runtime.Info{
|
|
TemplateSpec: runtime.TemplateSpec{
|
|
ObjApply: &jobsetv1alpha2ac.JobSetSpecApplyConfiguration{
|
|
ReplicatedJobs: []jobsetv1alpha2ac.ReplicatedJobApplyConfiguration{
|
|
{
|
|
Name: ptr.To(constants.Node),
|
|
Template: &batchv1ac.JobTemplateSpecApplyConfiguration{
|
|
Spec: &batchv1ac.JobSpecApplyConfiguration{
|
|
Template: &corev1ac.PodTemplateSpecApplyConfiguration{
|
|
Spec: &corev1ac.PodSpecApplyConfiguration{
|
|
Containers: []corev1ac.ContainerApplyConfiguration{
|
|
{
|
|
Name: ptr.To(constants.Node),
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
oldObj: utiltesting.MakeTrainJobWrapper(metav1.NamespaceDefault, "test").
|
|
RuntimePatches([]trainer.RuntimePatch{
|
|
{
|
|
Manager: "test.io/manager",
|
|
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{
|
|
ServiceAccountName: ptr.To("service-account"),
|
|
},
|
|
},
|
|
},
|
|
},
|
|
}},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
}).
|
|
Obj(),
|
|
newObj: utiltesting.MakeTrainJobWrapper(metav1.NamespaceDefault, "test").
|
|
RuntimePatches([]trainer.RuntimePatch{
|
|
{
|
|
Manager: "test.io/manager",
|
|
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{
|
|
ServiceAccountName: ptr.To("service-account"),
|
|
},
|
|
},
|
|
},
|
|
},
|
|
}},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
}).
|
|
Obj(),
|
|
wantError: nil,
|
|
},
|
|
"forbid changes to runtimePatches when trainJob is not suspended": {
|
|
info: &runtime.Info{
|
|
TemplateSpec: runtime.TemplateSpec{
|
|
ObjApply: &jobsetv1alpha2ac.JobSetSpecApplyConfiguration{
|
|
ReplicatedJobs: []jobsetv1alpha2ac.ReplicatedJobApplyConfiguration{
|
|
{
|
|
Name: ptr.To(constants.Node),
|
|
Template: &batchv1ac.JobTemplateSpecApplyConfiguration{
|
|
Spec: &batchv1ac.JobSpecApplyConfiguration{
|
|
Template: &corev1ac.PodTemplateSpecApplyConfiguration{
|
|
Spec: &corev1ac.PodSpecApplyConfiguration{
|
|
Containers: []corev1ac.ContainerApplyConfiguration{
|
|
{
|
|
Name: ptr.To(constants.Node),
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
oldObj: utiltesting.MakeTrainJobWrapper(metav1.NamespaceDefault, "test").
|
|
Suspend(false).
|
|
RuntimePatches([]trainer.RuntimePatch{
|
|
{
|
|
Manager: "test.io/manager",
|
|
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{
|
|
ServiceAccountName: ptr.To("service-account"),
|
|
},
|
|
},
|
|
},
|
|
},
|
|
}},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
}).
|
|
Obj(),
|
|
newObj: utiltesting.MakeTrainJobWrapper(metav1.NamespaceDefault, "test").
|
|
Suspend(false).
|
|
RuntimePatches([]trainer.RuntimePatch{
|
|
{
|
|
Manager: "test.io/manager",
|
|
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{
|
|
ServiceAccountName: ptr.To("service-account-updated"),
|
|
},
|
|
},
|
|
},
|
|
},
|
|
}},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
}).
|
|
Obj(),
|
|
wantError: field.ErrorList{
|
|
field.Forbidden(runtimePatchesPath, "RuntimePatches can only be modified when the TrainJob is suspended before or after the update"),
|
|
},
|
|
},
|
|
"allow changes to runtimePatches when trainJob is suspended and jobSet does not exist": {
|
|
info: &runtime.Info{
|
|
TemplateSpec: runtime.TemplateSpec{
|
|
ObjApply: &jobsetv1alpha2ac.JobSetSpecApplyConfiguration{
|
|
ReplicatedJobs: []jobsetv1alpha2ac.ReplicatedJobApplyConfiguration{
|
|
{
|
|
Name: ptr.To(constants.Node),
|
|
Template: &batchv1ac.JobTemplateSpecApplyConfiguration{
|
|
Spec: &batchv1ac.JobSpecApplyConfiguration{
|
|
Template: &corev1ac.PodTemplateSpecApplyConfiguration{
|
|
Spec: &corev1ac.PodSpecApplyConfiguration{
|
|
Containers: []corev1ac.ContainerApplyConfiguration{
|
|
{
|
|
Name: ptr.To(constants.Node),
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
oldObj: utiltesting.MakeTrainJobWrapper(metav1.NamespaceDefault, "test").
|
|
Suspend(true).
|
|
RuntimePatches([]trainer.RuntimePatch{
|
|
{
|
|
Manager: "test.io/manager",
|
|
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{
|
|
ServiceAccountName: ptr.To("service-account"),
|
|
},
|
|
},
|
|
},
|
|
},
|
|
}},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
}).
|
|
Obj(),
|
|
newObj: utiltesting.MakeTrainJobWrapper(metav1.NamespaceDefault, "test").
|
|
Suspend(true).
|
|
RuntimePatches([]trainer.RuntimePatch{
|
|
{
|
|
Manager: "test.io/manager",
|
|
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{
|
|
ServiceAccountName: ptr.To("service-account-updated"),
|
|
},
|
|
},
|
|
},
|
|
},
|
|
}},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
}).
|
|
Obj(),
|
|
clientErr: apierrors.NewNotFound(jobsetv1alpha2.Resource("jobset"), ""),
|
|
wantError: nil,
|
|
},
|
|
"allow atomic update: modify runtimePatches and unsuspend trainJob in a single request": {
|
|
info: &runtime.Info{
|
|
TemplateSpec: runtime.TemplateSpec{
|
|
ObjApply: &jobsetv1alpha2ac.JobSetSpecApplyConfiguration{
|
|
ReplicatedJobs: []jobsetv1alpha2ac.ReplicatedJobApplyConfiguration{
|
|
{
|
|
Name: ptr.To(constants.Node),
|
|
Template: &batchv1ac.JobTemplateSpecApplyConfiguration{
|
|
Spec: &batchv1ac.JobSpecApplyConfiguration{
|
|
Template: &corev1ac.PodTemplateSpecApplyConfiguration{
|
|
Spec: &corev1ac.PodSpecApplyConfiguration{
|
|
Containers: []corev1ac.ContainerApplyConfiguration{
|
|
{
|
|
Name: ptr.To(constants.Node),
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
oldObj: utiltesting.MakeTrainJobWrapper(metav1.NamespaceDefault, "test").
|
|
Suspend(true).
|
|
RuntimePatches([]trainer.RuntimePatch{
|
|
{
|
|
Manager: "test.io/manager",
|
|
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{
|
|
ServiceAccountName: ptr.To("service-account"),
|
|
},
|
|
},
|
|
},
|
|
},
|
|
}},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
}).
|
|
Obj(),
|
|
newObj: utiltesting.MakeTrainJobWrapper(metav1.NamespaceDefault, "test").
|
|
Suspend(false).
|
|
RuntimePatches([]trainer.RuntimePatch{
|
|
{
|
|
Manager: "test.io/manager",
|
|
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{"injected": "by-kueue"},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
}},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
}).
|
|
Obj(),
|
|
clientErr: apierrors.NewNotFound(jobsetv1alpha2.Resource("jobset"), ""),
|
|
wantError: nil,
|
|
},
|
|
"allow atomic update: modify runtimePatches and suspend trainJob in a single request": {
|
|
info: &runtime.Info{
|
|
TemplateSpec: runtime.TemplateSpec{
|
|
ObjApply: &jobsetv1alpha2ac.JobSetSpecApplyConfiguration{
|
|
ReplicatedJobs: []jobsetv1alpha2ac.ReplicatedJobApplyConfiguration{
|
|
{
|
|
Name: ptr.To(constants.Node),
|
|
Template: &batchv1ac.JobTemplateSpecApplyConfiguration{
|
|
Spec: &batchv1ac.JobSpecApplyConfiguration{
|
|
Template: &corev1ac.PodTemplateSpecApplyConfiguration{
|
|
Spec: &corev1ac.PodSpecApplyConfiguration{
|
|
Containers: []corev1ac.ContainerApplyConfiguration{
|
|
{
|
|
Name: ptr.To(constants.Node),
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
oldObj: utiltesting.MakeTrainJobWrapper(metav1.NamespaceDefault, "test").
|
|
Suspend(false).
|
|
RuntimePatches([]trainer.RuntimePatch{
|
|
{
|
|
Manager: "test.io/manager",
|
|
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{
|
|
ServiceAccountName: ptr.To("service-account"),
|
|
},
|
|
},
|
|
},
|
|
},
|
|
}},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
}).
|
|
Obj(),
|
|
newObj: utiltesting.MakeTrainJobWrapper(metav1.NamespaceDefault, "test").
|
|
Suspend(true).
|
|
RuntimePatches([]trainer.RuntimePatch{
|
|
{
|
|
Manager: "test.io/manager",
|
|
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{
|
|
ServiceAccountName: ptr.To("service-account-updated"),
|
|
},
|
|
},
|
|
},
|
|
},
|
|
}},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
}).
|
|
Obj(),
|
|
clientErr: apierrors.NewNotFound(jobsetv1alpha2.Resource("jobset"), ""),
|
|
wantError: nil,
|
|
},
|
|
"allow changes to runtimePatches when trainJob is suspended and jobSet exists but is inactive": {
|
|
info: &runtime.Info{
|
|
TemplateSpec: runtime.TemplateSpec{
|
|
ObjApply: &jobsetv1alpha2ac.JobSetSpecApplyConfiguration{
|
|
ReplicatedJobs: []jobsetv1alpha2ac.ReplicatedJobApplyConfiguration{
|
|
{
|
|
Name: ptr.To(constants.DatasetInitializer),
|
|
Template: &batchv1ac.JobTemplateSpecApplyConfiguration{
|
|
Spec: &batchv1ac.JobSpecApplyConfiguration{
|
|
Template: &corev1ac.PodTemplateSpecApplyConfiguration{
|
|
Spec: &corev1ac.PodSpecApplyConfiguration{
|
|
Containers: []corev1ac.ContainerApplyConfiguration{
|
|
{
|
|
Name: ptr.To(constants.DatasetInitializer),
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
{
|
|
Name: ptr.To(constants.Node),
|
|
Template: &batchv1ac.JobTemplateSpecApplyConfiguration{
|
|
Spec: &batchv1ac.JobSpecApplyConfiguration{
|
|
Template: &corev1ac.PodTemplateSpecApplyConfiguration{
|
|
Spec: &corev1ac.PodSpecApplyConfiguration{
|
|
Containers: []corev1ac.ContainerApplyConfiguration{
|
|
{
|
|
Name: ptr.To(constants.Node),
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
oldObj: utiltesting.MakeTrainJobWrapper(metav1.NamespaceDefault, "test").
|
|
Suspend(true).
|
|
RuntimePatches([]trainer.RuntimePatch{
|
|
{
|
|
Manager: "test.io/manager",
|
|
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{
|
|
ServiceAccountName: ptr.To("service-account"),
|
|
},
|
|
},
|
|
},
|
|
},
|
|
}},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
}).
|
|
Obj(),
|
|
newObj: utiltesting.MakeTrainJobWrapper(metav1.NamespaceDefault, "test").
|
|
Suspend(true).
|
|
RuntimePatches([]trainer.RuntimePatch{
|
|
{
|
|
Manager: "test.io/manager",
|
|
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{
|
|
ServiceAccountName: ptr.To("service-account-updated"),
|
|
},
|
|
},
|
|
},
|
|
},
|
|
}},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
}).
|
|
Obj(),
|
|
jobSet: &jobsetv1alpha2.JobSet{
|
|
ObjectMeta: metav1.ObjectMeta{
|
|
Name: "test",
|
|
Namespace: metav1.NamespaceDefault,
|
|
},
|
|
Status: jobsetv1alpha2.JobSetStatus{
|
|
ReplicatedJobsStatus: []jobsetv1alpha2.ReplicatedJobStatus{
|
|
{
|
|
Name: constants.DatasetInitializer,
|
|
Active: 0,
|
|
},
|
|
{
|
|
Name: constants.Node,
|
|
Active: 0,
|
|
},
|
|
},
|
|
},
|
|
},
|
|
wantError: nil,
|
|
},
|
|
"forbid changes to runtimePatches when trainJob is suspended but jobSet has an active replicatedJob": {
|
|
info: &runtime.Info{
|
|
TemplateSpec: runtime.TemplateSpec{
|
|
ObjApply: &jobsetv1alpha2ac.JobSetSpecApplyConfiguration{
|
|
ReplicatedJobs: []jobsetv1alpha2ac.ReplicatedJobApplyConfiguration{
|
|
{
|
|
Name: ptr.To(constants.DatasetInitializer),
|
|
Template: &batchv1ac.JobTemplateSpecApplyConfiguration{
|
|
Spec: &batchv1ac.JobSpecApplyConfiguration{
|
|
Template: &corev1ac.PodTemplateSpecApplyConfiguration{
|
|
Spec: &corev1ac.PodSpecApplyConfiguration{
|
|
Containers: []corev1ac.ContainerApplyConfiguration{
|
|
{
|
|
Name: ptr.To(constants.DatasetInitializer),
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
{
|
|
Name: ptr.To(constants.Node),
|
|
Template: &batchv1ac.JobTemplateSpecApplyConfiguration{
|
|
Spec: &batchv1ac.JobSpecApplyConfiguration{
|
|
Template: &corev1ac.PodTemplateSpecApplyConfiguration{
|
|
Spec: &corev1ac.PodSpecApplyConfiguration{
|
|
Containers: []corev1ac.ContainerApplyConfiguration{
|
|
{
|
|
Name: ptr.To(constants.Node),
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
oldObj: utiltesting.MakeTrainJobWrapper(metav1.NamespaceDefault, "test").
|
|
Suspend(true).
|
|
RuntimePatches([]trainer.RuntimePatch{
|
|
{
|
|
Manager: "test.io/manager",
|
|
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{
|
|
ServiceAccountName: ptr.To("service-account"),
|
|
},
|
|
},
|
|
},
|
|
},
|
|
}},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
}).
|
|
Obj(),
|
|
newObj: utiltesting.MakeTrainJobWrapper(metav1.NamespaceDefault, "test").
|
|
Suspend(true).
|
|
RuntimePatches([]trainer.RuntimePatch{
|
|
{
|
|
Manager: "test.io/manager",
|
|
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{
|
|
ServiceAccountName: ptr.To("service-account-updated"),
|
|
},
|
|
},
|
|
},
|
|
},
|
|
}},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
}).
|
|
Obj(),
|
|
jobSet: &jobsetv1alpha2.JobSet{
|
|
ObjectMeta: metav1.ObjectMeta{
|
|
Name: "test",
|
|
Namespace: metav1.NamespaceDefault,
|
|
},
|
|
Status: jobsetv1alpha2.JobSetStatus{
|
|
ReplicatedJobsStatus: []jobsetv1alpha2.ReplicatedJobStatus{
|
|
{
|
|
Name: constants.DatasetInitializer,
|
|
Active: 0,
|
|
},
|
|
{
|
|
Name: constants.Node,
|
|
Active: 2,
|
|
},
|
|
},
|
|
},
|
|
},
|
|
wantError: field.ErrorList{
|
|
field.Forbidden(runtimePatchesPath, "RuntimePatches cannot be modified when the JobSet's ReplicatedJob node is still active"),
|
|
},
|
|
},
|
|
"forbid atomic update: suspending trainJob with runtimePatches change is rejected when jobSet has an active replicatedJob": {
|
|
info: &runtime.Info{
|
|
TemplateSpec: runtime.TemplateSpec{
|
|
ObjApply: &jobsetv1alpha2ac.JobSetSpecApplyConfiguration{
|
|
ReplicatedJobs: []jobsetv1alpha2ac.ReplicatedJobApplyConfiguration{
|
|
{
|
|
Name: ptr.To(constants.Node),
|
|
Template: &batchv1ac.JobTemplateSpecApplyConfiguration{
|
|
Spec: &batchv1ac.JobSpecApplyConfiguration{
|
|
Template: &corev1ac.PodTemplateSpecApplyConfiguration{
|
|
Spec: &corev1ac.PodSpecApplyConfiguration{
|
|
Containers: []corev1ac.ContainerApplyConfiguration{
|
|
{
|
|
Name: ptr.To(constants.Node),
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
oldObj: utiltesting.MakeTrainJobWrapper(metav1.NamespaceDefault, "test").
|
|
Suspend(false).
|
|
RuntimePatches([]trainer.RuntimePatch{
|
|
{
|
|
Manager: "test.io/manager",
|
|
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{
|
|
ServiceAccountName: ptr.To("service-account"),
|
|
},
|
|
},
|
|
},
|
|
},
|
|
}},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
}).
|
|
Obj(),
|
|
newObj: utiltesting.MakeTrainJobWrapper(metav1.NamespaceDefault, "test").
|
|
Suspend(true).
|
|
RuntimePatches([]trainer.RuntimePatch{
|
|
{
|
|
Manager: "test.io/manager",
|
|
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{
|
|
ServiceAccountName: ptr.To("service-account-updated"),
|
|
},
|
|
},
|
|
},
|
|
},
|
|
}},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
}).
|
|
Obj(),
|
|
jobSet: &jobsetv1alpha2.JobSet{
|
|
ObjectMeta: metav1.ObjectMeta{
|
|
Name: "test",
|
|
Namespace: metav1.NamespaceDefault,
|
|
},
|
|
Status: jobsetv1alpha2.JobSetStatus{
|
|
ReplicatedJobsStatus: []jobsetv1alpha2.ReplicatedJobStatus{
|
|
{
|
|
Name: constants.Node,
|
|
Active: 2,
|
|
},
|
|
},
|
|
},
|
|
},
|
|
wantError: field.ErrorList{
|
|
field.Forbidden(runtimePatchesPath, "RuntimePatches cannot be modified when the JobSet's ReplicatedJob node is still active"),
|
|
},
|
|
},
|
|
"forbid changes to runtimePatches when trainJob is suspended but has multiple active replicatedJobs": {
|
|
info: &runtime.Info{
|
|
TemplateSpec: runtime.TemplateSpec{
|
|
ObjApply: &jobsetv1alpha2ac.JobSetSpecApplyConfiguration{
|
|
ReplicatedJobs: []jobsetv1alpha2ac.ReplicatedJobApplyConfiguration{
|
|
{
|
|
Name: ptr.To(constants.DatasetInitializer),
|
|
Template: &batchv1ac.JobTemplateSpecApplyConfiguration{
|
|
Spec: &batchv1ac.JobSpecApplyConfiguration{
|
|
Template: &corev1ac.PodTemplateSpecApplyConfiguration{
|
|
Spec: &corev1ac.PodSpecApplyConfiguration{
|
|
Containers: []corev1ac.ContainerApplyConfiguration{
|
|
{
|
|
Name: ptr.To(constants.DatasetInitializer),
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
{
|
|
Name: ptr.To(constants.Node),
|
|
Template: &batchv1ac.JobTemplateSpecApplyConfiguration{
|
|
Spec: &batchv1ac.JobSpecApplyConfiguration{
|
|
Template: &corev1ac.PodTemplateSpecApplyConfiguration{
|
|
Spec: &corev1ac.PodSpecApplyConfiguration{
|
|
Containers: []corev1ac.ContainerApplyConfiguration{
|
|
{
|
|
Name: ptr.To(constants.Node),
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
oldObj: utiltesting.MakeTrainJobWrapper(metav1.NamespaceDefault, "test").
|
|
Suspend(true).
|
|
RuntimePatches([]trainer.RuntimePatch{
|
|
{
|
|
Manager: "test.io/manager",
|
|
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{
|
|
ServiceAccountName: ptr.To("service-account"),
|
|
},
|
|
},
|
|
},
|
|
},
|
|
}},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
}).
|
|
Obj(),
|
|
newObj: utiltesting.MakeTrainJobWrapper(metav1.NamespaceDefault, "test").
|
|
Suspend(true).
|
|
RuntimePatches([]trainer.RuntimePatch{
|
|
{
|
|
Manager: "test.io/manager",
|
|
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{
|
|
ServiceAccountName: ptr.To("service-account-updated"),
|
|
},
|
|
},
|
|
},
|
|
},
|
|
}},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
}).
|
|
Obj(),
|
|
jobSet: &jobsetv1alpha2.JobSet{
|
|
ObjectMeta: metav1.ObjectMeta{
|
|
Name: "test",
|
|
Namespace: metav1.NamespaceDefault,
|
|
},
|
|
Status: jobsetv1alpha2.JobSetStatus{
|
|
ReplicatedJobsStatus: []jobsetv1alpha2.ReplicatedJobStatus{
|
|
{
|
|
Name: constants.DatasetInitializer,
|
|
Active: 1,
|
|
},
|
|
{
|
|
Name: constants.Node,
|
|
Active: 2,
|
|
},
|
|
},
|
|
},
|
|
},
|
|
wantError: field.ErrorList{
|
|
field.Forbidden(runtimePatchesPath, "RuntimePatches cannot be modified when the JobSet's ReplicatedJob dataset-initializer is still active"),
|
|
field.Forbidden(runtimePatchesPath, "RuntimePatches cannot be modified when the JobSet's ReplicatedJob node is still active"),
|
|
},
|
|
},
|
|
"forbid changes to runtimePatches when trainJob is suspended but jobSet cannot be checked due to a client error": {
|
|
info: &runtime.Info{
|
|
TemplateSpec: runtime.TemplateSpec{
|
|
ObjApply: &jobsetv1alpha2ac.JobSetSpecApplyConfiguration{
|
|
ReplicatedJobs: []jobsetv1alpha2ac.ReplicatedJobApplyConfiguration{
|
|
{
|
|
Name: ptr.To(constants.Node),
|
|
Template: &batchv1ac.JobTemplateSpecApplyConfiguration{
|
|
Spec: &batchv1ac.JobSpecApplyConfiguration{
|
|
Template: &corev1ac.PodTemplateSpecApplyConfiguration{
|
|
Spec: &corev1ac.PodSpecApplyConfiguration{
|
|
Containers: []corev1ac.ContainerApplyConfiguration{
|
|
{
|
|
Name: ptr.To(constants.Node),
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
oldObj: utiltesting.MakeTrainJobWrapper(metav1.NamespaceDefault, "test").
|
|
Suspend(true).
|
|
RuntimePatches([]trainer.RuntimePatch{
|
|
{
|
|
Manager: "test.io/manager",
|
|
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{
|
|
ServiceAccountName: ptr.To("service-account"),
|
|
},
|
|
},
|
|
},
|
|
},
|
|
}},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
}).
|
|
Obj(),
|
|
newObj: utiltesting.MakeTrainJobWrapper(metav1.NamespaceDefault, "test").
|
|
Suspend(true).
|
|
RuntimePatches([]trainer.RuntimePatch{
|
|
{
|
|
Manager: "test.io/manager",
|
|
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{
|
|
ServiceAccountName: ptr.To("service-account-updated"),
|
|
},
|
|
},
|
|
},
|
|
},
|
|
}},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
}).
|
|
Obj(),
|
|
clientErr: fmt.Errorf("client error"),
|
|
wantError: field.ErrorList{
|
|
field.InternalError(runtimePatchesPath, fmt.Errorf("client error")),
|
|
},
|
|
},
|
|
}
|
|
for name, tc := range cases {
|
|
t.Run(name, func(t *testing.T) {
|
|
_, ctx := ktesting.NewTestContext(t)
|
|
var cancel func()
|
|
ctx, cancel = context.WithCancel(ctx)
|
|
t.Cleanup(cancel)
|
|
|
|
clientBuilder := utiltesting.NewClientBuilder()
|
|
if tc.jobSet != nil {
|
|
clientBuilder = clientBuilder.WithObjects(tc.jobSet)
|
|
}
|
|
if tc.clientErr != nil {
|
|
clientBuilder = clientBuilder.WithInterceptorFuncs(interceptor.Funcs{
|
|
Get: func(ctx context.Context, cli client.WithWatch, key client.ObjectKey, obj client.Object, opts ...client.GetOption) error {
|
|
if _, ok := obj.(*jobsetv1alpha2.JobSet); ok {
|
|
return tc.clientErr
|
|
}
|
|
return cli.Get(ctx, key, obj, opts...)
|
|
},
|
|
})
|
|
}
|
|
cli := clientBuilder.Build()
|
|
|
|
p, err := New(ctx, cli, nil, nil)
|
|
if err != nil {
|
|
t.Fatalf("Failed to initialize JobSet plugin: %v", err)
|
|
}
|
|
|
|
warnings, errs := p.(framework.CustomValidationPlugin).Validate(ctx, tc.info, tc.oldObj, tc.newObj)
|
|
if diff := cmp.Diff(tc.wantError, errs, cmpopts.IgnoreFields(field.Error{}, "BadValue")); len(diff) != 0 {
|
|
t.Errorf("Unexpected error from Validate (-want, +got): %s", diff)
|
|
}
|
|
if diff := cmp.Diff(tc.wantWarnings, warnings); len(diff) != 0 {
|
|
t.Errorf("Unexpected warnings from Validate (-want, +got): %s", diff)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestBuild(t *testing.T) {
|
|
cases := map[string]struct {
|
|
info *runtime.Info
|
|
trainJob *trainer.TrainJob
|
|
wantObjs []apiruntime.Object
|
|
wantError error
|
|
}{
|
|
"init containers synced to JobSet replicated jobs": {
|
|
info: &runtime.Info{
|
|
Labels: make(map[string]string),
|
|
Annotations: make(map[string]string),
|
|
TemplateSpec: runtime.TemplateSpec{
|
|
PodSets: []runtime.PodSet{
|
|
{
|
|
Name: constants.Node,
|
|
Count: ptr.To[int32](2),
|
|
InitContainers: []runtime.Container{
|
|
{
|
|
Name: "preflight-check",
|
|
Image: "preflight:latest",
|
|
Command: []string{"/bin/sh", "-c"},
|
|
Env: []corev1ac.EnvVarApplyConfiguration{
|
|
{Name: ptr.To("PET_NNODES"), Value: ptr.To("2")},
|
|
{Name: ptr.To("PET_MASTER_ADDR"), Value: ptr.To("test-job-node-0-0.test-job")},
|
|
},
|
|
},
|
|
},
|
|
Containers: []runtime.Container{
|
|
{Name: constants.Node, Image: "pytorch:latest"},
|
|
},
|
|
},
|
|
},
|
|
ObjApply: jobsetv1alpha2ac.JobSetSpec().
|
|
WithReplicatedJobs(
|
|
jobsetv1alpha2ac.ReplicatedJob().
|
|
WithName(constants.Node).
|
|
WithTemplate(batchv1ac.JobTemplateSpec().
|
|
WithSpec(batchv1ac.JobSpec().
|
|
WithTemplate(corev1ac.PodTemplateSpec().
|
|
WithSpec(corev1ac.PodSpec().
|
|
WithInitContainers(
|
|
corev1ac.Container().WithName("preflight-check"),
|
|
).
|
|
WithContainers(
|
|
corev1ac.Container().WithName(constants.Node),
|
|
),
|
|
),
|
|
),
|
|
),
|
|
),
|
|
),
|
|
},
|
|
Scheduler: &runtime.Scheduler{PodLabels: make(map[string]string)},
|
|
},
|
|
trainJob: utiltesting.MakeTrainJobWrapper(metav1.NamespaceDefault, "test-job").
|
|
Obj(),
|
|
wantObjs: []apiruntime.Object{
|
|
&jobsetv1alpha2.JobSet{
|
|
ObjectMeta: metav1.ObjectMeta{
|
|
Name: "test-job",
|
|
Namespace: metav1.NamespaceDefault,
|
|
OwnerReferences: []metav1.OwnerReference{
|
|
{APIVersion: trainer.GroupVersion.String(), Kind: trainer.TrainJobKind, Name: "test-job", Controller: ptr.To(true)},
|
|
},
|
|
},
|
|
Spec: jobsetv1alpha2.JobSetSpec{
|
|
ReplicatedJobs: []jobsetv1alpha2.ReplicatedJob{
|
|
{
|
|
Name: constants.Node,
|
|
Replicas: 0,
|
|
Template: batchv1.JobTemplateSpec{
|
|
Spec: batchv1.JobSpec{
|
|
Parallelism: ptr.To[int32](2),
|
|
Completions: ptr.To[int32](2),
|
|
Template: corev1.PodTemplateSpec{
|
|
Spec: corev1.PodSpec{
|
|
InitContainers: []corev1.Container{
|
|
{
|
|
Name: "preflight-check",
|
|
Image: "preflight:latest",
|
|
Command: []string{"/bin/sh", "-c"},
|
|
Env: []corev1.EnvVar{
|
|
{Name: "PET_NNODES", Value: "2"},
|
|
{Name: "PET_MASTER_ADDR", Value: "test-job-node-0-0.test-job"},
|
|
},
|
|
},
|
|
},
|
|
Containers: []corev1.Container{
|
|
{Name: constants.Node, Image: "pytorch:latest"},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
"init container with ports and volume mounts synced": {
|
|
info: &runtime.Info{
|
|
Labels: make(map[string]string),
|
|
Annotations: make(map[string]string),
|
|
TemplateSpec: runtime.TemplateSpec{
|
|
PodSets: []runtime.PodSet{
|
|
{
|
|
Name: constants.Node,
|
|
Count: ptr.To[int32](1),
|
|
InitContainers: []runtime.Container{
|
|
{
|
|
Name: "nccl-check",
|
|
Image: "nccl-test:latest",
|
|
Command: []string{"/nccl-test"},
|
|
Ports: []corev1ac.ContainerPortApplyConfiguration{
|
|
{ContainerPort: ptr.To[int32](8080), Name: ptr.To("http")},
|
|
},
|
|
VolumeMounts: []corev1ac.VolumeMountApplyConfiguration{
|
|
{Name: ptr.To("mount"), MountPath: ptr.To("/data")},
|
|
},
|
|
},
|
|
},
|
|
Containers: []runtime.Container{
|
|
{Name: constants.Node, Image: "train:latest"},
|
|
},
|
|
},
|
|
},
|
|
ObjApply: jobsetv1alpha2ac.JobSetSpec().
|
|
WithReplicatedJobs(
|
|
jobsetv1alpha2ac.ReplicatedJob().
|
|
WithName(constants.Node).
|
|
WithTemplate(batchv1ac.JobTemplateSpec().
|
|
WithSpec(batchv1ac.JobSpec().
|
|
WithTemplate(corev1ac.PodTemplateSpec().
|
|
WithSpec(corev1ac.PodSpec().
|
|
WithInitContainers(
|
|
corev1ac.Container().WithName("nccl-check"),
|
|
).
|
|
WithContainers(
|
|
corev1ac.Container().WithName(constants.Node),
|
|
),
|
|
),
|
|
),
|
|
),
|
|
),
|
|
),
|
|
},
|
|
Scheduler: &runtime.Scheduler{PodLabels: make(map[string]string)},
|
|
},
|
|
trainJob: utiltesting.MakeTrainJobWrapper(metav1.NamespaceDefault, "test-job").
|
|
Obj(),
|
|
wantObjs: []apiruntime.Object{
|
|
&jobsetv1alpha2.JobSet{
|
|
ObjectMeta: metav1.ObjectMeta{
|
|
Name: "test-job",
|
|
Namespace: metav1.NamespaceDefault,
|
|
OwnerReferences: []metav1.OwnerReference{
|
|
{APIVersion: trainer.GroupVersion.String(), Kind: trainer.TrainJobKind, Name: "test-job", Controller: ptr.To(true)},
|
|
},
|
|
},
|
|
Spec: jobsetv1alpha2.JobSetSpec{
|
|
ReplicatedJobs: []jobsetv1alpha2.ReplicatedJob{
|
|
{
|
|
Name: constants.Node,
|
|
Replicas: 0,
|
|
Template: batchv1.JobTemplateSpec{
|
|
Spec: batchv1.JobSpec{
|
|
Parallelism: ptr.To[int32](1),
|
|
Completions: ptr.To[int32](1),
|
|
Template: corev1.PodTemplateSpec{
|
|
Spec: corev1.PodSpec{
|
|
InitContainers: []corev1.Container{
|
|
{
|
|
Name: "nccl-check",
|
|
Image: "nccl-test:latest",
|
|
Command: []string{"/nccl-test"},
|
|
Ports: []corev1.ContainerPort{{Name: "http", ContainerPort: 8080}},
|
|
VolumeMounts: []corev1.VolumeMount{
|
|
{Name: "mount", MountPath: "/data"},
|
|
},
|
|
},
|
|
},
|
|
Containers: []corev1.Container{
|
|
{Name: constants.Node, Image: "train:latest"},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
"multiple init containers synced": {
|
|
info: &runtime.Info{
|
|
Labels: make(map[string]string),
|
|
Annotations: make(map[string]string),
|
|
TemplateSpec: runtime.TemplateSpec{
|
|
PodSets: []runtime.PodSet{
|
|
{
|
|
Name: constants.Node,
|
|
Count: ptr.To[int32](3),
|
|
InitContainers: []runtime.Container{
|
|
{Name: "driver-check", Image: "driver:latest", Command: []string{"/check-driver"}},
|
|
{Name: "nccl-check", Image: "nccl:latest", Command: []string{"/check-nccl"}},
|
|
},
|
|
Containers: []runtime.Container{
|
|
{Name: constants.Node, Image: "train:latest"},
|
|
},
|
|
},
|
|
},
|
|
ObjApply: jobsetv1alpha2ac.JobSetSpec().
|
|
WithReplicatedJobs(
|
|
jobsetv1alpha2ac.ReplicatedJob().
|
|
WithName(constants.Node).
|
|
WithTemplate(batchv1ac.JobTemplateSpec().
|
|
WithSpec(batchv1ac.JobSpec().
|
|
WithTemplate(corev1ac.PodTemplateSpec().
|
|
WithSpec(corev1ac.PodSpec().
|
|
WithInitContainers(
|
|
corev1ac.Container().WithName("driver-check"),
|
|
corev1ac.Container().WithName("nccl-check"),
|
|
).
|
|
WithContainers(
|
|
corev1ac.Container().WithName(constants.Node),
|
|
),
|
|
),
|
|
),
|
|
),
|
|
),
|
|
),
|
|
},
|
|
Scheduler: &runtime.Scheduler{PodLabels: make(map[string]string)},
|
|
},
|
|
trainJob: utiltesting.MakeTrainJobWrapper(metav1.NamespaceDefault, "test-job").
|
|
Obj(),
|
|
wantObjs: []apiruntime.Object{
|
|
&jobsetv1alpha2.JobSet{
|
|
ObjectMeta: metav1.ObjectMeta{
|
|
Name: "test-job",
|
|
Namespace: metav1.NamespaceDefault,
|
|
OwnerReferences: []metav1.OwnerReference{
|
|
{APIVersion: trainer.GroupVersion.String(), Kind: trainer.TrainJobKind, Name: "test-job", Controller: ptr.To(true)},
|
|
},
|
|
},
|
|
Spec: jobsetv1alpha2.JobSetSpec{
|
|
ReplicatedJobs: []jobsetv1alpha2.ReplicatedJob{
|
|
{
|
|
Name: constants.Node,
|
|
Replicas: 0,
|
|
Template: batchv1.JobTemplateSpec{
|
|
Spec: batchv1.JobSpec{
|
|
Parallelism: ptr.To[int32](3),
|
|
Completions: ptr.To[int32](3),
|
|
Template: corev1.PodTemplateSpec{
|
|
Spec: corev1.PodSpec{
|
|
InitContainers: []corev1.Container{
|
|
{Name: "driver-check", Image: "driver:latest", Command: []string{"/check-driver"}},
|
|
{Name: "nccl-check", Image: "nccl:latest", Command: []string{"/check-nccl"}},
|
|
},
|
|
Containers: []corev1.Container{
|
|
{Name: constants.Node, Image: "train:latest"},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
"return error when podSet initContainers are missing from apply configuration": {
|
|
info: &runtime.Info{
|
|
TemplateSpec: runtime.TemplateSpec{
|
|
PodSets: []runtime.PodSet{{
|
|
Name: constants.Node,
|
|
InitContainers: []runtime.Container{{Name: "preflight-check"}},
|
|
}},
|
|
ObjApply: jobsetv1alpha2ac.JobSetSpec().
|
|
WithReplicatedJobs(jobsetv1alpha2ac.ReplicatedJob().
|
|
WithName(constants.Node).
|
|
WithTemplate(batchv1ac.JobTemplateSpec().
|
|
WithSpec(batchv1ac.JobSpec().
|
|
WithTemplate(corev1ac.PodTemplateSpec().
|
|
WithSpec(corev1ac.PodSpec().
|
|
WithContainers(corev1ac.Container().WithName(constants.Node)),
|
|
),
|
|
),
|
|
),
|
|
),
|
|
),
|
|
},
|
|
Scheduler: &runtime.Scheduler{PodLabels: make(map[string]string)},
|
|
},
|
|
trainJob: utiltesting.MakeTrainJobWrapper(metav1.NamespaceDefault, "trainJob").
|
|
Obj(),
|
|
wantError: fmt.Errorf("podSet %q initContainer %q does not have a matching initContainer in the runtime template", constants.Node, "preflight-check"),
|
|
},
|
|
}
|
|
for name, tc := range cases {
|
|
t.Run(name, func(t *testing.T) {
|
|
_, ctx := ktesting.NewTestContext(t)
|
|
var cancel func()
|
|
ctx, cancel = context.WithCancel(ctx)
|
|
t.Cleanup(cancel)
|
|
cli := utiltesting.NewClientBuilder().Build()
|
|
p, err := New(ctx, cli, nil, nil)
|
|
if err != nil {
|
|
t.Fatalf("Failed to initialize JobSet plugin: %v", err)
|
|
}
|
|
objs, err := p.(framework.ComponentBuilderPlugin).Build(ctx, tc.info, tc.trainJob)
|
|
if diff := cmp.Diff(tc.wantError, err, cmp.Comparer(func(x, y error) bool {
|
|
if x == nil || y == nil {
|
|
return x == y
|
|
}
|
|
return x.Error() == y.Error()
|
|
})); len(diff) != 0 {
|
|
t.Errorf("Unexpected error from Build (-want,+got):\n%s", diff)
|
|
}
|
|
if tc.wantError != nil {
|
|
return
|
|
}
|
|
typedObjs, err := utiltesting.ToObject(cli.Scheme(), objs...)
|
|
if err != nil {
|
|
t.Fatalf("Failed to convert objects: %v", err)
|
|
}
|
|
if diff := cmp.Diff(tc.wantObjs, typedObjs,
|
|
cmpopts.SortSlices(func(a, b string) bool { return a < b }),
|
|
cmpopts.SortMaps(func(a, b string) bool { return a < b }),
|
|
cmp.Transformer("", func(in jobsetv1alpha2.JobSet) jobsetv1alpha2.JobSet {
|
|
in.Kind = ""
|
|
in.APIVersion = ""
|
|
for i := range in.OwnerReferences {
|
|
in.OwnerReferences[i].BlockOwnerDeletion = nil
|
|
}
|
|
return in
|
|
}),
|
|
); len(diff) != 0 {
|
|
t.Errorf("Unexpected objects from Build (-want, +got): %s", diff)
|
|
}
|
|
})
|
|
}
|
|
}
|