Skip to content

Commit 0622ae7

Browse files
davidLifdavidLif
committed
feat(scheduler): enforce fractional GPU compute modes
Signed-off-by: davidLif <davidshani12@gmail.com>
1 parent bf8e68b commit 0622ae7

20 files changed

Lines changed: 586 additions & 111 deletions

File tree

‎.changes/unreleased/added-20260810-112224.yaml‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,3 @@
1+
kind: Added
2+
body: |-
3+
Enforce fractional GPU compute modes

‎pkg/scheduler/actions/allocate/allocateFractionalGpu_test.go‎

Lines changed: 99 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,13 +9,20 @@ import (
99
. "go.uber.org/mock/gomock"
1010
"gopkg.in/h2non/gock.v1"
1111
v1 "k8s.io/api/core/v1"
12+
metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
13+
"k8s.io/apimachinery/pkg/types"
1214
"k8s.io/utils/ptr"
1315

1416
kaiv1common "github.com/kai-scheduler/KAI-scheduler/pkg/apis/kai/v1/common"
17+
schedulingv1alpha2 "github.com/kai-scheduler/KAI-scheduler/pkg/apis/scheduling/v1alpha2"
1518
commonconstants "github.com/kai-scheduler/KAI-scheduler/pkg/common/constants"
1619
"github.com/kai-scheduler/KAI-scheduler/pkg/scheduler/actions/allocate"
1720
"github.com/kai-scheduler/KAI-scheduler/pkg/scheduler/actions/integration_tests/integration_tests_utils"
21+
"github.com/kai-scheduler/KAI-scheduler/pkg/scheduler/api/common_info"
22+
"github.com/kai-scheduler/KAI-scheduler/pkg/scheduler/api/node_info"
23+
"github.com/kai-scheduler/KAI-scheduler/pkg/scheduler/api/pod_info"
1824
"github.com/kai-scheduler/KAI-scheduler/pkg/scheduler/api/pod_status"
25+
"github.com/kai-scheduler/KAI-scheduler/pkg/scheduler/api/resource_info"
1926
"github.com/kai-scheduler/KAI-scheduler/pkg/scheduler/conf"
2027
"github.com/kai-scheduler/KAI-scheduler/pkg/scheduler/constants"
2128
"github.com/kai-scheduler/KAI-scheduler/pkg/scheduler/test_utils"
@@ -42,6 +49,69 @@ func TestHandleFractionalGPUAllocation(t *testing.T) {
4249
}
4350
}
4451

