kubeflow
/
trainer
mirror of https://github.com/kubeflow/trainer.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity
trainer/pkg/runtime/core/trainingruntime_test.go

806 lines
29 KiB
Go
Raw Permalink Blame History

/*
Copyright 2024 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 core
import (
"context"
"fmt"
"testing"
"github.com/google/go-cmp/cmp"
"github.com/google/go-cmp/cmp/cmpopts"
corev1 "k8s.io/api/core/v1"
"k8s.io/apimachinery/pkg/api/resource"
metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
"k8s.io/apimachinery/pkg/runtime"
"k8s.io/apimachinery/pkg/util/intstr"
"k8s.io/klog/v2/ktesting"
"k8s.io/utils/ptr"
jobsetv1alpha2 "sigs.k8s.io/jobset/api/jobset/v1alpha2"
schedulerpluginsv1alpha1 "sigs.k8s.io/scheduler-plugins/apis/scheduling/v1alpha1"
trainer "github.com/kubeflow/trainer/pkg/apis/trainer/v1alpha1"
"github.com/kubeflow/trainer/pkg/constants"
jobsetplgconsts "github.com/kubeflow/trainer/pkg/runtime/framework/plugins/jobset/constants"
testingutil "github.com/kubeflow/trainer/pkg/util/testing"
)
func TestTrainingRuntimeNewObjects(t *testing.T) {
resRequests := corev1.ResourceList{
corev1.ResourceCPU: resource.MustParse("1"),
}
// TODO (andreyvelich): Add more test cases.
cases := map[string]struct {
trainingRuntime *trainer.TrainingRuntime
trainJob *trainer.TrainJob
ObjCmpOpts []cmp.Option
wantObjs []runtime.Object
wantError error
}{
// Test cases for the PlainML MLPolicy.
"succeeded to build PodGroup and JobSet with NumNodes from the TrainJob and container from the Runtime.": {
trainingRuntime: testingutil.MakeTrainingRuntimeWrapper(metav1.NamespaceDefault, "test-runtime").
Label("conflictLabel", "overridden").
Annotation("conflictAnnotation", "overridden").
RuntimeSpec(
testingutil.MakeTrainingRuntimeSpecWrapper(testingutil.MakeTrainingRuntimeWrapper(metav1.NamespaceDefault, "test-runtime").Spec).
WithMLPolicy(
testingutil.MakeMLPolicyWrapper().
WithNumNodes(100).
Obj(),
).
PodGroupPolicyCoschedulingSchedulingTimeout(120).
Container(constants.DatasetInitializer, constants.DatasetInitializer, "test:runtime", []string{"runtime"}, []string{"runtime"}, resRequests).
Container(constants.ModelInitializer, constants.ModelInitializer, "test:runtime", []string{"runtime"}, []string{"runtime"}, resRequests).
Container(constants.Node, constants.Node, "test:runtime", []string{"runtime"}, []string{"runtime"}, resRequests).
Obj(),
).Obj(),
trainJob: testingutil.MakeTrainJobWrapper(metav1.NamespaceDefault, "test-job").
Suspend(true).
UID("uid").
RuntimeRef(trainer.SchemeGroupVersion.WithKind(trainer.TrainingRuntimeKind), "test-runtime").
SpecLabel("conflictLabel", "override").
SpecAnnotation("conflictAnnotation", "override").
Trainer(
testingutil.MakeTrainJobTrainerWrapper().
NumNodes(30).
Obj(),
).
Obj(),
wantObjs: []runtime.Object{
testingutil.MakeJobSetWrapper(metav1.NamespaceDefault, "test-job").
ControllerReference(trainer.SchemeGroupVersion.WithKind(trainer.TrainJobKind), "test-job", "uid").
Suspend(true).
Label("conflictLabel", "override").
Annotation("conflictAnnotation", "override").
PodLabel(schedulerpluginsv1alpha1.PodGroupLabel, "test-job").
Replicas(1, constants.DatasetInitializer, constants.ModelInitializer, constants.Node, constants.Launcher).
Parallelism(1, constants.DatasetInitializer, constants.ModelInitializer, constants.Launcher).
Completions(1, constants.DatasetInitializer, constants.ModelInitializer, constants.Launcher).
NumNodes(30).
Container(constants.DatasetInitializer, constants.DatasetInitializer, "test:runtime", []string{"runtime"}, []string{"runtime"}, resRequests).
Container(constants.ModelInitializer, constants.ModelInitializer, "test:runtime", []string{"runtime"}, []string{"runtime"}, resRequests).
Container(constants.Node, constants.Node, "test:runtime", []string{"runtime"}, []string{"runtime"}, resRequests).
Obj(),
testingutil.MakeSchedulerPluginsPodGroup(metav1.NamespaceDefault, "test-job").
ControllerReference(trainer.SchemeGroupVersion.WithKind(trainer.TrainJobKind), "test-job", "uid").
MinMember(32). // 32 replicas = 30 Trainer nodes + 2 Initializer.
MinResources(corev1.ResourceList{
// Trainer node has 30 CPUs + 2 CPUs from 2 initializer containers.
corev1.ResourceCPU: resource.MustParse("32"),
}).
SchedulingTimeout(120).
Obj(),
},
},
"succeeded to build JobSet with NumNodes from the Runtime and container from the TrainJob.": {
trainingRuntime: testingutil.MakeTrainingRuntimeWrapper(metav1.NamespaceDefault, "test-runtime").RuntimeSpec(
testingutil.MakeTrainingRuntimeSpecWrapper(testingutil.MakeTrainingRuntimeWrapper(metav1.NamespaceDefault, "test-runtime").Spec).
WithMLPolicy(
testingutil.MakeMLPolicyWrapper().
WithNumNodes(100).
Obj(),
).
Container(constants.Node, constants.Node, "test:runtime", []string{"runtime"}, []string{"runtime"}, resRequests).
Env(constants.Node, constants.Node,
[]corev1.EnvVar{
{
Name: "TRAIN_JOB",
Value: "original",
},
{
Name: "RUNTIME",
Value: "test:runtime",
},
}...,
).
Obj(),
).Obj(),
trainJob: testingutil.MakeTrainJobWrapper(metav1.NamespaceDefault, "test-job").
UID("uid").
RuntimeRef(trainer.SchemeGroupVersion.WithKind(trainer.TrainingRuntimeKind), "test-runtime").
Trainer(
testingutil.MakeTrainJobTrainerWrapper().
Container("test:trainjob", []string{"trainjob"}, []string{"trainjob"}, resRequests).
Env(
[]corev1.EnvVar{
{
Name: "TRAIN_JOB",
Value: "override",
},
{
Name: "TRAIN_JOB_CUSTOM",
Value: "test:trainjob",
},
}...,
).
Obj(),
).
Obj(),
wantObjs: []runtime.Object{
testingutil.MakeJobSetWrapper(metav1.NamespaceDefault, "test-job").
ControllerReference(trainer.SchemeGroupVersion.WithKind(trainer.TrainJobKind), "test-job", "uid").
Replicas(1, constants.DatasetInitializer, constants.ModelInitializer, constants.Node).
Parallelism(1, constants.DatasetInitializer, constants.ModelInitializer).
Completions(1, constants.DatasetInitializer, constants.ModelInitializer).
NumNodes(100).
Container(constants.Node, constants.Node, "test:trainjob", []string{"trainjob"}, []string{"trainjob"}, resRequests).
Env(constants.Node, constants.Node,
[]corev1.EnvVar{
{
Name: "TRAIN_JOB",
Value: "override",
},
{
Name: "TRAIN_JOB_CUSTOM",
Value: "test:trainjob",
},
{
Name: "RUNTIME",
Value: "test:runtime",
},
}...,
).
Obj(),
},
},
"succeeded to build JobSet with dataset and model initializer from the TrainJob.": {
trainingRuntime: testingutil.MakeTrainingRuntimeWrapper(metav1.NamespaceDefault, "test-runtime").RuntimeSpec(
testingutil.MakeTrainingRuntimeSpecWrapper(testingutil.MakeTrainingRuntimeWrapper(metav1.NamespaceDefault, "test-runtime").Spec).
WithMLPolicy(
testingutil.MakeMLPolicyWrapper().
WithNumNodes(100).
Obj(),
).
Container(constants.DatasetInitializer, constants.DatasetInitializer, "test:runtime", []string{"runtime"}, []string{"runtime"}, resRequests).
Container(constants.ModelInitializer, constants.DatasetInitializer, "test:runtime", []string{"runtime"}, []string{"runtime"}, resRequests).
Container(constants.Node, constants.Node, "test:runtime", []string{"runtime"}, []string{"runtime"}, resRequests).
Obj(),
).Obj(),
trainJob: testingutil.MakeTrainJobWrapper(metav1.NamespaceDefault, "test-job").
UID("uid").
RuntimeRef(trainer.SchemeGroupVersion.WithKind(trainer.TrainingRuntimeKind), "test-runtime").
Initializer(
testingutil.MakeTrainJobInitializerWrapper().
DatasetInitializer(
testingutil.MakeTrainJobDatasetInitializerWrapper().
StorageUri("hf://trainjob-dataset").
Env(
[]corev1.EnvVar{
{
Name: "TRAIN_JOB",
Value: "test:trainjob:dataset",
},
}...,
).
SecretRef(corev1.LocalObjectReference{Name: "trainjob-secret-dataset"}).
Obj(),
).
ModelInitializer(
testingutil.MakeTrainJobModelInitializerWrapper().
StorageUri("hf://trainjob-model").
Env(
[]corev1.EnvVar{
{
Name: "TRAIN_JOB",
Value: "test:trainjob:model",
},
}...,
).
SecretRef(corev1.LocalObjectReference{Name: "trainjob-secret-model"}).
Obj(),
).
Obj(),
).
Trainer(
testingutil.MakeTrainJobTrainerWrapper().
Obj(),
).
Obj(),
wantObjs: []runtime.Object{
testingutil.MakeJobSetWrapper(metav1.NamespaceDefault, "test-job").
ControllerReference(trainer.SchemeGroupVersion.WithKind(trainer.TrainJobKind), "test-job", "uid").
Replicas(1, constants.DatasetInitializer, constants.ModelInitializer, constants.Node).
Parallelism(1, constants.DatasetInitializer, constants.ModelInitializer).
Completions(1, constants.DatasetInitializer, constants.ModelInitializer).
NumNodes(100).
Container(constants.DatasetInitializer, constants.DatasetInitializer, "test:runtime", []string{"runtime"}, []string{"runtime"}, resRequests).
Container(constants.ModelInitializer, constants.DatasetInitializer, "test:runtime", []string{"runtime"}, []string{"runtime"}, resRequests).
Container(constants.Node, constants.Node, "test:runtime", []string{"runtime"}, []string{"runtime"}, resRequests).
Env(constants.DatasetInitializer, constants.DatasetInitializer,
[]corev1.EnvVar{
{
Name: jobsetplgconsts.InitializerEnvStorageUri,
Value: "hf://trainjob-dataset",
},
{
Name: "TRAIN_JOB",
Value: "test:trainjob:dataset",
},
}...,
).
EnvFrom(constants.DatasetInitializer, constants.DatasetInitializer,
[]corev1.EnvFromSource{
{
SecretRef: &corev1.SecretEnvSource{
LocalObjectReference: corev1.LocalObjectReference{
Name: "trainjob-secret-dataset",
},
},
},
}...,
).
Env(constants.ModelInitializer, constants.ModelInitializer,
[]corev1.EnvVar{
{
Name: jobsetplgconsts.InitializerEnvStorageUri,
Value: "hf://trainjob-model",
},
{
Name: "TRAIN_JOB",
Value: "test:trainjob:model",
},
}...,
).
EnvFrom(constants.ModelInitializer, constants.ModelInitializer,
[]corev1.EnvFromSource{
{
SecretRef: &corev1.SecretEnvSource{
LocalObjectReference: corev1.LocalObjectReference{
Name: "trainjob-secret-model",
},
},
},
}...,
).
Obj(),
},
},
// Test cases for the Torch MLPolicy.
"succeeded to build JobSet with Torch values from the TrainJob": {
trainingRuntime: testingutil.MakeTrainingRuntimeWrapper(metav1.NamespaceDefault, "test-runtime").RuntimeSpec(
testingutil.MakeTrainingRuntimeSpecWrapper(testingutil.MakeTrainingRuntimeWrapper(metav1.NamespaceDefault, "test-runtime").Spec).
WithMLPolicy(
testingutil.MakeMLPolicyWrapper().
WithNumNodes(100).
WithMLPolicySource(*testingutil.MakeMLPolicySourceWrapper().
TorchPolicy(ptr.To(intstr.FromString("auto")), nil).
Obj(),
).
Obj(),
).
JobSetSpec(
testingutil.MakeJobSetWrapper("", "").
DependsOn(constants.Node,
[]jobsetv1alpha2.DependsOn{
{
Name: constants.DatasetInitializer,
Status: jobsetv1alpha2.DependencyComplete,
},
{
Name: constants.ModelInitializer,
Status: jobsetv1alpha2.DependencyComplete,
},
}...,
).
Obj().
Spec,
).
Container(constants.Node, constants.Node, "test:runtime", []string{"runtime"}, []string{"runtime"}, resRequests).
Obj(),
).Obj(),
trainJob: testingutil.MakeTrainJobWrapper(metav1.NamespaceDefault, "test-job").
UID("uid").
RuntimeRef(trainer.SchemeGroupVersion.WithKind(trainer.TrainingRuntimeKind), "test-runtime").
Trainer(
testingutil.MakeTrainJobTrainerWrapper().
NumNodes(30).
NumProcPerNode(intstr.FromInt32(3)).
Obj(),
).
Obj(),
wantObjs: []runtime.Object{
testingutil.MakeJobSetWrapper(metav1.NamespaceDefault, "test-job").
ControllerReference(trainer.SchemeGroupVersion.WithKind(trainer.TrainJobKind), "test-job", "uid").
Replicas(1, constants.DatasetInitializer, constants.ModelInitializer, constants.Node, constants.Launcher).
Parallelism(1, constants.DatasetInitializer, constants.ModelInitializer).
Completions(1, constants.DatasetInitializer, constants.ModelInitializer).
NumNodes(30).
Container(constants.Node, constants.Node, "test:runtime", []string{"runtime"}, []string{"runtime"}, resRequests).
ContainerTrainerPorts([]corev1.ContainerPort{{ContainerPort: constants.ContainerTrainerPort}}).
Env(constants.Node, constants.Node,
[]corev1.EnvVar{
{
Name: constants.TorchEnvNumNodes,
Value: "30",
},
{
Name: constants.TorchEnvNumProcPerNode,
Value: "3",
},
{
Name: constants.TorchEnvNodeRank,
ValueFrom: &corev1.EnvVarSource{
FieldRef: &corev1.ObjectFieldSelector{
FieldPath: constants.JobCompletionIndexFieldPath,
},
},
},
{
Name: constants.TorchEnvMasterAddr,
Value: fmt.Sprintf("test-job-%s-0-0.test-job", constants.Node),
},
{
Name: constants.TorchEnvMasterPort,
Value: fmt.Sprintf("%d", constants.ContainerTrainerPort),
},
}...,
).
DependsOn(constants.Node,
[]jobsetv1alpha2.DependsOn{
{
Name: constants.DatasetInitializer,
Status: jobsetv1alpha2.DependencyComplete,
},
{
Name: constants.ModelInitializer,
Status: jobsetv1alpha2.DependencyComplete,
},
}...,
).
Obj(),
},
},
"succeeded to build JobSet with Torch values from the Runtime and envs.": {
trainingRuntime: testingutil.MakeTrainingRuntimeWrapper(metav1.NamespaceDefault, "test-runtime").RuntimeSpec(
testingutil.MakeTrainingRuntimeSpecWrapper(testingutil.MakeTrainingRuntimeWrapper(metav1.NamespaceDefault, "test-runtime").Spec).
WithMLPolicy(
testingutil.MakeMLPolicyWrapper().
WithNumNodes(100).
WithMLPolicySource(*testingutil.MakeMLPolicySourceWrapper().
TorchPolicy(ptr.To(intstr.FromString("auto")), nil).
Obj(),
).
Obj(),
).
Container(constants.Node, constants.Node, "test:runtime", []string{"runtime"}, []string{"runtime"}, resRequests).
Env(constants.Node, constants.Node,
[]corev1.EnvVar{
{
Name: "TRAIN_JOB",
Value: "original",
},
{
Name: "RUNTIME",
Value: "test:runtime",
},
}...,
).
Obj(),
).Obj(),
trainJob: testingutil.MakeTrainJobWrapper(metav1.NamespaceDefault, "test-job").
UID("uid").
RuntimeRef(trainer.SchemeGroupVersion.WithKind(trainer.TrainingRuntimeKind), "test-runtime").
Trainer(
testingutil.MakeTrainJobTrainerWrapper().
Container("test:trainjob", []string{"trainjob"}, []string{"trainjob"}, resRequests).
Env(
[]corev1.EnvVar{
{
Name: "TRAIN_JOB",
Value: "override",
},
{
Name: "TRAIN_JOB_CUSTOM",
Value: "test:trainjob",
},
}...,
).
Obj(),
).
Obj(),
wantObjs: []runtime.Object{
testingutil.MakeJobSetWrapper(metav1.NamespaceDefault, "test-job").
ControllerReference(trainer.SchemeGroupVersion.WithKind(trainer.TrainJobKind), "test-job", "uid").
Replicas(1, constants.DatasetInitializer, constants.ModelInitializer, constants.Node, constants.Launcher).
Parallelism(1, constants.DatasetInitializer, constants.ModelInitializer).
Completions(1, constants.DatasetInitializer, constants.ModelInitializer).
NumNodes(100).
Container(constants.Node, constants.Node, "test:trainjob", []string{"trainjob"}, []string{"trainjob"}, resRequests).
ContainerTrainerPorts([]corev1.ContainerPort{{ContainerPort: constants.ContainerTrainerPort}}).
Env(constants.Node, constants.Node,
[]corev1.EnvVar{
{
Name: "TRAIN_JOB",
Value: "override",
},
{
Name: "TRAIN_JOB_CUSTOM",
Value: "test:trainjob",
},
{
Name: constants.TorchEnvNumNodes,
Value: "100",
},
{
Name: constants.TorchEnvNumProcPerNode,
Value: "1",
},
{
Name: constants.TorchEnvNodeRank,
ValueFrom: &corev1.EnvVarSource{
FieldRef: &corev1.ObjectFieldSelector{
FieldPath: constants.JobCompletionIndexFieldPath,
},
},
},
{
Name: constants.TorchEnvMasterAddr,
Value: fmt.Sprintf("test-job-%s-0-0.test-job", constants.Node),
},
{
Name: constants.TorchEnvMasterPort,
Value: fmt.Sprintf("%d", constants.ContainerTrainerPort),
},
{
Name: "RUNTIME",
Value: "test:runtime",
},
}...,
).
Obj(),
},
},
"succeeded to build JobSet with TorchTune values from the TrainJob": {
trainingRuntime: testingutil.MakeTrainingRuntimeWrapper(metav1.NamespaceDefault, "torchtune-llama3.3-70b").RuntimeSpec(
testingutil.MakeTrainingRuntimeSpecWrapper(testingutil.MakeTrainingRuntimeWrapper(metav1.NamespaceDefault, "torchtune-llama3.3-70b").Spec).
WithMLPolicy(
testingutil.MakeMLPolicyWrapper().
WithNumNodes(100).
WithMLPolicySource(*testingutil.MakeMLPolicySourceWrapper().
TorchPolicy(ptr.To(intstr.FromString("auto")), nil).
Obj(),
).
Obj(),
).
JobSetSpec(
testingutil.MakeJobSetWrapper("", "").
DependsOn(constants.Node,
[]jobsetv1alpha2.DependsOn{
{
Name: constants.DatasetInitializer,
Status: jobsetv1alpha2.DependencyComplete,
},
{
Name: constants.ModelInitializer,
Status: jobsetv1alpha2.DependencyComplete,
},
}...,
).
Obj().
Spec,
).
Container(constants.Node, constants.Node, "test:runtime", []string{"runtime"}, []string{"runtime"}, resRequests).
Obj(),
).Obj(),
trainJob: testingutil.MakeTrainJobWrapper(metav1.NamespaceDefault, "test-job").
UID("uid").
RuntimeRef(trainer.SchemeGroupVersion.WithKind(trainer.TrainingRuntimeKind), "torchtune-llama3.3-70b").
Trainer(
testingutil.MakeTrainJobTrainerWrapper().
Container(
"test:trainjob",
[]string{"tune", "run"},
[]string{
"dtype=fp16",
"batch_size=32",
"epochs=10",
"loss=torchtune.modules.loss.CEWithChunkedOutputLoss",
},
resRequests).
NumNodes(30).
NumProcPerNode(intstr.FromInt32(3)).
Obj(),
).
Obj(),
wantObjs: []runtime.Object{
testingutil.MakeJobSetWrapper(metav1.NamespaceDefault, "test-job").
ControllerReference(trainer.SchemeGroupVersion.WithKind(trainer.TrainJobKind), "test-job", "uid").
Replicas(1, constants.DatasetInitializer, constants.ModelInitializer, constants.Node, constants.Launcher).
Parallelism(1, constants.DatasetInitializer, constants.ModelInitializer).
Completions(1, constants.DatasetInitializer, constants.ModelInitializer).
NumNodes(30).
Container(
constants.Node,
constants.Node,
"test:trainjob",
[]string{
"tune",
"run",
fmt.Sprintf("%s %s", constants.TorchTuneArgRdzvEndpoint, "test-job-node-0-0.test-job:29500"),
constants.TorchTuneFullFinetuneDistributed,
"--config llama3_3/70B_full_multinode.yaml",
},
[]string{
"dtype=fp16",
"batch_size=32",
"epochs=10",
"loss=torchtune.modules.loss.CEWithChunkedOutputLoss",
},
resRequests,
).
ContainerTrainerPorts([]corev1.ContainerPort{{ContainerPort: constants.ContainerTrainerPort}}).
Env(constants.Node, constants.Node,
[]corev1.EnvVar{
{
Name: constants.TorchEnvNumNodes,
Value: "30",
},
{
Name: constants.TorchEnvNumProcPerNode,
Value: "3",
},
{
Name: constants.TorchEnvNodeRank,
ValueFrom: &corev1.EnvVarSource{
FieldRef: &corev1.ObjectFieldSelector{
FieldPath: constants.JobCompletionIndexFieldPath,
},
},
},
}...,
).
DependsOn(constants.Node,
[]jobsetv1alpha2.DependsOn{
{
Name: constants.DatasetInitializer,
Status: jobsetv1alpha2.DependencyComplete,
},
{
Name: constants.ModelInitializer,
Status: jobsetv1alpha2.DependencyComplete,
},
}...,
).
Obj(),
},
},
"succeeded to build JobSet with OpenMPI values from the TrainJob": {
ObjCmpOpts: cmp.Options{
cmp.Comparer(testingutil.MPISecretDataComparer),
},
trainingRuntime: testingutil.MakeTrainingRuntimeWrapper(metav1.NamespaceDefault, "test-runtime").
RuntimeSpec(
testingutil.MakeTrainingRuntimeSpecWrapper(testingutil.MakeTrainingRuntimeWrapper(metav1.NamespaceDefault, "test-runtime").Spec).
WithMLPolicy(
testingutil.MakeMLPolicyWrapper().
WithNumNodes(1).
WithMLPolicySource(*testingutil.MakeMLPolicySourceWrapper().
MPIPolicy(ptr.To[int32](1), ptr.To(trainer.MPIImplementationOpenMPI), ptr.To("/root/.ssh"), ptr.To(false)).
Obj(),
).
Obj(),
).
LauncherReplica().
Replicas(1, constants.Launcher).
Container(constants.Node, constants.Node, "test:runtime", []string{"runtime"}, []string{"runtime"}, resRequests).
Obj(),
).
Obj(),
trainJob: testingutil.MakeTrainJobWrapper(metav1.NamespaceDefault, "test-job").
UID("uid").
RuntimeRef(trainer.SchemeGroupVersion.WithKind(trainer.TrainingRuntimeKind), "test-runtime").
Trainer(
testingutil.MakeTrainJobTrainerWrapper().
NumNodes(2).
NumProcPerNode(intstr.FromInt32(8)).
Obj(),
).
Obj(),
wantObjs: []runtime.Object{
testingutil.MakeConfigMapWrapper(fmt.Sprintf("test-job%s", constants.MPIHostfileConfigMapSuffix), metav1.NamespaceDefault).
ControllerReference(trainer.SchemeGroupVersion.WithKind(trainer.TrainJobKind), "test-job", "uid").
WithData(map[string]string{
constants.MPIHostfileName: `test-job-node-0-0.test-job slots=8
test-job-node-0-1.test-job slots=8
`,
}).
Obj(),
testingutil.MakeSecretWrapper(fmt.Sprintf("test-job%s", constants.MPISSHAuthSecretSuffix), metav1.NamespaceDefault).
ControllerReference(trainer.SchemeGroupVersion.WithKind(trainer.TrainJobKind), "test-job", "uid").
WithImmutable(true).
WithData(map[string][]byte{
corev1.SSHAuthPrivateKey: []byte("EXIST"),
constants.MPISSHPublicKey: []byte("EXIST"),
}).
WithType(corev1.SecretTypeSSHAuth).
Obj(),
testingutil.MakeJobSetWrapper(metav1.NamespaceDefault, "test-job").
ControllerReference(trainer.SchemeGroupVersion.WithKind(trainer.TrainJobKind), "test-job", "uid").
LauncherReplica().
Replicas(1, constants.DatasetInitializer, constants.ModelInitializer, constants.Node, constants.Launcher).
Parallelism(1, constants.DatasetInitializer, constants.ModelInitializer, constants.Launcher).
Completions(1, constants.DatasetInitializer, constants.ModelInitializer, constants.Launcher).
NumNodes(2).
Volumes(constants.Launcher,
corev1.Volume{
Name: constants.MPISSHAuthVolumeName,
VolumeSource: corev1.VolumeSource{
Secret: &corev1.SecretVolumeSource{
SecretName: fmt.Sprintf("test-job%s", constants.MPISSHAuthSecretSuffix),
Items: []corev1.KeyToPath{
{
Key: corev1.SSHAuthPrivateKey,
Path: constants.MPISSHPrivateKeyFile,
},
{
Key: constants.MPISSHPublicKey,
Path: constants.MPISSHPublicKeyFile,
},
{
Key: constants.MPISSHPublicKey,
Path: constants.MPISSHAuthorizedKeys,
},
},
},
},
},
corev1.Volume{
Name: constants.MPIHostfileVolumeName,
VolumeSource: corev1.VolumeSource{
ConfigMap: &corev1.ConfigMapVolumeSource{
LocalObjectReference: corev1.LocalObjectReference{
Name: fmt.Sprintf("test-job%s", constants.MPIHostfileConfigMapSuffix),
},
Items: []corev1.KeyToPath{{
Key: constants.MPIHostfileName,
Path: constants.MPIHostfileName,
Mode: ptr.To[int32](0444),
}},
},
},
},
).
Volumes(constants.Node,
corev1.Volume{
Name: constants.MPISSHAuthVolumeName,
VolumeSource: corev1.VolumeSource{
Secret: &corev1.SecretVolumeSource{
SecretName: fmt.Sprintf("test-job%s", constants.MPISSHAuthSecretSuffix),
Items: []corev1.KeyToPath{
{
Key: corev1.SSHAuthPrivateKey,
Path: constants.MPISSHPrivateKeyFile,
},
{
Key: constants.MPISSHPublicKey,
Path: constants.MPISSHPublicKeyFile,
},
{
Key: constants.MPISSHPublicKey,
Path: constants.MPISSHAuthorizedKeys,
},
},
},
},
},
).
VolumeMounts(constants.Launcher, constants.Node,
corev1.VolumeMount{Name: constants.MPISSHAuthVolumeName, MountPath: "/root/.ssh"},
corev1.VolumeMount{Name: constants.MPIHostfileVolumeName, MountPath: constants.MPIHostfileDir},
).
VolumeMounts(constants.Node, constants.Node,
corev1.VolumeMount{Name: constants.MPISSHAuthVolumeName, MountPath: "/root/.ssh"},
).
Env(constants.Launcher, constants.Node,
corev1.EnvVar{
Name: constants.OpenMPIEnvHostFileLocation,
Value: fmt.Sprintf("%s/%s", constants.MPIHostfileDir, constants.MPIHostfileName),
},
corev1.EnvVar{
Name: constants.OpenMPIEnvKeepFQDNHostNames,
Value: "true",
},
corev1.EnvVar{
Name: constants.OpenMPIEnvDefaultSlots,
Value: "8",
},
corev1.EnvVar{
Name: constants.OpenMPIEnvKeyRSHArgs,
Value: constants.OpenMPIEnvDefaultValueRSHArgs,
},
).
Container(constants.Node, constants.Node, "test:runtime", []string{"runtime"}, []string{"runtime"}, resRequests).
Obj(),
},
},
// Failed test cases.
"missing trainingRuntime resource": {
trainJob: testingutil.MakeTrainJobWrapper(metav1.NamespaceDefault, "test-job-3").
UID("uid").
RuntimeRef(trainer.SchemeGroupVersion.WithKind(trainer.TrainingRuntimeKind), "test-runtime-3").
Trainer(
testingutil.MakeTrainJobTrainerWrapper().
Obj(),
).
Obj(),
wantError: errorNotFoundSpecifiedTrainingRuntime,
},
}
cmpOpts := []cmp.Option{
cmpopts.SortSlices(func(a, b runtime.Object) bool {
return a.GetObjectKind().GroupVersionKind().String() < b.GetObjectKind().GroupVersionKind().String()
}),
cmpopts.SortSlices(func(a, b corev1.EnvVar) bool {
return a.Name < b.Name
}),
}
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 := testingutil.NewClientBuilder()
if tc.trainingRuntime != nil {
clientBuilder.WithObjects(tc.trainingRuntime)
}
c := clientBuilder.Build()
trainingRuntime, err := NewTrainingRuntime(ctx, c, testingutil.AsIndex(clientBuilder))
if err != nil {
t.Fatal(err)
}
objs, err := trainingRuntime.NewObjects(ctx, tc.trainJob)
if diff := cmp.Diff(tc.wantError, err, cmpopts.EquateErrors()); len(diff) != 0 {
t.Errorf("Unexpected error (-want,+got):\n%s", diff)
}
resultObjs, err := testingutil.ToObject(c.Scheme(), objs...)
if err != nil {
t.Errorf("Pipeline built unrecognizable objects: %v", err)
}
if diff := cmp.Diff(tc.wantObjs, resultObjs, append(cmpOpts, tc.ObjCmpOpts...)...); len(diff) != 0 {
t.Errorf("Unexpected objects (-want,+got):\n%s", diff)
}
})
}
}
Reference in New Issue View Git Blame Copy Permalink