Merge pull request #5 from OyutianO/main

add create queue unit function when  a new pytorchjob created

Author: denkensk

cmd/app/server.go

@@ -2,6 +2,7 @@ package app
 
 import (
 	v1 "github.com/kube-queue/pytorch-operator-extension/pkg/pytorch-operator/apis/pytorch/v1"
+	"k8s.io/client-go/kubernetes"
 	"k8s.io/client-go/rest"
 	"k8s.io/client-go/tools/cache"
 	"k8s.io/client-go/tools/clientcmd"
@@ -17,11 +18,6 @@ import (
 	"k8s.io/klog/v2"
 )
 
-const (
-	ConsumerRefKind       = v1.Kind
-	ConsumerRefAPIVersion = v1.GroupName + "/" + v1.GroupVersion
-)
-
 // Run runs the server.
 func Run(opt *options.ServerOption) error {
 	var restConfig *rest.Config
@@ -35,6 +31,11 @@ func Run(opt *options.ServerOption) error {
 		}
 	}
 
+	k8sClientSet, err := kubernetes.NewForConfig(restConfig)
+	if err != nil {
+		return err
+	}
+
 	queueClient, err := queueversioned.NewForConfig(restConfig)
 	if err != nil {
 		return err
@@ -50,7 +51,9 @@ func Run(opt *options.ServerOption) error {
 	pytorchJobInformerFactory := pytorchjobinformers.NewSharedInformerFactory(pytorchJobClient, 0)
 	pytorchJobInformer := pytorchJobInformerFactory.Kubeflow().V1().PyTorchJobs().Informer()
 
-	pytorchExtensionController := contorller.NewPyTorchExtensionController(queueInformerFactory.Scheduling().V1alpha1().QueueUnits(),
+	pytorchExtensionController := contorller.NewPyTorchExtensionController(
+		k8sClientSet,
+		queueInformerFactory.Scheduling().V1alpha1().QueueUnits(),
 		queueClient,
 		pytorchJobInformerFactory.Kubeflow().V1().PyTorchJobs(),
 		pytorchJobClient,
@@ -62,8 +65,8 @@ func Run(opt *options.ServerOption) error {
 				switch qu := obj.(type) {
 				case *v1alpha1.QueueUnit:
 					if qu.Spec.ConsumerRef != nil &&
-						qu.Spec.ConsumerRef.Kind == ConsumerRefKind &&
-						qu.Spec.ConsumerRef.APIVersion == ConsumerRefAPIVersion {
+						qu.Spec.ConsumerRef.Kind == contorller.ConsumerRefKind &&
+						qu.Spec.ConsumerRef.APIVersion == contorller.ConsumerRefAPIVersion {
 						return true
 					}
 					return false
@@ -90,9 +93,9 @@ func Run(opt *options.ServerOption) error {
 				}
 			},
 			Handler: cache.ResourceEventHandlerFuncs{
-				AddFunc:    pytorchExtensionController.AddPytorchJob,
-				UpdateFunc: pytorchExtensionController.UpdatePytorchJob,
-				DeleteFunc: pytorchExtensionController.DeletePytorchJob,
+				AddFunc:    pytorchExtensionController.AddPyTorchJob,
+				UpdateFunc: pytorchExtensionController.UpdatePyTorchJob,
+				DeleteFunc: pytorchExtensionController.DeletePyTorchJob,
 			},
 		},
 	)

go.mod

@@ -8,6 +8,7 @@ require (
 	k8s.io/apimachinery v0.21.2
 	k8s.io/client-go v0.21.2
 	k8s.io/code-generator v0.21.2
+	k8s.io/klog v1.0.0
 	k8s.io/klog/v2 v2.8.0
 	k8s.io/kube-openapi v0.0.0-20210305001622-591a79e4bda7
 	k8s.io/sample-controller v0.21.2