52+
func TestFractionalGPUAllocationDoesNotUseGpuGroupWithDifferentComputeMode(t *testing.T) {
53+
test_utils.InitTestingInfrastructure()
54+
controller := NewController(t)
55+
defer controller.Finish()
56+
defer gock.Off()
57+
58+
const gpuGroup = "time-slicing-group"
59+
topology := test_utils.TestTopologyBasic{
60+
Name: "sm-sharing pod does not allocate on time-slicing gpu group",
61+
Jobs: []*jobs_fake.TestJobBasic{
62+
{
63+
Name: "running_job0",
64+
RequiredGPUsPerTask: 0.5,
65+
Priority: constants.PriorityTrainNumber,
66+
QueueName: "queue0",
67+
Tasks: []*tasks_fake.TestTaskBasic{
68+
{
69+
NodeName: "node0",
70+
GPUGroups: []string{gpuGroup},
71+
State: pod_status.Running,
72+
},
73+
},
74+
},
75+
{
76+
Name: "pending_job0",
77+
RequiredGPUsPerTask: 0.5,
78+
Priority: constants.PriorityTrainNumber,
79+
QueueName: "queue0",
80+
Tasks: []*tasks_fake.TestTaskBasic{
81+
{
82+
State: pod_status.Pending,
83+
Annotations: map[string]string{
84+
commonconstants.GpuComputeSharingMode: string(schedulingv1alpha2.GPUComputeSharingModeSMSharing),
85+
},
86+
},
87+
},
88+
},
89+
},
90+
Nodes: map[string]nodes_fake.TestNodeBasic{
91+
"node0": {
92+
GPUs: 1,
93+
},
94+
},
95+
Queues: []test_utils.TestQueueBasic{
96+
{
97+
Name: "queue0",
98+
DeservedGPUs: 1,
99+
},
100+
},
101+
}
102+
103+
ssn := test_utils.BuildSession(topology, controller)
104+
addReservationPodToNodeForTest(ssn.ClusterInfo.Nodes["node0"], gpuGroup, schedulingv1alpha2.GPUComputeSharingModeTimeSlicing)
105+
106+
allocate.New().Execute(ssn)
107+
108+
pendingTask := ssn.ClusterInfo.PodGroupInfos["pending_job0"].GetAllPodsMap()[common_info.PodID("pending_job0-0")]
109+
if pendingTask.Status != pod_status.Pending {
110+
t.Fatalf("expected sm-sharing task to stay pending, got status %s on node %s with groups %v",
111+
pendingTask.Status, pendingTask.NodeName, pendingTask.GPUGroupIDs())
112+
}
113+
}
114+
45115
func TestFractionalGPUAllocationUsesNodeConditionOverride(t *testing.T) {
46116
test_utils.InitTestingInfrastructure()
47117
controller := NewController(t)
@@ -111,6 +181,35 @@ func TestFractionalGPUAllocationUsesNodeConditionOverride(t *testing.T) {
111181
test_utils.MatchExpectedAndRealTasks(t, 0, topology, ssn)
112182
}
113183

184+
func addReservationPodToNodeForTest(
185+
node *node_info.NodeInfo, gpuGroup string, mode schedulingv1alpha2.GPUComputeSharingMode,
186+
) {
187+
pod := &v1.Pod{
188+
ObjectMeta: metav1.ObjectMeta{
189+
UID: types.UID("reservation-" + gpuGroup),
190+
Name: commonconstants.GPUReservationPodPrefix + "-node0-test",
191+
Namespace: commonconstants.DefaultResourceReservationName,
192+
Labels: map[string]string{
193+
commonconstants.GPUGroup: gpuGroup,
194+
},
195+
Annotations: map[string]string{
196+
commonconstants.PodGroupAnnotationForPod: "reservation",
197+
commonconstants.GpuComputeSharingMode: string(mode),
198+
},
199+
},
200+
Spec: v1.PodSpec{
201+
NodeName: node.Name,
202+
Containers: []v1.Container{
203+
{Name: "reservation"},
204+
},
205+
},
206+
Status: v1.PodStatus{Phase: v1.PodRunning},
207+
}
208+
task := pod_info.NewTaskInfo(pod, resource_info.NewResourceVectorMap())
209+
task.Status = pod_status.Running
210+
node.PodInfos[task.UID] = task
211+
}
212+
114213
func getFractionalGPUTestsMetadata() []integration_tests_utils.TestTopologyMetadata {
115214
return []integration_tests_utils.TestTopologyMetadata{
116215
{

‎pkg/scheduler/actions/integration_tests/integration_tests_utils/integration_tests_utils.go‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -110,7 +110,7 @@ func runSchedulerOneRound(testMetadata *TestTopologyMetadata, controller *Contro
110110
case pod_status.Releasing:
111111
if jobMetadata.DeleteJobInTest {
112112
taskMetadata.NodeName = task.NodeName
113-
taskMetadata.GPUGroups = task.GPUGroups
113+
taskMetadata.GPUGroups = task.GPUGroupIDs()
114114
taskMetadata.State = pod_status.Releasing
115115
} else {
116116
taskMetadata.NodeName = ""
@@ -124,12 +124,12 @@ func runSchedulerOneRound(testMetadata *TestTopologyMetadata, controller *Contro
124124
case pod_status.Binding:
125125
taskMetadata.State = pod_status.Running
126126
taskMetadata.NodeName = task.NodeName
127-
taskMetadata.GPUGroups = task.GPUGroups
127+
taskMetadata.GPUGroups = task.GPUGroupIDs()
128128

129129
default:
130130
taskMetadata.State = task.Status
131131
taskMetadata.NodeName = task.NodeName
132-
taskMetadata.GPUGroups = task.GPUGroups
132+
taskMetadata.GPUGroups = task.GPUGroupIDs()
133133
}
134134

135135
}

‎pkg/scheduler/actions/integration_tests/reclaim/reclaimFractional_test.go‎

Lines changed: 134 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,8 +6,20 @@ package reclaim_test
66
import (
77
"testing"
88

9+
. "go.uber.org/mock/gomock"
10+
v1 "k8s.io/api/core/v1"
11+
metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
12+
"k8s.io/apimachinery/pkg/types"
13+
14+
schedulingv1alpha2 "github.com/kai-scheduler/KAI-scheduler/pkg/apis/scheduling/v1alpha2"
15+
commonconstants "github.com/kai-scheduler/KAI-scheduler/pkg/common/constants"
916
"github.com/kai-scheduler/KAI-scheduler/pkg/scheduler/actions/integration_tests/integration_tests_utils"
17+
"github.com/kai-scheduler/KAI-scheduler/pkg/scheduler/actions/reclaim"
18+
"github.com/kai-scheduler/KAI-scheduler/pkg/scheduler/api/common_info"
19+
"github.com/kai-scheduler/KAI-scheduler/pkg/scheduler/api/node_info"
20+
"github.com/kai-scheduler/KAI-scheduler/pkg/scheduler/api/pod_info"
1021
"github.com/kai-scheduler/KAI-scheduler/pkg/scheduler/api/pod_status"
22+
"github.com/kai-scheduler/KAI-scheduler/pkg/scheduler/api/resource_info"
1123
"github.com/kai-scheduler/KAI-scheduler/pkg/scheduler/constants"
1224
"github.com/kai-scheduler/KAI-scheduler/pkg/scheduler/test_utils"
1325
"github.com/kai-scheduler/KAI-scheduler/pkg/scheduler/test_utils/jobs_fake"
@@ -19,6 +31,128 @@ func TestReclaimFractionalIntegrationTest(t *testing.T) {
1931
integration_tests_utils.RunTests(t, getReclaimFractionalTestsMetadata())
2032
}
2133

34+
func TestReclaimCanUseFullyEvictedFractionalGpuForDifferentComputeMode(t *testing.T) {
35+
test_utils.InitTestingInfrastructure()
36+
controller := NewController(t)
37+
defer controller.Finish()
38+
39+
const gpuGroup = "time-slicing-group"
40+
topology := test_utils.TestTopologyBasic{
41+
Name: "reclaim all time-slicing fractions before using GPU for sm-sharing",
42+
Jobs: []*jobs_fake.TestJobBasic{
43+
{
44+
Name: "running_job0",
45+
RequiredGPUsPerTask: 0.5,
46+
Priority: constants.PriorityTrainNumber,
47+
QueueName: "queue0",
48+
Tasks: []*tasks_fake.TestTaskBasic{
49+
{
50+
NodeName: "node0",
51+
GPUGroups: []string{gpuGroup},
52+
State: pod_status.Running,
53+
},
54+
{
55+
NodeName: "node0",
56+
GPUGroups: []string{gpuGroup},
57+
State: pod_status.Running,
58+
},
59+
},
60+
},
61+
{
62+
Name: "pending_job0",
63+
RequiredGPUsPerTask: 0.5,
64+
Priority: constants.PriorityTrainNumber,
65+
QueueName: "queue1",
66+
Tasks: []*tasks_fake.TestTaskBasic{
67+
{
68+
State: pod_status.Pending,
69+
Annotations: map[string]string{
70+
commonconstants.GpuComputeSharingMode: string(schedulingv1alpha2.GPUComputeSharingModeSMSharing),
71+
},
72+
},
73+
},
74+
},
75+
},
76+
Nodes: map[string]nodes_fake.TestNodeBasic{
77+
"node0": {
78+
GPUs: 1,
79+
},
80+
},
81+
Queues: []test_utils.TestQueueBasic{
82+
{
83+
Name: "queue0",
84+
DeservedGPUs: 0,
85+
},
86+
{
87+
Name: "queue1",
88+
DeservedGPUs: 1,
89+
},
90+
},
91+
Mocks: &test_utils.TestMock{
92+
CacheRequirements: &test_utils.CacheMocking{
93+
NumberOfCacheEvictions: 2,
94+
NumberOfPipelineActions: 1,
95+
},
96+
},
97+
}
98+
99+
ssn := test_utils.BuildSession(topology, controller)
100+
addReservationPodToNodeForReclaimTest(ssn.ClusterInfo.Nodes["node0"], gpuGroup, schedulingv1alpha2.GPUComputeSharingModeTimeSlicing)
101+
102+
reclaim.New().Execute(ssn)
103+
104+
runningTasks := ssn.ClusterInfo.PodGroupInfos["running_job0"].GetAllPodsMap()
105+
for taskID, task := range runningTasks {
106+
if task.Status != pod_status.Releasing {
107+
t.Fatalf("expected time-slicing task %s to be releasing, got %s", taskID, task.Status)
108+
}
109+
}
110+
111+
pendingTask := ssn.ClusterInfo.PodGroupInfos["pending_job0"].GetAllPodsMap()[common_info.PodID("pending_job0-0")]
112+
if pendingTask.Status != pod_status.Pipelined {
113+
t.Fatalf("expected sm-sharing task to be pipelined, got status %s", pendingTask.Status)
114+
}
115+
if pendingTask.NodeName != "node0" {
116+
t.Fatalf("expected sm-sharing task on node0, got %s", pendingTask.NodeName)
117+
}
118+
gpuGroups := pendingTask.GPUGroupIDs()
119+
if len(gpuGroups) != 1 || gpuGroups[0] == gpuGroup {
120+
t.Fatalf("expected sm-sharing task to use a new gpu group, got %v", gpuGroups)
121+
}
122+
if pendingTask.FractionalGpuGroups[0].ComputeSharingMode != schedulingv1alpha2.GPUComputeSharingModeSMSharing {
123+
t.Fatalf("expected sm-sharing mode, got %s", pendingTask.FractionalGpuGroups[0].ComputeSharingMode)
124+
}
125+
}
126+
127+
func addReservationPodToNodeForReclaimTest(
128+
node *node_info.NodeInfo, gpuGroup string, mode schedulingv1alpha2.GPUComputeSharingMode,
129+
) {
130+
pod := &v1.Pod{
131+
ObjectMeta: metav1.ObjectMeta{
132+
UID: types.UID("reservation-" + gpuGroup),
133+
Name: commonconstants.GPUReservationPodPrefix + "-node0-test",
134+
Namespace: commonconstants.DefaultResourceReservationName,
135+
Labels: map[string]string{
136+
commonconstants.GPUGroup: gpuGroup,
137+
},
138+
Annotations: map[string]string{
139+
commonconstants.PodGroupAnnotationForPod: "reservation",
140+
commonconstants.GpuComputeSharingMode: string(mode),
141+
},
142+
},
143+
Spec: v1.PodSpec{
144+
NodeName: node.Name,
145+
Containers: []v1.Container{
146+
{Name: "reservation"},
147+
},
148+
},
149+
Status: v1.PodStatus{Phase: v1.PodRunning},
150+
}
151+
task := pod_info.NewTaskInfo(pod, resource_info.NewResourceVectorMap())
152+
task.Status = pod_status.Running
153+
node.PodInfos[task.UID] = task
154+
}
155+
22156
func getReclaimFractionalTestsMetadata() []integration_tests_utils.TestTopologyMetadata {
23157
return []integration_tests_utils.TestTopologyMetadata{
24158
{

0 commit comments

Comments
 (0)