diff --git a/flyteplugins/go/tasks/plugins/k8s/batchscheduler/scheduler/yunikorn/rayhandler.go b/flyteplugins/go/tasks/plugins/k8s/batchscheduler/scheduler/yunikorn/rayhandler.go index c39a5a937b..0884a68df1 100644 --- a/flyteplugins/go/tasks/plugins/k8s/batchscheduler/scheduler/yunikorn/rayhandler.go +++ b/flyteplugins/go/tasks/plugins/k8s/batchscheduler/scheduler/yunikorn/rayhandler.go @@ -19,14 +19,14 @@ func MutateRayJob(parameters string, app *rayv1.RayJob) error { 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, + 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 @@ -42,21 +42,17 @@ func MutateRayJob(parameters string, app *rayv1.RayJob) error { v1.ResourceCPU: resource.MustParse("500m"), v1.ResourceMemory: resource.MustParse("512Mi"), } - //tmp, _ := json.Marshal(res2) - //meta.Annotations["tmp"] = string(tmp) - //tmp, _ = json.Marshal(Add(res, res2)) - //meta.Annotations["Sum"] = string(tmp) 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, + 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)