fix: use configured RayCluster head group service account when autosc…

…aling closes #242
ray-project · May 12, 2022 · 4aa4af1 · 4aa4af1
1 parent c7356da
commit 4aa4af1
Showing 5 changed files with 132 additions and 13 deletions.
diff --git a/ray-operator/controllers/common/rbac.go b/ray-operator/controllers/common/rbac.go
@@ -7,11 +7,21 @@ import (
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 )
 
+// GetHeadGroupServiceAccountName returns the head group service account if it exists.
+// Otherwise, it returns the name of the cluster itself.
+func GetHeadGroupServiceAccountName(cluster *v1alpha1.RayCluster) string {
+	headGroupServiceAccountName := cluster.Spec.HeadGroupSpec.Template.Spec.ServiceAccountName
+	if headGroupServiceAccountName != "" {
+		return headGroupServiceAccountName
+	}
+	return cluster.Name
+}
+
 // BuildServiceAccount creates a new ServiceAccount for a head pod with autoscaler.
 func BuildServiceAccount(cluster *v1alpha1.RayCluster) (*v1.ServiceAccount, error) {
 	sa := &v1.ServiceAccount{
 		ObjectMeta: metav1.ObjectMeta{
-			Name:      cluster.Name,
+			Name:      GetHeadGroupServiceAccountName(cluster),
 			Namespace: cluster.Namespace,
 			Labels: map[string]string{
 				RayClusterLabelKey: cluster.Name,

diff --git a/ray-operator/controllers/common/rbac_test.go b/ray-operator/controllers/common/rbac_test.go
@@ -0,0 +1,61 @@
+package common
+
+import (
+	"testing"
+
+	rayiov1alpha1 "github.com/ray-project/kuberay/ray-operator/api/raycluster/v1alpha1"
+
+	corev1 "k8s.io/api/core/v1"
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+)
+
+func TestGetHeadGroupServiceAccountName(t *testing.T) {
+	tests := map[string]struct {
+		input *rayiov1alpha1.RayCluster
+		want  string
+	}{
+		"Ray cluster with head group service account": {
+			input: &rayiov1alpha1.RayCluster{
+				ObjectMeta: metav1.ObjectMeta{
+					Name:      "raycluster-sample",
+					Namespace: "default",
+				},
+				Spec: rayiov1alpha1.RayClusterSpec{
+					HeadGroupSpec: rayiov1alpha1.HeadGroupSpec{
+						Template: corev1.PodTemplateSpec{
+							Spec: corev1.PodSpec{
+								ServiceAccountName: "my-service-account",
+							},
+						},
+					},
+				},
+			},
+			want: "my-service-account",
+		},
+		"Ray cluster without head group service account": {
+			input: &rayiov1alpha1.RayCluster{
+				ObjectMeta: metav1.ObjectMeta{
+					Name:      "raycluster-sample",
+					Namespace: "default",
+				},
+				Spec: rayiov1alpha1.RayClusterSpec{
+					HeadGroupSpec: rayiov1alpha1.HeadGroupSpec{
+						Template: corev1.PodTemplateSpec{
+							Spec: corev1.PodSpec{},
+						},
+					},
+				},
+			},
+			want: "raycluster-sample",
+		},
+	}
+
+	for name, tc := range tests {
+		t.Run(name, func(t *testing.T) {
+			got := GetHeadGroupServiceAccountName(tc.input)
+			if got != tc.want {
+				t.Fatalf("got %s, want %s", got, tc.want)
+			}
+		})
+	}
+}
diff --git a/ray-operator/controllers/raycluster_controller.go b/ray-operator/controllers/raycluster_controller.go
@@ -95,7 +95,7 @@ func (r *RayClusterReconciler) Reconcile(ctx context.Context, request ctrl.Reque
 	}
 
 	if instance.DeletionTimestamp != nil && !instance.DeletionTimestamp.IsZero() {
-		log.Info("RayCluser is being deleted, just ignore", "cluster name", request.Name)
+		log.Info("RayCluster is being deleted, just ignore", "cluster name", request.Name)
 		return ctrl.Result{}, nil
 	}
 	if err := r.reconcileAutoscalerServiceAccount(instance); err != nil {
@@ -578,7 +578,8 @@ func (r *RayClusterReconciler) reconcileAutoscalerServiceAccount(instance *rayio
 	}
 
 	serviceAccount := &corev1.ServiceAccount{}
-	namespacedName := types.NamespacedName{Namespace: instance.Namespace, Name: instance.Name}
+	namespacedName := types.NamespacedName{Namespace: instance.Namespace, Name: common.GetHeadGroupServiceAccountName(instance)}
+
 	if err := r.Get(context.TODO(), namespacedName, serviceAccount); err != nil {
 		if !errors.IsNotFound(err) {
 			return err

diff --git a/ray-operator/controllers/raycluster_controller_fake_test.go b/ray-operator/controllers/raycluster_controller_fake_test.go
@@ -24,9 +24,11 @@ import (
 	"github.com/ray-project/kuberay/ray-operator/controllers/common"
 	"github.com/ray-project/kuberay/ray-operator/pkg/client/clientset/versioned/scheme"
 	"github.com/stretchr/testify/assert"
+	"k8s.io/apimachinery/pkg/api/errors"
 	"k8s.io/apimachinery/pkg/labels"
 	"k8s.io/apimachinery/pkg/runtime"
 	"k8s.io/apimachinery/pkg/selection"
+	"k8s.io/apimachinery/pkg/types"
 	"k8s.io/client-go/tools/record"
 	ctrl "sigs.k8s.io/controller-runtime"
 	clientFake "sigs.k8s.io/controller-runtime/pkg/client/fake"
@@ -43,15 +45,17 @@ import (
 )
 
 var (
-	namespaceStr     string
-	instanceName     string
-	headGroupNameStr string
-	groupNameStr     string
-	expectReplicaNum int32
-	testPods         []runtime.Object
-	testRayCluster   *rayiov1alpha1.RayCluster
-	workerSelector   labels.Selector
-	workersToDelete  []string
+	namespaceStr            string
+	instanceName            string
+	enableInTreeAutoscaling bool
+	headGroupNameStr        string
+	headGroupServiceAccount string
+	groupNameStr            string
+	expectReplicaNum        int32
+	testPods                []runtime.Object
+	testRayCluster          *rayiov1alpha1.RayCluster
+	workerSelector          labels.Selector
+	workersToDelete         []string
 )
 
 func setupTest(t *testing.T) {
@@ -61,7 +65,9 @@ func setupTest(t *testing.T) {
 
 	namespaceStr = "default"
 	instanceName = "raycluster-sample"
+	enableInTreeAutoscaling = true
 	headGroupNameStr = "head-group"
+	headGroupServiceAccount = "head-service-account"
 	groupNameStr = "small-group"
 	expectReplicaNum = 3
 	workersToDelete = []string{"pod1", "pod2"}
@@ -152,7 +158,8 @@ func setupTest(t *testing.T) {
 			Namespace: namespaceStr,
 		},
 		Spec: rayiov1alpha1.RayClusterSpec{
-			RayVersion: "1.0",
+			RayVersion:              "1.0",
+			EnableInTreeAutoscaling: &enableInTreeAutoscaling,
 			HeadGroupSpec: rayiov1alpha1.HeadGroupSpec{
 				ServiceType: "ClusterIP",
 				Replicas:    pointer.Int32Ptr(1),
@@ -166,6 +173,7 @@ func setupTest(t *testing.T) {
 				},
 				Template: corev1.PodTemplateSpec{
 					Spec: corev1.PodSpec{
+						ServiceAccountName: headGroupServiceAccount,
 						Containers: []corev1.Container{
 							{
 								Name:    "ray-head",
@@ -555,3 +563,33 @@ func getNotFailedPodItemNum(podList corev1.PodList) int {
 
 	return count
 }
+
+func TestReconcile_AutoscalerServiceAccount(t *testing.T) {
+	setupTest(t)
+	defer tearDown(t)
+
+	fakeClient := clientFake.NewClientBuilder().WithRuntimeObjects(testPods...).Build()
+
+	saNamespacedName := types.NamespacedName{
+		Name:      headGroupServiceAccount,
+		Namespace: namespaceStr,
+	}
+	sa := corev1.ServiceAccount{}
+	err := fakeClient.Get(context.Background(), saNamespacedName, &sa)
+
+	assert.True(t, errors.IsNotFound(err), "Head group service account should not exist yet")
+
+	testRayClusterReconciler := &RayClusterReconciler{
+		Client:   fakeClient,
+		Recorder: &record.FakeRecorder{},
+		Scheme:   scheme.Scheme,
+		Log:      ctrl.Log.WithName("controllers").WithName("RayCluster"),
+	}
+
+	err = testRayClusterReconciler.reconcileAutoscalerServiceAccount(testRayCluster)
+	assert.Nil(t, err, "Fail to reconcile autoscaler ServiceAccount")
+
+	err = fakeClient.Get(context.Background(), saNamespacedName, &sa)
+
+	assert.Nil(t, err, "Fail to get head group ServiceAccount after reconciliation")
+}
diff --git a/ray-operator/controllers/raycluster_controller_test.go b/ray-operator/controllers/raycluster_controller_test.go
@@ -67,6 +67,7 @@ var _ = Context("Inside the default namespace", func() {
 				},
 				Template: corev1.PodTemplateSpec{
 					Spec: corev1.PodSpec{
+						ServiceAccountName: "head-service-account",
 						Containers: []corev1.Container{
 							{
 								Name:    "ray-head",
@@ -174,6 +175,14 @@ var _ = Context("Inside the default namespace", func() {
 			Expect(pod.Status.Phase).Should(Or(Equal(v1.PodPending), Equal(v1.PodRunning)))
 		})
 
+		It("should create the head group's specified K8s ServiceAccount if it doesn't exist", func() {
+			saName := common.GetHeadGroupServiceAccountName(myRayCluster)
+			sa := &corev1.ServiceAccount{}
+			Eventually(
+				getResourceFunc(ctx, client.ObjectKey{Name: saName, Namespace: "default"}, sa),
+				time.Second*15, time.Millisecond*500).Should(BeNil(), "My head group ServiceAccount = %v", saName)
+		})
+
 		It("should re-create a deleted worker", func() {
 			Eventually(
 				listResourceFunc(ctx, &workerPods, filterLabels, &client.ListOptions{Namespace: "default"}),