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

test(runtime): add UT for torch runtime valid function. (#2560)

Browse Source 
Signed-off-by: Asklv <[email protected]>
This commit is contained in:
Asklv
2025-03-26 07:47:04 +00:00
committed by GitHub
parent 3da33088a8
commit ccbd58e4ee
6 changed files with 234 additions and 37 deletions
Download Patch File Download Diff File Expand all files Collapse all files
pkg/runtime/core/trainingruntime_test.go
+2 -2
View File
@@ -299,7 +299,7 @@ func TestTrainingRuntimeNewObjects(t *testing.T) {
testingutil.MakeMLPolicyWrapper().
WithNumNodes(100).
WithMLPolicySource(*testingutil.MakeMLPolicySourceWrapper().
TorchPolicy("auto", nil).
TorchPolicy(ptr.To(intstr.FromString("auto")), nil).
Obj(),
).
Obj(),
@@ -393,7 +393,7 @@ func TestTrainingRuntimeNewObjects(t *testing.T) {
testingutil.MakeMLPolicyWrapper().
WithNumNodes(100).
WithMLPolicySource(*testingutil.MakeMLPolicySourceWrapper().
TorchPolicy("auto", nil).
TorchPolicy(ptr.To(intstr.FromString("auto")), nil).
Obj(),
).
Obj(),
pkg/runtime/framework/plugins/plainml/plainml_test.go
+3 -2
View File
@@ -24,6 +24,7 @@ import (
"github.com/google/go-cmp/cmp/cmpopts"
corev1 "k8s.io/api/core/v1"
metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
"k8s.io/apimachinery/pkg/util/intstr"
corev1ac "k8s.io/client-go/applyconfigurations/core/v1"
"k8s.io/klog/v2/ktesting"
"k8s.io/utils/ptr"
@@ -48,7 +49,7 @@ func TestPlainML(t *testing.T) {
Labels: map[string]string{"key": "value"},
RuntimePolicy: runtime.RuntimePolicy{
MLPolicySource: utiltesting.MakeMLPolicySourceWrapper().
TorchPolicy("auto", nil).
TorchPolicy(ptr.To(intstr.FromString("auto")), nil).
Obj(),
},
},
@@ -56,7 +57,7 @@ func TestPlainML(t *testing.T) {
Labels: map[string]string{"key": "value"},
RuntimePolicy: runtime.RuntimePolicy{
MLPolicySource: utiltesting.MakeMLPolicySourceWrapper().
TorchPolicy("auto", nil).
TorchPolicy(ptr.To(intstr.FromString("auto")), nil).
Obj(),
},
},
pkg/runtime/framework/plugins/torch/torch.go
+2 -2
View File
@@ -66,7 +66,7 @@ func (t *Torch) Validate(runtimeInfo *runtime.Info, _, newObj *trainer.TrainJob)
if numProcPerNode.Type == intstr.String {
allowed := sets.New("auto", "cpu", "gpu")
if !allowed.Has(numProcPerNode.StrVal) {
allErrs = append(allErrs, field.Invalid(numProcPerNodePath, newObj.Spec.Trainer.NumProcPerNode, fmt.Sprintf("must have an int value or %v", allowed.UnsortedList())))
allErrs = append(allErrs, field.Invalid(numProcPerNodePath, numProcPerNode, fmt.Sprintf("must have an int value or %v", sets.List(allowed))))
}
}
@@ -79,7 +79,7 @@ func (t *Torch) Validate(runtimeInfo *runtime.Info, _, newObj *trainer.TrainJob)
if torchEnvs.Len() > 0 {
trainerEnvsPath := specPath.Child("trainer").Child("env")
allErrs = append(allErrs, field.Invalid(trainerEnvsPath, newObj.Spec.Trainer.Env, fmt.Sprintf("must not have reserved envs, invalid envs configured: %s", strings.Join(torchEnvs.UnsortedList(), ", "))))
allErrs = append(allErrs, field.Invalid(trainerEnvsPath, newObj.Spec.Trainer.Env, fmt.Sprintf("must not have reserved envs, invalid envs configured: %v", sets.List(torchEnvs))))
}
}
pkg/runtime/framework/plugins/torch/torch_test.go
+223 -28
View File
@@ -27,9 +27,12 @@ import (
"k8s.io/apimachinery/pkg/api/resource"
metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
"k8s.io/apimachinery/pkg/util/intstr"
"k8s.io/apimachinery/pkg/util/sets"
"k8s.io/apimachinery/pkg/util/validation/field"
corev1ac "k8s.io/client-go/applyconfigurations/core/v1"
"k8s.io/klog/v2/ktesting"
"k8s.io/utils/ptr"
"sigs.k8s.io/controller-runtime/pkg/webhook/admission"
trainer "github.com/kubeflow/trainer/pkg/apis/trainer/v1alpha1"
"github.com/kubeflow/trainer/pkg/constants"
@@ -71,7 +74,7 @@ func TestTorch(t *testing.T) {
runtime.WithMLPolicySource(
utiltesting.MakeMLPolicyWrapper().
WithMLPolicySource(*utiltesting.MakeMLPolicySourceWrapper().
TorchPolicy("auto", nil).
TorchPolicy(ptr.To(intstr.FromString("auto")), nil).
Obj(),
).
Obj(),
@@ -93,7 +96,7 @@ func TestTorch(t *testing.T) {
Annotations: make(map[string]string),
RuntimePolicy: runtime.RuntimePolicy{
MLPolicySource: utiltesting.MakeMLPolicySourceWrapper().
TorchPolicy("auto", nil).
TorchPolicy(ptr.To(intstr.FromString("auto")), nil).
Obj(),
},
TemplateSpec: runtime.TemplateSpec{
@@ -154,7 +157,7 @@ func TestTorch(t *testing.T) {
runtime.WithMLPolicySource(
utiltesting.MakeMLPolicyWrapper().
WithMLPolicySource(*utiltesting.MakeMLPolicySourceWrapper().
TorchPolicy("auto", nil).
TorchPolicy(ptr.To(intstr.FromString("auto")), nil).
Obj(),
).
Obj(),
@@ -168,7 +171,7 @@ func TestTorch(t *testing.T) {
Annotations: make(map[string]string),
RuntimePolicy: runtime.RuntimePolicy{
MLPolicySource: utiltesting.MakeMLPolicySourceWrapper().
TorchPolicy("auto", nil).
TorchPolicy(ptr.To(intstr.FromString("auto")), nil).
Obj(),
},
TemplateSpec: runtime.TemplateSpec{
@@ -227,7 +230,7 @@ func TestTorch(t *testing.T) {
runtime.WithMLPolicySource(
utiltesting.MakeMLPolicyWrapper().
WithMLPolicySource(*utiltesting.MakeMLPolicySourceWrapper().
TorchPolicy("auto", nil).
TorchPolicy(ptr.To(intstr.FromString("auto")), nil).
Obj(),
).
Obj(),
@@ -241,7 +244,7 @@ func TestTorch(t *testing.T) {
Annotations: make(map[string]string),
RuntimePolicy: runtime.RuntimePolicy{
MLPolicySource: utiltesting.MakeMLPolicySourceWrapper().
TorchPolicy("auto", nil).
TorchPolicy(ptr.To(intstr.FromString("auto")), nil).
Obj(),
},
TemplateSpec: runtime.TemplateSpec{
@@ -302,7 +305,7 @@ func TestTorch(t *testing.T) {
runtime.WithMLPolicySource(
utiltesting.MakeMLPolicyWrapper().
WithMLPolicySource(*utiltesting.MakeMLPolicySourceWrapper().
TorchPolicy("auto", nil).
TorchPolicy(ptr.To(intstr.FromString("auto")), nil).
Obj(),
).
Obj(),
@@ -316,7 +319,7 @@ func TestTorch(t *testing.T) {
Annotations: make(map[string]string),
RuntimePolicy: runtime.RuntimePolicy{
MLPolicySource: utiltesting.MakeMLPolicySourceWrapper().
TorchPolicy("auto", nil).
TorchPolicy(ptr.To(intstr.FromString("auto")), nil).
Obj(),
},
TemplateSpec: runtime.TemplateSpec{
@@ -377,7 +380,7 @@ func TestTorch(t *testing.T) {
runtime.WithMLPolicySource(
utiltesting.MakeMLPolicyWrapper().
WithMLPolicySource(*utiltesting.MakeMLPolicySourceWrapper().
TorchPolicy("auto", nil).
TorchPolicy(ptr.To(intstr.FromString("auto")), nil).
Obj(),
).
Obj(),
@@ -391,7 +394,7 @@ func TestTorch(t *testing.T) {
Annotations: make(map[string]string),
RuntimePolicy: runtime.RuntimePolicy{
MLPolicySource: utiltesting.MakeMLPolicySourceWrapper().
TorchPolicy("auto", nil).
TorchPolicy(ptr.To(intstr.FromString("auto")), nil).
Obj(),
},
TemplateSpec: runtime.TemplateSpec{
@@ -452,7 +455,7 @@ func TestTorch(t *testing.T) {
runtime.WithMLPolicySource(
utiltesting.MakeMLPolicyWrapper().
WithMLPolicySource(*utiltesting.MakeMLPolicySourceWrapper().
TorchPolicy("auto", nil).
TorchPolicy(ptr.To(intstr.FromString("auto")), nil).
Obj(),
).
Obj(),
@@ -466,7 +469,7 @@ func TestTorch(t *testing.T) {
Annotations: make(map[string]string),
RuntimePolicy: runtime.RuntimePolicy{
MLPolicySource: utiltesting.MakeMLPolicySourceWrapper().
TorchPolicy("auto", nil).
TorchPolicy(ptr.To(intstr.FromString("auto")), nil).
Obj(),
},
TemplateSpec: runtime.TemplateSpec{
@@ -527,7 +530,7 @@ func TestTorch(t *testing.T) {
runtime.WithMLPolicySource(
utiltesting.MakeMLPolicyWrapper().
WithMLPolicySource(*utiltesting.MakeMLPolicySourceWrapper().
TorchPolicy("auto", nil).
TorchPolicy(ptr.To(intstr.FromString("auto")), nil).
Obj(),
).
Obj(),
@@ -541,7 +544,7 @@ func TestTorch(t *testing.T) {
Annotations: make(map[string]string),
RuntimePolicy: runtime.RuntimePolicy{
MLPolicySource: utiltesting.MakeMLPolicySourceWrapper().
TorchPolicy("auto", nil).
TorchPolicy(ptr.To(intstr.FromString("auto")), nil).
Obj(),
},
TemplateSpec: runtime.TemplateSpec{
@@ -602,7 +605,7 @@ func TestTorch(t *testing.T) {
runtime.WithMLPolicySource(
utiltesting.MakeMLPolicyWrapper().
WithMLPolicySource(*utiltesting.MakeMLPolicySourceWrapper().
TorchPolicy("auto", nil).
TorchPolicy(ptr.To(intstr.FromString("auto")), nil).
Obj(),
).
Obj(),
@@ -616,7 +619,7 @@ func TestTorch(t *testing.T) {
Annotations: make(map[string]string),
RuntimePolicy: runtime.RuntimePolicy{
MLPolicySource: utiltesting.MakeMLPolicySourceWrapper().
TorchPolicy("auto", nil).
TorchPolicy(ptr.To(intstr.FromString("auto")), nil).
Obj(),
},
TemplateSpec: runtime.TemplateSpec{
@@ -677,7 +680,7 @@ func TestTorch(t *testing.T) {
runtime.WithMLPolicySource(
utiltesting.MakeMLPolicyWrapper().
WithMLPolicySource(*utiltesting.MakeMLPolicySourceWrapper().
TorchPolicy("auto", nil).
TorchPolicy(ptr.To(intstr.FromString("auto")), nil).
Obj(),
).
Obj(),
@@ -691,7 +694,7 @@ func TestTorch(t *testing.T) {
Annotations: make(map[string]string),
RuntimePolicy: runtime.RuntimePolicy{
MLPolicySource: utiltesting.MakeMLPolicySourceWrapper().
TorchPolicy("auto", nil).
TorchPolicy(ptr.To(intstr.FromString("auto")), nil).
Obj(),
},
TemplateSpec: runtime.TemplateSpec{
@@ -752,7 +755,7 @@ func TestTorch(t *testing.T) {
runtime.WithMLPolicySource(
utiltesting.MakeMLPolicyWrapper().
WithMLPolicySource(*utiltesting.MakeMLPolicySourceWrapper().
TorchPolicy("auto", nil).
TorchPolicy(ptr.To(intstr.FromString("auto")), nil).
Obj(),
).
Obj(),
@@ -766,7 +769,7 @@ func TestTorch(t *testing.T) {
Annotations: make(map[string]string),
RuntimePolicy: runtime.RuntimePolicy{
MLPolicySource: utiltesting.MakeMLPolicySourceWrapper().
TorchPolicy("auto", nil).
TorchPolicy(ptr.To(intstr.FromString("auto")), nil).
Obj(),
},
TemplateSpec: runtime.TemplateSpec{
@@ -827,7 +830,7 @@ func TestTorch(t *testing.T) {
runtime.WithMLPolicySource(
utiltesting.MakeMLPolicyWrapper().
WithMLPolicySource(*utiltesting.MakeMLPolicySourceWrapper().
TorchPolicy("auto", nil).
TorchPolicy(ptr.To(intstr.FromString("auto")), nil).
Obj(),
).
Obj(),
@@ -841,7 +844,7 @@ func TestTorch(t *testing.T) {
Annotations: make(map[string]string),
RuntimePolicy: runtime.RuntimePolicy{
MLPolicySource: utiltesting.MakeMLPolicySourceWrapper().
TorchPolicy("auto", nil).
TorchPolicy(ptr.To(intstr.FromString("auto")), nil).
Obj(),
},
TemplateSpec: runtime.TemplateSpec{
@@ -903,7 +906,7 @@ func TestTorch(t *testing.T) {
runtime.WithMLPolicySource(
utiltesting.MakeMLPolicyWrapper().
WithMLPolicySource(*utiltesting.MakeMLPolicySourceWrapper().
TorchPolicy("auto", nil).
TorchPolicy(ptr.To(intstr.FromString("auto")), nil).
Obj(),
).
Obj(),
@@ -917,7 +920,7 @@ func TestTorch(t *testing.T) {
Annotations: make(map[string]string),
RuntimePolicy: runtime.RuntimePolicy{
MLPolicySource: utiltesting.MakeMLPolicySourceWrapper().
TorchPolicy("auto", nil).
TorchPolicy(ptr.To(intstr.FromString("auto")), nil).
Obj(),
},
TemplateSpec: runtime.TemplateSpec{
@@ -978,7 +981,7 @@ func TestTorch(t *testing.T) {
runtime.WithMLPolicySource(
utiltesting.MakeMLPolicyWrapper().
WithMLPolicySource(*utiltesting.MakeMLPolicySourceWrapper().
TorchPolicy("auto", nil).
TorchPolicy(ptr.To(intstr.FromString("auto")), nil).
Obj(),
).
Obj(),
@@ -992,7 +995,7 @@ func TestTorch(t *testing.T) {
Annotations: make(map[string]string),
RuntimePolicy: runtime.RuntimePolicy{
MLPolicySource: utiltesting.MakeMLPolicySourceWrapper().
TorchPolicy("auto", nil).
TorchPolicy(ptr.To(intstr.FromString("auto")), nil).
Obj(),
},
TemplateSpec: runtime.TemplateSpec{
@@ -1056,7 +1059,7 @@ func TestTorch(t *testing.T) {
runtime.WithMLPolicySource(
utiltesting.MakeMLPolicyWrapper().
WithMLPolicySource(*utiltesting.MakeMLPolicySourceWrapper().
TorchPolicy("auto", nil).
TorchPolicy(ptr.To(intstr.FromString("auto")), nil).
Obj(),
).
Obj(),
@@ -1077,7 +1080,7 @@ func TestTorch(t *testing.T) {
Annotations: make(map[string]string),
RuntimePolicy: runtime.RuntimePolicy{
MLPolicySource: utiltesting.MakeMLPolicySourceWrapper().
TorchPolicy("auto", nil).
TorchPolicy(ptr.To(intstr.FromString("auto")), nil).
Obj(),
},
TemplateSpec: runtime.TemplateSpec{
@@ -1153,3 +1156,195 @@ func TestTorch(t *testing.T) {
})
}
}
func TestValidate(t *testing.T) {
cases := map[string]struct {
info *runtime.Info
oldObj *trainer.TrainJob
newObj *trainer.TrainJob
wantError field.ErrorList
wantWarnings admission.Warnings
}{
"no action when info is nil": {
info: nil,
oldObj: utiltesting.MakeTrainJobWrapper(metav1.NamespaceDefault, "test").
Obj(),
newObj: utiltesting.MakeTrainJobWrapper(metav1.NamespaceDefault, "test").
Obj(),
},
"no action when info does not have MLPolicySource": {
info: runtime.NewInfo(),
oldObj: utiltesting.MakeTrainJobWrapper(metav1.NamespaceDefault, "test").
Obj(),
newObj: utiltesting.MakeTrainJobWrapper(metav1.NamespaceDefault, "test").
Obj(),
},
"no action when info has MLPolicySource but no Torch policy": {
info: runtime.NewInfo(
runtime.WithMLPolicySource(utiltesting.MakeMLPolicyWrapper().Obj()),
),
oldObj: utiltesting.MakeTrainJobWrapper(metav1.NamespaceDefault, "test").
Obj(),
newObj: utiltesting.MakeTrainJobWrapper(metav1.NamespaceDefault, "test").
Obj(),
},
"no action when info does not have numProcPerNode": {
info: runtime.NewInfo(
runtime.WithMLPolicySource(utiltesting.MakeMLPolicyWrapper().
WithMLPolicySource(*utiltesting.MakeMLPolicySourceWrapper().
TorchPolicy(nil, nil).
Obj(),
).
Obj(),
),
),
oldObj: utiltesting.MakeTrainJobWrapper(metav1.NamespaceDefault, "test").
Trainer(utiltesting.MakeTrainJobTrainerWrapper().
Obj(),
).
Obj(),
newObj: utiltesting.MakeTrainJobWrapper(metav1.NamespaceDefault, "test").
Trainer(utiltesting.MakeTrainJobTrainerWrapper().
Obj(),
).
Obj(),
},
"numProcPerNode is string and invalid": {
info: runtime.NewInfo(
runtime.WithMLPolicySource(utiltesting.MakeMLPolicyWrapper().
WithMLPolicySource(*utiltesting.MakeMLPolicySourceWrapper().
TorchPolicy(ptr.To(intstr.FromString("npu")), nil).
Obj(),
).
Obj(),
),
),
newObj: utiltesting.MakeTrainJobWrapper(metav1.NamespaceDefault, "test").
Trainer(utiltesting.MakeTrainJobTrainerWrapper().
NumProcPerNode(intstr.FromString("npu")).
Obj(),
).
Obj(),
wantError: field.ErrorList{
field.Invalid(
field.NewPath("spec").Child("trainer").Child("numProcPerNode"),
intstr.FromString("npu"),
fmt.Sprintf("must have an int value or %v", []string{"auto", "cpu", "gpu"}),
),
},
},
"numProcPerNode is valid string": {
info: runtime.NewInfo(
runtime.WithMLPolicySource(utiltesting.MakeMLPolicyWrapper().
WithMLPolicySource(*utiltesting.MakeMLPolicySourceWrapper().
TorchPolicy(ptr.To(intstr.FromString("auto")), nil).
Obj(),
).
Obj(),
),
),
newObj: utiltesting.MakeTrainJobWrapper(metav1.NamespaceDefault, "test").
Trainer(utiltesting.MakeTrainJobTrainerWrapper().
NumProcPerNode(intstr.FromString("auto")).
Obj(),
).
Obj(),
},
"valid environment variable present": {
info: runtime.NewInfo(
runtime.WithMLPolicySource(utiltesting.MakeMLPolicyWrapper().
WithMLPolicySource(*utiltesting.MakeMLPolicySourceWrapper().
TorchPolicy(ptr.To(intstr.FromString("auto")), nil).
Obj(),
).
Obj(),
),
),
newObj: utiltesting.MakeTrainJobWrapper(metav1.NamespaceDefault, "test").
Trainer(utiltesting.MakeTrainJobTrainerWrapper().
Env(
[]corev1.EnvVar{
{
Name: "test",
Value: "value",
},
{
Name: "test2",
Value: "value",
},
}...,
).
Obj(),
).
Obj(),
},
"reserved environment variable present": {
info: runtime.NewInfo(
runtime.WithMLPolicySource(utiltesting.MakeMLPolicyWrapper().
WithMLPolicySource(*utiltesting.MakeMLPolicySourceWrapper().
TorchPolicy(ptr.To(intstr.FromString("auto")), nil).
Obj(),
).
Obj(),
),
),
newObj: utiltesting.MakeTrainJobWrapper(metav1.NamespaceDefault, "test").
Trainer(utiltesting.MakeTrainJobTrainerWrapper().
NumProcPerNode(intstr.FromString("auto")).
Env(
[]corev1.EnvVar{
{
Name: "test",
Value: "value",
},
{
Name: constants.TorchEnvNumProcPerNode,
Value: "value",
},
}...,
).
Obj(),
).
Obj(),
wantError: field.ErrorList{
field.Invalid(
field.NewPath("spec").Child("trainer").Child("env"),
[]corev1.EnvVar{
{
Name: "test",
Value: "value",
},
{
Name: constants.TorchEnvNumProcPerNode,
Value: "value",
},
},
fmt.Sprintf("must not have reserved envs, invalid envs configured: %v", func() []string {
torchEnvs := sets.New[string]()
torchEnvs.Insert(constants.TorchEnvNumProcPerNode)
return sets.List(torchEnvs)
}()),
),
},
},
}
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)
p, err := New(ctx, utiltesting.NewClientBuilder().Build(), nil)
if err != nil {
t.Fatalf("Failed to initialize Torch plugin: %v", err)
}
warnings, errs := p.(framework.CustomValidationPlugin).Validate(tc.info, tc.oldObj, tc.newObj)
if diff := cmp.Diff(tc.wantError, errs); 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)
}
})
}
}
pkg/util/testing/wrapper.go
+2 -2
View File
@@ -1014,12 +1014,12 @@ func MakeMLPolicySourceWrapper() *MLPolicySourceWrapper {
return &MLPolicySourceWrapper{}
}
func (m *MLPolicySourceWrapper) TorchPolicy(numProcPerNode string, elasticPolicy *trainer.TorchElasticPolicy) *MLPolicySourceWrapper {
func (m *MLPolicySourceWrapper) TorchPolicy(numProcPerNode *intstr.IntOrString, elasticPolicy *trainer.TorchElasticPolicy) *MLPolicySourceWrapper {
if m.Torch == nil {
m.Torch = &trainer.TorchMLPolicySource{}
}
m.Torch = &trainer.TorchMLPolicySource{
NumProcPerNode: ptr.To(intstr.FromString(numProcPerNode)),
NumProcPerNode: numProcPerNode,
ElasticPolicy: elasticPolicy,
}
return m
test/integration/controller/trainjob_controller_test.go
+2 -1
View File
@@ -26,6 +26,7 @@ import (
"k8s.io/apimachinery/pkg/api/meta"
"k8s.io/apimachinery/pkg/api/resource"
metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
"k8s.io/apimachinery/pkg/util/intstr"
"k8s.io/utils/ptr"
"sigs.k8s.io/controller-runtime/pkg/client"
jobsetv1alpha2 "sigs.k8s.io/jobset/api/jobset/v1alpha2"
@@ -332,7 +333,7 @@ var _ = ginkgo.Describe("TrainJob controller", ginkgo.Ordered, func() {
testingutil.MakeMLPolicyWrapper().
WithNumNodes(100).
WithMLPolicySource(*testingutil.MakeMLPolicySourceWrapper().
TorchPolicy("auto", nil).
TorchPolicy(ptr.To(intstr.FromString("auto")), nil).
Obj(),
).
Obj(),