Skip to content

Commit

Permalink
[YUNIKORN-2933] Don't add duplicated taskGroup to app
Browse files Browse the repository at this point in the history
  • Loading branch information
@XbaoWu
XbaoWu authored and wuxiaobao committed Oct 18, 2024
1 parent 753e4f8 commit 8e12958
Show file tree
Hide file tree
Showing 2 changed files with 56 additions and 17 deletions.
27 changes: 21 additions & 6 deletions pkg/cache/application.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -48,7 +48,7 @@ type Application struct {
groups []string
taskMap map[string]*Task
tags map[string]string
taskGroups []TaskGroup
taskGroups map[string]TaskGroup
taskGroupsDefinition string
schedulingParamsDefinition string
placeholderOwnerReferences []metav1.OwnerReference
Expand Down Expand Up @@ -80,7 +80,7 @@ func NewApplication(appID, queueName, user string, groups []string, tags map[str
taskMap: taskMap,
tags: tags,
sm: newAppState(),
taskGroups: make([]TaskGroup, 0),
taskGroups: make(map[string]TaskGroup),
lock: &locking.RWMutex{},
schedulerAPI: scheduler,
placeholderTimeoutInSec: 0,
Expand Down Expand Up @@ -166,9 +166,16 @@ func (app *Application) GetSchedulingParamsDefinition() string {
func (app *Application) setTaskGroups(taskGroups []TaskGroup) {
app.lock.Lock()
defer app.lock.Unlock()
app.taskGroups = taskGroups
for _, taskGroup := range app.taskGroups {
app.placeholderAsk = common.Add(app.placeholderAsk, common.GetTGResource(taskGroup.MinResource, int64(taskGroup.MinMember)))
for _, taskGroup := range taskGroups {
// for duplicated task-group, will no longer be added to the app
if _, exists := app.taskGroups[taskGroup.Name]; exists {
log.Log(log.ShimCacheApplication).Warn("duplicate task-group within the task-groups",
zap.String("appID", app.applicationID),
zap.String("groupName", taskGroup.Name))
} else {
app.taskGroups[taskGroup.Name] = taskGroup
app.placeholderAsk = common.Add(app.placeholderAsk, common.GetTGResource(taskGroup.MinResource, int64(taskGroup.MinMember)))
}
}
}

Expand All @@ -181,7 +188,15 @@ func (app *Application) getPlaceholderAsk() *si.Resource {
func (app *Application) getTaskGroups() []TaskGroup {
app.lock.RLock()
defer app.lock.RUnlock()
return app.taskGroups

if len(app.taskGroups) > 0 {
taskGroups := make([]TaskGroup, 0, len(app.taskGroups))
for _, taskGroup := range app.taskGroups {
taskGroups = append(taskGroups, taskGroup)
}
return taskGroups
}
return nil
}

func (app *Application) setPlaceholderOwnerReferences(ref []metav1.OwnerReference) {
Expand Down
46 changes: 35 additions & 11 deletions pkg/cache/placeholder_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -119,9 +119,10 @@ func TestNewPlaceholder(t *testing.T) {
assert.Equal(t, app.placeholderAsk.Resources[siCommon.Memory].Value, int64(10*1024*1000*1000))
assert.Equal(t, app.placeholderAsk.Resources["pods"].Value, int64(10))

holder := newPlaceholder("ph-name", app, app.taskGroups[0])
tgs := app.getTaskGroups()
holder := newPlaceholder("ph-name", app, tgs[0])
assert.Equal(t, holder.appID, appID)
assert.Equal(t, holder.taskGroupName, app.taskGroups[0].Name)
assert.Equal(t, holder.taskGroupName, tgs[0].Name)
assert.Equal(t, holder.pod.Spec.SchedulerName, constants.SchedulerName)
assert.Equal(t, holder.pod.Name, "ph-name")
assert.Equal(t, holder.pod.Namespace, namespace)
Expand All @@ -132,7 +133,7 @@ func TestNewPlaceholder(t *testing.T) {
"labelKey1": "labelKeyValue1",
})
assert.Equal(t, len(holder.pod.Annotations), 7, "unexpected number of annotations")
assert.Equal(t, holder.pod.Annotations[constants.AnnotationTaskGroupName], app.taskGroups[0].Name)
assert.Equal(t, holder.pod.Annotations[constants.AnnotationTaskGroupName], tgs[0].Name)
assert.Equal(t, holder.pod.Annotations[constants.AnnotationPlaceholderFlag], constants.True)
assert.Equal(t, holder.pod.Annotations["annotationKey0"], "annotationValue0")
assert.Equal(t, holder.pod.Annotations["annotationKey1"], "annotationValue1")
Expand Down Expand Up @@ -163,7 +164,8 @@ func TestNewPlaceholderWithNodeSelectors(t *testing.T) {
"bob", testGroups, map[string]string{constants.AppTagNamespace: namespace}, mockedSchedulerAPI)
app.setTaskGroups(taskGroups)

holder := newPlaceholder("ph-name", app, app.taskGroups[0])
tgs := app.getTaskGroups()
holder := newPlaceholder("ph-name", app, tgs[0])
assert.Equal(t, len(holder.pod.Spec.NodeSelector), 2)
assert.Equal(t, holder.pod.Spec.NodeSelector["nodeType"], "test")
assert.Equal(t, holder.pod.Spec.NodeSelector["nodeState"], "healthy")
Expand All @@ -178,7 +180,8 @@ func TestNewPlaceholderWithTolerations(t *testing.T) {
"bob", testGroups, map[string]string{constants.AppTagNamespace: namespace}, mockedSchedulerAPI)
app.setTaskGroups(taskGroups)

holder := newPlaceholder("ph-name", app, app.taskGroups[0])
tgs := app.getTaskGroups()
holder := newPlaceholder("ph-name", app, tgs[0])
assert.Equal(t, len(holder.pod.Spec.Tolerations), 1)
tlr := holder.pod.Spec.Tolerations[0]
assert.Equal(t, tlr.Key, "key1")
Expand All @@ -196,7 +199,8 @@ func TestNewPlaceholderWithAffinity(t *testing.T) {
"bob", testGroups, map[string]string{constants.AppTagNamespace: namespace}, mockedSchedulerAPI)
app.setTaskGroups(taskGroups)

holder := newPlaceholder("ph-name", app, app.taskGroups[0])
tgs := app.getTaskGroups()
holder := newPlaceholder("ph-name", app, tgs[0])
assert.Equal(t, len(holder.pod.Spec.Affinity.PodAffinity.RequiredDuringSchedulingIgnoredDuringExecution), 1)
term := holder.pod.Spec.Affinity.PodAffinity.RequiredDuringSchedulingIgnoredDuringExecution
assert.Equal(t, term[0].TopologyKey, "topologyKey")
Expand All @@ -215,14 +219,16 @@ func TestNewPlaceholderTaskGroupsDefinition(t *testing.T) {
app := NewApplication(appID, queue,
"bob", testGroups, map[string]string{constants.AppTagNamespace: namespace}, mockedSchedulerAPI)
app.setTaskGroups(taskGroups)
holder := newPlaceholder("ph-name", app, app.taskGroups[0])
tgs := app.getTaskGroups()
holder := newPlaceholder("ph-name", app, tgs[0])
assert.Equal(t, "", holder.pod.Annotations[constants.AnnotationTaskGroups])

app = NewApplication(appID, queue,
"bob", testGroups, map[string]string{constants.AppTagNamespace: namespace}, mockedSchedulerAPI)
app.setTaskGroups(taskGroups)
app.setTaskGroupsDefinition("taskGroupsDef")
holder = newPlaceholder("ph-name", app, app.taskGroups[0])
tgs = app.getTaskGroups()
holder = newPlaceholder("ph-name", app, tgs[0])
assert.Equal(t, "taskGroupsDef", holder.pod.Annotations[constants.AnnotationTaskGroups])
var priority *int32
assert.Equal(t, priority, holder.pod.Spec.Priority)
Expand All @@ -234,7 +240,9 @@ func TestNewPlaceholderExtendedResources(t *testing.T) {
app := NewApplication(appID, queue,
"bob", testGroups, map[string]string{constants.AppTagNamespace: namespace}, mockedSchedulerAPI)
app.setTaskGroups(taskGroups)
holder := newPlaceholder("ph-name", app, app.taskGroups[0])

tgs := app.getTaskGroups()
holder := newPlaceholder("ph-name", app, tgs[0])
assert.Equal(t, len(holder.pod.Spec.Containers[0].Resources.Requests), 5, "expected requests not found")
assert.Equal(t, len(holder.pod.Spec.Containers[0].Resources.Limits), 5, "expected limits not found")
assert.Equal(t, holder.pod.Spec.Containers[0].Resources.Limits[gpu], holder.pod.Spec.Containers[0].Resources.Requests[gpu], "gpu: expected same value for request and limit")
Expand Down Expand Up @@ -271,7 +279,8 @@ func TestNewPlaceholderWithPriorityClassName(t *testing.T) {
app.taskMap[taskID1] = task1
app.setOriginatingTask(task1)

holder := newPlaceholder("ph-name", app, app.taskGroups[0])
tgs := app.getTaskGroups()
holder := newPlaceholder("ph-name", app, tgs[0])
assert.Equal(t, len(holder.pod.Spec.Containers[0].Resources.Requests), 5, "expected requests not found")
assert.Equal(t, len(holder.pod.Spec.Containers[0].Resources.Limits), 5, "expected limits not found")
assert.Equal(t, holder.pod.Spec.Containers[0].Resources.Limits[gpu], holder.pod.Spec.Containers[0].Resources.Requests[gpu], "gpu: expected same value for request and limit")
Expand All @@ -287,7 +296,8 @@ func TestNewPlaceholderWithTopologySpreadConstraints(t *testing.T) {
"bob", testGroups, map[string]string{constants.AppTagNamespace: namespace}, mockedSchedulerAPI)
app.setTaskGroups(taskGroups)

holder := newPlaceholder("ph-name", app, app.taskGroups[0])
tgs := app.getTaskGroups()
holder := newPlaceholder("ph-name", app, tgs[0])
assert.Equal(t, len(holder.pod.Spec.TopologySpreadConstraints), 1)
assert.Equal(t, holder.pod.Spec.TopologySpreadConstraints[0].MaxSkew, int32(1))
assert.Equal(t, holder.pod.Spec.TopologySpreadConstraints[0].TopologyKey, v1.LabelTopologyZone)
Expand All @@ -297,3 +307,17 @@ func TestNewPlaceholderWithTopologySpreadConstraints(t *testing.T) {
"labelKey1": "labelKeyValue1",
})
}

func TestNewPlaceholderWithDuplicatedTaskGroup(t *testing.T) {
mockedSchedulerAPI := newMockSchedulerAPI()
app := NewApplication(appID, queue,
"bob", testGroups, map[string]string{constants.AppTagNamespace: namespace}, mockedSchedulerAPI)

var duplicatedTaskGroups = make([]TaskGroup, 0, 2)
duplicatedTaskGroups = append(duplicatedTaskGroups, taskGroups[0])
duplicatedTaskGroups = append(duplicatedTaskGroups, taskGroups[0])

app.setTaskGroups(duplicatedTaskGroups)
tgs := app.getTaskGroups()
assert.Equal(t, len(tgs), 1)
}

0 comments on commit 8e12958

Please sign in to comment.