Use podGroup instead of PDB in v1beta2

Gopkg.toml

@@ -21,6 +21,10 @@ required = [
   name = "github.com/stretchr/testify"
   version = "1.2.2"
 
+[[constraint]]
+  name = "github.com/kubernetes-sigs/kube-batch"
+  version = "v0.3"
+
 [[constraint]]
   name = "github.com/sirupsen/logrus"
   version = "v1.0.4"

cmd/tf-operator.v1beta2/app/server.go

@@ -25,8 +25,10 @@ import (
 	"github.com/kubeflow/tf-operator/pkg/client/clientset/versioned/scheme"
 	tfjobinformers "github.com/kubeflow/tf-operator/pkg/client/informers/externalversions"
 	controller "github.com/kubeflow/tf-operator/pkg/controller.v1beta2/tensorflow"
+	"github.com/kubeflow/tf-operator/pkg/util/kbutil"
 	"github.com/kubeflow/tf-operator/pkg/util/signals"
 	"github.com/kubeflow/tf-operator/pkg/version"
+	kubebatchclient "github.com/kubernetes-sigs/kube-batch/pkg/client/clientset/versioned"
 	log "github.com/sirupsen/logrus"
 	"k8s.io/api/core/v1"
 	crdclient "k8s.io/apiextensions-apiserver/pkg/client/clientset/clientset"
@@ -91,11 +93,14 @@ func Run(opt *options.ServerOption) error {
 	}
 
 	// Create clients.
-	kubeClientSet, leaderElectionClientSet, tfJobClientSet, err := createClientSets(kcfg)
+	kubeClientSet, leaderElectionClientSet, tfJobClientSet, kubeBatchClientSet, err := createClientSets(kcfg)
 	if err != nil {
 		return err
 	}
 
+	//Sets kube-batch Interface, used for accessing podGroup CRD APIs.
+	kbutil.SetKubeBatchClientInterface(kubeBatchClientSet)
+
 	// Create informer factory.
 	kubeInformerFactory := kubeinformers.NewFilteredSharedInformerFactory(kubeClientSet, resyncPeriod, opt.Namespace, nil)
 	tfJobInformerFactory := tfjobinformers.NewSharedInformerFactory(tfJobClientSet, resyncPeriod)
@@ -161,32 +166,36 @@ func Run(opt *options.ServerOption) error {
 	return nil
 }
 
-func createClientSets(config *restclientset.Config) (kubeclientset.Interface, kubeclientset.Interface, tfjobclientset.Interface, error) {
+func createClientSets(config *restclientset.Config) (kubeclientset.Interface, kubeclientset.Interface, tfjobclientset.Interface, kubebatchclient.Interface, error) {
 
 	crdClient, err := crdclient.NewForConfig(config)
 
 	if err != nil {
-		return nil, nil, nil, err
+		return nil, nil, nil, nil, err
 	}
 
 	checkCRDExists(crdClient, v1beta2.TFCRD)
 
 	kubeClientSet, err := kubeclientset.NewForConfig(restclientset.AddUserAgent(config, "tf-operator"))
 	if err != nil {
-		return nil, nil, nil, err
+		return nil, nil, nil, nil, err
 	}
 
 	leaderElectionClientSet, err := kubeclientset.NewForConfig(restclientset.AddUserAgent(config, "leader-election"))
 	if err != nil {
-		return nil, nil, nil, err
+		return nil, nil, nil, nil, err
 	}
 
 	tfJobClientSet, err := tfjobclientset.NewForConfig(config)
 	if err != nil {
-		return nil, nil, nil, err
+		return nil, nil, nil, nil, err
 	}
 
-	return kubeClientSet, leaderElectionClientSet, tfJobClientSet, nil
+	kubeBatchClientSet, err := kubebatchclient.NewForConfig(restclientset.AddUserAgent(config, "kube-batch"))
+	if err != nil {
+		return nil, nil, nil, nil, err
+	}
+	return kubeClientSet, leaderElectionClientSet, tfJobClientSet, kubeBatchClientSet, nil
 }
 
 func checkCRDExists(clientset crdclient.Interface, crdName string) {

examples/crd/crd-podgroup.yaml

@@ -0,0 +1,39 @@
+apiVersion: apiextensions.k8s.io/v1beta1
+kind: CustomResourceDefinition
+metadata:
+  name: podgroups.scheduling.incubator.k8s.io
+spec:
+  group: scheduling.incubator.k8s.io
+  names:
+    kind: PodGroup
+    plural: podgroups
+  scope: Namespaced
+  validation:
+    openAPIV3Schema:
+      properties:
+        apiVersion:
+          type: string
+        kind:
+          type: string
+        metadata:
+          type: object
+        spec:
+          properties:
+            minMember:
+              format: int32
+              type: integer
+          type: object
+        status:
+          properties:
+            succeeded:
+              format: int32
+              type: integer
+            failed:
+              format: int32
+              type: integer
+            running:
+              format: int32
+              type: integer
+          type: object
+      type: object
+  version: v1alpha1

examples/distribution_strategy/distributed_tfjob.yaml

@@ -10,6 +10,9 @@ spec:
       replicas: 3 
       restartPolicy: Never
       template:
+        metadata:
+              annotations:
+                scheduling.k8s.io/group-name: "distributed-training"
         spec:
           containers:
             - name: tensorflow

pkg/common/jobcontroller/jobcontroller.go

@@ -5,6 +5,7 @@ import (
 	"fmt"
 	"strings"
 
+	"github.com/kubernetes-sigs/kube-batch/pkg/apis/scheduling/v1alpha1"
 	log "github.com/sirupsen/logrus"
 	"k8s.io/api/core/v1"
 	"k8s.io/api/policy/v1beta1"
@@ -23,6 +24,7 @@ import (
 	"k8s.io/kubernetes/pkg/controller"
 
 	"github.com/kubeflow/tf-operator/pkg/control"
+	"github.com/kubeflow/tf-operator/pkg/util/kbutil"
 )
 
 // Common Interface to be implemented by all operators.
@@ -197,6 +199,31 @@ func (jc *JobController) GenLabels(jobName string) map[string]string {
 	}
 }
 
+func (jc *JobController) SyncPodGroup(job metav1.Object, minAvailableReplicas int32) (*v1alpha1.PodGroup, error) {
+
+	kubeBatchClientInterface := kbutil.GetKubeBatchClientInterface()
+	// Check whether podGroup exists or not
+	podGroup, err := kubeBatchClientInterface.SchedulingV1alpha1().PodGroups(job.GetNamespace()).Get(job.GetName(), metav1.GetOptions{})
+	if err == nil {
+		return podGroup, nil
+	}
+
+	// create podGroup for gang scheduling by kube-batch
+	minAvailable := intstr.FromInt(int(minAvailableReplicas))
+	createPodGroup := &v1alpha1.PodGroup{
+		ObjectMeta: metav1.ObjectMeta{
+			Name: job.GetName(),
+			OwnerReferences: []metav1.OwnerReference{
+				*jc.GenOwnerReference(job),
+			},
+		},
+		Spec: v1alpha1.PodGroupSpec{
+			MinMember: minAvailable.IntVal,
+		},
+	}
+	return kubeBatchClientInterface.SchedulingV1alpha1().PodGroups(job.GetNamespace()).Create(createPodGroup)
+}
+
 // SyncPdb will create a PDB for gang scheduling by kube-arbitrator.
 func (jc *JobController) SyncPdb(job metav1.Object, minAvailableReplicas int32) (*v1beta1.PodDisruptionBudget, error) {
 	labelJobName := jc.Controller.GetJobNameLabelKey()
@@ -231,6 +258,25 @@ func (jc *JobController) SyncPdb(job metav1.Object, minAvailableReplicas int32)
 	return jc.KubeClientSet.PolicyV1beta1().PodDisruptionBudgets(job.GetNamespace()).Create(createPdb)
 }
 
+func (jc *JobController) DeletePodGroup(job metav1.Object) error {
+	kubeBatchClientInterface := kbutil.GetKubeBatchClientInterface()
+
+	//check whether podGroup exists or not
+	_, err := kubeBatchClientInterface.SchedulingV1alpha1().PodGroups(job.GetNamespace()).Get(job.GetName(), metav1.GetOptions{})
+	if err != nil && k8serrors.IsNotFound(err) {
+		return nil
+	}
+
+	log.Infof("Deleting PodGroup %s", job.GetName())
+
+	//delete podGroup
+	err = kubeBatchClientInterface.SchedulingV1alpha1().PodGroups(job.GetNamespace()).Delete(job.GetName(), &metav1.DeleteOptions{})
+	if err != nil {
+		return fmt.Errorf("unable to delete PodGroup: %v", err)
+	}
+	return nil
+}
+
 func (jc *JobController) DeletePdb(job metav1.Object) error {
 
 	// Check the pdb exist or not

pkg/controller.v1beta2/tensorflow/controller.go

@@ -302,9 +302,9 @@ func (tc *TFController) syncTFJob(key string) (bool, error) {
 
 	if tc.Config.EnableGangScheduling {
 		minAvailableReplicas := getTotalReplicas(tfjob)
-		_, err := tc.SyncPdb(tfjob, minAvailableReplicas)
+		_, err := tc.SyncPodGroup(tfjob, minAvailableReplicas)
 		if err != nil {
-			logger.Warnf("Sync pdb %v: %v", tfjob.Name, err)
+			logger.Warnf("Sync PodGroup %v: %v", tfjob.Name, err)
 		}
 	}
 
@@ -362,12 +362,12 @@ func (tc *TFController) reconcileTFJobs(tfjob *tfv1beta2.TFJob) error {
 		}
 
 		if tc.Config.EnableGangScheduling {
-			tc.Recorder.Event(tfjob, v1.EventTypeNormal, "JobTerminated", "Job is terminated, deleting pdb")
-			if err := tc.DeletePdb(tfjob); err != nil {
-				tc.Recorder.Eventf(tfjob, v1.EventTypeWarning, "FailedDeletePdb", "Error deleting: %v", err)
+			tc.Recorder.Event(tfjob, v1.EventTypeNormal, "JobTerminated", "Job is terminated, deleting PodGroup")
+			if err := tc.DeletePodGroup(tfjob); err != nil {
+				tc.Recorder.Eventf(tfjob, v1.EventTypeWarning, "FailedDeletePodGroup", "Error deleting: %v", err)
 				return err
 			} else {
-				tc.Recorder.Eventf(tfjob, v1.EventTypeNormal, "SuccessfulDeletePdb", "Deleted pdb: %v", tfjob.Name)
+				tc.Recorder.Eventf(tfjob, v1.EventTypeNormal, "SuccessfulDeletePodGroup", "Deleted pdb: %v", tfjob.Name)
 
 			}
 		}

pkg/util/kbutil/kbutil.go

@@ -0,0 +1,31 @@
+// Copyright 2019 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 kbutil
+
+import (
+	kubebatchclient "github.com/kubernetes-sigs/kube-batch/pkg/client/clientset/versioned"
+)
+
+var kubeBatchClientSet kubebatchclient.Interface
+
+//SetKubeBatchClientSet sets kube-batch client set.
+func SetKubeBatchClientInterface(clientset kubebatchclient.Interface) {
+	kubeBatchClientSet = clientset
+}
+
+// GetKubeBatchClientSet gets kube-batch client set.
+func GetKubeBatchClientInterface() kubebatchclient.Interface {
+	return kubeBatchClientSet
+}