Skip to content

Commit

Permalink
Add XGBoost controller
Browse files Browse the repository at this point in the history
1. Move major controller logics from kubeflow/xgboost-operator to this repo.
2. Adapt to kubebuilder 3.0.0 change
3. Add xgboost train example
4. Leave some TODOs for code refactor later

Signed-off-by: Jiaxin Shan <[email protected]>
  • Loading branch information
Jeffwan committed Jul 6, 2021
1 parent d72641c commit bd6a9c7
Show file tree
Hide file tree
Showing 9 changed files with 913 additions and 24 deletions.
41 changes: 38 additions & 3 deletions config/samples/kubeflow.org_v1_xgboostjob.yaml
Original file line number Diff line number Diff line change
@@ -1,7 +1,42 @@
apiVersion: kubeflow.org/v1
kind: XGBoostJob
metadata:
name: xgboostjob-sample
name: xgboost-dist-iris-test-train
spec:
# Add fields here
foo: bar
xgbReplicaSpecs:
Master:
replicas: 1
restartPolicy: Never
template:
spec:
containers:
- name: xgboostjob
image: docker.io/merlintang/xgboost-dist-iris:1.1
ports:
- containerPort: 9991
name: xgboostjob-port
imagePullPolicy: Always
args:
- --job_type=Train
- --xgboost_parameter=objective:multi:softprob,num_class:3
- --n_estimators=10
- --learning_rate=0.1
- --model_path=/tmp/xgboost-model
- --model_storage_type=local
Worker:
replicas: 2
restartPolicy: ExitCode
template:
spec:
containers:
- name: xgboostjob
image: docker.io/merlintang/xgboost-dist-iris:1.1
ports:
- containerPort: 9991
name: xgboostjob-port
imagePullPolicy: Always
args:
- --job_type=Train
- --xgboost_parameter="objective:multi:softprob,num_class:3"
- --n_estimators=10
- --learning_rate=0.1
8 changes: 3 additions & 5 deletions main.go
Original file line number Diff line number Diff line change
Expand Up @@ -78,11 +78,9 @@ func main() {
os.Exit(1)
}

