Add Parent Queue Quota Checking to Scheduler.

[shchennu]

Add Parent Queue Quota and Limit Checking to KAI Scheduler

Description

This PR enhances the KAI Scheduler with comprehensive parent queue resource management, implementing both quota and limit checking at parent queue levels. The changes improve scheduler efficiency by preventing unnecessary scheduling attempts and provide better resource management for different job types.

Changes

• Added checkParentQueueLimits method to CapacityPolicy to check both limits and quotas at parent queue levels
• Added getFirstPendingPod helper function to identify the first pending pod in a job
• Modified OnSessionOpen to include early validation of resource limits
• Added comprehensive test cases for parent queue limit checking
• Added support for elastic jobs with minimum required resources
• Separated job priority from preemptibility handling

Implementation Details

• Resource checks are performed for GPU, CPU, and Memory resources
• Both limits and quotas are checked, using the more restrictive value
• Checks are done at each parent queue level in the hierarchy
• Early validation prevents unnecessary scheduling attempts
• Proper error messages indicate which resource limit/quota was exceeded
• Special handling for preemptible jobs (PriorityInferenceNumber)
• Support for elastic jobs with minimum resource requirements

Test Coverage

• Basic limit/quota checks for each resource type (GPU, CPU, Memory)
• Multi-level queue hierarchy checks
• Elastic job handling with minimum required resources
• Preemptible vs non-preemptible job behavior
• Edge cases (zero limits/quotas, missing values)
• Error message formatting
• Helper function validation

Files Changed

• Modified: pkg/scheduler/plugins/proportion/capacity_policy/capacity_policy.go
• Modified: pkg/scheduler/plugins/proportion/capacity_policy/parent_queue_test.go

Testing

All tests are passing:

• 37 specs from the original test suite
• 10 new test cases for parent queue limit checking
• 3 test cases for the helper function

Impact

• Improves scheduler efficiency by preventing unnecessary scheduling attempts
• Provides clear error messages when limits/quotas are exceeded
• Better resource management for different job types (preemptible, elastic)

[romanbaron]

Hi @shchennu 👋
Thanks for suggesting this change, I want to better understand the motivation here, more specifically what is the gap with IsJobOverQueueCapacity? From what I see it is called before PrePredicateFn so there is probably something that you identified as a gap in the implementation?

Another question, what is the problem with job allocation above parent's queue quota? If the job is preemptible it can be allowed.

[shchennu]

Hi @shchennu 👋 Thanks for suggesting this change, I want to better understand the motivation here, more specifically what is the gap with IsJobOverQueueCapacity? From what I see it is called before PrePredicateFn so there is probably something that you identified as a gap in the implementation?

Another question, what is the problem with job allocation above parent's queue quota? If the job is preemptible it can be allowed.

Hi @romanbaron 👋

Thank you for your questions! Let me address them:

1. Gap with IsJobOverQueueCapacity:

   • IsJobOverQueueCapacity only checks the immediate queue's capacity
   • Our new checkParentQueueQuotas function checks ALL parent queues up the hierarchy
   • This is important because a job might fit in its immediate queue but exceed parent queue quotas
   • By checking in PrePredicateFn, we fail fast and avoid unnecessary scheduling attempts

2. Preemptible Jobs and Parent Queue Quotas:

   • You're absolutely right! Preemptible jobs should be allowed to exceed parent queue quotas
   • I've updated the implementation to handle this:
     • Added a check for PriorityTrainNumber at the start of checkParentQueueQuotas
     • Preemptible jobs now skip quota checks entirely
     • Non-preemptible jobs still maintain strict quota enforcement
   • This change is reflected in the test cases:
     • preemptible_job_can_exceed_parent_queue_GPU_quota
     • non-preemptible_job_cannot_exceed_parent_queue_GPU_quota

The changes ensure that:

• Preemptible jobs can utilize resources beyond parent queue quotas
• Non-preemptible jobs maintain strict quota enforcement
• We still get early validation for jobs that would exceed quotas

[enoodle]

Hi,
I left a few comments, mostly about style.

I do have a question about the logic: how will this handle elastic jobs who doesn't have to schedule all pods to run? The job has a minimal number of pods that it needs to run, and you don't know how many will actually fit in the cluster / queue limits.

[enoodle]

Maybe I am missing something, but I don't think this is called by the changes in the current PR because this is not a plugin that is registered.

[enoodle]

Its a style comment, because this is a new project and all, but we don't have many obvious comments here (the comment really doesn't add anything that is not written in the next line)

[enoodle]

This together with line 239 can be one line of for with initialization, check and step

[enoodle]

this part seems duplicated 3 times, there are functions to avoid it in the code base. you can loop on the resource type and the . look at proportion/proportion.go for example.

[enoodle]

A queue can go over the quota, but not the limit. I would change the whole thing to address the limit.

[enoodle]

We aim to separate the job priority and preemptibility soon, and the Train priority is not the only one that is currently preemptible.
We do have several scalars to define the resource of a queue: quota, limit and over quota weight. We should look at limit here and not quota and then it will be correct for all jobs.

[enoodle]

Isn't the fact that the job needs scheduling (and then predicates are running on it) is not enough to understand that it has pending pods?

[shchennu]

Hi, I left a few comments, mostly about style.

I do have a question about the logic: how will this handle elastic jobs who doesn't have to schedule all pods to run? The job has a minimal number of pods that it needs to run, and you don't know how many will actually fit in the cluster / queue limits.

Thanks for the review! I've addressed all the comments:

• Added support for elastic jobs with minimum resource requirements
• Improved code style by removing redundant comments and combining loops
• Now checking both limits and quotas, using the more restrictive value
• Separated job priority from preemptibility
• Added test coverage

All tests are passing. Let me know if you'd like any further improvements!

[enoodle]

1. Gap with IsJobOverQueueCapacity:

   • IsJobOverQueueCapacity only checks the immediate queue's capacity
   • Our new checkParentQueueQuotas function checks ALL parent queues up the hierarchy
   • This is important because a job might fit in its immediate queue but exceed parent queue quotas
   • By checking in PrePredicateFn, we fail fast and avoid unnecessary scheduling attempts

@shchennu I looked at IsJobOverQueueCapacity and I don't agree with your observation - it is checking all the parent queues, so this addition is not needed.
https://github.com/NVIDIA/KAI-Scheduler/blob/main/pkg/scheduler/plugins/proportion/capacity_policy/quota_check.go#L34

[romanbaron]

@shchennu - all hierarchy levels are validated in these two places:
resultsOverLimit - checks that workload scheduling won't cause any of the queues (from leaf to top) to go over limit.
resultsWithNonPreemptibleOverQuota - checks that non-preemptible workload scheduling won't cause any of the queues (from leaf to top) to go over quota.

IsJobOverQueueCapacity calls these two functions here:

KAI-Scheduler/pkg/scheduler/plugins/proportion/capacity_policy/capacity_policy.go

Line 35 in 2a0bacd

checkFns := []capacityCheckFn{cp.resultsOverLimit, cp.resultsWithNonPreemptibleOverQuota}

Doesn't it cover it?

[romanbaron]

@shchennu - can you please update on your thoughts regarding the questions above?