[YUNIKORN-2982] Send event when preemption occurs

rhh777 · Nov 25, 2024 · 32824bf · 32824bf
1 parent ebb8b1c
commit 32824bf
Show file tree

Hide file tree

Showing 3 changed files with 30 additions and 0 deletions.
diff --git a/pkg/scheduler/objects/events/application_events.go b/pkg/scheduler/objects/events/application_events.go
@@ -79,6 +79,14 @@ func (ae *ApplicationEvents) SendRemoveAllocationEvent(appID, allocKey string, a
 	ae.eventSystem.AddEvent(event)
 }
 
+func (ae *ApplicationEvents) SendPreemptAllocationEvent(appID, allocKey string, allocated *resources.Resource, message string) {
+	if !ae.eventSystem.IsEventTrackingEnabled() {
+		return
+	}
+	event := events.CreateAppEventRecord(appID, message, allocKey, si.EventRecord_REMOVE, si.EventRecord_ALLOC_PREEMPT, allocated)
+	ae.eventSystem.AddEvent(event)
+}
+
 func (ae *ApplicationEvents) SendRemoveAskEvent(appID, allocKey string, allocated *resources.Resource, detail si.EventRecord_ChangeDetail) {
 	if !ae.eventSystem.IsEventTrackingEnabled() {
 		return

diff --git a/pkg/scheduler/objects/events/application_events_test.go b/pkg/scheduler/objects/events/application_events_test.go
@@ -168,6 +168,24 @@ func TestSendRemoveAllocationEvent(t *testing.T) {
 	}
 }
 
+func TestSendPreemptAllocationEvent(t *testing.T) {
+	eventSystem := mock.NewEventSystemDisabled()
+	appEvents := NewApplicationEvents(eventSystem)
+	appEvents.SendPreemptAllocationEvent(appID, allocKey, resources.NewResource(), "Preempted by pod-a from application yunikorn-test-autogen in root.test")
+	assert.Equal(t, 0, len(eventSystem.Events), "unexpected event")
+
+	eventSystem = mock.NewEventSystem()
+	appEvents = NewApplicationEvents(eventSystem)
+	appEvents.SendPreemptAllocationEvent(appID, allocKey, resources.NewResource(), "Preempted by pod-a from application yunikorn-test-autogen in root.test")
+	event := eventSystem.Events[0]
+	assert.Equal(t, si.EventRecord_APP, event.Type)
+	assert.Equal(t, si.EventRecord_REMOVE, event.EventChangeType)
+	assert.Equal(t, si.EventRecord_ALLOC_PREEMPT, event.EventChangeDetail)
+	assert.Equal(t, "app-0", event.ObjectID)
+	assert.Equal(t, "alloc-0", event.ReferenceID)
+	assert.Equal(t, "Preempted by pod-a from application yunikorn-test-autogen in root.test", event.Message)
+}
+
 func TestSendRemoveAskEvent(t *testing.T) {
 	eventSystem := mock.NewEventSystemDisabled()
 	appEvents := NewApplicationEvents(eventSystem)

diff --git a/pkg/scheduler/objects/preemption.go b/pkg/scheduler/objects/preemption.go
@@ -19,6 +19,7 @@
 package objects
 
 import (
+	"fmt"
 	"sort"
 	"strings"
 	"sync"
@@ -635,6 +636,9 @@ func (p *Preemptor) TryPreemption() (*AllocationResult, bool) {
 				zap.String("victimNodeID", victim.GetNodeID()),
 				zap.String("victimQueue", victimQueue.Name),
 			)
+			// send event
+			message := fmt.Sprintf("Preempted by %v from application %v in %v", p.ask.GetTag("kubernetes.io/meta/podName"), p.ask.applicationID, p.application.queuePath)
+			p.application.appEvents.SendPreemptAllocationEvent(victim.GetApplicationID(), victim.GetAllocationKey(), victim.GetAllocatedResource(), message)
 		} else {
 			log.Log(log.SchedPreemption).Warn("BUG: Queue not found for preemption victim",
 				zap.String("queue", p.queue.Name),