Updating annotations of pods belonging Ray cluster in order to adopting Yunikorn Gang scheduling

flyteplugins/go/tasks/pluginmachinery/k8s/plugin.go

@@ -5,6 +5,7 @@ import (
 
 	"sigs.k8s.io/controller-runtime/pkg/client"
 
+	"github.com/flyteorg/flyte/flyteidl/gen/pb-go/flyteidl/core"
 	pluginsCore "github.com/flyteorg/flyte/flyteplugins/go/tasks/pluginmachinery/core"
 	"github.com/flyteorg/flyte/flyteplugins/go/tasks/pluginmachinery/io"
 	"github.com/flyteorg/flyte/flytestdlib/storage"
@@ -187,3 +188,11 @@ func MaybeUpdatePhaseVersionFromPluginContext(phaseInfo *pluginsCore.PhaseInfo,
 	MaybeUpdatePhaseVersion(phaseInfo, &pluginState)
 	return nil
 }
+
+type YunikornScheduablePlugin interface {
+	MutateResourceForYunikorn(ctx context.Context, object client.Object, taskTmpl *core.TaskTemplate) (client.Object, error)
+}
+
+type KueueScheduablePlugin interface {
+	MutateResourceForKueue(ctx context.Context, object client.Object, taskTmpl *core.TaskTemplate) (client.Object, error)
+}

flyteplugins/go/tasks/plugins/k8s/batchscheduler/config.go

@@ -0,0 +1,23 @@
+package batchscheduler
+
+type Config struct {
+	Scheduler string                      `json:"scheduler,omitempty" pflag:", Specify batch scheduler to"`
+	Default   SchedulingConfig            `json:"default,omitempty" pflag:", Specify default scheduling config which batch scheduler adopts"`
+	NameSpace map[string]SchedulingConfig `json:"Namespace,omitempty" pflag:"-, Specify namespace scheduling config"`
+	Domain    map[string]SchedulingConfig `json:"Domain,omitempty" pflag:"-, Specify domain scheduling config"`
+}
+
+type SchedulingConfig struct {
+	KueueConfig    `json:"Kueue,omitempty" pflag:", Specify Kueue scheduling scheduling config"`
+	YunikornConfig `json:"Yunikorn,omitempty" pflag:", Yunikorn scheduling config"`
+}
+
+type KueueConfig struct {
+	PriorityClassName string `json:"Priority,omitempty" pflag:", Kueue Prioty class"`
+	Queue             string `json:"Queue,omitempty" pflag:", Specify batch scheduler to"`
+}
+
+type YunikornConfig struct {
+	Parameters string `json:"parameters,omitempty" pflag:", Specify gangscheduling policy"`
+	Queue      string `json:"queue,omitempty" pflag:", Specify leaf queue to submit to"`
+}

flyteplugins/go/tasks/plugins/k8s/batchscheduler/kueue/helper.go

@@ -0,0 +1,16 @@
+package kueue
+
+import (
+	rayv1 "github.com/ray-project/kuberay/ray-operator/apis/ray/v1"
+
+	"github.com/flyteorg/flyte/flyteplugins/go/tasks/plugins/k8s/batchscheduler/utils"
+)
+
+const (
+	QueueName         = "kueue.x-k8s.io/queue-name"
+	PriorityClassName = "kueue.x-k8s.io/priority-class"
+)
+
+func UpdateKueueLabels(labels map[string]string, app *rayv1.RayJob) {
+	utils.UpdateLabels(labels, &app.ObjectMeta)
+}

flyteplugins/go/tasks/plugins/k8s/batchscheduler/utils/helper.go

@@ -0,0 +1,30 @@
+package utils
+
+import (
+	v1 "k8s.io/api/core/v1"
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+)
+
+func UpdateLabels(wanted map[string]string, objectMeta *metav1.ObjectMeta) {
+	for key, value := range wanted {
+		if _, exist := objectMeta.Labels[key]; !exist {
+			objectMeta.Labels[key] = value
+		}
+	}
+}
+
+func UpdateAnnotations(wanted map[string]string, objectMeta *metav1.ObjectMeta) {
+	for key, value := range wanted {
+		if _, exist := objectMeta.Annotations[key]; !exist {
+			objectMeta.Annotations[key] = value
+		}
+	}
+}
+
+func UpdatePodTemplateAnnotatations(wanted map[string]string, pod *v1.PodTemplateSpec) {
+	UpdateAnnotations(wanted, &pod.ObjectMeta)
+}
+
+func UpdatePodTemplateLabels(wanted map[string]string, pod *v1.PodTemplateSpec) {
+	UpdateLabels(wanted, &pod.ObjectMeta)
+}

flyteplugins/go/tasks/plugins/k8s/batchscheduler/yunikorn/helper.go

@@ -0,0 +1,133 @@
+package yunikorn
+
+import (
+	"encoding/json"
+
+	rayv1 "github.com/ray-project/kuberay/ray-operator/apis/ray/v1"
+	v1 "k8s.io/api/core/v1"
+	"k8s.io/apimachinery/pkg/api/resource"
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+
+	"github.com/flyteorg/flyte/flyteplugins/go/tasks/plugins/k8s/batchscheduler/utils"
+)
+
+const (
+	Yunikorn            = "yunikorn"
+	AppID               = "yunikorn.apache.org/app-id"
+	Queue               = "yunikorn.apache.org/queue"
+	TaskGroupNameKey    = "yunikorn.apache.org/task-group-name"
+	TaskGroupsKey       = "yunikorn.apache.org/task-groups"
+	TaskGroupParameters = "yunikorn.apache.org/schedulingPolicyParameters"
+)
+
+func MutateRayJob(app *rayv1.RayJob) error {
+	appID := GenerateTaskGroupAppID()
+	rayjobSpec := &app.Spec
+	appSpec := rayjobSpec.RayClusterSpec
+	TaskGroups := make([]TaskGroup, 1)
+	for index := range appSpec.WorkerGroupSpecs {
+		worker := &appSpec.WorkerGroupSpecs[index]
+		worker.Template.Spec.SchedulerName = Yunikorn
+		meta := worker.Template.ObjectMeta
+		spec := worker.Template.Spec
+		name := GenerateTaskGroupName(false, index)
+		TaskGroups = append(TaskGroups, TaskGroup{
+			Name:                      name,
+			MinMember:                 *worker.Replicas,
+			Labels:                    meta.Labels,
+			Annotations:               meta.Annotations,
+			MinResource:               Allocation(spec.Containers),
+			NodeSelector:              spec.NodeSelector,
+			Affinity:                  spec.Affinity,
+			TopologySpreadConstraints: spec.TopologySpreadConstraints,
+		})
+		meta.Annotations[TaskGroupNameKey] = name
+		meta.Annotations[AppID] = appID
+	}
+	headSpec := &appSpec.HeadGroupSpec
+	headSpec.Template.Spec.SchedulerName = Yunikorn
+	meta := headSpec.Template.ObjectMeta
+	spec := headSpec.Template.Spec
+	headName := GenerateTaskGroupName(true, 0)
+	res := Allocation(spec.Containers)
+	if ok := *appSpec.EnableInTreeAutoscaling; ok {
+		res2 := v1.ResourceList{
+			v1.ResourceCPU:    resource.MustParse("500m"),
+			v1.ResourceMemory: resource.MustParse("512Mi"),
+		}
+		res = Add(res, res2)
+	}
+	TaskGroups[0] = TaskGroup{
+		Name:                      headName,
+		MinMember:                 1,
+		Labels:                    meta.Labels,
+		Annotations:               meta.Annotations,
+		MinResource:               res,
+		NodeSelector:              spec.NodeSelector,
+		Affinity:                  spec.Affinity,
+		TopologySpreadConstraints: spec.TopologySpreadConstraints,
+	}
+	meta.Annotations[TaskGroupNameKey] = headName
+	info, err := json.Marshal(TaskGroups)
+	if err != nil {
+		return err
+	}
+	meta.Annotations[TaskGroupsKey] = string(info[:])
+	meta.Annotations[AppID] = appID
+	return nil
+}
+
+func UpdateGangSchedulingParameters(parameters string, objectMeta *metav1.ObjectMeta) {
+	if len(parameters) == 0 {
+		return
+	}
+	utils.UpdateAnnotations(
+		map[string]string{TaskGroupParameters: parameters},
+		objectMeta,
+	)
+}
+
+func UpdateAnnotations(labels map[string]string, app *rayv1.RayJob) {
+	appSpec := app.Spec.RayClusterSpec
+	headSpec := appSpec.HeadGroupSpec
+	utils.UpdatePodTemplateAnnotatations(labels, &headSpec.Template)
+	for index := range appSpec.WorkerGroupSpecs {
+		worker := appSpec.WorkerGroupSpecs[index]
+		utils.UpdatePodTemplateAnnotatations(labels, &worker.Template)
+	}
+}
+
+func Allocation(containers []v1.Container) v1.ResourceList {
+	totalResources := v1.ResourceList{}
+	for _, c := range containers {
+		for name, q := range c.Resources.Limits {
+			if _, exists := totalResources[name]; !exists {
+				totalResources[name] = q.DeepCopy()
+				continue
+			}
+			total := totalResources[name]
+			total.Add(q)
+			totalResources[name] = total
+		}
+	}
+	return totalResources
+}
+
+func Add(left v1.ResourceList, right v1.ResourceList) v1.ResourceList {
+	result := left
+	for name, value := range left {
+		sum := value
+		if value2, ok := right[name]; ok {
+			sum.Add(value2)
+			result[name] = sum
+		} else {
+			result[name] = value
+		}
+	}
+	for name, value := range right {
+		if _, ok := left[name]; !ok {
+			result[name] = value
+		}
+	}
+	return result
+}

flyteplugins/go/tasks/plugins/k8s/batchscheduler/yunikorn/taskgroup.go

@@ -0,0 +1,24 @@
+package yunikorn
+
+import (
+	"encoding/json"
+
+	v1 "k8s.io/api/core/v1"
+)
+
+type TaskGroup struct {
+	Name                      string
+	MinMember                 int32
+	Labels                    map[string]string
+	Annotations               map[string]string
+	MinResource               v1.ResourceList
+	NodeSelector              map[string]string
+	Tolerations               []v1.Toleration
+	Affinity                  *v1.Affinity
+	TopologySpreadConstraints []v1.TopologySpreadConstraint
+}
+
+func Marshal(taskGroups []TaskGroup) ([]byte, error) {
+	info, err := json.Marshal(taskGroups)
+	return info, err
+}

flyteplugins/go/tasks/plugins/k8s/batchscheduler/yunikorn/taskgroup_test.go

@@ -0,0 +1,52 @@
+package yunikorn
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	v1 "k8s.io/api/core/v1"
+	"k8s.io/apimachinery/pkg/api/resource"
+)
+
+func TestMarshal(t *testing.T) {
+	res := v1.ResourceList{
+		v1.ResourceCPU:    resource.MustParse("500m"),
+		v1.ResourceMemory: resource.MustParse("512Mi"),
+	}
+	t1 := TaskGroup{
+		Name:                      "tg1",
+		MinMember:                 int32(1),
+		Labels:                    map[string]string{"attr": "value"},
+		Annotations:               map[string]string{"attr": "value"},
+		MinResource:               res,
+		NodeSelector:              map[string]string{"node": "gpunode"},
+		Tolerations:               nil,
+		Affinity:                  nil,
+		TopologySpreadConstraints: nil,
+	}
+	t2 := TaskGroup{
+		Name:                      "tg2",
+		MinMember:                 int32(1),
+		Labels:                    map[string]string{"attr": "value"},
+		Annotations:               map[string]string{"attr": "value"},
+		MinResource:               res,
+		NodeSelector:              map[string]string{"node": "gpunode"},
+		Tolerations:               nil,
+		Affinity:                  nil,
+		TopologySpreadConstraints: nil,
+	}
+	var tests = []struct {
+		input []TaskGroup
+	}{
+		{input: nil},
+		{input: []TaskGroup{}},
+		{input: []TaskGroup{t1}},
+		{input: []TaskGroup{t1, t2}},
+	}
+	t.Run("Serialize task groups", func(t *testing.T) {
+		for _, tt := range tests {
+			_, err := Marshal(tt.input)
+			assert.Nil(t, err)
+		}
+	})
+}

