update test in ray

Signed-off-by: yuteng <[email protected]>
flyteorg · Jul 29, 2024 · a202a11 · a202a11
1 parent a5c5ab0
commit a202a11
Showing 1 changed file with 32 additions and 0 deletions.
diff --git a/flyteplugins/go/tasks/plugins/k8s/ray/ray_test.go b/flyteplugins/go/tasks/plugins/k8s/ray/ray_test.go
@@ -27,6 +27,8 @@ import (
 	mocks2 "github.com/flyteorg/flyte/flyteplugins/go/tasks/pluginmachinery/k8s/mocks"
 	"github.com/flyteorg/flyte/flyteplugins/go/tasks/pluginmachinery/tasklog"
 	"github.com/flyteorg/flyte/flyteplugins/go/tasks/pluginmachinery/utils"
+	schedulerConfig "github.com/flyteorg/flyte/flyteplugins/go/tasks/plugins/k8s/ray/batchscheduler/config"
+	"github.com/flyteorg/flyte/flyteplugins/go/tasks/plugins/k8s/ray/batchscheduler/scheduler/yunikorn"
 )
 
 const (
@@ -474,6 +476,36 @@ func TestDefaultStartParameters(t *testing.T) {
 	assert.Equal(t, ray.Spec.RayClusterSpec.WorkerGroupSpecs[0].Template.Spec.Tolerations, toleration)
 }
 
+func TestYunikornAnnotationsCreate(t *testing.T) {
+	assert.NoError(t, SetConfig(&Config{
+		BatchScheduler: schedulerConfig.Config{
+			Scheduler:  "yunikorn",
+			Parameters: "gangSchedulingStyle=Soft",
+		},
+	}))
+	rayJobResourceHandler := rayJobResourceHandler{}
+	rayJob := &plugins.RayJob{
+		RayCluster: &plugins.RayCluster{
+			HeadGroupSpec:     &plugins.HeadGroupSpec{},
+			WorkerGroupSpec:   []*plugins.WorkerGroupSpec{{GroupName: workerGroupName, Replicas: 3, MinReplicas: 3, MaxReplicas: 3}},
+			EnableAutoscaling: true,
+		},
+		ShutdownAfterJobFinishes: true,
+		TtlSecondsAfterFinished:  120,
+	}
+
+	taskTemplate := dummyRayTaskTemplate("ray-id", rayJob)
+	RayResource, err := rayJobResourceHandler.BuildResource(context.TODO(), dummyRayTaskContext(taskTemplate, resourceRequirements, nil, "", serviceAccount))
+	assert.Nil(t, err)
+	ray, ok := RayResource.(*rayv1.RayJob)
+	assert.True(t, ok)
+	headAnnotations := ray.Spec.RayClusterSpec.HeadGroupSpec.Template.ObjectMeta.Annotations
+	workerAnnotations := ray.Spec.RayClusterSpec.WorkerGroupSpecs[0].Template.ObjectMeta.Annotations
+	assert.Equal(t, yunikorn.GenerateTaskGroupName(true, 0), headAnnotations[yunikorn.TaskGroupNameKey])
+	assert.Equal(t, "gangSchedulingStyle=Soft", headAnnotations[yunikorn.TaskGroupPrarameters])
+	assert.Equal(t, yunikorn.GenerateTaskGroupName(false, 0), workerAnnotations[yunikorn.TaskGroupNameKey])
+}
+
 func TestInjectLogsSidecar(t *testing.T) {
 	rayJobObj := transformRayJobToCustomObj(dummyRayCustomObj())
 	params := []struct {