From e8a0059cc5eee438319f498621214501f2fda56d Mon Sep 17 00:00:00 2001 From: cegao Date: Thu, 9 Dec 2021 15:25:13 +0800 Subject: [PATCH 1/6] feat: Add success policy Signed-off-by: cegao --- pkg/apis/pytorch/v1/common.go | 9 +++ pkg/apis/pytorch/v1/defaults.go | 12 +++- pkg/apis/pytorch/v1/types.go | 2 + pkg/common/util/util.go | 12 ++++ .../pytorch/pytorchjob_controller.go | 71 +++++++++++++++++-- .../pytorch/pytorchjob_controller_test.go | 12 +++- .../tensorflow/tfjob_controller.go | 14 ++-- pkg/controller.v1/tensorflow/util.go | 12 ---- 8 files changed, 115 insertions(+), 29 deletions(-) create mode 100644 pkg/apis/pytorch/v1/common.go diff --git a/pkg/apis/pytorch/v1/common.go b/pkg/apis/pytorch/v1/common.go new file mode 100644 index 0000000000..cc5c743198 --- /dev/null +++ b/pkg/apis/pytorch/v1/common.go @@ -0,0 +1,9 @@ +package v1 + +// SuccessPolicy is the success policy. +type SuccessPolicy string + +const ( + SuccessPolicyDefault SuccessPolicy = "" + SuccessPolicyAllWorkers SuccessPolicy = "AllWorkers" +) diff --git a/pkg/apis/pytorch/v1/defaults.go b/pkg/apis/pytorch/v1/defaults.go index 12d2274558..f423ba8bd2 100644 --- a/pkg/apis/pytorch/v1/defaults.go +++ b/pkg/apis/pytorch/v1/defaults.go @@ -106,13 +106,21 @@ func setTypeNameToCamelCase(job *PyTorchJob, typ common.ReplicaType) { } } -// SetDefaults_PyTorchJob sets any unspecified values to defaults. -func SetDefaults_PyTorchJob(job *PyTorchJob) { +func SetDefaultRunPolicy(job *PyTorchJob) { // Set default cleanpod policy to None. if job.Spec.RunPolicy.CleanPodPolicy == nil { policy := common.CleanPodPolicyNone job.Spec.RunPolicy.CleanPodPolicy = &policy } + if job.Spec.SuccessPolicy == nil { + policy := SuccessPolicyDefault + job.Spec.SuccessPolicy = &policy + } +} + +// SetDefaults_PyTorchJob sets any unspecified values to defaults. +func SetDefaults_PyTorchJob(job *PyTorchJob) { + SetDefaultRunPolicy(job) // Update the key of PyTorchReplicaSpecs to camel case. setTypeNamesToCamelCase(job) diff --git a/pkg/apis/pytorch/v1/types.go b/pkg/apis/pytorch/v1/types.go index 2f9e973922..ab2457ceb2 100644 --- a/pkg/apis/pytorch/v1/types.go +++ b/pkg/apis/pytorch/v1/types.go @@ -52,6 +52,8 @@ type PyTorchJobSpec struct { //+kubebuilder:validation:Optional RunPolicy common.RunPolicy `json:"runPolicy"` + SuccessPolicy *SuccessPolicy `json:"successPolicy,omitempty"` + ElasticPolicy *ElasticPolicy `json:"elasticPolicy,omitempty"` // A map of PyTorchReplicaType (type) to ReplicaSpec (value). Specifies the PyTorch cluster configuration. diff --git a/pkg/common/util/util.go b/pkg/common/util/util.go index f635f48f4b..e8a9251574 100644 --- a/pkg/common/util/util.go +++ b/pkg/common/util/util.go @@ -58,3 +58,15 @@ func GetSchedulerName(replicas map[commonv1.ReplicaType]*commonv1.ReplicaSpec) s } return "" } + +// GetContainerExitCode gets the container exit code from the given pod. +func GetContainerExitCode(pod *corev1.Pod, name string) int32 { + var exitCode int32 = 0xbeef // magic number + for _, status := range pod.Status.ContainerStatuses { + state := status.State + if status.Name == name && state.Terminated != nil { + exitCode = state.Terminated.ExitCode + } + } + return exitCode +} diff --git a/pkg/controller.v1/pytorch/pytorchjob_controller.go b/pkg/controller.v1/pytorch/pytorchjob_controller.go index c16d10f7e8..79bce7406b 100644 --- a/pkg/controller.v1/pytorch/pytorchjob_controller.go +++ b/pkg/controller.v1/pytorch/pytorchjob_controller.go @@ -17,6 +17,7 @@ package pytorch import ( "context" "fmt" + "strings" "github.com/go-logr/logr" commonv1 "github.com/kubeflow/common/pkg/apis/common/v1" @@ -26,6 +27,7 @@ import ( commonutil "github.com/kubeflow/common/pkg/util" "github.com/sirupsen/logrus" corev1 "k8s.io/api/core/v1" + v1 "k8s.io/api/core/v1" "k8s.io/apimachinery/pkg/api/equality" "k8s.io/apimachinery/pkg/api/errors" "k8s.io/apimachinery/pkg/api/meta" @@ -326,6 +328,8 @@ func (r *PyTorchJobReconciler) UpdateJobStatus(job interface{}, return fmt.Errorf("%+v is not a type of PyTorchJob", job) } + logger := commonutil.LoggerForJob(pytorchjob) + for rtype, spec := range replicas { status := jobStatus.ReplicaStatuses[rtype] if status.LabelSelector == nil { @@ -338,7 +342,7 @@ func (r *PyTorchJobReconciler) UpdateJobStatus(job interface{}, running := status.Active failed := status.Failed - logrus.Infof("PyTorchJob=%s, ReplicaType=%s expected=%d, running=%d, succeeded=%d , failed=%d", + logger.Infof("PyTorchJob=%s, ReplicaType=%s expected=%d, running=%d, succeeded=%d , failed=%d", pytorchjob.Name, rtype, expected, running, succeeded, failed) if ContainsMasterSpec(replicas) { @@ -347,14 +351,14 @@ func (r *PyTorchJobReconciler) UpdateJobStatus(job interface{}, msg := fmt.Sprintf("PyTorchJob %s is running.", pytorchjob.Name) err := commonutil.UpdateJobConditions(jobStatus, commonv1.JobRunning, commonutil.JobRunningReason, msg) if err != nil { - commonutil.LoggerForJob(pytorchjob).Infof("Append job condition error: %v", err) + logger.Infof("Append job condition error: %v", err) return err } } // when master is succeed, the job is finished. if expected == 0 { msg := fmt.Sprintf("PyTorchJob %s is successfully completed.", pytorchjob.Name) - logrus.Info(msg) + logger.Info(msg) r.Recorder.Event(pytorchjob, corev1.EventTypeNormal, commonutil.JobSucceededReason, msg) if jobStatus.CompletionTime == nil { now := metav1.Now() @@ -362,7 +366,7 @@ func (r *PyTorchJobReconciler) UpdateJobStatus(job interface{}, } err := commonutil.UpdateJobConditions(jobStatus, commonv1.JobSucceeded, commonutil.JobSucceededReason, msg) if err != nil { - commonutil.LoggerForJob(pytorchjob).Infof("Append job condition error: %v", err) + logger.Infof("Append job condition error: %v", err) return err } trainingoperatorcommon.SuccessfulJobsCounterInc(pytorchjob.Namespace, pytorchv1.FrameworkName) @@ -370,9 +374,17 @@ func (r *PyTorchJobReconciler) UpdateJobStatus(job interface{}, } } } else { + if rtype == pytorchv1.PyTorchReplicaTypeWorker { - // TODO(gaocegege): Support SuccessPolicy - if expected == 0 { + worker0Completed, err := r.IsWorker0Completed(pytorchjob, replicas) + if err != nil { + logger.Warnf("check if worker 0 completed error %v", err) + return err + } + // Leave a succeeded condition for the following two cases: + // 1. If default success policy is used and worker 0 has completed. + // 2. If `SuccessPolicyAllWorkers` success policy is used and all workers are succeeded. + if expected == 0 || (worker0Completed && *pytorchjob.Spec.SuccessPolicy != pytorchv1.SuccessPolicyAllWorkers) { msg := fmt.Sprintf("TFJob %s/%s successfully completed.", pytorchjob.Namespace, pytorchjob.Name) r.recorder.Event(pytorchjob, corev1.EventTypeNormal, commonutil.JobSucceededReason, msg) @@ -430,6 +442,53 @@ func (r *PyTorchJobReconciler) UpdateJobStatus(job interface{}, return nil } +// IsWorker0Completed returns true if pod of worker0 succeeded and exited with 0 +func (p *PyTorchJobReconciler) IsWorker0Completed(job *pytorchv1.PyTorchJob, + replicas map[commonv1.ReplicaType]*commonv1.ReplicaSpec) (bool, error) { + worker0Completed := false + _, ok := replicas[pytorchv1.PyTorchReplicaTypeWorker] + if !ok { + return true, nil + } + podSlices, err := p.getPodSlices(job, + replicas[pytorchv1.PyTorchReplicaTypeWorker].Replicas) + if err != nil { + return false, err + } + for index, podSlice := range podSlices { + if len(podSlice) == 1 { + pod := podSlice[0] + exitCode := util.GetContainerExitCode(pod, pytorchv1.DefaultContainerName) + if index == 0 && exitCode == 0 && pod.Status.Phase == v1.PodSucceeded { + worker0Completed = true + } + } + } + return worker0Completed, nil +} + +// getPodSlices returns a slice, which element is the slice of pod. +// It gives enough information to caller to make decision to up/down scale resources. +func (p *PyTorchJobReconciler) getPodSlices( + job *pytorchv1.PyTorchJob, replicasNum *int32) ([][]*v1.Pod, error) { + logger := commonutil.LoggerForReplica(job, strings.ToLower(string(pytorchv1.PyTorchReplicaTypeWorker))) + + pods, err := p.GetPodsForJob(job) + if err != nil { + commonutil.LoggerForJob(job).Warnf("getPodsForTFJob error %v", err) + return nil, err + } + + // Get all pods for the type rt. + pods, err = p.JobController.FilterPodsForReplicaType(pods, strings.ToLower(string(pytorchv1.PyTorchReplicaTypeWorker))) + if err != nil { + return nil, err + } + + podSlices := p.GetPodSlices(pods, int(*replicasNum), logger) + return podSlices, nil +} + // ContainsMasterSpec returns true if the tfjob contains master spec. func ContainsMasterSpec(replicas map[commonv1.ReplicaType]*commonv1.ReplicaSpec) bool { if _, ok := replicas[pytorchv1.PyTorchReplicaTypeMaster]; ok { diff --git a/pkg/controller.v1/pytorch/pytorchjob_controller_test.go b/pkg/controller.v1/pytorch/pytorchjob_controller_test.go index 23dd037ee2..b827f368d2 100644 --- a/pkg/controller.v1/pytorch/pytorchjob_controller_test.go +++ b/pkg/controller.v1/pytorch/pytorchjob_controller_test.go @@ -191,7 +191,7 @@ var _ = Describe("PyTorchJob controller", func() { } job.Spec.PyTorchReplicaSpecs = map[commonv1.ReplicaType]*commonv1.ReplicaSpec{ pytorchv1.PyTorchReplicaTypeWorker: { - Replicas: int32Ptr(1), + Replicas: int32Ptr(2), Template: corev1.PodTemplateSpec{ Spec: corev1.PodSpec{ Containers: []corev1.Container{ @@ -275,8 +275,15 @@ var _ = Describe("PyTorchJob controller", func() { BlockOwnerDeletion: &trueVal, })) - // Test job status. + // Set the worker 0 succeeded. pod.Status.Phase = corev1.PodSucceeded + pod.Status.ContainerStatuses = make([]corev1.ContainerStatus, 1) + pod.Status.ContainerStatuses[0].Name = pytorchv1.DefaultContainerName + pod.Status.ContainerStatuses[0].State = corev1.ContainerState{ + Terminated: &corev1.ContainerStateTerminated{ + ExitCode: 0, + }, + } pod.ResourceVersion = "" Expect(testK8sClient.Status().Update(ctx, pod)).Should(Succeed()) Eventually(func() bool { @@ -289,6 +296,7 @@ var _ = Describe("PyTorchJob controller", func() { }, timeout, interval).Should(BeTrue()) // Check if the job is succeeded. cond := getCondition(created.Status, commonv1.JobSucceeded) + Expect(cond).NotTo(BeNil()) Expect(cond.Status).To(Equal(corev1.ConditionTrue)) By("Deleting the PyTorchJob") Expect(testK8sClient.Delete(ctx, job)).Should(Succeed()) diff --git a/pkg/controller.v1/tensorflow/tfjob_controller.go b/pkg/controller.v1/tensorflow/tfjob_controller.go index 6e7f105d35..3e7c8767b7 100644 --- a/pkg/controller.v1/tensorflow/tfjob_controller.go +++ b/pkg/controller.v1/tensorflow/tfjob_controller.go @@ -394,12 +394,6 @@ func (r *TFJobReconciler) UpdateJobStatus(job interface{}, replicas map[commonv1 logger := commonutil.LoggerForJob(tfJob) - worker0Completed, err := r.IsWorker0Completed(tfJob, replicas) - if err != nil { - logger.Warnf("check if worker 0 completed error %v", err) - return err - } - // Set StartTime. if jobStatus.StartTime == nil { now := metav1.Now() @@ -469,6 +463,12 @@ func (r *TFJobReconciler) UpdateJobStatus(job interface{}, replicas map[commonv1 } } else { if rtype == tensorflowv1.TFReplicaTypeWorker { + worker0Completed, err := r.IsWorker0Completed(tfJob, replicas) + if err != nil { + logger.Warnf("check if worker 0 completed error %v", err) + return err + } + // Leave a succeeded condition for the following two cases: // 1. If default success policy is used and worker 0 has completed. // 2. If `SuccessPolicyAllWorkers` success policy is used and all workers are succeeded. @@ -640,7 +640,7 @@ func (r *TFJobReconciler) IsWorker0Completed(tfjob *tensorflowv1.TFJob, replicas for index, podSlice := range podSlices { if len(podSlice) == 1 { pod := podSlice[0] - exitCode := getContainerExitCode(pod) + exitCode := util.GetContainerExitCode(pod, tfv1.DefaultContainerName) if index == 0 && exitCode == 0 && pod.Status.Phase == v1.PodSucceeded { worker0Completed = true } diff --git a/pkg/controller.v1/tensorflow/util.go b/pkg/controller.v1/tensorflow/util.go index 3bf0f727fc..6e60cd6f1b 100644 --- a/pkg/controller.v1/tensorflow/util.go +++ b/pkg/controller.v1/tensorflow/util.go @@ -46,18 +46,6 @@ func ContainsChiefOrMasterSpec(replicas map[commonv1.ReplicaType]*commonv1.Repli return false } -// originally from pkg/controller.v1/tensorflow/pod.go (deleted) -func getContainerExitCode(pod *corev1.Pod) int32 { - var exitCode int32 = 0xbeef // magic number - for _, status := range pod.Status.ContainerStatuses { - state := status.State - if status.Name == tfv1.DefaultContainerName && state.Terminated != nil { - exitCode = state.Terminated.ExitCode - } - } - return exitCode -} - // originally from pkg/controller.v1/tensorflow/pod.go (deleted) func setRestartPolicy(podTemplateSpec *corev1.PodTemplateSpec, spec *commonv1.ReplicaSpec) { // This is necessary since restartPolicyExitCode is not supported in v1.PodTemplateSpec From afeff31ae18ee8b1ab90715ebaebad9f179db765 Mon Sep 17 00:00:00 2001 From: cegao Date: Fri, 10 Dec 2021 15:35:36 +0800 Subject: [PATCH 2/6] fix: Update Signed-off-by: cegao --- pkg/apis/pytorch/v1/openapi_generated.go | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/pkg/apis/pytorch/v1/openapi_generated.go b/pkg/apis/pytorch/v1/openapi_generated.go index 943f4d170c..9626a9c1a6 100644 --- a/pkg/apis/pytorch/v1/openapi_generated.go +++ b/pkg/apis/pytorch/v1/openapi_generated.go @@ -410,7 +410,7 @@ func schema_pkg_apis_pytorch_v1_ElasticPolicy(ref common.ReferenceCallback) comm }, "metrics": { SchemaProps: spec.SchemaProps{ - Description: "metrics contains the specifications for which to use to calculate the desired replica count (the maximum replica count across all metrics will be used). The desired replica count is calculated multiplying the ratio between the target value and the current value by the current number of pods. Ergo, metrics used must decrease as the pod count is increased, and vice-versa. See the individual metric source types for more information about how each type of metric must respond. If not set, the default metric will be set to 80% average CPU utilization.", + Description: "Metrics contains the specifications which are used to calculate the desired replica count (the maximum replica count across all metrics will be used). The desired replica count is calculated with multiplying the ratio between the target value and the current value by the current number of pods. Ergo, metrics used must decrease as the pod count is increased, and vice-versa. See the individual metric source types for more information about how each type of metric must respond. If not set, the HPA will not be created.", Type: []string{"array"}, Items: &spec.SchemaOrArray{ Schema: &spec.Schema{ @@ -537,6 +537,12 @@ func schema_pkg_apis_pytorch_v1_PyTorchJobSpec(ref common.ReferenceCallback) com Ref: ref("github.com/kubeflow/common/pkg/apis/common/v1.RunPolicy"), }, }, + "successPolicy": { + SchemaProps: spec.SchemaProps{ + Type: []string{"string"}, + Format: "", + }, + }, "elasticPolicy": { SchemaProps: spec.SchemaProps{ Ref: ref("github.com/kubeflow/training-operator/pkg/apis/pytorch/v1.ElasticPolicy"), From 26a48cc462724cc29a65671012077c8a1e64bd8d Mon Sep 17 00:00:00 2001 From: cegao Date: Fri, 10 Dec 2021 15:37:54 +0800 Subject: [PATCH 3/6] Update Signed-off-by: cegao --- manifests/base/crds/kubeflow.org_pytorchjobs.yaml | 3 +++ 1 file changed, 3 insertions(+) diff --git a/manifests/base/crds/kubeflow.org_pytorchjobs.yaml b/manifests/base/crds/kubeflow.org_pytorchjobs.yaml index daedf9b93b..0904abd711 100644 --- a/manifests/base/crds/kubeflow.org_pytorchjobs.yaml +++ b/manifests/base/crds/kubeflow.org_pytorchjobs.yaml @@ -7324,6 +7324,9 @@ spec: format: int32 type: integer type: object + successPolicy: + description: SuccessPolicy is the success policy. + type: string required: - pytorchReplicaSpecs type: object From 89d046d4c293cc6398b0c25a7d77d437e4e07841 Mon Sep 17 00:00:00 2001 From: cegao Date: Fri, 10 Dec 2021 15:40:52 +0800 Subject: [PATCH 4/6] Update Signed-off-by: cegao --- pkg/apis/pytorch/v1/zz_generated.deepcopy.go | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/pkg/apis/pytorch/v1/zz_generated.deepcopy.go b/pkg/apis/pytorch/v1/zz_generated.deepcopy.go index 1fc845ff7b..5d078eb5f6 100644 --- a/pkg/apis/pytorch/v1/zz_generated.deepcopy.go +++ b/pkg/apis/pytorch/v1/zz_generated.deepcopy.go @@ -160,6 +160,11 @@ func (in *PyTorchJobList) DeepCopyObject() runtime.Object { func (in *PyTorchJobSpec) DeepCopyInto(out *PyTorchJobSpec) { *out = *in in.RunPolicy.DeepCopyInto(&out.RunPolicy) + if in.SuccessPolicy != nil { + in, out := &in.SuccessPolicy, &out.SuccessPolicy + *out = new(SuccessPolicy) + **out = **in + } if in.ElasticPolicy != nil { in, out := &in.ElasticPolicy, &out.ElasticPolicy *out = new(ElasticPolicy) From e45128ec93de470aa6fc134feb9009a42f1e320c Mon Sep 17 00:00:00 2001 From: cegao Date: Mon, 27 Dec 2021 16:04:49 +0800 Subject: [PATCH 5/6] fix: Remove v1 Signed-off-by: cegao --- pkg/controller.v1/pytorch/pytorchjob_controller.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pkg/controller.v1/pytorch/pytorchjob_controller.go b/pkg/controller.v1/pytorch/pytorchjob_controller.go index 79bce7406b..c5b1dd44db 100644 --- a/pkg/controller.v1/pytorch/pytorchjob_controller.go +++ b/pkg/controller.v1/pytorch/pytorchjob_controller.go @@ -470,7 +470,7 @@ func (p *PyTorchJobReconciler) IsWorker0Completed(job *pytorchv1.PyTorchJob, // getPodSlices returns a slice, which element is the slice of pod. // It gives enough information to caller to make decision to up/down scale resources. func (p *PyTorchJobReconciler) getPodSlices( - job *pytorchv1.PyTorchJob, replicasNum *int32) ([][]*v1.Pod, error) { + job *pytorchv1.PyTorchJob, replicasNum *int32) ([][]*corev1.Pod, error) { logger := commonutil.LoggerForReplica(job, strings.ToLower(string(pytorchv1.PyTorchReplicaTypeWorker))) pods, err := p.GetPodsForJob(job) From 5b9b1cc5baf19e88325b4322dfd0139e917874f8 Mon Sep 17 00:00:00 2001 From: cegao Date: Mon, 27 Dec 2021 16:10:52 +0800 Subject: [PATCH 6/6] fix: Remove v1 Signed-off-by: cegao --- pkg/controller.v1/pytorch/pytorchjob_controller.go | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/pkg/controller.v1/pytorch/pytorchjob_controller.go b/pkg/controller.v1/pytorch/pytorchjob_controller.go index c5b1dd44db..a0453d4fb6 100644 --- a/pkg/controller.v1/pytorch/pytorchjob_controller.go +++ b/pkg/controller.v1/pytorch/pytorchjob_controller.go @@ -27,7 +27,6 @@ import ( commonutil "github.com/kubeflow/common/pkg/util" "github.com/sirupsen/logrus" corev1 "k8s.io/api/core/v1" - v1 "k8s.io/api/core/v1" "k8s.io/apimachinery/pkg/api/equality" "k8s.io/apimachinery/pkg/api/errors" "k8s.io/apimachinery/pkg/api/meta" @@ -459,7 +458,7 @@ func (p *PyTorchJobReconciler) IsWorker0Completed(job *pytorchv1.PyTorchJob, if len(podSlice) == 1 { pod := podSlice[0] exitCode := util.GetContainerExitCode(pod, pytorchv1.DefaultContainerName) - if index == 0 && exitCode == 0 && pod.Status.Phase == v1.PodSucceeded { + if index == 0 && exitCode == 0 && pod.Status.Phase == corev1.PodSucceeded { worker0Completed = true } }