flyteplugins/go/tasks/plugins/k8s/batchscheduler/yunikorn/utils.go

@@ -0,0 +1,23 @@
+package yunikorn
+
+import (
+	"fmt"
+
+	"github.com/google/uuid"
+)
+
+const (
+	TaskGroupGenericName = "task-group"
+)
+
+func GenerateTaskGroupName(master bool, index int) string {
+	if master {
+		return fmt.Sprintf("%s-%s", TaskGroupGenericName, "head")
+	}
+	return fmt.Sprintf("%s-%s-%d", TaskGroupGenericName, "worker", index)
+}
+
+func GenerateTaskGroupAppID() string {
+	uid := uuid.New().String()
+	return fmt.Sprintf("%s-%s", TaskGroupGenericName, uid)
+}

flyteplugins/go/tasks/plugins/k8s/batchscheduler/yunikorn/utils_test.go

@@ -0,0 +1,51 @@
+package yunikorn
+
+import (
+	"fmt"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestGenerateTaskGroupName(t *testing.T) {
+	type inputFormat struct {
+		isMaster bool
+		index    int
+	}
+	var tests = []struct {
+		input  inputFormat
+		expect string
+	}{
+		{
+			input:  inputFormat{isMaster: true, index: 0},
+			expect: fmt.Sprintf("%s-%s", TaskGroupGenericName, "head"),
+		},
+		{
+			input:  inputFormat{isMaster: true, index: 1},
+			expect: fmt.Sprintf("%s-%s", TaskGroupGenericName, "head"),
+		},
+		{
+			input:  inputFormat{isMaster: false, index: 0},
+			expect: fmt.Sprintf("%s-%s-%d", TaskGroupGenericName, "worker", 0),
+		},
+		{
+			input:  inputFormat{isMaster: false, index: 1},
+			expect: fmt.Sprintf("%s-%s-%d", TaskGroupGenericName, "worker", 1),
+		},
+	}
+	t.Run("Generate ray task group name", func(t *testing.T) {
+		for _, tt := range tests {
+			got := GenerateTaskGroupName(tt.input.isMaster, tt.input.index)
+			assert.Equal(t, tt.expect, got)
+		}
+	})
+}
+
+func TestGenerateTaskGroupAppID(t *testing.T) {
+	t.Run("Generate ray app ID", func(t *testing.T) {
+		got := GenerateTaskGroupAppID()
+		if len(got) <= 0 {
+			t.Error("Ray app ID is empty")
+		}
+	})
+}