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

flyteplugins/go/tasks/pluginmachinery/flytek8s/config/config.go

@@ -6,6 +6,7 @@
 package config
 
 import (
+	schedulerConfig "github.com/flyteorg/flyte/flyteplugins/go/tasks/plugins/k8s/batchscheduler"
 	"time"
 
 	v1 "k8s.io/api/core/v1"
@@ -206,6 +207,8 @@ type K8sPluginConfig struct {
 
 	// SendObjectEvents indicates whether to send k8s object events in TaskExecutionEvent updates (similar to kubectl get events).
 	SendObjectEvents bool `json:"send-object-events" pflag:",If true, will send k8s object events in TaskExecutionEvent updates."`
+
+	BatchScheduler schedulerConfig.Config `json:"batchScheduler,omitempty"`
 }
 
 // FlyteCoPilotConfig specifies configuration for the Flyte CoPilot system. FlyteCoPilot, allows running flytekit-less containers

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

@@ -2,7 +2,6 @@ package k8s
 
 import (
 	"context"
-
 	"sigs.k8s.io/controller-runtime/pkg/client"
 
 	pluginsCore "github.com/flyteorg/flyte/flyteplugins/go/tasks/pluginmachinery/core"

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

@@ -0,0 +1,21 @@
+package batchscheduler
+
+type Config struct {
+	Scheduler  string `json:"scheduler,omitempty" pflag:", Specify batch scheduler to"`
+	Parameters string `json:"parameters,omitempty" pflag:", Specify static parameters"`
+}
+
+func NewConfig() Config {
+	return Config{
+		Scheduler:  "",
+		Parameters: "",
+	}
+}
+
+func (b *Config) GetScheduler() string {
+	return b.Scheduler
+}
+
+func (b *Config) GetParameters() string {
+	return b.Parameters
+}

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

@@ -0,0 +1,15 @@
+package batchscheduler
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestNewConfig(t *testing.T) {
+	t.Run("New scheduler plugin config", func(t *testing.T) {
+		config := NewConfig()
+		assert.Equal(t, "", config.GetScheduler())
+		assert.Equal(t, "", config.GetParameters())
+	})
+}

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

@@ -0,0 +1,23 @@
+package batchscheduler
+
+import (
+	"context"
+	"github.com/flyteorg/flyte/flyteplugins/go/tasks/plugins/k8s/batchscheduler/scheduler"
+	"github.com/flyteorg/flyte/flyteplugins/go/tasks/plugins/k8s/batchscheduler/scheduler/yunikorn"
+	"sigs.k8s.io/controller-runtime/pkg/client"
+)
+
+type SchedulerManager interface {
+	// Mutate is responsible for mutating the object to be scheduled.
+	// It will add the necessary annotations, labels, etc. to the object.
+	Mutate(ctx context.Context, object client.Object) error
+}
+
+func NewSchedulerManager(cfg *Config) SchedulerManager {
+	switch cfg.GetScheduler() {
+	case yunikorn.Yunikorn:
+		return yunikorn.NewSchedulerManager(cfg)
+	default:
+		return scheduler.NewNoopSchedulerManager()
+	}
+}

flyteplugins/go/tasks/plugins/k8s/batchscheduler/scheduler/noop_scheduler.go

@@ -0,0 +1,16 @@
+package scheduler
+
+import (
+	"context"
+	"sigs.k8s.io/controller-runtime/pkg/client"
+)
+
+type NoopSchedulerManager struct{}
+
+func NewNoopSchedulerManager() *NoopSchedulerManager {
+	return &NoopSchedulerManager{}
+}
+
+func (p *NoopSchedulerManager) Mutate(ctx context.Context, object client.Object) error {
+	return nil
+}

flyteplugins/go/tasks/plugins/k8s/batchscheduler/scheduler/yunikorn/rayhandler.go

@@ -0,0 +1,101 @@
+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"
+)
+
+func MutateRayJob(parameters string, 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[TaskGroupParameters] = parameters
+	meta.Annotations[AppID] = appID
+	return nil
+}
+
+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(a v1.ResourceList, b v1.ResourceList) v1.ResourceList {
+	result := a
+	for name, value := range a {
+		sum := &value
+		if value2, ok := b[name]; ok {
+			sum.Add(value2)
+			result[name] = *sum
+		} else {
+			result[name] = value
+		}
+	}
+	for name, value := range b {
+		if _, ok := a[name]; !ok {
+			result[name] = value
+		}
+	}
+	return result
+}

flyteplugins/go/tasks/plugins/k8s/batchscheduler/scheduler/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/scheduler/yunikorn/taskgroup_test.go

@@ -0,0 +1,46 @@
+package yunikorn
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestMarshal(t *testing.T) {
+	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/scheduler/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/scheduler/yunikorn/utils_test.go

@@ -0,0 +1,53 @@
+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) {
+		for _, tt := range tests {
+			got := GenerateTaskGroupAppID()
+			if len(got) <= 0 {
+				t.Error("Ray app ID is empty")
+			} 
+		}
+	})
+}