if err = (&xgboostcontroller.XGBoostJobReconciler{
Client: mgr.GetClient(),
Log: ctrl.Log.WithName("controllers").WithName("XGBoostJob"),
Scheme: mgr.GetScheme(),
}).SetupWithManager(mgr); err != nil {
// TODO: We need a general manager. all rest reconciler addsToManager
// Based on the user configuration, we start different controllers
if err = xgboostcontroller.NewReconciler(mgr).SetupWithManager(mgr); err != nil {
setupLog.Error(err, "unable to create controller", "controller", "XGBoostJob")
os.Exit(1)
}
Expand Down
110 changes: 110 additions & 0 deletions pkg/controller.v1/xgboost/expectation.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,110 @@
/*
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 xgboost

import (
"fmt"
commonv1 "github.com/kubeflow/common/pkg/apis/common/v1"
"github.com/kubeflow/common/pkg/controller.v1/common"
"github.com/kubeflow/common/pkg/controller.v1/expectation"
v1 "github.com/kubeflow/tf-operator/pkg/apis/xgboost/v1"
"github.com/sirupsen/logrus"
corev1 "k8s.io/api/core/v1"
metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
utilruntime "k8s.io/apimachinery/pkg/util/runtime"
"sigs.k8s.io/controller-runtime/pkg/event"
"sigs.k8s.io/controller-runtime/pkg/reconcile"
)

// satisfiedExpectations returns true if the required adds/dels for the given job have been observed.
// Add/del counts are established by the controller at sync time, and updated as controllees are observed by the controller
// manager.
func (r *XGBoostJobReconciler) satisfiedExpectations(xgbJob *v1.XGBoostJob) bool {
satisfied := false
key, err := common.KeyFunc(xgbJob)
if err != nil {
utilruntime.HandleError(fmt.Errorf("couldn't get key for job object %#v: %v", xgbJob, err))
return false
}
for rtype := range xgbJob.Spec.XGBReplicaSpecs {
// Check the expectations of the pods.
expectationPodsKey := expectation.GenExpectationPodsKey(key, string(rtype))
satisfied = satisfied || r.Expectations.SatisfiedExpectations(expectationPodsKey)
// Check the expectations of the services.
expectationServicesKey := expectation.GenExpectationServicesKey(key, string(rtype))
satisfied = satisfied || r.Expectations.SatisfiedExpectations(expectationServicesKey)
}
return satisfied
}

// onDependentCreateFunc modify expectations when dependent (pod/service) creation observed.
func onDependentCreateFunc(r reconcile.Reconciler) func(event.CreateEvent) bool {
return func(e event.CreateEvent) bool {
xgbr, ok := r.(*XGBoostJobReconciler)
if !ok {
return true
}
rtype := e.Object.GetLabels()[commonv1.ReplicaTypeLabel]
if len(rtype) == 0 {
return false
}

logrus.Info("Update on create function ", xgbr.ControllerName(), " create object ", e.Object.GetName())
if controllerRef := metav1.GetControllerOf(e.Object); controllerRef != nil {
var expectKey string
if _, ok := e.Object.(*corev1.Pod); ok {
expectKey = expectation.GenExpectationPodsKey(e.Object.GetNamespace()+"/"+controllerRef.Name, rtype)
}

if _, ok := e.Object.(*corev1.Service); ok {
expectKey = expectation.GenExpectationServicesKey(e.Object.GetNamespace()+"/"+controllerRef.Name, rtype)
}
xgbr.Expectations.CreationObserved(expectKey)
return true
}

return true
}
}

// onDependentDeleteFunc modify expectations when dependent (pod/service) deletion observed.
func onDependentDeleteFunc(r reconcile.Reconciler) func(event.DeleteEvent) bool {
return func(e event.DeleteEvent) bool {
xgbr, ok := r.(*XGBoostJobReconciler)
if !ok {
return true
}

rtype := e.Object.GetLabels()[commonv1.ReplicaTypeLabel]
if len(rtype) == 0 {
return false
}

logrus.Info("Update on deleting function ", xgbr.ControllerName(), " delete object ", e.Object.GetName())
if controllerRef := metav1.GetControllerOf(e.Object); controllerRef != nil {
var expectKey string
if _, ok := e.Object.(*corev1.Pod); ok {
expectKey = expectation.GenExpectationPodsKey(e.Object.GetNamespace()+"/"+controllerRef.Name, rtype)
}

if _, ok := e.Object.(*corev1.Service); ok {
expectKey = expectation.GenExpectationServicesKey(e.Object.GetNamespace()+"/"+controllerRef.Name, rtype)
}

xgbr.Expectations.DeletionObserved(expectKey)
return true
}

return true
}
}
225 changes: 225 additions & 0 deletions pkg/controller.v1/xgboost/job.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,225 @@
/*
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 xgboost

import (
"context"
"fmt"
"reflect"
"sigs.k8s.io/controller-runtime/pkg/log"

commonv1 "github.com/kubeflow/common/pkg/apis/common/v1"
commonutil "github.com/kubeflow/common/pkg/util"
logger "github.com/kubeflow/common/pkg/util"
xgboostv1 "github.com/kubeflow/tf-operator/pkg/apis/xgboost/v1"
"github.com/sirupsen/logrus"
corev1 "k8s.io/api/core/v1"
k8sv1 "k8s.io/api/core/v1"
"k8s.io/apimachinery/pkg/api/errors"
metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
"k8s.io/apimachinery/pkg/types"
"k8s.io/client-go/kubernetes/scheme"
"sigs.k8s.io/controller-runtime/pkg/event"
"sigs.k8s.io/controller-runtime/pkg/reconcile"
)

// Reasons for job events.
const (
FailedDeleteJobReason = "FailedDeleteJob"
SuccessfulDeleteJobReason = "SuccessfulDeleteJob"
// xgboostJobCreatedReason is added in a job when it is created.
xgboostJobCreatedReason = "XGBoostJobCreated"

xgboostJobSucceededReason = "XGBoostJobSucceeded"
xgboostJobRunningReason = "XGBoostJobRunning"
xgboostJobFailedReason = "XGBoostJobFailed"
xgboostJobRestartingReason = "XGBoostJobRestarting"
)

// DeleteJob deletes the job
func (r *XGBoostJobReconciler) DeleteJob(job interface{}) error {
xgboostjob, ok := job.(*xgboostv1.XGBoostJob)
if !ok {
return fmt.Errorf("%+v is not a type of XGBoostJob", xgboostjob)
}
if err := r.Delete(context.Background(), xgboostjob); err != nil {
r.recorder.Eventf(xgboostjob, corev1.EventTypeWarning, FailedDeleteJobReason, "Error deleting: %v", err)
r.Log.Error(err, "failed to delete job", "namespace", xgboostjob.Namespace, "name", xgboostjob.Name)
return err
}
r.recorder.Eventf(xgboostjob, corev1.EventTypeNormal, SuccessfulDeleteJobReason, "Deleted job: %v", xgboostjob.Name)
r.Log.Info("job deleted", "namespace", xgboostjob.Namespace, "name", xgboostjob.Name)
return nil
}

// GetJobFromInformerCache returns the Job from Informer Cache
func (r *XGBoostJobReconciler) GetJobFromInformerCache(namespace, name string) (metav1.Object, error) {
job := &xgboostv1.XGBoostJob{}
// Default reader for XGBoostJob is cache reader.
err := r.Get(context.Background(), types.NamespacedName{Namespace: namespace, Name: name}, job)
if err != nil {
if errors.IsNotFound(err) {
r.Log.Error(err, "xgboost job not found", "namespace", namespace, "name", name)
} else {
r.Log.Error(err, "failed to get job from api-server", "namespace", namespace, "name", name)
}
return nil, err
}
return job, nil
}

// GetJobFromAPIClient returns the Job from API server
func (r *XGBoostJobReconciler) GetJobFromAPIClient(namespace, name string) (metav1.Object, error) {
job := &xgboostv1.XGBoostJob{}

// TODO (Jeffwan@): consider to read from apiserver directly.
err := r.Client.Get(context.Background(), types.NamespacedName{Namespace: namespace, Name: name}, job)
if err != nil {
if errors.IsNotFound(err) {
r.Log.Error(err, "xgboost job not found", "namespace", namespace, "name", name)
} else {
r.Log.Error(err, "failed to get job from api-server", "namespace", namespace, "name", name)
}
return nil, err
}
return job, nil
}

// UpdateJobStatus updates the job status and job conditions
func (r *XGBoostJobReconciler) UpdateJobStatus(job interface{}, replicas map[commonv1.ReplicaType]*commonv1.ReplicaSpec, jobStatus *commonv1.JobStatus) error {
xgboostJob, ok := job.(*xgboostv1.XGBoostJob)
if !ok {
return fmt.Errorf("%+v is not a type of xgboostJob", xgboostJob)
}

for rtype, spec := range replicas {
status := jobStatus.ReplicaStatuses[rtype]

succeeded := status.Succeeded
expected := *(spec.Replicas) - succeeded
running := status.Active
failed := status.Failed

logrus.Infof("XGBoostJob=%s, ReplicaType=%s expected=%d, running=%d, succeeded=%d , failed=%d",
xgboostJob.Name, rtype, expected, running, succeeded, failed)

if rtype == commonv1.ReplicaType(xgboostv1.XGBoostReplicaTypeMaster) {
if running > 0 {
msg := fmt.Sprintf("XGBoostJob %s is running.", xgboostJob.Name)
err := commonutil.UpdateJobConditions(jobStatus, commonv1.JobRunning, xgboostJobRunningReason, msg)
if err != nil {
logger.LoggerForJob(xgboostJob).Infof("Append job condition error: %v", err)
return err
}
}
// when master is succeed, the job is finished.
if expected == 0 {
msg := fmt.Sprintf("XGBoostJob %s is successfully completed.", xgboostJob.Name)
logrus.Info(msg)
r.Recorder.Event(xgboostJob, k8sv1.EventTypeNormal, xgboostJobSucceededReason, msg)
if jobStatus.CompletionTime == nil {
now := metav1.Now()
jobStatus.CompletionTime = &now
}
err := commonutil.UpdateJobConditions(jobStatus, commonv1.JobSucceeded, xgboostJobSucceededReason, msg)
if err != nil {
logger.LoggerForJob(xgboostJob).Infof("Append job condition error: %v", err)
return err
}
return nil
}
}
if failed > 0 {
if spec.RestartPolicy == commonv1.RestartPolicyExitCode {
msg := fmt.Sprintf("XGBoostJob %s is restarting because %d %s replica(s) failed.", xgboostJob.Name, failed, rtype)
r.Recorder.Event(xgboostJob, k8sv1.EventTypeWarning, xgboostJobRestartingReason, msg)
err := commonutil.UpdateJobConditions(jobStatus, commonv1.JobRestarting, xgboostJobRestartingReason, msg)
if err != nil {
logger.LoggerForJob(xgboostJob).Infof("Append job condition error: %v", err)
return err
}
} else {
msg := fmt.Sprintf("XGBoostJob %s is failed because %d %s replica(s) failed.", xgboostJob.Name, failed, rtype)
r.Recorder.Event(xgboostJob, k8sv1.EventTypeNormal, xgboostJobFailedReason, msg)
if xgboostJob.Status.CompletionTime == nil {
now := metav1.Now()
xgboostJob.Status.CompletionTime = &now
}
err := commonutil.UpdateJobConditions(jobStatus, commonv1.JobFailed, xgboostJobFailedReason, msg)
if err != nil {
logger.LoggerForJob(xgboostJob).Infof("Append job condition error: %v", err)
return err
}
}
}
}

// Some workers are still running, leave a running condition.
msg := fmt.Sprintf("XGBoostJob %s is running.", xgboostJob.Name)
logger.LoggerForJob(xgboostJob).Infof(msg)

if err := commonutil.UpdateJobConditions(jobStatus, commonv1.JobRunning, xgboostJobRunningReason, msg); err != nil {
logger.LoggerForJob(xgboostJob).Error(err, "failed to update XGBoost Job conditions")
return err
}

return nil
}

// UpdateJobStatusInApiServer updates the job status in to cluster.
func (r *XGBoostJobReconciler) UpdateJobStatusInApiServer(job interface{}, jobStatus *commonv1.JobStatus) error {
xgboostjob, ok := job.(*xgboostv1.XGBoostJob)
if !ok {
return fmt.Errorf("%+v is not a type of XGBoostJob", xgboostjob)
}

// Job status passed in differs with status in job, update in basis of the passed in one.
if !reflect.DeepEqual(&xgboostjob.Status.JobStatus, jobStatus) {
xgboostjob = xgboostjob.DeepCopy()
xgboostjob.Status.JobStatus = *jobStatus.DeepCopy()
}

result := r.Update(context.Background(), xgboostjob)

if result != nil {
logger.LoggerForJob(xgboostjob).Error(result, "failed to update XGBoost Job conditions in the API server")
return result
}

return nil
}

// onOwnerCreateFunc modify creation condition.
func onOwnerCreateFunc(r reconcile.Reconciler) func(event.CreateEvent) bool {
return func(e event.CreateEvent) bool {
xgboostJob, ok := e.Object.(*xgboostv1.XGBoostJob)
if !ok {
return true
}
scheme.Scheme.Default(xgboostJob)
msg := fmt.Sprintf("xgboostJob %s is created.", e.Object.GetName())
logrus.Info(msg)
//specific the run policy

if xgboostJob.Spec.RunPolicy.CleanPodPolicy == nil {
xgboostJob.Spec.RunPolicy.CleanPodPolicy = new(commonv1.CleanPodPolicy)
xgboostJob.Spec.RunPolicy.CleanPodPolicy = &defaultCleanPodPolicy
}

if err := commonutil.UpdateJobConditions(&xgboostJob.Status.JobStatus, commonv1.JobCreated, xgboostJobCreatedReason, msg); err != nil {
log.Log.Error(err, "append job condition error")
return false
}
return true
}
}
Loading

0 comments on commit bd6a9c7

Please sign in to comment.