diff --git a/api/openapi-spec/swagger.json b/api/openapi-spec/swagger.json index 0ed4e99fe9..cfa8180412 100644 --- a/api/openapi-spec/swagger.json +++ b/api/openapi-spec/swagger.json @@ -517,7 +517,7 @@ }, "numProcPerNode": { "description": "Number of processes per node. This value is inserted into the `--nproc-per-node` argument of the `torchrun` CLI. Supported values: `auto`, `cpu`, `gpu`, or int value. Defaults to `auto`.", - "type": "string" + "$ref": "#/definitions/k8s.io.apimachinery.pkg.util.intstr.IntOrString" } } }, @@ -716,7 +716,7 @@ }, "numProcPerNode": { "description": "Number of processes/workers/slots on every training node. For the Torch runtime: `auto`, `cpu`, `gpu`, or int value can be set. For the MPI runtime only int value can be set.", - "type": "string" + "$ref": "#/definitions/k8s.io.apimachinery.pkg.util.intstr.IntOrString" }, "resourcesPerNode": { "description": "Compute resources for each training node.", diff --git a/manifests/base/crds/trainer.kubeflow.org_clustertrainingruntimes.yaml b/manifests/base/crds/trainer.kubeflow.org_clustertrainingruntimes.yaml index 3aaa6efdc3..0b350bb353 100644 --- a/manifests/base/crds/trainer.kubeflow.org_clustertrainingruntimes.yaml +++ b/manifests/base/crds/trainer.kubeflow.org_clustertrainingruntimes.yaml @@ -583,12 +583,17 @@ spec: type: integer type: object numProcPerNode: + anyOf: + - type: integer + - type: string description: |- Number of processes per node. This value is inserted into the `--nproc-per-node` argument of the `torchrun` CLI. Supported values: `auto`, `cpu`, `gpu`, or int value. Defaults to `auto`. - type: string + x-kubernetes-int-or-string: true + x-kubernetes-validations: + - rule: self > 0 || self in ['auto', 'cpu', 'gpu'] type: object type: object podGroupPolicy: diff --git a/manifests/base/crds/trainer.kubeflow.org_trainingruntimes.yaml b/manifests/base/crds/trainer.kubeflow.org_trainingruntimes.yaml index 30d6a445c4..589bfb68bc 100644 --- a/manifests/base/crds/trainer.kubeflow.org_trainingruntimes.yaml +++ b/manifests/base/crds/trainer.kubeflow.org_trainingruntimes.yaml @@ -583,12 +583,17 @@ spec: type: integer type: object numProcPerNode: + anyOf: + - type: integer + - type: string description: |- Number of processes per node. This value is inserted into the `--nproc-per-node` argument of the `torchrun` CLI. Supported values: `auto`, `cpu`, `gpu`, or int value. Defaults to `auto`. - type: string + x-kubernetes-int-or-string: true + x-kubernetes-validations: + - rule: self > 0 || self in ['auto', 'cpu', 'gpu'] type: object type: object podGroupPolicy: diff --git a/manifests/base/crds/trainer.kubeflow.org_trainjobs.yaml b/manifests/base/crds/trainer.kubeflow.org_trainjobs.yaml index 5c4c7cb7f4..16a2294cb9 100644 --- a/manifests/base/crds/trainer.kubeflow.org_trainjobs.yaml +++ b/manifests/base/crds/trainer.kubeflow.org_trainjobs.yaml @@ -3138,11 +3138,14 @@ spec: format: int32 type: integer numProcPerNode: + anyOf: + - type: integer + - type: string description: |- Number of processes/workers/slots on every training node. For the Torch runtime: `auto`, `cpu`, `gpu`, or int value can be set. For the MPI runtime only int value can be set. - type: string + x-kubernetes-int-or-string: true resourcesPerNode: description: Compute resources for each training node. properties: diff --git a/pkg/apis/trainer/v1alpha1/trainingruntime_types.go b/pkg/apis/trainer/v1alpha1/trainingruntime_types.go index b47bd9284b..9f907e28cb 100644 --- a/pkg/apis/trainer/v1alpha1/trainingruntime_types.go +++ b/pkg/apis/trainer/v1alpha1/trainingruntime_types.go @@ -19,6 +19,7 @@ package v1alpha1 import ( autoscalingv2 "k8s.io/api/autoscaling/v2" metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" + "k8s.io/apimachinery/pkg/util/intstr" jobsetv1alpha2 "sigs.k8s.io/jobset/api/jobset/v1alpha2" ) @@ -171,9 +172,9 @@ type TorchMLPolicySource struct { // Number of processes per node. // This value is inserted into the `--nproc-per-node` argument of the `torchrun` CLI. // Supported values: `auto`, `cpu`, `gpu`, or int value. - // TODO (andreyvelich): Add kubebuilder validation. // Defaults to `auto`. - NumProcPerNode *string `json:"numProcPerNode,omitempty"` + // +kubebuilder:validation:XValidation:rule="self > 0 || self in ['auto', 'cpu', 'gpu']" + NumProcPerNode *intstr.IntOrString `json:"numProcPerNode,omitempty"` // Elastic policy for the PyTorch training. ElasticPolicy *TorchElasticPolicy `json:"elasticPolicy,omitempty"` diff --git a/pkg/apis/trainer/v1alpha1/trainjob_types.go b/pkg/apis/trainer/v1alpha1/trainjob_types.go index 03ef0816ad..0ea6ddbbf4 100644 --- a/pkg/apis/trainer/v1alpha1/trainjob_types.go +++ b/pkg/apis/trainer/v1alpha1/trainjob_types.go @@ -19,6 +19,7 @@ package v1alpha1 import ( corev1 "k8s.io/api/core/v1" metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" + "k8s.io/apimachinery/pkg/util/intstr" ) const ( @@ -194,7 +195,7 @@ type Trainer struct { // Number of processes/workers/slots on every training node. // For the Torch runtime: `auto`, `cpu`, `gpu`, or int value can be set. // For the MPI runtime only int value can be set. - NumProcPerNode *string `json:"numProcPerNode,omitempty"` + NumProcPerNode *intstr.IntOrString `json:"numProcPerNode,omitempty"` } // DatasetConfig represents the desired dataset configuration. diff --git a/pkg/apis/trainer/v1alpha1/zz_generated.deepcopy.go b/pkg/apis/trainer/v1alpha1/zz_generated.deepcopy.go index 298d478216..96bab87de7 100644 --- a/pkg/apis/trainer/v1alpha1/zz_generated.deepcopy.go +++ b/pkg/apis/trainer/v1alpha1/zz_generated.deepcopy.go @@ -24,6 +24,7 @@ import ( v1 "k8s.io/api/core/v1" metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" runtime "k8s.io/apimachinery/pkg/runtime" + intstr "k8s.io/apimachinery/pkg/util/intstr" ) // DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil. @@ -576,7 +577,7 @@ func (in *TorchMLPolicySource) DeepCopyInto(out *TorchMLPolicySource) { *out = *in if in.NumProcPerNode != nil { in, out := &in.NumProcPerNode, &out.NumProcPerNode - *out = new(string) + *out = new(intstr.IntOrString) **out = **in } if in.ElasticPolicy != nil { @@ -786,7 +787,7 @@ func (in *Trainer) DeepCopyInto(out *Trainer) { } if in.NumProcPerNode != nil { in, out := &in.NumProcPerNode, &out.NumProcPerNode - *out = new(string) + *out = new(intstr.IntOrString) **out = **in } return diff --git a/pkg/apis/trainer/v1alpha1/zz_generated.openapi.go b/pkg/apis/trainer/v1alpha1/zz_generated.openapi.go index 27cc25c68c..6ae643bdbe 100644 --- a/pkg/apis/trainer/v1alpha1/zz_generated.openapi.go +++ b/pkg/apis/trainer/v1alpha1/zz_generated.openapi.go @@ -974,8 +974,7 @@ func schema_pkg_apis_trainer_v1alpha1_TorchMLPolicySource(ref common.ReferenceCa "numProcPerNode": { SchemaProps: spec.SchemaProps{ Description: "Number of processes per node. This value is inserted into the `--nproc-per-node` argument of the `torchrun` CLI. Supported values: `auto`, `cpu`, `gpu`, or int value. Defaults to `auto`.", - Type: []string{"string"}, - Format: "", + Ref: ref("k8s.io/apimachinery/pkg/util/intstr.IntOrString"), }, }, "elasticPolicy": { @@ -988,7 +987,7 @@ func schema_pkg_apis_trainer_v1alpha1_TorchMLPolicySource(ref common.ReferenceCa }, }, Dependencies: []string{ - "github.com/kubeflow/trainer/pkg/apis/trainer/v1alpha1.TorchElasticPolicy"}, + "github.com/kubeflow/trainer/pkg/apis/trainer/v1alpha1.TorchElasticPolicy", "k8s.io/apimachinery/pkg/util/intstr.IntOrString"}, } } @@ -1352,15 +1351,14 @@ func schema_pkg_apis_trainer_v1alpha1_Trainer(ref common.ReferenceCallback) comm "numProcPerNode": { SchemaProps: spec.SchemaProps{ Description: "Number of processes/workers/slots on every training node. For the Torch runtime: `auto`, `cpu`, `gpu`, or int value can be set. For the MPI runtime only int value can be set.", - Type: []string{"string"}, - Format: "", + Ref: ref("k8s.io/apimachinery/pkg/util/intstr.IntOrString"), }, }, }, }, }, Dependencies: []string{ - "k8s.io/api/core/v1.EnvVar", "k8s.io/api/core/v1.ResourceRequirements"}, + "k8s.io/api/core/v1.EnvVar", "k8s.io/api/core/v1.ResourceRequirements", "k8s.io/apimachinery/pkg/util/intstr.IntOrString"}, } } diff --git a/pkg/client/applyconfiguration/trainer/v1alpha1/torchmlpolicysource.go b/pkg/client/applyconfiguration/trainer/v1alpha1/torchmlpolicysource.go index c9d14b1ec1..6c0bae78dc 100644 --- a/pkg/client/applyconfiguration/trainer/v1alpha1/torchmlpolicysource.go +++ b/pkg/client/applyconfiguration/trainer/v1alpha1/torchmlpolicysource.go @@ -16,10 +16,14 @@ package v1alpha1 +import ( + intstr "k8s.io/apimachinery/pkg/util/intstr" +) + // TorchMLPolicySourceApplyConfiguration represents a declarative configuration of the TorchMLPolicySource type for use // with apply. type TorchMLPolicySourceApplyConfiguration struct { - NumProcPerNode *string `json:"numProcPerNode,omitempty"` + NumProcPerNode *intstr.IntOrString `json:"numProcPerNode,omitempty"` ElasticPolicy *TorchElasticPolicyApplyConfiguration `json:"elasticPolicy,omitempty"` } @@ -32,7 +36,7 @@ func TorchMLPolicySource() *TorchMLPolicySourceApplyConfiguration { // WithNumProcPerNode sets the NumProcPerNode field in the declarative configuration to the given value // and returns the receiver, so that objects can be built by chaining "With" function invocations. // If called multiple times, the NumProcPerNode field is set to the value of the last call. -func (b *TorchMLPolicySourceApplyConfiguration) WithNumProcPerNode(value string) *TorchMLPolicySourceApplyConfiguration { +func (b *TorchMLPolicySourceApplyConfiguration) WithNumProcPerNode(value intstr.IntOrString) *TorchMLPolicySourceApplyConfiguration { b.NumProcPerNode = &value return b } diff --git a/pkg/client/applyconfiguration/trainer/v1alpha1/trainer.go b/pkg/client/applyconfiguration/trainer/v1alpha1/trainer.go index 628757be34..464d641bc6 100644 --- a/pkg/client/applyconfiguration/trainer/v1alpha1/trainer.go +++ b/pkg/client/applyconfiguration/trainer/v1alpha1/trainer.go @@ -18,6 +18,7 @@ package v1alpha1 import ( v1 "k8s.io/api/core/v1" + intstr "k8s.io/apimachinery/pkg/util/intstr" ) // TrainerApplyConfiguration represents a declarative configuration of the Trainer type for use @@ -29,7 +30,7 @@ type TrainerApplyConfiguration struct { Env []v1.EnvVar `json:"env,omitempty"` NumNodes *int32 `json:"numNodes,omitempty"` ResourcesPerNode *v1.ResourceRequirements `json:"resourcesPerNode,omitempty"` - NumProcPerNode *string `json:"numProcPerNode,omitempty"` + NumProcPerNode *intstr.IntOrString `json:"numProcPerNode,omitempty"` } // TrainerApplyConfiguration constructs a declarative configuration of the Trainer type for use with @@ -95,7 +96,7 @@ func (b *TrainerApplyConfiguration) WithResourcesPerNode(value v1.ResourceRequir // WithNumProcPerNode sets the NumProcPerNode field in the declarative configuration to the given value // and returns the receiver, so that objects can be built by chaining "With" function invocations. // If called multiple times, the NumProcPerNode field is set to the value of the last call. -func (b *TrainerApplyConfiguration) WithNumProcPerNode(value string) *TrainerApplyConfiguration { +func (b *TrainerApplyConfiguration) WithNumProcPerNode(value intstr.IntOrString) *TrainerApplyConfiguration { b.NumProcPerNode = &value return b } diff --git a/pkg/runtime/core/trainingruntime_test.go b/pkg/runtime/core/trainingruntime_test.go index 7eb75de381..63301d168a 100644 --- a/pkg/runtime/core/trainingruntime_test.go +++ b/pkg/runtime/core/trainingruntime_test.go @@ -19,6 +19,7 @@ package core import ( "context" "fmt" + "k8s.io/apimachinery/pkg/util/intstr" "testing" "github.com/google/go-cmp/cmp" @@ -263,7 +264,7 @@ func TestTrainingRuntimeNewObjects(t *testing.T) { "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). - TorchPolicy(100, "auto"). + TorchPolicy(100, intstr.FromString("auto")). ContainerTrainer("test:runtime", []string{"runtime"}, []string{"runtime"}, resRequests). Obj(), ).Obj(), @@ -273,7 +274,7 @@ func TestTrainingRuntimeNewObjects(t *testing.T) { Trainer( testingutil.MakeTrainJobTrainerWrapper(). NumNodes(30). - NumProcPerNode("3"). + NumProcPerNode(intstr.FromInt32(3)). Obj(), ). Obj(), @@ -317,7 +318,7 @@ func TestTrainingRuntimeNewObjects(t *testing.T) { "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). - TorchPolicy(100, "auto"). + TorchPolicy(100, intstr.FromString("auto")). ContainerTrainer("test:runtime", []string{"runtime"}, []string{"runtime"}, resRequests). ContainerTrainerEnv( []corev1.EnvVar{ diff --git a/pkg/runtime/framework/plugins/torch/torch.go b/pkg/runtime/framework/plugins/torch/torch.go index 9f4d35c8e6..fbba91978b 100644 --- a/pkg/runtime/framework/plugins/torch/torch.go +++ b/pkg/runtime/framework/plugins/torch/torch.go @@ -19,6 +19,7 @@ package torch import ( "context" "fmt" + "k8s.io/apimachinery/pkg/util/intstr" corev1 "k8s.io/api/core/v1" "k8s.io/apimachinery/pkg/util/sets" @@ -61,9 +62,9 @@ func (t *Torch) EnforceMLPolicy(info *runtime.Info, trainJob *trainer.TrainJob) } info.Trainer.NumNodes = numNodes - numProcPerNode := info.RuntimePolicy.MLPolicy.Torch.NumProcPerNode + numProcPerNode := ptr.Deref(info.RuntimePolicy.MLPolicy.Torch.NumProcPerNode, intstr.FromString("auto")) if trainJob.Spec.Trainer != nil && trainJob.Spec.Trainer.NumProcPerNode != nil { - numProcPerNode = trainJob.Spec.Trainer.NumProcPerNode + numProcPerNode = ptr.Deref(trainJob.Spec.Trainer.NumProcPerNode, intstr.FromString("auto")) } // Update envs for Info object. @@ -78,7 +79,7 @@ func (t *Torch) EnforceMLPolicy(info *runtime.Info, trainJob *trainer.TrainJob) }, { Name: constants.TorchEnvNumProcPerNode, - Value: ptr.Deref(numProcPerNode, "auto"), + Value: numProcPerNode.String(), }, { Name: constants.TorchEnvNodeRank, diff --git a/pkg/util/testing/wrapper.go b/pkg/util/testing/wrapper.go index 2002f66a4e..f2d735b12f 100644 --- a/pkg/util/testing/wrapper.go +++ b/pkg/util/testing/wrapper.go @@ -22,6 +22,7 @@ import ( metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" "k8s.io/apimachinery/pkg/runtime/schema" "k8s.io/apimachinery/pkg/types" + "k8s.io/apimachinery/pkg/util/intstr" "k8s.io/utils/ptr" jobsetv1alpha2 "sigs.k8s.io/jobset/api/jobset/v1alpha2" schedulerpluginsv1alpha1 "sigs.k8s.io/scheduler-plugins/apis/scheduling/v1alpha1" @@ -392,7 +393,7 @@ func (t *TrainJobTrainerWrapper) NumNodes(numNodes int32) *TrainJobTrainerWrappe return t } -func (t *TrainJobTrainerWrapper) NumProcPerNode(numProcPerNode string) *TrainJobTrainerWrapper { +func (t *TrainJobTrainerWrapper) NumProcPerNode(numProcPerNode intstr.IntOrString) *TrainJobTrainerWrapper { t.Trainer.NumProcPerNode = &numProcPerNode return t } @@ -689,7 +690,7 @@ func (s *TrainingRuntimeSpecWrapper) NumNodes(numNodes int32) *TrainingRuntimeSp return s } -func (s *TrainingRuntimeSpecWrapper) TorchPolicy(numNodes int32, numProcPerNode string) *TrainingRuntimeSpecWrapper { +func (s *TrainingRuntimeSpecWrapper) TorchPolicy(numNodes int32, numProcPerNode intstr.IntOrString) *TrainingRuntimeSpecWrapper { s.MLPolicy = &trainer.MLPolicy{ NumNodes: &numNodes, MLPolicySource: trainer.MLPolicySource{ diff --git a/test/integration/controller/trainjob_controller_test.go b/test/integration/controller/trainjob_controller_test.go index e3715c160a..793effb367 100644 --- a/test/integration/controller/trainjob_controller_test.go +++ b/test/integration/controller/trainjob_controller_test.go @@ -26,6 +26,7 @@ import ( "k8s.io/apimachinery/pkg/api/resource" metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" "k8s.io/apimachinery/pkg/runtime/schema" + "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" @@ -278,7 +279,7 @@ var _ = ginkgo.Describe("TrainJob controller", ginkgo.Ordered, func() { trainingRuntime = testingutil.MakeTrainingRuntimeWrapper(ns.Name, "alpha"). RuntimeSpec( testingutil.MakeTrainingRuntimeSpecWrapper(testingutil.MakeTrainingRuntimeWrapper(metav1.NamespaceDefault, "alpha").Spec). - TorchPolicy(100, "auto"). + TorchPolicy(100, intstr.FromString("auto")). ContainerTrainer("test:runtime", []string{"runtime"}, []string{"runtime"}, resRequests). Obj()). Obj()