@@ -30,6 +30,7 @@ const (
3030 NodePartitionTypeStateKey = "NodePartitionType"
3131 PodLauncherStateKey = "PodLauncher"
3232 PodResourceTypeStateKey = "PodResourceType"
33+ PodSchedulingCtxKey = "PodSchedulingCtx"
3334 PodTraceStateKey = "PodTrace"
3435 NodeGroupStateKey = "NodeGroup"
3536 PotentialVictimsKey = "PotentialVictims"
@@ -45,6 +46,7 @@ const (
4546 NodePartitionTypeMissedErrorString = "failed to get NodePartitionType, supposed to be set in cycle state"
4647 PodLauncherMissedErrorString = "failed to get PodLauncher, supposed to be set in cycle state"
4748 PodResourceTypeMissingErrorString = "failed to get PodResourceType, supposed to be set in cycle state"
49+ PodSchedulingCtxMissingErrorString = "failed to get PodSchedulingCtx, supposed to be set in cycle state"
4850 PodTraceMissingErrorString = "failed to get PodTrace, supposed to be set in cycle state"
4951 NodeGroupMissedErrorString = "failed to get NodeGroup, supposed to be set in cycle state"
5052
@@ -91,6 +93,17 @@ func SetPodLauncherState(podLauncher podutil.PodLauncher, state *CycleState) err
9193 return podutil .PodLauncherUnsupportError
9294}
9395
96+ func SetPodSchedulingCtxKey (schedulingCtx * PodSchedulingCtx , state * CycleState ) error {
97+ if schedulingCtx == nil {
98+ schedulingCtx = & PodSchedulingCtx {}
99+ }
100+ data := & stateData {
101+ data : schedulingCtx ,
102+ }
103+ state .Write (PodSchedulingCtxKey , data )
104+ return nil
105+ }
106+
94107func SetPodTrace (podTrace tracing.SchedulingTrace , state * CycleState ) error {
95108 data := & stateData {
96109 data : podTrace ,
@@ -242,6 +255,17 @@ func GetPodLauncher(state *CycleState) (podutil.PodLauncher, error) {
242255 return "" , PodLauncherMissedError
243256}
244257
258+ var PodSchedulingCtxMissingError = fmt .Errorf (PodSchedulingCtxMissingErrorString )
259+
260+ func GetPodSchedulingCtx (state * CycleState ) (* PodSchedulingCtx , error ) {
261+ if data , err := state .Read (PodSchedulingCtxKey ); err == nil {
262+ if s , ok := data .(* stateData ); ok {
263+ return s .data .(* PodSchedulingCtx ), nil
264+ }
265+ }
266+ return nil , PodSchedulingCtxMissingError
267+ }
268+
245269var PodTraceMissingError = fmt .Errorf (PodTraceMissingErrorString )
246270
247271func GetPodTrace (state * CycleState ) (tracing.SchedulingTrace , error ) {
0 commit comments